Implemented it for all architectures and added tests

Author: pfebrer

src/metatrain/composition/model.py

@@ -1,6 +1,6 @@
 import logging
 import warnings
-from typing import Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 import metatensor.torch as mts
 import torch
@@ -22,6 +22,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.metadata import merge_metadata
 
 from . import checkpoints
@@ -204,7 +205,9 @@ def train_model(
             checkpoint_dir="",
         )
 
-    def restart(self, dataset_info: DatasetInfo) -> "CompositionModel":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "CompositionModel":
         """
         Update the model to continue training, possibly with new targets.
 
@@ -214,8 +217,13 @@ def restart(self, dataset_info: DatasetInfo) -> "CompositionModel":
 
         :param dataset_info: Information about the new dataset, including the
             targets that will be used for training.
+        :param model_hypers: New hyperparameters for the model. They must match
+            the existing hyperparameters, otherwise an error is raised.
         :return: The updated model.
         """
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         raw_targets = {}
         for target_name in dataset_info.targets:
             target_info = dataset_info.targets[target_name]

src/metatrain/experimental/classifier/model.py

@@ -123,7 +123,9 @@ def build_mlp(self, feature_size: int, num_classes: int) -> None:
         # Final classification layer
         self.linear = torch.nn.Linear(current_size, num_classes, bias=False)
 
-    def restart(self, dataset_info: DatasetInfo) -> "Classifier":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "Classifier":
         raise ValueError("Restarting from a Classifier model is not supported.")
 
     def forward(

src/metatrain/experimental/dpa3/model.py

@@ -25,6 +25,7 @@
 from metatrain.utils.data.atom_pair_helpers import check_no_atom_pair_targets
 from metatrain.utils.data.dataset import DatasetInfo
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
 from metatrain.utils.sum_over_atoms import sum_over_atoms
@@ -391,7 +392,13 @@ def forward(
 
         return return_dict
 
-    def restart(self, dataset_info: DatasetInfo) -> "DPA3":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "DPA3":
+
+        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/experimental/flashmd/model.py

@@ -28,6 +28,7 @@
 from metatrain.utils.data import DatasetInfo, TargetInfo
 from metatrain.utils.data.atom_pair_helpers import check_no_atom_pair_targets
 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
@@ -232,7 +233,13 @@ def set_timestep(self, timestep: float):
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "FlashMD":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "FlashMD":
+
+        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/experimental/flashmd_symplectic/model.py

@@ -31,6 +31,7 @@
 from metatrain.utils.data.atom_pair_helpers import check_no_atom_pair_targets
 from metatrain.utils.data.target_info import get_energy_target_info
 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
@@ -222,7 +223,13 @@ def set_timestep(self, timestep: float):
     def supported_outputs(self) -> Dict[str, ModelOutput]:
         return self.outputs
 
-    def restart(self, dataset_info: DatasetInfo) -> "FlashMDSymplectic":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "FlashMDSymplectic":
+
+        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/experimental/mace/model.py

@@ -30,6 +30,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.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
 from metatrain.utils.sum_over_atoms import sum_over_atoms
@@ -306,7 +307,13 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
 
         self.finetune_config: Dict[str, Any] = {}
 
-    def restart(self, dataset_info: DatasetInfo) -> "MetaMACE":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "MetaMACE":
+
+        if model_hypers is not None:
+            raise_if_hypers_mismatch(self.hypers, model_hypers)
+
         # Check that the new dataset info does not contain new atomic types
         if new_atomic_types := set(dataset_info.atomic_types) - set(
             self.dataset_info.atomic_types

src/metatrain/experimental/space/model.py

@@ -35,6 +35,7 @@
 )
 from metatrain.utils.data.dataset import DatasetInfo, TargetInfo
 from metatrain.utils.dtype import dtype_to_str
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.scaler import Scaler
 
@@ -166,7 +167,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) -> "SPACE":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "SPACE":
+
+        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/experimental/space/tests/test_basic.py

@@ -11,6 +11,7 @@
 from metatrain.utils.testing import (
     ArchitectureTests,
     CheckpointTests,
+    InputTests,
     OutputTests,
     TorchscriptTests,
 )
@@ -31,6 +32,9 @@ def minimal_model_hypers(self):
         return hypers
 
 
+class TestInput(InputTests, SPACETests): ...
+
+
 class TestOutput(OutputTests, SPACETests):
     is_equivariant_reflections = False
     equivariance_error_tolerance = 1e-4  # due to many layers in the default hypers

src/metatrain/gap/model.py

@@ -169,7 +169,9 @@ 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) -> "GAP":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "GAP":
         raise NotImplementedError("GAP does not allow restarting training")
 
     @classmethod

src/metatrain/llpr/model.py

@@ -26,6 +26,7 @@
 from metatrain.utils.data.target_info import (
     is_auxiliary_output,
 )
+from metatrain.utils.hypers import raise_if_hypers_mismatch
 from metatrain.utils.io import model_from_checkpoint
 from metatrain.utils.metadata import merge_metadata
 from metatrain.utils.neighbor_lists import (
@@ -250,7 +251,13 @@ def set_wrapped_model(self, model: ModelInterface) -> None:
                 bias=False,
             )
 
-    def restart(self, dataset_info: DatasetInfo) -> "LLPRUncertaintyModel":
+    def restart(
+        self, dataset_info: DatasetInfo, model_hypers: Optional[dict[str, Any]] = None
+    ) -> "LLPRUncertaintyModel":
+
+        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 = [