def __init__(self,
calc,
executor,
name="default",
task_dir="tasks",
scheduler="slurm",
allow_submit=True,
**kwargs):
if calc.name not in self.implemented_calculators:
raise CalculatorSetupError(
"Wrapper is not implemented for calculator {}.".format(
calc.name))
if scheduler not in self.implemented_schedulers:
raise CalculatorSetupError(
"Scheduler {} is not implemeted.".format(scheduler))
Calculator.__init__(self)
self.executor = executor
self.implemented_properties = calc.implemented_properties
self.default_parameters = calc.default_parameters
self.calc = calc
self.directory = calc.directory
self.name = name
self.task_dir = task_dir
self.scheduler = scheduler
self.allow_submit = allow_submit
self.kwargs = kwargs
self.job = None
self.path = None
# Here we correct a typo in old version of VASP2 calculator where was
# 'verison' in initialization instead of 'version'. It didn't allow to
# save calculator's state when it is not finished.
self.calc.version = None
def submit(self):
"""Submits the job to the cluster"""
if self.job is not None:
if self.allow_submit == True:
self.executor.mkdirs(self.path)
for filename in os.listdir(self.calc.directory):
self.executor.put(
os.path.join(self.calc.directory, filename),
os.path.join(self.path, filename),
)
stdout, stderr = self.executor.execute(
"cd {}; {} {}.{}".format(self.path,
self.job.scheduler.cmd_run,
self.name, self.scheduler))
if len(stderr) != 0:
raise CalculatorSetupError(stderr[0])
response = stdout[0]
if "Submitted batch job " in response:
print("Job {} is submitted to the cluster.".format(
self.name))
self.job.status = "submitted"
self.job.id = int(
response.replace("Submitted batch job ", ""))
self.executor.jobs.append(self.job)
else:
print("Failed to submit the job {} to the cluster.".format(
self.name))
else:
print(
"Job {} was created but not submitted (allow_submit = False)."
.format(self.name))
else:
raise ValueError("There is no job to submit.")
def calculate(self,
atoms=None,
properties=['energy'],
system_changes=all_changes,
symmetry='c1'):
Calculator.calculate(self, atoms=atoms)
if self.atoms is None:
raise CalculatorSetupError('An Atoms object must be provided to '
'perform a calculation')
atoms = self.atoms
pos = torch.tensor(atoms.positions).double().to("cpu")
cell = atoms.cell.tolist()
cell = torch.tensor([cell[0][0], cell[1][1],
cell[2][2]]).double().to("cpu")
energy = self.evaluator.compute(pos, cell)
# Do the calculations
if 'forces' in properties:
# energy comes for free
self.results['energy'] = energy
# convert to eV/A
# also note that the gradient is -1 * forces
self.results['forces'] = self.evaluator.forces.cpu().numpy()
elif 'energy' in properties:
# convert to eV
self.results['energy'] = energy
def check_atoms_type(atoms: Atoms) -> None:
"""Check that the passed atoms object is in fact an Atoms object.
Raises CalculatorSetupError.
"""
if not isinstance(atoms, Atoms):
raise CalculatorSetupError(
('Expected an Atoms object, '
'instead got object of type {}'.format(type(atoms))))
def calculate(self,
atoms=None,
properties=['energy'],
system_changes=all_changes,
symmetry='c1'):
Calculator.calculate(self, atoms=atoms)
if self.atoms is None:
raise CalculatorSetupError(
'An Atoms object must be provided to perform a calculation')
atoms = self.atoms
if atoms.get_initial_magnetic_moments().any():
self.parameters['reference'] = 'uhf'
self.parameters['multiplicity'] = None
# this inputs all the settings into psi4
self.set_psi4(atoms=atoms)
self.psi4.core.set_output_file(self.label + '.dat', False)
# Set up the method
method = self.parameters['method']
basis = self.parameters['basis']
# Do the calculations
if 'forces' in properties:
grad, wf = self.psi4.driver.gradient('{}/{}'.format(method, basis),
return_wfn=True)
# energy comes for free
energy = wf.energy()
self.results['energy'] = energy * Hartree
# convert to eV/A
# also note that the gradient is -1 * forces
self.results['forces'] = -1 * np.array(grad) * Hartree / Bohr
elif 'energy' in properties:
energy = self.psi4.energy('{}/{}'.format(method, basis),
molecule=self.molecule)
# convert to eV
self.results['energy'] = energy * Hartree
# dump the calculator info to the psi4 file
save_atoms = self.atoms.copy()
#del save_atoms.calc
# use io.write to encode atoms
with StringIO() as f:
io.write(f, save_atoms, format='json')
json_atoms = f.getvalue()
# convert forces to list for json storage
save_results = self.results.copy()
if 'forces' in save_results:
save_results['forces'] = save_results['forces'].tolist()
save_dict = {
'parameters': self.parameters,
'results': save_results,
'atoms': json_atoms
}
self.psi4.core.print_out('!ASE Information\n')
self.psi4.core.print_out(json.dumps(save_dict))
self.psi4.core.print_out('!')
def check_pbc(atoms: Atoms) -> None:
"""Check if any boundaries are not PBC, as VASP
cannot handle non-PBC.
Raises CalculatorSetupError.
"""
if not atoms.pbc.all():
raise CalculatorSetupError(
"Vasp cannot handle non-periodic boundaries. "
"Please enable all PBC, e.g. atoms.pbc=True")
def check_cell(atoms: Atoms) -> None:
"""Check if there is a zero unit cell.
Raises CalculatorSetupError if the cell is wrong.
"""
if atoms.cell.rank < 3:
raise CalculatorSetupError(
"The lattice vectors are zero! "
"This is the default value - please specify a "
"unit cell.")
def calculate(self, atoms=None, properties=['energy'],
system_changes=all_changes):
"""Actual calculation."""
Calculator.calculate(self, atoms, properties, system_changes)
if self.library is None:
raise CalculatorSetupError(
"No shared object bound to {} calculator!"
.format(self.__class__.__name__))
def set(self, **kwargs):
"""Set parameters like set(key1=value1, key2=value2, ...)."""
msg = '"%s" is not a known keyword for the CP2K calculator. ' \
'To access all features of CP2K by means of an input ' \
'template, consider using the "inp" keyword instead.'
for key in kwargs:
if key not in self.default_parameters:
raise CalculatorSetupError(msg % key)
changed_parameters = Calculator.set(self, **kwargs)
if changed_parameters:
self.reset()
def get_status(self):
"""Gets status of job on the server"""
stdout, stderr = self.executor.execute(
"sacct -X -p -j {} --format=state,elapsed,start,end".format(
self.job.id))
if len(stderr) != 0:
raise CalculatorSetupError(stderr[0])
status, dur_time, start, finish = stdout[1].strip()[:-1].split("|")
if "cancelled" in status.lower():
status = "cancelled"
if dur_time == "00:00:00":
dur_time = None
if start == "Unknown":
start = None
if finish == "Unknown":
finish = None
return status.lower(), dur_time, start, finish
def make_command(self, command=None):
"""Return command if one is passed, otherwise try to find
ASE_VASP_COMMAND, VASP_COMMAND or VASP_SCRIPT.
If none are set, a CalculatorSetupError is raised"""
if command:
cmd = command
else:
# Search for the environment commands
for env in self.env_commands:
if env in os.environ:
cmd = os.environ[env].replace('PREFIX', self.prefix)
if env == 'VASP_SCRIPT':
# Make the system python exe run $VASP_SCRIPT
exe = sys.executable
cmd = ' '.join([exe, cmd])
break
else:
msg = ('Please set either command in calculator'
' or one of the following environment '
'variables (prioritized as follows): {}').format(
', '.join(self.env_commands))
raise CalculatorSetupError(msg)
return cmd
def write(self, label=None):
if label is None:
label = self.label
if self.atoms is None:
raise CalculatorSetupError("An Atoms object must be present "
"in the calculator to use write")
# Normalise a few parameters, otherwise json serialisation has trouble
if isinstance(self.parameters["kpts"], np.ndarray):
self.parameters["kpts"] = self.parameters["kpts"].tolist()
with io.StringIO() as fp:
ase.io.write(fp, self.atoms.copy(), format="json")
atoms_json = fp.getvalue()
with open(label + ".json", "w") as fp:
save_dict = {
"parameters": self.parameters,
"results": self.results,
"atoms": atoms_json,
"scfres": self.scfres,
}
json.dump(save_dict, fp)
def calculate(self,
atoms=None,
properties=["energy"],
system_changes=all_changes):
"""Do a calculation in the specified directory on server.
This will generate the necessary input files, send them to the cluster
and then execute. After execution, the output files are received
back from the server and the energy, forces. etc. are read."""
if self.job is None:
# Note: in order to do the structural changes to the initial Atoms object
# we can write:
#
# initial_atoms.calc = calc # optional, but allows to calculate
# # properties afterwards
# calc.calculate(initial_atoms)
#
# If we are not going to change the structure it is better to use
# the following construction:
#
# initial_atoms.calc = calc
# initial_atoms.get_energy()
#
Calculator.calculate(self.calc, atoms, properties, system_changes)
# We intentianally keep self.atoms = None until the calculation is done.
# It makes get_property method to go back to the current method because
# check_state will always return all_changes.
#
# self.atoms = self.calc.atoms
#
self.calc.check_cell() # Check for zero-length lattice vectors
self.calc._xml_data = None # Reset the stored data
self.calc.write_input(self.calc.atoms, properties, system_changes)
start_time = datetime.now()
curr_task_dir = self.name + start_time.strftime("-%Y-%m-%d-")
abs_task_dir = os.path.join(self.executor.home, self.task_dir)
self.executor.mkdirs(abs_task_dir)
stdout, stderr = self.executor.execute(
"cd {}; ls -d {}* | sort -V | tail -1".format(
abs_task_dir, curr_task_dir))
if len(stderr) != 0:
if "No" in stderr[0]: # file pattern doesn't exist
curr_task_dir += "1"
else:
raise CalculatorSetupError(stderr[0])
if len(stdout) != 0:
curr_task_dir += str(int(stdout[0].split("-")[-1]) + 1)
self.path = os.path.join(self.executor.home, self.task_dir,
curr_task_dir)
# we attempt to set some parameters automatically
if "modulelist" not in self.kwargs:
self.kwargs["modulelist"] = self.executor.find_modules(
["mkl", "impi", "vasp"])
if "command" not in self.kwargs:
self.kwargs["command"] = "mpirun vasp_std"
if "njobs" not in self.kwargs:
self.kwargs["njobs"] = 8
if "memory" not in self.kwargs:
self.kwargs["memory"] = 10 * GiB
self.job = Job(name=self.name,
path=self.path,
scheduler=self.scheduler,
host=self.executor.host,
status="not submitted",
**self.kwargs)
self.job.write_job(
os.path.join(self.calc.directory,
self.name + "." + self.scheduler))
self.submit()
elif self.job.status == "not submitted":
self.submit()
elif self.job.status != "completed":
status, dur_time, start, finish = self.get_status()
self.job.status = status
self.job.dur_time = dur_time
self.job.start = start
self.job.finish = finish
if self.job.status != "completed":
print(
"Job {} is {} on the cluster.".format(
self.name, self.job.status),
end="",
)
if self.job.status == "running":
print(" ({})".format(self.job.dur_time))
else:
print()
else:
self.executor.get(self.path, self.calc.directory)
# we can loose some parameters when we save results and job is not finished
if not hasattr(self.calc, "resort"):
self.initialize(atoms)
self.calc.update_atoms(atoms)
self.calc.read_results()
self.results = self.calc.results
# Now we are not going back to the current method
self.atoms = self.calc.atoms
# Here we handle the case where calculator is loaded from json and the atoms object,
# which is passed to the calculator, needs to be changed
elif self.job.status == "completed":
if atoms is not None and self.atoms is not None:
atoms.positions = self.atoms.positions
atoms.cell = self.atoms.cell
# Hacky way to ignore the 'not present' error if calculation is not ready
if self.job is None or (self.job is not None
and self.job.status != "completed"):
for prop in properties:
self.results[prop] = None
# all of these properties are present anyway
vasp_properties = [
"energy",
"forces",
"stress",
"free_energy",
"magmom",
"magmoms",
"dipole",
"nbands",
]
for prop in vasp_properties:
self.results[prop] = None
def check_parameters(self):
"""sanity check for the parameters to be verified"""
check_failed = [key for key in self.verify if self.parameters[key] < 0]
if check_failed:
raise CalculatorSetupError('Method parameters {} cannot be negative!'
.format(check_failed))