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

72 statements  

« prev     ^ index     » next       coverage.py v7.6.4, created at 2024-11-08 12:59 +0000

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

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

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

16 

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

18 

19 Args: 

20 filename: POSCAR input file path/name. 

21 

22 Returns: 

23 Atom species, positions, and cell vectors. 

24 """ 

25 if "POSCAR" not in filename: 

26 filename += ".POSCAR" 

27 

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

29 lines = fh.readlines() 

30 

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

32 comment = lines[0].strip() 

33 if comment: 

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

35 

36 # The second line contains scaling factors 

37 scaling = np.float64(lines[1].strip().split()) 

38 

39 # Followed by unscaled lattice coordinates, scale them here 

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

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

42 a[i] = scaling * np.float64(line.strip().split()) 

43 

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

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

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

47 raise NotImplementedError(msg) 

48 # Otherwise the atom species are given with their amount 

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

50 Natom = np.int64(lines[6].strip().split()) 

51 # Extend the atoms by their respective amount 

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

53 

54 # Ignore the dynamics line if available 

55 skip = 0 

56 if "dynamics" in lines[7]: 

57 skip += 1 

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

59 

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

61 # Following lines contain atom positions 

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

63 if mode == "direct": 

64 pos[i] = np.sum(a * np.float64(line.strip().split()[:3]), axis=0) 

65 if mode == "cartesian": 

66 pos[i] = scaling * np.float64(line.strip().split()[:3]) 

67 # Skip all the properties afterwards 

68 

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

70 pos = ang2bohr(pos) 

71 a = ang2bohr(a) 

72 return atom, pos, a 

73 

74 

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

76 """Generate POSCAR files from atoms objects. 

77 

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

79 

80 Args: 

81 obj: Atoms or SCF object. 

82 filename: POSCAR output file path/name. 

83 

84 Keyword Args: 

85 fods: FOD coordinates to write. 

86 elec_symbols: Identifier for up and down FODs. 

87 """ 

88 atoms = obj._atoms 

89 

90 if "POSCAR" not in filename: 

91 filename += ".POSCAR" 

92 

93 # Convert the coordinates from atomic units to Angstrom 

94 pos = bohr2ang(atoms.pos) 

95 a = bohr2ang(atoms.a) 

96 if fods is not None: 

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

98 

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

100 log.warning( 

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

102 "polarized case." 

103 ) 

104 

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

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

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

108 

109 # We have scaled vectors and coordinates 

110 fp.write("1.0\n") 

111 

112 # Write lattice 

113 fp.write( 

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

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

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

117 ) 

118 

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

120 sort = np.argsort(atoms.atom) 

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

122 pos = pos[sort] 

123 

124 # Write the sorted species 

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

126 if fods is not None: 

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

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

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

130 fp.write("\n") 

131 

132 # Write the number per (sorted) species 

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

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

135 if fods is not None: 

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

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

138 fp.write(f" {len(fods[s])}") 

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

140 

141 # Write the coordinates 

142 for ia in range(atoms.Natoms): 

143 fp.write(f"{pos[ia, 0]: .6f} {pos[ia, 1]: .6f} {pos[ia, 2]: .6f}\n") 

144 

145 # Add FOD coordinates if desired 

146 if fods is not None: 

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

148 for ie in fods[s]: 

149 fp.write(f"{ie[0]: .6f} {ie[1]: .6f} {ie[2]: .6f}\n")