diff --git a/pyproject.toml b/pyproject.toml index 5fdc064..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,36 +27,34 @@ 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" +"chemshell.file_conversion.mlip_training" = "aiida_chemshell.parsers.file_conversion:CreateJanusTrainingInputsParser" [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"] @@ -68,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 @@ -82,9 +81,7 @@ select = [ # pyupgrade "UP", ] -ignore = [ - "C901", -] +ignore = ["C901"] [tool.ruff.lint.pydocstyle] convention = "numpy" diff --git a/src/aiida_chemshell/calculations/base.py b/src/aiida_chemshell/calculations/base.py index 5cd40bf..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 @@ -23,6 +24,8 @@ class ChemShellCalculation(CalcJob): FILE_DLFIND = "_dl_find.cjson" FILE_TMP_STRUCTURE = "input_structure.xyz" FILE_RESULTS = "result.json" + FILE_TRJPTH = "path.xyz" + FILE_TRJFRC = "path_force.xyz" @classmethod def define(cls, spec: CalcJobProcessSpec) -> None: @@ -41,8 +44,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." ), ) @@ -131,6 +134,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, @@ -160,10 +171,26 @@ 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, - "num_mpiprocs_per_machine": 1, + "num_mpiprocs_per_machine": 4, } spec.inputs["metadata"]["options"]["parser_name"].default = "chemshell" @@ -327,6 +354,7 @@ def get_valid_optimisation_parameter_keys(cls) -> tuple[str]: "delta", "tsrelative", "thermal", + "save_path", ) @classmethod @@ -565,7 +593,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." ) @@ -664,22 +692,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 @@ -700,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 @@ -829,9 +881,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 @@ -871,5 +935,12 @@ 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( + "_dl_find/" + ChemShellCalculation.FILE_TRJPTH + ) + calc_info.retrieve_list.append( + "_dl_find/" + ChemShellCalculation.FILE_TRJFRC + ) return calc_info diff --git a/src/aiida_chemshell/calculations/file_conversion.py b/src/aiida_chemshell/calculations/file_conversion.py new file mode 100644 index 0000000..6664d9f --- /dev/null +++ b/src/aiida_chemshell/calculations/file_conversion.py @@ -0,0 +1,206 @@ +"""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, PortNamespace +from aiida.orm import ArrayData, Dict, SinglefileData, Str + +from aiida_chemshell.units import UnitsConverter + + +class CreateJanusTrainingInputsCalcJob(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( + "energies", + valid_type=ArrayData, + required=True, + 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( + "filename", + valid_type=Str, + required=False, + help="The name to give the directory of output files.", + ) + + spec.output( + "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 validation data set in extended XYZ format.", + ) + spec.output( + "test_input", + valid_type=SinglefileData, + required=True, + help="The testing data set in extended XYZ format.", + ) + + ## Metadata + spec.inputs["metadata"]["options"]["resources"].default = { + "num_machines": 1, + "num_mpiprocs_per_machine": 2, + } + spec.inputs["metadata"]["options"][ + "parser_name" + ].default = "chemshell.file_conversion.mlip_training" + + # Exit Codes + spec.exit_code( + 300, + "ERROR_NO_TRAJECTORY_FILES", + message="No trajectory files have been produced.", + ) + + 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: + print(lines) + 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 = "" + 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() + 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"{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 + 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.provenance_exclude_list = [] + calc_info.retrieve_temporary_list = ["train.xyz", "valid.xyz", "test.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() +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 +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(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") + " ") + 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/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/parsers/base.py b/src/aiida_chemshell/parsers/base.py index ed41848..8d94bf1 100644 --- a/src/aiida_chemshell/parsers/base.py +++ b/src/aiida_chemshell/parsers/base.py @@ -81,19 +81,58 @@ 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( file=f, filename=ChemShellCalculation.FILE_DLFIND, - label="Structure File", + label="CJSON 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() + ): + 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 + return ExitCode(0) def parse_vibrational_analysis(self, stdout: str) -> None: @@ -135,3 +174,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/src/aiida_chemshell/parsers/file_conversion.py b/src/aiida_chemshell/parsers/file_conversion.py new file mode 100644 index 0000000..fef13d0 --- /dev/null +++ b/src/aiida_chemshell/parsers/file_conversion.py @@ -0,0 +1,54 @@ +"""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 + + +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." + retrieved_temporary_folder = kwargs.get("retrieved_temporary_folder", None) + + # 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( + file=f, + filename="train.xyz", + label="MLIP training data.", + description=description_str.format("training"), + ), + ) + + # 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( + file=f, + filename="test.xyz", + label="MLIP test data.", + description=description_str.format("testing"), + ), + ) + + # 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( + file=f, + filename="valid.xyz", + label="MLIP validation data.", + description=description_str.format("validation"), + ), + ) + + return ExitCode(0) 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 new file mode 100644 index 0000000..080f19e --- /dev/null +++ b/src/aiida_chemshell/units.py @@ -0,0 +1,28 @@ +"""Unit conversion utilities.""" + + +class UnitsConverter: + """Unit conversion utility.""" + + ANGSTROM = 0.529177 + EV = 27.2114 + + @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 + + @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/utils.py b/src/aiida_chemshell/utils.py index 33c8f10..26cce95 100644 --- a/src/aiida_chemshell/utils.py +++ b/src/aiida_chemshell/utils.py @@ -85,3 +85,42 @@ 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, + } diff --git a/src/aiida_chemshell/workflows/isolated_atoms.py b/src/aiida_chemshell/workflows/isolated_atoms.py new file mode 100644 index 0000000..bcd5f3d --- /dev/null +++ b/src/aiida_chemshell/workflows/isolated_atoms.py @@ -0,0 +1,135 @@ +"""Workflows for isolating atomic species and calculating SP energies.""" + +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") + + +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) + else: + self._atom_types_from_file() + return + + 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) + structure.set_pbc((False, False, False)) + structure.label = f"{atom_symbol} atom" + structure.description = ( + f"Isolated {atom_symbol} atom extracted from Node: " + f"{self.inputs.structure.pk}" + ) + 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) + + def result(self) -> None: + """Collect the results into a dictionary.""" + results_dict = create_dictionary( + 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 + + def _atom_types_from_structuredata(self, structure: StructureData) -> None: + """Determine the unique atom types from a StructureData object.""" + 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.""" + 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 = list( + {PeriodicTable.atom_z_to_symbol(number) for number in atom_numbers} + ) + else: + self.unique_atoms = list(set(atom_symbols)) + return + + def _atom_types_from_pun(self) -> None: + """Determine the unique atom types from a .pun structure file.""" + return diff --git a/src/aiida_chemshell/workflows/optimisation.py b/src/aiida_chemshell/workflows/optimisation.py index acabd5c..3b1ea53 100644 --- a/src/aiida_chemshell/workflows/optimisation.py +++ b/src/aiida_chemshell/workflows/optimisation.py @@ -1,9 +1,12 @@ """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.orm import ArrayData, Bool, Code, Dict, Float, SinglefileData +from aiida.plugins.factories import CalculationFactory from aiida_chemshell.calculations.base import ChemShellCalculation +from aiida_chemshell.workflows.isolated_atoms import IsolatedAtomicEnergiesWorkChain class GeometryOptimisationWorkChain(WorkChain): @@ -23,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( @@ -57,8 +53,46 @@ def define(cls, spec) -> None: help="The calculated vibrational modes for the optimised structure.", ) + # 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.result) + spec.outline( + cls.optimise, + cls.energy, + cls.isolated_atom_energies, + cls.generate_mlip_training_inputs, + cls.train_mlip, + cls.result, + ) return @@ -71,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"}) @@ -92,8 +117,14 @@ 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) + future.description = ( + f"Geometry optimisation step from WorkChainNode pk: {self.node.pk}" + ) return ToContext(optimise=future) def energy(self): @@ -118,12 +149,136 @@ def energy(self): ) inputs["optimisation_parameters"]["thermal"] = 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 + 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.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 + + def generate_mlip_training_inputs(self): + """Convert the optimisation path files to Janus compatible inputs.""" + if "mlip_model" in self.inputs: + inputs = { + "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": { + "resources": {"num_mpiprocs_per_machine": 2, "num_machines": 1}, + # "withmpi": True, + } + }, + } + from aiida_chemshell.calculations.file_conversion import ( + CreateJanusTrainingInputsCalcJob, + ) + + 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 + + def train_mlip(self): + """Train a given MLIP model.""" + try: + mlip_train_calc = CalculationFactory("mlip.train") + except MissingEntryPointError: + pass + else: + 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 = ( + f"MLIP fine-tuning step from WorkChainNode pk: {self.node.pk}" + ) + 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) + 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( @@ -133,37 +288,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 "" 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_calculations.py b/tests/test_calculations.py index 7e407e4..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" ) @@ -178,6 +177,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.""" 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 0183a12..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,11 +21,7 @@ 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-10, ( + assert abs(results.get("final_energy") - -75.585959742867) < 1e-9, ( "Incorrect final energy for geometry optimisation workflow." ) @@ -40,3 +35,61 @@ 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" + + +## 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