from scm.plams.core.functions import log
from scm.plams.interfaces.adfsuite.ams import AMSJob
from scm.plams.core.settings import Settings
from scm.plams.recipes.adffragment import ADFFragmentJob, ADFFragmentResults
from scm.plams.recipes._utils import to_plams_settings
from scm.plams.tools.units import Units
from scm.plams.core.errors import FileError
from scm.plams.core.functions import add_to_instance
from scm.plams.core.enums import JobStatus
from os.path import join as opj
from os.path import abspath, isdir, basename, relpath
from os import symlink
from typing import Dict, List, Union, Optional
__all__ = ["BANDFragmentJob", "BANDFragmentResults"]
[docs]class BANDFragmentResults(ADFFragmentResults):
"""Subclass of |ADFFragmentResults| for BAND calculations."""
[docs] def get_energy_decomposition(self, unit="au") -> Dict[str, float]:
"""Get the energy decomposition of the fragment calculation.
Args:
unit (str, optional): The unit of the energy. Defaults to 'au'.
Returns:
Dict[str, float]: The energy decomposition.
"""
res = self.job.full.results
res1 = self.job.f1.results
res2 = self.job.f2.results
ret = {}
def try_grep_result(key, pattern, match_position=-1, split_position=2):
match = res.grep_output(pattern)
if match:
match = Units.convert(float(match[match_position].split()[split_position]), "au", unit)
ret[key] = match
try_grep_result("E_int", "E_int", match_position=-2)
try_grep_result("E_int_disp", "E_disp")
try_grep_result("E_Pauli", "E_Pauli")
try_grep_result("E_elstat", "E_elstat")
try_grep_result("E_orb", "E_orb")
ret["E_1"] = res1.get_energy(unit=unit)
ret["E_2"] = res2.get_energy(unit=unit)
if hasattr(self.job, "f1_opt"):
ret["E_1_opt"] = self.job.f1_opt.results.get_energy(unit=unit)
ret["E_prep_1"] = ret["E_1"] - ret["E_1_opt"]
if hasattr(self.job, "f2_opt"):
ret["E_2_opt"] = self.job.f2_opt.results.get_energy(unit=unit)
ret["E_prep_2"] = ret["E_2"] - ret["E_2_opt"]
if ("E_1_opt" in ret) and ("E_2_opt" in ret):
ret["E_prep"] = ret["E_prep_1"] + ret["E_prep_2"]
ret["E_bond"] = ret["E_int"] + ret["E_prep"]
return ret
[docs]class BANDFragmentJob(ADFFragmentJob):
_result_type = BANDFragmentResults
[docs] def prerun(self) -> None:
"""Creates the fragment jobs."""
self.f1 = AMSJob(name="frag1", molecule=self.fragment1, settings=self.frag1_settings + self.settings)
self.f2 = AMSJob(name="frag2", molecule=self.fragment2, settings=self.frag2_settings + self.settings)
self.children += [self.f1, self.f2]
[docs] def new_children(self) -> Union[None, List[AMSJob]]:
"""After the first round, add the full job to the children list."""
if hasattr(self, "full"):
return None
else:
# create the correct mapping settings for the full job
if "fragment" in self.full_settings.input.band:
log(
"Fragment already present in full_settings. Assuming that the user has already set up the mapping. Skipping the mapping setup.",
level=1,
)
else:
# first fragment 1 then fragment 2
set1 = Settings()
set1.atommapping = {str(i + 1): str(i + 1) for i in range(len(self.fragment1))}
set2 = Settings()
set2.atommapping = {str(i + 1): str(i + 1 + len(self.fragment1)) for i in range(len(self.fragment2))}
set1.filename = (self.f1, "band")
set2.filename = (self.f2, "band")
self.full_settings.input.band.fragment = [set1, set2]
# create the full job
self.full = AMSJob(
name="full",
molecule=self.fragment1 + self.fragment2,
settings=self.full_settings + self.settings,
)
return [self.full]
[docs] @classmethod
def load_external(cls, path: str, jobname: Optional[str] = None) -> "BANDFragmentJob":
"""Load the results of the BANDFragmentJob job from an external path.
Args:
path (str): The path to the job. It should at least have the
subfolders 'frag1', 'frag2' and 'full'.
jobname (str, optional): The name of the job. Defaults to None.
Returns:
BANDFragmentJob: The job with the loaded results.
"""
if not isdir(path):
raise FileError("Path {} does not exist, cannot load from it.".format(path))
path = abspath(path)
jobname = basename(path) if jobname is None else str(jobname)
job = cls(name=jobname)
job.path = path
job.status = JobStatus.COPIED
job.f1 = AMSJob.load_external(opj(path, "frag1"))
job.f2 = AMSJob.load_external(opj(path, "frag2"))
job.full = AMSJob.load_external(opj(path, "full"))
job.children = [job.f1, job.f2, job.full]
if isdir(opj(path, "frag1_opt")):
job.f1_opt = AMSJob.load_external(opj(path, "frag1_opt"))
job.children.append(job.f1_opt)
if isdir(opj(path, "frag2_opt")):
job.f2_opt = AMSJob.load_external(opj(path, "frag2_opt"))
job.children.append(job.f2_opt)
return job
[docs]class NOCVBandFragmentJob(BANDFragmentJob):
_result_type = BANDFragmentResults
def __init__(self, nocv_settings=None, **kwargs) -> None:
"""Create a BANDFragmentJob with final NOCV calculation.
Args:
nocv_settings (Settings or scm.inputs.AMS, optional): Settings for the NOCV calculation. Defaults to None.
"""
super().__init__(**kwargs)
self.nocv_settings = to_plams_settings(nocv_settings, type(self).__name__) or Settings()
[docs] def new_children(self) -> Union[None, List[AMSJob]]:
"""After the first round, add the full job to the children list.
After the second round, add the NOCV job to the children list."""
# new_children of BANDFragmentJob creates the full job
ret = super().new_children()
if ret is not None:
return ret
if hasattr(self, "nocv"):
return None
else:
# add NOCV run
# settings for the NOCV calculation
set = self.nocv_settings + self.full_settings + self.settings
# set the restart file
set.input.band.restart.file = "full.rkf"
# copy the fragment settings
set.input.band.fragment = [fset.copy() for fset in self.full.settings.input.band.fragment]
self.nocv = AMSJob(name="NOCV", molecule=self.fragment1 + self.fragment2, settings=set)
self.nocv.frag_paths = []
for job in [self.f1, self.f2, self.full]:
self.nocv.frag_paths.append(job.path)
# edit NOCV prerun to create symlinks
@add_to_instance(self.nocv)
def prerun(self):
"""Create symlinks for the restart files."""
for i, job in enumerate(["frag1", "frag2", "full"]):
rel_path = relpath(self.frag_paths[i], self.path)
symlink(opj(rel_path, "band.rkf"), opj(self.path, f"{job}.rkf"))
return [self.nocv]
[docs] @classmethod
def load_external(cls, path: str, jobname: Optional[str] = None) -> "NOCVBandFragmentJob":
"""Load the results of the BANDFragmentJob job from an external path.
Args:
path (str): The path to the job. It should at least have the
subfolders 'frag1', 'frag2', 'full' and 'NOCV'.
jobname (str, optional): The name of the job. Defaults to None.
Returns:
NOCVBandFragmentJob: The job with the loaded results.
"""
job = super().load_external(path, jobname)
job.nocv = AMSJob.load_external(opj(path, "NOCV"))
job.children.append(job.nocv)
return job
