diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index e96cf15..7618670 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -31,6 +31,7 @@ class ChemShellCalculation(CalcJob): FILE_STDOUT = "output.log" FILE_DLFIND = "_dl_find.cjson" FILE_TMP_STRUCTURE = "input_structure.xyz" + FILE_TMP_STRUCTURE2 = "input_structure_2.xyz" FILE_RESULTS = "result.json" FILE_TRJPTH = "path.xyz" FILE_TRJFRC = "path_force.xyz" @@ -66,6 +67,18 @@ def define(cls, spec: CalcJobProcessSpec) -> None: "contains multiple structures, such as if it is a TrajectoryData node." ), ) + spec.input( + "structure2", + valid_type=(SinglefileData, StructureData), + validator=cls.validate_structure_file, + required=False, + help=( + "An additional input structure for the ChemShell calculation either" + "contained within an '.xyz', '.pun' or '.cjson' file or as a " + "StructureData instance. This is used in jobs such as NEB " + "optimisations as the final structure." + ), + ) ## Task object parameters spec.input( @@ -204,6 +217,21 @@ def inputs_validator_wrapper(inputs, namespace): ), ) + spec.output( + "neb_path", + valid_type=TrajectoryData, + required=False, + help="The pathway determined by a Nudged Elastic Band calculation.", + ) + spec.output( + "neb_info", + valid_type=ArrayData, + required=False, + help=( + "Information generated at each point along a Nudged Elastic Band path." + ), + ) + ## Metadata spec.inputs["metadata"]["options"]["resources"].default = { "num_machines": 1, @@ -784,6 +812,15 @@ def chemsh_script_generator(self) -> str: f"structure = Fragment(coords='{self.inputs.structure.filename:s}')\n" ) + # Create second fragment object if requested + if "structure2" in self.inputs: + if isinstance(self.inputs.structure2, SinglefileData): + fname = self.inputs.structure2.filename + else: + fname = ChemShellCalculation.FILE_TMP_STRUCTURE2 + print(fname) + script += f"frg2 = Fragment(coords='{fname:s}')\n" + ## Setup Theory objects if "qm_parameters" in self.inputs: @@ -857,6 +894,8 @@ def chemsh_script_generator(self) -> str: # Run a geometry optimisation using DL_FIND script += "from chemsh import Opt\n" opt_str = f"job = Opt(theory={theory_str:s}" + if "structure2" in self.inputs: + opt_str += ", " + "frag2=frg2" for key in self.inputs.optimisation_parameters.keys(): if isinstance(self.inputs.optimisation_parameters.get(key), str): opt_str += ", " + key + "='" @@ -881,7 +920,14 @@ def chemsh_script_generator(self) -> str: script += "job.run()\njob.result.save()\n" if "optimisation_parameters" in self.inputs: - script += f'structure.save("{ChemShellCalculation.FILE_DLFIND}")\n' + if not self.inputs.optimisation_parameters.get( + "thermal", False + ) and self.inputs.optimisation_parameters.get("neb", "no") not in [ + "free", + "frozen", + "perpendicular", + ]: + script += f'structure.save("{ChemShellCalculation.FILE_DLFIND}")\n' return script @@ -950,9 +996,23 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: self.inputs.structure.uuid, self.inputs.structure.filename, self.inputs.structure.filename, - ) + ), ) + # Copy data for second input fragment if required + if "structure2" in self.inputs: + if isinstance(self.inputs.structure2, StructureData): + with folder.open(ChemShellCalculation.FILE_TMP_STRUCTURE, "wb") as f: + f.write(self.inputs.structure2._prepare_xyz()[0]) + else: + calc_info.local_copy_list.append( + ( + self.inputs.structure2.uuid, + self.inputs.structure2.filename, + self.inputs.structure2.filename, + ), + ) + # If running with an MM theory a force field file is required and copied if "force_field_file" in self.inputs: calc_info.local_copy_list.append( @@ -966,7 +1026,14 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: # If performing a geometry optimisation retrieve the generated _dl_find.pun # file containing the optimised structure if "optimisation_parameters" in self.inputs: - calc_info.retrieve_list.append(ChemShellCalculation.FILE_DLFIND) + if not self.inputs.optimisation_parameters.get( + "thermal", False + ) and self.inputs.optimisation_parameters.get("neb", "no") not in [ + "free", + "frozen", + "perpendicular", + ]: + calc_info.retrieve_list.append(ChemShellCalculation.FILE_DLFIND) if self.inputs.optimisation_parameters.get("save_path", False): calc_info.retrieve_list.append( "_dl_find/" + ChemShellCalculation.FILE_TRJPTH @@ -974,5 +1041,12 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: calc_info.retrieve_list.append( "_dl_find/" + ChemShellCalculation.FILE_TRJFRC ) + if self.inputs.optimisation_parameters.get("neb", "no") in [ + "free", + "frozen", + "perpendicular", + ]: + calc_info.retrieve_temporary_list.append("nebinfo") + calc_info.retrieve_temporary_list.append("nebpath.xyz") return calc_info diff --git a/src/aiida_chemshell/calculations/solvant_removal.py b/src/aiida_chemshell/calculations/solvant_removal.py new file mode 100644 index 0000000..798ac3a --- /dev/null +++ b/src/aiida_chemshell/calculations/solvant_removal.py @@ -0,0 +1,115 @@ +"""Module for performing a solvant removal CalcJob.""" + +from aiida.common import CalcInfo, CodeInfo +from aiida.common.folders import Folder +from aiida.engine import CalcJob, CalcJobProcessSpec +from aiida.orm import Int, SinglefileData + +from aiida_chemshell.utils import xyz_file_validator + + +class SolvantRemovalCalcJob(CalcJob): + """CalcJob to extract a solvant from a complex for NEB calculations.""" + + _SCRIPT_NAME = "separate_solvant.py" + + @classmethod + def define(cls, spec: CalcJobProcessSpec) -> None: + """Define the CalcJob spec for solvant removal.""" + super().define(spec) + spec.input( + "structure", + valid_type=SinglefileData, + required=True, + validator=xyz_file_validator, + help=( + "Input structure file containing ligand+solvant gas phase complex " + "in XYZ format." + ), + ) + spec.input( + "num_ligand_atoms", + valid_type=Int, + required=True, + help="Number of atoms within the core ligand molecule.", + ) + spec.output( + "unbound_structure", + valid_type=SinglefileData, + required=True, + help="Gas phase structure for the unboun ligand+solvant complex.", + ) + + spec.inputs["metadata"]["options"]["resources"].default = { + "num_machines": 1, + "num_mpiprocs_per_machine": 1, + } + spec.inputs["metadata"]["options"]["parser_name"].default = "chemshell.unbind" + + def prepare_for_submission(self, folder: Folder) -> CalcInfo: + """Prepare the CalcJob for submission.""" + py_scipt = self._generate_python_script() + with folder.open(SolvantRemovalCalcJob._SCRIPT_NAME, "w") as handle: + handle.write(py_scipt) + + code_info = CodeInfo() + code_info.code_uuid = self.inputs.code.uuid + code_info.cmdline_params = [SolvantRemovalCalcJob._SCRIPT_NAME] + + calc_info = CalcInfo() + calc_info.codes_info = [code_info] + calc_info.retrieve_temporary_list = [] + calc_info.provenance_exclude_list = [] + calc_info.retrieve_list = [self._generate_output_filename()] + calc_info.local_copy_list = [ + ( + self.inputs.structure.uuid, + self.inputs.structure.filename, + self.inputs.structure.filename, + ), + ] + + return calc_info + + @classmethod + def generate_output_filename(cls, input_filename) -> str: + """Generate a filename based on the name of the initial structure file.""" + return input_filename.replace(".xyz", "_unbound.xyz") + + def _generate_output_filename(self) -> str: + """Generate the output filename.""" + return self.generate_output_filename(self.inputs.structure.filename) + + def _generate_python_script(self) -> str: + """Generate the Python script to be executed.""" + return f""" +import numpy +with open("{self.inputs.structure.filename}", "r") as f: + lines = f.readlines() +natoms = int(lines[0].strip()) +ligand = [] +for i in range(2, {int(self.inputs.num_ligand_atoms)} + 2): + line = lines[i].split() + ligand.append(numpy.array([float(line[1]), float(line[2]), float(line[3])])) +solvant = [] +for i in range({int(self.inputs.num_ligand_atoms)} + 2, natoms + 2): + line = lines[i].split() + solvant.append(numpy.array([float(line[1]), float(line[2]), float(line[3])])) +ligand_center = numpy.mean(ligand, axis=0) +solvant_center = numpy.mean(solvant, axis=0) +vector = solvant_center - ligand_center +distance = numpy.linalg.norm(vector) +vector /= distance +displacement = vector * 10.0 # Move solvant 10 Angstroms away +with open("{self._generate_output_filename()}", "w") as f: + f.write(f"{{natoms}}\\n") + f.write("Unbound ligand+solvant complex\\n") + for i in range(2, {int(self.inputs.num_ligand_atoms)} + 2): + f.write(lines[i]) + for i in range({int(self.inputs.num_ligand_atoms)} + 2, natoms + 2): + line = lines[i].split() + x = float(line[1]) + displacement[0] + y = float(line[2]) + displacement[1] + z = float(line[3]) + displacement[2] + f.write(f"{{line[0]}} {{x:.8f}} {{y:.8f}} {{z:.8f}}\\n") +""" diff --git a/src/aiida_chemshell/parsers/base.py b/src/aiida_chemshell/parsers/base.py index 8d94bf1..0ccf8a5 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -1,11 +1,12 @@ """Defines the calculation parsers for the ChemShell AiiDA plugin.""" import json +from pathlib import Path import numpy from aiida.common import ModificationNotAllowed from aiida.engine import ExitCode -from aiida.orm import ArrayData, Dict, Float, SinglefileData +from aiida.orm import ArrayData, Dict, Float, SinglefileData, TrajectoryData from aiida.parsers.parser import Parser from aiida_chemshell.calculations.base import ChemShellCalculation @@ -16,6 +17,8 @@ class ChemShellParser(Parser): def parse(self, **kwargs): """Parse the output of a ChemShell calculation.""" + retrieved_tmp_folder = kwargs.get("retrieved_temporary_folder", None) + if ChemShellCalculation.FILE_STDOUT not in self.retrieved.list_object_names(): return self.exit_codes.ERROR_STDOUT_NOT_FOUND if ChemShellCalculation.FILE_RESULTS not in self.retrieved.list_object_names(): @@ -73,6 +76,13 @@ def parse(self, **kwargs): ChemShellCalculation.FILE_STDOUT, "r" ) ) + elif self.node.inputs.optimisation_parameters.get("neb", "no") in [ + "free", + "frozen", + "perpendicular", + ]: + self.parse_neb_path(Path(retrieved_tmp_folder) / "nebpath.xyz") + self.parse_neb_info(Path(retrieved_tmp_folder) / "nebinfo") elif ChemShellCalculation.FILE_DLFIND in self.retrieved.list_object_names(): descrip = "Optimised structure from a ChemShell optimisation" input_pk = self.node.inputs.structure.pk @@ -189,3 +199,57 @@ def parse_optimisation_path(self, stdout: str) -> None: results.set_array("energies", numpy.asarray(energies)) self.out("optimisation_path", results) return + + def parse_neb_path(self, file_path: Path) -> None: + """Parse the NEB pathway into an AiiDA TrajectoryData node.""" + with open(file_path) as f: + lines = f.readlines() + natoms = int(lines[0]) + symbols = [] + positions = [] + step = 0 + i = 2 + while i < len(lines): + step_positions = numpy.zeros((natoms, 3), dtype=float) + for atm_index, atm_line in enumerate(lines[i : i + natoms]): + line = atm_line.split() + if step == 0: + symbols.append(line[0]) + step_positions[atm_index][0] = float(line[1]) + step_positions[atm_index][1] = float(line[2]) + step_positions[atm_index][2] = float(line[3]) + positions.append(step_positions) + step += 1 + i += natoms + 2 + path = TrajectoryData() + path.set_trajectory(symbols=symbols, positions=numpy.asarray(positions)) + self.out("neb_path", path) + return + + def parse_neb_info(self, file_path: Path) -> None: + """Parse the NEB info file into an AiiDA ArrayData noe.""" + output = ArrayData( + label="Step information from an NEB calculation.", + description=( + "Calculated step values for an ChemShell NEB calculation from node: " + f"{self.node.pk}" + ), + ) + with open(file_path) as f: + lines = f.readlines() + length = [] + energy = [] + work = [] + mass = [] + for line in lines[1:]: + vals = line.split() + length.append(float(vals[0])) + energy.append(float(vals[1])) + work.append(float(vals[2])) + mass.append(float(vals[3])) + output.set_array("path_length", numpy.asarray(length)) + output.set_array("energy", numpy.asarray(energy)) + output.set_array("work", numpy.asarray(work)) + output.set_array("effective_mass", numpy.asarray(mass)) + self.out("neb_info", output) + return diff --git a/src/aiida_chemshell/utils.py b/src/aiida_chemshell/utils.py index 26cce95..a980662 100644 --- a/src/aiida_chemshell/utils.py +++ b/src/aiida_chemshell/utils.py @@ -2,7 +2,7 @@ from enum import Enum, auto -from aiida.orm import StructureData +from aiida.orm import SinglefileData, StructureData class ChemShellQMTheory(Enum): @@ -124,3 +124,22 @@ def generate_default_mlip_fine_tune_config(): "eval_interval": 1, # "enable_cueq": True, } + + +def xyz_file_validator(value: SinglefileData) -> str | None: + """Check if a file is a valid XYZ file.""" + contents = value.get_content(mode="r").splitlines() + try: + natoms = int(contents[0].strip()) + except ValueError: + return ( + "The first line of the XYZ file must be an integer" + "representing the number of atoms." + ) + else: + if natoms != len(contents) - 2: + return ( + f"The number of atoms specified ({natoms}) does not" + f"match the number of atom lines ({len(contents) - 2})." + ) + return None diff --git a/tests/data/h2o_dimer_2.cjson b/tests/data/h2o_dimer_2.cjson new file mode 100644 index 0000000..8f8af40 --- /dev/null +++ b/tests/data/h2o_dimer_2.cjson @@ -0,0 +1,18 @@ +{ + "name": "h2o dimer", + "atoms": { + "elements": { + "number": [ 8, 1, 1, 8, 1, 1 ] + }, + "coords": { + "unit": "angstrom", + "3d": [ 0.000000, 0.000000, 0.000000, + 0.000000, -0.751841, 0.568201, + -0.000000, 0.751841, 0.568201, + -10.511474, 0.000000, -1.450000, + -10.651960, 0.000000, -1.063537, + -10.374291, 0.000000, -2.382362 ] + }, + "charges": [ -0.834, 0.417, 0.417, -0.834, 0.417, 0.417 ] + } +} diff --git a/tests/test_calculations.py b/tests/test_calculations.py index 4a0ddb3..73f4407 100644 --- a/tests/test_calculations.py +++ b/tests/test_calculations.py @@ -1,7 +1,7 @@ """Tests for performing calculation processes with aiida_chemshell.""" from aiida.engine import run -from aiida.orm import Dict +from aiida.orm import Dict, TrajectoryData from numpy.linalg import norm from aiida_chemshell.calculations.base import ChemShellCalculation @@ -296,6 +296,54 @@ def test_structure_from_trajectorydata(chemsh_code, water_trajectory_object): ) +def test_neb_calculation(chemsh_code, get_test_data_file): + """QM test for neb calculation and second structure input.""" + code = chemsh_code() + builder = code.get_builder() + builder.structure = get_test_data_file("h2o_dimer.cjson") + builder.structure2 = get_test_data_file("h2o_dimer_2.cjson") + builder.qm_parameters = Dict({"theory": "PySCF", "method": "hf", "basis": "3-21G"}) + builder.optimisation_parameters = Dict({"neb": "frozen"}) + + results, node = run.get_node(builder) + + assert node.is_finished_ok, "CalcJob failed for `test_neb_calculationF`" + + ofiles = results.get("retrieved").list_object_names() + assert ChemShellCalculation.FILE_STDOUT in ofiles + assert ChemShellCalculation.FILE_DLFIND not in ofiles + assert ChemShellCalculation.FILE_RESULTS in ofiles + + eref = -149.56942655605 + assert abs(results.get("energy") - eref) < 1e-8 + assert isinstance(results.get("neb_path", None), TrajectoryData) + assert results.get("neb_path").numsteps == 9, ( + "Incorrect number of steps in NEB path TrajectoryData" + ) + assert results.get("neb_path").numsites == 6, ( + "Incorrect number of sites in NEB path TrajectoryData" + ) + + neb_info = results.get("neb_info", None) + assert "Step information from an NEB calculation." in neb_info.label, ( + f"Incorrect node label for NEB info array output node. {neb_info.label}" + ) + assert neb_info.description != "", "Missing NEB info array node description" + + assert neb_info.get_arraynames() == [ + "path_length", + "energy", + "work", + "effective_mass", + ], "Incorrect NEB info array names in output node." + + assert neb_info.get_shape("path_length") == (9,), ( + "Incorrect number of steps in path_length array." + ) + + return + + # def test_opt_calculation_qmmm(chemsh_code, get_test_data_file): # """QM/MM geometry optimisation test.""" # code = chemsh_code()