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

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""TRAJ file handling.""" 

4 

5import numpy as np 

6 

7from ..logger import log 

8from ..units import ang2bohr 

9from .xyz import write_xyz 

10 

11 

12def read_traj(filename): 

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

14 

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

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

17 

18 Args: 

19 filename: TRAJ input file path/name. 

20 

21 Returns: 

22 Atom species and positions. 

23 """ 

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

25 filename += ".traj" 

26 

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

28 lines = fh.readlines() 

29 Nlines = len(lines) 

30 

31 # The first line contains the number of atoms 

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

33 

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

35 comment = lines[1].strip() 

36 if comment: 

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

38 

39 traj = [] 

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

41 atom = [] 

42 pos = [] 

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

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

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

46 atom.append(line_split[0]) 

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

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

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

50 traj.append((atom, pos)) 

51 return traj 

52 

53 

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

55 """Generate TRAJ files from atoms objects. 

56 

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

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

59 

60 Args: 

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

62 filename: TRAJ output file path/name. 

63 

64 Keyword Args: 

65 fods: FOD coordinates to write. 

66 elec_symbols: Identifier for up and down FODs. 

67 """ 

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

69 filename += ".traj" 

70 

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

72 for iobj in obj: 

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

74 else: 

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