TST(flows): cover density consistency and checkpoint restoration

tests/test_flowjax.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import pickle
+from pathlib import Path
+
 import numpy as np
 import pytest
 from orng import RandomGenerator
@@ -13,6 +16,27 @@
 from sequince_density.flows.flowjax import FlowJax  # noqa: E402
 
 
+@pytest.fixture
+def fitted_flow() -> FlowJax:
+    particles = np.random.default_rng(21).normal(size=(64, 2))
+    particles = particles * [2.0, 3.0] + [4.0, -2.0]
+    flow = FlowJax(
+        dims=2,
+        particle_transform=AffineTransform(),
+        seed=22,
+        flow_layers=1,
+        nn_width=8,
+        nn_depth=1,
+    )
+    return flow.fit(
+        particles,
+        rng=23,
+        max_epochs=1,
+        batch_size=16,
+        show_progress=False,
+    )
+
+
 def test_flowjax_fits_samples_and_round_trips() -> None:
     key = jax.random.key(12)
     particles = jax.random.normal(key, (64, 2))
@@ -48,3 +72,48 @@ def test_flowjax_fits_samples_and_round_trips() -> None:
         0.0,
         atol=1e-5,
     )
+
+
+@pytest.mark.parametrize("n_samples", [1, 8])
+def test_sample_and_log_prob_matches_log_prob(
+    fitted_flow: FlowJax, n_samples: int
+) -> None:
+    samples, log_prob = fitted_flow.sample_and_log_prob(n_samples, rng=24)
+
+    assert samples.shape == (n_samples, 2)
+    assert log_prob.shape == (n_samples,)
+    assert np.all(np.isfinite(samples))
+    assert np.all(np.isfinite(log_prob))
+    # Inverting the float32 flow introduces round-off in density evaluation.
+    np.testing.assert_allclose(
+        log_prob, fitted_flow.log_prob(samples), rtol=1e-5, atol=1e-5
+    )
+
+
+def test_checkpoint_restores_fitted_density_and_rng_sequence(
+    fitted_flow: FlowJax, tmp_path: Path
+) -> None:
+    # Save an advanced internal key so restoring constructor state is insufficient.
+    fitted_flow.sample(3)
+    checkpoint = tmp_path / "flow.pkl"
+    checkpoint.write_bytes(pickle.dumps(fitted_flow.state_dict()))
+    restored = FlowJax(
+        dims=2,
+        seed=99,
+        flow_layers=1,
+        nn_width=8,
+        nn_depth=1,
+    )
+    restored.load_state_dict(pickle.loads(checkpoint.read_bytes()))
+
+    points = jnp.array([[4.0, -2.0], [6.0, 1.0]], dtype=jnp.float32)
+    np.testing.assert_array_equal(
+        restored.log_prob(points), fitted_flow.log_prob(points)
+    )
+    for _ in range(2):
+        np.testing.assert_array_equal(
+            restored.sample(8), fitted_flow.sample(8)
+        )
+    assert fitted_flow.fit_history["training_loss"]
+    assert fitted_flow.fit_history["validation_loss"]
+    assert restored.fit_history == fitted_flow.fit_history

tests/test_zuko.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import pickle
+from pathlib import Path
+
 import numpy as np
 import pytest
 from orng import RandomGenerator
@@ -9,12 +12,27 @@
     FlowTransform,
 )
 
-pytest.importorskip("torch")
+torch = pytest.importorskip("torch")
 pytest.importorskip("zuko")
 
 from sequince_density.flows.zuko import ZukoFlow  # noqa: E402
 
 
+@pytest.fixture
+def fitted_flow() -> ZukoFlow:
+    particles = np.random.default_rng(21).normal(size=(64, 2))
+    particles = particles * [2.0, 3.0] + [4.0, -2.0]
+    flow = ZukoFlow(
+        dims=2,
+        hidden_features=[8, 8],
+        particle_transform=AffineTransform(),
+        seed=22,
+        device="cpu",
+        dtype=torch.float32,
+    )
+    return flow.fit(particles, n_epochs=2, batch_size=16, rng=23)
+
+
 def test_zuko_flow_transform_fits_and_round_trips() -> None:
     rng = np.random.default_rng(1234)
     particles = rng.normal(size=(64, 2))
@@ -51,3 +69,49 @@ def test_zuko_flow_transform_fits_and_round_trips() -> None:
     assert log_det.shape == (particles.shape[0],)
     np.testing.assert_allclose(restored, particles, atol=1e-5)
     np.testing.assert_allclose(log_det + inverse_log_det, 0.0, atol=1e-5)
+
+
+@pytest.mark.parametrize("n_samples", [1, 8])
+def test_sample_and_log_prob_matches_log_prob(
+    fitted_flow: ZukoFlow, n_samples: int
+) -> None:
+    samples, log_prob = fitted_flow.sample_and_log_prob(n_samples, rng=24)
+
+    assert samples.shape == (n_samples, 2)
+    assert log_prob.shape == (n_samples,)
+    assert torch.isfinite(samples).all()
+    assert torch.isfinite(log_prob).all()
+    # Inverting the float32 flow introduces round-off in density evaluation.
+    torch.testing.assert_close(
+        log_prob, fitted_flow.log_prob(samples), rtol=1e-5, atol=1e-5
+    )
+
+
+def test_checkpoint_restores_fitted_density_and_seeded_samples(
+    fitted_flow: ZukoFlow, tmp_path: Path
+) -> None:
+    checkpoint = tmp_path / "flow.pkl"
+    checkpoint.write_bytes(pickle.dumps(fitted_flow.state_dict()))
+    restored = ZukoFlow(
+        dims=2,
+        hidden_features=[8, 8],
+        seed=99,
+        device="cpu",
+        dtype=torch.float32,
+    )
+    restored.load_state_dict(pickle.loads(checkpoint.read_bytes()))
+
+    points = torch.tensor([[4.0, -2.0], [6.0, 1.0]], dtype=torch.float32)
+    torch.testing.assert_close(
+        restored.log_prob(points), fitted_flow.log_prob(points), rtol=0, atol=0
+    )
+    for seed in (25, 26):
+        torch.testing.assert_close(
+            restored.sample(8, rng=seed),
+            fitted_flow.sample(8, rng=seed),
+            rtol=0,
+            atol=0,
+        )
+    assert fitted_flow.fit_history["training_loss"]
+    assert fitted_flow.fit_history["validation_loss"]
+    assert restored.fit_history == fitted_flow.fit_history