fix: test serialization roundtrip with dataset

tests/models/test_eof.py

@@ -1,3 +1,5 @@
+from copy import deepcopy
+
 import numpy as np
 import xarray as xr
 import pytest
@@ -518,3 +520,41 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
         original.inverse_transform(original.scores()),
         loaded.inverse_transform(loaded.scores()),
     )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataarray(dim, mock_data_array):
+    """Test roundtrip serialization when the model is fit on a DataArray."""
+    model = EOF()
+    model.fit(mock_data_array, dim)
+    dt = model.serialize()
+    rebuilt_model = EOF.deserialize(dt)
+    assert np.allclose(
+        model.transform(mock_data_array), rebuilt_model.transform(mock_data_array)
+    )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataset(dim, mock_dataset):
+    """Test roundtrip serialization when the model is fit on a Dataset."""
+    model = EOF()
+    model.fit(mock_dataset, dim)
+    dt = model.serialize()
+    rebuilt_model = EOF.deserialize(dt)
+    assert np.allclose(
+        model.transform(mock_dataset), rebuilt_model.transform(mock_dataset)
+    )

tests/models/test_eof_rotator.py

@@ -244,3 +244,45 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
         original.inverse_transform(original.scores()),
         loaded.inverse_transform(loaded.scores()),
     )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataarray(dim, mock_data_array):
+    """Test roundtrip serialization when the model is fit on a DataArray."""
+    model = EOF()
+    model.fit(mock_data_array, dim)
+    rotator = EOFRotator()
+    rotator.fit(model)
+    dt = rotator.serialize()
+    rebuilt_rotator = EOFRotator.deserialize(dt)
+    assert np.allclose(
+        rotator.transform(mock_data_array), rebuilt_rotator.transform(mock_data_array)
+    )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataset(dim, mock_dataset):
+    """Test roundtrip serialization when the model is fit on a Dataset."""
+    model = EOF()
+    model.fit(mock_dataset, dim)
+    rotator = EOFRotator()
+    rotator.fit(model)
+    dt = rotator.serialize()
+    rebuilt_rotator = EOFRotator.deserialize(dt)
+    assert np.allclose(
+        rotator.transform(mock_dataset), rebuilt_rotator.transform(mock_dataset)
+    )

tests/models/test_mca.py

@@ -432,3 +432,41 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
         original.inverse_transform(*original.scores()),
         loaded.inverse_transform(*loaded.scores()),
     )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataarray(dim, mock_data_array):
+    """Test roundtrip serialization when the model is fit on a DataArray."""
+    model = MCA()
+    model.fit(mock_data_array, mock_data_array, dim)
+    dt = model.serialize()
+    rebuilt_model = MCA.deserialize(dt)
+    assert np.allclose(
+        model.transform(mock_data_array), rebuilt_model.transform(mock_data_array)
+    )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataset(dim, mock_dataset):
+    """Test roundtrip serialization when the model is fit on a Dataset."""
+    model = MCA()
+    model.fit(mock_dataset, mock_dataset, dim)
+    dt = model.serialize()
+    rebuilt_model = MCA.deserialize(dt)
+    assert np.allclose(
+        model.transform(mock_dataset), rebuilt_model.transform(mock_dataset)
+    )

tests/models/test_mca_rotator.py

@@ -304,3 +304,45 @@ def test_save_load(dim, mock_data_array, tmp_path, engine):
         original.inverse_transform(*original.scores()),
         loaded.inverse_transform(*loaded.scores()),
     )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataarray(dim, mock_data_array):
+    """Test roundtrip serialization when the model is fit on a DataArray."""
+    model = MCA()
+    model.fit(mock_data_array, mock_data_array, dim)
+    rotator = MCARotator()
+    rotator.fit(model)
+    dt = rotator.serialize()
+    rebuilt_rotator = MCARotator.deserialize(dt)
+    assert np.allclose(
+        rotator.transform(mock_data_array), rebuilt_rotator.transform(mock_data_array)
+    )
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        (("time",)),
+        (("lat", "lon")),
+        (("lon", "lat")),
+    ],
+)
+def test_serialize_deserialize_dataset(dim, mock_dataset):
+    """Test roundtrip serialization when the model is fit on a Dataset."""
+    model = MCA()
+    model.fit(mock_dataset, mock_dataset, dim)
+    rotator = MCARotator()
+    rotator.fit(model)
+    dt = rotator.serialize()
+    rebuilt_rotator = MCARotator.deserialize(dt)
+    assert np.allclose(
+        rotator.transform(mock_dataset), rebuilt_rotator.transform(mock_dataset)
+    )