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

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""XYZ file handling.""" 

4 

5import time 

6 

7import numpy as np 

8 

9from ..logger import log 

10from ..units import ang2bohr, bohr2ang 

11from ..version import __version__ 

12 

13 

14def read_xyz(filename): 

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

16 

17 File format definition: 

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

19 

20 Args: 

21 filename: XYZ input file path/name. 

22 

23 Returns: 

24 Atom species and positions. 

25 """ 

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

27 filename += ".xyz" 

28 

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

30 lines = fh.readlines() 

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

39 

40 atom = [] 

41 pos = [] 

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

43 for line in lines[2 : 2 + Natoms]: 

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

45 atom.append(line_split[0]) 

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

47 

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

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

50 return atom, pos 

51 

52 

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

54 """Generate XYZ files from atoms objects. 

55 

56 File format definition: 

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

58 

59 Args: 

60 obj: Atoms or SCF object. 

61 filename: XYZ output file path/name. 

62 

63 Keyword Args: 

64 fods: FOD coordinates to write. 

65 elec_symbols: Identifier for up and down FODs. 

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

67 """ 

68 atoms = obj._atoms 

69 

70 # The trajectory write calls this function 

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

72 filename += ".xyz" 

73 

74 # Convert the coordinates from atomic units to Angstrom 

75 pos = bohr2ang(atoms.pos) 

76 if fods is not None: 

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

78 

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

80 log.warning( 

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

82 "polarized case." 

83 ) 

84 

85 # Append to a file when using the trajectory keyword 

86 if trajectory: 

87 mode = "a" 

88 else: 

89 mode = "w" 

90 

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

92 # The first line contains the number of atoms 

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

94 if fods is None: 

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

96 else: 

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

98 # The second line can contain a comment 

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

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

101 for ia in range(atoms.Natoms): 

102 fp.write( 

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

104 ) 

105 # Add FOD coordinates if desired 

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

107 if fods is not None: 

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

109 for ie in fods[s]: 

110 fp.write(f"{elec_symbols[s]:<2s} {ie[0]: .6f} {ie[1]: .6f} {ie[2]: .6f}\n")