Coverage for calorine/nep/model.py: 100%

1from dataclasses import dataclass 

2from itertools import product 

3 

4import numpy as np 

5 

6NetworkWeights = dict[str, dict[str, np.ndarray]] 

7DescriptorWeights = dict[tuple[str, str], np.ndarray] 

8RestartParameters = dict[str, dict[str, dict[str, np.ndarray]]] 

9 

10 

11def _get_restart_contents(filename: str) -> tuple[list[float], list[float]]: 

12 """Parses a ``nep.restart`` file, and returns an unformatted list of the 

13 mean and standard deviation for all model parameters. 

14 Intended to be used by the py:meth:`~Model.read_restart` function. 

15 

16 Parameters 

17 ---------- 

18 filename 

19 input file name 

20 """ 

21 mu = [] # Mean 

22 sigma = [] # Standard deviation 

23 with open(filename) as f: 

24 for k, line in enumerate(f.readlines()): 

25 flds = line.split() 

26 assert len(flds) != 0, f'Empty line number {k}' 

27 if len(flds) == 2: 

28 mu.append(float(flds[0])) 

29 sigma.append(float(flds[1])) 

30 else: 

31 raise IOError(f'Failed to parse line {k} from {filename}') 

32 return mu, sigma 

33 

34 

35def _get_model_type(first_row: list[str]) -> str: 

36 """Parses a the first row of a ``nep.txt`` file, and returns the 

37 type of NEP model. Available types are `potential`, `potential_with_charges`, 

38 `dipole`, and `polarizability`. 

39 

40 Parameters 

41 ---------- 

42 first_row 

43 First row of a NEP file, split by white space. 

44 """ 

45 model_type = first_row[0] 

46 if 'charge' in model_type: 

47 return 'potential_with_charges' 

48 elif 'dipole' in model_type: 

49 return 'dipole' 

50 elif 'polarizability' in model_type: 

51 return 'polarizability' 

52 return 'potential' 

53 

54 

55def _get_nep_contents(filename: str) -> tuple[dict, list[float]]: 

56 """Parses a ``nep.txt`` file, and returns a dict describing the header 

57 and an unformatted list of all model parameters. 

58 Intended to be used by the :func:`read_model <calorine.nep.read_model>` function. 

59 

60 Parameters 

61 ---------- 

62 filename 

63 input file name 

64 """ 

65 # parse file and split header and parameters 

66 header = [] 

67 parameters = [] 

68 nheader = 5 # 5 rows for NEP2, 6-7 rows for NEP3 onwards 

69 base_line = 3 

70 with open(filename) as f: 

71 for k, line in enumerate(f.readlines()): 

72 flds = line.split() 

73 assert len(flds) != 0, f'Empty line number {k}' 

74 if k == 0 and 'zbl' in flds[0]: 

75 base_line += 1 

76 nheader += 1 

77 if k == base_line and 'basis_size' in flds[0]: 

78 # Introduced in nep.txt after GPUMD v3.2 

79 nheader += 1 

80 if k < nheader: 

81 header.append(tuple(flds)) 

82 elif len(flds) == 1: 

83 parameters.append(float(flds[0])) 

84 else: 

85 raise IOError(f'Failed to parse line {k} from {filename}') 

86 # compile data from the header into a dict 

87 data = {} 

88 for flds in header: 

89 if flds[0] in ['cutoff', 'zbl']: 

90 data[flds[0]] = tuple(map(float, flds[1:])) 

91 elif flds[0] in ['n_max', 'l_max', 'ANN', 'basis_size']: 

92 data[flds[0]] = tuple(map(int, flds[1:])) 

93 elif flds[0].startswith('nep'): 

94 version = flds[0].replace('nep', '').split('_')[0] 

95 version = int(version) 

96 data['version'] = version 

97 data['types'] = flds[2:] 

98 data['model_type'] = _get_model_type(flds) 

99 else: 

100 raise ValueError(f'Unknown field: {flds[0]}') 

101 return data, parameters 

102 

103 

104def _sort_descriptor_parameters(parameters: list[float], 

105 types: list[str], 

106 n_max_radial: int, 

107 n_basis_radial: int, 

108 n_max_angular: int, 

109 n_basis_angular: int) -> tuple[DescriptorWeights, 

110 DescriptorWeights]: 

111 """Reads a list of descriptors parameters and sorts them into two 

112 appropriately structured `dicts`, one for radial and one for angular descriptor weights. 

113 Intended to be used by the :func:`read_model <calorine.nep.read_model>` function. 

114 """ 

115 # split up descriptor by chemical species and radial/angular 

116 n_types = len(types) 

117 n = len(parameters) / (n_types * n_types) 

118 assert n.is_integer(), 'number of descriptor groups must be an integer' 

119 n = int(n) 

120 

121 m = (n_max_radial + 1) * (n_basis_radial + 1) 

122 descriptor_weights = parameters.reshape((n, n_types * n_types)).T 

123 descriptor_weights_radial = descriptor_weights[:, :m] 

124 descriptor_weights_angular = descriptor_weights[:, m:] 

125 

126 # add descriptors to data dict 

127 radial_descriptor_weights = {} 

128 angular_descriptor_weights = {} 

129 m = -1 

130 for i, j in product(range(n_types), repeat=2): 

131 m += 1 

132 s1, s2 = types[i], types[j] 

133 radial_descriptor_weights[(s1, s2)] = descriptor_weights_radial[m, :].reshape( 

134 (n_max_radial + 1, n_basis_radial + 1) 

135 ) 

136 angular_descriptor_weights[(s1, s2)] = descriptor_weights_angular[m, :].reshape( 

137 (n_max_angular + 1, n_basis_angular + 1) 

138 ) 

139 return radial_descriptor_weights, angular_descriptor_weights 

140 

141 

142def _sort_ann_parameters(parameters: list[float], 

143 ann_groupings: list[str], 

144 n_neuron: int, 

145 n_networks: int, 

146 n_bias: int, 

147 n_descriptor: int, 

148 is_polarizability_model: bool, 

149 is_model_with_charges: bool 

150 ) -> NetworkWeights: 

151 """Reads a list of model parameters and sorts them into an appropriately structured `dict`. 

152 Intended to be used by the :func:`read_model <calorine.nep.read_model>` function. 

153 """ 

154 n_ann_input_weights = (n_descriptor + 1) * n_neuron # weights + bias 

155 n_ann_output_weights = 2*n_neuron if is_model_with_charges else n_neuron # only weights 

156 n_ann_parameters = ( 

157 n_ann_input_weights + n_ann_output_weights 

158 ) * n_networks + n_bias 

159 

160 # Group ANN parameters 

161 pars = {} 

162 n1 = 0 

163 n_network_params = n_ann_input_weights + n_ann_output_weights # except last bias(es) 

164 

165 n_count = 2 if is_polarizability_model else 1 

166 n_outputs = 2 if is_model_with_charges else 1 

167 for count in range(n_count): 

168 # if polarizability model, all parameters including bias are repeated 

169 # need to offset n1 by +1 to handle bias 

170 n1 += count 

171 for s in ann_groupings: 

172 # Get the parameters for the ANN; in the case of NEP4, there is effectively 

173 # one network per atomic species. 

174 ann_parameters = parameters[n1 : n1 + n_network_params] 

175 ann_input_weights = ann_parameters[:n_ann_input_weights] 

176 w0 = np.zeros((n_neuron, n_descriptor)) 

177 w0[...] = np.nan 

178 b0 = np.zeros((n_neuron, 1)) 

179 b0[...] = np.nan 

180 for n in range(n_neuron): 

181 for nu in range(n_descriptor): 

182 w0[n, nu] = ann_input_weights[n * n_descriptor + nu] 

183 b0[:, 0] = ann_input_weights[n_neuron * n_descriptor :] 

184 

185 assert np.all( 

186 w0.shape == (n_neuron, n_descriptor) 

187 ), f'w0 has invalid shape for key {s}; please submit a bug report' 

188 assert np.all( 

189 b0.shape == (n_neuron, 1) 

190 ), f'b0 has invalid shape for key {s}; please submit a bug report' 

191 assert not np.any( 

192 np.isnan(w0) 

193 ), f'some weights in w0 are nan for key {s}; please submit a bug report' 

194 assert not np.any( 

195 np.isnan(b0) 

196 ), f'some weights in b0 are nan for key {s}; please submit a bug report' 

197 

198 ann_output_weights = ann_parameters[ 

199 n_ann_input_weights : n_ann_input_weights + n_ann_output_weights 

200 ] 

201 w1 = np.zeros((1, n_neuron * n_outputs)) 

202 w1[0, :] = ann_output_weights[:] 

203 assert np.all( 

204 w1.shape == (1, n_neuron * n_outputs) 

205 ), f'w1 has invalid shape for key {s}; please submit a bug report' 

206 assert not np.any( 

207 np.isnan(w1) 

208 ), f'some weights in w1 are nan for key {s}; please submit a bug report' 

209 

210 if count == 0 and n_outputs == 1: 

211 pars[s] = dict(w0=w0, b0=b0, w1=w1) 

212 elif count == 0 and n_outputs == 2: 

213 pars[s] = dict(w0=w0, b0=b0, w1=w1[0, :n_neuron], w1_charge=w1[0, n_neuron:]) 

214 else: 

215 pars[s].update({'w0_polar': w0, 'b0_polar': b0, 'w1_polar': w1}) 

216 # Jump to bias 

217 n1 += n_network_params 

218 if n_bias > 1 and not is_model_with_charges: 

219 # For NEP5 models we additionally have one bias term per species. 

220 # Currently NEP5 only exists for potential models, but we'll 

221 # keep it here in case it gets added down the line. 

222 bias_label = 'b1' if count == 0 else 'b1_polar' 

223 pars[s][bias_label] = parameters[n1] 

224 n1 += 1 

225 # For NEP3 and NEP4 we only have one bias. 

226 # For NEP4 with charges we have two biases. 

227 # For NEP5 we have one bias per species, and one global. 

228 if count == 0 and n_outputs == 1: 

229 pars['b1'] = parameters[n1] 

230 elif count == 0 and n_outputs == 2: 

231 pars['sqrt_epsilon_infinity'] = parameters[n1] 

232 pars['b1'] = parameters[n1+1] 

233 else: 

234 pars['b1_polar'] = parameters[n1] 

235 sum = 0 

236 for s in pars.keys(): 

237 if s.startswith('b1') or s.startswith('sqrt'): 

238 sum += 1 

239 else: 

240 sum += np.sum([np.array(p).size for p in pars[s].values()]) 

241 assert sum == n_ann_parameters * n_count, ( 

242 'Inconsistent number of parameters accounted for; please submit a bug report\n' 

243 f'{sum} != {n_ann_parameters}' 

244 ) 

245 return pars 

246 

247 

248@dataclass 

249class Model: 

250 r"""Objects of this class represent a NEP model in a form suitable for 

251 inspection and manipulation. Typically a :class:`Model` object is instantiated 

252 by calling the :func:`read_model <calorine.nep.read_model>` function. 

253 

254 Attributes 

255 ---------- 

256 version : int 

257 NEP version. 

258 model_type: str 

259 One of ``potential``, ``dipole`` or ``polarizability``. 

260 types : tuple[str, ...] 

261 Chemical species that this model represents. 

262 radial_cutoff : float 

263 The radial cutoff parameter in Å. 

264 angular_cutoff : float 

265 The angular cutoff parameter in Å. 

266 max_neighbors_radial : int 

267 Maximum number of neighbors in neighbor list for radial terms. 

268 max_neighbors_angular : int 

269 Maximum number of neighbors in neighbor list for angular terms. 

270 radial_typewise_cutoff_factor : float 

271 The radial cutoff factor if use_typewise_cutoff is used. 

272 angular_typewise_cutoff_factor : float 

273 The angular cutoff factor if use_typewise_cutoff is used. 

274 zbl : tuple[float, float] 

275 Inner and outer cutoff for transition to ZBL potential. 

276 zbl_typewise_cutoff_factor : float 

277 Typewise cutoff when use_typewise_cutoff_zbl is used. 

278 n_basis_radial : int 

279 Number of radial basis functions :math:`n_\mathrm{basis}^\mathrm{R}`. 

280 n_basis_angular : int 

281 Number of angular basis functions :math:`n_\mathrm{basis}^\mathrm{A}`. 

282 n_max_radial : int 

283 Maximum order of Chebyshev polymonials included in 

284 radial expansion :math:`n_\mathrm{max}^\mathrm{R}`. 

285 n_max_angular : int 

286 Maximum order of Chebyshev polymonials included in 

287 angular expansion :math:`n_\mathrm{max}^\mathrm{A}`. 

288 l_max_3b : int 

289 Maximum expansion order for three-body terms :math:`l_\mathrm{max}^\mathrm{3b}`. 

290 l_max_4b : int 

291 Maximum expansion order for four-body terms :math:`l_\mathrm{max}^\mathrm{4b}`. 

292 l_max_5b : int 

293 Maximum expansion order for five-body terms :math:`l_\mathrm{max}^\mathrm{5b}`. 

294 n_descriptor_radial : int 

295 Dimension of radial part of descriptor. 

296 n_descriptor_angular : int 

297 Dimension of angular part of descriptor. 

298 n_neuron : int 

299 Number of neurons in hidden layer. 

300 n_parameters : int 

301 Total number of parameters including scalers (which are not fit parameters). 

302 n_descriptor_parameters : int 

303 Number of parameters in descriptor. 

304 n_ann_parameters : int 

305 Number of neural network weights. 

306 ann_parameters : dict[tuple[str, dict[str, np.darray]]] 

307 Neural network weights. 

308 q_scaler : List[float] 

309 Scaling parameters. 

310 radial_descriptor_weights : dict[tuple[str, str], np.ndarray] 

311 Radial descriptor weights by combination of species; the array for each combination 

312 has dimensions of 

313 :math:`(n_\mathrm{max}^\mathrm{R}+1) \times (n_\mathrm{basis}^\mathrm{R}+1)`. 

314 angular_descriptor_weights : dict[tuple[str, str], np.ndarray] 

315 Angular descriptor weights by combination of species; the array for each combination 

316 has dimensions of 

317 :math:`(n_\mathrm{max}^\mathrm{A}+1) \times (n_\mathrm{basis}^\mathrm{A}+1)`. 

318 sqrt_epsilon_infinity : Optional[float] 

319 Square root of epsilon infinity $\epsilon_\infty$ (only for NEP models with charges). 

320 restart_parameters : dict[str, dict[str, dict[str, np.ndarray]]] 

321 NEP restart parameters. A nested dictionary that contains the mean (mu) and standard 

322 deviation (sigma) for the ANN and descriptor parameters. Is set using the 

323 py:meth:`~Model.read_restart` method. Defaults to None. 

324 """ 

325 

326 version: int 

327 model_type: str 

328 types: tuple[str, ...] 

329 

330 radial_cutoff: float 

331 angular_cutoff: float 

332 

333 n_basis_radial: int 

334 n_basis_angular: int 

335 n_max_radial: int 

336 n_max_angular: int 

337 l_max_3b: int 

338 l_max_4b: int 

339 l_max_5b: int 

340 n_descriptor_radial: int 

341 n_descriptor_angular: int 

342 

343 n_neuron: int 

344 n_parameters: int 

345 n_descriptor_parameters: int 

346 n_ann_parameters: int 

347 ann_parameters: NetworkWeights 

348 q_scaler: list[float] 

349 radial_descriptor_weights: DescriptorWeights 

350 angular_descriptor_weights: DescriptorWeights 

351 sqrt_epsilon_infinity: float = None 

352 restart_parameters: RestartParameters = None 

353 

354 zbl: tuple[float, float] = None 

355 zbl_typewise_cutoff_factor: float = None 

356 max_neighbors_radial: int = None 

357 max_neighbors_angular: int = None 

358 radial_typewise_cutoff_factor: float = None 

359 angular_typewise_cutoff_factor: float = None 

360 

361 _special_fields = [ 

362 'ann_parameters', 

363 'q_scaler', 

364 'radial_descriptor_weights', 

365 'angular_descriptor_weights', 

366 ] 

367 

368 def __str__(self) -> str: 

369 s = [] 

370 for fld in self.__dataclass_fields__: 

371 if fld not in self._special_fields: 

372 s += [f'{fld:22} : {getattr(self, fld)}'] 

373 return '\n'.join(s) 

374 

375 def _repr_html_(self) -> str: 

376 s = [] 

377 s += ['<table border="1" class="dataframe"'] 

378 s += [ 

379 '<thead><tr><th style="text-align: left;">Field</th><th>Value</th></tr></thead>' 

380 ] 

381 s += ['<tbody>'] 

382 for fld in self.__dataclass_fields__: 

383 if fld not in self._special_fields: 

384 s += [ 

385 f'<tr><td style="text-align: left;">{fld:22}</td>' 

386 f'<td>{getattr(self, fld)}</td><tr>' 

387 ] 

388 for fld in self._special_fields: 

389 d = getattr(self, fld) 

390 # print('xxx', fld, d) 

391 if fld.endswith('descriptor_weights'): 

392 dim = list(d.values())[0].shape 

393 elif fld == 'ann_parameters' and self.version == 4: 

394 dim = (len(self.types), len(list(d.values())[0])) 

395 else: 

396 dim = len(d) 

397 s += [ 

398 f'<tr><td style="text-align: left;">Dimension of {fld:22}</td><td>{dim}</td><tr>' 

399 ] 

400 s += ['</tbody>'] 

401 s += ['</table>'] 

402 return ''.join(s) 

403 

404 def remove_species(self, species: list[str]): 

405 """Removes one or more species from the model. 

406 

407 This method modifies the model in-place by removing all parameters 

408 associated with the specified chemical species. It prunes the species 

409 list, the Artificial Neural Network (ANN) parameters, and the 

410 descriptor weights. It also recalculates the total number of 

411 parameters in the model. 

412 

413 Parameters 

414 ---------- 

415 species 

416 A list of species names (str) to remove from the model. 

417 

418 Raises 

419 ------ 

420 ValueError 

421 If any of the provided species is not found in the model. 

422 """ 

423 for s in species: 

424 if s not in self.types: 

425 raise ValueError(f'{s} is not a species supported by the NEP model') 

426 

427 # --- Prune attributes based on species --- 

428 types_to_keep = [t for t in self.types if t not in species] 

429 self.types = tuple(types_to_keep) 

430 

431 # Prune ANN parameters (for NEP4 and NEP5) 

432 if self.version in [4, 5]: 

433 self.ann_parameters = { 

434 key: value for key, value in self.ann_parameters.items() 

435 if key in types_to_keep or key.startswith('b1') 

436 } 

437 

438 # Prune descriptor weights 

439 # key is here a tuple, (species1, species2) 

440 self.radial_descriptor_weights = { 

441 key: value for key, value in self.radial_descriptor_weights.items() 

442 if key[0] in types_to_keep and key[1] in types_to_keep 

443 } 

444 self.angular_descriptor_weights = { 

445 key: value for key, value in self.angular_descriptor_weights.items() 

446 if key[0] in types_to_keep and key[1] in types_to_keep 

447 } 

448 

449 # Prune restart parameters if they have been loaded 

450 if self.restart_parameters is not None: 

451 for param_type in ['mu', 'sigma']: 

452 # Prune ANN restart parameters 

453 ann_key = f'ann_{param_type}' 

454 if self.version in [4, 5]: 

455 self.restart_parameters[ann_key] = { 

456 key: value for key, value in self.restart_parameters[ann_key].items() 

457 if key in types_to_keep or key.startswith('b1') 

458 } 

459 

460 # Prune descriptor restart parameters 

461 for desc_type in ['radial', 'angular']: 

462 key = f'{desc_type}_descriptor_{param_type}' 

463 self.restart_parameters[key] = { 

464 k: v for k, v in self.restart_parameters[key].items() 

465 if k[0] in types_to_keep and k[1] in types_to_keep 

466 } 

467 

468 # --- Recalculate parameter counts --- 

469 n_types = len(self.types) 

470 n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular 

471 

472 # Recalculate descriptor parameter count 

473 self.n_descriptor_parameters = n_types**2 * ( 

474 (self.n_max_radial + 1) * (self.n_basis_radial + 1) 

475 + (self.n_max_angular + 1) * (self.n_basis_angular + 1) 

476 ) 

477 

478 # Recalculate ANN parameter count 

479 if self.version == 3: 

480 n_networks = 1 

481 n_bias = 1 

482 elif self.version == 4: 

483 n_networks = n_types 

484 n_bias = 1 

485 else: # NEP5 

486 n_networks = n_types 

487 n_bias = 1 + n_types 

488 

489 n_ann_input_weights = (n_descriptor + 1) * self.n_neuron 

490 n_ann_output_weights = self.n_neuron 

491 self.n_ann_parameters = ( 

492 n_ann_input_weights + n_ann_output_weights 

493 ) * n_networks + n_bias 

494 

495 # Recalculate total parameter count 

496 self.n_parameters = ( 

497 self.n_ann_parameters 

498 + self.n_descriptor_parameters 

499 + n_descriptor # q_scaler parameters 

500 ) 

501 if self.model_type == 'polarizability': 

502 self.n_parameters += self.n_ann_parameters 

503 

504 def write(self, filename: str) -> None: 

505 """Write NEP model to file in `nep.txt` format.""" 

506 with open(filename, 'w') as f: 

507 # header 

508 version_name = f'nep{self.version}' 

509 if self.zbl is not None: 

510 version_name += '_zbl' 

511 elif self.model_type != 'potential': 

512 version_name += f'_{self.model_type}' 

513 f.write(f'{version_name} {len(self.types)} {" ".join(self.types)}\n') 

514 if self.zbl is not None: 

515 f.write(f'zbl {" ".join(map(str, self.zbl))}\n') 

516 f.write(f'cutoff {self.radial_cutoff} {self.angular_cutoff}') 

517 f.write(f' {self.max_neighbors_radial} {self.max_neighbors_angular}') 

518 if ( 

519 self.radial_typewise_cutoff_factor is not None 

520 and self.angular_typewise_cutoff_factor is not None 

521 ): 

522 f.write(f' {self.radial_typewise_cutoff_factor}' 

523 f' {self.angular_typewise_cutoff_factor}') 

524 if self.zbl_typewise_cutoff_factor: 

525 f.write(f' {self.zbl_typewise_cutoff_factor}') 

526 f.write('\n') 

527 f.write(f'n_max {self.n_max_radial} {self.n_max_angular}\n') 

528 f.write(f'basis_size {self.n_basis_radial} {self.n_basis_angular}\n') 

529 f.write(f'l_max {self.l_max_3b} {self.l_max_4b} {self.l_max_5b}\n') 

530 f.write(f'ANN {self.n_neuron} 0\n') 

531 

532 # neural network weights 

533 keys = self.types if self.version in (4, 5) else ['all_species'] 

534 suffixes = ['', '_polar'] if self.model_type == 'polarizability' else [''] 

535 for suffix in suffixes: 

536 for s in keys: 

537 # Order: w0, b0, w1 (, b1 if NEP5) 

538 # w0 indexed as: n*N_descriptor + nu 

539 w0 = self.ann_parameters[s][f'w0{suffix}'] 

540 b0 = self.ann_parameters[s][f'b0{suffix}'] 

541 w1 = self.ann_parameters[s][f'w1{suffix}'] 

542 for n in range(self.n_neuron): 

543 for nu in range( 

544 self.n_descriptor_radial + self.n_descriptor_angular 

545 ): 

546 f.write(f'{w0[n, nu]:15.7e}\n') 

547 for b in b0[:, 0]: 

548 f.write(f'{b:15.7e}\n') 

549 for v in w1[0, :]: 

550 f.write(f'{v:15.7e}\n') 

551 if self.version == 5: 

552 b1 = self.ann_parameters[s][f'b1{suffix}'] 

553 f.write(f'{b1:15.7e}\n') 

554 b1 = self.ann_parameters[f'b1{suffix}'] 

555 f.write(f'{b1:15.7e}\n') 

556 

557 # descriptor weights 

558 mat = [] 

559 for s1 in self.types: 

560 for s2 in self.types: 

561 mat = np.hstack( 

562 [mat, self.radial_descriptor_weights[(s1, s2)].flatten()] 

563 ) 

564 mat = np.hstack( 

565 [mat, self.angular_descriptor_weights[(s1, s2)].flatten()] 

566 ) 

567 n_types = len(self.types) 

568 n = int(len(mat) / (n_types * n_types)) 

569 mat = mat.reshape((n_types * n_types, n)).T 

570 for v in mat.flatten(): 

571 f.write(f'{v:15.7e}\n') 

572 

573 # scaler 

574 for v in self.q_scaler: 

575 f.write(f'{v:15.7e}\n') 

576 

577 def read_restart(self, filename: str): 

578 """Parses a file in `nep.restart` format and saves the 

579 content in the form of mean and standard deviation for each 

580 parameter in the corresponding NEP model. 

581 

582 Parameters 

583 ---------- 

584 filename 

585 Input file name. 

586 """ 

587 mu, sigma = _get_restart_contents(filename) 

588 restart_parameters = np.array([mu, sigma]).T 

589 

590 is_polarizability_model = self.model_type == 'polarizability' 

591 is_charged_model = self.model_type == 'potential_with_charges' 

592 

593 n1 = self.n_ann_parameters 

594 n1 *= 2 if is_polarizability_model else 1 

595 n2 = n1 + self.n_descriptor_parameters 

596 ann_parameters = restart_parameters[:n1] 

597 descriptor_parameters = np.array(restart_parameters[n1:n2]) 

598 

599 if self.version == 3: 

600 n_networks = 1 

601 n_bias = 1 

602 elif self.version == 4: 

603 # one hidden layer per atomic species 

604 n_networks = len(self.types) 

605 n_bias = 1 

606 else: 

607 raise ValueError(f'Cannot load nep.restart for NEP model version {self.version}') 

608 

609 ann_groups = [s for s in self.ann_parameters.keys() if not s.startswith('b1')] 

610 n_bias = len([s for s in self.ann_parameters.keys() if s.startswith('b1')]) 

611 n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular 

612 restart = {} 

613 

614 for i, content_type in enumerate(['mu', 'sigma']): 

615 ann = _sort_ann_parameters(ann_parameters[:, i], 

616 ann_groups, 

617 self.n_neuron, 

618 n_networks, 

619 n_bias, 

620 n_descriptor, 

621 is_polarizability_model, 

622 is_charged_model) 

623 radial, angular = _sort_descriptor_parameters(descriptor_parameters[:, i], 

624 self.types, 

625 self.n_max_radial, 

626 self.n_basis_radial, 

627 self.n_max_angular, 

628 self.n_basis_angular) 

629 

630 restart[f'ann_{content_type}'] = ann 

631 restart[f'radial_descriptor_{content_type}'] = radial 

632 restart[f'angular_descriptor_{content_type}'] = angular 

633 self.restart_parameters = restart 

634 

635 def write_restart(self, filename: str): 

636 """Write NEP restart parameters to file in `nep.restart` format.""" 

637 keys = self.types if self.version in (4, 5) else ['all_species'] 

638 suffixes = ['', '_polar'] if self.model_type == 'polarizability' else [''] 

639 columns = [] 

640 for i, parameter in enumerate(['mu', 'sigma']): 

641 # neural network weights 

642 ann_parameters = self.restart_parameters[f'ann_{parameter}'] 

643 column = [] 

644 for suffix in suffixes: 

645 for s in keys: 

646 # Order: w0, b0, w1 (, b1 if NEP5) 

647 # w0 indexed as: n*N_descriptor + nu 

648 w0 = ann_parameters[s][f'w0{suffix}'] 

649 b0 = ann_parameters[s][f'b0{suffix}'] 

650 w1 = ann_parameters[s][f'w1{suffix}'] 

651 for n in range(self.n_neuron): 

652 for nu in range( 

653 self.n_descriptor_radial + self.n_descriptor_angular 

654 ): 

655 column.append(f'{w0[n, nu]:15.7e}') 

656 for b in b0[:, 0]: 

657 column.append(f'{b:15.7e}') 

658 for v in w1[0, :]: 

659 column.append(f'{v:15.7e}') 

660 b1 = ann_parameters[f'b1{suffix}'] 

661 column.append(f'{b1:15.7e}') 

662 columns.append(column) 

663 

664 # descriptor weights 

665 radial_descriptor_parameters = self.restart_parameters[f'radial_descriptor_{parameter}'] 

666 angular_descriptor_parameters = self.restart_parameters[ 

667 f'angular_descriptor_{parameter}'] 

668 mat = [] 

669 for s1 in self.types: 

670 for s2 in self.types: 

671 mat = np.hstack( 

672 [mat, radial_descriptor_parameters[(s1, s2)].flatten()] 

673 ) 

674 mat = np.hstack( 

675 [mat, angular_descriptor_parameters[(s1, s2)].flatten()] 

676 ) 

677 n_types = len(self.types) 

678 n = int(len(mat) / (n_types * n_types)) 

679 mat = mat.reshape((n_types * n_types, n)).T 

680 for v in mat.flatten(): 

681 column.append(f'{v:15.7e}') 

682 

683 # Join the mean and standard deviation columns 

684 assert len(columns[0]) == len(columns[1]), 'Length of means must match standard deviation' 

685 joined = [f'{s1} {s2}\n' for s1, s2 in zip(*columns)] 

686 with open(filename, 'w') as f: 

687 f.writelines(joined) 

688 

689 

690def read_model(filename: str) -> Model: 

691 """Parses a file in ``nep.txt`` format and returns the 

692 content in the form of a :class:`Model <calorine.nep.model.Model>` 

693 object. 

694 

695 Parameters 

696 ---------- 

697 filename 

698 Input file name. 

699 """ 

700 data, parameters = _get_nep_contents(filename) 

701 

702 # sanity checks 

703 for fld in ['cutoff', 'basis_size', 'n_max', 'l_max', 'ANN']: 

704 assert fld in data, f'Invalid model file; {fld} line is missing' 

705 assert data['version'] in [ 

706 3, 

707 4, 

708 5, 

709 ], 'Invalid model file; only NEP versions 3, 4 and 5 are currently supported' 

710 

711 # split up cutoff tuple 

712 assert len(data['cutoff']) in [4, 6, 7] 

713 data['radial_cutoff'] = data['cutoff'][0] 

714 data['angular_cutoff'] = data['cutoff'][1] 

715 data['max_neighbors_radial'] = int(data['cutoff'][2]) 

716 data['max_neighbors_angular'] = int(data['cutoff'][3]) 

717 if len(data['cutoff']) >= 6: 

718 data['radial_typewise_cutoff_factor'] = data['cutoff'][4] 

719 data['angular_typewise_cutoff_factor'] = data['cutoff'][5] 

720 if len(data['cutoff']) == 7: 

721 data['zbl_typewise_cutoff_factor'] = data['cutoff'][6] 

722 del data['cutoff'] 

723 

724 # split up basis_size tuple 

725 assert len(data['basis_size']) == 2 

726 data['n_basis_radial'] = data['basis_size'][0] 

727 data['n_basis_angular'] = data['basis_size'][1] 

728 del data['basis_size'] 

729 

730 # split up n_max tuple 

731 assert len(data['n_max']) == 2 

732 data['n_max_radial'] = data['n_max'][0] 

733 data['n_max_angular'] = data['n_max'][1] 

734 del data['n_max'] 

735 

736 # split up nl_max tuple 

737 len_l = len(data['l_max']) 

738 assert len_l in [1, 2, 3] 

739 data['l_max_3b'] = data['l_max'][0] 

740 data['l_max_4b'] = data['l_max'][1] if len_l > 1 else 0 

741 data['l_max_5b'] = data['l_max'][2] if len_l > 2 else 0 

742 del data['l_max'] 

743 

744 # compute dimensions of descriptor components 

745 data['n_descriptor_radial'] = data['n_max_radial'] + 1 

746 l_max_enh = data['l_max_3b'] + (data['l_max_4b'] > 0) + (data['l_max_5b'] > 0)