"""
Tools for calculations
"""
import warnings
from aiida.common import InputValidationError
from aiida.common.links import LinkType
from aiida.engine import CalcJob, ProcessBuilder
from aiida.orm import CalcJobNode, Dict
from aiida.plugins import DataFactory
from aiida.tools import CalculationTools
from aiida_castep.common import INPUT_LINKNAMES, OUTPUT_LINKNAMES
__all__ = [
"CastepCalcTools",
"create_restart",
"castep_input_summary",
"update_parameters",
"use_pseudos_from_family",
]
[docs]def use_pseudos_from_family(builder, family_name):
"""
Set the pseudos port namespace for a builder using pseudo family name
:note: The structure must already be set in the builder.
:param builder: ProcessBuilder instance to be processed, it must have a structure
:param family_name: the name of the group containing the pseudos
:returns: The same builder with the pseudopotential set
"""
from collections import defaultdict
from aiida_castep.data import get_pseudos_from_structure
# A dict {kind_name: pseudo_object}
# But we want to run with use_pseudo(pseudo, kinds)
structure = builder.get(INPUT_LINKNAMES["structure"], None)
if structure is None:
raise RuntimeError("The builder must have a StructureData")
kind_pseudo_dict = get_pseudos_from_structure(structure, family_name)
for kind, pseudo in kind_pseudo_dict.items():
builder.pseudos.__setattr__(kind, pseudo)
return builder
[docs]def update_parameters(inputs, force=False, delete=None, **kwargs):
"""
Convenient function to update the parameters of the calculation.
Will atomiatically set the PARAM or CELL field in unstored
ParaemterData linked to the calculation.
If no ``Dict`` is linked to the calculation, a new node will be
created.
..note:
This method relies on the help information to check and assign
keywords to PARAM or CELL field of the Dict
(i.e for generating .param and .cell file)
calc.update_parameters(task="singlepoint")
:param force: flag to force the update even if the Dict node is stored.
:param delete: A list of the keywords to be deleted.
"""
param_node = inputs.get(INPUT_LINKNAMES["parameters"])
# Create the node if none is found
if param_node is None:
warnings.warn("No existing Dict node found, creating a new one.")
param_node = Dict(dict={"CELL": {}, "PARAM": {}})
inputs[INPUT_LINKNAMES["parameters"]] = param_node
if isinstance(param_node, Dict) and param_node.is_stored:
if force:
# Create a new node if the existing node is stored
param_node = Dict(dict=param_node.get_dict())
inputs[INPUT_LINKNAMES["parameters"]] = param_node
else:
raise RuntimeError(f"The input Dict<{param_node.pk}> is already stored")
# If the `node` is just a plain dict, we keep it that way
if isinstance(param_node, Dict):
param_dict = param_node.get_dict()
py_dict = False
else:
param_dict = param_node
py_dict = True
# Update the dictionary
from .helper import CastepHelper, HelperCheckError
helper = CastepHelper()
dict_update, not_found = helper._from_flat_dict(kwargs)
if not_found:
suggest = [helper.get_suggestion(i) for i in not_found]
error_string = "Following keys are invalid -- "
for error_key, sug in zip(not_found, suggest):
error_string += f"{error_key}: {sug}; "
raise HelperCheckError(error_string)
else:
param_dict["PARAM"].update(dict_update["PARAM"])
param_dict["CELL"].update(dict_update["CELL"])
# Delete any keys as requested
if delete:
for key in delete:
tmp1 = param_dict["PARAM"].pop(key, None)
tmp2 = param_dict["CELL"].pop(key, None)
if (tmp1 is None) and (tmp2 is None):
warnings.warn(f"Key '{key}' not found")
# Apply the change to the node
if py_dict:
inputs[INPUT_LINKNAMES["parameters"]] = param_dict
else:
param_node.set_dict(param_dict)
return inputs
[docs]def create_restart(
inputs,
entry_point="castep.castep",
calcjob=None,
param_update=None,
param_delete=None,
restart_mode="restart",
use_castep_bin=False,
parent_folder=None,
reuse=False,
):
"""
Function to create a restart for a calculation.
:param inputs: A builder or nested dictionary
:param entry_point: Name of the entry points
:param param_update: Update the parameters
:param param_delete: A list of parameters to be deleted
:param restart_mode: Mode of the restart, 'continuation' or 'restart'
:param use_castep_bin: Use hte 'castep_bin' file instead of check
:param parent_folder: Remote folder to be used for restart
:param reuse: Use the reuse mode
"""
from aiida.engine import ProcessBuilder
from aiida.plugins import CalculationFactory
# Create the builder, in any case
if isinstance(inputs, dict):
processclass = CalculationFactory(entry_point)
builder = processclass.get_builder()
elif isinstance(inputs, ProcessBuilder):
builder = inputs._process_class.get_builder()
builder._update(inputs)
# Update list
update = {}
delete = []
# Set the restart tag
suffix = ".check" if not use_castep_bin else ".castep_bin"
if restart_mode == "continuation":
update["continuation"] = "parent/" + builder.metadata.seedname + suffix
delete.append("reuse")
elif restart_mode == "restart" and reuse:
update["reuse"] = "parent/" + builder.metadata.seedname + suffix
delete.append("continuation")
elif restart_mode is None:
delete.extend(["continuation", "reuse"])
elif restart_mode != "restart":
raise RuntimeError("Unknown restart mode: " + restart_mode)
if param_update:
update.update(param_update)
if param_delete:
delete.extend(param_delete)
new_builder = update_parameters(builder, force=True, delete=delete, **update)
# Set the parent folder
if parent_folder is not None:
new_builder[INPUT_LINKNAMES["parent_calc_folder"]] = parent_folder
return new_builder
def validate_input_param(input_dict, allow_flat=False):
"""
Validate inputs parameters
:param input_dict: A Dict instance or python dict instance
"""
from .helper import CastepHelper
if isinstance(input_dict, Dict):
py_dict = input_dict.get_dict()
else:
py_dict = input_dict
helper = CastepHelper()
helper.check_dict(py_dict, auto_fix=False, allow_flat=allow_flat)
def input_param_validator(input_dict, port=None):
"""
Validator used for input ports
"""
from .helper import HelperCheckError
try:
validate_input_param(input_dict)
except HelperCheckError as error:
return error.args[0]
def flat_input_param_validator(input_dict, port=None):
"""
Validator that allows allow_flat parameter format
"""
from .helper import HelperCheckError
try:
validate_input_param(input_dict, allow_flat=True)
except HelperCheckError as error:
return error.args[0]
def check_restart(builder, verbose=False):
"""
Check the RemoteData reference by the builder is satisfied
:returns: True if OK
:raises: InputValidationError if error is found
"""
import os
from .utils import _lowercase_dict
def _print(inp):
if verbose:
print(inp)
paramdict = builder[INPUT_LINKNAMES["parameters"]].get_dict()["PARAM"]
paramdict = _lowercase_dict(paramdict, "paramdict")
stemp = paramdict.get("reuse", None)
if not stemp:
stemp = paramdict.get("continuation", None)
if stemp is not None:
fname = os.path.split(stemp)[-1]
_print(f"This calculation requires a restart file: '{fname}'")
else:
# No restart file needed
_print("This calculation does not require a restart file.")
return True
# Now check if the remote folder has this file
remote_data = builder.get(INPUT_LINKNAMES["parent_calc_folder"])
if not remote_data:
raise InputValidationError(
"Restart requires " "parent_folder to be specified".format(fname)
)
else:
_print("Checking remote directory")
folder_list = remote_data.listdir()
if fname not in folder_list:
raise InputValidationError(
"Restart file {}" " is not in the remote folder".format(fname)
)
else:
_print(f"Check finished, restart file '{fname}' exists.")
return True