/
base.py
485 lines (396 loc) · 16.4 KB
/
base.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
"""Module defining base FHI-aims input set and generator."""
from __future__ import annotations
import copy
import json
import logging
from collections.abc import Iterable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from warnings import warn
import numpy as np
from monty.json import MontyDecoder, MontyEncoder
from pymatgen.core import Molecule, Structure
from pymatgen.io.aims.inputs import AimsControlIn, AimsGeometryIn
from pymatgen.io.aims.parsers import AimsParseError, read_aims_output
from pymatgen.io.core import InputFile, InputGenerator, InputSet
if TYPE_CHECKING:
from collections.abc import Sequence
from pathlib import Path
TMPDIR_NAME: str = "tmpdir"
OUTPUT_FILE_NAME: str = "aims.out"
CONTROL_FILE_NAME: str = "control.in"
PARAMS_JSON_FILE_NAME: str = "parameters.json"
GEOMETRY_FILE_NAME: str = "geometry.in"
DEFAULT_AIMS_PROPERTIES = ("energy", "free_energy", "forces", "stress", "stresses", "dipole", "magmom")
logger = logging.getLogger(__name__)
class AimsInputSet(InputSet):
"""A class to represent a set of Aims inputs."""
def __init__(
self,
parameters: dict[str, Any],
structure: Structure | Molecule,
properties: Sequence[str] = ("energy", "free_energy"),
) -> None:
"""Construct the AimsInputSet.
Args:
parameters (dict[str, Any]): The ASE parameters object for the calculation
structure (Structure or Molecule): The Structure/Molecule objects to
create the inputs for
properties (Sequence[str]): The properties to calculate for the calculation
"""
self._parameters = parameters
self._structure = structure
self._properties = properties
aims_control_in, aims_geometry_in = self.get_input_files()
super().__init__(
inputs={
CONTROL_FILE_NAME: aims_control_in,
GEOMETRY_FILE_NAME: aims_geometry_in,
PARAMS_JSON_FILE_NAME: json.dumps(self._parameters, cls=MontyEncoder),
}
)
def get_input_files(self) -> tuple[str, str]:
"""Get the input file contents for the calculation.
Returns:
tuple[str, str]: The contents of the control.in and geometry.in file
"""
property_flags = {
"forces": "compute_forces",
"stress": "compute_analytical_stress",
"stresses": "compute_heat_flux",
}
updated_params = dict(**self._parameters)
for prop in self._properties:
aims_name = property_flags.get(prop)
if aims_name is not None:
updated_params[aims_name] = True
aims_geometry_in = AimsGeometryIn.from_structure(self._structure)
aims_control_in = AimsControlIn(updated_params)
return (
aims_control_in.get_content(structure=self._structure),
aims_geometry_in.content,
)
@property
def control_in(self) -> str | slice | InputFile:
"""Get the control.in file contents."""
return self[CONTROL_FILE_NAME]
@property
def geometry_in(self) -> str | slice | InputFile:
"""Get the geometry.in file contents."""
return self[GEOMETRY_FILE_NAME]
@property
def params_json(self) -> str | slice | InputFile:
"""Get the JSON representation of the parameters dict."""
return self[PARAMS_JSON_FILE_NAME]
def set_parameters(self, *args, **kwargs) -> dict[str, Any]:
"""Set the parameters object for the AimsTemplate.
This sets the parameters object that is passed to an AimsTemplate and
resets the control.in file
One can pass a dictionary mapping the aims variables to their values or
the aims variables as keyword arguments. A combination of the two
options is also allowed.
Returns:
dict[str, Any]: dictionary with the variables that have been added.
"""
self._parameters.clear()
for arg in args:
self._parameters.update(arg)
self._parameters.update(kwargs)
aims_control_in, _ = self.get_input_files()
self.inputs[CONTROL_FILE_NAME] = aims_control_in
self.inputs[PARAMS_JSON_FILE_NAME] = json.dumps(self._parameters, cls=MontyEncoder)
inputs = {str(key): val for key, val in self.inputs.items()}
self.__dict__.update(inputs)
return self._parameters
def remove_parameters(self, keys: Iterable[str] | str, strict: bool = True) -> dict[str, Any]:
"""Remove the aims parameters listed in keys.
This removes the aims variables from the parameters object.
Args:
keys (Iterable[str] or str): string or list of strings with the names of
the aims parameters to be removed.
strict (bool): whether to raise a KeyError if one of the aims parameters
to be removed is not present.
Returns:
dict[str, Any]: Dictionary with the variables that have been removed.
"""
if isinstance(keys, str):
keys = [keys]
for key in keys:
if key not in self._parameters:
if strict:
raise ValueError(f"{key=} not in {list(self._parameters)=}")
continue
del self._parameters[key]
return self.set_parameters(**self._parameters)
def set_structure(self, structure: Structure | Molecule):
"""Set the structure object for this input set.
Args:
structure (Structure or Molecule): The new Structure or Molecule
for the calculation
"""
self._structure = structure
aims_control_in, aims_geometry_in = self.get_input_files()
self.inputs[GEOMETRY_FILE_NAME] = aims_geometry_in
self.inputs[CONTROL_FILE_NAME] = aims_control_in
inputs = {str(key): val for key, val in self.inputs.items()}
self.__dict__.update(inputs)
@dataclass
class AimsInputGenerator(InputGenerator):
"""
A class to generate Aims input sets.
Parameters
----------
user_params: dict[str, Any]
Updates the default parameters for the FHI-aims calculator
user_kpoints_settings: dict[str, Any]
The settings used to create the k-grid parameters for FHI-aims
"""
user_params: dict[str, Any] = field(default_factory=dict)
user_kpoints_settings: dict[str, Any] = field(default_factory=dict)
def get_input_set( # type: ignore
self,
structure: Structure | Molecule | None = None,
prev_dir: str | Path | None = None,
properties: list[str] | None = None,
) -> AimsInputSet:
"""Generate an AimsInputSet object.
Parameters
----------
structure : Structure or Molecule
Structure or Molecule to generate the input set for.
prev_dir: str or Path
Path to the previous working directory
properties: list[str]
System properties that are being calculated
Returns
-------
The input set for the calculation of structure
"""
prev_structure, prev_parameters, _ = self._read_previous(prev_dir)
structure = structure or prev_structure
if structure is None:
raise ValueError("No structure can be determined to generate the input set")
parameters = self._get_input_parameters(structure, prev_parameters)
properties = self._get_properties(properties, parameters)
return AimsInputSet(parameters=parameters, structure=structure, properties=properties)
@staticmethod
def _read_previous(
prev_dir: str | Path | None = None,
) -> tuple[Structure | Molecule | None, dict[str, Any], dict[str, Any]]:
"""Read in previous results.
Parameters
----------
prev_dir: str or Path
The previous directory for the calculation
"""
prev_structure: Structure | Molecule | None = None
prev_parameters = {}
prev_results: dict[str, Any] = {}
if prev_dir:
# strip hostname from the directory (not good, works only with run_locally.
# Should be checked with Fireworks, will not for sure work with
# jobflow_remote)
split_prev_dir = str(prev_dir).split(":")[-1]
with open(f"{split_prev_dir}/parameters.json") as param_file:
prev_parameters = json.load(param_file, cls=MontyDecoder)
try:
aims_output: Sequence[Structure | Molecule] = read_aims_output(
f"{split_prev_dir}/aims.out", index=slice(-1, None)
)
prev_structure = aims_output[0]
prev_results = prev_structure.properties
prev_results.update(prev_structure.site_properties)
except (IndexError, AimsParseError):
pass
return prev_structure, prev_parameters, prev_results
@staticmethod
def _get_properties(
properties: list[str] | None = None,
parameters: dict[str, Any] | None = None,
) -> list[str]:
"""Get the properties to calculate.
Parameters
----------
properties: list[str]
The currently requested properties
parameters: dict[str, Any]
The parameters for this calculation
Returns
-------
The list of properties to calculate
"""
if properties is None:
properties = ["energy", "free_energy"]
if parameters is None:
return properties
if "compute_forces" in parameters and "forces" not in properties:
properties.append("forces")
if "compute_heat_flux" in parameters and "stresses" not in properties:
properties.append("stress")
properties.append("stresses")
if "stress" not in properties and (
("compute_analytical_stress" in parameters)
or ("compute_numerical_stress" in parameters)
or ("compute_heat_flux" in parameters)
):
properties.append("stress")
return properties
def _get_input_parameters(
self, structure: Structure | Molecule, prev_parameters: dict[str, Any] | None = None
) -> dict[str, Any]:
"""Create the input parameters.
Parameters
----------
structure: Structure | Molecule
The structure or molecule for the system
prev_parameters: dict[str, Any]
The previous calculation's calculation parameters
Returns
-------
The input object
"""
# Get the default configuration
# FHI-aims recommends using their defaults so bare-bones default parameters
parameters: dict[str, Any] = {
"xc": "pbe",
"relativistic": "atomic_zora scalar",
}
# Override default parameters with previous parameters
prev_parameters = {} if prev_parameters is None else copy.deepcopy(prev_parameters)
prev_parameters.pop("relax_geometry", None)
prev_parameters.pop("relax_unit_cell", None)
kpt_settings = copy.deepcopy(self.user_kpoints_settings)
if isinstance(structure, Structure) and "k_grid" in prev_parameters:
density = self.k2d(structure, prev_parameters.pop("k_grid"))
if "density" not in kpt_settings:
kpt_settings["density"] = density
parameter_updates = self.get_parameter_updates(structure, prev_parameters)
parameters = recursive_update(parameters, parameter_updates)
# Override default parameters with user_params
parameters = recursive_update(parameters, self.user_params)
if ("k_grid" in parameters) and ("density" in kpt_settings):
warn(
"WARNING: the k_grid is set in user_params and in the kpt_settings,"
" using the one passed in user_params.",
stacklevel=1,
)
elif isinstance(structure, Structure) and ("k_grid" not in parameters):
density = kpt_settings.get("density", 5.0)
even = kpt_settings.get("even", True)
parameters["k_grid"] = self.d2k(structure, density, even)
elif isinstance(structure, Molecule) and "k_grid" in parameters:
warn("WARNING: removing unnecessary k_grid information", stacklevel=1)
del parameters["k_grid"]
return parameters
def get_parameter_updates(self, structure: Structure | Molecule, prev_parameters: dict[str, Any]) -> dict[str, Any]:
"""Update the parameters for a given calculation type.
Parameters
----------
structure : Structure or Molecule
The system to run
prev_parameters: dict[str, Any]
Previous calculation parameters.
Returns
-------
A dictionary of updates to apply.
"""
return prev_parameters
def d2k(
self,
structure: Structure,
kptdensity: float | list[float] = 5.0,
even: bool = True,
) -> Iterable[float]:
"""Convert k-point density to Monkhorst-Pack grid size.
inspired by [ase.calculators.calculator.kptdensity2monkhorstpack]
Parameters
----------
structure: Structure
Contains unit cell and information about boundary conditions.
kptdensity: float or list of floats
Required k-point density. Default value is 5.0 point per Ang^-1.
even: bool
Round up to even numbers.
Returns
-------
Monkhorst-Pack grid size in all directions
"""
recipcell = structure.lattice.inv_matrix
return self.d2k_recipcell(recipcell, structure.lattice.pbc, kptdensity, even)
def k2d(self, structure: Structure, k_grid: np.ndarray[int]):
"""Generate the kpoint density in each direction from given k_grid.
Parameters
----------
structure: Structure
Contains unit cell and information about boundary conditions.
k_grid: np.ndarray[int]
k_grid that was used.
Returns
-------
Density of kpoints in each direction. result.mean() computes average density
"""
recipcell = structure.lattice.inv_matrix
densities = k_grid / (2 * np.pi * np.sqrt((recipcell**2).sum(axis=1)))
return np.array(densities)
@staticmethod
def d2k_recipcell(
recipcell: np.ndarray,
pbc: Sequence[bool],
kptdensity: float | Sequence[float] = 5.0,
even: bool = True,
) -> Sequence[int]:
"""Convert k-point density to Monkhorst-Pack grid size.
Parameters
----------
recipcell: Cell
The reciprocal cell
pbc: Sequence[bool]
If element of pbc is True then system is periodic in that direction
kptdensity: float or list[floats]
Required k-point density. Default value is 3.5 point per Ang^-1.
even: bool
Round up to even numbers.
Returns
-------
Monkhorst-Pack grid size in all directions
"""
if not isinstance(kptdensity, Iterable):
kptdensity = 3 * [float(kptdensity)]
kpts: list[int] = []
for i in range(3):
if pbc[i]:
k = 2 * np.pi * np.sqrt((recipcell[i] ** 2).sum()) * float(kptdensity[i])
if even:
kpts.append(2 * int(np.ceil(k / 2)))
else:
kpts.append(int(np.ceil(k)))
else:
kpts.append(1)
return kpts
def recursive_update(dct: dict, up: dict) -> dict:
"""
Update a dictionary recursively and return it.
Parameters
----------
dct: Dict
Input dictionary to modify
up: Dict
Dictionary of updates to apply
Returns
-------
The updated dictionary.
Example
-------
d = {'activate_hybrid': {"hybrid_functional": "HSE06"}}
u = {'activate_hybrid': {"cutoff_radius": 8}}
yields {'activate_hybrid': {"hybrid_functional": "HSE06", "cutoff_radius": 8}}}
"""
for key, val in up.items():
if isinstance(val, dict):
dct[key] = recursive_update(dct.get(key, {}), val)
elif key == "output" and isinstance(val, list): # for all other keys the list addition is not needed (I guess)
old_v = dct.get(key, [])
dct[key] = old_v + val
else:
dct[key] = val
return dct