-
Notifications
You must be signed in to change notification settings - Fork 840
/
core.py
281 lines (221 loc) · 9.54 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
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
"""
This module defines the abstract interface for reading and writing calculation
inputs in pymatgen. The interface comprises a 3-tiered hierarchy of classes.
1. An InputFile object represents the contents of a single input file, e.g.
the INCAR. This class standardizes file read and write operations.
2. An InputSet is a dict-like container that maps filenames (keys) to file
contents (either strings or InputFile objects). This class provides a standard
write_input() method.
3. InputGenerator classes implement a get_input_set method that, when provided
with a structure, return an InputSet object with all parameters set correctly.
Calculation input files can be written to disk with the write_inputs method.
If you want to implement a new InputGenerator, please take note of the following:
1. You must implement a get_input_set method that returns an InputSet
2. All customization of calculation parameters should be done in the __init__
method of the InputGenerator. The idea is that the generator contains
the "recipe", but nothing that is specific to a particular system. get_input_set
takes system-specific information (such as structure) and applies the recipe.
3. All InputGenerator must save all supplied args and kwargs as instance variables.
E.g., self.my_arg = my_arg and self.kwargs = kwargs in the __init__. This
ensures the as_dict and from_dict work correctly.
"""
from __future__ import annotations
import abc
import os
from collections.abc import Iterator, MutableMapping
from pathlib import Path
from zipfile import ZipFile
import numpy as np
from monty.io import zopen
from monty.json import MSONable
__author__ = "Ryan Kingsbury"
__email__ = "RKingsbury@lbl.gov"
__status__ = "Development"
__date__ = "October 2021"
class InputFile(MSONable):
"""
Abstract base class to represent a single input file. Note that use of this class
is optional; it is possible create an InputSet that does not rely on underlying
InputFile objects.
All InputFile classes must implement a get_string method, which is called by
write_file.
If InputFile classes implement an __init__ method, they must assign all arguments
to __init__ as attributes.
"""
@abc.abstractmethod
def get_str(self) -> str:
"""Return a string representation of an entire input file."""
@np.deprecate(message="Use get_str instead")
@abc.abstractmethod
def get_string(self) -> str:
"""Return a string representation of an entire input file."""
def write_file(self, filename: str | Path) -> None:
"""
Write the input file.
Args:
filename: The filename to output to, including path.
"""
filename = filename if isinstance(filename, Path) else Path(filename)
with zopen(filename, "wt") as file:
file.write(self.get_str())
@classmethod
@np.deprecate(message="Use from_str instead")
@abc.abstractmethod
def from_string(cls, contents: str) -> InputFile:
"""
Create an InputFile object from a string.
Args:
contents: The contents of the file as a single string
Returns:
InputFile
"""
@classmethod
@abc.abstractmethod
def from_str(cls, contents: str) -> InputFile:
"""
Create an InputFile object from a string.
Args:
contents: The contents of the file as a single string
Returns:
InputFile
"""
@classmethod
def from_file(cls, path: str | Path):
"""
Creates an InputFile object from a file.
Args:
path: Filename to read, including path.
Returns:
InputFile
"""
filename = path if isinstance(path, Path) else Path(path)
with zopen(filename, "rt") as f:
return cls.from_str(f.read())
def __str__(self) -> str:
return self.get_str()
class InputSet(MSONable, MutableMapping):
"""
Abstract base class for all InputSet classes. InputSet are dict-like
containers for all calculation input data.
Since InputSet inherits dict, it can be instantiated in the same manner,
or a custom __init__ can be provided. Either way, `self` should be
populated with keys that are filenames to be written, and values that are
InputFile objects or strings representing the entire contents of the file.
All InputSet must implement from_directory. Implementing the validate method
is optional.
"""
def __init__(self, inputs: dict[str | Path, str | InputFile] | None = None, **kwargs):
"""
Instantiate an InputSet.
Args:
inputs: The core mapping of filename: file contents that defines the InputSet data.
This should be a dict where keys are filenames and values are InputFile objects
or strings representing the entire contents of the file. If a value is not an
InputFile object nor a str, but has a __str__ method, this str representation
of the object will be written to the corresponding file. This mapping will
become the .inputs attribute of the InputSet.
**kwargs: Any kwargs passed will be set as class attributes e.g.
InputSet(inputs={}, foo='bar') will make InputSet.foo == 'bar'.
"""
self.inputs = inputs or {}
self._kwargs = kwargs
self.__dict__.update(**kwargs)
def __getattr__(self, key):
# allow accessing keys as attributes
if key in self._kwargs:
return self.get(key)
raise AttributeError(f"'{type(self).__name__}' object has no attribute {key!r}")
def __copy__(self) -> InputSet:
cls = self.__class__
new_instance = cls.__new__(cls)
for k, v in self.__dict__.items():
setattr(new_instance, k, v)
return new_instance
def __deepcopy__(self, memo: dict[int, InputSet]) -> InputSet:
import copy
cls = self.__class__
new_instance = cls.__new__(cls)
memo[id(self)] = new_instance
for k, v in self.__dict__.items():
setattr(new_instance, k, copy.deepcopy(v, memo))
return new_instance
def __len__(self) -> int:
return len(self.inputs)
def __iter__(self) -> Iterator[str | Path]:
return iter(self.inputs)
def __getitem__(self, key) -> str | InputFile | slice:
return self.inputs[key]
def __setitem__(self, key: str | Path, value: str | InputFile) -> None:
self.inputs[key] = value
def __delitem__(self, key: str | Path) -> None:
del self.inputs[key]
def write_input(
self,
directory: str | Path,
make_dir: bool = True,
overwrite: bool = True,
zip_inputs: bool = False,
):
"""
Write Inputs to one or more files.
Args:
directory: Directory to write input files to
make_dir: Whether to create the directory if it does not already exist.
overwrite: Whether to overwrite an input file if it already exists.
Additional kwargs are passed to generate_inputs
zip_inputs: If True, inputs will be zipped into a file with the
same name as the InputSet (e.g., InputSet.zip)
"""
path = directory if isinstance(directory, Path) else Path(directory)
for fname, contents in self.inputs.items():
file_path = path / fname
if not path.exists() and make_dir:
path.mkdir(parents=True, exist_ok=True)
if file_path.exists() and not overwrite:
raise FileExistsError(fname)
file_path.touch()
# write the file
if isinstance(contents, InputFile):
contents.write_file(file_path)
else:
with zopen(file_path, "wt") as f:
f.write(str(contents))
if zip_inputs:
filename = path / f"{type(self).__name__}.zip"
with ZipFile(filename, "w") as zip_file:
for fname in self.inputs:
file_path = path / fname
try:
zip_file.write(file_path)
os.remove(file_path)
except FileNotFoundError:
pass
@classmethod
def from_directory(cls, directory: str | Path):
"""
Construct an InputSet from a directory of one or more files.
Args:
directory: Directory to read input files from
"""
raise NotImplementedError(f"from_directory has not been implemented in {cls}")
def validate(self) -> bool:
"""
A place to implement basic checks to verify the validity of an
input set. Can be as simple or as complex as desired.
Will raise a NotImplementedError unless overloaded by the inheriting class.
"""
raise NotImplementedError(f".validate() has not been implemented in {self.__class__}")
class InputGenerator(MSONable):
"""
InputGenerator classes serve as generators for Input objects. They contain
settings or sets of instructions for how to create Input from a set of
coordinates or a previous calculation directory.
"""
@abc.abstractmethod
def get_input_set(self) -> InputSet:
"""
Generate an InputSet object. Typically the first argument to this method
will be a Structure or other form of atomic coordinates.
"""
class ParseError(SyntaxError):
"""This exception indicates a problem was encountered during parsing due to unexpected formatting."""