Coverage for eminus/io/traj.py: 100.00%

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""TRAJ file handling.""" 

4 

5import numpy as np 

6 

7from eminus.logger import log 

8from eminus.units import ang2bohr 

9 

10from .xyz import write_xyz 

11 

12 

13def read_traj(filename): 

14 """Load atom species and positions from TRAJ files. 

15 

16 TRAJ files are just multiple XYZ files appended to one file. 

17 See :func:`~eminus.io.xyz.read_xyz` for more information about the XYZ file format. 

18 

19 Args: 

20 filename: TRAJ input file path/name. 

21 

22 Returns: 

23 Atom species and positions. 

24 """ 

25 if not filename.endswith((".trj", ".traj")): 

26 filename += ".traj" 

27 

28 with open(filename, encoding="utf-8") as fh: 

29 lines = fh.readlines() 

30 Nlines = len(lines) 

31 

32 # The first line contains the number of atoms 

33 Natoms = int(lines[0].strip()) 

34 

35 # The second line can contain a comment, print it if available 

36 comment = lines[1].strip() 

37 if comment: 

38 log.info(f'TRAJ file comment: "{comment}"') 

39 

40 traj = [] 

41 for frame in range(Nlines // (2 + Natoms)): 

42 atom = [] 

43 pos = [] 

44 # Following lines contain atom positions with the format: Atom x-pos y-pos z-pos 

45 for line in lines[(2 + Natoms) * frame + 2 : (2 + Natoms) * (frame + 1)]: 

46 line_split = line.strip().split() 

47 atom.append(line_split[0]) 

48 pos.append(np.float64(line_split[1:4])) 

49 # XYZ files are in Angstrom, so convert to Bohr 

50 pos = ang2bohr(np.asarray(pos)) 

51 traj.append((atom, pos)) 

52 return traj 

53 

54 

55def write_traj(obj, filename, fods=None, elec_symbols=("X", "He")): 

56 """Generate TRAJ files from atoms objects. 

57 

58 TRAJ files are just multiple XYZ files appended to one file. 

59 See :func:`~eminus.io.xyz.write_xyz` for more information about the XYZ file format. 

60 

61 Args: 

62 obj: Atoms or SCF object or list/tuple of these objects. 

63 filename: TRAJ output file path/name. 

64 

65 Keyword Args: 

66 fods: FOD coordinates to write. 

67 elec_symbols: Identifier for up and down FODs. 

68 """ 

69 if not filename.endswith((".trj", ".traj")): 

70 filename += ".traj" 

71 

72 if isinstance(obj, (list, tuple)): 

73 for iobj in obj: 

74 write_xyz(iobj, filename, fods=fods, elec_symbols=elec_symbols, trajectory=True) 

75 else: 

76 write_xyz(obj, filename, fods=fods, elec_symbols=elec_symbols, trajectory=True)