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 | list[float] 

263 The radial cutoff parameter in Å. 

264 Is a list of radial cutoffs ordered after ``types`` in the case of typewise cutoffs. 

265 angular_cutoff : float | list[float] 

266 The angular cutoff parameter in Å. 

267 Is a list of angular cutoffs ordered after ``types`` in the case of typewise cutoffs. 

268 max_neighbors_radial : int 

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

270 max_neighbors_angular : int 

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

272 radial_typewise_cutoff_factor : float 

273 The radial cutoff factor if use_typewise_cutoff is used. 

274 angular_typewise_cutoff_factor : float 

275 The angular cutoff factor if use_typewise_cutoff is used. 

276 zbl : tuple[float, float] 

277 Inner and outer cutoff for transition to ZBL potential. 

278 zbl_typewise_cutoff_factor : float 

279 Typewise cutoff when use_typewise_cutoff_zbl is used. 

280 n_basis_radial : int 

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

282 n_basis_angular : int 

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

284 n_max_radial : int 

285 Maximum order of Chebyshev polymonials included in 

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

287 n_max_angular : int 

288 Maximum order of Chebyshev polymonials included in 

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

290 l_max_3b : int 

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

292 l_max_4b : int 

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

294 l_max_5b : int 

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

296 n_descriptor_radial : int 

297 Dimension of radial part of descriptor. 

298 n_descriptor_angular : int 

299 Dimension of angular part of descriptor. 

300 n_neuron : int 

301 Number of neurons in hidden layer. 

302 n_parameters : int 

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

304 n_descriptor_parameters : int 

305 Number of parameters in descriptor. 

306 n_ann_parameters : int 

307 Number of neural network weights. 

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

309 Neural network weights. 

310 q_scaler : List[float] 

311 Scaling parameters. 

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

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

314 has dimensions of 

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

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

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

318 has dimensions of 

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

320 sqrt_epsilon_infinity : Optional[float] 

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

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

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

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

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

326 """ 

327 

328 version: int 

329 model_type: str 

330 types: tuple[str, ...] 

331 

332 radial_cutoff: float | list[float] 

333 angular_cutoff: float | list[float] 

334 

335 n_basis_radial: int 

336 n_basis_angular: int 

337 n_max_radial: int 

338 n_max_angular: int 

339 l_max_3b: int 

340 l_max_4b: int 

341 l_max_5b: int 

342 n_descriptor_radial: int 

343 n_descriptor_angular: int 

344 

345 n_neuron: int 

346 n_parameters: int 

347 n_descriptor_parameters: int 

348 n_ann_parameters: int 

349 ann_parameters: NetworkWeights 

350 q_scaler: list[float] 

351 radial_descriptor_weights: DescriptorWeights 

352 angular_descriptor_weights: DescriptorWeights 

353 sqrt_epsilon_infinity: float = None 

354 restart_parameters: RestartParameters = None 

355 

356 zbl: tuple[float, float] = None 

357 zbl_typewise_cutoff_factor: float = None 

358 max_neighbors_radial: int = None 

359 max_neighbors_angular: int = None 

360 radial_typewise_cutoff_factor: float = None 

361 angular_typewise_cutoff_factor: float = None 

362 

363 _special_fields = [ 

364 'ann_parameters', 

365 'q_scaler', 

366 'radial_descriptor_weights', 

367 'angular_descriptor_weights', 

368 ] 

369 

370 def __str__(self) -> str: 

371 s = [] 

372 for fld in self.__dataclass_fields__: 

373 if fld not in self._special_fields: 

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

375 return '\n'.join(s) 

376 

377 def _repr_html_(self) -> str: 

378 s = [] 

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

380 s += [ 

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

382 ] 

383 s += ['<tbody>'] 

384 for fld in self.__dataclass_fields__: 

385 if fld not in self._special_fields: 

386 s += [ 

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

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

389 ] 

390 for fld in self._special_fields: 

391 d = getattr(self, fld) 

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

393 if fld.endswith('descriptor_weights'): 

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

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

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

397 else: 

398 dim = len(d) 

399 s += [ 

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

401 ] 

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

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

404 return ''.join(s) 

405 

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

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

408 

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

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

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

412 descriptor weights. It also recalculates the total number of 

413 parameters in the model. 

414 

415 Parameters 

416 ---------- 

417 species 

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

419 

420 Raises 

421 ------ 

422 ValueError 

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

424 """ 

425 for s in species: 

426 if s not in self.types: 

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

428 

429 # --- Prune attributes based on species --- 

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

431 self.types = tuple(types_to_keep) 

432 

433 # Prune ANN parameters (for NEP4 and NEP5) 

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

435 self.ann_parameters = { 

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

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

438 } 

439 

440 # Prune descriptor weights 

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

442 self.radial_descriptor_weights = { 

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

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

445 } 

446 self.angular_descriptor_weights = { 

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

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

449 } 

450 

451 # Prune restart parameters if they have been loaded 

452 if self.restart_parameters is not None: 

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

454 # Prune ANN restart parameters 

455 ann_key = f'ann_{param_type}' 

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

457 self.restart_parameters[ann_key] = { 

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

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

460 } 

461 

462 # Prune descriptor restart parameters 

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

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

465 self.restart_parameters[key] = { 

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

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

468 } 

469 

470 # --- Recalculate parameter counts --- 

471 n_types = len(self.types) 

472 n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular 

473 

474 # Recalculate descriptor parameter count 

475 self.n_descriptor_parameters = n_types**2 * ( 

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

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

478 ) 

479 

480 # Recalculate ANN parameter count 

481 if self.version == 3: 

482 n_networks = 1 

483 n_bias = 1 

484 elif self.version == 4: 

485 n_networks = n_types 

486 n_bias = 1 

487 else: # NEP5 

488 n_networks = n_types 

489 n_bias = 1 + n_types 

490 

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

492 n_ann_output_weights = self.n_neuron 

493 self.n_ann_parameters = ( 

494 n_ann_input_weights + n_ann_output_weights 

495 ) * n_networks + n_bias 

496 

497 # Recalculate total parameter count 

498 self.n_parameters = ( 

499 self.n_ann_parameters 

500 + self.n_descriptor_parameters 

501 + n_descriptor # q_scaler parameters 

502 ) 

503 if self.model_type == 'polarizability': 

504 self.n_parameters += self.n_ann_parameters 

505 

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

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

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

509 # header 

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

511 if self.zbl is not None: 

512 version_name += '_zbl' 

513 elif self.model_type != 'potential': 

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

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

516 if self.zbl is not None: 

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

518 f.write('cutoff') 

519 if isinstance(self.radial_cutoff, float) and isinstance(self.angular_cutoff, float): 

520 f.write(f' {self.radial_cutoff} {self.angular_cutoff}') 

521 else: 

522 # Typewise cutoffs: one set of cutoffs per type 

523 for i in range(len(self.types)): 

524 f.write(f' {self.radial_cutoff[i]} {self.angular_cutoff[i]}') 

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

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 N_types = len(data['types']) 

713 # Either global cutoffs + max neighbirs, or typewise cutoffs + max_neighbors 

714 assert len(data['cutoff']) in [4, 2*N_types+2] 

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

716 data['max_neighbors_angular'] = int(data['cutoff'][-1]) 

717 if len(data['cutoff']) == 2*N_types+2: 

718 # Typewise cutoffs: radial are even, angular are odd 

719 data['radial_cutoff'] = [data['cutoff'][i*2] for i in range(N_types)] 

720 data['angular_cutoff'] = [data['cutoff'][i*2+1] for i in range(N_types)] 

721 else: 

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

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

724 del data['cutoff'] 

725 

726 # split up basis_size tuple 

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

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

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

730 del data['basis_size'] 

731 

732 # split up n_max tuple 

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

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

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

736 del data['n_max'] 

737 

738 # split up nl_max tuple 

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

740 assert len_l in [1, 2, 3] 

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

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

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