1""" 

2ASE Calculator for interatomic models compatible with the Knowledgebase 

3of Interatomic Models (KIM) application programming interface (API). 

4Written by: 

5 

6Mingjian Wen 

7Daniel S. Karls 

8University of Minnesota 

9""" 

10import numpy as np 

11 

12from ase.calculators.calculator import Calculator 

13from ase.calculators.calculator import compare_atoms 

14 

15from . import kimpy_wrappers 

16from . import neighborlist 

17 

18 

19class KIMModelData: 

20 """Initializes and subsequently stores the KIM API Portable Model 

21 object, KIM API ComputeArguments object, and the neighbor list 

22 object used by instances of KIMModelCalculator. Also stores the 

23 arrays which are registered in the KIM API and which are used to 

24 communicate with the model. 

25 """ 

26 

27 def __init__(self, model_name, ase_neigh, neigh_skin_ratio, debug=False): 

28 self.model_name = model_name 

29 self.ase_neigh = ase_neigh 

30 self.debug = debug 

31 

32 # Initialize KIM API Portable Model object and ComputeArguments object 

33 self.init_kim() 

34 

35 # Set cutoff 

36 model_influence_dist = self.kim_model.get_influence_distance() 

37 model_cutoffs, padding_not_require_neigh = ( 

38 self.kim_model.get_neighbor_list_cutoffs_and_hints() 

39 ) 

40 

41 self.species_map = self.create_species_map() 

42 

43 # Initialize neighbor list object 

44 self.init_neigh( 

45 neigh_skin_ratio, 

46 model_influence_dist, 

47 model_cutoffs, 

48 padding_not_require_neigh, 

49 ) 

50 

51 def init_kim(self): 

52 """Create the KIM API Portable Model object and KIM API ComputeArguments 

53 object 

54 """ 

55 if self.kim_initialized: 

56 return 

57 

58 self.kim_model = kimpy_wrappers.PortableModel(self.model_name, self.debug) 

59 

60 # KIM API model object is what actually creates/destroys the ComputeArguments 

61 # object, so we must pass it as a parameter 

62 self.compute_args = self.kim_model.compute_arguments_create() 

63 

64 def init_neigh( 

65 self, 

66 neigh_skin_ratio, 

67 model_influence_dist, 

68 model_cutoffs, 

69 padding_not_require_neigh, 

70 ): 

71 """Initialize neighbor list, either an ASE-native neighborlist 

72 or one created using the neighlist module in kimpy 

73 """ 

74 neigh_list_object_type = ( 

75 neighborlist.ASENeighborList 

76 if self.ase_neigh 

77 else neighborlist.KimpyNeighborList 

78 ) 

79 self.neigh = neigh_list_object_type( 

80 self.compute_args, 

81 neigh_skin_ratio, 

82 model_influence_dist, 

83 model_cutoffs, 

84 padding_not_require_neigh, 

85 self.debug, 

86 ) 

87 

88 def update_compute_args_pointers(self, energy, forces): 

89 self.compute_args.update( 

90 self.num_particles, 

91 self.species_code, 

92 self.particle_contributing, 

93 self.coords, 

94 energy, 

95 forces, 

96 ) 

97 

98 def create_species_map(self): 

99 """Get all the supported species of the KIM model and the 

100 corresponding integer codes used by the model 

101 

102 Returns 

103 ------- 

104 species_map : dict 

105 key : str 

106 chemical symbols (e.g. "Ar") 

107 value : int 

108 species integer code (e.g. 1) 

109 """ 

110 supported_species, codes = self.get_model_supported_species_and_codes() 

111 species_map = dict() 

112 for i, spec in enumerate(supported_species): 

113 species_map[spec] = codes[i] 

114 if self.debug: 

115 print( 

116 "Species {} is supported and its code is: {}".format(spec, codes[i]) 

117 ) 

118 

119 return species_map 

120 

121 @property 

122 def padding_image_of(self): 

123 return self.neigh.padding_image_of 

124 

125 @property 

126 def num_particles(self): 

127 return self.neigh.num_particles 

128 

129 @property 

130 def coords(self): 

131 return self.neigh.coords 

132 

133 @property 

134 def particle_contributing(self): 

135 return self.neigh.particle_contributing 

136 

137 @property 

138 def species_code(self): 

139 return self.neigh.species_code 

140 

141 @property 

142 def kim_initialized(self): 

143 return hasattr(self, "kim_model") 

144 

145 @property 

146 def neigh_initialized(self): 

147 return hasattr(self, "neigh") 

148 

149 @property 

150 def get_model_supported_species_and_codes(self): 

151 return self.kim_model.get_model_supported_species_and_codes 

152 

153 

154class KIMModelCalculator(Calculator): 

155 """Calculator that works with KIM Portable Models (PMs). 

156 

157 Calculator that carries out direct communication between ASE and a 

158 KIM Portable Model (PM) through the kimpy library (which provides a 

159 set of python bindings to the KIM API). 

160 

161 Parameters 

162 ---------- 

163 model_name : str 

164 The unique identifier assigned to the interatomic model (for 

165 details, see https://openkim.org/doc/schema/kim-ids) 

166 

167 ase_neigh : bool, optional 

168 False (default): Use kimpy's neighbor list library 

169 

170 True: Use ASE's internal neighbor list mechanism (usually slower 

171 than the kimpy neighlist library) 

172 

173 neigh_skin_ratio : float, optional 

174 Used to determine the neighbor list cutoff distance, r_neigh, 

175 through the relation r_neigh = (1 + neigh_skin_ratio) * rcut, 

176 where rcut is the model's influence distance. (Default: 0.2) 

177 

178 release_GIL : bool, optional 

179 Whether to release python GIL. Releasing the GIL allows a KIM 

180 model to run with multiple concurrent threads. (Default: False) 

181 

182 debug : bool, optional 

183 If True, detailed information is printed to stdout. (Default: 

184 False) 

185 """ 

186 

187 implemented_properties = ["energy", "forces", "stress"] 

188 

189 def __init__( 

190 self, 

191 model_name, 

192 ase_neigh=False, 

193 neigh_skin_ratio=0.2, 

194 release_GIL=False, 

195 debug=False, 

196 *args, 

197 **kwargs 

198 ): 

199 super().__init__(*args, **kwargs) 

200 

201 self.model_name = model_name 

202 self.release_GIL = release_GIL 

203 self.debug = debug 

204 

205 if neigh_skin_ratio < 0: 

206 raise ValueError('Argument "neigh_skin_ratio" must be non-negative') 

207 

208 # Model output 

209 self.energy = None 

210 self.forces = None 

211 

212 # Create KIMModelData object. This will take care of creating and storing the KIM 

213 # API Portable Model object, KIM API ComputeArguments object, and the neighbor 

214 # list object that our calculator needs 

215 self.kimmodeldata = KIMModelData( 

216 self.model_name, ase_neigh, neigh_skin_ratio, self.debug 

217 ) 

218 

219 def __enter__(self): 

220 return self 

221 

222 def __exit__(self, exc_type, value, traceback): 

223 pass 

224 

225 def __repr__(self): 

226 return "KIMModelCalculator(model_name={})".format(self.model_name) 

227 

228 def calculate( 

229 self, 

230 atoms=None, 

231 properties=["energy", "forces", "stress"], 

232 system_changes=["positions", "numbers", "cell", "pbc"], 

233 ): 

234 """ 

235 Inherited method from the ase Calculator class that is called by 

236 get_property() 

237 

238 Parameters 

239 ---------- 

240 atoms : Atoms 

241 Atoms object whose properties are desired 

242 

243 properties : list of str 

244 List of what needs to be calculated. Can be any combination 

245 of 'energy', 'forces' and 'stress'. 

246 

247 system_changes : list of str 

248 List of what has changed since last calculation. Can be any 

249 combination of these six: 'positions', 'numbers', 'cell', 

250 and 'pbc'. 

251 """ 

252 

253 Calculator.calculate(self, atoms, properties, system_changes) 

254 

255 # Update KIM API input data and neighbor list, if necessary 

256 if system_changes: 

257 if self.need_neigh_update(atoms, system_changes): 

258 self.update_neigh(atoms, self.species_map) 

259 self.energy = np.array([0.0], dtype=np.double) 

260 self.forces = np.zeros([self.num_particles[0], 3], dtype=np.double) 

261 self.update_compute_args_pointers(self.energy, self.forces) 

262 else: 

263 self.update_kim_coords(atoms) 

264 

265 self.kim_model.compute(self.compute_args, self.release_GIL) 

266 

267 energy = self.energy[0] 

268 forces = self.assemble_padding_forces() 

269 

270 try: 

271 volume = atoms.get_volume() 

272 stress = self.compute_virial_stress(self.forces, self.coords, volume) 

273 except ValueError: # Volume cannot be computed 

274 stress = None 

275 

276 # Quantities passed back to ASE 

277 self.results["energy"] = energy 

278 self.results["free_energy"] = energy 

279 self.results["forces"] = forces 

280 self.results["stress"] = stress 

281 

282 def check_state(self, atoms, tol=1e-15): 

283 return compare_atoms(self.atoms, atoms, excluded_properties={'initial_charges', 

284 'initial_magmoms'}) 

285 

286 def assemble_padding_forces(self): 

287 """ 

288 Assemble forces on padding atoms back to contributing atoms. 

289 

290 Parameters 

291 ---------- 

292 forces : 2D array of doubles 

293 Forces on both contributing and padding atoms 

294 

295 num_contrib: int 

296 Number of contributing atoms 

297 

298 padding_image_of : 1D array of int 

299 Atom number, of which the padding atom is an image 

300 

301 

302 Returns 

303 ------- 

304 Total forces on contributing atoms. 

305 """ 

306 

307 total_forces = np.array(self.forces[: self.num_contributing_particles]) 

308 

309 if self.padding_image_of.size != 0: 

310 pad_forces = self.forces[self.num_contributing_particles:] 

311 for f, org_index in zip(pad_forces, self.padding_image_of): 

312 total_forces[org_index] += f 

313 

314 return total_forces 

315 

316 @staticmethod 

317 def compute_virial_stress(forces, coords, volume): 

318 """Compute the virial stress in Voigt notation. 

319 

320 Parameters 

321 ---------- 

322 forces : 2D array 

323 Partial forces on all atoms (padding included) 

324 

325 coords : 2D array 

326 Coordinates of all atoms (padding included) 

327 

328 volume : float 

329 Volume of cell 

330 

331 Returns 

332 ------- 

333 stress : 1D array 

334 stress in Voigt order (xx, yy, zz, yz, xz, xy) 

335 """ 

336 stress = np.zeros(6) 

337 stress[0] = -np.dot(forces[:, 0], coords[:, 0]) / volume 

338 stress[1] = -np.dot(forces[:, 1], coords[:, 1]) / volume 

339 stress[2] = -np.dot(forces[:, 2], coords[:, 2]) / volume 

340 stress[3] = -np.dot(forces[:, 1], coords[:, 2]) / volume 

341 stress[4] = -np.dot(forces[:, 0], coords[:, 2]) / volume 

342 stress[5] = -np.dot(forces[:, 0], coords[:, 1]) / volume 

343 

344 return stress 

345 

346 def get_model_supported_species_and_codes(self): 

347 return self.kimmodeldata.get_model_supported_species_and_codes 

348 

349 @property 

350 def update_compute_args_pointers(self): 

351 return self.kimmodeldata.update_compute_args_pointers 

352 

353 @property 

354 def kim_model(self): 

355 return self.kimmodeldata.kim_model 

356 

357 @property 

358 def compute_args(self): 

359 return self.kimmodeldata.compute_args 

360 

361 @property 

362 def num_particles(self): 

363 return self.kimmodeldata.num_particles 

364 

365 @property 

366 def coords(self): 

367 return self.kimmodeldata.coords 

368 

369 @property 

370 def padding_image_of(self): 

371 return self.kimmodeldata.padding_image_of 

372 

373 @property 

374 def species_map(self): 

375 return self.kimmodeldata.species_map 

376 

377 @property 

378 def neigh(self): 

379 return self.kimmodeldata.neigh 

380 

381 @property 

382 def num_contributing_particles(self): 

383 return self.neigh.num_contributing_particles 

384 

385 @property 

386 def update_kim_coords(self): 

387 return self.neigh.update_kim_coords 

388 

389 @property 

390 def need_neigh_update(self): 

391 return self.neigh.need_neigh_update 

392 

393 @property 

394 def update_neigh(self): 

395 return self.neigh.update