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 import backend as xp 

8from eminus.logger import log 

9from eminus.units import ang2bohr 

10 

11from .xyz import write_xyz 

12 

13 

14def read_traj(filename): 

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

16 

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

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

19 

20 Args: 

21 filename: TRAJ input file path/name. 

22 

23 Returns: 

24 Atom species and positions. 

25 """ 

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

27 filename += ".traj" 

28 

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

30 lines = fh.readlines() 

31 Nlines = len(lines) 

32 

33 # The first line contains the number of atoms 

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

35 

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

37 comment = lines[1].strip() 

38 if comment: 

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

40 

41 traj = [] 

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

43 atom = [] 

44 pos = [] 

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

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

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

48 atom.append(line_split[0]) 

49 pos.append(np.asarray(line_split[1:4], dtype=float)) 

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

51 pos = xp.asarray(ang2bohr(np.asarray(pos))) 

52 traj.append((atom, pos)) 

53 return traj 

54 

55 

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

57 """Generate TRAJ files from atoms objects. 

58 

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

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

61 

62 Args: 

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

64 filename: TRAJ output file path/name. 

65 

66 Keyword Args: 

67 fods: FOD coordinates to write. 

68 elec_symbols: Identifier for up and down FODs. 

69 """ 

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

71 filename += ".traj" 

72 

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

74 for iobj in obj: 

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

76 else: 

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