/
core.py
123 lines (96 loc) · 3.41 KB
/
core.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
"""Module defining core FHI-aims input set generators."""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
from pymatgen.core import Structure
from pymatgen.io.aims.sets.base import AimsInputGenerator
if TYPE_CHECKING:
from pymatgen.core import Molecule
@dataclass
class StaticSetGenerator(AimsInputGenerator):
"""Common class for ground-state generators.
Parameters
----------
calc_type: str
The type of calculation
"""
calc_type: str = "static"
def get_parameter_updates(self, structure: Structure | Molecule, prev_parameters: dict[str, Any]) -> dict[str, Any]:
"""Get the parameter updates for the calculation.
Parameters
----------
structure: Structure or Molecule
The structure to calculate the bands for
prev_parameters: Dict[str, Any]
The previous parameters
Returns
-------
The updated for the parameters for the output section of FHI-aims
"""
return prev_parameters
@dataclass
class RelaxSetGenerator(AimsInputGenerator):
"""
Class to generate FHI-aims relax sets.
I.e., sets for optimization of internal and lattice coordinates.
Parameters
----------
calc_type: str
The type of calculation
relax_cell: bool
If True then relax the unit cell from the structure
max_force: float
Maximum allowed force in the calculation
method: str
Method used for the geometry optimization
"""
calc_type: str = "relaxation"
relax_cell: bool = True
max_force: float = 1e-3
method: str = "trm"
def get_parameter_updates(self, structure: Structure | Molecule, prev_parameters: dict[str, Any]) -> dict:
"""Get the parameter updates for the calculation.
Parameters
----------
structure: Structure or Molecule
The structure to calculate the bands for
prev_parameters: Dict[str, Any]
The previous parameters
Returns
-------
The updated for the parameters for the output section of FHI-aims
"""
updates = {"relax_geometry": f"{self.method} {self.max_force:e}"}
if isinstance(structure, Structure) and self.relax_cell:
updates["relax_unit_cell"] = "full"
elif isinstance(structure, Structure):
updates["relax_unit_cell"] = "none"
return updates
@dataclass
class SocketIOSetGenerator(AimsInputGenerator):
"""Class to generate FHI-aims input sets for running with the socket.
Parameters
----------
calc_type: str
The type of calculation
host: str
The hostname for the server the socket is on
port: int
The port the socket server is listening on
"""
calc_type: str = "multi_scf"
host: str = "localhost"
port: int = 12345
def get_parameter_updates(self, structure: Structure | Molecule, prev_parameters: dict[str, Any]) -> dict:
"""Get the parameter updates for the calculation.
Parameters
----------
structure: Structure or Molecule
The structure to calculate the bands for
prev_parameters: Dict[str, Any]
The previous parameters
Returns
-------
The updated for the parameters for the output section of FHI-aims
"""
return {"use_pimd_wrapper": (self.host, self.port)}