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

1import copy

2import re

3from dataclasses import dataclass

4from itertools import product

6import numpy as np

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

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

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

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

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

15 mean and standard deviation for all model parameters.

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

18 Parameters

19 ----------

20 filename

21 input file name

22 """

23 mu = [] # Mean

24 sigma = [] # Standard deviation

25 with open(filename) as f:

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

27 flds = line.split()

28 if len(flds) == 0:

29 raise IOError(f'Empty line number {k}')

30 if len(flds) == 2:

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

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

33 else:

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

35 return mu, sigma

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

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

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

41 `dipole`, and `polarizability`.

43 Parameters

44 ----------

45 first_row

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

47 """

48 model_type = first_row[0]

49 if 'charge' in model_type:

50 return 'potential_with_charges'

51 elif 'dipole' in model_type:

52 return 'dipole'

53 elif 'polarizability' in model_type:

54 return 'polarizability'

55 return 'potential'

58def _get_charge_mode(model_type_token: str) -> int:

59 """Parses the charge_mode (0, 1, or 2) from the first token of a ``nep.txt``

60 header line, e.g. ``nep4_charge1``, ``nep4_zbl_charge2``. Returns 0 for

61 non-charge models.

63 Parameters

64 ----------

65 model_type_token

66 First token of the first row of a NEP file (``flds[0]``).

67 """

68 match = re.search(r'charge(\d+)', model_type_token)

69 return int(match.group(1)) if match else 0

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

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

74 and an unformatted list of all model parameters.

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

77 Parameters

78 ----------

79 filename

80 input file name

81 """

82 # parse file and split header and parameters

83 header = []

84 parameters = []

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

86 base_line = 3

87 with open(filename) as f:

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

89 flds = line.split()

90 if len(flds) == 0:

91 raise IOError(f'Empty line number {k}')

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

93 base_line += 1

94 nheader += 1

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

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

97 nheader += 1

98 if k < nheader:

99 header.append(tuple(flds))

100 elif len(flds) == 1:

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

102 else:

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

104 # compile data from the header into a dict

105 data = {}

106 for flds in header:

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

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

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

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

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

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

113 version = int(version)

114 data['version'] = version

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

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

117 data['charge_mode'] = _get_charge_mode(flds[0])

118 else:

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

120 return data, parameters

121

122

123def _sort_descriptor_parameters(parameters: list[float],

124 types: list[str],

125 n_max_radial: int,

126 n_basis_radial: int,

127 n_max_angular: int,

128 n_basis_angular: int) -> tuple[DescriptorWeights,

129 DescriptorWeights]:

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

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

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

133 """

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

135 n_types = len(types)

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

137

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

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

140 descriptor_weights_radial = descriptor_weights[:, :m]

141 descriptor_weights_angular = descriptor_weights[:, m:]

142

143 # add descriptors to data dict

144 radial_descriptor_weights = {}

145 angular_descriptor_weights = {}

146 m = -1

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

148 m += 1

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

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

151 (n_max_radial + 1, n_basis_radial + 1)

152 )

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

154 (n_max_angular + 1, n_basis_angular + 1)

155 )

156 return radial_descriptor_weights, angular_descriptor_weights

157

158

159def _sort_ann_parameters(parameters: list[float],

160 ann_groupings: list[str],

161 n_neuron: int,

162 n_networks: int,

163 n_bias: int,

164 n_descriptor: int,

165 is_polarizability_model: bool,

166 is_model_with_charges: bool

167 ) -> NetworkWeights:

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

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

170 """

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

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

173 n_ann_parameters = (

174 n_ann_input_weights + n_ann_output_weights

175 ) * n_networks + n_bias

176

177 # Group ANN parameters

178 pars = {}

179 n1 = 0

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

181

182 n_count = 2 if is_polarizability_model else 1

183 n_outputs = 2 if is_model_with_charges else 1

184 for count in range(n_count):

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

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

187 n1 += count

188 for s in ann_groupings:

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

190 # one network per atomic species.

191 ann_parameters = parameters[n1 : n1 + n_network_params]

192 ann_input_weights = ann_parameters[:n_ann_input_weights]

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

194 w0[...] = np.nan

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

196 b0[...] = np.nan

197 for n in range(n_neuron):

198 for nu in range(n_descriptor):

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

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

201

202 assert np.all(

203 w0.shape == (n_neuron, n_descriptor)

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

205 assert np.all(

206 b0.shape == (n_neuron, 1)

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

208 assert not np.any(

209 np.isnan(w0)

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

211 assert not np.any(

212 np.isnan(b0)

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

214

215 ann_output_weights = ann_parameters[

216 n_ann_input_weights : n_ann_input_weights + n_ann_output_weights

217 ]

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

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

220 assert np.all(

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

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

223 assert not np.any(

224 np.isnan(w1)

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

226

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

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

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

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

231 else:

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

233 # Jump to bias

234 n1 += n_network_params

235 if n_bias > 1 and not is_model_with_charges:

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

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

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

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

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

241 n1 += 1

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

243 # For NEP4 with charges we have two biases.

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

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

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

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

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

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

250 else:

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

252 sum = 0

253 for s in pars.keys():

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

255 sum += 1

256 else:

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

258 assert sum == n_ann_parameters * n_count, (

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

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

261 )

262 return pars

263

264

265def _adaptive_sigma(mu_arr, sigma_factor: float, sigma_floor: float) -> np.ndarray:

266 """Return adaptive SNES sigma: ``max(sigma_floor, sigma_factor * |mu|)``."""

267 return np.maximum(sigma_floor, sigma_factor * np.abs(mu_arr))

268

269

270_RESTART_COMPONENTS = ('network_weights', 'descriptor', 'charge_head')

271

272

273def _restart_leaves(model, restart_params, component=None, species=None):

274 """Yield ``(mu, sigma_container, sigma_key)`` for every leaf entry of

275 *restart_params* that matches the requested *component*/*species* filters.

276

277 ``sigma_container[sigma_key]`` is either a numpy array or a scalar float;

278 together with ``mu`` (same shape/type) this is everything

279 :func:`_apply_sigma_strategy` needs to read and update one leaf.

280

281 *component* selects among ``'network_weights'`` (w0, b0, w1, and the global

282 b1 bias), ``'descriptor'`` (radial/angular descriptor weight pairs), and

283 ``'charge_head'`` (w1_charge and sqrt_epsilon_infinity). ``None`` means all

284 three. *species* restricts per-species entries (and descriptor pairs

285 involving that species) to the given species; global scalars (b1,

286 sqrt_epsilon_infinity) are only included when *species* is ``None``, since

287 they are not owned by a single species.

288 """

289 if component is None:

290 wanted = set(_RESTART_COMPONENTS)

291 else:

292 wanted = {component} if isinstance(component, str) else set(component)

293 unknown = wanted - set(_RESTART_COMPONENTS)

294 if unknown:

295 raise ValueError(

296 f'Unknown component(s) {sorted(unknown)}; expected any of '

297 f'{_RESTART_COMPONENTS}'

298 )

299

300 if species is None:

301 species_filter = None

302 else:

303 species_filter = {species} if isinstance(species, str) else set(species)

304

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

306 ann_mu, ann_sigma = restart_params['ann_mu'], restart_params['ann_sigma']

307

308 if 'network_weights' in wanted:

309 for s in keys:

310 if species_filter is not None and s not in species_filter:

311 continue

312 for pname in ('w0', 'b0', 'w1', 'w0_polar', 'b0_polar', 'w1_polar'):

313 if pname in ann_mu[s]:

314 yield ann_mu[s][pname], ann_sigma[s], pname

315 if species_filter is None:

316 for pname in ('b1', 'b1_polar'):

317 if pname in ann_mu:

318 yield ann_mu[pname], ann_sigma, pname

319

320 if 'charge_head' in wanted:

321 for s in keys:

322 if species_filter is not None and s not in species_filter:

323 continue

324 if 'w1_charge' in ann_mu[s]:

325 yield ann_mu[s]['w1_charge'], ann_sigma[s], 'w1_charge'

326 if species_filter is None and 'sqrt_epsilon_infinity' in ann_mu:

327 yield ann_mu['sqrt_epsilon_infinity'], ann_sigma, 'sqrt_epsilon_infinity'

328

329 if 'descriptor' in wanted:

330 for desc_type in ('radial', 'angular'):

331 mu_dict = restart_params[f'{desc_type}_descriptor_mu']

332 sigma_dict = restart_params[f'{desc_type}_descriptor_sigma']

333 for pair, mu_val in mu_dict.items():

334 if species_filter is not None and not (species_filter & set(pair)):

335 continue

336 yield mu_val, sigma_dict, pair

337

338

339def _apply_sigma_strategy(mu, sigma_container, sigma_key, strategy, target, rng, **kw):

340 """Update ``sigma_container[sigma_key]`` in place at the positions selected

341 by *target* (``'unset'`` -> NaN entries, ``'set'`` -> non-NaN entries,

342 ``'all'`` -> everything), using *mu* and *strategy* to compute new values.

343 """

344 sigma = sigma_container[sigma_key]

345 if np.isscalar(sigma) or isinstance(sigma, (float, int)):

346 current = float(sigma)

347 is_unset = np.isnan(current)

348 apply_here = (

349 target == 'all' or (target == 'unset' and is_unset)

350 or (target == 'set' and not is_unset)

351 )

352 if not apply_here:

353 return

354 mu_val = float(mu)

355 if strategy == 'constant':

356 new_val = kw['value']

357 elif strategy == 'scale_mu':

358 new_val = float(_adaptive_sigma(np.array(mu_val), kw['factor'], kw['floor']))

359 elif strategy == 'scale_sigma':

360 if is_unset: 360 ↛ 361line 360 didn't jump to line 361 because the condition on line 360 was never true

361 raise ValueError(

362 f"strategy='scale_sigma' requires an existing sigma value, but "

363 f'{sigma_key!r} is unset (NaN); set it first, e.g. with '

364 "target='unset'."

365 )

366 new_val = current * kw['factor']

367 elif strategy == 'uniform':

368 new_val = float(rng.uniform(kw['low'], kw['high']))

369 elif strategy == 'normal': 369 ↛ 372line 369 didn't jump to line 372 because the condition on line 369 was always true

370 new_val = float(abs(rng.normal(kw['mean'], kw['std'])))

371 else:

372 raise ValueError(f'Unknown strategy {strategy!r}')

373 sigma_container[sigma_key] = float(new_val)

374 return

375

376 sigma_arr = sigma_container[sigma_key]

377 mu_arr = np.asarray(mu, dtype=float)

378 if target == 'unset':

379 mask = np.isnan(sigma_arr)

380 elif target == 'set':

381 mask = ~np.isnan(sigma_arr)

382 else:

383 mask = np.ones_like(sigma_arr, dtype=bool)

384 if not np.any(mask):

385 return

386

387 if strategy == 'constant':

388 sigma_arr[mask] = kw['value']

389 elif strategy == 'scale_mu':

390 sigma_arr[mask] = _adaptive_sigma(mu_arr[mask], kw['factor'], kw['floor'])

391 elif strategy == 'scale_sigma':

392 if np.any(np.isnan(sigma_arr[mask])):

393 raise ValueError(

394 f"strategy='scale_sigma' requires existing sigma values, but "

395 f'{sigma_key!r} has unset (NaN) entries within the selected '

396 "target; set them first, e.g. with target='unset'."

397 )

398 sigma_arr[mask] = sigma_arr[mask] * kw['factor']

399 elif strategy == 'uniform':

400 sigma_arr[mask] = rng.uniform(kw['low'], kw['high'], size=int(np.sum(mask)))

401 elif strategy == 'normal': 401 ↛ 404line 401 didn't jump to line 404 because the condition on line 401 was always true

402 sigma_arr[mask] = np.abs(rng.normal(kw['mean'], kw['std'], size=int(np.sum(mask))))

403 else:

404 raise ValueError(f'Unknown strategy {strategy!r}')

405

406

407def _new_restart_parameters_from_model(model) -> RestartParameters:

408 """Build a fresh restart-parameters dict from a model's current (trained)

409 parameters: ``mu`` is copied from the model, ``sigma`` is set to ``NaN``

410 everywhere (unset), to be filled in via :meth:`Model.set_restart_sigma`.

411 """

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

413 suffixes = ['', '_polar'] if model.model_type == 'polarizability' else ['']

414

415 ann_mu, ann_sigma = {}, {}

416 for s in keys:

417 params = model.ann_parameters[s]

418 mu_entry, sigma_entry = {}, {}

419 for suffix in suffixes:

420 for base in ('w0', 'b0', 'w1', 'w1_charge'):

421 pname = f'{base}{suffix}'

422 if pname in params:

423 arr = np.array(params[pname], dtype=float)

424 mu_entry[pname] = arr.copy()

425 sigma_entry[pname] = np.full(arr.shape, np.nan)

426 ann_mu[s] = mu_entry

427 ann_sigma[s] = sigma_entry

428

429 for suffix in suffixes:

430 b1_key = f'b1{suffix}'

431 if b1_key in model.ann_parameters: 431 ↛ 429line 431 didn't jump to line 429 because the condition on line 431 was always true

432 ann_mu[b1_key] = float(model.ann_parameters[b1_key])

433 ann_sigma[b1_key] = float('nan')

434 if model.sqrt_epsilon_infinity is not None: 434 ↛ 435line 434 didn't jump to line 435 because the condition on line 434 was never true

435 ann_mu['sqrt_epsilon_infinity'] = float(model.sqrt_epsilon_infinity)

436 ann_sigma['sqrt_epsilon_infinity'] = float('nan')

437

438 radial_mu = {

439 k: np.array(v, dtype=float).copy() for k, v in model.radial_descriptor_weights.items()

440 }

441 radial_sigma = {k: np.full(v.shape, np.nan) for k, v in radial_mu.items()}

442 angular_mu = {

443 k: np.array(v, dtype=float).copy() for k, v in model.angular_descriptor_weights.items()

444 }

445 angular_sigma = {k: np.full(v.shape, np.nan) for k, v in angular_mu.items()}

446

447 return {

448 'ann_mu': ann_mu,

449 'ann_sigma': ann_sigma,

450 'radial_descriptor_mu': radial_mu,

451 'radial_descriptor_sigma': radial_sigma,

452 'angular_descriptor_mu': angular_mu,

453 'angular_descriptor_sigma': angular_sigma,

454 }

455

456

457def _recalculate_parameter_counts(new) -> None:

458 """Recompute n_ann_parameters, n_descriptor_parameters, and n_parameters on *new*.

459

460 Reads all architectural state from *new* directly, so callers must update

461 new.n_neuron, new.n_descriptor_radial/angular, new.model_type, and new.types

462 before calling this function.

463 """

464 n_types = len(new.types)

465 n_desc = new.n_descriptor_radial + new.n_descriptor_angular

466 is_charged = new.model_type == 'potential_with_charges'

467 n_networks = n_types if new.version in (4, 5) else 1

468 n_bias = 2 if is_charged else (1 + n_types if new.version == 5 else 1)

469 n_ann_input_weights = (n_desc + 1) * new.n_neuron

470 n_ann_output_weights = 2 * new.n_neuron if is_charged else new.n_neuron

471 new.n_ann_parameters = (n_ann_input_weights + n_ann_output_weights) * n_networks + n_bias

472 new.n_descriptor_parameters = n_types ** 2 * (

473 (new.n_max_radial + 1) * (new.n_basis_radial + 1)

474 + (new.n_max_angular + 1) * (new.n_basis_angular + 1)

475 )

476 new.n_parameters = new.n_ann_parameters + new.n_descriptor_parameters + n_desc

477 if new.model_type == 'polarizability':

478 new.n_parameters += new.n_ann_parameters

479

480

481@dataclass

482class Model:

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

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

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

486

487 Attributes

488 ----------

489 version : int

490 NEP version.

491 model_type: str

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

493 types : tuple[str, ...]

494 Chemical species that this model represents.

495 radial_cutoff : float | list[float]

496 The radial cutoff parameter in Å.

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

498 angular_cutoff : float | list[float]

499 The angular cutoff parameter in Å.

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

501 max_neighbors_radial : int

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

503 max_neighbors_angular : int

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

505 zbl : tuple[float, float]

506 Inner and outer cutoff for transition to ZBL potential.

507 zbl_typewise_cutoff_factor : float

508 Optional typewise cutoff factor for the ZBL potential, corresponding to an

509 optional third value on the ``zbl`` line in ``nep.txt`` when

510 ``use_typewise_cutoff_zbl`` is enabled during training. ``None`` if not set.

511 n_basis_radial : int

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

513 n_basis_angular : int

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

515 n_max_radial : int

516 Maximum order of Chebyshev polymonials included in

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

518 n_max_angular : int

519 Maximum order of Chebyshev polymonials included in

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

521 l_max_3b : int

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

523 l_max_4b : int

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

525 l_max_5b : int

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

527 has_q_112 : int

528 Flag enabling the 5-body :math:`q_{112}` descriptor (0 or 1).

529 has_q_123 : int

530 Flag enabling the 5-body :math:`q_{123}` descriptor (0 or 1).

531 has_q_233 : int

532 Flag enabling the 5-body :math:`q_{233}` descriptor (0 or 1).

533 has_q_134 : int

534 Flag enabling the higher-body :math:`q_{134}` descriptor (0 or 1).

535 n_descriptor_radial : int

536 Dimension of radial part of descriptor.

537 n_descriptor_angular : int

538 Dimension of angular part of descriptor.

539 n_neuron : int

540 Number of neurons in hidden layer.

541 n_parameters : int

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

543 n_descriptor_parameters : int

544 Number of parameters in descriptor.

545 n_ann_parameters : int

546 Number of neural network weights.

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

548 Neural network weights.

549 q_scaler : List[float]

550 Scaling parameters.

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

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

553 has dimensions of

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

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

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

557 has dimensions of

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

559 sqrt_epsilon_infinity : Optional[float]

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

561 charge_mode : int

562 Charge algorithm variant for ``potential_with_charges`` models; 0 for

563 non-charge-aware models. 1 corresponds to a qNEP model including both real- and

564 reciprocal-space contributions. 2 corresponds to a qNEP model, including the

565 reciprocal-space contribution only.

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

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

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

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

570 """

571

572 version: int

573 model_type: str

574 types: tuple[str, ...]

575

576 radial_cutoff: float | list[float]

577 angular_cutoff: float | list[float]

578

579 n_basis_radial: int

580 n_basis_angular: int

581 n_max_radial: int

582 n_max_angular: int

583 l_max_3b: int

584 l_max_4b: int

585 l_max_5b: int

586 has_q_112: int

587 has_q_123: int

588 has_q_233: int

589 has_q_134: int

590 n_descriptor_radial: int

591 n_descriptor_angular: int

592

593 n_neuron: int

594 n_parameters: int

595 n_descriptor_parameters: int

596 n_ann_parameters: int

597 ann_parameters: NetworkWeights

598 q_scaler: list[float]

599 radial_descriptor_weights: DescriptorWeights

600 angular_descriptor_weights: DescriptorWeights

601 sqrt_epsilon_infinity: float = None

602 charge_mode: int = 0

603 restart_parameters: RestartParameters = None

604

605 zbl: tuple[float, float] = None

606 zbl_typewise_cutoff_factor: float = None

607 max_neighbors_radial: int = None

608 max_neighbors_angular: int = None

609

610 _special_fields = [

611 'ann_parameters',

612 'q_scaler',

613 'radial_descriptor_weights',

614 'angular_descriptor_weights',

615 ]

616

617 def __str__(self) -> str:

618 s = []

619 for fld in self.__dataclass_fields__:

620 if fld not in self._special_fields:

621 value = getattr(self, fld)

622 if fld == 'restart_parameters':

623 value = 'available' if value is not None else 'not available'

624 s += [f'{fld:22} : {value}']

625 return '\n'.join(s)

626

627 def _repr_html_(self) -> str:

628 s = []

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

630 s += [

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

632 ]

633 s += ['<tbody>']

634 for fld in self.__dataclass_fields__:

635 if fld not in self._special_fields:

636 value = getattr(self, fld)

637 if fld == 'restart_parameters':

638 value = 'available' if value is not None else 'not available'

639 s += [

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

641 f'<td>{value}</td><tr>'

642 ]

643 for fld in self._special_fields:

644 d = getattr(self, fld)

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

646 if fld.endswith('descriptor_weights'):

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

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

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

650 else:

651 dim = len(d)

652 s += [

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

654 ]

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

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

657 return ''.join(s)

658

659 @property

660 def training_parameters(self) -> dict:

661 """Return model hyperparameters in the format accepted by :func:`write_nepfile

662 <calorine.nep.write_nepfile>`.

663

664 Use this after any model modification (:meth:`augment`, :meth:`add_species`,

665 :meth:`remove_species`, :meth:`keep_species`) to produce the architecture fields

666 that must go into the new ``nep.in`` before training. Merge the result with your

667 existing training-specific parameters (``lambda_*``, ``generation``, ``batch``,

668 etc.) before calling :func:`write_nepfile <calorine.nep.write_nepfile>`.

669

670 Returns

671 -------

672 dict

673 Keys ``version``, ``type``, ``cutoff``, ``n_max``, ``basis_size``, ``l_max``,

674 and ``neuron`` (plus ``zbl`` when applicable) with values in the format

675 expected by :func:`write_nepfile <calorine.nep.write_nepfile>`. ``zbl`` is the

676 single outer cutoff value the ``nep.in`` ``zbl`` keyword expects (the inner

677 cutoff is always half of it), not the ``(inner, outer)`` pair stored in

678 :attr:`zbl`.

679

680 """

681 l_max = [self.l_max_3b, self.l_max_4b, self.l_max_5b,

682 self.has_q_112, self.has_q_123, self.has_q_233, self.has_q_134]

683 while len(l_max) > 1 and l_max[-1] == 0:

684 l_max = l_max[:-1]

685

686 if isinstance(self.radial_cutoff, list):

687 cutoff = []

688 for rc, ac in zip(self.radial_cutoff, self.angular_cutoff):

689 cutoff += [rc, ac]

690 else:

691 cutoff = [self.radial_cutoff, self.angular_cutoff]

692

693 params = {

694 'version': self.version,

695 'type': [len(self.types)] + list(self.types),

696 'cutoff': cutoff,

697 'n_max': [self.n_max_radial, self.n_max_angular],

698 'basis_size': [self.n_basis_radial, self.n_basis_angular],

699 'l_max': l_max,

700 'neuron': self.n_neuron,

701 }

702 if self.zbl is not None:

703 params['zbl'] = self.zbl[1]

704 return params

705

706 def remove_species(self, species: list[str]) -> 'Model':

707 """Remove one or more species from the model.

708

709 Returns a new :class:`Model` with the specified species removed.

710 The source model is not modified.

711

712 If ``restart_parameters`` are loaded, they are pruned to match (the

713 entries for the removed species/pairs are dropped); the surviving

714 entries are left exactly as they were. Use :meth:`set_restart_sigma`

715 explicitly afterwards if you want to re-open the SNES search width for

716 the surviving parameters before continuing training.

717

718 Parameters

719 ----------

720 species

721 Species names to remove.

722

723 Returns

724 -------

725 Model

726 New model with the specified species removed.

727

728 Raises

729 ------

730 ValueError

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

732 """

733 for s in species:

734 if s not in self.types:

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

736

737 new = copy.deepcopy(self)

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

739 new.types = tuple(types_to_keep)

740

741 # Prune ANN parameters (for NEP4 and NEP5)

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

743 new.ann_parameters = {

744 key: value for key, value in new.ann_parameters.items()

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

746 }

747

748 # Prune descriptor weights; key is a (species1, species2) tuple

749 new.radial_descriptor_weights = {

750 key: value for key, value in new.radial_descriptor_weights.items()

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

752 }

753 new.angular_descriptor_weights = {

754 key: value for key, value in new.angular_descriptor_weights.items()

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

756 }

757

758 # Prune typewise cutoff lists so remaining species map to correct cutoffs

759 if isinstance(self.radial_cutoff, list):

760 indices = [i for i, t in enumerate(self.types) if t not in species]

761 new.radial_cutoff = [self.radial_cutoff[i] for i in indices]

762 new.angular_cutoff = [self.angular_cutoff[i] for i in indices]

763

764 # Prune restart parameters to match; survivors are left untouched

765 if new.restart_parameters is not None:

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

767 ann_key = f'ann_{param_type}'

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

769 # Keep per-species keys for survivors, global bias keys, and

770 # sqrt_epsilon_infinity (charge models)

771 new.restart_parameters[ann_key] = {

772 key: value for key, value in new.restart_parameters[ann_key].items()

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

774 or key == 'sqrt_epsilon_infinity')

775 }

776

777 # Prune descriptor restart parameters

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

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

780 new.restart_parameters[key] = {

781 k: v for k, v in new.restart_parameters[key].items()

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

783 }

784

785 # Recalculate parameter counts

786 _recalculate_parameter_counts(new)

787

788 return new

789

790 def keep_species(self, species: list[str]) -> 'Model':

791 """Retain only the specified species, removing all others.

792

793 Convenience complement to :meth:`remove_species`. Useful when the set

794 of species to drop is large (e.g. isolating two elements from a

795 foundation model with dozens of species).

796

797 Parameters

798 ----------

799 species

800 Species names to keep. All other species are removed.

801

802 Returns

803 -------

804 Model

805 New model containing only the requested species.

806

807 Raises

808 ------

809 ValueError

810 If any of the requested species is not in the model.

811 """

812 unknown = [s for s in species if s not in self.types]

813 if unknown:

814 raise ValueError(

815 f'Species not in model: {unknown}'

816 )

817 to_remove = [s for s in self.types if s not in species]

818 return self.remove_species(to_remove)

819

820 def reorder(self, order: list[str]) -> 'Model':

821 """Reorder the species in the model.

822

823 Returns a new :class:`Model` with species permuted according to

824 ``order``. This is useful for aligning the species order of two

825 models that must share the same order when used jointly by GPUMD,

826 e.g. a NEP potential and a TNEP dipole/polarizability model

827 referenced together via two ``potential`` lines in ``run.in`` and

828 ``dump_dipole`` or ``dump_polarizability``.

829

830 The source model is not modified. Since ``ann_parameters``,

831 ``radial_descriptor_weights``, ``angular_descriptor_weights``, and

832 ``restart_parameters`` are keyed by species name (or species-pair)

833 rather than position, reordering only requires updating ``types``

834 and, if typewise cutoffs are in use, the positional

835 ``radial_cutoff`` and ``angular_cutoff`` lists.

836

837 Parameters

838 ----------

839 order

840 New species order. Must be a permutation of ``self.types``.

841

842 Returns

843 -------

844 Model

845 New model with species reordered.

846

847 Raises

848 ------

849 ValueError

850 If ``order`` is not a permutation of the current species.

851 """

852 if sorted(order) != sorted(self.types):

853 raise ValueError(

854 f'order must be a permutation of the current species {self.types}, '

855 f'got {list(order)}'

856 )

857

858 new = copy.deepcopy(self)

859 new.types = tuple(order)

860

861 if isinstance(self.radial_cutoff, list):

862 indices = [self.types.index(t) for t in order]

863 new.radial_cutoff = [self.radial_cutoff[i] for i in indices]

864 new.angular_cutoff = [self.angular_cutoff[i] for i in indices]

865

866 return new

867

868 def add_species(self,

869 species: list[str],

870 radial_cutoff: float | list[float] = None,

871 angular_cutoff: float | list[float] = None,

872 seed: int | None = None) -> 'Model':

873 """Add one or more species to the model.

874

875 Returns a new :class:`Model` with the requested species added. New ANN

876 sub-networks and descriptor weight pairs are initialised by drawing

877 ``mu`` uniformly from [-1, 1] (matching the GPUMD fresh-model

878 initialisation); the corresponding restart sigma entries are left

879 unset (``NaN``) — call :meth:`set_restart_sigma` afterwards to

880 initialize them (e.g. ``model.add_species(['X']).set_restart_sigma()``

881 fills only the new entries by default). Charge-specific parameters

882 (``w1_charge``) are kept at ``mu = 0`` to preserve stability, also

883 matching GPUMD. Existing parameters (``mu`` and ``sigma``) are left

884 untouched.

885

886 Only supported for NEP4 models. For NEP3 the ANN is shared across all

887 species and adding a per-species sub-network is not meaningful.

888

889 Parameters

890 ----------

891 species

892 New species names to add. Appended to ``types`` in the order given.

893 radial_cutoff

894 Radial cutoff(s) for the new species, in Å. Required when the model

895 uses typewise cutoffs (i.e. ``isinstance(model.radial_cutoff, list)``

896 is ``True``). Pass a single float or a list with one value per new

897 species.

898 angular_cutoff

899 Angular cutoff(s) for the new species, in Å. Same requirements as

900 ``radial_cutoff``.

901 seed

902 Seed for the random number generator used to draw the initial ``mu``

903 values. Pass an integer for reproducible initialisation.

904

905 Returns

906 -------

907 Model

908 New model with updated structure, weights, and restart statistics.

909

910 Raises

911 ------

912 ValueError

913 If the model version is not 4, if ``restart_parameters`` are not

914 loaded, if any species is already in the model, or if typewise

915 cutoffs are used and ``radial_cutoff``/``angular_cutoff`` are not

916 provided.

917 """

918 if self.version != 4:

919 raise ValueError(

920 f'add_species() only supports NEP4 models; got version {self.version}.'

921 )

922 for s in species:

923 if s in self.types:

924 raise ValueError(f'{s!r} is already in the model.')

925 if self.restart_parameters is None:

926 raise ValueError(

927 'restart_parameters must be loaded before calling add_species(). '

928 'Pass restart_file= to read_model() or call model.read_restart() first.'

929 )

930

931 uses_typewise = isinstance(self.radial_cutoff, list)

932 if uses_typewise:

933 if radial_cutoff is None or angular_cutoff is None:

934 raise ValueError(

935 'Model uses typewise cutoffs; provide radial_cutoff and angular_cutoff '

936 'for the new species.'

937 )

938 rc_list = ([radial_cutoff] * len(species)

939 if isinstance(radial_cutoff, (int, float)) else list(radial_cutoff))

940 ac_list = ([angular_cutoff] * len(species)

941 if isinstance(angular_cutoff, (int, float)) else list(angular_cutoff))

942 if len(rc_list) != len(species) or len(ac_list) != len(species):

943 raise ValueError(

944 'Length of radial_cutoff/angular_cutoff must match the number of new species.'

945 )

946

947 new = copy.deepcopy(self)

948

949 n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular

950 n_neuron = self.n_neuron

951 is_charged = self.model_type == 'potential_with_charges'

952 all_types_after = list(self.types) + list(species)

953 rng = np.random.default_rng(seed)

954

955 def _rand(shape):

956 return rng.uniform(-1.0, 1.0, size=shape)

957

958 # Step 1: New ANN sub-networks

959 w1_shape = (n_neuron,) if is_charged else (1, n_neuron)

960 for s_new in species:

961 w0_vals = _rand((n_neuron, n_descriptor))

962 b0_vals = _rand((n_neuron, 1))

963 w1_vals = _rand(w1_shape)

964 s_params = {'w0': w0_vals.copy(), 'b0': b0_vals.copy(), 'w1': w1_vals.copy()}

965 if is_charged:

966 s_params['w1_charge'] = np.zeros(n_neuron)

967 new.ann_parameters[s_new] = s_params

968

969 mu_entry = {'w0': w0_vals, 'b0': b0_vals, 'w1': w1_vals}

970 sigma_entry = {

971 'w0': np.full((n_neuron, n_descriptor), np.nan),

972 'b0': np.full((n_neuron, 1), np.nan),

973 'w1': np.full(w1_shape, np.nan),

974 }

975 if is_charged:

976 mu_entry['w1_charge'] = np.zeros(n_neuron)

977 sigma_entry['w1_charge'] = np.full(n_neuron, np.nan)

978 new.restart_parameters['ann_mu'][s_new] = mu_entry

979 new.restart_parameters['ann_sigma'][s_new] = sigma_entry

980

981 # Step 2: New descriptor weight pairs

982 n_r = (self.n_max_radial + 1, self.n_basis_radial + 1)

983 n_a = (self.n_max_angular + 1, self.n_basis_angular + 1)

984 existing_pairs = set(self.radial_descriptor_weights)

985 new_pairs = {

986 (s1, s2)

987 for s1 in all_types_after for s2 in all_types_after

988 if (s1, s2) not in existing_pairs

989 }

990 for pair in new_pairs:

991 r_vals = _rand(n_r)

992 a_vals = _rand(n_a)

993 new.radial_descriptor_weights[pair] = r_vals.copy()

994 new.angular_descriptor_weights[pair] = a_vals.copy()

995 new.restart_parameters['radial_descriptor_mu'][pair] = r_vals

996 new.restart_parameters['angular_descriptor_mu'][pair] = a_vals

997 new.restart_parameters['radial_descriptor_sigma'][pair] = np.full(n_r, np.nan)

998 new.restart_parameters['angular_descriptor_sigma'][pair] = np.full(n_a, np.nan)

999

1000 # Step 3: Update types and typewise cutoffs

1001 new.types = tuple(all_types_after)

1002 if uses_typewise:

1003 new.radial_cutoff = list(self.radial_cutoff) + rc_list

1004 new.angular_cutoff = list(self.angular_cutoff) + ac_list

1005

1006 # Step 4: Recalculate parameter counts

1007 _recalculate_parameter_counts(new)

1008

1009 return new

1010

1011 def write(self, filename: str, restart_file: str = None) -> None:

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

1013

1014 Parameters

1015 ----------

1016 filename

1017 Output file name for the NEP model.

1018 restart_file

1019 If provided, also write restart parameters to this file in

1020 `nep.restart` format. Defaults to None.

1021 """

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

1023 # header

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

1025 if self.zbl is not None:

1026 version_name += '_zbl'

1027 if self.model_type == 'potential_with_charges':

1028 version_name += f'_charge{self.charge_mode}'

1029 elif self.model_type != 'potential':

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

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

1032 if self.zbl is not None:

1033 zbl_tokens = list(self.zbl)

1034 if self.zbl_typewise_cutoff_factor is not None:

1035 zbl_tokens.append(self.zbl_typewise_cutoff_factor)

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

1037 f.write('cutoff')

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

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

1040 else:

1041 # Typewise cutoffs: one set of cutoffs per type

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

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

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

1045 f.write('\n')

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

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

1048 l_max_line = f'l_max {self.l_max_3b} {self.l_max_4b} {self.l_max_5b}'

1049 if self.has_q_112 or self.has_q_123 or self.has_q_233 or self.has_q_134:

1050 l_max_line += f' {self.has_q_112}'

1051 if self.has_q_123 or self.has_q_233 or self.has_q_134:

1052 l_max_line += f' {self.has_q_123}'

1053 if self.has_q_233 or self.has_q_134:

1054 l_max_line += f' {self.has_q_233}'

1055 if self.has_q_134:

1056 l_max_line += f' {self.has_q_134}'

1057 f.write(l_max_line + '\n')

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

1059

1060 # neural network weights

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

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

1063 for suffix in suffixes:

1064 for s in keys:

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

1066 # w0 indexed as: n*N_descriptor + nu

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

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

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

1070 for n in range(self.n_neuron):

1071 for nu in range(

1072 self.n_descriptor_radial + self.n_descriptor_angular

1073 ):

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

1075 for b in b0[:, 0]:

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

1077 for v in (w1[0, :] if w1.ndim == 2 else w1):

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

1079 if f'w1_charge{suffix}' in self.ann_parameters[s]:

1080 for v in self.ann_parameters[s][f'w1_charge{suffix}']:

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

1082 if self.version == 5:

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

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

1085 if self.sqrt_epsilon_infinity is not None:

1086 f.write(f'{self.sqrt_epsilon_infinity:15.7e}\n')

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

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

1089

1090 # descriptor weights

1091 mat = []

1092 for s1 in self.types:

1093 for s2 in self.types:

1094 mat = np.hstack(

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

1096 )

1097 mat = np.hstack(

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

1099 )

1100 n_types = len(self.types)

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

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

1103 for v in mat.flatten():

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

1105

1106 # scaler

1107 for v in self.q_scaler:

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

1109

1110 if restart_file is not None:

1111 self.write_restart(restart_file)

1112

1113 def read_restart(self, filename: str):

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

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

1116 parameter in the corresponding NEP model.

1117

1118 Parameters

1119 ----------

1120 filename

1121 Input file name.

1122 """

1123 mu, sigma = _get_restart_contents(filename)

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

1125

1126 is_polarizability_model = self.model_type == 'polarizability'

1127 is_charged_model = self.model_type == 'potential_with_charges'

1128

1129 n1 = self.n_ann_parameters

1130 n1 *= 2 if is_polarizability_model else 1

1131 n2 = n1 + self.n_descriptor_parameters

1132 ann_parameters = restart_parameters[:n1]

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

1134

1135 if self.version == 3:

1136 n_networks = 1

1137 n_bias = 1

1138 elif self.version == 4:

1139 # one hidden layer per atomic species

1140 n_networks = len(self.types)

1141 n_bias = 1

1142 else:

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

1144

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

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

1147 if self.sqrt_epsilon_infinity is not None:

1148 n_bias += 1 # charge models have sqrt_epsilon_infinity before b1

1149 n_descriptor = self.n_descriptor_radial + self.n_descriptor_angular

1150 restart = {}

1151

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

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

1154 ann_groups,

1155 self.n_neuron,

1156 n_networks,

1157 n_bias,

1158 n_descriptor,

1159 is_polarizability_model,

1160 is_charged_model)

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

1162 self.types,

1163 self.n_max_radial,

1164 self.n_basis_radial,

1165 self.n_max_angular,

1166 self.n_basis_angular)

1167

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

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

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

1171 self.restart_parameters = restart

1172

1173 def write_restart(self, filename: str):

1174 """Write the restart parameters to file in `nep.restart` format.

1175

1176 Parameters

1177 ----------

1178 filename

1179 Output file name.

1180

1181 Raises

1182 ------

1183 ValueError

1184 If ``restart_parameters`` is not loaded, or if any restart sigma

1185 value is unset (``NaN``), e.g. because :meth:`add_species` or

1186 :meth:`augment` were called without a follow-up

1187 :meth:`set_restart_sigma` to initialize the sigma of the newly

1188 created parameters.

1189 """

1190 if self.restart_parameters is None:

1191 raise ValueError(

1192 'restart_parameters is not loaded; nothing to write. Pass restart_file= '

1193 'to read_model(), call Model.read_restart(), or call '

1194 'Model.set_restart_sigma() to bootstrap one from the current model '

1195 'parameters before write_restart().'

1196 )

1197 for _, sigma_container, sigma_key in _restart_leaves(self, self.restart_parameters):

1198 if np.any(np.isnan(np.asarray(sigma_container[sigma_key], dtype=float))):

1199 raise ValueError(

1200 f'restart_parameters contains an unset (NaN) sigma value for '

1201 f'{sigma_key!r}. Call Model.set_restart_sigma() to initialize it '

1202 'before write_restart().'

1203 )

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

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

1206 columns = []

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

1208 # neural network weights

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

1210 column = []

1211 for suffix in suffixes:

1212 for s in keys:

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

1214 # w0 indexed as: n*N_descriptor + nu

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

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

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

1218 for n in range(self.n_neuron):

1219 for nu in range(

1220 self.n_descriptor_radial + self.n_descriptor_angular

1221 ):

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

1223 for b in b0[:, 0]:

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

1225 for v in (w1[0, :] if w1.ndim == 2 else w1):

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

1227 if f'w1_charge{suffix}' in ann_parameters[s]:

1228 for v in ann_parameters[s][f'w1_charge{suffix}']:

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

1230 if f'sqrt_epsilon_infinity{suffix}' in ann_parameters:

1231 column.append(f'{ann_parameters[f"sqrt_epsilon_infinity{suffix}"]:15.7e}')

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

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

1234 columns.append(column)

1235

1236 # descriptor weights

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

1238 angular_descriptor_parameters = self.restart_parameters[

1239 f'angular_descriptor_{parameter}']

1240 mat = []

1241 for s1 in self.types:

1242 for s2 in self.types:

1243 mat = np.hstack(

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

1245 )

1246 mat = np.hstack(

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

1248 )

1249 n_types = len(self.types)

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

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

1252 for v in mat.flatten():

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

1254

1255 # Join the mean and standard deviation columns

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

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

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

1259 f.writelines(joined)

1260

1261 def set_restart_sigma(self,

1262 strategy: str = 'scale_mu',

1263 *,

1264 value: float = None,

1265 factor: float = None,

1266 floor: float = 1e-6,

1267 low: float = None,

1268 high: float = None,

1269 mean: float = None,

1270 std: float = None,

1271 species: str | list[str] = None,

1272 component: str | list[str] = None,

1273 target: str = 'unset',

1274 seed: int | None = None) -> 'Model':

1275 """Assign SNES restart sigma values.

1276

1277 Returns a new :class:`Model` with sigma values updated according to

1278 *strategy*, at the positions selected by *target*/*species*/*component*.

1279 ``mu`` and every other field are left unchanged. This is the only

1280 method that ever assigns sigma values; the structural methods

1281 (:meth:`remove_species`, :meth:`keep_species`, :meth:`add_species`,

1282 :meth:`augment`, :meth:`prune`) leave survivors' sigma untouched and

1283 mark newly created parameters' sigma as unset (``NaN``) rather than

1284 computing a value inline.

1285

1286 If ``restart_parameters`` is not loaded, it is created first: ``mu`` is

1287 copied from the model's current (trained) parameters, and every sigma

1288 is initialized as unset (``NaN``). This makes it possible to bootstrap

1289 a ``nep.restart`` file "from scratch" for a plain ``nep.txt`` model.

1290

1291 Parameters

1292 ----------

1293 strategy

1294 How to compute new sigma values at the selected positions:

1295

1296 - ``'constant'``: ``sigma = value``.

1297 - ``'scale_mu'`` (default): ``sigma = max(floor, factor * |mu|)``,

1298 re-opening the SNES search width in proportion to each

1299 parameter's magnitude. ``factor`` defaults to ``0.1`` for this

1300 strategy.

1301 - ``'scale_sigma'``: ``sigma = sigma * factor``. Requires the

1302 selected sigma values to already be set (not ``NaN``).

1303 - ``'uniform'``: draw ``sigma ~ U(low, high)``.

1304 - ``'normal'``: draw ``sigma = |N(mean, std)|``.

1305 value

1306 Sigma value for ``strategy='constant'``.

1307 factor

1308 Scale factor for ``strategy='scale_mu'`` (default ``0.1`` if not

1309 given) or ``strategy='scale_sigma'`` (required).

1310 floor

1311 Minimum sigma for ``strategy='scale_mu'``.

1312 low, high

1313 Bounds for ``strategy='uniform'``.

1314 mean, std

1315 Parameters of the normal distribution for ``strategy='normal'``.

1316 species

1317 Restrict the update to one or more species (and descriptor pairs

1318 involving them). ``None`` (default) applies to all species; global

1319 parameters (the shared bias, ``sqrt_epsilon_infinity``) are only

1320 included when ``species`` is ``None``.

1321 component

1322 Restrict the update to one or more of ``'network_weights'``,

1323 ``'descriptor'``, ``'charge_head'``. ``None`` (default) applies to

1324 all three.

1325 target

1326 Which existing sigma values to update: ``'unset'`` (default) only

1327 fills in ``NaN`` entries (e.g. those left by :meth:`add_species`/

1328 :meth:`augment`); ``'set'`` only updates already-set entries;

1329 ``'all'`` updates every selected entry regardless of its current

1330 value.

1331 seed

1332 Seed for the random number generator used by the ``'uniform'`` and

1333 ``'normal'`` strategies. Pass an integer for reproducibility.

1334

1335 Returns

1336 -------

1337 Model

1338 New model with updated restart sigma values.

1339

1340 Raises

1341 ------

1342 ValueError

1343 If ``strategy``/``target``/``component`` is not recognized, if a

1344 strategy-specific required argument is missing, or if

1345 ``strategy='scale_sigma'`` is applied to a still-unset (``NaN``)

1346 sigma value.

1347 """

1348 valid_strategies = {'constant', 'scale_mu', 'scale_sigma', 'uniform', 'normal'}

1349 if strategy not in valid_strategies:

1350 raise ValueError(

1351 f'strategy must be one of {sorted(valid_strategies)}; got {strategy!r}'

1352 )

1353 if target not in ('unset', 'set', 'all'):

1354 raise ValueError(f"target must be 'unset', 'set', or 'all'; got {target!r}")

1355 if strategy == 'constant' and value is None:

1356 raise ValueError("strategy='constant' requires value.")

1357 if strategy == 'scale_mu' and factor is None:

1358 factor = 0.1

1359 if strategy == 'scale_sigma' and factor is None:

1360 raise ValueError("strategy='scale_sigma' requires factor.")

1361 if strategy == 'uniform' and (low is None or high is None):

1362 raise ValueError("strategy='uniform' requires low and high.")

1363 if strategy == 'normal' and (mean is None or std is None):

1364 raise ValueError("strategy='normal' requires mean and std.")

1365

1366 new = copy.deepcopy(self)

1367 if new.restart_parameters is None:

1368 new.restart_parameters = _new_restart_parameters_from_model(new)

1369

1370 rng = np.random.default_rng(seed)

1371 kw = dict(value=value, factor=factor, floor=floor, low=low, high=high, mean=mean, std=std)

1372 for mu, sigma_container, sigma_key in _restart_leaves(

1373 new, new.restart_parameters, component, species

1374 ):

1375 _apply_sigma_strategy(mu, sigma_container, sigma_key, strategy, target, rng, **kw)

1376

1377 return new

1378

1379 def augment(self,

1380 n_neuron: int = None,

1381 l_max_4b: int = None,

1382 l_max_5b: int = None,

1383 has_q_112: bool = None,

1384 has_q_123: bool = None,

1385 has_q_233: bool = None,

1386 has_q_134: bool = None,

1387 charge_head: bool = False,

1388 charge_mode: int = 1) -> 'Model':

1389 """Augment the model by adding neurons, descriptor terms, or a charge output head.

1390

1391 Returns a new :class:`Model` with the requested structural changes applied.

1392 The source model is not modified. Existing parameter values (``mu`` and

1393 ``sigma``) are preserved exactly; new parameters are initialized to

1394 ``mu = 0``, with the corresponding restart sigma left unset (``NaN``).

1395 Call :meth:`set_restart_sigma` afterwards to initialize the new sigma

1396 entries (e.g. ``model.augment(n_neuron=40).set_restart_sigma()`` fills

1397 only the new entries by default).

1398

1399 Parameters

1400 ----------

1401 n_neuron

1402 Target neuron count; must be >= current. ``None`` leaves unchanged.

1403 l_max_4b

1404 Target 4-body l_max value; must be >= current. ``None`` leaves unchanged.

1405 l_max_5b

1406 Target 5-body l_max value; must be >= current. ``None`` leaves unchanged.

1407 has_q_112

1408 ``True`` enables the q_112 5-body descriptor; ``None`` or ``False`` leaves

1409 the current state unchanged (disabling an already-enabled term raises).

1410 has_q_123

1411 Same as ``has_q_112`` but for the q_123 term.

1412 has_q_233

1413 Same as ``has_q_112`` but for the q_233 term.

1414 has_q_134

1415 Same as ``has_q_112`` but for the q_134 term.

1416 charge_head

1417 If ``True``, promote a ``potential`` model to ``potential_with_charges`` by

1418 adding a charge output head (w1_charge per species and sqrt_epsilon_infinity).

1419 charge_mode

1420 Charge algorithm variant to record for the new charge head; must be 1 or 2.

1421 1 corresponds to a qNEP model, including both real- and reciprocal-space

1422 contributions. 2 corresponds to a qNEP model, including the reciprocal-space

1423 contribution only. Only meaningful when ``charge_head=True``.

1424

1425 Returns

1426 -------

1427 Model

1428 New model with updated structure, weights, and restart statistics.

1429

1430 Raises

1431 ------

1432 ValueError

1433 If ``restart_parameters`` is not loaded, if ``n_neuron`` or an ``l_max_*``

1434 target is smaller than the current value, if a ``has_q_*`` flag attempts to

1435 disable an already-enabled term, or if ``charge_head=True`` on a model that

1436 is not of type ``potential``.

1437 """

1438 # Structural checks (independent of restart)

1439 if self.version not in (3, 4):

1440 raise ValueError(

1441 f'augment() only supports NEP versions 3 and 4; got version {self.version}.'

1442 )

1443 if n_neuron is not None and n_neuron < self.n_neuron:

1444 raise ValueError(

1445 f'n_neuron ({n_neuron}) must be >= current n_neuron ({self.n_neuron}); '

1446 'use prune() to reduce.'

1447 )

1448 if l_max_4b is not None and l_max_4b < self.l_max_4b:

1449 raise ValueError(

1450 f'l_max_4b ({l_max_4b}) must be >= current l_max_4b ({self.l_max_4b}); '

1451 'use prune() to disable.'

1452 )

1453 if l_max_5b is not None and l_max_5b < self.l_max_5b:

1454 raise ValueError(

1455 f'l_max_5b ({l_max_5b}) must be >= current l_max_5b ({self.l_max_5b}); '

1456 'use prune() to disable.'

1457 )

1458 for flag_val, name in [

1459 (has_q_112, 'has_q_112'), (has_q_123, 'has_q_123'), (has_q_233, 'has_q_233'),

1460 (has_q_134, 'has_q_134')

1461 ]:

1462 if flag_val is False and getattr(self, name):

1463 raise ValueError(

1464 f'Cannot disable {name} via augment(); '

1465 'use prune() to disable descriptor terms.'

1466 )

1467 if charge_head and self.model_type != 'potential':

1468 raise ValueError(

1469 f'charge_head=True requires model_type="potential"; '

1470 f'got "{self.model_type}".'

1471 )

1472 if charge_head and charge_mode not in (1, 2):

1473 raise ValueError(f'charge_mode must be 1 or 2; got {charge_mode}.')

1474 if self.restart_parameters is None:

1475 raise ValueError(

1476 'restart_parameters must be loaded before calling augment(). '

1477 'Pass restart_file= to read_model() or call model.read_restart() first.'

1478 )

1479

1480 new = copy.deepcopy(self)

1481

1482 # Resolve new structural parameters

1483 new_l_max_4b = l_max_4b if l_max_4b is not None else self.l_max_4b

1484 new_l_max_5b = l_max_5b if l_max_5b is not None else self.l_max_5b

1485 new_has_q_112 = int(has_q_112) if has_q_112 is not None else self.has_q_112

1486 new_has_q_123 = int(has_q_123) if has_q_123 is not None else self.has_q_123

1487 new_has_q_233 = int(has_q_233) if has_q_233 is not None else self.has_q_233

1488 new_has_q_134 = int(has_q_134) if has_q_134 is not None else self.has_q_134

1489 new_n_neuron = n_neuron if n_neuron is not None else self.n_neuron

1490

1491 new_l_max_enh = (self.l_max_3b

1492 + (new_l_max_4b > 0) + (new_l_max_5b > 0)

1493 + (new_has_q_112 > 0) + (new_has_q_123 > 0) + (new_has_q_233 > 0)

1494 + (new_has_q_134 > 0))

1495 new_n_desc_angular = (self.n_max_angular + 1) * new_l_max_enh

1496 old_n_desc = self.n_descriptor_radial + self.n_descriptor_angular

1497 new_n_desc = self.n_descriptor_radial + new_n_desc_angular

1498 delta_desc = new_n_desc - old_n_desc

1499 delta_neuron = new_n_neuron - self.n_neuron

1500

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

1502

1503 # Step 1: Expand descriptor dimensions (new columns in w0, new q_scaler entries)

1504 if delta_desc > 0:

1505 for s in keys:

1506 old_w0 = new.ann_parameters[s]['w0'] # (n_neuron_old, old_n_desc)

1507 new.ann_parameters[s]['w0'] = np.hstack(

1508 [old_w0, np.zeros((self.n_neuron, delta_desc))]

1509 )

1510 old_mu_w0 = new.restart_parameters['ann_mu'][s]['w0']

1511 new.restart_parameters['ann_mu'][s]['w0'] = np.hstack(

1512 [old_mu_w0, np.zeros((self.n_neuron, delta_desc))]

1513 )

1514 old_sigma_w0 = new.restart_parameters['ann_sigma'][s]['w0']

1515 new.restart_parameters['ann_sigma'][s]['w0'] = np.hstack(

1516 [old_sigma_w0, np.full((self.n_neuron, delta_desc), np.nan)]

1517 )

1518 new.q_scaler = list(new.q_scaler) + [1.0] * delta_desc

1519

1520 # Step 2: Expand neuron count (new rows in w0/b0, new columns in w1)

1521 if delta_neuron > 0:

1522 for s in keys:

1523 # w0: append new rows

1524 cur_w0 = new.ann_parameters[s]['w0'] # (n_old, new_n_desc)

1525 new.ann_parameters[s]['w0'] = np.vstack(

1526 [cur_w0, np.zeros((delta_neuron, new_n_desc))]

1527 )

1528 # b0: append new rows

1529 cur_b0 = new.ann_parameters[s]['b0']

1530 new.ann_parameters[s]['b0'] = np.vstack(

1531 [cur_b0, np.zeros((delta_neuron, 1))]

1532 )

1533 # w1: append new columns; handle both 2D (standard) and 1D (charge)

1534 cur_w1 = new.ann_parameters[s]['w1']

1535 zeros_w1 = (np.zeros(delta_neuron) if cur_w1.ndim == 1

1536 else np.zeros((1, delta_neuron)))

1537 new.ann_parameters[s]['w1'] = np.hstack([cur_w1, zeros_w1])

1538 if 'w1_charge' in new.ann_parameters[s]:

1539 cur_wc = new.ann_parameters[s]['w1_charge']

1540 new.ann_parameters[s]['w1_charge'] = np.hstack([cur_wc, np.zeros(delta_neuron)])

1541

1542 # restart w0

1543 cur_mu_w0 = new.restart_parameters['ann_mu'][s]['w0']

1544 new.restart_parameters['ann_mu'][s]['w0'] = np.vstack(

1545 [cur_mu_w0, np.zeros((delta_neuron, new_n_desc))]

1546 )

1547 cur_sigma_w0 = new.restart_parameters['ann_sigma'][s]['w0']

1548 new.restart_parameters['ann_sigma'][s]['w0'] = np.vstack(

1549 [cur_sigma_w0, np.full((delta_neuron, new_n_desc), np.nan)]

1550 )

1551 # restart b0

1552 cur_mu_b0 = new.restart_parameters['ann_mu'][s]['b0']

1553 new.restart_parameters['ann_mu'][s]['b0'] = np.vstack(

1554 [cur_mu_b0, np.zeros((delta_neuron, 1))]

1555 )

1556 cur_sigma_b0 = new.restart_parameters['ann_sigma'][s]['b0']

1557 new.restart_parameters['ann_sigma'][s]['b0'] = np.vstack(

1558 [cur_sigma_b0, np.full((delta_neuron, 1), np.nan)]

1559 )

1560 # restart w1

1561 cur_mu_w1 = new.restart_parameters['ann_mu'][s]['w1']

1562 zeros_w1 = (np.zeros(delta_neuron) if cur_mu_w1.ndim == 1

1563 else np.zeros((1, delta_neuron)))

1564 new.restart_parameters['ann_mu'][s]['w1'] = np.hstack([cur_mu_w1, zeros_w1])

1565 cur_sigma_w1 = new.restart_parameters['ann_sigma'][s]['w1']

1566 nan_w1 = (np.full(delta_neuron, np.nan) if cur_sigma_w1.ndim == 1

1567 else np.full((1, delta_neuron), np.nan))

1568 new.restart_parameters['ann_sigma'][s]['w1'] = np.hstack([cur_sigma_w1, nan_w1])

1569 if 'w1_charge' in new.restart_parameters['ann_mu'][s]:

1570 cur = new.restart_parameters['ann_mu'][s]['w1_charge']

1571 new.restart_parameters['ann_mu'][s]['w1_charge'] = np.hstack(

1572 [cur, np.zeros(delta_neuron)]

1573 )

1574 cur = new.restart_parameters['ann_sigma'][s]['w1_charge']

1575 new.restart_parameters['ann_sigma'][s]['w1_charge'] = np.hstack(

1576 [cur, np.full(delta_neuron, np.nan)]

1577 )

1578

1579 # Step 3: Add charge output head

1580 if charge_head:

1581 new.model_type = 'potential_with_charges'

1582 new.charge_mode = charge_mode

1583 new.sqrt_epsilon_infinity = 1.0

1584 for s in keys:

1585 cur_w1 = new.ann_parameters[s]['w1'] # (1, new_n_neuron)

1586 new.ann_parameters[s]['w1'] = cur_w1[0, :] # flatten to 1D

1587 new.ann_parameters[s]['w1_charge'] = np.zeros(new_n_neuron)

1588

1589 cur_mu_w1 = new.restart_parameters['ann_mu'][s]['w1']

1590 new.restart_parameters['ann_mu'][s]['w1'] = cur_mu_w1[0, :]

1591 new.restart_parameters['ann_mu'][s]['w1_charge'] = np.zeros(new_n_neuron)

1592

1593 cur_sigma_w1 = new.restart_parameters['ann_sigma'][s]['w1']

1594 new.restart_parameters['ann_sigma'][s]['w1'] = cur_sigma_w1[0, :]

1595 new.restart_parameters['ann_sigma'][s]['w1_charge'] = np.full(

1596 new_n_neuron, np.nan

1597 )

1598

1599 new.restart_parameters['ann_mu']['sqrt_epsilon_infinity'] = 1.0

1600 new.restart_parameters['ann_sigma']['sqrt_epsilon_infinity'] = float('nan')

1601

1602 # Step 4: Update header metadata

1603 new.l_max_4b = new_l_max_4b

1604 new.l_max_5b = new_l_max_5b

1605 new.has_q_112 = new_has_q_112

1606 new.has_q_123 = new_has_q_123

1607 new.has_q_233 = new_has_q_233

1608 new.has_q_134 = new_has_q_134

1609 new.n_descriptor_angular = new_n_desc_angular

1610 new.n_neuron = new_n_neuron

1611

1612 # Step 5: Recalculate parameter counts

1613 _recalculate_parameter_counts(new)

1614

1615 return new

1616

1617 def prune(self,

1618 n_neuron: int = None,

1619 l_max_4b: int = None,

1620 l_max_5b: int = None,

1621 has_q_112: bool = None,

1622 has_q_123: bool = None,

1623 has_q_233: bool = None,

1624 has_q_134: bool = None,

1625 charge_head: bool = False) -> 'Model':

1626 """Prune the model by removing neurons, disabling descriptor terms, or removing

1627 the charge output head.

1628

1629 Returns a new :class:`Model` with the requested structural changes applied.

1630 The source model is not modified. When reducing ``n_neuron``, neurons are

1631 selected by importance score averaged over species:

1632 ``importance[n] = mean_s(||w0_s[n,:]||_2 * |w1_s[n]|)``.

1633

1634 Surviving parameters (``mu`` and ``sigma``) are left exactly as they

1635 were. Use :meth:`set_restart_sigma` explicitly afterwards if you want

1636 to re-open the SNES search width for the survivors before continuing

1637 training.

1638

1639 Parameters

1640 ----------

1641 n_neuron

1642 Target neuron count; must be <= current. ``None`` leaves unchanged.

1643 l_max_4b

1644 Target 4-body l_max; must be <= current. Setting to ``0`` removes the

1645 4-body angular descriptor block. Reducing to a lower non-zero value is

1646 a header-only change (descriptor dimensions unchanged). ``None`` leaves

1647 unchanged.

1648 l_max_5b

1649 Same as ``l_max_4b`` but for five-body terms.

1650 has_q_112

1651 ``False`` disables and removes the q_112 descriptor block. ``None``

1652 leaves unchanged. ``True`` is not valid; use :meth:`augment` instead.

1653 has_q_123

1654 Same as ``has_q_112`` but for the q_123 term.

1655 has_q_233

1656 Same as ``has_q_112`` but for the q_233 term.

1657 has_q_134

1658 Same as ``has_q_112`` but for the q_134 term.

1659 charge_head

1660 If ``True``, remove the charge output head from a

1661 ``potential_with_charges`` model, converting it back to ``potential``.

1662 Removes ``w1_charge`` per species and ``sqrt_epsilon_infinity`` from

1663 the restart.

1664

1665 Returns

1666 -------

1667 Model

1668 New model with reduced structure, weights, and restart statistics.

1669

1670 Raises

1671 ------

1672 ValueError

1673 If ``restart_parameters`` is not loaded, if any target value would

1674 expand the model (use :meth:`augment` instead), if a ``has_q_*``

1675 flag is set to ``True``, or if ``charge_head=True`` on a model

1676 without charges.

1677 """

1678 # --- Resolve target values ---

1679 new_n_neuron = n_neuron if n_neuron is not None else self.n_neuron

1680 new_l_max_4b = l_max_4b if l_max_4b is not None else self.l_max_4b

1681 new_l_max_5b = l_max_5b if l_max_5b is not None else self.l_max_5b

1682 new_has_q_112 = 0 if has_q_112 is False else self.has_q_112

1683 new_has_q_123 = 0 if has_q_123 is False else self.has_q_123

1684 new_has_q_233 = 0 if has_q_233 is False else self.has_q_233

1685 new_has_q_134 = 0 if has_q_134 is False else self.has_q_134

1686

1687 # --- Validate ---

1688 if self.version not in (3, 4):

1689 raise ValueError(

1690 f'prune() only supports NEP versions 3 and 4; got version {self.version}.'

1691 )

1692 if new_n_neuron > self.n_neuron:

1693 raise ValueError(

1694 f'n_neuron ({new_n_neuron}) must be <= current n_neuron ({self.n_neuron}); '

1695 'use augment() to increase.'

1696 )

1697 if new_l_max_4b > self.l_max_4b:

1698 raise ValueError(

1699 f'l_max_4b ({new_l_max_4b}) must be <= current l_max_4b ({self.l_max_4b}); '

1700 'use augment() to increase.'

1701 )

1702 if new_l_max_5b > self.l_max_5b:

1703 raise ValueError(

1704 f'l_max_5b ({new_l_max_5b}) must be <= current l_max_5b ({self.l_max_5b}); '

1705 'use augment() to increase.'

1706 )

1707 for flag_val, name in [

1708 (has_q_112, 'has_q_112'), (has_q_123, 'has_q_123'),

1709 (has_q_233, 'has_q_233'), (has_q_134, 'has_q_134')

1710 ]:

1711 if flag_val is True:

1712 raise ValueError(

1713 f'Cannot enable {name} via prune(); '

1714 'use augment() to enable descriptor terms.'

1715 )

1716 if charge_head and self.model_type != 'potential_with_charges':

1717 raise ValueError(

1718 f'charge_head=True requires model_type="potential_with_charges"; '

1719 f'got "{self.model_type}".'

1720 )

1721 if self.restart_parameters is None:

1722 raise ValueError(

1723 'restart_parameters must be loaded before calling prune(). '

1724 'Pass restart_file= to read_model() or call model.read_restart() first.'

1725 )

1726

1727 new = copy.deepcopy(self)

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

1729

1730 # Step 1: Neuron pruning — keep the most important neurons

1731 if new_n_neuron < self.n_neuron:

1732 importances = []

1733 for s in keys:

1734 w0 = self.ann_parameters[s]['w0'] # (n_neuron, n_desc)

1735 w1_flat = self.ann_parameters[s]['w1'].ravel()

1736 if 'w1_charge' in self.ann_parameters[s]:

1737 output_norm = np.abs(w1_flat) + np.abs(self.ann_parameters[s]['w1_charge'])

1738 else:

1739 output_norm = np.abs(w1_flat)

1740 importances.append(np.linalg.norm(w0, axis=1) * output_norm)

1741

1742 keep_idx = np.sort(np.argsort(np.mean(importances, axis=0))[-new_n_neuron:])

1743

1744 for s in keys:

1745 new.ann_parameters[s]['w0'] = new.ann_parameters[s]['w0'][keep_idx, :]

1746 new.ann_parameters[s]['b0'] = new.ann_parameters[s]['b0'][keep_idx, :]

1747 w1 = new.ann_parameters[s]['w1']

1748 new.ann_parameters[s]['w1'] = w1[:, keep_idx] if w1.ndim == 2 else w1[keep_idx]

1749 if 'w1_charge' in new.ann_parameters[s]:

1750 new.ann_parameters[s]['w1_charge'] = (

1751 new.ann_parameters[s]['w1_charge'][keep_idx]

1752 )

1753 for pk in ['ann_mu', 'ann_sigma']:

1754 rp = new.restart_parameters[pk][s]

1755 rp['w0'] = rp['w0'][keep_idx, :]

1756 rp['b0'] = rp['b0'][keep_idx, :]

1757 w1 = rp['w1']

1758 rp['w1'] = w1[:, keep_idx] if w1.ndim == 2 else w1[keep_idx]

1759 if 'w1_charge' in rp:

1760 rp['w1_charge'] = rp['w1_charge'][keep_idx]

1761

1762 # Step 2: Descriptor column pruning (disabling higher-body terms)

1763 n_per = self.n_max_angular + 1

1764 hb_terms = [

1765 (self.l_max_4b, new_l_max_4b),

1766 (self.l_max_5b, new_l_max_5b),

1767 (self.has_q_112, new_has_q_112),

1768 (self.has_q_123, new_has_q_123),

1769 (self.has_q_233, new_has_q_233),

1770 (self.has_q_134, new_has_q_134),

1771 ]

1772 keep_cols = list(range(self.n_descriptor_radial + n_per * self.l_max_3b))

1773 col_offset = len(keep_cols)

1774 for old_val, new_val in hb_terms:

1775 if old_val > 0:

1776 if new_val > 0:

1777 keep_cols.extend(range(col_offset, col_offset + n_per))

1778 col_offset += n_per

1779

1780 old_n_desc = self.n_descriptor_radial + self.n_descriptor_angular

1781 if len(keep_cols) < old_n_desc:

1782 keep_cols = np.array(keep_cols, dtype=int)

1783 for s in keys:

1784 new.ann_parameters[s]['w0'] = new.ann_parameters[s]['w0'][:, keep_cols]

1785 for pk in ['ann_mu', 'ann_sigma']:

1786 rp = new.restart_parameters[pk][s]

1787 rp['w0'] = rp['w0'][:, keep_cols]

1788 new.q_scaler = [new.q_scaler[i] for i in keep_cols]

1789

1790 # Step 3: Charge head removal

1791 if charge_head:

1792 new.model_type = 'potential'

1793 new.charge_mode = 0

1794 new.sqrt_epsilon_infinity = None

1795 for s in keys:

1796 w1 = new.ann_parameters[s]['w1'] # 1D (n_neuron,)

1797 new.ann_parameters[s]['w1'] = w1.reshape(1, -1)

1798 del new.ann_parameters[s]['w1_charge']

1799 for pk in ['ann_mu', 'ann_sigma']:

1800 rp = new.restart_parameters[pk][s]

1801 rp['w1'] = rp['w1'].reshape(1, -1)

1802 del rp['w1_charge']

1803 del new.restart_parameters['ann_mu']['sqrt_epsilon_infinity']

1804 del new.restart_parameters['ann_sigma']['sqrt_epsilon_infinity']

1805

1806 # Step 4: Update header fields

1807 new.n_neuron = new_n_neuron

1808 new.l_max_4b = new_l_max_4b

1809 new.l_max_5b = new_l_max_5b

1810 new.has_q_112 = new_has_q_112

1811 new.has_q_123 = new_has_q_123

1812 new.has_q_233 = new_has_q_233

1813 new.has_q_134 = new_has_q_134

1814

1815 new_l_max_enh = (self.l_max_3b

1816 + (new_l_max_4b > 0) + (new_l_max_5b > 0)

1817 + (new_has_q_112 > 0) + (new_has_q_123 > 0) + (new_has_q_233 > 0)

1818 + (new_has_q_134 > 0))

1819 new.n_descriptor_angular = (self.n_max_angular + 1) * new_l_max_enh

1820

1821 # Step 5: Recalculate parameter counts

1822 _recalculate_parameter_counts(new)

1823

1824 return new

1825

1826

1827def read_model(filename: str, restart_file: str = None) -> Model:

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

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

1830 object.

1831

1832 Parameters

1833 ----------

1834 filename

1835 Input file name.

1836 restart_file

1837 If provided, also read restart parameters from this file in

1838 `nep.restart` format and attach them to the returned model.

1839 Defaults to None.

1840 """

1841 data, parameters = _get_nep_contents(filename)

1842

1843 # sanity checks

1844 for fld in ['version', 'types', 'model_type', 'cutoff', 'basis_size', 'n_max', 'l_max', 'ANN']:

1845 if fld not in data:

1846 raise ValueError(f'Invalid model file; {fld} line is missing')

1847 if data['version'] not in [3, 4, 5]:

1848 raise ValueError('Invalid model file; only NEP versions 3, 4 and 5 are currently supported')

1849

1850 # split up zbl tuple (optional typewise cutoff factor as a third entry)

1851 if 'zbl' in data:

1852 if len(data['zbl']) == 3:

1853 data['zbl_typewise_cutoff_factor'] = data['zbl'][2]

1854 data['zbl'] = data['zbl'][:2]

1855 elif len(data['zbl']) != 2:

1856 raise ValueError(

1857 f'Invalid model file; zbl line must have 2 or 3 entries, got {len(data["zbl"])}'

1858 )

1859

1860 # split up cutoff tuple

1861 N_types = len(data['types'])

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

1863 if len(data['cutoff']) not in [4, 2*N_types+2]:

1864 raise ValueError(

1865 'Invalid model file; cutoff line must have 4 entries (global cutoffs) or '

1866 f'{2*N_types+2} entries (typewise cutoffs for {N_types} types), '

1867 f'got {len(data["cutoff"])}'

1868 )

1869 if not all(np.isfinite(data['cutoff'])):

1870 raise ValueError('Invalid model file; cutoff values must be finite')

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

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

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

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

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

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

1877 else:

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

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

1880 del data['cutoff']

1881

1882 # split up basis_size tuple

1883 if len(data['basis_size']) != 2:

1884 raise ValueError(

1885 f'Invalid model file; basis_size line must have 2 entries, '

1886 f'got {len(data["basis_size"])}'

1887 )

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

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

1890 del data['basis_size']

1891

1892 # split up n_max tuple

1893 if len(data['n_max']) != 2:

1894 raise ValueError(

1895 f'Invalid model file; n_max line must have 2 entries, got {len(data["n_max"])}'

1896 )

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

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

1899 del data['n_max']

1900

1901 # split up nl_max tuple

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

1903 if len_l not in [1, 2, 3, 4, 5, 6, 7]:

1904 raise ValueError(

1905 f'Invalid model file; l_max line must have between 1 and 7 entries, got {len_l}'

1906 )

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

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

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

1910 data['has_q_112'] = data['l_max'][3] if len_l > 3 else 0

1911 data['has_q_123'] = data['l_max'][4] if len_l > 4 else 0

1912 data['has_q_233'] = data['l_max'][5] if len_l > 5 else 0

1913 data['has_q_134'] = data['l_max'][6] if len_l > 6 else 0

1914 del data['l_max']

1915

1916 # compute dimensions of descriptor components

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

1918 l_max_enh = (data['l_max_3b']

1919 + (data['l_max_4b'] > 0)

1920 + (data['l_max_5b'] > 0)

1921 + (data['has_q_112'] > 0)

1922 + (data['has_q_123'] > 0)

1923 + (data['has_q_233'] > 0)

1924 + (data['has_q_134'] > 0))

1925 data['n_descriptor_angular'] = (data['n_max_angular'] + 1) * l_max_enh

1926 n_descriptor = data['n_descriptor_radial'] + data['n_descriptor_angular']

1927

1928 is_charged_model = data['model_type'] == 'potential_with_charges'

1929 # compute number of parameters

1930 data['n_neuron'] = data['ANN'][0]

1931 del data['ANN']

1932 n_types = len(data['types'])

1933 if data['version'] == 3:

1934 n = 1

1935 n_bias = 1

1936 elif data['version'] == 4 and is_charged_model:

1937 # one hidden layer per atomic species, but two output nodes

1938 n = n_types

1939 n_bias = 2

1940 elif data['version'] == 4:

1941 # one hidden layer per atomic species

1942 n = n_types

1943 n_bias = 1

1944 else: # NEP5

1945 # like nep4, but additionally has an

1946 # individual bias term in the output

1947 # layer for each species.

1948 n = n_types

1949 n_bias = 1 + n_types # one global bias + one per species

1950

1951 n_ann_input_weights = (n_descriptor + 1) * data['n_neuron'] # weights + bias

1952 n_ann_output_weights = 2*data['n_neuron'] if is_charged_model else data['n_neuron'] # weights

1953 n_ann_parameters = (

1954 n_ann_input_weights + n_ann_output_weights

1955 ) * n + n_bias

1956

1957 n_descriptor_weights = n_types**2 * (

1958 (data['n_max_radial'] + 1) * (data['n_basis_radial'] + 1)

1959 + (data['n_max_angular'] + 1) * (data['n_basis_angular'] + 1)

1960 )

1961 data['n_parameters'] = n_ann_parameters + n_descriptor_weights + n_descriptor

1962 is_polarizability_model = data['model_type'] == 'polarizability'

1963 if data['n_parameters'] + n_ann_parameters == len(parameters):

1964 data['n_parameters'] += n_ann_parameters

1965 if not is_polarizability_model:

1966 raise ValueError(

1967 'Model is not labelled as a polarizability model, but the number of '

1968 'parameters matches a polarizability model.\n'

1969 'If this is a polarizability model trained with GPUMD <=v3.8, please '

1970 'modify the header in the nep.txt file to enable parsing '

1971 f'`nep{data["version"]}_polarizability`.\n'

1972 )

1973 if len(parameters) < data['n_parameters']:

1974 raise ValueError(

1975 'Invalid model file; expected '

1976 f'{data["n_parameters"]} parameter values, found {len(parameters)} '

1977 '-- file may be truncated'

1978 )

1979 elif len(parameters) > data['n_parameters']:

1980 raise ValueError(

1981 'Invalid model file; expected '

1982 f'{data["n_parameters"]} parameter values, found {len(parameters)} '

1983 '-- file may contain extra or corrupted data'

1984 )

1985 data['n_ann_parameters'] = n_ann_parameters

1986

1987 # split up parameters into the ANN weights, descriptor weights, and scaling parameters

1988 n1 = n_ann_parameters

1989 n1 *= 2 if is_polarizability_model else 1

1990 n2 = n1 + n_descriptor_weights

1991 data['ann_parameters'] = parameters[:n1]

1992 descriptor_weights = np.array(parameters[n1:n2])

1993 data['q_scaler'] = parameters[n2:]

1994

1995 # add ann parameters to data dict

1996 ann_groups = data['types'] if data['version'] in (4, 5) else ['all_species']

1997 sorted_ann_parameters = _sort_ann_parameters(data['ann_parameters'],

1998 ann_groups,

1999 data['n_neuron'],

2000 n,

2001 n_bias,

2002 n_descriptor,

2003 is_polarizability_model,

2004 is_charged_model)

2005

2006 data['ann_parameters'] = sorted_ann_parameters

2007 if 'sqrt_epsilon_infinity' in sorted_ann_parameters.keys():

2008 data['sqrt_epsilon_infinity'] = sorted_ann_parameters['sqrt_epsilon_infinity']

2009 sorted_ann_parameters.pop('sqrt_epsilon_infinity')

2010 data['ann_parameters'] = sorted_ann_parameters

2011

2012 # add descriptors to data dict

2013 data['n_descriptor_parameters'] = len(descriptor_weights)

2014 radial, angular = _sort_descriptor_parameters(descriptor_weights,

2015 data['types'],

2016 data['n_max_radial'],

2017 data['n_basis_radial'],

2018 data['n_max_angular'],

2019 data['n_basis_angular'])

2020 data['radial_descriptor_weights'] = radial

2021 data['angular_descriptor_weights'] = angular

2022

2023 model = Model(**data)

2024 if restart_file is not None:

2025 model.read_restart(restart_file)

2026 return model