Source code for aiida_mlip.calculations.train

"""Class for training machine learning models."""

from __future__ import annotations

from pathlib import Path
import shutil

from aiida.common import InputValidationError, datastructures
import aiida.common.folders
from aiida.engine import CalcJob, CalcJobProcessSpec
import aiida.engine.processes
from aiida.orm import Bool, Dict, FolderData, SinglefileData, Str

from aiida_mlip.data.config import JanusConfigfile
from aiida_mlip.data.model import ModelData
from aiida_mlip.helpers.converters import kwarg_to_param




def validate_inputs(
    inputs: dict, port_namespace: aiida.engine.processes.ports.PortNamespace
):
    """
    Check if the inputs are valid.

    Parameters
    ----------
    inputs : dict
        The inputs dictionary.

    port_namespace : `aiida.engine.processes.ports.PortNamespace`
        An instance of aiida's `PortNameSpace`.

    Raises
    ------
    InputValidationError
        Error message if validation fails, None otherwise.
    """
    if "mlip_config" in port_namespace:
        # Check if a config file is given
        if "mlip_config" not in inputs:
            raise InputValidationError("No config file given")
        config_file = inputs["mlip_config"]
        # Check if 'name' keyword is given
        if "name" not in config_file:
            raise InputValidationError("key 'name' must be defined in the config file")
        # Check if the xyz files paths are given
        required_keys = ("train_file", "valid_file", "test_file")
        for key in required_keys:
            if key not in config_file:
                raise InputValidationError(f"Mandatory key {key} not in config file")
            # Check if the keys actually correspond to a path
            if not ((Path(config_file.as_dictionary[key])).resolve()).exists():
                raise InputValidationError(f"Path given for {key} does not exist")
        # Check if fine-tuning is enabled and validate accordingly
        if (
            inputs["fine_tune"]
            and "foundation_model" not in config_file
            and "foundation_model" not in inputs
        ):
            raise InputValidationError(
                "Undefined Model to fine-tune in inputs or config file"
            )

    if "arch" in port_namespace:
        model_arch = getattr(inputs.get("foundation_model"), "architecture", None)
        # `arch` is a positional argument of `janus train`, so it must be determined
        if "arch" not in inputs and model_arch is None:
            raise InputValidationError(
                "'arch' must be specified in inputs or ModelData"
            )
        if (
            "arch" in inputs
            and model_arch is not None
            and inputs["arch"].value != model_arch
        ):
            raise InputValidationError(
                "'arch' in ModelData and in inputs must be the same, "
                f"but they are {model_arch} and {inputs['arch'].value}"
            )





class Train(CalcJob):  # numpydoc ignore=PR01
    """
    Calcjob implementation to train mlips.

    Attributes
    ----------
    DEFAULT_OUTPUT_FILE : str
        Default stdout file name.

    Methods
    -------
    define(spec: CalcJobProcessSpec) -> None:
        Define the process specification, its inputs, outputs and exit codes.
    validate_inputs(value: dict, port_namespace: PortNamespace) -> str | None:
        Check if the inputs are valid.
    prepare_for_submission(folder: Folder) -> CalcInfo:
        Create the input files for the `CalcJob`.
    """

    DEFAULT_OUTPUT_FILE = "aiida-stdout.txt"



    @classmethod
    def define(cls, spec: CalcJobProcessSpec) -> None:
        """
        Define the process specification, its inputs, outputs and exit codes.

        Parameters
        ----------
        spec : `aiida.engine.CalcJobProcessSpec`
            The calculation job process spec to define.
        """
        super().define(spec)

        # Define inputs
        spec.input(
            "mlip_config",
            valid_type=JanusConfigfile,
            required=True,
            help="Config file with parameters for training",
        )

        spec.input(
            "arch",
            valid_type=Str,
            required=False,
            help="MLIP architecture to train. Inferred from `foundation_model` if "
            "not set.",
        )

        spec.input(
            "fine_tune",
            valid_type=Bool,
            required=False,
            default=lambda: Bool(False),
            help="Whether fine-tuning a model",
        )
        spec.input(
            "foundation_model",
            valid_type=ModelData,
            required=False,
            help="Model to fine-tune",
        )

        spec.input(
            "metadata.options.output_filename",
            valid_type=str,
            default=cls.DEFAULT_OUTPUT_FILE,
        )

        spec.input(
            "metadata.options.scheduler_stdout",
            valid_type=str,
            default="_scheduler-stdout.txt",
            help="Filename to which the content of stdout of the scheduler is written.",
        )
        spec.inputs["metadata"]["options"]["parser_name"].default = "mlip.train_parser"
        spec.inputs.validator = validate_inputs
        spec.output("model", valid_type=ModelData)
        spec.output("compiled_model", valid_type=SinglefileData)
        spec.output(
            "results_dict",
            valid_type=Dict,
            help="The `results_dict` output node of the training.",
        )
        spec.output("logs", valid_type=FolderData)
        spec.output("checkpoints", valid_type=FolderData)
        spec.default_output_node = "results_dict"
        # Exit codes
        spec.exit_code(
            305,
            "ERROR_MISSING_OUTPUT_FILES",
            message="Some output files missing or cannot be read",
        )




    def prepare_for_submission(
        self, folder: aiida.common.folders.Folder
    ) -> datastructures.CalcInfo:
        """
        Create the input files for the `Calcjob`.

        Parameters
        ----------
        folder : aiida.common.folders.Folder
            Folder where the calculation is run.

        Returns
        -------
        aiida.common.datastructures.CalcInfo
            An instance of `aiida.common.datastructures.CalcInfo`.
        """
        # The config file needs to be copied in the working folder
        # Read content
        mlip_dict = self.inputs.mlip_config.as_dictionary
        config_parse = self.inputs.mlip_config.get_content()

        # Extract paths from the config
        for file in ("train_file", "test_file", "valid_file"):
            abs_path = Path(mlip_dict[file]).resolve()

            # Update the config file with absolute paths
            config_parse = config_parse.replace(mlip_dict[file], str(abs_path))

        # Add foundation_model to the config file if fine-tuning is enabled
        if self.inputs.fine_tune and "foundation_model" in self.inputs:
            with self.inputs.foundation_model.open(mode="rb") as source:
                with folder.open("mlff.model", mode="wb") as target:
                    shutil.copyfileobj(source, target)
            config_parse += "\nfoundation_model: mlff.model"

        # Copy config file content inside the folder where the calculation is run
        config_copy = "mlip_train.yml"
        with folder.open(config_copy, "w", encoding="utf-8") as configfile:
            configfile.write(config_parse)

        codeinfo = datastructures.CodeInfo()

        # Initialize cmdline_params with train command
        # Create the rest of the command line
        cmd_line = {
            "mlip-config": config_copy,
            "fine-tune": bool(self.inputs.fine_tune),
            # janus writes to ./janus_results by default, but the parser resolves the
            # output directories relative to the remote working directory
            "file-prefix": ".",
        }

        # `arch` is a positional argument of `janus train`, so it cannot be passed
        # through `kwarg_to_param`
        codeinfo.cmdline_params = [
            "train",
            self._get_arch(),
            *kwarg_to_param(cmd_line),
        ]

        # Node where the code is saved
        codeinfo.code_uuid = self.inputs.code.uuid
        # Save name of output as you need it for running the code
        codeinfo.stdout_name = self.metadata.options.output_filename

        calcinfo = datastructures.CalcInfo()
        calcinfo.codes_info = [codeinfo]
        # Save the info about the node where the calc is stored
        calcinfo.uuid = str(self.uuid)
        # Retrieve output files
        model_dir = Path(mlip_dict.get("model_dir", "."))
        model_output = model_dir / f"{mlip_dict['name']}.model"
        compiled_model_output = model_dir / f"{mlip_dict['name']}_compiled.model"
        calcinfo.retrieve_list = [
            self.metadata.options.output_filename,
            self.uuid,
            mlip_dict.get("log_dir", "logs"),
            mlip_dict.get("result_dir", "results"),
            mlip_dict.get("checkpoint_dir", "checkpoints"),
            str(model_output),
            str(compiled_model_output),
        ]

        return calcinfo




    def _get_arch(self) -> str:
        """
        Find the architecture to train, from the inputs or the model to fine-tune.

        Returns
        -------
        str
            The MLIP architecture to pass to `janus train`.
        """
        if "arch" in self.inputs:
            return str(self.inputs.arch.value)
        return str(self.inputs.foundation_model.architecture)