Coverage for eminus/io/xyz.py: 97.73%

1# SPDX-FileCopyrightText: 2023 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""XYZ file handling.""" 

4 

5import time 

6 

7import numpy as np 

8 

9from eminus import backend as xp 

10from eminus.logger import log 

11from eminus.units import ang2bohr, bohr2ang 

12from eminus.version import __version__ 

13 

14 

15def read_xyz(filename): 

16 """Load atom species and positions from XYZ files. 

17 

18 File format definition: 

19 https://open-babel.readthedocs.io/en/latest/FileFormats/XYZ_cartesian_coordinates_format.html 

20 

21 Args: 

22 filename: XYZ input file path/name. 

23 

24 Returns: 

25 Atom species and positions. 

26 """ 

27 if not filename.endswith(".xyz"): 

28 filename += ".xyz" 

29 

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

31 lines = fh.readlines() 

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'XYZ file comment: "{comment}"') 

40 

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 : 2 + Natoms]: 

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

46 atom.append(line_split[0]) 

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

48 

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

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

51 return atom, pos 

52 

53 

54def write_xyz(obj, filename, fods=None, elec_symbols=("X", "He"), trajectory=False): 

55 """Generate XYZ files from atoms objects. 

56 

57 File format definition: 

58 https://open-babel.readthedocs.io/en/latest/FileFormats/XYZ_cartesian_coordinates_format.html 

59 

60 Args: 

61 obj: Atoms or SCF object. 

62 filename: XYZ output file path/name. 

63 

64 Keyword Args: 

65 fods: FOD coordinates to write. 

66 elec_symbols: Identifier for up and down FODs. 

67 trajectory: Allow appending to a file to create trajectories. 

68 """ 

69 atoms = obj._atoms 

70 

71 # The trajectory write calls this function 

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

73 filename += ".xyz" 

74 

75 # Convert the coordinates from atomic units to Angstrom 

76 pos = bohr2ang(atoms.pos) 

77 if fods is not None: 

78 fods = [bohr2ang(i) for i in fods] 

79 

80 if "He" in atoms.atom and atoms.unrestricted: 

81 log.warning( 

82 'You need to modify "elec_symbols" to write helium with FODs in the spin-' 

83 "polarized case." 

84 ) 

85 

86 # Append to a file when using the trajectory keyword 

87 if trajectory: 

88 mode = "a" 

89 else: 

90 mode = "w" 

91 

92 with open(filename, mode, encoding="utf-8") as fp: 

93 # The first line contains the number of atoms 

94 # If we add FOD coordinates, add them to the count 

95 if fods is None: 

96 fp.write(f"{atoms.Natoms}\n") 

97 else: 

98 fp.write(f"{atoms.Natoms + sum(len(i) for i in fods)}\n") 

99 # The second line can contain a comment 

100 # Print information about the file and program, and the file creation time 

101 fp.write(f"File generated with eminus {__version__} on {time.ctime()}\n") 

102 fp.writelines( 

103 f"{atoms.atom[ia]:<2s} {pos[ia, 0]: .6f} {pos[ia, 1]: .6f} {pos[ia, 2]: .6f}\n" 

104 for ia in range(atoms.Natoms) 

105 ) 

106 # Add FOD coordinates if needed 

107 # The atom symbol will default to pos (no atom type) 

108 if fods is not None: 

109 for s in range(len(fods)): 

110 fp.writelines( 

111 f"{elec_symbols[s]:<2s} {ie[0]: .6f} {ie[1]: .6f} {ie[2]: .6f}\n" 

112 for ie in fods[s] 

113 )