Preserve DCL dynamics parameters in checkpoints

cebra/solver/single_session.py

@@ -22,6 +22,7 @@
 """Single session solvers embed a single pair of time series."""
 
 import copy
+import warnings
 from typing import Optional
 
 import literate_dataclasses as dataclasses
@@ -224,9 +225,35 @@ def _inference(self, batch: cebra.data.Batch) -> cebra.data.Batch:
         return cebra.data.Batch(ref, pos, neg)
 
 
+class _DynamicsStateMixin:
+    """Include dynamics parameters in solver checkpoints."""
+
+    def state_dict(self) -> dict:
+        """Save the encoder, optimizer, criterion, and dynamics model."""
+        state = super().state_dict()
+        state["dynamics_model"] = self.dynamics_model.state_dict()
+        return state
+
+    def load_state_dict(self, state_dict: dict, strict: bool = True):
+        """Restore dynamics; reject incomplete legacy checkpoints in strict mode."""
+        if "dynamics_model" not in state_dict:
+            message = (
+                "Checkpoint is missing dynamics_model parameters. "
+                "Legacy DCL checkpoints did not save them. "
+                "Use strict=False only to retain the initialized dynamics.")
+            if strict:
+                raise KeyError(message)
+            warnings.warn(message, UserWarning, stacklevel=2)
+        else:
+            self.dynamics_model.load_state_dict(state_dict["dynamics_model"],
+                                                strict=strict)
+        super().load_state_dict(state_dict, strict=strict)
+
+
 @register("single-session-dcl")
 @dataclasses.dataclass
-class SingleSessionDCLSolver(SingleSessionAuxVariableSolver):
+class SingleSessionDCLSolver(_DynamicsStateMixin,
+                             SingleSessionAuxVariableSolver):
     dynamics_model: torch.nn.Module = None
 
     def __post_init__(self):
@@ -368,7 +395,7 @@ def _inference(self, batch: cebra.data.Batch) -> cebra.data.Batch:
 
 @register("single-session-full-dcl")
 @dataclasses.dataclass
-class BatchSingleSessionDCLSolver(BatchSingleSessionSolver,
+class BatchSingleSessionDCLSolver(_DynamicsStateMixin, BatchSingleSessionSolver,
                                   SingleSessionAuxVariableSolver):
     """Optimize a DCL model with batch gradient descent.
 

docs/source/api/pytorch/dynamics.rst

@@ -23,4 +23,14 @@ Models
 
 .. automodule:: cebra.dynamics.linear
    :members:
-   :show-inheritance:
+   :show-inheritance:
+
+Checkpoint behavior
+~~~~~~~~~~~~~~~~~~~
+
+Both minibatch and full-batch DCL solver checkpoints include the dynamics
+parameters. Historical DCL checkpoints that omitted those parameters cannot
+restore the trained dynamics. Strict loading raises an explanatory error.
+Low-level ``solver.load_state_dict(checkpoint, strict=False)`` loads the
+available components and warns that the dynamics retain their initialization.
+Ordinary CEBRA checkpoints have no dynamics requirement.

tests/test_dcl_checkpoint.py

@@ -0,0 +1,62 @@
+#
+# CEBRA: Consistent EmBeddings of high-dimensional Recordings using Auxiliary variables
+# © Mackenzie W. Mathis & Steffen Schneider (v0.4.0+)
+# Source code:
+# https://github.com/AdaptiveMotorControlLab/CEBRA
+#
+# Please see LICENSE.md for the full license document:
+# https://github.com/AdaptiveMotorControlLab/CEBRA/blob/main/LICENSE.md
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""DCL checkpoints preserve dynamics as well as encoder parameters."""
+
+import numpy as np
+import pytest
+import torch
+
+import cebra
+
+
+@pytest.mark.parametrize("architecture", ["linear", "orthogonal-linear"])
+@pytest.mark.parametrize("batch_size", [8, None])
+def test_existing_dcl_fit_save_load_and_resume(tmp_path, architecture,
+                                               batch_size):
+    from cebra.integrations.sklearn.cebra import DCL
+    torch.manual_seed(4)
+    data = np.random.default_rng(4).normal(size=(100, 5)).astype("float32")
+    model = DCL(device="cpu",
+                dynamics_model_architecture=architecture,
+                output_dimension=3,
+                batch_size=batch_size,
+                batch_size_negatives=11,
+                max_iterations=2,
+                distance="euclidean")
+    model.fit(data)
+    path = tmp_path / "flow.pt"
+    model.save(path, backend="sklearn")
+    loaded = cebra.CEBRA.load(path, backend="sklearn")
+    np.testing.assert_allclose(model.transform(data), loaded.transform(data))
+    state = torch.randn(7, 3)
+    torch.testing.assert_close(model.dynamics_model_(state),
+                               loaded.dynamics_model_(state))
+    before = loaded.dynamics_model_(state).detach().clone()
+    loaded.partial_fit(data)
+    assert torch.isfinite(loaded.dynamics_model_(state)).all()
+    assert not torch.equal(before, loaded.dynamics_model_(state))
+    legacy = model.solver_.state_dict()
+    del legacy["dynamics_model"]
+    with pytest.raises(KeyError, match="missing dynamics_model"):
+        loaded.solver_.load_state_dict(legacy)
+    with pytest.warns(UserWarning, match="Legacy DCL"):
+        loaded.solver_.load_state_dict(legacy, strict=False)