Check input model hypers when restarting trainings

src/metatrain/cli/train.py

@@ -298,6 +298,7 @@ def train_model(
     # MERGE OPTIONS ###########
     ###########################
 
+    input_options = copy.deepcopy(options)
     options = OmegaConf.merge(
         BASE_OPTIONS,
         {"architecture": get_default_hypers(architecture_name)},
@@ -631,6 +632,8 @@ def train_model(
     else:
         training_context = None
 
+    new_model_hypers = input_options.get("architecture", {}).get("model", {})
+
     try:
         if training_context == "restart" and restart_from is not None:
             logging.info(f"Restarting training from '{restart_from}'")
@@ -644,7 +647,7 @@ def train_model(
                     f"The file {restart_from} does not contain a valid checkpoint for "
                     f"the '{architecture_name}' architecture"
                 ) from e
-            model = model.restart(dataset_info)
+            model = model.restart(dataset_info, model_hypers=new_model_hypers)
             try:
                 trainer = trainer_from_checkpoint(
                     checkpoint=checkpoint,
@@ -668,7 +671,7 @@ def train_model(
                     f"The file {restart_from} does not contain a valid checkpoint for "
                     f"the '{architecture_name}' architecture"
                 ) from e
-            model = model.restart(dataset_info)
+            model = model.restart(dataset_info, model_hypers=new_model_hypers)
             trainer = Trainer(hypers["training"])
         else:
             logging.info("Starting training from scratch")

src/metatrain/pet/model.py

@@ -29,6 +29,7 @@
     sparsify_atomic_basis_target,
 )
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.long_range import DummyLongRangeFeaturizer, LongRangeFeaturizer
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
@@ -202,7 +203,13 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "PET":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict] = None
+    ) -> "PET":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
         new_atomic_types = [

src/metatrain/utils/abc.py

@@ -135,7 +135,9 @@ def supported_outputs(self) -> Dict[str, ModelOutput]:
         """
 
     @abstractmethod
-    def restart(self, dataset_info: DatasetInfo) -> "ModelInterface":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict] = None
+    ) -> "ModelInterface":
         """
         Update a model to restart training, potentially with different dataset and/or
         targets.
@@ -147,6 +149,8 @@ def restart(self, dataset_info: DatasetInfo) -> "ModelInterface":
         :param dataset_info: Information about the new dataset, including the targets
             that will be used for training.
 
+        :param model_hypers: The new hyperparameters for the model.
+
         :return: The updated model, or a new instance of the model, that is able to
             handle the new dataset.
         """

src/metatrain/utils/hypers.py

@@ -1,4 +1,4 @@
-from typing import Type, TypedDict, TypeVar
+from typing import Any, Type, TypedDict, TypeVar
 
 from typing_extensions import TypedDict as TE_TypedDict
 
@@ -97,3 +97,60 @@ def overwrite_defaults(
     :param new_defaults: A dict with the new default hyperparameters.
     """
     _OVERWRITTEN_DEFAULTS[hypers_cls] = new_defaults
+
+
+def get_hypers_diff(
+    old_hypers: dict,
+    new_hypers: dict,
+) -> dict[str, tuple[Any, Any]]:
+    """Get the difference between two hypers dictionaries.
+
+    :param old_hypers: The old hyperparameters.
+    :param new_hypers: The new hyperparameters.
+
+    :return: A dict with the hyperparameters that are different in the new
+        hypers compared to the old hypers. It is assumed that every key in
+        the new hypers is also present in the old hypers.
+    """
+    diff = {}
+    for key, new_value in new_hypers.items():
+        old_value = old_hypers[key]
+        if old_value != new_value:
+            diff[key] = (old_value, new_value)
+    return diff
+
+
+def raise_hypers_mismatch(
+    hypers_diff: dict[str, tuple[Any, Any]],
+) -> None:
+    """Raise an error if the hypers diff is not empty.
+
+    The hypers diff can be computed using :func:`get_hypers_diff`.
+    """
+    if hypers_diff:
+        n_mismatches = len(hypers_diff)
+        raise ValueError(
+            f"Found {n_mismatches} mismatch{(n_mismatches != 1) * 'es'} in model hyperparameters.\n"
+            f"Mismatched hypers: {list(hypers_diff.keys())}\n"
+            "\n-------- Mismatches --------\n\n"
+            + "\n".join(
+                f"[Mismatch {i + 1}] {key}\n Previous: {old}\n New: {new}"
+                for i, (key, (old, new)) in enumerate(hypers_diff.items())
+            )
+        )
+
+
+def raise_if_hypers_mismatch(
+    old_hypers: dict,
+    new_hypers: dict,
+) -> None:
+    """Raise an error if the new hypers do not match the old hypers.
+
+    :param old_hypers: The old hyperparameters.
+    :param new_hypers: The new hyperparameters.
+    """
+    # Gather mismatchs
+    mismatches = get_hypers_diff(old_hypers, new_hypers)
+
+    if mismatches:
+        raise_hypers_mismatch(mismatches)