add serialization tests, minor fix to cosine schedule config

Author: vpratz

bayesflow/experimental/diffusion_model/cosine_noise_schedule.py

@@ -41,6 +41,7 @@ def __init__(
         """
         super().__init__(name="cosine_noise_schedule", variance_type="preserving", weighting=weighting)
         self._shift = shift
+        self._weighting = weighting
         self.log_snr_min = min_log_snr
         self.log_snr_max = max_log_snr
 
@@ -75,7 +76,9 @@ def derivative_log_snr(self, log_snr_t: Tensor, training: bool) -> Tensor:
         return -factor * dsnr_dt
 
     def get_config(self):
-        return dict(min_log_snr=self.log_snr_min, max_log_snr=self.log_snr_max, shift=self._shift)
+        return dict(
+            min_log_snr=self.log_snr_min, max_log_snr=self.log_snr_max, shift=self._shift, weighting=self._weighting
+        )
 
     @classmethod
     def from_config(cls, config, custom_objects=None):

bayesflow/experimental/diffusion_model/noise_schedule.py

@@ -138,7 +138,7 @@ def get_weights_for_snr(self, log_snr_t: Tensor) -> Tensor:
             raise TypeError(f"Unknown weighting type: {self._weighting}")
 
     def get_config(self):
-        return dict(name=self.name, variance_type=self._variance_type)
+        return dict(name=self.name, variance_type=self._variance_type, weighting=self._weighting)
 
     @classmethod
     def from_config(cls, config, custom_objects=None):

tests/test_networks/test_diffusion_model/__init__.py

@@ -0,0 +1,23 @@
+import pytest
+
+
+@pytest.fixture()
+def cosine_noise_schedule():
+    from bayesflow.experimental.diffusion_model import CosineNoiseSchedule
+
+    return CosineNoiseSchedule(min_log_snr=-12, max_log_snr=12, shift=0.1, weighting="likelihood_weighting")
+
+
+@pytest.fixture()
+def edm_noise_schedule():
+    from bayesflow.experimental.diffusion_model import EDMNoiseSchedule
+
+    return EDMNoiseSchedule(sigma_data=10.0, sigma_min=1e-5, sigma_max=85.0)
+
+
+@pytest.fixture(
+    params=["cosine_noise_schedule", "edm_noise_schedule"],
+    scope="function",
+)
+def noise_schedule(request):
+    return request.getfixturevalue(request.param)

tests/test_networks/test_diffusion_model/conftest.py

@@ -0,0 +1,25 @@
+def test_serialize_deserialize_noise_schedule(noise_schedule):
+    from bayesflow.utils.serialization import serialize, deserialize
+
+    serialized = serialize(noise_schedule)
+    deserialized = deserialize(serialized)
+    reserialized = serialize(deserialized)
+
+    assert serialized == reserialized
+    t = 0.251
+    x = 0.5
+    training = True
+    assert noise_schedule.get_log_snr(t, training=training) == deserialized.get_log_snr(t, training=training)
+    assert noise_schedule.get_t_from_log_snr(t, training=training) == deserialized.get_t_from_log_snr(
+        t, training=training
+    )
+    assert noise_schedule.derivative_log_snr(t, training=False) == deserialized.derivative_log_snr(t, training=False)
+    assert noise_schedule.get_drift_diffusion(t, x, training=False) == deserialized.get_drift_diffusion(
+        t, x, training=False
+    )
+    assert noise_schedule.get_alpha_sigma(t, training=training) == deserialized.get_alpha_sigma(t, training=training)
+    assert noise_schedule.get_weights_for_snr(t) == deserialized.get_weights_for_snr(t)
+
+
+def test_validate_noise_schedule(noise_schedule):
+    noise_schedule.validate()