From 702f795aea3ab5b81dc399322e03ecd9ea3d74d8 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 15:20:29 +0100 Subject: [PATCH 1/9] Update base CalcJob to handle a TrajectoryData input object --- src/aiida_chemshell/calculations/base.py | 60 ++++++++++++++++++------ 1 file changed, 46 insertions(+), 14 deletions(-) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 6602798..e96cf15 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -3,7 +3,15 @@ from aiida.common import CalcInfo, CodeInfo from aiida.common.folders import Folder from aiida.engine import CalcJob, CalcJobProcessSpec, PortNamespace -from aiida.orm import ArrayData, Dict, Float, SinglefileData, StructureData +from aiida.orm import ( + ArrayData, + Dict, + Float, + Int, + SinglefileData, + StructureData, + TrajectoryData, +) from aiida_chemshell.units import UnitsConverter from aiida_chemshell.utils import ChemShellMMTheory, ChemShellQMTheory @@ -40,7 +48,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None: super().define(spec) spec.input( "structure", - valid_type=(SinglefileData, StructureData), + valid_type=(SinglefileData, StructureData, TrajectoryData), validator=cls.validate_structure_file, required=True, help=( @@ -49,6 +57,15 @@ def define(cls, spec: CalcJobProcessSpec) -> None: "instance." ), ) + spec.input( + "structure_index", + valid_type=Int, + required=False, + help=( + "An index to extract a specific structure if the input 'structure' " + "contains multiple structures, such as if it is a TrajectoryData node." + ), + ) ## Task object parameters spec.input( @@ -742,20 +759,30 @@ def chemsh_script_generator(self) -> str: qmmm_chk = "qm_parameters" in self.inputs and "mm_parameters" in self.inputs script = "from chemsh import Fragment\n" - if isinstance(self.inputs.structure, SinglefileData): + if isinstance(self.inputs.structure, StructureData): + if len(self.inputs.structure.sites) < 2: + atom_names = [site.kind_name for site in self.inputs.structure.sites] + coords = [ + [UnitsConverter.angstrom_to_bohr(r) for r in site.position] + for site in self.inputs.structure.sites + ] + script += ( + f"structure = Fragment(coords={str(coords):s}, names=" + f"{str(atom_names):s})\n" + ) + else: + script += "structure = Fragment(coords=" + script += f"'{ChemShellCalculation.FILE_TMP_STRUCTURE:s}')\n" + elif isinstance(self.inputs.structure, TrajectoryData): + script += "structure = Fragment(coords=" + script += f"'{ChemShellCalculation.FILE_TMP_STRUCTURE:s}')\n" + elif "structure_index" in self.inputs: + print("Not yet supported.") + raise Exception("SinglefileData trajectories not yet supported.") + else: # SinglefileData script += ( f"structure = Fragment(coords='{self.inputs.structure.filename:s}')\n" ) - else: - atom_names = [site.kind_name for site in self.inputs.structure.sites] - coords = [ - [UnitsConverter.angstrom_to_bohr(r) for r in site.position] - for site in self.inputs.structure.sites - ] - script += ( - f"structure = Fragment(coords={str(coords):s}, names=" - f"{str(atom_names):s})\n" - ) ## Setup Theory objects @@ -912,13 +939,18 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: if isinstance(self.inputs.structure, StructureData): with folder.open(ChemShellCalculation.FILE_TMP_STRUCTURE, "wb") as f: f.write(self.inputs.structure._prepare_xyz()[0]) + elif isinstance(self.inputs.structure, TrajectoryData): + with folder.open(ChemShellCalculation.FILE_TMP_STRUCTURE, "wb") as f: + index = self.inputs.structure_index.value + structure = self.inputs.structure.get_step_structure(index=index) + f.write(structure._prepare_xyz()[0]) else: calc_info.local_copy_list.append( ( self.inputs.structure.uuid, self.inputs.structure.filename, self.inputs.structure.filename, - ), + ) ) # If running with an MM theory a force field file is required and copied From 20b8f484ed04877d93bb20ae39ec6c2be8151a9c Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 15:21:10 +0100 Subject: [PATCH 2/9] Add test for TrajectoryData input to base CalcJob --- tests/conftest.py | 23 ++++++++++++++++++++- tests/test_calculations.py | 42 +++++++++++++++++++++++++++++++++++++- 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index f651160..72518ee 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,12 +3,13 @@ import os import pathlib +import numpy import pytest from aiida.common.folders import Folder from aiida.engine import CalcJob from aiida.engine.utils import instantiate_process from aiida.manage.manager import get_manager -from aiida.orm import Dict, InstalledCode, SinglefileData, StructureData +from aiida.orm import Dict, InstalledCode, SinglefileData, StructureData, TrajectoryData pytest_plugins = "aiida.tools.pytest_fixtures" @@ -73,6 +74,26 @@ def water_structure_object() -> StructureData: return structure +@pytest.fixture +def water_trajectory_object() -> TrajectoryData: + """Return a AiiDA StructureData object of a water molecule.""" + trajectory = TrajectoryData() + symbols = ["O", "H", "H"] + positions = numpy.array( + [ + [[0.0, 0.0, 0.0], [-0.9, 0.590032355, 0.0], [0.9, 0.590032355, 0.0]], + [ + [0.0, 0.0, 0.0], + [-0.754606402, 0.590032355, 0.0], + [0.754606402, 0.590032355, 0.0], + ], + [[0.0, 0.0, 0.0], [-0.5, 0.590032355, 0.0], [0.5, 0.590032355, 0.0]], + ] + ) + trajectory.set_trajectory(symbols=symbols, positions=positions) + return trajectory + + @pytest.fixture def generate_inputs(chemsh_code, get_test_data_file): """Return a dictionary of inputs for the ChemShellCalculation.""" diff --git a/tests/test_calculations.py b/tests/test_calculations.py index abea831..4a0ddb3 100644 --- a/tests/test_calculations.py +++ b/tests/test_calculations.py @@ -81,7 +81,8 @@ def test_sp_calculation_qm_dft(chemsh_code, get_test_data_file, water_structure_ assert ChemShellCalculation.FILE_STDOUT in ofiles assert ChemShellCalculation.FILE_RESULTS in ofiles - eref = -75.946889377347 + # eref = -75.946889377347 # If using conversion to bohr + eref = -75.946889436563 assert abs(results.get("energy") - eref) < 1e-8, ( "Incorrect energy result for PySCF based SP calculation" ) @@ -256,6 +257,45 @@ def test_vibrational_calculation(chemsh_code, get_test_data_file): assert results.get("vibrational_energies").get("Entropy / J/mol/K") == 0.01313 +def test_structure_from_trajectorydata(chemsh_code, water_trajectory_object): + """DFT based single point test.""" + code = chemsh_code() + builder = code.get_builder() + builder.structure = water_trajectory_object + builder.structure_index = 1 + builder.qm_parameters = Dict( + { + "theory": "PySCF", + "method": "DFT", + "functional": "BLYP", + "charge": 0, + # "scftype": "uks", + } + ) + + results, node = run.get_node(builder) + + assert node.is_finished_ok, "CalcJob failed for `test_SPCalculation_nwchem_DFT`" + + assert "Single_Point" in node.process_label + assert "QM" in node.process_label + + ofiles = results.get("retrieved").list_object_names() + assert ChemShellCalculation.FILE_STDOUT in ofiles + assert ChemShellCalculation.FILE_RESULTS in ofiles + + # eref = -75.946889377347 # Use if inputs are in bohr + eref = -75.946889436563 + assert abs(results.get("energy") - eref) < 1e-8, ( + "Incorrect energy result for PySCF based SP calculation" + ) + + assert "gradients" not in results, ( + "Gradients have been returned for a SP calculation, \ + but they were not requested in the inputs." + ) + + # def test_opt_calculation_qmmm(chemsh_code, get_test_data_file): # """QM/MM geometry optimisation test.""" # code = chemsh_code() From f5bfe7e3be07766ec9a98e43ef3e0e9476116feb Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 15:21:44 +0100 Subject: [PATCH 3/9] Update to test_inputs to setup a test for when the input StructureData is atomic (handled slightly differently in base CalcJob) --- tests/test_inputs.py | 33 +++++++++++++++++++++++++++++---- 1 file changed, 29 insertions(+), 4 deletions(-) diff --git a/tests/test_inputs.py b/tests/test_inputs.py index 640d7cf..96e0437 100644 --- a/tests/test_inputs.py +++ b/tests/test_inputs.py @@ -1,5 +1,7 @@ """Tests for ChemShell input script generation based on various input parameters.""" +from aiida.orm import StructureData + from aiida_chemshell.calculations.base import ChemShellCalculation @@ -149,7 +151,30 @@ def test_structure_as_structuredata_object( script_txt = script_file.read_text() assert "from chemsh import Fragment\n" in script_txt - assert ( - "structure = Fragment(coords=[[0.0, 0.0, 0.0], [-1.4259993290233388, 1.1149994" - in script_txt - ) + assert ChemShellCalculation.FILE_TMP_STRUCTURE in script_txt + + +def test_atom_as_structuredata_object( + generate_calcjob, generate_inputs, water_structure_object +): + """Test taking a StructureData object as an input.""" + structure = StructureData() + structure.pbc = [False, False, False] + structure.append_atom(symbols="O", position=[0.0, 0.0, 0.0]) + inputs = generate_inputs(structure_fname=structure) + + tmp_pth, calc_info = generate_calcjob(ChemShellCalculation, inputs) + + structure_file = tmp_pth / ChemShellCalculation.FILE_TMP_STRUCTURE + assert structure_file.exists() + chk_str = """1 +Lattice="0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" pbc="False False False" +O 0.0000000000 0.0000000000 0.0000000000""" + assert structure_file.read_text() == chk_str + + script_file = tmp_pth / ChemShellCalculation.FILE_SCRIPT + assert script_file.exists() + + script_txt = script_file.read_text() + assert "from chemsh import Fragment\n" in script_txt + assert "Fragment(coords=[[0.0, 0.0, 0.0]]," in script_txt From ec874feabcf2a89c35b0d5191c2958416f7e5705 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 15:23:02 +0100 Subject: [PATCH 4/9] Add test for isolated atomic energy workchain (marked xfail as test container doesn't currently support NWChem which is required) --- tests/test_isolated_atoms_workchain.py | 27 ++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 tests/test_isolated_atoms_workchain.py diff --git a/tests/test_isolated_atoms_workchain.py b/tests/test_isolated_atoms_workchain.py new file mode 100644 index 0000000..377c4cd --- /dev/null +++ b/tests/test_isolated_atoms_workchain.py @@ -0,0 +1,27 @@ +"""Tests for carrying the pre-defined IsolatedAtoms workflows with aiida-chemshell.""" + +import pytest +from aiida.engine import run_get_node + +from aiida_chemshell.workflows.isolated_atoms import IsolatedAtomicEnergiesWorkChain + + +@pytest.mark.xfail(reason="Will fail if NWChem not properly configured.") +def test_geometry_optimisation_workflow(chemsh_code, get_test_data_file): + """Test a geometry optimisation workflow with vibrational analysis.""" + inputs = { + "structure": get_test_data_file(), + "code": chemsh_code(), + "qm_parameters": {"theory": "NWChem", "method": "HF"}, + } + results, node = run_get_node(IsolatedAtomicEnergiesWorkChain, **inputs) + + # print(results) + assert node.is_finished_ok, ( + "WorkChain node failed for IsolatedAtomicEnergiesWorkChain" + ) + assert len(list(results.get("atom_energies").get_dict().keys())) == 2 + h_energy_ref = -0.496198609381 + assert abs(results.get("atom_energies").get_dict()["H"] - h_energy_ref) < 1e-10 + o_energy_ref = -74.267449889229 + assert abs(results.get("atom_energies").get_dict()["O"] - o_energy_ref) < 1e-10 From ad9768cb8e0df9db1e08bc4122b300a94a43d7fb Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 15:59:31 +0100 Subject: [PATCH 5/9] Create a workchain that carries out a given calculation on each structure within a TrajectoryData object --- .../workflows/batch_calculation.py | 66 +++++++++++++++++++ tests/test_batch_workcahin.py | 28 ++++++++ 2 files changed, 94 insertions(+) create mode 100644 src/aiida_chemshell/workflows/batch_calculation.py create mode 100644 tests/test_batch_workcahin.py diff --git a/src/aiida_chemshell/workflows/batch_calculation.py b/src/aiida_chemshell/workflows/batch_calculation.py new file mode 100644 index 0000000..01ceb13 --- /dev/null +++ b/src/aiida_chemshell/workflows/batch_calculation.py @@ -0,0 +1,66 @@ +"""Workflow for processing a series of structures from a single input.""" + +from aiida.engine import ProcessSpec, ToContext, WorkChain +from aiida.orm import ProcessNode, SinglefileData, TrajectoryData +from aiida.plugins.factories import CalculationFactory + +ChemShellCalculation = CalculationFactory("chemshell") + + +class BatchProcessWorkChain(WorkChain): + """Process a series of structures with the same inputs.""" + + @classmethod + def define(cls, spec: ProcessSpec) -> None: + """Define the AiiDA process specification.""" + super().define(spec) + + # Expose Chemshell inputs + spec.expose_inputs(ChemShellCalculation, exclude=("structure", "metadata")) + + # Input structure series + spec.input( + "structure", + valid_type=(SinglefileData, TrajectoryData), + required=True, + help="The series of structures to process.", + ) + + spec.outline(cls.submit_jobs, cls.collate_results) + + def submit_jobs(self): + """Extract all individual structures and submit their calculations.""" + futures: dict[str, ProcessNode] = {} + if isinstance(self.inputs.structure, TrajectoryData): + for i in range(self.inputs.structure.numsteps): + inputs = { + "code": self.inputs.code, + "structure": self.inputs.structure, + "structure_index": i, + } + if "qm_parameters" in self.inputs: + inputs["qm_parameters"] = self.inputs.qm_parameters + if "mm_parameters" in self.inputs: + inputs["mm_parameters"] = self.inputs.mm_parameters + inputs["force_field_file"] = self.inputs.force_field_file + if "qmmm_parameters" in self.inputs: + inputs["qmmm_parameters"] = self.inputs.qmmm_parameters + if "calculation_parameters" in self.inputs: + inputs["calculation_parameters"] = ( + self.inputs.calculation_parameters + ) + if "optimisation_parameters" in self.inputs: + inputs["optimisation_parameters"] = ( + self.inputs.optimisation_parameters + ) + future = self.submit(ChemShellCalculation, **inputs) + futures[f"structure_{i}"] = future + return ToContext(**futures) + + def collate_results(self) -> None: + """Collect the WorkChain's results.""" + return + + +# def read_structures_from_file(filname, contents) -> list(): +# return [] diff --git a/tests/test_batch_workcahin.py b/tests/test_batch_workcahin.py new file mode 100644 index 0000000..0d08534 --- /dev/null +++ b/tests/test_batch_workcahin.py @@ -0,0 +1,28 @@ +"""Tests for the BatchProcessWorkChain.""" + +from aiida.engine import run_get_node +from aiida.orm import Dict + +from aiida_chemshell.workflows.batch_calculation import BatchProcessWorkChain + + +def test_batch_from_trajectorydata(chemsh_code, water_trajectory_object): + """DFT based single point test.""" + inputs = { + "code": chemsh_code(), + "structure": water_trajectory_object, + "qm_parameters": Dict( + { + "theory": "PySCF", + "method": "hf", + } + ), + } + results, node = run_get_node(BatchProcessWorkChain, **inputs) + + sub_nodes = node.called + + final_energies = [-75.565560193461, -75.585287771819, -75.426355430539] + + for i, sub_node in enumerate(sub_nodes): + assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10 From d3e59e23c2a2a0fc81e6909b956aff6be5f8db4b Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 16:43:11 +0100 Subject: [PATCH 6/9] Allow BatchProcessWorkChain to accept a dictionary of StructureData objects as an input --- pyproject.toml | 2 +- .../workflows/batch_calculation.py | 77 ++++++++++++------- tests/conftest.py | 4 +- tests/test_batch_workcahin.py | 34 +++++++- 4 files changed, 85 insertions(+), 32 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 02b8096..925971d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,7 +42,7 @@ where = ["src"] [project.optional-dependencies] dev = [ - "aiida-core>=2.6", # Required to support aiida.tools.pytest_fixtures + "aiida-core>=2.8", "pytest>=8.0", "pytest-cov>=6.0", "ruff>=0.11.0", diff --git a/src/aiida_chemshell/workflows/batch_calculation.py b/src/aiida_chemshell/workflows/batch_calculation.py index 01ceb13..f0eb287 100644 --- a/src/aiida_chemshell/workflows/batch_calculation.py +++ b/src/aiida_chemshell/workflows/batch_calculation.py @@ -1,7 +1,7 @@ """Workflow for processing a series of structures from a single input.""" from aiida.engine import ProcessSpec, ToContext, WorkChain -from aiida.orm import ProcessNode, SinglefileData, TrajectoryData +from aiida.orm import ProcessNode, StructureData, TrajectoryData from aiida.plugins.factories import CalculationFactory ChemShellCalculation = CalculationFactory("chemshell") @@ -20,41 +20,60 @@ def define(cls, spec: ProcessSpec) -> None: # Input structure series spec.input( - "structure", - valid_type=(SinglefileData, TrajectoryData), - required=True, - help="The series of structures to process.", + "trajectory", + valid_type=TrajectoryData, + required=False, + help="The series of structures to process as a TrajectoryData node.", + ) + spec.input_namespace( + "structures", + valid_type=StructureData, + required=False, + help="A dictionary of StructureData nodes to batch process.", + ) + + spec.exit_code( + 350, + "ERROR_NO_INPUTS", + message=("Must specify either 'trajectory' or 'structure' input."), ) - spec.outline(cls.submit_jobs, cls.collate_results) + spec.outline(cls.validate_inputs, cls.submit_jobs, cls.collate_results) + + def validate_inputs(self): + """Validate the inputs provided to the WorkChain.""" + has_trajectory = "trajectory" in self.inputs + has_structures = "structures" in self.inputs + if not has_trajectory and not has_structures: + return self.exit_codes.ERROR_NO_INPUTS + return None def submit_jobs(self): """Extract all individual structures and submit their calculations.""" futures: dict[str, ProcessNode] = {} - if isinstance(self.inputs.structure, TrajectoryData): - for i in range(self.inputs.structure.numsteps): - inputs = { - "code": self.inputs.code, - "structure": self.inputs.structure, - "structure_index": i, - } - if "qm_parameters" in self.inputs: - inputs["qm_parameters"] = self.inputs.qm_parameters - if "mm_parameters" in self.inputs: - inputs["mm_parameters"] = self.inputs.mm_parameters - inputs["force_field_file"] = self.inputs.force_field_file - if "qmmm_parameters" in self.inputs: - inputs["qmmm_parameters"] = self.inputs.qmmm_parameters - if "calculation_parameters" in self.inputs: - inputs["calculation_parameters"] = ( - self.inputs.calculation_parameters - ) - if "optimisation_parameters" in self.inputs: - inputs["optimisation_parameters"] = ( - self.inputs.optimisation_parameters - ) + inputs = {"code": self.inputs.code} + if "qm_parameters" in self.inputs: + inputs["qm_parameters"] = self.inputs.qm_parameters + if "mm_parameters" in self.inputs: + inputs["mm_parameters"] = self.inputs.mm_parameters + inputs["force_field_file"] = self.inputs.force_field_file + if "qmmm_parameters" in self.inputs: + inputs["qmmm_parameters"] = self.inputs.qmmm_parameters + if "calculation_parameters" in self.inputs: + inputs["calculation_parameters"] = self.inputs.calculation_parameters + if "optimisation_parameters" in self.inputs: + inputs["optimisation_parameters"] = self.inputs.optimisation_parameters + if "trajectory" in self.inputs: + for i in range(self.inputs.trajectory.numsteps): + inputs["structure"] = self.inputs.trajectory + inputs["structure_index"] = i + future = self.submit(ChemShellCalculation, **inputs) + futures[f"trajectory_frame_{i}"] = future + if "structures" in self.inputs: + for key, structure in self.inputs.structures.items(): + inputs["structure"] = structure future = self.submit(ChemShellCalculation, **inputs) - futures[f"structure_{i}"] = future + futures[key] = future return ToContext(**futures) def collate_results(self) -> None: diff --git a/tests/conftest.py b/tests/conftest.py index 72518ee..61f7f98 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -90,7 +90,9 @@ def water_trajectory_object() -> TrajectoryData: [[0.0, 0.0, 0.0], [-0.5, 0.590032355, 0.0], [0.5, 0.590032355, 0.0]], ] ) - trajectory.set_trajectory(symbols=symbols, positions=positions) + trajectory.set_trajectory( + symbols=symbols, positions=positions, pbc=[False, False, False] + ) return trajectory diff --git a/tests/test_batch_workcahin.py b/tests/test_batch_workcahin.py index 0d08534..a2f2c59 100644 --- a/tests/test_batch_workcahin.py +++ b/tests/test_batch_workcahin.py @@ -1,16 +1,48 @@ """Tests for the BatchProcessWorkChain.""" +from aiida import __version__ as aiida_core_version from aiida.engine import run_get_node from aiida.orm import Dict +from packaging.version import parse as parse_version from aiida_chemshell.workflows.batch_calculation import BatchProcessWorkChain +AIIDA_LESS_THAN_2_8 = parse_version(aiida_core_version) < parse_version("2.8.0") + def test_batch_from_trajectorydata(chemsh_code, water_trajectory_object): """DFT based single point test.""" inputs = { "code": chemsh_code(), - "structure": water_trajectory_object, + "trajectory": water_trajectory_object, + "qm_parameters": Dict( + { + "theory": "PySCF", + "method": "hf", + } + ), + } + results, node = run_get_node(BatchProcessWorkChain, **inputs) + + sub_nodes = node.called + + final_energies = [-75.565560193461, -75.585287771819, -75.426355430539] + + for i, sub_node in enumerate(sub_nodes): + assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10 + + +def test_batch_from_structuredata(chemsh_code, water_trajectory_object): + """DFT based single point test.""" + trajectory = water_trajectory_object + + inputs = { + "code": chemsh_code(), + "structures": { + "structure_1": trajectory.get_step_structure(0), + "structure_2": trajectory.get_step_structure(1), + "structure_3": trajectory.get_step_structure(2), + }, "qm_parameters": Dict( { "theory": "PySCF", From a2f58cdb44fcb86e18d81aa38a47b3929c841f04 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 16:57:39 +0100 Subject: [PATCH 7/9] Update test for BatchProcessWorkChain with StructureData dictionary as input --- tests/conftest.py | 33 ++++++++++++++++++++++++++++++--- tests/test_batch_workcahin.py | 6 ++---- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 61f7f98..5e635ae 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,15 +5,39 @@ import numpy import pytest +from aiida import __version__ as aiida_core_version from aiida.common.folders import Folder from aiida.engine import CalcJob from aiida.engine.utils import instantiate_process from aiida.manage.manager import get_manager from aiida.orm import Dict, InstalledCode, SinglefileData, StructureData, TrajectoryData +from packaging.version import parse as parse_version pytest_plugins = "aiida.tools.pytest_fixtures" +def pytest_configure(config): + """Dynamically register a custom xfail condition based on AiiDA version.""" + config.addinivalue_line( + "markers", + "xfail_aiida_2_8: mark test as expected failure if aiida-core is < 2.8", + ) + + +def pytest_runtest_setup(item): + """Evaluate the custom marker before running the test.""" + marker = item.get_closest_marker("xfail_aiida_2_8") + if marker: + if parse_version(aiida_core_version) < parse_version("2.8.0"): + # Dynamically apply the xfail to this specific test instance + item.add_marker( + pytest.mark.xfail( + reason="This test requires features present in aiida-core >= 2.8", + raises=ValueError, + ) + ) + + @pytest.fixture def get_data_filepath() -> pathlib.Path: """Return the path to the tests data folder.""" @@ -90,9 +114,12 @@ def water_trajectory_object() -> TrajectoryData: [[0.0, 0.0, 0.0], [-0.5, 0.590032355, 0.0], [0.5, 0.590032355, 0.0]], ] ) - trajectory.set_trajectory( - symbols=symbols, positions=positions, pbc=[False, False, False] - ) + if parse_version(aiida_core_version) < parse_version("2.8.0"): + trajectory.set_trajectory(symbols=symbols, positions=positions) + else: + trajectory.set_trajectory( + symbols=symbols, positions=positions, pbc=[False, False, False] + ) return trajectory diff --git a/tests/test_batch_workcahin.py b/tests/test_batch_workcahin.py index a2f2c59..58328f8 100644 --- a/tests/test_batch_workcahin.py +++ b/tests/test_batch_workcahin.py @@ -1,14 +1,11 @@ """Tests for the BatchProcessWorkChain.""" -from aiida import __version__ as aiida_core_version +import pytest from aiida.engine import run_get_node from aiida.orm import Dict -from packaging.version import parse as parse_version from aiida_chemshell.workflows.batch_calculation import BatchProcessWorkChain -AIIDA_LESS_THAN_2_8 = parse_version(aiida_core_version) < parse_version("2.8.0") - def test_batch_from_trajectorydata(chemsh_code, water_trajectory_object): """DFT based single point test.""" @@ -32,6 +29,7 @@ def test_batch_from_trajectorydata(chemsh_code, water_trajectory_object): assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10 +@pytest.mark.xfail_aiida_2_8 def test_batch_from_structuredata(chemsh_code, water_trajectory_object): """DFT based single point test.""" trajectory = water_trajectory_object From 97b99b3c7d9bfab779308f332e113670933d268e Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 30 Jun 2026 17:03:15 +0100 Subject: [PATCH 8/9] Create test for BatchProcessWorkChain where both "trajectory" and "structures" inputs are specified --- tests/test_batch_workcahin.py | 44 +++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/tests/test_batch_workcahin.py b/tests/test_batch_workcahin.py index 58328f8..0a084e9 100644 --- a/tests/test_batch_workcahin.py +++ b/tests/test_batch_workcahin.py @@ -21,11 +21,15 @@ def test_batch_from_trajectorydata(chemsh_code, water_trajectory_object): } results, node = run_get_node(BatchProcessWorkChain, **inputs) + assert node.is_finished_ok, "WorkChain Failed" + sub_nodes = node.called + assert len(sub_nodes) == 3, "Incorrect number of sub processes created." final_energies = [-75.565560193461, -75.585287771819, -75.426355430539] for i, sub_node in enumerate(sub_nodes): + assert sub_node.is_finished_ok, "Sub Process Failed" assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10 @@ -50,9 +54,49 @@ def test_batch_from_structuredata(chemsh_code, water_trajectory_object): } results, node = run_get_node(BatchProcessWorkChain, **inputs) + assert node.is_finished_ok, "WorkChain Failed" + sub_nodes = node.called + assert len(sub_nodes) == 3, "Incorrect number of sub processes created." final_energies = [-75.565560193461, -75.585287771819, -75.426355430539] for i, sub_node in enumerate(sub_nodes): + assert sub_node.is_finished_ok, "Sub Process Failed" + assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10 + + +def test_batch_from_structuredata_and_trajectorydata( + chemsh_code, water_trajectory_object, water_structure_object +): + """DFT based single point test.""" + inputs = { + "code": chemsh_code(), + "trajectory": water_trajectory_object, + "structures": { + "Structure_1": water_structure_object, + }, + "qm_parameters": Dict( + { + "theory": "PySCF", + "method": "hf", + } + ), + } + results, node = run_get_node(BatchProcessWorkChain, **inputs) + + assert node.is_finished_ok, "WorkChain Failed" + + sub_nodes = node.called + assert len(sub_nodes) == 4, "Incorrect number of sub processes created." + + final_energies = [ + -75.565560193461, + -75.585287771819, + -75.426355430539, + -75.585287771819, + ] + + for i, sub_node in enumerate(sub_nodes): + assert sub_node.is_finished_ok, "Sub Process Failed" assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10 From 4fe0f9f544486041e1e6c22553a19d08ac96a814 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 1 Jul 2026 10:50:21 +0100 Subject: [PATCH 9/9] Add support for supplying an .xyz based trajectory file to the BatchProcessWorkChain --- .../workflows/batch_calculation.py | 125 +++++++++++++++++- tests/data/trajectory.xyz | 40 ++++++ tests/test_batch_workcahin.py | 34 +++++ 3 files changed, 192 insertions(+), 7 deletions(-) create mode 100644 tests/data/trajectory.xyz diff --git a/src/aiida_chemshell/workflows/batch_calculation.py b/src/aiida_chemshell/workflows/batch_calculation.py index f0eb287..8185f9a 100644 --- a/src/aiida_chemshell/workflows/batch_calculation.py +++ b/src/aiida_chemshell/workflows/batch_calculation.py @@ -1,7 +1,9 @@ """Workflow for processing a series of structures from a single input.""" -from aiida.engine import ProcessSpec, ToContext, WorkChain -from aiida.orm import ProcessNode, StructureData, TrajectoryData +import re + +from aiida.engine import ProcessSpec, ToContext, WorkChain, calcfunction +from aiida.orm import ProcessNode, SinglefileData, StructureData, TrajectoryData from aiida.plugins.factories import CalculationFactory ChemShellCalculation = CalculationFactory("chemshell") @@ -31,23 +33,59 @@ def define(cls, spec: ProcessSpec) -> None: required=False, help="A dictionary of StructureData nodes to batch process.", ) + spec.input_namespace( + "structure_files", + valid_type=SinglefileData, + required=False, + help=( + "A dictionary of SinglefileData objects containing the series of " + "input structures. The dictionary keys are not used, the node labels" + "are determined by the filename." + ), + ) spec.exit_code( 350, "ERROR_NO_INPUTS", - message=("Must specify either 'trajectory' or 'structure' input."), + message=( + "Must specify either 'trajectory', 'structures' or " + "'structure_files' input." + ), ) - spec.outline(cls.validate_inputs, cls.submit_jobs, cls.collate_results) + spec.outline( + cls.validate_inputs, + cls.extract_structures_from_files, + cls.submit_jobs, + cls.collate_results, + ) def validate_inputs(self): """Validate the inputs provided to the WorkChain.""" has_trajectory = "trajectory" in self.inputs has_structures = "structures" in self.inputs - if not has_trajectory and not has_structures: + has_files = "structure_files" in self.inputs + if not has_trajectory and not has_structures and not has_files: return self.exit_codes.ERROR_NO_INPUTS return None + def extract_structures_from_files(self) -> None: + """Extract the structures from file like input objects.""" + self.structures_from_files = {} + if "structure_files" in self.inputs: + for _key, file in self.inputs.structure_files.items(): + if file.filename[-4:] != ".xyz": + self.report( + "Only XYZ structured trajectory files are currently " + f"supported, {file.filename} will be skipped..." + ) + else: + self.report(f"Parsing structures from {file.filename}") + self.structures_from_files = ( + self.structures_from_files | extract_structures_from_xyz(file) + ) + return + def submit_jobs(self): """Extract all individual structures and submit their calculations.""" futures: dict[str, ProcessNode] = {} @@ -74,6 +112,11 @@ def submit_jobs(self): inputs["structure"] = structure future = self.submit(ChemShellCalculation, **inputs) futures[key] = future + if "structure_files" in self.inputs: + for key, structure in self.structures_from_files.items(): + inputs["structure"] = structure + future = self.submit(ChemShellCalculation, **inputs) + futures[key] = future return ToContext(**futures) def collate_results(self) -> None: @@ -81,5 +124,73 @@ def collate_results(self) -> None: return -# def read_structures_from_file(filname, contents) -> list(): -# return [] +@calcfunction +def extract_structures_from_xyz(file: SinglefileData): + """Parse a SinglefileData XYZ trajectory into individual StructureData nodes.""" + with file.open(mode="r") as f: + lines = f.readlines() + + structures = {} + line_count = len(lines) + i = 0 + frame_idx = 0 + + while i < line_count: + line = lines[i].strip() + try: + natoms = int(line) + except ValueError as e: + raise Exception("Invalid XYZ format detected.") from e + if (i + 2 + natoms) > line_count: + raise Exception("XYZ file truncation detected.") + + # Create the base StructureData object + structure = StructureData(pbc=(False, False, False)) + + # Read the comment line + i += 1 + line = lines[i].strip() + cell = None + if "Lattice=" in line: + match = re.search(r'Lattice="([^"]+)"', line) + if match: + lat_vals = [float(x) for x in match.group(1).split()] + if len(lat_vals) == 9: + cell = [lat_vals[0:3], lat_vals[3:6], lat_vals[6:9]] + pbc = [True, True, True] + if "pbc=" in line: + match_pbc = re.search(r'pbc="([^"]+)"', line) + if match_pbc: + pbc_vals = match_pbc.group(1).split() + if len(pbc_vals) == 3: + # Robust check: converts 'T', 'True', or '1' to True + pbc = [val.upper() in ["T", "TRUE", "1"] for val in pbc_vals] + + # Assign the parse cell parameters to the StructureData object + structure.cell = cell + structure.pbc = pbc + + i += 1 + for atmi in range(natoms): + atom_line = lines[i + atmi].strip().split() + if len(atom_line) < 4: + raise Exception( + f"Invalid atom entry in xyz file: {line[i + atmi].strip()}" + ) + + structure.append_atom( + position=[ + float(atom_line[1]), + float(atom_line[2]), + float(atom_line[3]), + ], + symbols=atom_line[0], + ) + + structures[ + f"{file.filename.replace(' ', '_').strip('.xyz')}_frame_{frame_idx}" + ] = structure + frame_idx += 1 + i += natoms + + return structures diff --git a/tests/data/trajectory.xyz b/tests/data/trajectory.xyz new file mode 100644 index 0000000..36a031a --- /dev/null +++ b/tests/data/trajectory.xyz @@ -0,0 +1,40 @@ +3 +Step=0 Properties="species:S:1:pos:R:3:force:R:3" Lattice="10.0 0.0 0.0 0.0 10.0 0.0 0.0 0.0 10.0" pbc="F F F" dft_energy=-75.58528778 +O 0.0000000 0.0000000 0.0000000 0.0000000 -0.0010112 -0.0000000 +H -0.7546067 0.5900326 0.0000000 -0.0076697 0.0005056 -0.0000000 +H 0.7546067 0.5900326 0.0000000 0.0076697 0.0005056 0.0000000 +3 +Step=1 Properties="species:S:1:pos:R:3:force:R:3" Lattice="10.0 0.0 0.0 0.0 10.0 0.0 0.0 0.0 10.0" pbc="F F F" dft_energy=-75.5855946 +O 0.0000000 -0.0009803 -0.0000000 -0.0000000 0.0025996 0.0000000 +H -0.7620421 0.5905228 -0.0000000 -0.0040707 -0.0012998 -0.0000000 +H 0.7620421 0.5905228 0.0000000 0.0040707 -0.0012998 0.0000000 +3 +Step=3 Properties="species:S:1:pos:R:3:force:R:3" Lattice="10.0 0.0 0.0 0.0 10.0 0.0 0.0 0.0 10.0" pbc="F F F" dft_energy=-75.58595962 +O 0.0000000 0.0133814 -0.0000000 -0.0000000 -0.0001765 0.0000000 +H -0.7804552 0.5833419 -0.0000000 -0.0000765 0.0000882 -0.0000000 +H 0.7804552 0.5833419 0.0000000 0.0000765 0.0000882 0.0000000 +23 + +C 0.01210000 -0.50580000 0.02050000 0.00000000 +C -1.25610000 0.35070000 0.00910000 0.00000000 +C 1.26960000 0.36640000 0.02690000 0.00000000 +C -2.54460000 -0.47400000 -0.02160000 0.00000000 +C 2.53790000 -0.48640000 -0.00030000 0.00000000 +C -3.80810000 0.37230000 -0.00710000 0.00000000 +C 3.78910000 0.37680000 -0.02750000 0.00000000 +H 0.00850000 -1.15620000 0.90340000 0.00000000 +H 0.01940000 -1.15710000 -0.86180000 0.00000000 +H -1.23760000 1.01530000 -0.86360000 0.00000000 +H -1.26390000 0.99510000 0.89710000 0.00000000 +H 1.27220000 1.00300000 0.91990000 0.00000000 +H 1.25800000 1.03230000 -0.84460000 0.00000000 +H -2.55050000 -1.10440000 -0.91870000 0.00000000 +H -2.56310000 -1.14990000 0.84160000 0.00000000 +H 2.56740000 -1.13740000 0.88120000 0.00000000 +H 2.53170000 -1.13590000 -0.88340000 0.00000000 +H -4.69530000 -0.26790000 -0.03070000 0.00000000 +H -3.84270000 1.03600000 -0.87680000 0.00000000 +H -3.85500000 0.98770000 0.89680000 0.00000000 +H 3.80590000 1.01570000 -0.91630000 0.00000000 +H 3.84250000 1.01800000 0.85810000 0.00000000 +H 4.68390000 -0.25330000 -0.04530000 0.00000000 \ No newline at end of file diff --git a/tests/test_batch_workcahin.py b/tests/test_batch_workcahin.py index 0a084e9..e2a92e6 100644 --- a/tests/test_batch_workcahin.py +++ b/tests/test_batch_workcahin.py @@ -100,3 +100,37 @@ def test_batch_from_structuredata_and_trajectorydata( for i, sub_node in enumerate(sub_nodes): assert sub_node.is_finished_ok, "Sub Process Failed" assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10 + + +def test_batch_from_file(chemsh_code, get_test_data_file): + """DFT based single point test.""" + structure_file = get_test_data_file("trajectory.xyz") + inputs = { + "code": chemsh_code(), + "structure_files": { + structure_file.filename.strip(".xyz").replace(" ", "_"): structure_file + }, + "qm_parameters": Dict( + { + "theory": "PySCF", + "method": "hf", + } + ), + } + results, node = run_get_node(BatchProcessWorkChain, **inputs) + + assert node.is_finished_ok, "WorkChain Failed" + + sub_nodes = node.called + assert len(sub_nodes) == 5, "Incorrect number of sub processes created." + + final_energies = [ + -75.585287789025, + -75.585594607649, + -75.585959615566, + -272.88756364993, + ] + + for i, sub_node in enumerate(sub_nodes[1:]): + assert sub_node.is_finished_ok, "Sub Process Failed" + assert abs(sub_node.outputs.energy - final_energies[i]) < 1e-10