def __init__(self,
atoms=None,
restart=None,
directory='.',
label='vasp',
ignore_bad_restart_file=False,
command=None,
txt='vasp.out',
**kwargs):
self._atoms = None
self.results = {}
# Initialize parameter dictionaries
GenerateVaspInput.__init__(self)
self._store_param_state() # Initialize an empty parameter state
# Store atoms objects from vasprun.xml here - None => uninitialized
self._xml_data = None
Calculator.__init__(self,
restart=restart,
ignore_bad_restart_file=ignore_bad_restart_file,
label=label,
directory=directory,
atoms=atoms,
**kwargs)
self.command = command
self._txt = None
self.txt = txt # Set the output txt stream
self.version = None
def __init__(self,
atoms=None,
restart=None,
directory='',
label='vasp',
ignore_bad_restart_file=False,
command=None,
txt=None,
**kwargs):
# Initialize parameter dictionaries
GenerateVaspInput.__init__(self)
self._store_param_state() # Initialize an empty parameter state
# Store atoms objects from vasprun.xml here - None => uninitialized
self._xml_data = None
label = os.path.join(directory, label)
if restart is True:
# We restart in the label directory
restart = label
Calculator.__init__(self,
restart=restart,
ignore_bad_restart_file=ignore_bad_restart_file,
label=label,
atoms=atoms,
**kwargs)
self.command = command
self.set_txt(txt) # Set the output txt stream
self.verison = None
def write_input(self, atoms, properties=None, system_changes=None):
"""Write VASP inputfiles, INCAR, KPOINTS and POTCAR"""
# Create the folders where we write the files, if we aren't in the
# current working directory.
if self.directory != os.curdir and not os.path.isdir(self.directory):
os.makedirs(self.directory)
self.initialize(atoms)
GenerateVaspInput.write_input(self, atoms, directory=self.directory)
def write_input(self, atoms, properties=['energies'],
system_changes=all_changes):
"""Write VASP inputfiles, INCAR, KPOINTS and POTCAR"""
# Create the folders where we write the files, if we aren't in the
# current working directory.
FileIOCalculator.write_input(self, atoms, properties, system_changes)
self.initialize(atoms)
GenerateVaspInput.write_input(self, atoms, directory=self.directory)
def __init__(self,
atoms=None,
restart=None,
directory='.',
label='vasp',
ignore_bad_restart_file=Calculator._deprecated,
command=None,
txt='vasp.out',
**kwargs):
self._atoms = None
self.results = {}
# Initialize parameter dictionaries
GenerateVaspInput.__init__(self)
self._store_param_state() # Initialize an empty parameter state
# Store calculator from vasprun.xml here - None => uninitialized
self._xml_calc = None
# Set directory and label
self.directory = directory
if '/' in label:
warn(
'Specifying directory in "label" is deprecated, use "directory" instead.',
np.VisibleDeprecationWarning)
if self.directory != '.':
raise ValueError('Directory redundantly specified though '
'directory="{}" and label="{}". '
'Please omit "/" in label.'.format(
self.directory, label))
self.label = label
else:
self.prefix = label # The label should only contain the prefix
if isinstance(restart, bool):
if restart is True:
restart = self.label
else:
restart = None
Calculator.__init__(
self,
restart=restart,
ignore_bad_restart_file=ignore_bad_restart_file,
# We already, manually, created the label
label=self.label,
atoms=atoms,
**kwargs)
self.command = command
self._txt = None
self.txt = txt # Set the output txt stream
self.version = None
def _vaspinput_factory(atoms, **kwargs) -> GenerateVaspInput:
mocker = mock.Mock()
inputs = GenerateVaspInput()
inputs.set(**kwargs)
inputs._build_pp_list = mocker(return_value=None) # type: ignore
inputs.initialize(atoms)
return inputs
def set(self, **kwargs):
"""Override the set function, to test for changes in the
Vasp Calculator, then call the create_input.set()
on remaining inputs for VASP specific keys.
Allows for setting ``label``, ``directory`` and ``txt``
without resetting the results in the calculator.
"""
changed_parameters = {}
if 'label' in kwargs:
label = kwargs.pop('label')
self.set_label(label)
if 'directory' in kwargs:
# If we explicitly set directory, overwrite the one in label.
# XXX: Should we just raise an error here if clash?
directory = kwargs.pop('directory')
label = os.path.join(directory, self.prefix)
self.set_label(label)
if 'txt' in kwargs:
txt = kwargs.pop('txt')
self.set_txt(txt)
if 'atoms' in kwargs:
atoms = kwargs.pop('atoms')
self.set_atoms(atoms) # Resets results
if 'command' in kwargs:
self.command = kwargs.pop('command')
changed_parameters.update(Calculator.set(self, **kwargs))
# We might at some point add more to changed parameters, or use it
if changed_parameters:
self.results.clear() # We don't want to clear atoms
if kwargs:
# If we make any changes to Vasp input, we always reset
GenerateVaspInput.set(self, **kwargs)
self.results.clear()
def set(self, **kwargs):
"""Override the set function, to test for changes in the
Vasp Calculator, then call the create_input.set()
on remaining inputs for VASP specific keys.
Allows for setting ``label``, ``directory`` and ``txt``
without resetting the results in the calculator.
"""
changed_parameters = {}
if 'label' in kwargs:
self.label = kwargs.pop('label')
if 'directory' in kwargs:
# str() call to deal with pathlib objects
self.directory = str(kwargs.pop('directory'))
if 'txt' in kwargs:
self.txt = kwargs.pop('txt')
if 'atoms' in kwargs:
atoms = kwargs.pop('atoms')
self.atoms = atoms # Resets results
if 'command' in kwargs:
self.command = kwargs.pop('command')
changed_parameters.update(Calculator.set(self, **kwargs))
# We might at some point add more to changed parameters, or use it
if changed_parameters:
self.results.clear() # We don't want to clear atoms
if kwargs:
# If we make any changes to Vasp input, we always reset
GenerateVaspInput.set(self, **kwargs)
self.results.clear()
def create_job(
given_params,
main_path="/home/al_mlp/examples/main/al_main.py",
sample_job_spec_path="sample_job_spec.yml",
sample_config_path=None,
images=None,
local=False,
):
# generate the params dictionary from the given config if applicable
params = given_params
if sample_config_path is not None:
with open(sample_config_path, "r") as sample_config:
base_params = yaml.safe_load(sample_config)
params = merge_dict(base_params, given_params)
# get the job name
job_name = (
params.get("learner", {}).get("wandb_init", {}).get("name", "default_job_name")
) # TODO
basedir = os.getcwd()
subdir = "/"
subdir += job_name
subdir += "_0"
# change to new directory created inside the current working directory
if not os.path.isdir(basedir + subdir):
os.mkdir(basedir + subdir)
else:
while os.path.isdir(basedir + subdir):
i = int(subdir[subdir.rindex("_") + 1 :])
subdir = subdir[: subdir.rindex("_") + 1] + str(i + 1)
os.mkdir(basedir + subdir)
os.chdir(basedir + subdir)
# if given some images to pretrain on, make an ase_db and save the link in the config
if images is not None:
images_path = basedir + subdir + "/" + "pretrain_images.db"
params["links"]["images_path"] = images_path
with connect(images_path) as pretrain_db:
for image in images:
pretrain_db.write(image)
# if given incar in links, change the vasp params to match the incar
if "incar" in params["links"]:
vasp_input = GenerateVaspInput()
vasp_input.atoms = None
vasp_input.read_incar(params["links"]["incar"])
if "kpoints" in params["links"]:
vasp_input.read_kpoints(params["links"]["kpoints"])
params["vasp"] = vasp_input.todict()
if "kpts" in params["vasp"]:
kpts = [float(i) for i in params["vasp"]["kpts"]]
params["vasp"]["kpts"] = kpts
else:
params["vasp"].pop("kpts")
if "gga" not in params["vasp"]:
params["vasp"]["gga"] = "PE" # defaults to PE for oxide
params["vasp"]["nsw"] = 0
params["vasp"]["ibrion"] = -1
params["vasp"]["lreal"] = "Auto"
# create the new config in the new directory
config_path = basedir + subdir + "/" + job_name + "_config.yml"
with open(config_path, "w") as config_file:
yaml.dump(params, config_file, default_flow_style=False)
# load the sample_job_spec.yml
with open(sample_job_spec_path, "r") as sample_job_spec:
job_spec = yaml.safe_load(sample_job_spec)
# create the new job_spec dictionary from the sample_job_spec.yml
job_spec["metadata"]["name"] = "job-" + job_name.replace("_", "-")
job_spec["spec"]["template"]["spec"]["containers"][0]["name"] = job_name.replace(
"_", "-"
)
if "NAMESPACE" in os.environ:
namespace = os.environ["NAMESPACE"]
job_spec["metadata"]["namespace"] = namespace
if "VOLUME" in os.environ:
volume = os.environ["VOLUME"]
job_spec["spec"]["template"]["spec"]["containers"][0]["volumeMounts"][0][
"name"
] = volume
job_spec["spec"]["template"]["spec"]["volumes"][0]["name"] = volume
job_spec["spec"]["template"]["spec"]["volumes"][0]["persistentVolumeClaim"][
"claimName"
] = volume
args_string = job_spec["spec"]["template"]["spec"]["containers"][0]["args"][0]
args_string = (
args_string[: args_string.rindex("python")]
+ "python "
+ main_path
+ " --config-yml "
+ config_path
+ " 2>&1 | tee "
+ basedir
+ subdir
+ "/run_logs.txt"
)
job_spec["spec"]["template"]["spec"]["containers"][0]["args"][0] = args_string
# create the new job_spec.yml from the job_spec dictionary
job_spec_path = basedir + subdir + "/" + job_name + "_spec.yml"
with open(job_spec_path, "w") as job_spec_file:
yaml.dump(job_spec, job_spec_file, default_flow_style=False)
if local is False:
# call kubectl on the job_spec.yml
run_result = subprocess.run(["kubectl", "apply", "-f", job_spec_path])
print(
"Executed job " + job_name + " with exit code " + str(run_result.returncode)
)
# change back to original working directory
os.chdir(basedir)
return config_path