Coverage for sparc/sparc_parsers/static.py: 90%

1""" 

2Created on Thu Oct 18 14:16:21 2018 

3 

4Ben Comer (Georgia Tech) 

5 

6This file has been heavily modified since SPARC 0.1 

7 

8TODO: more descriptions about this file io parser 

9""" 

10from warnings import warn 

11 

12import numpy as np 

13from ase.units import Bohr, GPa, Hartree 

14 

15# Safe wrappers for both string and fd 

16from ase.utils import reader, writer 

17 

18from ..api import SparcAPI 

19from .utils import strip_comments 

20 

21# TODO: should allow user to select the api 

22# defaultAPI = SparcAPI() 

23 

24 

25@reader 

26def _read_static(fileobj): 

27 """ 

28 Read the .static file content 

29 

30 Each .static file should only host 1 or more images 

31 (if socket mode is enabled), but the output may vary 

32 a lot depending on the flags (e.g. PRINT_ATOMS, PRINT_FORCES etc) 

33 """ 

34 contents = fileobj.read() 

35 data, comments = strip_comments(contents) 

36 # Most static files should containe the Atom positions lines 

37 # this can happen for both a single- or multi-image file 

38 step_bounds = [i for i, x in enumerate(data) if "Atom positions" in x] + [len(data)] 

39 raw_static_steps = [ 

40 data[start:end] for start, end in zip(step_bounds[:-1], step_bounds[1:]) 

41 ] 

42 # In some cases (e.g. PRINT_ATOMS=0), the static file may not contain Atom positions 

43 # All existing lines will be regarded as one step 

44 if len(raw_static_steps) == 0: 

45 raw_static_steps = [data] 

46 static_steps = [_read_static_step(step) for step in raw_static_steps] 

47 return {"static": static_steps} 

48 

49 

50def _read_static_block(raw_block): 

51 """Parse ONE static data block, this will return dict with keys: 

52 {"name": PARAM, "value": value} 

53 

54 Arguments: 

55 raw_block: un-parsed block as list of strings 

56 """ 

57 header, body = raw_block[0], raw_block[1:] 

58 header_name, header_rest = header.split(":") 

59 if len(header_rest.strip()) > 0: 

60 body = [header_rest.strip()] + body 

61 

62 try: 

63 value = np.genfromtxt(body, dtype=float) 

64 if np.isnan(value).any(): 

65 warn( 

66 ( 

67 f"Field contains data that are not parsable by numpy!\n" 

68 f"Contents are: {body}" 

69 ) 

70 ) 

71 except Exception: 

72 value = body 

73 

74 # name = None 

75 if "Total free energy" in header_name: 

76 name = "free energy" 

77 elif "Atomic forces" in header_name: 

78 name = "forces" 

79 elif "Net magnetization" in header_name: 

80 name = "net_magnetization" 

81 elif "Atomic magnetization" in header_name: 

82 name = "atomic_magnetization" 

83 elif "Stress (GPa)" in header_name: 

84 name = "stress" 

85 elif "Stress equiv." in header_name: 

86 name = "stress_equiv" 

87 elif "Stress (Ha/Bohr)" in header_name: 

88 name = "stress_1d" 

89 elif "Stress (Ha/Bohr**2)" in header_name: 

90 name = "stress_2d" 

91 elif "Fractional coordinates" in header_name: 

92 # Fractional coordinates of Si -- > name=coord_frac symbol="Si" 

93 name = "coord_frac" 

94 symbol = header_name.split("of")[1].strip() 

95 clean_array = value.reshape((-1, 3)) 

96 value = {"value": clean_array, "symbol": symbol} 

97 # Exclusive to the socket mode 

98 elif "Lattice (Bohr)" in header_name: 

99 name = "lattice" 

100 value = value.reshape((3, 3)) 

101 else: 

102 name = header_name.strip() 

103 

104 return {"name": name, "value": value} 

105 

106 

107def _read_static_step(step): 

108 """Parse all the lines in one step and compose a dict containing sanitized blocks 

109 Args: 

110 step (list): Lines of raw lines in one step 

111 """ 

112 separator = "*" * 60 # Make the separator long enough 

113 # Clean up boundary lines 

114 data = [ 

115 line 

116 for line in step 

117 if ("Atom positions" not in line) and (separator not in line) 

118 ] 

119 block_bounds = [i for i, x in enumerate(data) if ":" in x] + [len(data)] 

120 raw_blocks = [ 

121 data[start:end] for start, end in zip(block_bounds[:-1], block_bounds[1:]) 

122 ] 

123 # import pdb; pdb.set_trace() 

124 block_dict = {} 

125 coord_dict = {} 

126 block_contents = [_read_static_block(block) for block in raw_blocks] 

127 for bc in block_contents: 

128 name, raw_value = bc["name"], bc["value"] 

129 # Coord frac needs to be collected in all positions 

130 if name == "coord_frac": 

131 value = None 

132 symbol, coord_f = raw_value["symbol"], raw_value["value"] 

133 pos_count = len(coord_f) 

134 if coord_dict == {}: 

135 coord_dict.update( 

136 { 

137 "symbols": [ 

138 symbol, 

139 ] 

140 * pos_count, 

141 "coord_frac": coord_f, 

142 } 

143 ) 

144 else: 

145 coord_dict["symbols"] += [ 

146 symbol, 

147 ] * pos_count 

148 # import pdb; pdb.set_trace() 

149 coord_dict["coord_frac"] = np.vstack( 

150 [coord_dict["coord_frac"], coord_f] 

151 ) 

152 

153 elif name == "free energy": 

154 value = raw_value * Hartree 

155 elif name == "forces": 

156 value = raw_value * Hartree / Bohr 

157 elif name == "atomic_magnetization": 

158 value = raw_value 

159 elif name == "net_magnetization": 

160 value = raw_value 

161 elif name == "stress": 

162 # Stress is in eV/Ang^3, may need to convert to Virial later when cell is known 

163 # For low-dimension stress info, use stress_equiv 

164 stress_ev_a3 = raw_value * GPa 

165 if stress_ev_a3.shape != (3, 3): 

166 raise ValueError("Stress from static file is not a 3x3 matrix!") 

167 # make the stress in voigt notation 

168 # TODO: check the order! 

169 value = np.array( 

170 [ 

171 stress_ev_a3[0, 0], 

172 stress_ev_a3[1, 1], 

173 stress_ev_a3[2, 2], 

174 stress_ev_a3[1, 2], 

175 stress_ev_a3[0, 2], 

176 stress_ev_a3[0, 1], 

177 ] 

178 ) 

179 elif name == "stress_equiv": 

180 # Only store the size up to the max. periodic directions, 

181 # let the atom parser decide how to transform the matrix 

182 value = raw_value * GPa 

183 elif name == "stress_1d": 

184 value = raw_value * Hartree / Bohr 

185 elif name == "stress_2d": 

186 value = raw_value * Hartree / (Bohr**2) 

187 elif name == "lattice": 

188 value = raw_value * Bohr 

189 

190 # Non-frac coord 

191 if value is not None: 

192 block_dict[name] = value 

193 # Finally, update the atomic positions 

194 # TODO: should we keep a default? 

195 if coord_dict != {}: 

196 block_dict["atoms"] = coord_dict 

197 

198 return block_dict 

199 

200 

201def _add_cell_info(static_steps, cell=None): 

202 """Use the cell information to convert positions 

203 if lattice exists in each step, use it to convert coord_frac 

204 else use the external cell (for example from inpt file) 

205 Args: 

206 static_steps: raw list of steps 

207 cell: external lattice information 

208 """ 

209 new_steps = [] 

210 for step in static_steps: 

211 new_step = step.copy() 

212 if "lattice" in step: 

213 lat = step["lattice"] 

214 elif cell is not None: 

215 lat = cell 

216 else: 

217 lat = None 

218 

219 if (lat is not None) and (step.get("atoms", None) is not None): 

220 coord_frac = new_step["atoms"]["coord_frac"] 

221 coord_cart = np.dot(coord_frac, lat) 

222 new_step["atoms"]["coord"] = coord_cart 

223 new_steps.append(new_step) 

224 return new_steps 

225 

226 

227@writer 

228def _write_static( 

229 fileobj, 

230 data_dict, 

231): 

232 raise NotImplementedError("Writing static file from SPARC-X-API not supported!")