Enable NEB Calculations

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,13 +1026,27 @@ 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
                 )
                 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

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")
+"""

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

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