Coverage for eminus/io/poscar.py: 95.77%

1# SPDX-FileCopyrightText: 2024 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""POSCAR 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_poscar(filename): 

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

17 

18 File format definition: https://www.vasp.at/wiki/index.php/POSCAR 

19 

20 Args: 

21 filename: POSCAR input file path/name. 

22 

23 Returns: 

24 Atom species, positions, and cell vectors. 

25 """ 

26 if "POSCAR" not in filename: 

27 filename += ".POSCAR" 

28 

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

30 lines = fh.readlines() 

31 

32 # The first line contains a comment, print it if available 

33 comment = lines[0].strip() 

34 if comment: 

35 log.info(f'POSCAR file comment: "{comment}"') 

36 

37 # The second line contains scaling factors 

38 scaling = np.asarray(lines[1].strip().split(), dtype=float) 

39 

40 # Followed by unscaled lattice coordinates, scale them here 

41 a = np.empty((3, 3)) 

42 for i, line in enumerate(lines[2:5]): 

43 a[i] = scaling * np.asarray(line.strip().split(), dtype=float) 

44 

45 # If line number six contains numbers a POTCAR file is required 

46 if lines[5].strip().split()[0].isnumeric(): 

47 msg = "POTCAR files are not supported, provide species in the POSCAR file" 

48 raise NotImplementedError(msg) 

49 # Otherwise the atom species are given with their amount 

50 atom = lines[5].strip().split() 

51 Natom = np.asarray(lines[6].strip().split(), dtype=int) 

52 # Extend the atoms by their respective amount 

53 atom = [a for a, N in zip(atom, Natom) for _ in range(N)] 

54 

55 # Ignore the dynamics line if available 

56 skip = 0 

57 if "dynamics" in lines[7]: 

58 skip += 1 

59 mode = lines[7 + skip].strip().lower() 

60 

61 pos = np.empty((np.sum(Natom), 3)) 

62 # Following lines contain atom positions 

63 for i, line in enumerate(lines[8 + skip : 8 + skip + np.sum(Natom)]): 

64 if mode == "direct": 

65 pos[i] = np.sum(a * np.asarray(line.strip().split()[:3], dtype=float), axis=0) 

66 if mode == "cartesian": 

67 pos[i] = scaling * np.asarray(line.strip().split()[:3], dtype=float) 

68 # Skip all the properties afterwards 

69 

70 # POSCAR files are in Angstrom, so convert to Bohr 

71 pos = xp.asarray(ang2bohr(pos)) 

72 a = xp.asarray(ang2bohr(a)) 

73 return atom, pos, a 

74 

75 

76def write_poscar(obj, filename, fods=None, elec_symbols=("X", "He")): 

77 """Generate POSCAR files from atoms objects. 

78 

79 File format definition: https://www.vasp.at/wiki/index.php/POSCAR 

80 

81 Args: 

82 obj: Atoms or SCF object. 

83 filename: POSCAR output file path/name. 

84 

85 Keyword Args: 

86 fods: FOD coordinates to write. 

87 elec_symbols: Identifier for up and down FODs. 

88 """ 

89 atoms = obj._atoms 

90 

91 if "POSCAR" not in filename: 

92 filename += ".POSCAR" 

93 

94 # Convert the coordinates from atomic units to Angstrom 

95 pos = bohr2ang(atoms.pos) 

96 a = bohr2ang(atoms.a) 

97 if fods is not None: 

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

99 

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

101 log.warning( 

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

103 "polarized case." 

104 ) 

105 

106 with open(filename, "w", encoding="utf-8") as fp: 

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

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

109 

110 # We have scaled vectors and coordinates 

111 fp.write("1.0\n") 

112 

113 # Write lattice 

114 fp.write( 

115 f"{a[0, 0]:.6f} {a[0, 1]:.6f} {a[0, 2]:.6f}\n" 

116 f"{a[1, 0]:.6f} {a[1, 1]:.6f} {a[1, 2]:.6f}\n" 

117 f"{a[2, 0]:.6f} {a[2, 1]:.6f} {a[2, 2]:.6f}\n" 

118 ) 

119 

120 # We need to sort the atoms in the POSCAR file 

121 sort = np.argsort(atoms.atom) 

122 atom = np.asarray(atoms.atom)[sort] 

123 pos = pos[sort] 

124 

125 # Write the sorted species 

126 fp.write(f"{' '.join(set(atom))}") 

127 if fods is not None: 

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

129 if len(fods[s]) > 0: 

130 fp.write(f" {elec_symbols[s]}") 

131 fp.write("\n") 

132 

133 # Write the number per (sorted) species 

134 _, counts = np.unique(atom, return_counts=True) 

135 fp.write(f"{' '.join(map(str, counts))}") 

136 if fods is not None: 

137 for fod in fods: 

138 if len(fod) > 0: 

139 fp.write(f" {len(fod)}") 

140 fp.write("\nCartesian\n") 

141 

142 # Write the coordinates 

143 fp.writelines( 

144 f"{pos[ia, 0]: .6f} {pos[ia, 1]: .6f} {pos[ia, 2]: .6f}\n" 

145 for ia in range(atoms.Natoms) 

146 ) 

147 

148 # Add FOD coordinates if needed 

149 if fods is not None: 

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

151 fp.writelines(f"{ie[0]: .6f} {ie[1]: .6f} {ie[2]: .6f}\n" for ie in fods[s])