From 179f0e744ae5e16c5075b4837f76aabb1a29a82c Mon Sep 17 00:00:00 2001 From: Benjamin Speake Date: Wed, 5 Nov 2025 15:33:12 +0000 Subject: [PATCH 01/35] Add output fields to the base calculation for storing the (optional) trajectory files produced during geometry optimisation --- src/aiida_chemshell/calculations/base.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 5cd40bf..7ef48ac 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -160,6 +160,20 @@ def inputs_validator_wrapper(inputs, namespace): spec.inputs.validator = inputs_validator_wrapper + spec.output( + "trajectory_path", + valid_type=SinglefileData, + required=False, + help="XYZ trajectory file for the geometry optimisation", + ) + spec.output( + "trajectory_force", + valid_type=SinglefileData, + required=False, + help="XYZ trajectory containing forces at each step of a geometry " + "optimisation", + ) + ## Metadata spec.inputs["metadata"]["options"]["resources"].default = { "num_machines": 1, @@ -871,5 +885,8 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: # file containing the optimised structure if "optimisation_parameters" in self.inputs: calc_info.retrieve_list.append(ChemShellCalculation.FILE_DLFIND) + if self.inputs.optimisation_parameters.get("save_path", False): + calc_info.retrieve_list.append(ChemShellCalculation.FILE_TRJPTH) + calc_info.retrieve_list.append(ChemShellCalculation.FILE_TRJFRC) return calc_info From 01a18806b52d6cd881943122243b2f7879bb5569 Mon Sep 17 00:00:00 2001 From: Benjamin Speake Date: Wed, 5 Nov 2025 15:41:15 +0000 Subject: [PATCH 02/35] Fix for string formatting in base calculation --- src/aiida_chemshell/calculations/base.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 7ef48ac..8efd733 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -41,8 +41,8 @@ def define(cls, spec: CalcJobProcessSpec) -> None: validator=cls.validate_structure_file, required=True, help=( - "The input structure for the ChemShell calculation either contained" - "within an '.xyz', '.pun' or '.cjson' file or as a StructureData" + "The input structure for the ChemShell calculation either contained " + "within an '.xyz', '.pun' or '.cjson' file or as a StructureData " "instance." ), ) @@ -170,8 +170,10 @@ def inputs_validator_wrapper(inputs, namespace): "trajectory_force", valid_type=SinglefileData, required=False, - help="XYZ trajectory containing forces at each step of a geometry " - "optimisation", + help=( + "XYZ trajectory containing forces at each step of a geometry " + "optimisation" + ), ) ## Metadata @@ -579,7 +581,7 @@ def validate_mm_parameters(cls, value: Dict | None, _) -> str | None: theory = value.get("theory", "").upper() if theory not in ChemShellMMTheory.__members__: return ( - "The specified MM theory '{theory:s}' is not a " + f"The specified MM theory '{theory:s}' is not a " "valid ChemShell MM interface within the AiiDA-ChemShell workflow." ) From 9fa980dac6c26f594635841556d7d80266beb45e Mon Sep 17 00:00:00 2001 From: Benjamin Speake Date: Wed, 5 Nov 2025 15:47:15 +0000 Subject: [PATCH 03/35] Add labels and descriptions to output nodes --- src/aiida_chemshell/parsers/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aiida_chemshell/parsers/base.py b/src/aiida_chemshell/parsers/base.py index ed41848..3890423 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -87,7 +87,7 @@ def parse(self, **kwargs): SinglefileData( file=f, filename=ChemShellCalculation.FILE_DLFIND, - label="Structure File", + label="ChemShell Punch Structure File", description=descrip, ), ) From eea2c2b584c03cd3a903bf4943b3238a78058c74 Mon Sep 17 00:00:00 2001 From: Benjamin Speake Date: Wed, 5 Nov 2025 16:20:44 +0000 Subject: [PATCH 04/35] Correctly add 'save_path' as a valid optimisation parameter --- src/aiida_chemshell/calculations/base.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 8efd733..936741d 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -23,6 +23,8 @@ class ChemShellCalculation(CalcJob): FILE_DLFIND = "_dl_find.cjson" FILE_TMP_STRUCTURE = "input_structure.xyz" FILE_RESULTS = "result.json" + FILE_TRJPTH = "_dl_find/path.xyz" + FILE_TRJFRC = "_dl_find/force.xyz" @classmethod def define(cls, spec: CalcJobProcessSpec) -> None: @@ -343,6 +345,7 @@ def get_valid_optimisation_parameter_keys(cls) -> tuple[str]: "delta", "tsrelative", "thermal", + "save_path", ) @classmethod From 233d50bacc56539218e9acb488bbdbc4b40a5307 Mon Sep 17 00:00:00 2001 From: Benjamin Speake Date: Wed, 5 Nov 2025 16:25:11 +0000 Subject: [PATCH 05/35] Correct names for trajectory path files --- src/aiida_chemshell/calculations/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 936741d..852a8a7 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -23,8 +23,8 @@ class ChemShellCalculation(CalcJob): FILE_DLFIND = "_dl_find.cjson" FILE_TMP_STRUCTURE = "input_structure.xyz" FILE_RESULTS = "result.json" - FILE_TRJPTH = "_dl_find/path.xyz" - FILE_TRJFRC = "_dl_find/force.xyz" + FILE_TRJPTH = "path.xyz" + FILE_TRJFRC = "path_force.xyz" @classmethod def define(cls, spec: CalcJobProcessSpec) -> None: From 6886e19ef4225124ba264c4e3d6d6b96df4f1fe9 Mon Sep 17 00:00:00 2001 From: Benjamin Speake Date: Tue, 18 Nov 2025 08:52:49 +0000 Subject: [PATCH 06/35] Create outline for clf workchain --- src/aiida_chemshell/workflows/clf_ultra.py | 92 ++++++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 src/aiida_chemshell/workflows/clf_ultra.py diff --git a/src/aiida_chemshell/workflows/clf_ultra.py b/src/aiida_chemshell/workflows/clf_ultra.py new file mode 100644 index 0000000..949b0e6 --- /dev/null +++ b/src/aiida_chemshell/workflows/clf_ultra.py @@ -0,0 +1,92 @@ +"""CLF-ULTRA ChemShell + Janus workflow.""" + +from aiida.engine import ToContext, WorkChain +from aiida.orm import Code, Dict, Float, Int, List, SinglefileData, StructureData +from aiida.plugins.factories import CalculationFactory + +ChemShellCalculation = CalculationFactory("chemshell") + + +class CLFULTRAOptimisationWorkChain(WorkChain): + """ + CLF-ULTRA Geometry Optimisation WorkChain. + + This workchain will perform a QM or QM/MM based geometry optimisation + on a given structure and generate etxXYZ formatted trajectory files for + every step in the optimisation. It uses the NWChem and DL_POLY backends + and the B3LYP dft functional for the QM calculations. + """ + + @classmethod + def define(cls, spec) -> None: + """Define the AiiDA process specification for the WorkChain.""" + super().define(spec) + spec.input( + "structure", + valid_type=(SinglefileData, StructureData), + required=True, + help="The structure on which to perform the geometry optimisation", + ) + spec.input("chemshell", valid_type=Code, required=True, help="") + + spec.input( + "force_field_file", + valid_type=SinglefileData, + required=False, + help="File defining the MM force field for QM/MM calculation.", + ) + spec.input( + "qm_region", + valid_type=List, + required=False, + help=( + "A list of atom indexes to apply the QM portion of a QM/MM " + "calculation to." + ), + ) + + spec.output("final_energy", valid_type=Float) + spec.output("natms", valid_type=Int) + + spec.outline(cls.optimise, cls.generate_xyz_files, cls.result) + return + + def optimise(self): + """Perform the Geometry Optimisation Task.""" + inputs = { + "qm_parameters": Dict( + { + "theory": "NWChem", + "method": "dft", + "functional": "b3lyp", + "basis": "cc-pvtz", + # Add D3 correction + } + ), + "optimisation_parameters": Dict({"save_path": True}), + "structure": self.inputs.structure, + "code": self.inputs.chemshell, + "metadata": { + "options": { + "resources": {"num_mpiprocs_per_machine": 4, "num_machines": 1}, + "withmpi": True, + } + }, + } + future = self.submit(ChemShellCalculation, **inputs) + return ToContext(optimise=future) + + def generate_xyz_files(self): + """Convert the paths to individual extended XYZ files.""" + trjp = self.ctx.optimise.outputs.trajectory_path.get_content(mode="r") + trjf = self.ctx.optimise.outputs.trajectory_force.get_content(mode="r") + trjp = trjp.split("\n") + trjf = trjf.split("\n") + natms = int(trjp[0]) + self.out("natms", natms) + return + + def result(self) -> None: + """Report the final results.""" + self.out("final_energy", self.ctx.optimise.outputs.energy) + return From 7ffdc586b8fcd5992bf1421b48915b96d89918d2 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 18 Nov 2025 12:42:09 +0000 Subject: [PATCH 07/35] Fix for correct return of trajectory path outputs when 'save_path' option used in optimisation_parameters --- src/aiida_chemshell/calculations/base.py | 8 ++++-- src/aiida_chemshell/parsers/base.py | 33 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 852a8a7..90b82e8 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -891,7 +891,11 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: if "optimisation_parameters" in self.inputs: calc_info.retrieve_list.append(ChemShellCalculation.FILE_DLFIND) if self.inputs.optimisation_parameters.get("save_path", False): - calc_info.retrieve_list.append(ChemShellCalculation.FILE_TRJPTH) - calc_info.retrieve_list.append(ChemShellCalculation.FILE_TRJFRC) + calc_info.retrieve_list.append( + "_dl_find/" + ChemShellCalculation.FILE_TRJPTH + ) + calc_info.retrieve_list.append( + "_dl_find/" + ChemShellCalculation.FILE_TRJFRC + ) return calc_info diff --git a/src/aiida_chemshell/parsers/base.py b/src/aiida_chemshell/parsers/base.py index 3890423..5f2da1f 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -73,6 +73,39 @@ def parse(self, **kwargs): ChemShellCalculation.FILE_STDOUT, "r" ) ) + elif self.node.inputs.optimisation_parameters.get("save_path", False): + if ( + ChemShellCalculation.FILE_TRJPTH + in self.retrieved.list_object_names() + ): + with self.retrieved.open( + ChemShellCalculation.FILE_TRJPTH, "r" + ) as f: + self.out( + "trajectory_path", + SinglefileData( + file=f, + filename=ChemShellCalculation.FILE_TRJPTH.replace( + "/", "_" + ), + label="ChemShell optimisation trajectory.", + ), + ) + with self.retrieved.open( + ChemShellCalculation.FILE_TRJFRC, "r" + ) as f: + self.out( + "trajectory_force", + SinglefileData( + file=f, + filename=ChemShellCalculation.FILE_TRJFRC.replace( + "/", "_" + ), + label="ChemShell optimisation trajectory.", + ), + ) + else: + return self.exit_codes.ERROR_MISSING_OPTIMISED_STRUCTURE_FILE elif ChemShellCalculation.FILE_DLFIND in self.retrieved.list_object_names(): descrip = "Optimised structure from a ChemShell optimisation" input_pk = self.node.inputs.structure.pk From 859ae18d0e92fe169e14b3890a3eebdb60aa3bfd Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 18 Nov 2025 14:06:12 +0000 Subject: [PATCH 08/35] Optionally port command line parameters for mpi running with the chemsh wrapper script (instead of the chemsh.x binary) --- src/aiida_chemshell/calculations/base.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 90b82e8..9bad61c 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -181,7 +181,7 @@ def inputs_validator_wrapper(inputs, namespace): ## Metadata spec.inputs["metadata"]["options"]["resources"].default = { "num_machines": 1, - "num_mpiprocs_per_machine": 1, + "num_mpiprocs_per_machine": 4, } spec.inputs["metadata"]["options"]["parser_name"].default = "chemshell" @@ -848,9 +848,21 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: # Define the AiiDA code parameters code_info = CodeInfo() code_info.code_uuid = self.inputs.code.uuid - code_info.cmdline_params = [ - ChemShellCalculation.FILE_SCRIPT, - ] + if "chemsh.x" in str(self.inputs.code.filepath_executable): + code_info.cmdline_params = [ + ChemShellCalculation.FILE_SCRIPT, + ] + else: + n_machines = self.inputs.metadata.options.resources.get("num_machines") + n_mpi_pm = self.inputs.metadata.options.resources.get( + "num_mpiprocs_per_machine" + ) + tot_mpi = n_machines * n_mpi_pm + code_info.cmdline_params = [ + "-np", + self.inputs.metadata.options.resources.get("tot_num_mpiprocs", tot_mpi), + ChemShellCalculation.FILE_SCRIPT, + ] code_info.stdout_name = ChemShellCalculation.FILE_STDOUT # Setup the calculation information object From 9908b20ab9b7a941b4b74f730ec05c74f9a8ec85 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 18 Nov 2025 14:21:50 +0000 Subject: [PATCH 09/35] Create calc job that runs a python script to split ChemShell path files into individual extXYZ files for each step in the path --- pyproject.toml | 1 + .../calculations/splt_trajectory.py | 134 ++++++++++++++++++ .../parsers/split_trajectory.py | 34 +++++ 3 files changed, 169 insertions(+) create mode 100644 src/aiida_chemshell/calculations/splt_trajectory.py create mode 100644 src/aiida_chemshell/parsers/split_trajectory.py diff --git a/pyproject.toml b/pyproject.toml index 5fdc064..4c6ea72 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ Source = "https://github.com/stfc/aiida-chemshell" [project.entry-points."aiida.parsers"] "chemshell" = "aiida_chemshell.parsers.base:ChemShellParser" +"chemshell.split_traj" = "aiida_chemshell.parsers.split_trajectory:SplitTrajectoryParser" [project.entry-points."aiida.workflows"] "chemshell.opt" = "aiida_chemshell.workflows.optimisation:GeometryOptimisationWorkChain" diff --git a/src/aiida_chemshell/calculations/splt_trajectory.py b/src/aiida_chemshell/calculations/splt_trajectory.py new file mode 100644 index 0000000..1d33468 --- /dev/null +++ b/src/aiida_chemshell/calculations/splt_trajectory.py @@ -0,0 +1,134 @@ +"""CalcJob to create individual extended XYZ files for steps in an optimisation.""" + +from aiida.common import CalcInfo, CodeInfo +from aiida.common.folders import Folder +from aiida.engine import CalcJob, CalcJobProcessSpec +from aiida.orm import FolderData, SinglefileData, Str + + +class SplitTrajectory(CalcJob): + """CalcJob to split an XYZ trajectory into individual ext XYZ files.""" + + @classmethod + def define(cls, spec: CalcJobProcessSpec) -> None: + """Define the inputs, outputs and metadata for the CalcJob.""" + super().define(spec) + spec.input( + "path", + valid_type=SinglefileData, + required=True, + help="An XYZ trajectory file containing the positions of each step.", + ) + spec.input( + "force", + valid_type=SinglefileData, + required=True, + help="An XYZ trajectory file containing the forces at each step.", + ) + + spec.input( + "filename", + valid_type=Str, + required=False, + help="The name to give the directory of output files.", + ) + + spec.output( + "trajectory_folder", + valid_type=FolderData, + required=True, + help="The directory containing the resulting extXYZ files.", + ) + + ## Metadata + spec.inputs["metadata"]["options"]["resources"].default = { + "num_machines": 1, + "num_mpiprocs_per_machine": 2, + } + spec.inputs["metadata"]["options"][ + "parser_name" + ].default = "chemshell.split_traj" + + return + + def prepare_for_submission(self, folder: Folder) -> CalcInfo: + """Perform the python task to split the input trajectories.""" + script = self.generate_script() + with folder.open("input.py", "w") as f: + f.write(script) + + code_info = CodeInfo() + code_info.code_uuid = self.inputs.code.uuid + if "chemsh.x" in str(self.inputs.code.filepath_executable): + code_info.cmdline_params = [ + "input.py", + ] + else: + n_machines = self.inputs.metadata.options.resources.get("num_machines") + n_mpi_pm = self.inputs.metadata.options.resources.get( + "num_mpiprocs_per_machine" + ) + tot_mpi = n_machines * n_mpi_pm + code_info.cmdline_params = [ + "-np", + self.inputs.metadata.options.resources.get("tot_num_mpiprocs", tot_mpi), + "input.py", + ] + + calc_info = CalcInfo() + calc_info.codes_info = [code_info] + calc_info.retrieve_temporary_list = [] + calc_info.provenance_exclude_list = [] + calc_info.retrieve_temporary_list = [ + "trajectory*xyz", + ] + calc_info.local_copy_list = [ + ( + self.inputs.path.uuid, + self.inputs.path.filename, + self.inputs.path.filename, + ), + ( + self.inputs.force.uuid, + self.inputs.force.filename, + self.inputs.force.filename, + ), + ] + + return calc_info + + def generate_script(self) -> str: + """Generate the python script for splitting the trajectory files.""" + return """ +with open("path.xyz", "r") as f: + path = f.readlines() +with open("path_force.xyz", 'r') as f: + force = f.readlines() +natms = int(path[0]) +xyz_str = "\\n{0:8s} {1:12.8f} {2:12.8f} {3:12.8f} {4:12.8f} {5:12.8f} {6:12.8f}" +i = 2 +step = 0 +while i < len(path): + # Write a new xyz file for given step + with open(f"trajectory_{step:03d}.xyz", 'w') as f: + # write the number of atoms + f.write(f"{natms:10d}") + # Write the header line + f.write(f"\\nProperties=\\"species:S:1:pos:R:3:force:R:3\\" Step={step}") + for j in range(natms): + path_line = path[i + j].split() + force_line = force[i + j].split() + f.write( + xyz_str.format( + path_line[0], + float(path_line[1]), + float(path_line[2]), + float(path_line[3]), + float(force_line[1]), + float(force_line[2]), + float(force_line[3]) + ) + ) + i += natms + 2 + step += 1 +""" diff --git a/src/aiida_chemshell/parsers/split_trajectory.py b/src/aiida_chemshell/parsers/split_trajectory.py new file mode 100644 index 0000000..d67a63b --- /dev/null +++ b/src/aiida_chemshell/parsers/split_trajectory.py @@ -0,0 +1,34 @@ +"""Defines the parser for the split trajectory CalcJob.""" + +from pathlib import Path + +from aiida.engine import ExitCode +from aiida.orm import FolderData +from aiida.parsers.parser import Parser + + +class SplitTrajectoryParser(Parser): + """AiiDA parser plugin for SplitTrajectory CalcJob.""" + + def parse(self, **kwargs): + """AiiDA parser plugin for SplitTrajectory CalcJob.""" + traj_folder = FolderData( + label="ExtXYZ Trajectory Folder", + description=( + "Folder contianing extXYZ files for each step in from a given " + "trajectory." + ), + ) + tmp_folder = Path(kwargs["retrieved_temporary_folder"]) + for file in tmp_folder.iterdir(): + traj_folder.put_object_from_file(file, file.name) + # for file_name in self.retrieved.list_object_names(): + # if "trajectory" in file_name: + # traj_folder.put_object_from_bytes( + # self.retrieved.get_object_content(file_name, 'rb'), + # file_name + # ) + if len(traj_folder.list_object_names()) < 1: + print("ERROR :: no returned trajectory files") + self.out("trajectory_folder", traj_folder) + return ExitCode(0) From 12289b35e5795e41a659d96edae6df584be4c470 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 18 Nov 2025 14:44:10 +0000 Subject: [PATCH 10/35] Add exit code for when no trajectory files have been produced --- src/aiida_chemshell/calculations/splt_trajectory.py | 7 +++++++ src/aiida_chemshell/parsers/split_trajectory.py | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/aiida_chemshell/calculations/splt_trajectory.py b/src/aiida_chemshell/calculations/splt_trajectory.py index 1d33468..6f2a656 100644 --- a/src/aiida_chemshell/calculations/splt_trajectory.py +++ b/src/aiida_chemshell/calculations/splt_trajectory.py @@ -49,6 +49,13 @@ def define(cls, spec: CalcJobProcessSpec) -> None: "parser_name" ].default = "chemshell.split_traj" + # Exit Codes + spec.exit_code( + 300, + "ERROR_NO_TRAJECTORY_FILES", + message="No trajectory files have been produced.", + ) + return def prepare_for_submission(self, folder: Folder) -> CalcInfo: diff --git a/src/aiida_chemshell/parsers/split_trajectory.py b/src/aiida_chemshell/parsers/split_trajectory.py index d67a63b..07da71a 100644 --- a/src/aiida_chemshell/parsers/split_trajectory.py +++ b/src/aiida_chemshell/parsers/split_trajectory.py @@ -29,6 +29,6 @@ def parse(self, **kwargs): # file_name # ) if len(traj_folder.list_object_names()) < 1: - print("ERROR :: no returned trajectory files") + return self.exit_codes.ERROR_NO_TRAJECTORY_FILES self.out("trajectory_folder", traj_folder) return ExitCode(0) From 17c1923489b8026bedb1026e4ad3eabb92bdfe88 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Fri, 21 Nov 2025 15:47:51 +0000 Subject: [PATCH 11/35] Update basic workflow outline to include a call to aiida-mlip for single point calculation example --- src/aiida_chemshell/workflows/clf_ultra.py | 132 +++++++++++++++++---- 1 file changed, 112 insertions(+), 20 deletions(-) diff --git a/src/aiida_chemshell/workflows/clf_ultra.py b/src/aiida_chemshell/workflows/clf_ultra.py index 949b0e6..3ce300e 100644 --- a/src/aiida_chemshell/workflows/clf_ultra.py +++ b/src/aiida_chemshell/workflows/clf_ultra.py @@ -1,7 +1,16 @@ """CLF-ULTRA ChemShell + Janus workflow.""" -from aiida.engine import ToContext, WorkChain -from aiida.orm import Code, Dict, Float, Int, List, SinglefileData, StructureData +from aiida.engine import ToContext, WorkChain, calcfunction +from aiida.orm import ( + Code, + Dict, + Float, + FolderData, + List, + SinglefileData, + Str, + StructureData, +) from aiida.plugins.factories import CalculationFactory ChemShellCalculation = CalculationFactory("chemshell") @@ -28,6 +37,7 @@ def define(cls, spec) -> None: help="The structure on which to perform the geometry optimisation", ) spec.input("chemshell", valid_type=Code, required=True, help="") + spec.input("janus", valid_type=Code, required=True, help="") spec.input( "force_field_file", @@ -45,10 +55,19 @@ def define(cls, spec) -> None: ), ) - spec.output("final_energy", valid_type=Float) - spec.output("natms", valid_type=Int) + spec.output("final_qm_energy", valid_type=Float) + # spec.output("mlip_outputs", valid_type=Dict) + spec.output("trajectory_files", valid_type=FolderData) + spec.output("energy_difference", valid_type=Float) - spec.outline(cls.optimise, cls.generate_xyz_files, cls.result) + spec.outline( + cls.optimise, + cls.generate_xyz_files, + cls.extract_final_structure, + cls.mlip_sp, + cls.calculate_energy_difference, + cls.result, + ) return def optimise(self): @@ -56,11 +75,11 @@ def optimise(self): inputs = { "qm_parameters": Dict( { - "theory": "NWChem", - "method": "dft", - "functional": "b3lyp", - "basis": "cc-pvtz", - # Add D3 correction + "theory": "NWChem", # NWChem + "method": "dft", # DFT + "functional": "B3LYP", # B3LYP + "basis": "cc-pvtz", # cc=pvtz + # "d3": True # Add D3 correction } ), "optimisation_parameters": Dict({"save_path": True}), @@ -68,8 +87,8 @@ def optimise(self): "code": self.inputs.chemshell, "metadata": { "options": { - "resources": {"num_mpiprocs_per_machine": 4, "num_machines": 1}, - "withmpi": True, + "resources": {"num_mpiprocs_per_machine": 8, "num_machines": 1}, + # "withmpi": True, } }, } @@ -77,16 +96,89 @@ def optimise(self): return ToContext(optimise=future) def generate_xyz_files(self): - """Convert the paths to individual extended XYZ files.""" - trjp = self.ctx.optimise.outputs.trajectory_path.get_content(mode="r") - trjf = self.ctx.optimise.outputs.trajectory_force.get_content(mode="r") - trjp = trjp.split("\n") - trjf = trjf.split("\n") - natms = int(trjp[0]) - self.out("natms", natms) + """Convert the paths to individual extended XYZ filesn for each step.""" + inputs = { + "path": self.ctx.optimise.outputs.trajectory_path, + "force": self.ctx.optimise.outputs.trajectory_force, + "code": self.inputs.chemshell, + "metadata": { + "options": { + "resources": {"num_mpiprocs_per_machine": 2, "num_machines": 1}, + # "withmpi": True, + } + }, + } + from aiida_chemshell.calculations.splt_trajectory import SplitTrajectory + + future = self.submit(SplitTrajectory, **inputs) + return ToContext(split_trajectory=future) + + def extract_final_structure(self): + """Extract the optimised structure from the XYZ trajectory folder.""" + self.ctx.final_structure = extract_final_structure( + self.ctx.split_trajectory.outputs.trajectory_folder + ) + return + + def mlip_sp(self): + """Run a single point energy with Janus-Core on the optimised structure.""" + structure = self.ctx.final_structure + from aiida_mlip.helpers.help_load import load_model + + model = load_model(None, "mace_mp") + inputs = { + "metadata": {"options": {"resources": {"num_machines": 1}}}, + "code": self.inputs.janus, + # "arch": Str(model.architecture), + "struct": structure, + "model": model, + "device": Str("cpu"), + # "calc_kwargs": Dict({"dispersion": True}), + } + mlip_sp = CalculationFactory("mlip.sp") + future = self.submit(mlip_sp, **inputs) + return ToContext(mlip_sp=future) + + def calculate_energy_difference(self): + """Calculate the difference in final energies.""" + self.ctx.energy_difference = calculate_difference( + self.ctx.optimise.outputs.energy, self.ctx.mlip_sp.outputs.results_dict + ) return def result(self) -> None: """Report the final results.""" - self.out("final_energy", self.ctx.optimise.outputs.energy) + self.out("final_qm_energy", self.ctx.optimise.outputs.energy) + self.out( + "trajectory_files", self.ctx.split_trajectory.outputs.trajectory_folder + ) + # self.out( + # "mlip_outputs", + # self.ctx.mlip_sp.outputs.results_dict + # ) + self.out("energy_difference", self.ctx.energy_difference) return + + +@calcfunction +def extract_final_structure(folder: FolderData) -> StructureData: + """Extract the final optimised structure from a folder of xyz files.""" + structure_file = folder.list_object_names()[-1] + structure = StructureData() + structure._parse_xyz(folder.get_object_content(structure_file, mode="r") + "\n") + return structure + + +@calcfunction +def calculate_difference(qm_energy: Float, ml_results: Dict) -> Float: + """Calculate the difference between the final energies.""" + ml_energy = ml_results["info"]["mace_mp_energy"] + diff = qm_energy - ml_energy + return Float( + diff, + label="Final Energy Difference", + description=( + "Difference in final energies between DFT calculation and " + "MACE_mp energy calculation using Janus-core." + ), + ) From 9df4c706cbc3a46066b00224a1f4db6e47f1410a Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 20 May 2026 15:12:56 +0100 Subject: [PATCH 12/35] Update split_trajectory job to correctly generate required inputs for janus mlip training input (still WIP) --- pyproject.toml | 2 +- ...{splt_trajectory.py => file_conversion.py} | 78 ++++++++++--------- .../parsers/file_conversion.py | 48 ++++++++++++ .../parsers/split_trajectory.py | 34 -------- src/aiida_chemshell/workflows/optimisation.py | 71 ++++++++++++++++- tests/conftest.py | 15 ++++ tests/test_workflows.py | 40 ++++++++++ 7 files changed, 216 insertions(+), 72 deletions(-) rename src/aiida_chemshell/calculations/{splt_trajectory.py => file_conversion.py} (68%) create mode 100644 src/aiida_chemshell/parsers/file_conversion.py delete mode 100644 src/aiida_chemshell/parsers/split_trajectory.py diff --git a/pyproject.toml b/pyproject.toml index 4c6ea72..d571b86 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,7 +32,7 @@ Source = "https://github.com/stfc/aiida-chemshell" [project.entry-points."aiida.parsers"] "chemshell" = "aiida_chemshell.parsers.base:ChemShellParser" -"chemshell.split_traj" = "aiida_chemshell.parsers.split_trajectory:SplitTrajectoryParser" +"chemshell.file_conversion.mlip_training" = "aiida_chemshell.parsers.file_conversion:CreateJanusTrainingInputsParser" [project.entry-points."aiida.workflows"] "chemshell.opt" = "aiida_chemshell.workflows.optimisation:GeometryOptimisationWorkChain" diff --git a/src/aiida_chemshell/calculations/splt_trajectory.py b/src/aiida_chemshell/calculations/file_conversion.py similarity index 68% rename from src/aiida_chemshell/calculations/splt_trajectory.py rename to src/aiida_chemshell/calculations/file_conversion.py index 6f2a656..45029a4 100644 --- a/src/aiida_chemshell/calculations/splt_trajectory.py +++ b/src/aiida_chemshell/calculations/file_conversion.py @@ -3,10 +3,10 @@ from aiida.common import CalcInfo, CodeInfo from aiida.common.folders import Folder from aiida.engine import CalcJob, CalcJobProcessSpec -from aiida.orm import FolderData, SinglefileData, Str +from aiida.orm import SinglefileData, Str -class SplitTrajectory(CalcJob): +class CreateJanusTrainingInputsCalcJob(CalcJob): """CalcJob to split an XYZ trajectory into individual ext XYZ files.""" @classmethod @@ -34,10 +34,22 @@ def define(cls, spec: CalcJobProcessSpec) -> None: ) spec.output( - "trajectory_folder", - valid_type=FolderData, + "training_input", + valid_type=SinglefileData, + required=True, + help="The main training data set in extended XYZ format.", + ) + spec.output( + "validation_input", + valid_type=SinglefileData, required=True, - help="The directory containing the resulting extXYZ files.", + help="The validation data set in extended XYZ format.", + ) + spec.output( + "test_input", + valid_type=SinglefileData, + requried=True, + help="The testing data set in extended XYZ format.", ) ## Metadata @@ -47,7 +59,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None: } spec.inputs["metadata"]["options"][ "parser_name" - ].default = "chemshell.split_traj" + ].default = "chemshell.file_conversion.mlip_training" # Exit Codes spec.exit_code( @@ -86,9 +98,7 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: calc_info.codes_info = [code_info] calc_info.retrieve_temporary_list = [] calc_info.provenance_exclude_list = [] - calc_info.retrieve_temporary_list = [ - "trajectory*xyz", - ] + calc_info.retrieve_temporary_list = ["train.xyz", "valid.xyz", "test.xyz"] calc_info.local_copy_list = [ ( self.inputs.path.uuid, @@ -109,33 +119,29 @@ def generate_script(self) -> str: return """ with open("path.xyz", "r") as f: path = f.readlines() -with open("path_force.xyz", 'r') as f: +with open("path_force.xyz", "r") as f: force = f.readlines() + natms = int(path[0]) -xyz_str = "\\n{0:8s} {1:12.8f} {2:12.8f} {3:12.8f} {4:12.8f} {5:12.8f} {6:12.8f}" -i = 2 -step = 0 -while i < len(path): - # Write a new xyz file for given step - with open(f"trajectory_{step:03d}.xyz", 'w') as f: - # write the number of atoms - f.write(f"{natms:10d}") - # Write the header line - f.write(f"\\nProperties=\\"species:S:1:pos:R:3:force:R:3\\" Step={step}") - for j in range(natms): - path_line = path[i + j].split() - force_line = force[i + j].split() - f.write( - xyz_str.format( - path_line[0], - float(path_line[1]), - float(path_line[2]), - float(path_line[3]), - float(force_line[1]), - float(force_line[2]), - float(force_line[3]) - ) - ) - i += natms + 2 - step += 1 +nsteps = len(path) // (natms + 2) +valid_interval = 5 +test_interval = 10 +if nsteps < test_interval: + test_interval = nsteps + valid_interval = (nsteps // 2) + 1 +for step in range(nsteps): + if (step + 1) % test_interval == 0: + fname = "test.xyz" + elif (step + 1) % valid_interval == 0: + fname = "valid.xyz" + else: + fname = "train.xyz" + with open(fname, "a") as f: + f.write(f"{natms}\\n") + f.write('Properties="species:S:1:pos:R:3:force:R:3"\\n') + for i in range(2, natms + 2): + index = (step * (natms + 2)) + i + f.write(path[index].strip("\\n") + " ") + f.write(" ".join(force[index].split()[1:])) + f.write("\\n") """ diff --git a/src/aiida_chemshell/parsers/file_conversion.py b/src/aiida_chemshell/parsers/file_conversion.py new file mode 100644 index 0000000..f18ace4 --- /dev/null +++ b/src/aiida_chemshell/parsers/file_conversion.py @@ -0,0 +1,48 @@ +"""Defines the parser for the split trajectory CalcJob.""" + +from aiida.engine import ExitCode +from aiida.orm import SinglefileData +from aiida.parsers.parser import Parser + + +class CreateJanusTrainingInputsParser(Parser): + """AiiDA parser plugin for SplitTrajectory CalcJob.""" + + def parse(self, **kwargs): + """AiiDA parser plugin for SplitTrajectory CalcJob.""" + description_str = "MLIP {} data set extracted from ChemShell calculation." + + with self.retrieved.open("train.xyz", "r") as f: + self.out( + "training_input", + SinglefileData( + file=f, + filename="train.xyz", + label="MLIP training data.", + description=description_str.format("training"), + ), + ) + + with self.retrieved.open("test.xyz", "r") as f: + self.out( + "test_input", + SinglefileData( + file=f, + filename="test.xyz", + label="MLIP test data.", + description=description_str.format("testing"), + ), + ) + + with self.retrieved.open("valid.xyz", "r") as f: + self.out( + "validation_input", + SinglefileData( + file=f, + filename="valid.xyz", + label="MLIP validation data.", + description=description_str.format("validation"), + ), + ) + + return ExitCode(0) diff --git a/src/aiida_chemshell/parsers/split_trajectory.py b/src/aiida_chemshell/parsers/split_trajectory.py deleted file mode 100644 index 07da71a..0000000 --- a/src/aiida_chemshell/parsers/split_trajectory.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Defines the parser for the split trajectory CalcJob.""" - -from pathlib import Path - -from aiida.engine import ExitCode -from aiida.orm import FolderData -from aiida.parsers.parser import Parser - - -class SplitTrajectoryParser(Parser): - """AiiDA parser plugin for SplitTrajectory CalcJob.""" - - def parse(self, **kwargs): - """AiiDA parser plugin for SplitTrajectory CalcJob.""" - traj_folder = FolderData( - label="ExtXYZ Trajectory Folder", - description=( - "Folder contianing extXYZ files for each step in from a given " - "trajectory." - ), - ) - tmp_folder = Path(kwargs["retrieved_temporary_folder"]) - for file in tmp_folder.iterdir(): - traj_folder.put_object_from_file(file, file.name) - # for file_name in self.retrieved.list_object_names(): - # if "trajectory" in file_name: - # traj_folder.put_object_from_bytes( - # self.retrieved.get_object_content(file_name, 'rb'), - # file_name - # ) - if len(traj_folder.list_object_names()) < 1: - return self.exit_codes.ERROR_NO_TRAJECTORY_FILES - self.out("trajectory_folder", traj_folder) - return ExitCode(0) diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index acabd5c..5daf456 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -1,7 +1,9 @@ """Workflows for geometry optimisation based taks.""" +from aiida.common.exceptions import MissingEntryPointError from aiida.engine import ToContext, WorkChain from aiida.orm import ArrayData, Bool, Dict, Float, SinglefileData, Str +from aiida.plugins.factories import CalculationFactory from aiida_chemshell.calculations.base import ChemShellCalculation @@ -31,6 +33,36 @@ def define(cls, spec) -> None: help="Set basis set quality for QM calculation based on defined options.", ) + try: + mlip_train_calc = CalculationFactory("mlip.train") + except MissingEntryPointError: + pass + else: + print("Exposiing the mlip inputs...") + spec.expose_inputs( + mlip_train_calc, + namespace="mlip", + namespace_options={ + "required": False, + "populate_defaults": False, + "help": ( + "Optional inputs to fine-tune the MLIP model using aiida-mlip." + ), + }, + ) + spec.expose_outputs( + mlip_train_calc, + namespace="mlip", + namespace_options={ + "required": False, + "populate_defaults": False, + "help": ( + "Optional outputs from fine-tuning the MLIP model using " + "aiida-mlip." + ), + }, + ) + ## Outputs ## spec.output( "final_energy", @@ -58,7 +90,7 @@ def define(cls, spec) -> None: ) ## Workflow ## - spec.outline(cls.optimise, cls.energy, cls.result) + spec.outline(cls.optimise, cls.energy, cls.train_mlip, cls.result) return @@ -117,10 +149,47 @@ def energy(self): self.ctx.optimise.inputs.optimisation_parameters.get_dict() ) inputs["optimisation_parameters"]["thermal"] = True + if self.inputs.get("mlip", None): + inputs["optimisation_parameters"]["save_path"] = True future = self.submit(ChemShellCalculation, **inputs) return ToContext(energy=future) return None + def generate_mlip_training_inputs(self): + """Convert the optimisation path files to Janus compatible inputs.""" + if "mlip" in self.inputs: + inputs = { + "path": self.ctx.optimise.outputs.trajectory_path, + "force": self.ctx.optimise.outputs.trajectory_force, + "code": self.inputs.chemshell, + "metadata": { + "options": { + "resources": {"num_mpiprocs_per_machine": 2, "num_machines": 1}, + # "withmpi": True, + } + }, + } + from aiida_chemshell.calculations.file_conversion import ( + CreateJanusTrainingInputsCalcJob, + ) + + future = self.submit(CreateJanusTrainingInputsCalcJob, **inputs) + return ToContext(create_mlip_inputs=future) + return None + + def train_mlip(self): + """Train a given MLIP model.""" + try: + mlip_train_calc = CalculationFactory("mlip.train") + except MissingEntryPointError: + pass + else: + if "mlip" in self.inputs: + mlip_inputs = self.exposed_inputs(mlip_train_calc, namespace="mlip") + future = self.submit(mlip_train_calc, **mlip_inputs) + return ToContext(mlip_training=future) + return None + def result(self): """Extract the final workflow results.""" self.out("optimised_structure", self.ctx.optimise.outputs.optimised_structure) diff --git a/tests/conftest.py b/tests/conftest.py index 10accce..f651160 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,6 +44,21 @@ def factory(plugin: str = "chemshell") -> InstalledCode: return factory +@pytest.fixture(scope="function") +def janus_code(aiida_code_installed): + """Return a Janus AiiDA code instance.""" + import os + import shutil + + janus_path = shutil.which("janus") or os.environ.get("JANUS_PATH") + + return aiida_code_installed( + label="janus", + default_calc_job_plugin="mlip.sp", + filepath_executable=janus_path, + ) + + @pytest.fixture def water_structure_object() -> StructureData: """Return a AiiDA StructureData object of a water molecule.""" diff --git a/tests/test_workflows.py b/tests/test_workflows.py index 0183a12..e52032d 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -40,3 +40,43 @@ def test_geometry_optimisation_workflow(chemsh_code, get_test_data_file): modes = results.get("vibrational_modes").get_array("Modes") assert (modes[0][0] - 1799.584) < 1e-10, "Incorrect frequency reported for mode 1" assert (modes[2][2] - 0.0089900284) < 1e-10, "Incorrect ZPE reported for mode 3" + + +def test_optimisation_workflow_mlip_training( + chemsh_code, get_test_data_file, janus_code +): + """Test a geometry optimisation workflow with vibrational analysis.""" + from aiida_mlip.data.config import JanusConfigfile + from aiida_mlip.helpers.help_load import load_model + + inputs = { + "chemsh": { + "code": chemsh_code(), + "structure": get_test_data_file(), + "qm_parameters": {"theory": "PySCF", "method": "HF"}, + }, + "mlip": { + "mlip_config": JanusConfigfile( + "/opt/aiida-chemshell/.local/janus_config.yml" + ), + "code": janus_code, + "fine_tune": True, + "foundation_model": load_model(None, "mace_mp"), + "metadata": {"options": {"resources": {"num_machines": 1}}}, + }, + "basis_quality": "fast", + "vibrational_analysis": False, + } + results, node = run_get_node(GeometryOptimisationWorkChain, **inputs) + + assert node.is_finished_ok, f"WorkChain failed with exit status {node.exit_status}" + + assert len(node.called) > 0, "WorkChain did not launch any subprocesses" + + assert node.called[0].inputs.qm_parameters.get( + "basis", "" + ) == GeometryOptimisationWorkChain.get_basis_set_label("fast") + + assert abs(results.get("final_energy") - -75.585959742867) < 1e-10, ( + "Incorrect final energy for geometry optimisation workflow." + ) From bfed4166af4a0dccc6de05662486311887ab6b07 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Thu, 21 May 2026 10:19:13 +0100 Subject: [PATCH 13/35] Add the energy of each step in an optimisation as a result of the base optimisation CalcJob --- src/aiida_chemshell/calculations/base.py | 8 ++++++++ src/aiida_chemshell/parsers/base.py | 25 ++++++++++++++++++++++++ tests/test_calculations.py | 19 ++++++++++++++++++ 3 files changed, 52 insertions(+) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 9bad61c..2ad3e24 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -133,6 +133,14 @@ def define(cls, spec: CalcJobProcessSpec) -> None: "structure is contained within a ChemShell '.pun' file." ), ) + spec.output( + "optimisation_path", + valid_type=ArrayData, + required=False, + help=( + "Values calculated at each step of an optimisation based calculation." + ), + ) spec.output( "vibrational_energies", valid_type=Dict, diff --git a/src/aiida_chemshell/parsers/base.py b/src/aiida_chemshell/parsers/base.py index 5f2da1f..3001c5b 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -104,6 +104,11 @@ def parse(self, **kwargs): label="ChemShell optimisation trajectory.", ), ) + self.parse_optimisation_path( + self.retrieved.get_object_content( + ChemShellCalculation.FILE_STDOUT, "r" + ) + ) else: return self.exit_codes.ERROR_MISSING_OPTIMISED_STRUCTURE_FILE elif ChemShellCalculation.FILE_DLFIND in self.retrieved.list_object_names(): @@ -124,6 +129,11 @@ def parse(self, **kwargs): description=descrip, ), ) + self.parse_optimisation_path( + self.retrieved.get_object_content( + ChemShellCalculation.FILE_STDOUT, "r" + ) + ) else: return self.exit_codes.ERROR_MISSING_OPTIMISED_STRUCTURE_FILE @@ -168,3 +178,18 @@ def parse_vibrational_analysis(self, stdout: str) -> None: modes_data_node.set_array("Modes", modes) self.out("vibrational_modes", modes_data_node) return + + def parse_optimisation_path(self, stdout: str) -> None: + """Extract per step values from the optimisation job.""" + energies = [] + for line in stdout.split("\n"): + if "Energy calculation finished" in line: + energies.append(float(line.split()[-1])) + + results = ArrayData( + label="Optimisation Path Properties", + description="Values calculated at each step of an optimisation.", + ) + results.set_array("energies", numpy.asarray(energies)) + self.out("optimisation_path", results) + return diff --git a/tests/test_calculations.py b/tests/test_calculations.py index 7e407e4..35cf5a5 100644 --- a/tests/test_calculations.py +++ b/tests/test_calculations.py @@ -178,6 +178,25 @@ def test_opt_calculation_qm_dft(chemsh_code, get_test_data_file): "Incorrect energy result for PySCF based optimisation calculation." ) + assert "optimisation_path" in results, "Missing optimisation_path output node." + optimisation_path = results.get("optimisation_path") + + assert optimisation_path.description != "", ( + "Missing optimisation_path array node description" + ) + + assert optimisation_path.get_arraynames() == ["energies"], ( + "Optimisation path energies entry missing." + ) + + assert optimisation_path.get_shape("energies")[0] == 7, ( + "Incorrect number of entries for the optimisation_path energies." + ) + + assert abs(optimisation_path.get_array("energies")[-1] - eref) < 1e-7, ( + "Incorrect final energy in optimisation_path energy series." + ) + def test_opt_calculation_dlpoly(chemsh_code, get_test_data_file): """MM based geometry optimisation test.""" From 254e5a134416438359a2da7adb7d90cf94878dba Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Tue, 26 May 2026 12:06:50 +0100 Subject: [PATCH 14/35] Add default labelling of subprocesses in a geometry optimisation workflow --- src/aiida_chemshell/calculations/base.py | 23 +++++++++++++++---- src/aiida_chemshell/workflows/optimisation.py | 18 ++++++++++++++- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 2ad3e24..ad22eca 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -691,22 +691,37 @@ def _build_process_label(self) -> str: Defines the process label to be associated with the created ProcessNode stored in the AiiDA database. + Returns + ------- + str + The process label based on what inputs have been provided. + """ + return ChemShellCalculation.default_process_label(self) + + @classmethod + def default_process_label(cls, node) -> str: + """ + AiiDA Process label definition (Class Method). + + Defines the process label to be associated with the created ProcessNode + stored in the AiiDA database. + Returns ------- str The process label based on what inputs have been provided. """ theory_key = "" - if "qm_parameters" in self.inputs: - if "mm_parameters" in self.inputs: + if "qm_parameters" in node.inputs: + if "mm_parameters" in node.inputs: theory_key = "_(QM/MM)" else: theory_key = "_(QM)" else: theory_key = "_(MM)" - if "optimisation_parameters" in self.inputs: - if self.inputs.optimisation_parameters.get("thermal", False): + if "optimisation_parameters" in node.inputs: + if node.inputs.optimisation_parameters.get("thermal", False): return "ChemShell_Vibrational_Frequencies" + theory_key return "ChemShell_Geometry_Optimisation" + theory_key diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index 5daf456..bc8dc4f 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -38,7 +38,6 @@ def define(cls, spec) -> None: except MissingEntryPointError: pass else: - print("Exposiing the mlip inputs...") spec.expose_inputs( mlip_train_calc, namespace="mlip", @@ -126,6 +125,10 @@ def optimise(self): inputs["optimisation_parameters"] = Dict({}) future = self.submit(ChemShellCalculation, **inputs) + future.label = ChemShellCalculation.default_process_label(future) + future.description = ( + f"Geometry optimisation step from WorkChainNode pk: {self.node.pk}" + ) return ToContext(optimise=future) def energy(self): @@ -152,6 +155,11 @@ def energy(self): if self.inputs.get("mlip", None): inputs["optimisation_parameters"]["save_path"] = True future = self.submit(ChemShellCalculation, **inputs) + future.label = ChemShellCalculation.default_process_label(future) + future.description = ( + f"Vibrational frequency calculation step from WorkChainNode " + f"pk: {self.node.pk}" + ) return ToContext(energy=future) return None @@ -174,6 +182,10 @@ def generate_mlip_training_inputs(self): ) future = self.submit(CreateJanusTrainingInputsCalcJob, **inputs) + future.label = "Generate MLIP training data set from geometry optimisation." + future.description = ( + f"Data extraction step from WorkChainNode pk: {self.node.pk}" + ) return ToContext(create_mlip_inputs=future) return None @@ -187,6 +199,10 @@ def train_mlip(self): if "mlip" in self.inputs: mlip_inputs = self.exposed_inputs(mlip_train_calc, namespace="mlip") future = self.submit(mlip_train_calc, **mlip_inputs) + future.label = "MLIP Fine-Tuning." + future.description = ( + f"MLIP fine-tuning step from WorkChainNode pk: {self.node.pk}" + ) return ToContext(mlip_training=future) return None From 9171b6c78f5cf43ae0b2859e0f2f57e3fd599cb1 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Thu, 11 Jun 2026 10:04:43 +0100 Subject: [PATCH 15/35] Enable fine-tuning of MLIP model as additional step to an optimisation workflow --- .../calculations/file_conversion.py | 24 ++- .../parsers/file_conversion.py | 12 +- src/aiida_chemshell/utils.py | 40 +++++ src/aiida_chemshell/workflows/clf_ultra.py | 30 ++-- src/aiida_chemshell/workflows/optimisation.py | 148 +++++++++++++----- tests/test_workflows.py | 56 ++++--- 6 files changed, 229 insertions(+), 81 deletions(-) diff --git a/src/aiida_chemshell/calculations/file_conversion.py b/src/aiida_chemshell/calculations/file_conversion.py index 45029a4..ee75b7d 100644 --- a/src/aiida_chemshell/calculations/file_conversion.py +++ b/src/aiida_chemshell/calculations/file_conversion.py @@ -3,7 +3,7 @@ from aiida.common import CalcInfo, CodeInfo from aiida.common.folders import Folder from aiida.engine import CalcJob, CalcJobProcessSpec -from aiida.orm import SinglefileData, Str +from aiida.orm import ArrayData, SinglefileData, Str class CreateJanusTrainingInputsCalcJob(CalcJob): @@ -25,6 +25,12 @@ def define(cls, spec: CalcJobProcessSpec) -> None: required=True, help="An XYZ trajectory file containing the forces at each step.", ) + spec.input( + "energies", + valid_type=ArrayData, + required=True, + help="The calculated dft energy of each step in the data series.", + ) spec.input( "filename", @@ -48,7 +54,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None: spec.output( "test_input", valid_type=SinglefileData, - requried=True, + required=True, help="The testing data set in extended XYZ format.", ) @@ -76,6 +82,10 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: with folder.open("input.py", "w") as f: f.write(script) + with folder.open("energies.txt", "w") as f: + for val in self.inputs.energies.get_array("energies"): + f.write(f"{val:.10f}\n") + code_info = CodeInfo() code_info.code_uuid = self.inputs.code.uuid if "chemsh.x" in str(self.inputs.code.filepath_executable): @@ -96,7 +106,6 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: calc_info = CalcInfo() calc_info.codes_info = [code_info] - calc_info.retrieve_temporary_list = [] calc_info.provenance_exclude_list = [] calc_info.retrieve_temporary_list = ["train.xyz", "valid.xyz", "test.xyz"] calc_info.local_copy_list = [ @@ -121,7 +130,10 @@ def generate_script(self) -> str: path = f.readlines() with open("path_force.xyz", "r") as f: force = f.readlines() - +energies = [] +with open("energies.txt", 'r') as f: + for line in f: + energies.append(float(line)) natms = int(path[0]) nsteps = len(path) // (natms + 2) valid_interval = 5 @@ -138,7 +150,9 @@ def generate_script(self) -> str: fname = "train.xyz" with open(fname, "a") as f: f.write(f"{natms}\\n") - f.write('Properties="species:S:1:pos:R:3:force:R:3"\\n') + f.write(f'Step={step} ') + f.write(f'Properties=species:S:1:pos:R:3:dft_forces:R:3 pbc="F F F" ') + f.write(f'energy={energies[step]}\\n') for i in range(2, natms + 2): index = (step * (natms + 2)) + i f.write(path[index].strip("\\n") + " ") diff --git a/src/aiida_chemshell/parsers/file_conversion.py b/src/aiida_chemshell/parsers/file_conversion.py index f18ace4..fef13d0 100644 --- a/src/aiida_chemshell/parsers/file_conversion.py +++ b/src/aiida_chemshell/parsers/file_conversion.py @@ -1,5 +1,7 @@ """Defines the parser for the split trajectory CalcJob.""" +import os + from aiida.engine import ExitCode from aiida.orm import SinglefileData from aiida.parsers.parser import Parser @@ -11,8 +13,10 @@ class CreateJanusTrainingInputsParser(Parser): def parse(self, **kwargs): """AiiDA parser plugin for SplitTrajectory CalcJob.""" description_str = "MLIP {} data set extracted from ChemShell calculation." + retrieved_temporary_folder = kwargs.get("retrieved_temporary_folder", None) - with self.retrieved.open("train.xyz", "r") as f: + # with self.retrieved.open("train.xyz", "r") as f: + with open(os.path.join(retrieved_temporary_folder, "train.xyz"), "rb") as f: self.out( "training_input", SinglefileData( @@ -23,7 +27,8 @@ def parse(self, **kwargs): ), ) - with self.retrieved.open("test.xyz", "r") as f: + # with self.retrieved.open("test.xyz", "r") as f: + with open(os.path.join(retrieved_temporary_folder, "test.xyz"), "rb") as f: self.out( "test_input", SinglefileData( @@ -34,7 +39,8 @@ def parse(self, **kwargs): ), ) - with self.retrieved.open("valid.xyz", "r") as f: + # with self.retrieved.open("valid.xyz", "r") as f: + with open(os.path.join(retrieved_temporary_folder, "valid.xyz"), "rb") as f: self.out( "validation_input", SinglefileData( diff --git a/src/aiida_chemshell/utils.py b/src/aiida_chemshell/utils.py index 33c8f10..dc50927 100644 --- a/src/aiida_chemshell/utils.py +++ b/src/aiida_chemshell/utils.py @@ -85,3 +85,43 @@ def generate_parameter_string(params: dict) -> str: else: s += f"{key}={params[key]}, " return s.rstrip(", ") + + +def generate_default_mlip_fine_tune_config(): + """Generate a default configuration for mlip fine-tuning via Janus.""" + return { + # "multiheads_finetuning": True, + "foundation_filter_elements": True, + "foundation_model_readout": True, + "foundation_model_elements": False, + "loss": "universal", + "weight_pt_head": 10.0, + "energy_weight": 1.0, + "forces_weight": 10.0, + "stress_weight": 10.0, + "stress_key": "stress", + "energy_key": "energy", + "forces_key": "forces", + "compute_stress": False, + "compute_forces": True, + "clip_grad": 10, + "error_table": "PerAtomRMSE", + "scaling": "rms_forces_scaling", + "force_mh_ft_lr": True, + "lr": 0.0001, + "batch_size": 2, + "max_num_epochs": 10, + "ema": True, + "ema_decay": 0.99999, + "amsgrad": True, + "default_dtype": "float64", + "device": "cpu", + "restart_latest": True, + # "seed": 2024, + "keep_isolated_atoms": True, + "save_cpu": True, + "weight_decay": 1e-8, + "eval_interval": 1, + # "enable_cueq": True, + "E0s": {1: -13.664311914087541, 6: -1030.2235576550245, 8: -2043.4891174666404}, + } diff --git a/src/aiida_chemshell/workflows/clf_ultra.py b/src/aiida_chemshell/workflows/clf_ultra.py index 3ce300e..9bb2c27 100644 --- a/src/aiida_chemshell/workflows/clf_ultra.py +++ b/src/aiida_chemshell/workflows/clf_ultra.py @@ -97,21 +97,21 @@ def optimise(self): def generate_xyz_files(self): """Convert the paths to individual extended XYZ filesn for each step.""" - inputs = { - "path": self.ctx.optimise.outputs.trajectory_path, - "force": self.ctx.optimise.outputs.trajectory_force, - "code": self.inputs.chemshell, - "metadata": { - "options": { - "resources": {"num_mpiprocs_per_machine": 2, "num_machines": 1}, - # "withmpi": True, - } - }, - } - from aiida_chemshell.calculations.splt_trajectory import SplitTrajectory - - future = self.submit(SplitTrajectory, **inputs) - return ToContext(split_trajectory=future) + # inputs = { + # "path": self.ctx.optimise.outputs.trajectory_path, + # "force": self.ctx.optimise.outputs.trajectory_force, + # "code": self.inputs.chemshell, + # "metadata": { + # "options": { + # "resources": {"num_mpiprocs_per_machine": 2, "num_machines": 1}, + # # "withmpi": True, + # } + # }, + # } + + # future = self.submit(SplitTrajectory, **inputs) + # return ToContext(split_trajectory=future) + return def extract_final_structure(self): """Extract the optimised structure from the XYZ trajectory folder.""" diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index bc8dc4f..f6d7e53 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -2,7 +2,7 @@ from aiida.common.exceptions import MissingEntryPointError from aiida.engine import ToContext, WorkChain -from aiida.orm import ArrayData, Bool, Dict, Float, SinglefileData, Str +from aiida.orm import ArrayData, Bool, Code, Dict, Float, SinglefileData, Str from aiida.plugins.factories import CalculationFactory from aiida_chemshell.calculations.base import ChemShellCalculation @@ -33,35 +33,6 @@ def define(cls, spec) -> None: help="Set basis set quality for QM calculation based on defined options.", ) - try: - mlip_train_calc = CalculationFactory("mlip.train") - except MissingEntryPointError: - pass - else: - spec.expose_inputs( - mlip_train_calc, - namespace="mlip", - namespace_options={ - "required": False, - "populate_defaults": False, - "help": ( - "Optional inputs to fine-tune the MLIP model using aiida-mlip." - ), - }, - ) - spec.expose_outputs( - mlip_train_calc, - namespace="mlip", - namespace_options={ - "required": False, - "populate_defaults": False, - "help": ( - "Optional outputs from fine-tuning the MLIP model using " - "aiida-mlip." - ), - }, - ) - ## Outputs ## spec.output( "final_energy", @@ -88,8 +59,52 @@ def define(cls, spec) -> None: help="The calculated vibrational modes for the optimised structure.", ) + spec.output( + "test_file", + valid_type=SinglefileData, + required=False, + help="", + ) + + # Optional inputs/outputs for using the results to fine tune a MLIP model + # using the janus-core project via the aiida-mlip plugin. + try: + CalculationFactory("mlip.train") + except MissingEntryPointError: + pass + else: + from aiida_mlip.data.model import ModelData + + spec.input( + "mlip_model", + valid_type=ModelData, + required=False, + help="The MLIP foundation model to apply fine-tuning to.", + ) + spec.input( + "mlip_code", + valid_type=Code, + required=False, + help="The Janus-core AiiDA code instance.", + ) + spec.output( + "fine_tuned_model", valid_type=ModelData, required=False, help="" + ) + spec.output( + "fine_tuned_model_compiled", + valid_type=SinglefileData, + required=False, + help="", + ) + ## Workflow ## - spec.outline(cls.optimise, cls.energy, cls.train_mlip, cls.result) + spec.outline( + cls.optimise, + cls.energy, + cls.generate_mlip_training_inputs, + cls.train_mlip, + cls.result, + ) return @@ -123,6 +138,8 @@ def optimise(self): return None if "optimisation_parameters" not in inputs: inputs["optimisation_parameters"] = Dict({}) + if "mlip_model" in self.inputs: + inputs["optimisation_parameters"]["save_path"] = True future = self.submit(ChemShellCalculation, **inputs) future.label = ChemShellCalculation.default_process_label(future) @@ -152,8 +169,6 @@ def energy(self): self.ctx.optimise.inputs.optimisation_parameters.get_dict() ) inputs["optimisation_parameters"]["thermal"] = True - if self.inputs.get("mlip", None): - inputs["optimisation_parameters"]["save_path"] = True future = self.submit(ChemShellCalculation, **inputs) future.label = ChemShellCalculation.default_process_label(future) future.description = ( @@ -165,11 +180,12 @@ def energy(self): def generate_mlip_training_inputs(self): """Convert the optimisation path files to Janus compatible inputs.""" - if "mlip" in self.inputs: + if "mlip_model" in self.inputs: inputs = { "path": self.ctx.optimise.outputs.trajectory_path, "force": self.ctx.optimise.outputs.trajectory_force, - "code": self.inputs.chemshell, + "energies": self.ctx.optimise.outputs.optimisation_path, + "code": self.inputs.chemsh.code, "metadata": { "options": { "resources": {"num_mpiprocs_per_machine": 2, "num_machines": 1}, @@ -196,8 +212,54 @@ def train_mlip(self): except MissingEntryPointError: pass else: - if "mlip" in self.inputs: - mlip_inputs = self.exposed_inputs(mlip_train_calc, namespace="mlip") + pass + if "mlip_model" in self.inputs: + # This needs to be properly addressed within aiida-mlip + computer = self.inputs.get("mlip_code", None).computer + work_dir = computer.get_workdir() + with open("train.xyz", mode="wb") as f: + f.write(self.ctx.create_mlip_inputs.outputs.training_input.content) + with open("test.xyz", mode="wb") as f: + f.write(self.ctx.create_mlip_inputs.outputs.test_input.content) + with open("valid.xyz", mode="wb") as f: + f.write( + self.ctx.create_mlip_inputs.outputs.validation_input.content + ) + with computer.get_transport() as transport: + from pathlib import Path + + transport.putfile( + Path("train.xyz").absolute(), f"{work_dir}/train.xyz" + ) + transport.putfile( + Path("test.xyz").absolute(), f"{work_dir}/test.xyz" + ) + transport.putfile( + Path("valid.xyz").absolute(), f"{work_dir}/valid.xyz" + ) + + import yaml + from aiida_mlip.data.config import JanusConfigfile + + from aiida_chemshell.utils import generate_default_mlip_fine_tune_config + + config_dict = generate_default_mlip_fine_tune_config() + config_dict["train_file"] = f"{work_dir}/train.xyz" + config_dict["test_file"] = f"{work_dir}/test.xyz" + config_dict["valid_file"] = f"{work_dir}/valid.xyz" + config_dict["name"] = "ChemShell_Workflow_Test" + + with open("mlip_config.yml", "w+") as f: + yaml.dump(config_dict, f, default_flow_style=False) + + mlip_inputs = { + "mlip_config": JanusConfigfile(Path("mlip_config.yml").absolute()), + "code": self.inputs.get("mlip_code", None), + "fine_tune": True, + "foundation_model": self.inputs.get("mlip_model", None), + "metadata": {"options": {"resources": {"num_machines": 1}}}, + } + # Submit the mlip training job future = self.submit(mlip_train_calc, **mlip_inputs) future.label = "MLIP Fine-Tuning." future.description = ( @@ -208,7 +270,17 @@ def train_mlip(self): def result(self): """Extract the final workflow results.""" - self.out("optimised_structure", self.ctx.optimise.outputs.optimised_structure) + if "mlip_model" in self.inputs: + self.out("optimised_structure", self.ctx.optimise.outputs.trajectory_path) + self.out("fine_tuned_model", self.ctx.mlip_training.outputs.model) + self.out( + "fine_tuned_model_compiled", + self.ctx.mlip_training.outputs.compiled_model, + ) + else: + 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( diff --git a/tests/test_workflows.py b/tests/test_workflows.py index e52032d..ab8015c 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -26,7 +26,7 @@ def test_geometry_optimisation_workflow(chemsh_code, get_test_data_file): "basis", "" ) == GeometryOptimisationWorkChain.get_basis_set_label("fast") - assert abs(results.get("final_energy") - -75.585959742867) < 1e-10, ( + assert abs(results.get("final_energy") - -75.585959742867) < 1e-9, ( "Incorrect final energy for geometry optimisation workflow." ) @@ -46,24 +46,20 @@ def test_optimisation_workflow_mlip_training( chemsh_code, get_test_data_file, janus_code ): """Test a geometry optimisation workflow with vibrational analysis.""" - from aiida_mlip.data.config import JanusConfigfile from aiida_mlip.helpers.help_load import load_model inputs = { "chemsh": { "code": chemsh_code(), - "structure": get_test_data_file(), - "qm_parameters": {"theory": "PySCF", "method": "HF"}, - }, - "mlip": { - "mlip_config": JanusConfigfile( - "/opt/aiida-chemshell/.local/janus_config.yml" - ), - "code": janus_code, - "fine_tune": True, - "foundation_model": load_model(None, "mace_mp"), - "metadata": {"options": {"resources": {"num_machines": 1}}}, + "structure": get_test_data_file("butanol.cjson"), + "qm_parameters": { + "theory": "PySCF", + "method": "hf", + "functional": "blyp", + }, }, + "mlip_model": load_model(None, "mace_mp"), + "mlip_code": janus_code, "basis_quality": "fast", "vibrational_analysis": False, } @@ -73,10 +69,30 @@ def test_optimisation_workflow_mlip_training( assert len(node.called) > 0, "WorkChain did not launch any subprocesses" - assert node.called[0].inputs.qm_parameters.get( - "basis", "" - ) == GeometryOptimisationWorkChain.get_basis_set_label("fast") - - assert abs(results.get("final_energy") - -75.585959742867) < 1e-10, ( - "Incorrect final energy for geometry optimisation workflow." - ) + # assert node.called[0].inputs.qm_parameters.get( + # "basis", "" + # ) == GeometryOptimisationWorkChain.get_basis_set_label("fast") + # + # assert abs(results.get("final_energy") - -75.585959742867) < 1e-9, ( + # "Incorrect final energy for geometry optimisation workflow." + # ) + + subs = node.called + for sub in subs: + # if "Generate MLIP training" in sub.label: + # print(sub.outputs.training_input.content) + if "MLIP Fine-Tuning" in sub.label: + print(sub.outputs.retrieved.list_object_names()) + print("ERROR") + print(sub.outputs.retrieved.get_object_content("_scheduler-stderr.txt")) + print("OUTPUT") + print(sub.outputs.retrieved.get_object_content("_scheduler-stdout.txt")) + print("AIIDA OUTPUT") + print(sub.outputs.retrieved.get_object_content("aiida-stdout.txt")) + print("results") + print(sub.outputs.retrieved.list_object_names(path="results")) + print("LOG") + # print(sub.outputs.retrieved.get_object_content("logs/test_run-123.log")) + + assert sub.is_finished_ok, f"Node '{sub.label}' failed to finish correctly." + # assert False From f032f3949aa65fc2600cc5b9b7539e7b17cf9527 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 17 Jun 2026 16:50:20 +0100 Subject: [PATCH 16/35] Create a workchain for extracting and calculating isolated atomic energies from a given structure input --- src/aiida_chemshell/calculations/base.py | 16 +- src/aiida_chemshell/parsers/base.py | 2 +- src/aiida_chemshell/units.py | 17 + .../workflows/isolated_atoms.py | 409 ++++++++++++++++++ 4 files changed, 440 insertions(+), 4 deletions(-) create mode 100644 src/aiida_chemshell/units.py create mode 100644 src/aiida_chemshell/workflows/isolated_atoms.py diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index ad22eca..6602798 100644 --- a/src/aiida_chemshell/calculations/base.py +++ b/src/aiida_chemshell/calculations/base.py @@ -5,6 +5,7 @@ from aiida.engine import CalcJob, CalcJobProcessSpec, PortNamespace from aiida.orm import ArrayData, Dict, Float, SinglefileData, StructureData +from aiida_chemshell.units import UnitsConverter from aiida_chemshell.utils import ChemShellMMTheory, ChemShellQMTheory @@ -742,10 +743,19 @@ def chemsh_script_generator(self) -> str: script = "from chemsh import Fragment\n" if isinstance(self.inputs.structure, SinglefileData): - fname = self.inputs.structure.filename + script += ( + f"structure = Fragment(coords='{self.inputs.structure.filename:s}')\n" + ) else: - fname = ChemShellCalculation.FILE_TMP_STRUCTURE - script += f"structure = Fragment(coords='{fname:s}')\n" + 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 diff --git a/src/aiida_chemshell/parsers/base.py b/src/aiida_chemshell/parsers/base.py index 3001c5b..81ce002 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -119,7 +119,7 @@ def parse(self, **kwargs): input_fname = self.node.inputs.structure.filename descrip += f" ({input_fname})" # Store the optimised structure file - with self.retrieved.open(ChemShellCalculation.FILE_DLFIND, "r") as f: + with self.retrieved.open(ChemShellCalculation.FILE_DLFIND, "rb") as f: self.out( "optimised_structure", SinglefileData( diff --git a/src/aiida_chemshell/units.py b/src/aiida_chemshell/units.py new file mode 100644 index 0000000..fe298e7 --- /dev/null +++ b/src/aiida_chemshell/units.py @@ -0,0 +1,17 @@ +"""Unit conversion utilities.""" + + +class UnitsConverter: + """Unit conversion utility.""" + + ANGSTROM = 0.529177 + + @classmethod + def angstrom_to_bohr(cls, val: float) -> float: + """Convert Angstrom to Bohr (a.u.).""" + return val * 1.8897259886 + + @classmethod + def bohr_to_angstrom(cls, val: float) -> float: + """Convert bohr (a.u.) to Angstrom.""" + return val * UnitsConverter.ANGSTROM diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py new file mode 100644 index 0000000..1b53ac7 --- /dev/null +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -0,0 +1,409 @@ +"""Workflows for isolating atomic species and calculating SP energies.""" + +from aiida.engine import ToContext, WorkChain, calcfunction +from aiida.orm import Dict, StructureData +from aiida.plugins.factories import CalculationFactory + +ChemShellCalculation = CalculationFactory("chemshell") + + +class IsolatedAtomicEnergiesWorkChain(WorkChain): + """AiiDA workflow for extracting isolated atomic energies from a given structure.""" + + @classmethod + def define(cls, spec) -> None: + """Define the AiiDA process specification for the WorkChain.""" + super().define(spec) + + spec.expose_inputs( + ChemShellCalculation, include=("structure", "qm_parameters", "code") + ) + + spec.output( + "atom_energies", + valid_type=Dict, + required=False, + help=( + "The individual isolated atomic energies for every unique atom type in " + "the given system." + ), + ) + + spec.outline(cls.determine_unique_atom_types, cls.atom_energies, cls.result) + + return + + def determine_unique_atom_types(self) -> None: + """Determine all unique atom types within the given structure.""" + self.unique_atoms = {} + if isinstance(self.inputs.structure, StructureData): + self._atom_types_from_structuredata(self.inputs.structure) + # elif isinstance(self.inputs.structure, SinglefileData): + else: + self._atom_types_from_file() + return + + def atom_energies(self): + """Run ChemShell single point calculations for each atom type.""" + calculations = {} + for atom_symbol in self.unique_atoms.keys(): + structure = StructureData() + structure.append_atom(position=(0.0, 0.0, 0.0), symbols=atom_symbol) + structure.set_pbc((False, False, False)) + builder = self.inputs.code.get_builder() + builder.structure = structure + builder.qm_parameters = self.inputs.qm_parameters + future = self.submit(builder) + calculations[atom_symbol] = future + return ToContext(**calculations) + + def result(self) -> None: + """Collect the results into a dictionary.""" + results_dict = collate_energies( + list(self.unique_atoms.keys()), + [self.ctx.get(atom).outputs.energy for atom in self.unique_atoms.keys()], + ) + self.out("atom_energies", results_dict) + return + + def _atom_types_from_structuredata(self, structure: StructureData) -> None: + """Determine the unique atom types from a StructureData object.""" + for site in structure.sites: + self.unique_atoms[site.kind_name] = atom_symbol_to_z(site.kind_name) + return + + def _atom_types_from_file(self) -> None: + """Determine the unique atom types from a SinglefileData object.""" + if self.inputs.structure.filename[-4:] == ".xyz": + self._atom_types_from_xyz() + elif self.inputs.structure.filename[-6:] == ".cjson": + self._atom_types_from_cjson() + elif self.inputs.structure.filename[-4:] == ".pun": + self._atom_types_from_pun() + return + + def _atom_types_from_xyz(self) -> None: + """Determine the unique atom types from an xyz structure file.""" + structure = StructureData() + structure._parse_xyz(self.inputs.structure.content.decode("utf-8")) + self._atom_types_from_structuredata(structure) + return + + def _atom_types_from_cjson(self) -> None: + """Determine the unique atom types from a cjson structure file.""" + import json + + data = json.loads(self.inputs.structure.content).get("atoms").get("elements") + atom_symbols: list[str] = data.get("symbol", []) + atom_numbers: list[int] = data.get("number", []) + if len(atom_symbols) == 0: + self.unique_atoms = { + atom_z_to_symbol(number): number for number in atom_numbers + } + elif len(atom_numbers) == 0: + self.unique_atoms = { + symbol: atom_symbol_to_z(symbol) for symbol in atom_symbols + } + else: + self.unique_atoms = dict(zip(atom_symbols, atom_numbers, strict=False)) + return + + def _atom_types_from_pun(self) -> None: + """Determine the unique atom types from a .pun structure file.""" + return + + +@calcfunction +def collate_energies(atoms, energies) -> Dict: + """Collate a series of isolated atom energies into a dictionary output.""" + if len(atoms) != len(energies): + raise ValueError( + f"Mismatched lengths: Got {len(atoms)} atoms and {len(energies)} energies." + ) + return Dict(dict(zip(atoms, energies, strict=False))) + + +def atom_symbol_to_z(at_symb: str) -> int: + """ + Get the corresponding nuclear charge for a given atom. + + Parameters + ---------- + at_symb : str + The atom label. + + Returns + ------- + Z : int + The nuclear charge for the given atom. + """ + atom_dict = { + "H": 1, + "HE": 2, + "LI": 3, + "BE": 4, + "B": 5, + "C": 6, + "N": 7, + "O": 8, + "F": 9, + "NE": 10, + "NA": 11, + "MG": 12, + "AL": 13, + "SI": 14, + "P": 15, + "S": 16, + "CL": 17, + "AR": 18, + "K": 19, + "CA": 20, + "SC": 21, + "TI": 22, + "V": 23, + "CR": 24, + "MN": 25, + "FE": 26, + "CO": 27, + "NI": 28, + "CU": 29, + "ZN": 30, + "GA": 31, + "GE": 32, + "AS": 33, + "SE": 34, + "BR": 35, + "KR": 36, + "RB": 37, + "SR": 38, + "Y": 39, + "ZR": 40, + "NB": 41, + "MO": 42, + "TC": 43, + "RU": 44, + "RH": 45, + "PD": 46, + "AG": 47, + "CD": 48, + "IN": 49, + "SN": 50, + "SB": 51, + "TE": 52, + "I": 53, + "XE": 54, + "CS": 55, + "BA": 56, + "LA": 57, + "CE": 58, + "PR": 59, + "ND": 60, + "PM": 61, + "SM": 62, + "EU": 63, + "GD": 64, + "TB": 65, + "DY": 66, + "HO": 67, + "ER": 68, + "TM": 69, + "YB": 70, + "LU": 71, + "HF": 72, + "TA": 73, + "W": 74, + "RE": 75, + "OS": 76, + "IR": 77, + "PT": 78, + "AU": 79, + "HG": 80, + "TL": 81, + "PB": 82, + "BI": 83, + "PO": 84, + "AT": 85, + "RN": 86, + "FR": 87, + "RA": 88, + "AC": 89, + "TH": 90, + "PA": 91, + "U": 92, + "NP": 93, + "PU": 94, + "AM": 95, + "CM": 96, + "BK": 97, + "CF": 98, + "ES": 99, + "FM": 100, + "MD": 101, + "NO": 102, + "LR": 103, + "RF": 104, + "DB": 105, + "DG": 106, + "BH": 107, + "HS": 108, + "MT": 109, + "DS": 110, + "RG": 111, + "UUB": 112, + "UUT": 113, + "UUQ": 114, + "UUP": 115, + "UUH": 116, + "UUS": 117, + "UUO": 118, + "GH": 0, + "X": -1, + } # X denotes dummy atom in zmatrix + try: + z = atom_dict[at_symb.upper()] + except KeyError: + raise KeyError(f"Invalid atomic symbol {at_symb}.") from None + return z + + +def atom_z_to_symbol(z: int) -> str: + """ + Get the corresponding atomic symbol for a given nuclear charge. + + Parameters + ---------- + Z : int + The nuclear charge for the given atom. + + Returns + ------- + at_symb : str + The atom label. + """ + atom_dict = { + 1: "H", + 2: "HE", + 3: "LI", + 4: "BE", + 5: "B", + 6: "C", + 7: "N", + 8: "O", + 9: "F", + 10: "NE", + 11: "NA", + 12: "MG", + 13: "AL", + 14: "SI", + 15: "P", + 16: "S", + 17: "CL", + 18: "AR", + 19: "K", + 20: "CA", + 21: "SC", + 22: "TI", + 23: "V", + 24: "CR", + 25: "MN", + 26: "FE", + 27: "CO", + 28: "NI", + 29: "CU", + 30: "ZN", + 31: "GA", + 32: "GE", + 33: "AS", + 34: "SE", + 35: "BR", + 36: "KR", + 37: "RB", + 38: "SR", + 39: "Y", + 40: "ZR", + 41: "NB", + 42: "MO", + 43: "TC", + 44: "RU", + 45: "RH", + 46: "PD", + 47: "AG", + 48: "CD", + 49: "IN", + 50: "SN", + 51: "SB", + 52: "TE", + 53: "I", + 54: "XE", + 55: "CS", + 56: "BA", + 57: "LA", + 58: "CE", + 59: "PR", + 60: "ND", + 61: "PM", + 62: "SM", + 63: "EU", + 64: "GD", + 65: "TB", + 66: "DY", + 67: "HO", + 68: "ER", + 69: "TM", + 70: "YB", + 71: "LU", + 72: "HF", + 73: "TA", + 74: "W", + 75: "RE", + 76: "OS", + 77: "IR", + 78: "PT", + 79: "AU", + 80: "HG", + 81: "TL", + 82: "PB", + 83: "BI", + 84: "PO", + 85: "AT", + 86: "RN", + 87: "FR", + 88: "RA", + 89: "AC", + 90: "TH", + 91: "PA", + 92: "U", + 93: "NP", + 94: "PU", + 95: "AM", + 96: "CM", + 97: "BK", + 98: "CF", + 99: "ES", + 100: "FM", + 101: "MD", + 102: "NO", + 103: "LR", + 104: "RF", + 105: "DB", + 106: "DG", + 107: "BH", + 108: "HS", + 109: "MT", + 110: "DS", + 111: "RG", + 112: "UUB", + 113: "UUT", + 114: "UUQ", + 115: "UUP", + 116: "UUH", + 117: "UUS", + 118: "UUO", + 0: "GH", + -1: "X", + } # X denotes dummy atom in zmatrix + try: + atm_symb = atom_dict[z] + except KeyError: + raise KeyError(f"Invalid atomic number {z}.") from None + return atm_symb From c33f0c58ea6a164080223594d34e330577eade2b Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Thu, 18 Jun 2026 09:04:06 +0100 Subject: [PATCH 17/35] Simplify isolated atoms workflow and extract re-usable parts --- src/aiida_chemshell/calculations/utils.py | 14 + src/aiida_chemshell/periodic_table.py | 291 ++++++++++++++++ src/aiida_chemshell/units.py | 11 + .../workflows/isolated_atoms.py | 328 +----------------- 4 files changed, 330 insertions(+), 314 deletions(-) create mode 100644 src/aiida_chemshell/calculations/utils.py create mode 100644 src/aiida_chemshell/periodic_table.py diff --git a/src/aiida_chemshell/calculations/utils.py b/src/aiida_chemshell/calculations/utils.py new file mode 100644 index 0000000..6ff05aa --- /dev/null +++ b/src/aiida_chemshell/calculations/utils.py @@ -0,0 +1,14 @@ +"""A collection of smaller utility AiiDA CalcFunctions.""" + +from aiida.engine import calcfunction +from aiida.orm import Dict + + +@calcfunction +def create_dictionary(atoms, energies) -> Dict: + """Collate a series of isolated atom energies into a dictionary output.""" + if len(atoms) != len(energies): + raise ValueError( + f"Mismatched lengths: Got {len(atoms)} atoms and {len(energies)} energies." + ) + return Dict(dict(zip(atoms, energies, strict=False))) diff --git a/src/aiida_chemshell/periodic_table.py b/src/aiida_chemshell/periodic_table.py new file mode 100644 index 0000000..246a792 --- /dev/null +++ b/src/aiida_chemshell/periodic_table.py @@ -0,0 +1,291 @@ +"""Utilities based on determining atomistic properties.""" + + +class PeriodicTable: + """Defines methods for determining atomistic properties.""" + + @classmethod + def atom_symbol_to_z(cls, at_symb: str) -> int: + """ + Get the corresponding nuclear charge for a given atom. + + Parameters + ---------- + at_symb : str + The atom label. + + Returns + ------- + Z : int + The nuclear charge for the given atom. + """ + atom_dict = { + "H": 1, + "HE": 2, + "LI": 3, + "BE": 4, + "B": 5, + "C": 6, + "N": 7, + "O": 8, + "F": 9, + "NE": 10, + "NA": 11, + "MG": 12, + "AL": 13, + "SI": 14, + "P": 15, + "S": 16, + "CL": 17, + "AR": 18, + "K": 19, + "CA": 20, + "SC": 21, + "TI": 22, + "V": 23, + "CR": 24, + "MN": 25, + "FE": 26, + "CO": 27, + "NI": 28, + "CU": 29, + "ZN": 30, + "GA": 31, + "GE": 32, + "AS": 33, + "SE": 34, + "BR": 35, + "KR": 36, + "RB": 37, + "SR": 38, + "Y": 39, + "ZR": 40, + "NB": 41, + "MO": 42, + "TC": 43, + "RU": 44, + "RH": 45, + "PD": 46, + "AG": 47, + "CD": 48, + "IN": 49, + "SN": 50, + "SB": 51, + "TE": 52, + "I": 53, + "XE": 54, + "CS": 55, + "BA": 56, + "LA": 57, + "CE": 58, + "PR": 59, + "ND": 60, + "PM": 61, + "SM": 62, + "EU": 63, + "GD": 64, + "TB": 65, + "DY": 66, + "HO": 67, + "ER": 68, + "TM": 69, + "YB": 70, + "LU": 71, + "HF": 72, + "TA": 73, + "W": 74, + "RE": 75, + "OS": 76, + "IR": 77, + "PT": 78, + "AU": 79, + "HG": 80, + "TL": 81, + "PB": 82, + "BI": 83, + "PO": 84, + "AT": 85, + "RN": 86, + "FR": 87, + "RA": 88, + "AC": 89, + "TH": 90, + "PA": 91, + "U": 92, + "NP": 93, + "PU": 94, + "AM": 95, + "CM": 96, + "BK": 97, + "CF": 98, + "ES": 99, + "FM": 100, + "MD": 101, + "NO": 102, + "LR": 103, + "RF": 104, + "DB": 105, + "DG": 106, + "BH": 107, + "HS": 108, + "MT": 109, + "DS": 110, + "RG": 111, + "UUB": 112, + "UUT": 113, + "UUQ": 114, + "UUP": 115, + "UUH": 116, + "UUS": 117, + "UUO": 118, + "GH": 0, + "X": -1, + } # X denotes dummy atom in zmatrix + try: + z = atom_dict[at_symb.upper()] + except KeyError: + raise KeyError(f"Invalid atomic symbol {at_symb}.") from None + return z + + @classmethod + def atom_z_to_symbol(cls, z: int) -> str: + """ + Get the corresponding atomic symbol for a given nuclear charge. + + Parameters + ---------- + Z : int + The nuclear charge for the given atom. + + Returns + ------- + at_symb : str + The atom label. + """ + atom_dict = { + 1: "H", + 2: "HE", + 3: "LI", + 4: "BE", + 5: "B", + 6: "C", + 7: "N", + 8: "O", + 9: "F", + 10: "NE", + 11: "NA", + 12: "MG", + 13: "AL", + 14: "SI", + 15: "P", + 16: "S", + 17: "CL", + 18: "AR", + 19: "K", + 20: "CA", + 21: "SC", + 22: "TI", + 23: "V", + 24: "CR", + 25: "MN", + 26: "FE", + 27: "CO", + 28: "NI", + 29: "CU", + 30: "ZN", + 31: "GA", + 32: "GE", + 33: "AS", + 34: "SE", + 35: "BR", + 36: "KR", + 37: "RB", + 38: "SR", + 39: "Y", + 40: "ZR", + 41: "NB", + 42: "MO", + 43: "TC", + 44: "RU", + 45: "RH", + 46: "PD", + 47: "AG", + 48: "CD", + 49: "IN", + 50: "SN", + 51: "SB", + 52: "TE", + 53: "I", + 54: "XE", + 55: "CS", + 56: "BA", + 57: "LA", + 58: "CE", + 59: "PR", + 60: "ND", + 61: "PM", + 62: "SM", + 63: "EU", + 64: "GD", + 65: "TB", + 66: "DY", + 67: "HO", + 68: "ER", + 69: "TM", + 70: "YB", + 71: "LU", + 72: "HF", + 73: "TA", + 74: "W", + 75: "RE", + 76: "OS", + 77: "IR", + 78: "PT", + 79: "AU", + 80: "HG", + 81: "TL", + 82: "PB", + 83: "BI", + 84: "PO", + 85: "AT", + 86: "RN", + 87: "FR", + 88: "RA", + 89: "AC", + 90: "TH", + 91: "PA", + 92: "U", + 93: "NP", + 94: "PU", + 95: "AM", + 96: "CM", + 97: "BK", + 98: "CF", + 99: "ES", + 100: "FM", + 101: "MD", + 102: "NO", + 103: "LR", + 104: "RF", + 105: "DB", + 106: "DG", + 107: "BH", + 108: "HS", + 109: "MT", + 110: "DS", + 111: "RG", + 112: "UUB", + 113: "UUT", + 114: "UUQ", + 115: "UUP", + 116: "UUH", + 117: "UUS", + 118: "UUO", + 0: "GH", + -1: "X", + } # X denotes dummy atom in zmatrix + try: + atm_symb = atom_dict[z] + except KeyError: + raise KeyError(f"Invalid atomic number {z}.") from None + return atm_symb diff --git a/src/aiida_chemshell/units.py b/src/aiida_chemshell/units.py index fe298e7..080f19e 100644 --- a/src/aiida_chemshell/units.py +++ b/src/aiida_chemshell/units.py @@ -5,6 +5,7 @@ class UnitsConverter: """Unit conversion utility.""" ANGSTROM = 0.529177 + EV = 27.2114 @classmethod def angstrom_to_bohr(cls, val: float) -> float: @@ -15,3 +16,13 @@ def angstrom_to_bohr(cls, val: float) -> float: def bohr_to_angstrom(cls, val: float) -> float: """Convert bohr (a.u.) to Angstrom.""" return val * UnitsConverter.ANGSTROM + + @classmethod + def hartree_to_ev(cls, val: float) -> float: + """Convert Hartree (a.u.) to eV.""" + return val * UnitsConverter.EV + + @classmethod + def ev_to_hartree(cls, val: float) -> float: + """Convert eV to Hartree (a.u.).""" + return val / UnitsConverter.EV diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py index 1b53ac7..91fa8b4 100644 --- a/src/aiida_chemshell/workflows/isolated_atoms.py +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -1,9 +1,12 @@ """Workflows for isolating atomic species and calculating SP energies.""" -from aiida.engine import ToContext, WorkChain, calcfunction +from aiida.engine import ToContext, WorkChain from aiida.orm import Dict, StructureData from aiida.plugins.factories import CalculationFactory +from aiida_chemshell.calculations.utils import create_dictionary +from aiida_chemshell.periodic_table import PeriodicTable + ChemShellCalculation = CalculationFactory("chemshell") @@ -35,10 +38,9 @@ def define(cls, spec) -> None: def determine_unique_atom_types(self) -> None: """Determine all unique atom types within the given structure.""" - self.unique_atoms = {} + self.unique_atoms = [] if isinstance(self.inputs.structure, StructureData): self._atom_types_from_structuredata(self.inputs.structure) - # elif isinstance(self.inputs.structure, SinglefileData): else: self._atom_types_from_file() return @@ -46,7 +48,7 @@ def determine_unique_atom_types(self) -> None: def atom_energies(self): """Run ChemShell single point calculations for each atom type.""" calculations = {} - for atom_symbol in self.unique_atoms.keys(): + for atom_symbol in self.unique_atoms: structure = StructureData() structure.append_atom(position=(0.0, 0.0, 0.0), symbols=atom_symbol) structure.set_pbc((False, False, False)) @@ -59,18 +61,16 @@ def atom_energies(self): def result(self) -> None: """Collect the results into a dictionary.""" - results_dict = collate_energies( - list(self.unique_atoms.keys()), - [self.ctx.get(atom).outputs.energy for atom in self.unique_atoms.keys()], + results_dict = create_dictionary( + list(self.unique_atoms), + [self.ctx.get(atom).outputs.energy for atom in self.unique_atoms], ) self.out("atom_energies", results_dict) return def _atom_types_from_structuredata(self, structure: StructureData) -> None: """Determine the unique atom types from a StructureData object.""" - for site in structure.sites: - self.unique_atoms[site.kind_name] = atom_symbol_to_z(site.kind_name) - return + self.unique_atoms = [site.kind_name for site in structure.sites] def _atom_types_from_file(self) -> None: """Determine the unique atom types from a SinglefileData object.""" @@ -97,313 +97,13 @@ def _atom_types_from_cjson(self) -> None: atom_symbols: list[str] = data.get("symbol", []) atom_numbers: list[int] = data.get("number", []) if len(atom_symbols) == 0: - self.unique_atoms = { - atom_z_to_symbol(number): number for number in atom_numbers - } - elif len(atom_numbers) == 0: - self.unique_atoms = { - symbol: atom_symbol_to_z(symbol) for symbol in atom_symbols - } + self.unique_atoms = [ + PeriodicTable.atom_z_to_symbol(number) for number in atom_numbers + ] else: - self.unique_atoms = dict(zip(atom_symbols, atom_numbers, strict=False)) + self.unique_atoms = atom_symbols return def _atom_types_from_pun(self) -> None: """Determine the unique atom types from a .pun structure file.""" return - - -@calcfunction -def collate_energies(atoms, energies) -> Dict: - """Collate a series of isolated atom energies into a dictionary output.""" - if len(atoms) != len(energies): - raise ValueError( - f"Mismatched lengths: Got {len(atoms)} atoms and {len(energies)} energies." - ) - return Dict(dict(zip(atoms, energies, strict=False))) - - -def atom_symbol_to_z(at_symb: str) -> int: - """ - Get the corresponding nuclear charge for a given atom. - - Parameters - ---------- - at_symb : str - The atom label. - - Returns - ------- - Z : int - The nuclear charge for the given atom. - """ - atom_dict = { - "H": 1, - "HE": 2, - "LI": 3, - "BE": 4, - "B": 5, - "C": 6, - "N": 7, - "O": 8, - "F": 9, - "NE": 10, - "NA": 11, - "MG": 12, - "AL": 13, - "SI": 14, - "P": 15, - "S": 16, - "CL": 17, - "AR": 18, - "K": 19, - "CA": 20, - "SC": 21, - "TI": 22, - "V": 23, - "CR": 24, - "MN": 25, - "FE": 26, - "CO": 27, - "NI": 28, - "CU": 29, - "ZN": 30, - "GA": 31, - "GE": 32, - "AS": 33, - "SE": 34, - "BR": 35, - "KR": 36, - "RB": 37, - "SR": 38, - "Y": 39, - "ZR": 40, - "NB": 41, - "MO": 42, - "TC": 43, - "RU": 44, - "RH": 45, - "PD": 46, - "AG": 47, - "CD": 48, - "IN": 49, - "SN": 50, - "SB": 51, - "TE": 52, - "I": 53, - "XE": 54, - "CS": 55, - "BA": 56, - "LA": 57, - "CE": 58, - "PR": 59, - "ND": 60, - "PM": 61, - "SM": 62, - "EU": 63, - "GD": 64, - "TB": 65, - "DY": 66, - "HO": 67, - "ER": 68, - "TM": 69, - "YB": 70, - "LU": 71, - "HF": 72, - "TA": 73, - "W": 74, - "RE": 75, - "OS": 76, - "IR": 77, - "PT": 78, - "AU": 79, - "HG": 80, - "TL": 81, - "PB": 82, - "BI": 83, - "PO": 84, - "AT": 85, - "RN": 86, - "FR": 87, - "RA": 88, - "AC": 89, - "TH": 90, - "PA": 91, - "U": 92, - "NP": 93, - "PU": 94, - "AM": 95, - "CM": 96, - "BK": 97, - "CF": 98, - "ES": 99, - "FM": 100, - "MD": 101, - "NO": 102, - "LR": 103, - "RF": 104, - "DB": 105, - "DG": 106, - "BH": 107, - "HS": 108, - "MT": 109, - "DS": 110, - "RG": 111, - "UUB": 112, - "UUT": 113, - "UUQ": 114, - "UUP": 115, - "UUH": 116, - "UUS": 117, - "UUO": 118, - "GH": 0, - "X": -1, - } # X denotes dummy atom in zmatrix - try: - z = atom_dict[at_symb.upper()] - except KeyError: - raise KeyError(f"Invalid atomic symbol {at_symb}.") from None - return z - - -def atom_z_to_symbol(z: int) -> str: - """ - Get the corresponding atomic symbol for a given nuclear charge. - - Parameters - ---------- - Z : int - The nuclear charge for the given atom. - - Returns - ------- - at_symb : str - The atom label. - """ - atom_dict = { - 1: "H", - 2: "HE", - 3: "LI", - 4: "BE", - 5: "B", - 6: "C", - 7: "N", - 8: "O", - 9: "F", - 10: "NE", - 11: "NA", - 12: "MG", - 13: "AL", - 14: "SI", - 15: "P", - 16: "S", - 17: "CL", - 18: "AR", - 19: "K", - 20: "CA", - 21: "SC", - 22: "TI", - 23: "V", - 24: "CR", - 25: "MN", - 26: "FE", - 27: "CO", - 28: "NI", - 29: "CU", - 30: "ZN", - 31: "GA", - 32: "GE", - 33: "AS", - 34: "SE", - 35: "BR", - 36: "KR", - 37: "RB", - 38: "SR", - 39: "Y", - 40: "ZR", - 41: "NB", - 42: "MO", - 43: "TC", - 44: "RU", - 45: "RH", - 46: "PD", - 47: "AG", - 48: "CD", - 49: "IN", - 50: "SN", - 51: "SB", - 52: "TE", - 53: "I", - 54: "XE", - 55: "CS", - 56: "BA", - 57: "LA", - 58: "CE", - 59: "PR", - 60: "ND", - 61: "PM", - 62: "SM", - 63: "EU", - 64: "GD", - 65: "TB", - 66: "DY", - 67: "HO", - 68: "ER", - 69: "TM", - 70: "YB", - 71: "LU", - 72: "HF", - 73: "TA", - 74: "W", - 75: "RE", - 76: "OS", - 77: "IR", - 78: "PT", - 79: "AU", - 80: "HG", - 81: "TL", - 82: "PB", - 83: "BI", - 84: "PO", - 85: "AT", - 86: "RN", - 87: "FR", - 88: "RA", - 89: "AC", - 90: "TH", - 91: "PA", - 92: "U", - 93: "NP", - 94: "PU", - 95: "AM", - 96: "CM", - 97: "BK", - 98: "CF", - 99: "ES", - 100: "FM", - 101: "MD", - 102: "NO", - 103: "LR", - 104: "RF", - 105: "DB", - 106: "DG", - 107: "BH", - 108: "HS", - 109: "MT", - 110: "DS", - 111: "RG", - 112: "UUB", - 113: "UUT", - 114: "UUQ", - 115: "UUP", - 116: "UUH", - 117: "UUS", - 118: "UUO", - 0: "GH", - -1: "X", - } # X denotes dummy atom in zmatrix - try: - atm_symb = atom_dict[z] - except KeyError: - raise KeyError(f"Invalid atomic number {z}.") from None - return atm_symb From ba2354bfbbf8807fdc7ebefc23dcd597e6caf3a3 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Thu, 18 Jun 2026 14:20:39 +0100 Subject: [PATCH 18/35] Fully integrated mlip options in ChemShell optimisation workflow --- .../calculations/file_conversion.py | 34 +++++++++++++++++-- src/aiida_chemshell/utils.py | 1 - src/aiida_chemshell/workflows/optimisation.py | 15 ++++++++ 3 files changed, 46 insertions(+), 4 deletions(-) diff --git a/src/aiida_chemshell/calculations/file_conversion.py b/src/aiida_chemshell/calculations/file_conversion.py index ee75b7d..cf50556 100644 --- a/src/aiida_chemshell/calculations/file_conversion.py +++ b/src/aiida_chemshell/calculations/file_conversion.py @@ -3,7 +3,9 @@ from aiida.common import CalcInfo, CodeInfo from aiida.common.folders import Folder from aiida.engine import CalcJob, CalcJobProcessSpec -from aiida.orm import ArrayData, SinglefileData, Str +from aiida.orm import ArrayData, Dict, SinglefileData, Str + +from aiida_chemshell.units import UnitsConverter class CreateJanusTrainingInputsCalcJob(CalcJob): @@ -29,7 +31,13 @@ def define(cls, spec: CalcJobProcessSpec) -> None: "energies", valid_type=ArrayData, required=True, - help="The calculated dft energy of each step in the data series.", + help="The calculated dft energy of each step in the data series in a.u.", + ) + spec.input( + "atom_energies", + valid_type=Dict, + required=True, + help="The isolated atomic energies for the training set.", ) spec.input( @@ -76,6 +84,19 @@ def define(cls, spec: CalcJobProcessSpec) -> None: return + def create_isolated_atom_energy_xyz(self) -> str: + """Create the isolated atomistic energies in the required XYZ format.""" + xyz_str = "" + zero = 0.0 + for atom in self.inputs.atom_energies.keys(): + xyz_str += "1\nProperties=species:S:1:pos:R:3:dft_forces:R:3 " + energy = UnitsConverter.hartree_to_ev(self.inputs.atom_energies[atom]) + xyz_str += f"dft_energy={energy} " + xyz_str += 'config_type=IsolatedAtom pbc="F F F"\n' + xyz_str += f"{atom:8s} {zero:.4f} {zero:.4f} {zero:.4f} " + xyz_str += f"{zero:.4f} {zero:.4f} {zero:.4f}\n" + return xyz_str + def prepare_for_submission(self, folder: Folder) -> CalcInfo: """Perform the python task to split the input trajectories.""" script = self.generate_script() @@ -84,7 +105,10 @@ def prepare_for_submission(self, folder: Folder) -> CalcInfo: with folder.open("energies.txt", "w") as f: for val in self.inputs.energies.get_array("energies"): - f.write(f"{val:.10f}\n") + f.write(f"{UnitsConverter.hartree_to_ev(val):.10f}\n") + + with folder.open("isolated_atoms.xyz", "w") as f: + f.write(self.create_isolated_atom_energy_xyz()) code_info = CodeInfo() code_info.code_uuid = self.inputs.code.uuid @@ -158,4 +182,8 @@ def generate_script(self) -> str: f.write(path[index].strip("\\n") + " ") f.write(" ".join(force[index].split()[1:])) f.write("\\n") +with open("isolated_atoms.xyz", "r") as f: + atoms = f.read() +with open("train.xyz", "a") as f: + f.write(atoms) """ diff --git a/src/aiida_chemshell/utils.py b/src/aiida_chemshell/utils.py index dc50927..26cce95 100644 --- a/src/aiida_chemshell/utils.py +++ b/src/aiida_chemshell/utils.py @@ -123,5 +123,4 @@ def generate_default_mlip_fine_tune_config(): "weight_decay": 1e-8, "eval_interval": 1, # "enable_cueq": True, - "E0s": {1: -13.664311914087541, 6: -1030.2235576550245, 8: -2043.4891174666404}, } diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index f6d7e53..efc7e59 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -6,6 +6,7 @@ from aiida.plugins.factories import CalculationFactory from aiida_chemshell.calculations.base import ChemShellCalculation +from aiida_chemshell.workflows.isolated_atoms import IsolatedAtomicEnergiesWorkChain class GeometryOptimisationWorkChain(WorkChain): @@ -101,6 +102,7 @@ def define(cls, spec) -> None: spec.outline( cls.optimise, cls.energy, + cls.isolated_atom_energies, cls.generate_mlip_training_inputs, cls.train_mlip, cls.result, @@ -178,6 +180,18 @@ def energy(self): return ToContext(energy=future) return None + def isolated_atom_energies(self): + """Calculate isolated atomic energies for all species in the input structure.""" + if "mlip_model" in self.inputs: + inputs = { + "structure": self.inputs.chemsh.structure, + "code": self.inputs.chemsh.code, + "qm_parameters": self.inputs.chemsh.qm_parameters, + } + future = self.submit(IsolatedAtomicEnergiesWorkChain, **inputs) + return ToContext(isolated_atoms=future) + return None + def generate_mlip_training_inputs(self): """Convert the optimisation path files to Janus compatible inputs.""" if "mlip_model" in self.inputs: @@ -185,6 +199,7 @@ def generate_mlip_training_inputs(self): "path": self.ctx.optimise.outputs.trajectory_path, "force": self.ctx.optimise.outputs.trajectory_force, "energies": self.ctx.optimise.outputs.optimisation_path, + "atom_energies": self.ctx.isolated_atoms.outputs.atom_energies, "code": self.inputs.chemsh.code, "metadata": { "options": { From 5114fb4b081322637a901443413bac3df57dc90e Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Mon, 22 Jun 2026 14:18:08 +0100 Subject: [PATCH 19/35] Fix to make sure isolated atoms workchain only does the unique atom types not all atoms --- src/aiida_chemshell/workflows/isolated_atoms.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py index 91fa8b4..65f5c4c 100644 --- a/src/aiida_chemshell/workflows/isolated_atoms.py +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -70,7 +70,11 @@ def result(self) -> None: def _atom_types_from_structuredata(self, structure: StructureData) -> None: """Determine the unique atom types from a StructureData object.""" - self.unique_atoms = [site.kind_name for site in structure.sites] + self.unique_atoms = [] + for site in structure.sites: + if site.kind_name not in self.unique_atoms: + self.unique_atoms.append(site.kind_name) + return def _atom_types_from_file(self) -> None: """Determine the unique atom types from a SinglefileData object.""" From 53e321aa53ec584ce501abc11b9cead13b21874c Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Mon, 22 Jun 2026 14:28:04 +0100 Subject: [PATCH 20/35] Fix for base calculation parser when collecting both the optimisation path and resulting optimised structure --- src/aiida_chemshell/parsers/base.py | 58 ++++++++++++++--------------- 1 file changed, 27 insertions(+), 31 deletions(-) diff --git a/src/aiida_chemshell/parsers/base.py b/src/aiida_chemshell/parsers/base.py index 81ce002..e3b2783 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -73,7 +73,33 @@ def parse(self, **kwargs): ChemShellCalculation.FILE_STDOUT, "r" ) ) - elif self.node.inputs.optimisation_parameters.get("save_path", False): + 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}" + if isinstance(self.node.inputs.structure, SinglefileData): + input_fname = self.node.inputs.structure.filename + descrip += f" ({input_fname})" + # Store the optimised structure file + with self.retrieved.open(ChemShellCalculation.FILE_DLFIND, "rb") as f: + self.out( + "optimised_structure", + SinglefileData( + file=f, + filename=ChemShellCalculation.FILE_DLFIND, + label="ChemShell Punch Structure File", + description=descrip, + ), + ) + self.parse_optimisation_path( + self.retrieved.get_object_content( + ChemShellCalculation.FILE_STDOUT, "r" + ) + ) + else: + return self.exit_codes.ERROR_MISSING_OPTIMISED_STRUCTURE_FILE + + if self.node.inputs.optimisation_parameters.get("save_path", False): if ( ChemShellCalculation.FILE_TRJPTH in self.retrieved.list_object_names() @@ -104,38 +130,8 @@ def parse(self, **kwargs): label="ChemShell optimisation trajectory.", ), ) - self.parse_optimisation_path( - self.retrieved.get_object_content( - ChemShellCalculation.FILE_STDOUT, "r" - ) - ) else: return self.exit_codes.ERROR_MISSING_OPTIMISED_STRUCTURE_FILE - 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}" - if isinstance(self.node.inputs.structure, SinglefileData): - input_fname = self.node.inputs.structure.filename - descrip += f" ({input_fname})" - # Store the optimised structure file - with self.retrieved.open(ChemShellCalculation.FILE_DLFIND, "rb") as f: - self.out( - "optimised_structure", - SinglefileData( - file=f, - filename=ChemShellCalculation.FILE_DLFIND, - label="ChemShell Punch Structure File", - description=descrip, - ), - ) - self.parse_optimisation_path( - self.retrieved.get_object_content( - ChemShellCalculation.FILE_STDOUT, "r" - ) - ) - else: - return self.exit_codes.ERROR_MISSING_OPTIMISED_STRUCTURE_FILE return ExitCode(0) From e543d60355bf1ea0b3c17598db67c1d97fa595e6 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Mon, 22 Jun 2026 14:29:18 +0100 Subject: [PATCH 21/35] Clarify auto-labelling of optimised structure file, it is now a cjson file not a chemshell punch file --- src/aiida_chemshell/parsers/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aiida_chemshell/parsers/base.py b/src/aiida_chemshell/parsers/base.py index e3b2783..8d94bf1 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -87,7 +87,7 @@ def parse(self, **kwargs): SinglefileData( file=f, filename=ChemShellCalculation.FILE_DLFIND, - label="ChemShell Punch Structure File", + label="CJSON Structure File", description=descrip, ), ) From 12966a66eaa6d238fdc67780a9a01b2d8a01024a Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Mon, 22 Jun 2026 15:32:04 +0100 Subject: [PATCH 22/35] Add additional validation for xyz path files when creating Janus traingin calculation inputs --- .../calculations/file_conversion.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/src/aiida_chemshell/calculations/file_conversion.py b/src/aiida_chemshell/calculations/file_conversion.py index cf50556..01edcfe 100644 --- a/src/aiida_chemshell/calculations/file_conversion.py +++ b/src/aiida_chemshell/calculations/file_conversion.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, SinglefileData, Str from aiida_chemshell.units import UnitsConverter @@ -50,12 +50,14 @@ def define(cls, spec: CalcJobProcessSpec) -> None: spec.output( "training_input", valid_type=SinglefileData, + validator=cls.validate_path_file, required=True, help="The main training data set in extended XYZ format.", ) spec.output( "validation_input", valid_type=SinglefileData, + vlidatory=cls.validate_path_file, required=True, help="The validation data set in extended XYZ format.", ) @@ -84,6 +86,22 @@ def define(cls, spec: CalcJobProcessSpec) -> None: return + @classmethod + def validate_path_file( + cls, value: SinglefileData, namepsace: PortNamespace + ) -> str | None: + """Perform validation checks on the input path xyz files.""" + if not isinstance(value, SinglefileData): + return "Input trajectory needs to be a SinglefileData node." + content = value.get_content("r") + lines = content.split("\n") + natoms = int(lines[0]) + if len(lines) % (natoms + 2) != 0: + return "Invalid XYZ trajectory structure detected." + if len(lines) // (natoms + 2) < 5: + return "Not enough individual configurations within input trajectory." + return None + def create_isolated_atom_energy_xyz(self) -> str: """Create the isolated atomistic energies in the required XYZ format.""" xyz_str = "" From c70d528cd095eb463d1016d586163856fd670db3 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 10:29:50 +0100 Subject: [PATCH 23/35] Add error if pyscf used for isolated atom energy workflow --- src/aiida_chemshell/workflows/isolated_atoms.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py index 65f5c4c..25f4246 100644 --- a/src/aiida_chemshell/workflows/isolated_atoms.py +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -48,6 +48,10 @@ def determine_unique_atom_types(self) -> None: def atom_energies(self): """Run ChemShell single point calculations for each atom type.""" calculations = {} + if self.inputs.qm_parameters["theory"] == "PySCF": + raise Exception( + "Isolated atom calculations not supported by PySCF QM backend." + ) for atom_symbol in self.unique_atoms: structure = StructureData() structure.append_atom(position=(0.0, 0.0, 0.0), symbols=atom_symbol) From 579cf75406bb65da47f721a28143c2be18f24f87 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 10:30:15 +0100 Subject: [PATCH 24/35] Make isolated atom workflow call in optimisation workflow use the optimisation steps qm inputs --- src/aiida_chemshell/workflows/optimisation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index efc7e59..7b36825 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -186,7 +186,7 @@ def isolated_atom_energies(self): inputs = { "structure": self.inputs.chemsh.structure, "code": self.inputs.chemsh.code, - "qm_parameters": self.inputs.chemsh.qm_parameters, + "qm_parameters": self.ctx.optimise.inputs.qm_parameters, } future = self.submit(IsolatedAtomicEnergiesWorkChain, **inputs) return ToContext(isolated_atoms=future) From 0298c1db4a4232c6d9dc2acc95c18a8a012854f9 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 10:32:34 +0100 Subject: [PATCH 25/35] Remove old clf-ultra.py workflow file --- src/aiida_chemshell/workflows/clf_ultra.py | 184 --------------------- 1 file changed, 184 deletions(-) delete mode 100644 src/aiida_chemshell/workflows/clf_ultra.py diff --git a/src/aiida_chemshell/workflows/clf_ultra.py b/src/aiida_chemshell/workflows/clf_ultra.py deleted file mode 100644 index 9bb2c27..0000000 --- a/src/aiida_chemshell/workflows/clf_ultra.py +++ /dev/null @@ -1,184 +0,0 @@ -"""CLF-ULTRA ChemShell + Janus workflow.""" - -from aiida.engine import ToContext, WorkChain, calcfunction -from aiida.orm import ( - Code, - Dict, - Float, - FolderData, - List, - SinglefileData, - Str, - StructureData, -) -from aiida.plugins.factories import CalculationFactory - -ChemShellCalculation = CalculationFactory("chemshell") - - -class CLFULTRAOptimisationWorkChain(WorkChain): - """ - CLF-ULTRA Geometry Optimisation WorkChain. - - This workchain will perform a QM or QM/MM based geometry optimisation - on a given structure and generate etxXYZ formatted trajectory files for - every step in the optimisation. It uses the NWChem and DL_POLY backends - and the B3LYP dft functional for the QM calculations. - """ - - @classmethod - def define(cls, spec) -> None: - """Define the AiiDA process specification for the WorkChain.""" - super().define(spec) - spec.input( - "structure", - valid_type=(SinglefileData, StructureData), - required=True, - help="The structure on which to perform the geometry optimisation", - ) - spec.input("chemshell", valid_type=Code, required=True, help="") - spec.input("janus", valid_type=Code, required=True, help="") - - spec.input( - "force_field_file", - valid_type=SinglefileData, - required=False, - help="File defining the MM force field for QM/MM calculation.", - ) - spec.input( - "qm_region", - valid_type=List, - required=False, - help=( - "A list of atom indexes to apply the QM portion of a QM/MM " - "calculation to." - ), - ) - - spec.output("final_qm_energy", valid_type=Float) - # spec.output("mlip_outputs", valid_type=Dict) - spec.output("trajectory_files", valid_type=FolderData) - spec.output("energy_difference", valid_type=Float) - - spec.outline( - cls.optimise, - cls.generate_xyz_files, - cls.extract_final_structure, - cls.mlip_sp, - cls.calculate_energy_difference, - cls.result, - ) - return - - def optimise(self): - """Perform the Geometry Optimisation Task.""" - inputs = { - "qm_parameters": Dict( - { - "theory": "NWChem", # NWChem - "method": "dft", # DFT - "functional": "B3LYP", # B3LYP - "basis": "cc-pvtz", # cc=pvtz - # "d3": True # Add D3 correction - } - ), - "optimisation_parameters": Dict({"save_path": True}), - "structure": self.inputs.structure, - "code": self.inputs.chemshell, - "metadata": { - "options": { - "resources": {"num_mpiprocs_per_machine": 8, "num_machines": 1}, - # "withmpi": True, - } - }, - } - future = self.submit(ChemShellCalculation, **inputs) - return ToContext(optimise=future) - - def generate_xyz_files(self): - """Convert the paths to individual extended XYZ filesn for each step.""" - # inputs = { - # "path": self.ctx.optimise.outputs.trajectory_path, - # "force": self.ctx.optimise.outputs.trajectory_force, - # "code": self.inputs.chemshell, - # "metadata": { - # "options": { - # "resources": {"num_mpiprocs_per_machine": 2, "num_machines": 1}, - # # "withmpi": True, - # } - # }, - # } - - # future = self.submit(SplitTrajectory, **inputs) - # return ToContext(split_trajectory=future) - return - - def extract_final_structure(self): - """Extract the optimised structure from the XYZ trajectory folder.""" - self.ctx.final_structure = extract_final_structure( - self.ctx.split_trajectory.outputs.trajectory_folder - ) - return - - def mlip_sp(self): - """Run a single point energy with Janus-Core on the optimised structure.""" - structure = self.ctx.final_structure - from aiida_mlip.helpers.help_load import load_model - - model = load_model(None, "mace_mp") - inputs = { - "metadata": {"options": {"resources": {"num_machines": 1}}}, - "code": self.inputs.janus, - # "arch": Str(model.architecture), - "struct": structure, - "model": model, - "device": Str("cpu"), - # "calc_kwargs": Dict({"dispersion": True}), - } - mlip_sp = CalculationFactory("mlip.sp") - future = self.submit(mlip_sp, **inputs) - return ToContext(mlip_sp=future) - - def calculate_energy_difference(self): - """Calculate the difference in final energies.""" - self.ctx.energy_difference = calculate_difference( - self.ctx.optimise.outputs.energy, self.ctx.mlip_sp.outputs.results_dict - ) - return - - def result(self) -> None: - """Report the final results.""" - self.out("final_qm_energy", self.ctx.optimise.outputs.energy) - self.out( - "trajectory_files", self.ctx.split_trajectory.outputs.trajectory_folder - ) - # self.out( - # "mlip_outputs", - # self.ctx.mlip_sp.outputs.results_dict - # ) - self.out("energy_difference", self.ctx.energy_difference) - return - - -@calcfunction -def extract_final_structure(folder: FolderData) -> StructureData: - """Extract the final optimised structure from a folder of xyz files.""" - structure_file = folder.list_object_names()[-1] - structure = StructureData() - structure._parse_xyz(folder.get_object_content(structure_file, mode="r") + "\n") - return structure - - -@calcfunction -def calculate_difference(qm_energy: Float, ml_results: Dict) -> Float: - """Calculate the difference between the final energies.""" - ml_energy = ml_results["info"]["mace_mp_energy"] - diff = qm_energy - ml_energy - return Float( - diff, - label="Final Energy Difference", - description=( - "Difference in final energies between DFT calculation and " - "MACE_mp energy calculation using Janus-core." - ), - ) From 29cb1b4a2b06e8c1861fa899f75a46a24b091663 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 10:33:00 +0100 Subject: [PATCH 26/35] Expose new calculations/workflows as aiida plugins --- pyproject.toml | 50 +++++++++++++++++++++++--------------------------- 1 file changed, 23 insertions(+), 27 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d571b86..0bd643d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools",] +requires = ["setuptools"] build-backend = "setuptools.build_meta" [project] @@ -7,20 +7,18 @@ version = "0.1.10" name = "aiida-chemshell" description = "AiiDA workflow plugin for the ChemShell chemical modelling software package" authors = [ - {name = "Dr. Benjamin T. Speake", email = "benjamin.speake@stfc.ac.uk"}, -] -dependencies = [ - "aiida-core >= 2.5", + { name = "Dr. Benjamin T. Speake", email = "benjamin.speake@stfc.ac.uk" }, ] +dependencies = ["aiida-core >= 2.5"] readme = "README.md" -license = {file = "LICENSE"} -keywords = ["aiida", "workflows", "chemshell", ] +license = { file = "LICENSE" } +keywords = ["aiida", "workflows", "chemshell"] classifiers = [ "Intended Audience :: Science/Research", "Programming Language :: Python :: 3", - "Topic :: Scientific/Engineering", + "Topic :: Scientific/Engineering", "Framework :: AiiDA", - "Development Status :: 5 - Production/Stable" + "Development Status :: 5 - Production/Stable", ] requires-python = ">=3.10" @@ -29,6 +27,7 @@ Source = "https://github.com/stfc/aiida-chemshell" [project.entry-points."aiida.calculations"] "chemshell" = "aiida_chemshell.calculations.base:ChemShellCalculation" +"chemshell.file_conversion.mlip_training" = "aiida_chemshell.calculations.file_conversion:CreateJanusTrainingInputsCalcJob" [project.entry-points."aiida.parsers"] "chemshell" = "aiida_chemshell.parsers.base:ChemShellParser" @@ -36,30 +35,26 @@ Source = "https://github.com/stfc/aiida-chemshell" [project.entry-points."aiida.workflows"] "chemshell.opt" = "aiida_chemshell.workflows.optimisation:GeometryOptimisationWorkChain" +"chemshell.atomic_energies" = "aiida_chemshell.workflows.isolated_atoms:IsolatedAtomicEnergiesWorkChain" [tool.setuptools.packages.find] -where = ["src",] +where = ["src"] [project.optional-dependencies] dev = [ "aiida-core>=2.6", # Required to support aiida.tools.pytest_fixtures - "pytest>=8.0", + "pytest>=8.0", "pytest-cov>=6.0", "ruff>=0.11.0", - "pre-commit" -] -docs = [ - "sphinx>=8.0.2", - "piccolo-theme>=0.14.0", - "sphinx-toolbox", - "numpydoc", + "pre-commit", ] +docs = ["sphinx>=8.0.2", "piccolo-theme>=0.14.0", "sphinx-toolbox", "numpydoc"] [tool.ruff] -target-version = "py310" -line-length = 88 +target-version = "py310" +line-length = 88 indent-width = 4 -exclude = ["conf.py",] +exclude = ["conf.py"] [tool.ruff.lint.per-file-ignores] "tests/conftest.py" = ["B008"] @@ -69,13 +64,16 @@ select = [ # flake8-bugbear "B", # pylint - "C", "R", + "C", + "R", # pydocstyle "D", # pycodestyle - "E", "W", + "E", + "W", # Pyflakes - "F", "FA", + "F", + "FA", # isort "I", # pep8-naming @@ -83,9 +81,7 @@ select = [ # pyupgrade "UP", ] -ignore = [ - "C901", -] +ignore = ["C901"] [tool.ruff.lint.pydocstyle] convention = "numpy" From f5679ad0f2ea2d53b0d8c5e851aa5340cd87e7ac Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 14:16:07 +0100 Subject: [PATCH 27/35] Remove basis quality input from optimisation workchain, will be handled better in AiiDAlab UI --- src/aiida_chemshell/workflows/optimisation.py | 61 +------------------ 1 file changed, 2 insertions(+), 59 deletions(-) diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index 7b36825..f3034e9 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -2,7 +2,7 @@ from aiida.common.exceptions import MissingEntryPointError from aiida.engine import ToContext, WorkChain -from aiida.orm import ArrayData, Bool, Code, Dict, Float, SinglefileData, Str +from aiida.orm import ArrayData, Bool, Code, Dict, Float, SinglefileData from aiida.plugins.factories import CalculationFactory from aiida_chemshell.calculations.base import ChemShellCalculation @@ -26,13 +26,6 @@ def define(cls, spec) -> None: 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( @@ -60,13 +53,6 @@ def define(cls, spec) -> None: help="The calculated vibrational modes for the optimised structure.", ) - spec.output( - "test_file", - valid_type=SinglefileData, - required=False, - help="", - ) - # Optional inputs/outputs for using the results to fine tune a MLIP model # using the janus-core project via the aiida-mlip plugin. try: @@ -119,18 +105,9 @@ def optimise(self): "theory": "NWChem", "method": "dft", "functional": "B3LYP", - # "d3": True, + "basis": "cc-pvdz", } ) - 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"}) @@ -305,37 +282,3 @@ def result(self): 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 "" From 6ccfd186adfbbaa4ba5d9285818e0f1e5fd877c5 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 14:30:47 +0100 Subject: [PATCH 28/35] Correctly identify only unique elements from cjson based input structure --- src/aiida_chemshell/workflows/isolated_atoms.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py index 25f4246..246c3b5 100644 --- a/src/aiida_chemshell/workflows/isolated_atoms.py +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -105,11 +105,11 @@ def _atom_types_from_cjson(self) -> None: atom_symbols: list[str] = data.get("symbol", []) atom_numbers: list[int] = data.get("number", []) if len(atom_symbols) == 0: - self.unique_atoms = [ - PeriodicTable.atom_z_to_symbol(number) for number in atom_numbers - ] + self.unique_atoms = list( + {PeriodicTable.atom_z_to_symbol(number) for number in atom_numbers} + ) else: - self.unique_atoms = atom_symbols + self.unique_atoms = list(set(atom_symbols)) return def _atom_types_from_pun(self) -> None: From c740ff824793d4f02884d2c5d5c22d0bf40eaa38 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 14:37:20 +0100 Subject: [PATCH 29/35] Fix spelling mistake --- src/aiida_chemshell/calculations/file_conversion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aiida_chemshell/calculations/file_conversion.py b/src/aiida_chemshell/calculations/file_conversion.py index 01edcfe..05238b4 100644 --- a/src/aiida_chemshell/calculations/file_conversion.py +++ b/src/aiida_chemshell/calculations/file_conversion.py @@ -57,7 +57,7 @@ def define(cls, spec: CalcJobProcessSpec) -> None: spec.output( "validation_input", valid_type=SinglefileData, - vlidatory=cls.validate_path_file, + validator=cls.validate_path_file, required=True, help="The validation data set in extended XYZ format.", ) From a18649a2ccd12059cc3a0dc5fb148280c3d729f4 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 14:37:44 +0100 Subject: [PATCH 30/35] Add label and description to create atomistic StructureData nodes in isolated atom workchain --- src/aiida_chemshell/workflows/isolated_atoms.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py index 246c3b5..0cb8c5d 100644 --- a/src/aiida_chemshell/workflows/isolated_atoms.py +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -56,6 +56,11 @@ def atom_energies(self): structure = StructureData() structure.append_atom(position=(0.0, 0.0, 0.0), symbols=atom_symbol) structure.set_pbc((False, False, False)) + structure.label = f"{atom_symbol} atom" + structure.description = ( + f"Isolated {atom_symbol} atom extracted from " + f"{self.inputs.structure.filename}" + ) builder = self.inputs.code.get_builder() builder.structure = structure builder.qm_parameters = self.inputs.qm_parameters From d2c6716b941e1221e26ef0fac1778e35f287e19b Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 15:14:40 +0100 Subject: [PATCH 31/35] Add node labels to created nodes during isolated atoms workflow --- src/aiida_chemshell/workflows/isolated_atoms.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py index 0cb8c5d..4c3a9c4 100644 --- a/src/aiida_chemshell/workflows/isolated_atoms.py +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -58,8 +58,8 @@ def atom_energies(self): structure.set_pbc((False, False, False)) structure.label = f"{atom_symbol} atom" structure.description = ( - f"Isolated {atom_symbol} atom extracted from " - f"{self.inputs.structure.filename}" + f"Isolated {atom_symbol} atom extracted from Node: " + f"{self.inputs.structure.pk}" ) builder = self.inputs.code.get_builder() builder.structure = structure @@ -74,6 +74,16 @@ def result(self) -> None: list(self.unique_atoms), [self.ctx.get(atom).outputs.energy for atom in self.unique_atoms], ) + if isinstance(self.inputs.structure, StructureData): + results_dict.label = ( + f"Atomistic energies for all unique atoms extracted from Node: " + f"{self.inputs.structure.pk}" + ) + else: + results_dict.label = ( + "Atomistic energies for all unique atoms extracted from " + f"{self.inputs.structure.filename} (Node: {self.inputs.structure.pk})" + ) self.out("atom_energies", results_dict) return From bd5ce6785328ffe9d22297be68caf922ffd9177f Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 15:14:57 +0100 Subject: [PATCH 32/35] Remove validators for outputs which are breaking workchain --- src/aiida_chemshell/calculations/file_conversion.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/aiida_chemshell/calculations/file_conversion.py b/src/aiida_chemshell/calculations/file_conversion.py index 05238b4..6664d9f 100644 --- a/src/aiida_chemshell/calculations/file_conversion.py +++ b/src/aiida_chemshell/calculations/file_conversion.py @@ -50,14 +50,12 @@ def define(cls, spec: CalcJobProcessSpec) -> None: spec.output( "training_input", valid_type=SinglefileData, - validator=cls.validate_path_file, required=True, help="The main training data set in extended XYZ format.", ) spec.output( "validation_input", valid_type=SinglefileData, - validator=cls.validate_path_file, required=True, help="The validation data set in extended XYZ format.", ) @@ -97,6 +95,7 @@ def validate_path_file( lines = content.split("\n") natoms = int(lines[0]) if len(lines) % (natoms + 2) != 0: + print(lines) return "Invalid XYZ trajectory structure detected." if len(lines) // (natoms + 2) < 5: return "Not enough individual configurations within input trajectory." From eb5a6af70e384019f35906866fa6fc94346ef80e Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 15:35:18 +0100 Subject: [PATCH 33/35] Improve calling structure in Isolated Atom workchain --- src/aiida_chemshell/workflows/isolated_atoms.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py index 4c3a9c4..bcd5f3d 100644 --- a/src/aiida_chemshell/workflows/isolated_atoms.py +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -61,10 +61,13 @@ def atom_energies(self): f"Isolated {atom_symbol} atom extracted from Node: " f"{self.inputs.structure.pk}" ) - builder = self.inputs.code.get_builder() - builder.structure = structure - builder.qm_parameters = self.inputs.qm_parameters - future = self.submit(builder) + inputs = { + "structure": structure, + "qm_parameters": self.inputs.qm_parameters, + "code": self.inputs.code, + } + future = self.submit(ChemShellCalculation, **inputs) + future.description = f"Single point energy for {atom_symbol} atom." calculations[atom_symbol] = future return ToContext(**calculations) From b6b73abe5dce1e09e2f0f58ab1cf4aeaaca102e0 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 15:37:50 +0100 Subject: [PATCH 34/35] Improve step labelling in optimisation workchain --- src/aiida_chemshell/workflows/optimisation.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index f3034e9..3b1ea53 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -166,6 +166,12 @@ def isolated_atom_energies(self): "qm_parameters": self.ctx.optimise.inputs.qm_parameters, } future = self.submit(IsolatedAtomicEnergiesWorkChain, **inputs) + future.label = "Isolated Atomic Energy WorkChain" + future.description = ( + f"Isolated atom energies extracted from Node: " + f"{self.inputs.structure.pk} for ChemShell optimisation " + f"WorkChain: {self.node.pk} to be used for MLIP fine-tuning." + ) return ToContext(isolated_atoms=future) return None From a1e9bde4fe6499db6151d7655ac860c55561de60 Mon Sep 17 00:00:00 2001 From: "Benjamin T. Speake" Date: Wed, 24 Jun 2026 16:03:44 +0100 Subject: [PATCH 35/35] Fix tests --- tests/test_calculations.py | 3 +- tests/test_inputs.py | 10 ++-- tests/test_workflows.py | 117 ++++++++++++++++++------------------- 3 files changed, 63 insertions(+), 67 deletions(-) diff --git a/tests/test_calculations.py b/tests/test_calculations.py index 35cf5a5..abea831 100644 --- a/tests/test_calculations.py +++ b/tests/test_calculations.py @@ -81,8 +81,7 @@ 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.9468895533 - eref = -75.946889436563 + eref = -75.946889377347 assert abs(results.get("energy") - eref) < 1e-8, ( "Incorrect energy result for PySCF based SP calculation" ) diff --git a/tests/test_inputs.py b/tests/test_inputs.py index f86e197..640d7cf 100644 --- a/tests/test_inputs.py +++ b/tests/test_inputs.py @@ -12,9 +12,7 @@ def test_defaults(generate_calcjob): ChemShellCalculation.FILE_RESULTS, ] code_info = calc_info.codes_info[0] - assert code_info.cmdline_params == [ - ChemShellCalculation.FILE_SCRIPT, - ] + assert ChemShellCalculation.FILE_SCRIPT in code_info.cmdline_params assert code_info.stdout_name == ChemShellCalculation.FILE_STDOUT script_file = tmp_pth / ChemShellCalculation.FILE_SCRIPT @@ -151,5 +149,7 @@ def test_structure_as_structuredata_object( script_txt = script_file.read_text() assert "from chemsh import Fragment\n" in script_txt - tmp_structure_file = ChemShellCalculation.FILE_TMP_STRUCTURE - assert f"structure = Fragment(coords='{tmp_structure_file:s}')\n" in script_txt + assert ( + "structure = Fragment(coords=[[0.0, 0.0, 0.0], [-1.4259993290233388, 1.1149994" + in script_txt + ) diff --git a/tests/test_workflows.py b/tests/test_workflows.py index ab8015c..e3eb6f6 100644 --- a/tests/test_workflows.py +++ b/tests/test_workflows.py @@ -11,9 +11,8 @@ def test_geometry_optimisation_workflow(chemsh_code, get_test_data_file): "chemsh": { "code": chemsh_code(), "structure": get_test_data_file(), - "qm_parameters": {"theory": "PySCF", "method": "HF"}, + "qm_parameters": {"theory": "PySCF", "method": "HF", "basis": "3-21G"}, }, - "basis_quality": "fast", "vibrational_analysis": True, } results, node = run_get_node(GeometryOptimisationWorkChain, **inputs) @@ -22,10 +21,6 @@ def test_geometry_optimisation_workflow(chemsh_code, get_test_data_file): assert len(node.called) > 0, "WorkChain did not launch any subprocesses" - assert node.called[0].inputs.qm_parameters.get( - "basis", "" - ) == GeometryOptimisationWorkChain.get_basis_set_label("fast") - assert abs(results.get("final_energy") - -75.585959742867) < 1e-9, ( "Incorrect final energy for geometry optimisation workflow." ) @@ -42,57 +37,59 @@ def test_geometry_optimisation_workflow(chemsh_code, get_test_data_file): assert (modes[2][2] - 0.0089900284) < 1e-10, "Incorrect ZPE reported for mode 3" -def test_optimisation_workflow_mlip_training( - chemsh_code, get_test_data_file, janus_code -): - """Test a geometry optimisation workflow with vibrational analysis.""" - from aiida_mlip.helpers.help_load import load_model - - inputs = { - "chemsh": { - "code": chemsh_code(), - "structure": get_test_data_file("butanol.cjson"), - "qm_parameters": { - "theory": "PySCF", - "method": "hf", - "functional": "blyp", - }, - }, - "mlip_model": load_model(None, "mace_mp"), - "mlip_code": janus_code, - "basis_quality": "fast", - "vibrational_analysis": False, - } - results, node = run_get_node(GeometryOptimisationWorkChain, **inputs) - - assert node.is_finished_ok, f"WorkChain failed with exit status {node.exit_status}" - - assert len(node.called) > 0, "WorkChain did not launch any subprocesses" - - # assert node.called[0].inputs.qm_parameters.get( - # "basis", "" - # ) == GeometryOptimisationWorkChain.get_basis_set_label("fast") - # - # assert abs(results.get("final_energy") - -75.585959742867) < 1e-9, ( - # "Incorrect final energy for geometry optimisation workflow." - # ) - - subs = node.called - for sub in subs: - # if "Generate MLIP training" in sub.label: - # print(sub.outputs.training_input.content) - if "MLIP Fine-Tuning" in sub.label: - print(sub.outputs.retrieved.list_object_names()) - print("ERROR") - print(sub.outputs.retrieved.get_object_content("_scheduler-stderr.txt")) - print("OUTPUT") - print(sub.outputs.retrieved.get_object_content("_scheduler-stdout.txt")) - print("AIIDA OUTPUT") - print(sub.outputs.retrieved.get_object_content("aiida-stdout.txt")) - print("results") - print(sub.outputs.retrieved.list_object_names(path="results")) - print("LOG") - # print(sub.outputs.retrieved.get_object_content("logs/test_run-123.log")) - - assert sub.is_finished_ok, f"Node '{sub.label}' failed to finish correctly." - # assert False +## Test doesn't work with PySCF which is the only QM backend currently available +# def test_optimisation_workflow_mlip_training( +# chemsh_code, get_test_data_file, janus_code +# ): +# """Test a geometry optimisation workflow with vibrational analysis.""" +# from aiida_mlip.helpers.help_load import load_model + +# inputs = { +# "chemsh": { +# "code": chemsh_code(), +# "structure": get_test_data_file("butanol.cjson"), +# "qm_parameters": { +# "theory": "PySCF", +# "method": "hf", +# "functional": "blyp", +# }, +# }, +# "mlip_model": load_model(None, "mace_mp"), +# "mlip_code": janus_code, +# "basis_quality": "fast", +# "vibrational_analysis": False, +# } +# results, node = run_get_node(GeometryOptimisationWorkChain, **inputs) + +# assert node.is_finished_ok, \ +# f"WorkChain failed with exit status {node.exit_status}" + +# assert len(node.called) > 0, "WorkChain did not launch any subprocesses" + +# # assert node.called[0].inputs.qm_parameters.get( +# # "basis", "" +# # ) == GeometryOptimisationWorkChain.get_basis_set_label("fast") +# # +# # assert abs(results.get("final_energy") - -75.585959742867) < 1e-9, ( +# # "Incorrect final energy for geometry optimisation workflow." +# # ) + +# subs = node.called +# for sub in subs: +# # if "Generate MLIP training" in sub.label: +# # print(sub.outputs.training_input.content) +# if "MLIP Fine-Tuning" in sub.label: +# print(sub.outputs.retrieved.list_object_names()) +# print("ERROR") +# print(sub.outputs.retrieved.get_object_content("_scheduler-stderr.txt")) +# print("OUTPUT") +# print(sub.outputs.retrieved.get_object_content("_scheduler-stdout.txt")) +# print("AIIDA OUTPUT") +# print(sub.outputs.retrieved.get_object_content("aiida-stdout.txt")) +# print("results") +# print(sub.outputs.retrieved.list_object_names(path="results")) +# print("LOG") +# # print(sub.outputs.retrieved.get_object_content("logs/test_run-123.log")) + +# assert sub.is_finished_ok, f"Node '{sub.label}' failed to finish correctly." +# # assert False