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

1from os.path import exists

2from os.path import join as join_path

3from typing import Any, Iterable, NamedTuple, TextIO

4from warnings import warn

6import numpy as np

7from ase import Atoms

8from ase.io import read, write

9from ase.stress import voigt_6_to_full_3x3_stress

10from pandas import DataFrame

12from calorine.nep.tensor_conventions import ASE_VOIGT6_ORDER, BEC_FULL9_ORDER, \

13 nep_reduced6_to_ase_voigt6

16def read_loss(filename: str) -> DataFrame:

17 """Parses a file in `loss.out` format from GPUMD and returns the

18 content as a data frame. More information concerning file format,

19 content and units can be found `here

20 <https://gpumd.org/nep/output_files/loss_out.html>`__.

22 Parameters

23 ----------

24 filename

25 input file name

27 """

28 data = np.loadtxt(filename)

29 if isinstance(data[0], np.float64):

30 # If only a single row in loss.out, append a dimension

31 data = data.reshape(1, -1)

32 if len(data[0]) == 6:

33 tags = 'total_loss L1 L2'

34 tags += ' RMSE_P_train'

35 tags += ' RMSE_P_test'

36 elif len(data[0]) == 10:

37 tags = 'total_loss L1 L2'

38 tags += ' RMSE_E_train RMSE_F_train RMSE_V_train'

39 tags += ' RMSE_E_test RMSE_F_test RMSE_V_test'

40 elif len(data[0]) == 14:

41 tags = 'total_loss L1 L2'

42 tags += ' RMSE_E_train RMSE_F_train RMSE_V_train RMSE_Q_train RMSE_Z_train'

43 tags += ' RMSE_E_test RMSE_F_test RMSE_V_test RMSE_Q_test RMSE_Z_test'

44 else:

45 raise ValueError(

46 f'Input file contains {len(data[0])} data columns. Expected 6 or 10 columns.'

47 )

48 generations = range(100, len(data) * 100 + 1, 100)

49 df = DataFrame(data=data[:, 1:], columns=tags.split(), index=generations)

50 return df

53def _write_structure_in_nep_format(structure: Atoms, f: TextIO) -> None:

54 """Write structure block into a file-like object in format readable by nep executable.

56 Parameters

57 ----------

58 structure

59 input structure; must hold information regarding energy and forces

60 f

61 file-like object to which to write

62 """

64 # Allowed keyword=value pairs. Use ASEs extyz write functionality.:

65 # lattice="ax ay az bx by bz cx cy cz" (mandatory)

66 # energy=energy_value (mandatory)

67 # virial="vxx vxy vxz vyx vyy vyz vzx vzy vzz" (optional)

68 # weight=relative_weight (optional)

69 # properties=property_name:data_type:number_of_columns

70 # species:S:1 (mandatory)

71 # pos:R:3 (mandatory)

72 # force:R:3 or forces:R:3 (mandatory)

74 # If a structure is to be used for training, it needs to either have target

75 # 1. energies and forces,

76 # 2. dipole, denoted `dipole="dx dy dz"` in the info string, or

77 # 3. polarizability/susceptibility, denoted `pol="pxx pxy pxz pyx pyy pyz pzx pzy pzz"`

78 # in the info string.

79 has_energies_and_forces = True

80 try:

81 structure.get_potential_energy()

82 structure.get_forces() # calculate forces to have them on the Atoms object

83 except RuntimeError:

84 has_energies_and_forces = False

86 has_dipole = 'dipole' in structure.info.keys()

87 has_pol = 'pol' in structure.info.keys()

89 if not has_energies_and_forces and not has_dipole and not has_pol:

90 raise RuntimeError('Failed to retrieve target energies/forces,'

91 ' dipoles, or polarizabilities for structure')

92 if np.isclose(structure.get_volume(), 0):

93 raise ValueError('Structure cell must have a non-zero volume!')

94 try:

95 structure.get_stress()

96 except RuntimeError:

97 warn('Failed to retrieve stresses for structure')

98 write(filename=f, images=structure, write_info=True, format='extxyz')

100

101def write_structures(outfile: str, structures: list[Atoms]) -> None:

102 """Writes structures for training/testing in format readable by nep executable.

103

104 Parameters

105 ----------

106 outfile

107 output filename

108 structures

109 list of structures with energy, forces, and (possibly) stresses

110 """

111 with open(outfile, 'w') as f:

112 for structure in structures:

113 _write_structure_in_nep_format(structure, f)

114

115

116def write_nepfile(parameters: NamedTuple, dirname: str) -> None:

117 """Writes parameters file for NEP construction.

118

119 Parameters

120 ----------

121 parameters

122 input parameters; see `here <https://gpumd.org/nep/input_parameters/index.html>`__

123 dirname

124 directory in which to place input file and links

125 """

126 with open(join_path(dirname, 'nep.in'), 'w') as f:

127 for key, val in parameters.items():

128 f.write(f'{key} ')

129 if isinstance(val, Iterable) and not isinstance(val, str):

130 f.write(' '.join([f'{v}' for v in val]))

131 else:

132 f.write(f'{val}')

133 f.write('\n')

134

135

136def read_nepfile(filename: str) -> dict[str, Any]:

137 """Returns the content of a configuration file (`nep.in`) as a dictionary.

138

139 Parameters

140 ----------

141 filename

142 input file name

143 """

144 int_vals = ['version', 'neuron', 'generation', 'batch', 'population',

145 'mode', 'model_type']

146 float_vals = ['lambda_1', 'lambda_2', 'lambda_e', 'lambda_f', 'lambda_v',

147 'lambda_q', 'lambda_shear', 'force_delta', 'atomic_v', 'zbl']

148 settings = {}

149 with open(filename) as f:

150 for line in f.readlines():

151 # remove comments - throw away everything after a '#'

152 cleaned = line.split('#', 1)[0].strip()

153 flds = cleaned.split()

154 if len(flds) == 0:

155 continue

156 settings[flds[0]] = ' '.join(flds[1:])

157 for key, val in settings.items():

158 if key in int_vals:

159 settings[key] = int(val)

160 elif key in float_vals:

161 settings[key] = float(val)

162 elif key in ['n_max', 'l_max', 'basis_size', 'charge_mode']:

163 settings[key] = [int(v) for v in val.split()]

164 elif key in ['cutoff', 'type_weight']:

165 settings[key] = [float(v) for v in val.split()]

166 elif key == 'type':

167 types = val.split()

168 types[0] = int(types[0])

169 settings[key] = types

170 return settings

171

172

173def read_structures(dirname: str) -> tuple[list[Atoms], list[Atoms]]:

174 """Parses the output files with training and test data from a nep run and returns their

175 content as two lists of structures, representing training and test data, respectively.

176 Target and predicted data are included in the :attr:`info` dict of the :class:`Atoms`

177 objects.

178

179 Parameters

180 ----------

181 dirname

182 Directory from which to read output files.

183

184 """

185 path = join_path(dirname)

186 if not exists(path):

187 raise FileNotFoundError(f'Directory {path} does not exist')

188

189 # fetch model type from nep input file

190 nep_info = read_nepfile(f'{path}/nep.in')

191 model_type = nep_info.get('model_type', 0)

192

193 # set up which files to parse, what dimensions to expect etc

194 # depending on the type of model that is parsed

195 #

196 # The `nep` executable's virial_*.out/stress_*.out/polarizability_*.out

197 # output files use a reduced-6 component order (xx,yy,zz,xy,yz,xz) that

198 # differs from ASE's own Voigt-6 convention (xx,yy,zz,yz,xz,xy) -- see

199 # GPUMD's main_nep/structure.cu `reduced_index` array (shared by the

200 # `virial=` and `pol=` extxyz parsers). Below, virial/stress/

201 # polarizability columns are converted to ASE order immediately upon

202 # parsing, so `structure.info`/`.arrays` always hold ASE-Voigt-ordered

203 # data, matching what `atoms.get_stress()` returns elsewhere in

204 # calorine. BEC's 9 columns are left as is (row-major, see

205 # BEC_FULL9_ORDER) -- BEC is not symmetric, so no Voigt form applies.

206 if model_type == 0:

207 charge_mode = nep_info.get('charge_mode', [0])[0]

208 if charge_mode not in [0, 1, 2]:

209 raise ValueError(f'Unknown charge_mode: {charge_mode}')

210 # files to parse: (sname, size, mandatory, includes_target, per_atom)

211 files_to_parse = [

212 ('energy', 1, True, True, False),

213 ('force', 3, True, True, True),

214 ('virial', 6, True, True, False),

215 ('stress', 6, True, True, False),

216 ]

217 if charge_mode in [1, 2]:

218 # files to parse: (sname, size, includes_target, per_atom)

219 files_to_parse += [

220 ('charge', 1, True, False, True),

221 ('bec', 9, False, True, True),

222 ]

223 elif model_type == 1:

224 # files to parse: (sname, size, includes_target, per_atom)

225 files_to_parse = [('dipole', 3, True, True, False)]

226 if nep_info.get('atomic_v', 0) == 1:

227 files_to_parse = [('dipole', 3, True, True, True)]

228 elif model_type == 2:

229 # files to parse: (sname, size, includes_target, per_atom)

230 files_to_parse = [('polarizability', 6, True, True, False)]

231 if nep_info.get('atomic_v', 0) == 1:

232 files_to_parse = [('polarizability', 6, True, True, True)]

233 else:

234 raise ValueError(f'Unknown model_type: {model_type}')

235

236 # read training and test data

237 structures = {}

238 for stype in ['train', 'test']:

239 filename = join_path(dirname, f'{stype}.xyz')

240 try:

241 structures[stype] = read(filename, format='extxyz', index=':')

242 except FileNotFoundError:

243 warn(f'File {filename} not found.')

244 structures[stype] = []

245 continue

246

247 n_structures = len(structures[stype])

248

249 # loop over files from which to read target data and predictions

250 for sname, size, mandatory, includes_target, per_atom in files_to_parse:

251 infile = f'{sname}_{stype}.out'

252 path = join_path(dirname, infile)

253 if not exists(path):

254 if mandatory:

255 raise FileNotFoundError(f'File {path} does not exist')

256 else:

257 continue

258 ts, ps = _read_data_file(path, includes_target=includes_target)

259

260 if ts is not None:

261 if ts.shape[1] != size:

262 raise ValueError(f'Target data in {infile} has unexpected shape:'

263 f' {ts.shape} (expected: (-1, {size}))')

264 if ps.shape[1] != size:

265 raise ValueError(f'Predicted data in {infile} has unexpected shape:'

266 f' {ps.shape} (expected: (-1, {size}))')

267

268 if sname in ('virial', 'stress', 'polarizability'):

269 ts = nep_reduced6_to_ase_voigt6(ts)

270 ps = nep_reduced6_to_ase_voigt6(ps)

271

272 if per_atom:

273 # data per-atom, e.g., forces, per-atom-virials, Born effective charges ...

274 n_atoms_total = sum([len(s) for s in structures[stype]])

275 if len(ps) != n_atoms_total:

276 raise ValueError(f'Number of atoms in {infile} ({len(ps)})'

277 f' and {stype}.xyz ({n_atoms_total}) inconsistent.')

278 n = 0

279 for structure in structures[stype]:

280 nat = len(structure)

281 if ts is not None:

282 t = np.array(ts[n: n + nat]).reshape(nat, size)

283 structure.new_array(f'{sname}_target', t)

284 p = np.array(ps[n: n + nat]).reshape(nat, size)

285 structure.new_array(f'{sname}_predicted', p)

286 n += nat

287 else:

288 # data per structure, e.g., energy, virials, stress

289 if len(ps) != n_structures:

290 raise ValueError(f'Number of structures in {infile} ({len(ps)})'

291 f' and {stype}.xyz ({n_structures}) inconsistent.')

292 for k, structure in enumerate(structures[stype]):

293 assert ts is not None, 'This should not occur. Please report.'

294 t = ts[k]

295 assert np.shape(t) == (size,)

296 structure.info[f'{sname}_target'] = t

297 p = ps[k]

298 assert np.shape(p) == (size,)

299 structure.info[f'{sname}_predicted'] = p

300

301 # special handling of target data for BECs

302 # If a structure has no 'bec' array in the xyz file, no target BEC data was provided.

303 # In that case nep writes zeros for both predicted and target columns. Replace both

304 # with NaN so callers can easily identify and filter out structures without BEC targets.

305 for s in structures[stype]:

306 if 'bec_target' in s.arrays and 'bec' not in s.arrays:

307 nat = len(s)

308 s.arrays['bec_target'] = np.full((nat, 9), np.nan)

309 s.arrays['bec_predicted'] = np.full((nat, 9), np.nan)

310

311 # special handling of per-atom TNEP

312 # Data has to be loaded as dipole/polarizability since NEP saves them in dipole_*.out

313 # Dipole/polarizability arrays are therefore moved to atomic_v here

314 if nep_info.get('atomic_v', 0) == 1:

315 if model_type == 1:

316 s.new_array('atomic_v_target', s.arrays['dipole_target'])

317 s.new_array('atomic_v_predicted', s.arrays['dipole_predicted'])

318 del s.arrays['dipole_target']

319 del s.arrays['dipole_predicted']

320 if model_type == 2:

321 s.new_array('atomic_v_target', s.arrays['polarizability_target'])

322 s.new_array('atomic_v_predicted', s.arrays['polarizability_predicted'])

323 del s.arrays['polarizability_target']

324 del s.arrays['polarizability_predicted']

325

326 return structures['train'], structures['test']

327

328

329def _read_data_file(

330 path: str,

331 includes_target: bool = True,

332):

333 """Private function that parses *.out files and

334 returns their content for further processing.

335 """

336 with open(path, 'r') as f:

337 lines = f.readlines()

338 target, predicted = [], []

339 for line in lines:

340 flds = line.split()

341 if includes_target:

342 if len(flds) % 2 != 0:

343 raise ValueError(f'Incorrect number of columns in {path} ({len(flds)}).')

344 n = len(flds) // 2

345 predicted.append([float(s) for s in flds[:n]])

346 target.append([float(s) for s in flds[n:]])

347 else:

348 predicted.append([float(s) for s in flds])

349 target = None

350 if target is not None:

351 target = np.array(target)

352 predicted = np.array(predicted)

353 return target, predicted

354

355

356# Maps x/y/z (vectors) and reduced-6 symmetric-tensor component names to

357# their index. virial/stress/polarizability (and per-atom atomic_v when it

358# holds a reduced-6 polarizability) are converted to ASE-Voigt order by

359# read_structures() above, so this mapping uses ASE_VOIGT6_ORDER directly.

360_REDUCED6_MAPPING = {

361 'x': 0, 'y': 1, 'z': 2,

362 **{label: i for i, label in enumerate(ASE_VOIGT6_ORDER)},

363}

364# BEC is a raw, unreduced, generally-asymmetric 9-component per-atom tensor

365# (row-major, see BEC_FULL9_ORDER) -- not a Voigt-reduced 6-component form,

366# so it needs its own, different index scheme.

367_BEC_MAPPING = {label: i for i, label in enumerate(BEC_FULL9_ORDER)}

368

369

370def get_parity_data(

371 structures: list[Atoms],

372 property: str,

373 selection: list[str] = None,

374 flatten: bool = True,

375) -> DataFrame:

376 """Returns the predicted and target energies, forces, virials or stresses

377 from a list of structures in a format suitable for generating parity plots.

378

379 The structures should have been read using :func:`read_structures

380 <calorine.nep.read_structures>`, such that the `info` object is

381 populated with keys of the form `<property>_<type>` where `<property>`

382 is, e.g., `energy` or `force` and `<type>` is one of `predicted` or `target`.

383

384 The resulting parity data is returned as a tuple of dicts, where each entry

385 corresponds to a list.

386

387 Parameters

388 ----------

389 structures

390 List of structures as read with :func:`read_structures <calorine.nep.read_structures>`.

391 property

392 One of `energy`, `force`, `virial`, `stress`, `bec`, `dipole`,

393 `polarizability`, or `atomic_v`.

394 selection

395 A list containing which components to return, and/or the norm.

396 For `force`, `atomic_v` (dipole), `virial`, `stress`, `polarizability`,

397 and `dipole`, possible values are `x`, `y`, `z`, `xx`, `yy`, `zz`,

398 `yz`, `xz`, `xy`, `norm`, `pressure` (the latter only for `stress`).

399 For `bec`, all nine of `xx`, `xy`, `xz`, `yx`, `yy`, `yz`, `zx`,

400 `zy`, `zz` are available (BEC is not symmetric, so unlike the other

401 properties it has no reduced/Voigt form).

402 flatten

403 if True return flattened lists; this is useful for flattening

404 the components of force or virials into a simple list

405 """

406 global_properties = ['energy', 'virial', 'stress', 'polarizability', 'dipole']

407 per_atom_properties = ['force', 'bec', 'atomic_v']

408 if property not in global_properties + per_atom_properties:

409 raise ValueError(

410 '`property` must be one of the following:'

411 + ', '.join(global_properties + per_atom_properties)

412 )

413 if property in ['energy'] and selection:

414 raise ValueError('Selection cannot be applied to scalars.')

415 if property != 'stress' and selection and 'pressure' in selection:

416 raise ValueError(f'Cannot calculate pressure for `{property}`.')

417

418 data = {'predicted': [], 'target': []}

419 if property in ['force', 'bec'] and flatten:

420 size = 3 if property == 'force' else 9

421 data['species'] = []

422 for structure in structures:

423 if 'species' in data:

424 data['species'].extend(np.repeat(structure.symbols, size).tolist())

425 for stype in ['predicted', 'target']:

426 property_with_stype = f'{property}_{stype}'

427 if property in global_properties:

428 if property_with_stype not in structure.info.keys():

429 raise KeyError(f'{property_with_stype} not'

430 ' available in info field of structure')

431 extracted_property = np.array(structure.info[property_with_stype])

432 else:

433 if property_with_stype not in structure.arrays:

434 raise KeyError(f'{property_with_stype} not available in arrays of structure')

435 extracted_property = np.array(structure.arrays[property_with_stype])

436

437 if selection is None or len(selection) == 0:

438 data[stype].append(extracted_property)

439 continue

440

441 if property in ['force', 'bec', 'atomic_v']:

442 extracted_property = extracted_property.T

443 selected_values = []

444 for select in selection:

445 if select == 'norm':

446 if property == 'force':

447 selected_values.append(np.linalg.norm(extracted_property, axis=0))

448 elif property == 'atomic_v':

449 if extracted_property.shape[0] == 3:

450 # per-atom dipole-like vector

451 selected_values.append(

452 np.linalg.norm(extracted_property, axis=0))

453 elif extracted_property.shape[0] == 6:

454 # per-atom polarizability-like tensor in reduced Voigt notation

455 full_tensor = voigt_6_to_full_3x3_stress(extracted_property.T)

456 selected_values.append(

457 np.linalg.norm(full_tensor, axis=(1, 2)))

458 else:

459 raise ValueError(

460 'Cannot handle selection=`norm` with property=`atomic_v`'

461 f' of size {extracted_property.shape[0]}.'

462 )

463 elif property in ['virial', 'stress']:

464 full_tensor = voigt_6_to_full_3x3_stress(extracted_property)

465 selected_values.append(np.linalg.norm(full_tensor))

466 elif property in ['dipole']:

467 selected_values.append(np.linalg.norm(extracted_property))

468 else:

469 raise ValueError(

470 f'Cannot handle selection=`norm` with property=`{property}`.')

471 continue

472

473 if select == 'pressure' and property == 'stress':

474 total_stress = extracted_property

475 selected_values.append(-np.sum(total_stress[:3]) / 3)

476 continue

477

478 mapping = _BEC_MAPPING if property == 'bec' else _REDUCED6_MAPPING

479 if select not in mapping:

480 raise ValueError(f'Selection `{select}` is not allowed.')

481 index = mapping[select]

482 if index >= extracted_property.shape[0]:

483 raise ValueError(

484 f'Selection `{select}` is not compatible with property `{property}`.'

485 )

486 selected_values.append(extracted_property[index])

487

488 data[stype].append(selected_values)

489 if flatten:

490 for stype in ['target', 'predicted']:

491 value = data[stype]

492 if len(np.shape(value[0])) > 0:

493 data[stype] = np.concatenate(value).ravel().tolist()

494 if property in ['force', 'atomic_v']:

495 default_labels = ['x', 'y', 'z']

496 labels = selection if selection else default_labels

497 n = len(data['target']) // len(labels)

498 data['component'] = list(labels) * n

499 elif property in ['virial', 'stress', 'polarizability']:

500 default_labels = list(ASE_VOIGT6_ORDER)

501 labels = selection if selection else default_labels

502 n = len(data['target']) // len(labels)

503 data['component'] = list(labels) * n

504 elif property in ['bec']:

505 default_labels = list(BEC_FULL9_ORDER)

506 labels = selection if selection else default_labels

507 n = len(data['target']) // len(labels)

508 data['component'] = list(labels) * n

509 df = DataFrame(data)

510 # In case of flatten, cast to float64 for compatibility

511 # with e.g. seaborn.

512 # Casting in this way breaks tensorial properties though,

513 # so skip it there.

514 if flatten:

515 df['target'] = df.target.astype('float64')

516 df['predicted'] = df.predicted.astype('float64')

517 return df