Create Optimisation Workflow

src/aiida_chemshell/calculations/base.py

@@ -2,7 +2,7 @@
 
 from aiida.common import CalcInfo, CodeInfo
 from aiida.common.folders import Folder
-from aiida.engine import CalcJob, CalcJobProcessSpec
+from aiida.engine import CalcJob, CalcJobProcessSpec, PortNamespace
 from aiida.orm import ArrayData, Dict, Float, SinglefileData, StructureData
 
 from aiida_chemshell.utils import ChemShellMMTheory, ChemShellQMTheory
@@ -131,6 +131,34 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
                 "structure is contained within a ChemShell '.pun' file."
             ),
         )
+        spec.output(
+            "vibrational_energies",
+            valid_type=Dict,
+            required=False,
+            help="The calculated thermochemical properties of the system.",
+        )
+        spec.output(
+            "vibrational_modes",
+            valid_type=ArrayData,
+            required=False,
+            help="The calculated vibrational modes of the system.",
+        )
+
+        # Validate inputs namespace
+        existing_validator = spec.inputs.validator
+
+        def inputs_validator_wrapper(inputs, namespace):
+            """Wrap the existing and custom input namespace validators."""
+            if existing_validator:
+                error = existing_validator(inputs, namespace)
+                if error:
+                    return error
+            error = cls.validate_inputs_namespace(inputs, namespace)
+            if error:
+                return error
+            return None
+
+        spec.inputs.validator = inputs_validator_wrapper
 
         ## Metadata
         spec.inputs["metadata"]["options"]["resources"].default = {
@@ -178,6 +206,22 @@ def define(cls, spec: CalcJobProcessSpec) -> None:
 
         return
 
+    @classmethod
+    def validate_inputs_namespace(
+        cls, value: Dict, namespace: PortNamespace
+    ) -> str | None:
+        """Perform additional validation checks on the total inputs namespace."""
+        if "mm_parameters" in value and "force_field_file" not in value:
+            return "A force field must be specified to use molecular mechanics."
+        if "force_field_file" in value and "mm_parameters" not in value:
+            return "A MM theory code must be specified to use molecular mechanics."
+        if "qmmm_parameters" in value:
+            if "qm_parameters" not in value:
+                return "Missing QM parameters for QM/MM calculation"
+            if "mm_parameters" not in value:
+                return "Missing MM parameters for QM/MM calculation"
+        return None
+
     @classmethod
     def validate_structure_file(
         cls, value: SinglefileData | StructureData | None, _
@@ -282,6 +326,7 @@ def get_valid_optimisation_parameter_keys(cls) -> tuple[str]:
             "dimer",
             "delta",
             "tsrelative",
+            "thermal",
         )
 
     @classmethod

src/aiida_chemshell/parsers/base.py

@@ -5,7 +5,7 @@
 import numpy
 from aiida.common import ModificationNotAllowed
 from aiida.engine import ExitCode
-from aiida.orm import ArrayData, Float, SinglefileData
+from aiida.orm import ArrayData, Dict, Float, SinglefileData
 from aiida.parsers.parser import Parser
 
 from aiida_chemshell.calculations.base import ChemShellCalculation
@@ -23,7 +23,7 @@ def parse(self, **kwargs):
 
         # Read the 'json' formatted results file
         results = json.loads(
-            self.retrieved.get_object_content(ChemShellCalculation.FILE_RESULTS)
+            self.retrieved.get_object_content(ChemShellCalculation.FILE_RESULTS, "rb")
         )
 
         # Extract the final energy
@@ -67,7 +67,13 @@ def parse(self, **kwargs):
 
         # If the calculation was a geometry optimisation, store the optimised structure
         if "optimisation_parameters" in self.node.inputs:
-            if ChemShellCalculation.FILE_DLFIND in self.retrieved.list_object_names():
+            if self.node.inputs.optimisation_parameters.get("thermal", False):
+                self.parse_vibrational_analysis(
+                    self.retrieved.get_object_content(
+                        ChemShellCalculation.FILE_STDOUT, "r"
+                    )
+                )
+            elif ChemShellCalculation.FILE_DLFIND in self.retrieved.list_object_names():
                 descrip = "Optimised structure from a ChemShell optimisation"
                 input_pk = self.node.inputs.structure.pk
                 descrip += f" of node {input_pk}"
@@ -89,3 +95,43 @@ def parse(self, **kwargs):
                 return self.exit_codes.ERROR_MISSING_OPTIMISED_STRUCTURE_FILE
 
         return ExitCode(0)
+
+    def parse_vibrational_analysis(self, stdout: str) -> None:
+        """Extract the vibrational analysis from ChemShell output log."""
+        read = False
+        energies = {}
+        modes = []
+        for line in stdout.split("\n"):
+            if read:
+                line_vals = line.split()
+                if "Temperature:" in line:
+                    energies[f"Temperature / {line_vals[2]}"] = float(line_vals[1])
+                elif "E_electronic" in line:
+                    energies[f"E_electronic correction / {line_vals[7]}"] = float(
+                        line_vals[6]
+                    )
+                elif "total ZPE" in line:
+                    energies[f"ZPE / {line_vals[3]}"] = float(line_vals[2])
+                elif "total E vib" in line:
+                    energies[f"Enthalpy / {line_vals[4]}"] = float(line_vals[3])
+                elif "total S vib" in line:
+                    energies[f"Entropy / {line_vals[4]}"] = float(line_vals[3])
+                elif "Mode" in line or "total" in line:
+                    pass
+                else:
+                    # All remaining lines should be part of the modes table
+                    modes.append(numpy.array([float(x) for x in line_vals[2:]]))
+
+            if "Thermochemical analysis" in line:
+                read = True
+            elif "total S vib" in line:
+                read = False
+        self.out("vibrational_energies", Dict(energies))
+        modes = numpy.asarray(modes)
+        modes_data_node = ArrayData(
+            label="Vibrational Modes",
+            description="Calculated vibrational modes for the system.",
+        )
+        modes_data_node.set_array("Modes", modes)
+        self.out("vibrational_modes", modes_data_node)
+        return

src/aiida_chemshell/workflows/optimisation.py

@@ -0,0 +1,164 @@
+"""Workflows for geometry optimisation based taks."""
+
+from aiida.engine import ToContext, WorkChain
+from aiida.orm import ArrayData, Bool, Dict, Float, SinglefileData, Str
+
+from aiida_chemshell.calculations.base import ChemShellCalculation
+
+
+class GeometryOptimisationWorkChain(WorkChain):
+    """Geometry optimisation calculation with extended optional calculation options."""
+
+    @classmethod
+    def define(cls, spec) -> None:
+        """Define the AiiDA process specification for the WorkChain."""
+        super().define(spec)
+
+        ## Inputs ##
+        spec.expose_inputs(ChemShellCalculation, exclude=("metadata"))
+
+        spec.input(
+            "vibrational_analysis",
+            valid_type=Bool,
+            required=False,
+            help="Calculate vibrational modes of resulting structure. (default=True)",
+        )
+        spec.input(
+            "basis_quality",
+            valid_type=Str,
+            validator=cls.validate_basis_quality_input,
+            required=False,
+            help="Set basis set quality for QM calculation based on defined options.",
+        )
+
+        ## Outputs ##
+        spec.output(
+            "final_energy",
+            valid_type=Float,
+            required=True,
+            help="The final energy for the optimised structure.",
+        )
+        spec.output(
+            "optimised_structure",
+            valid_type=SinglefileData,
+            required=True,
+            help="The final optimised geometry of the given structure.",
+        )
+        spec.output(
+            "vibrational_energies",
+            valid_type=Dict,
+            required=False,
+            help="The calculated thermochemical properties of the optimised structure",
+        )
+        spec.output(
+            "vibrational_modes",
+            valid_type=ArrayData,
+            required=False,
+            help="The calculated vibrational modes for the optimised structure.",
+        )
+
+        ## Workflow ##
+        spec.outline(cls.optimise, cls.energy, cls.result)
+
+        return
+
+    def optimise(self):
+        """Perform the geometry optimisation."""
+        inputs = self.exposed_inputs(ChemShellCalculation)
+        if "qm_parameters" not in self.inputs:
+            inputs["qm_parameters"] = Dict(
+                {
+                    "theory": "NWChem",
+                    "method": "dft",
+                    "functional": "B3LYP",
+                    "d3": True,
+                }
+            )
+            if "basis_quality" in self.inputs:
+                inputs["qm_parameters"]["basis"] = self.get_basis_set_label(
+                    self.inputs.basis_quality.value
+                )
+        elif "basis_quality" in self.inputs:
+            inputs["qm_parameters"] = inputs.get("qm_parameters").get_dict()
+            inputs["qm_parameters"]["basis"] = self.get_basis_set_label(
+                self.inputs.basis_quality.value
+            )
+        if "force_field_file" in inputs:
+            if "mm_parameters" not in inputs:
+                inputs["mm_parameters"] = Dict({"theory": "DL_POLY"})
+            if "qmmm_parameters" not in inputs:
+                inputs["qmmm_parameters"] = Dict({"qm_region": []})
+        elif "mm_parameters" in inputs:
+            return None
+        if "optimisation_parameters" not in inputs:
+            inputs["optimisation_parameters"] = Dict({})
+
+        future = self.submit(ChemShellCalculation, **inputs)
+        return ToContext(optimise=future)
+
+    def energy(self):
+        """Perform a single point energy calculation on the optimised structure."""
+        if self.inputs.get("vibrational_analysis", False):
+            inputs = {
+                "code": self.exposed_inputs(ChemShellCalculation)["code"],
+                "structure": self.ctx.optimise.outputs.optimised_structure,
+                "qm_parameters": self.ctx.optimise.inputs.qm_parameters,
+            }
+            if "force_field_file" in self.ctx.optimise.inputs:
+                inputs["force_field_file"] = self.ctx.optimise.inputs.force_field_file
+                inputs["mm_parameters"] = self.ctx.optimise.inputs.mm_parameters
+                inputs["qmmm_parameters"] = self.ctx.optimise.inputs.qmmm_parameters
+            inputs["optimisation_parameters"] = (
+                self.ctx.optimise.inputs.optimisation_parameters.get_dict()
+            )
+            inputs["optimisation_parameters"]["thermal"] = True
+            future = self.submit(ChemShellCalculation, **inputs)
+            return ToContext(energy=future)
+        return None
+
+    def result(self):
+        """Extract the final workflow results."""
+        self.out("optimised_structure", self.ctx.optimise.outputs.optimised_structure)
+        if self.inputs.get("vibrational_analysis", False):
+            self.out("final_energy", self.ctx.energy.outputs.energy)
+            self.out(
+                "vibrational_energies", self.ctx.energy.outputs.vibrational_energies
+            )
+            self.out("vibrational_modes", self.ctx.energy.outputs.vibrational_modes)
+        else:
+            self.out("final_energy", self.ctx.optimise.outputs.energy)
+        return
+
+    @classmethod
+    def validate_basis_quality_input(cls, value: Str, _) -> str | None:
+        """
+        Validate the basis set quality key input.
+
+        Parameters
+        ----------
+        value : Str | None
+            The input key for the basis set quality mapping.
+
+        Returns
+        -------
+        str | None
+            Returns `None` if input is valid otherwise returns an error message.
+        """
+        if value.value.lower() not in ["fast", "balanced", "quality"]:
+            return (
+                "Invalid basis set quality key, valid keys are: 'fast', 'balanced' or "
+                "'quality'."
+            )
+        return None
+
+    @classmethod
+    def get_basis_set_label(cls, key: str) -> str:
+        """Define a custom dictionary of basis set labels for user convenience."""
+        match key.lower():
+            case "fast":
+                return "3-21G"
+            case "balanced":
+                return "cc-pvdz"
+            case "quality":
+                return "aug-cc-pvtz"
+        return ""

tests/conftest.py

@@ -69,7 +69,7 @@ def factory(
         structure_fname: str | StructureData = "water.cjson",
         ff_fname: str | None = None,
         opt: dict | None = None,
-    ):
+    ) -> dict:
         if isinstance(structure_fname, str):
             structure = get_test_data_file(structure_fname)
         else:

tests/test_calculations.py

@@ -216,6 +216,26 @@ def test_opt_calculation_dlpoly(chemsh_code, get_test_data_file):
     )
 
 
+def test_vibrational_calculation(chemsh_code, get_test_data_file):
+    """MM based geometry optimisation test."""
+    code = chemsh_code()
+    builder = code.get_builder()
+    builder.structure = get_test_data_file()
+    builder.qm_parameters = Dict({"theory": "PySCF", "method": "hf", "basis": "3-21G"})
+    builder.optimisation_parameters = Dict({"thermal": True})
+
+    results, node = run.get_node(builder)
+
+    assert node.is_finished_ok, "CalcJob failed for `test_OptCalculation_dlpoly`"
+
+    assert results.get("vibrational_modes").get_shape("Modes") == (3, 5)
+
+    assert results.get("vibrational_energies").get("Temperature / Kelvin") == 300.0
+    assert results.get("vibrational_energies").get("ZPE / J/mol") == 58728.93143
+    assert results.get("vibrational_energies").get("Enthalpy / J/mol") == 3.53277
+    assert results.get("vibrational_energies").get("Entropy / J/mol/K") == 0.01313
+
+
 # def test_opt_calculation_qmmm(chemsh_code, get_test_data_file):
 #     """QM/MM geometry optimisation test."""
 #     code = chemsh_code()