ENH: add Zuko Flow Matching

src/sequince_density/flows/zuko.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import array_api_compat.torch as torch_api
 import torch
@@ -12,6 +12,9 @@
 from ..utils import torch_random_context
 from .base import Flow
 
+if TYPE_CHECKING:
+    from sequince.transforms import ParticleTransform
+
 logger = logging.getLogger(__name__)
 
 
@@ -24,9 +27,8 @@ def __init__(
         self,
         dims: int,
         flow_class: str | Callable = "MAF",
-        particle_transform=None,
-        seed: int = 1234,
-        device: str = "cpu",
+        particle_transform: ParticleTransform | None = None,
+        device: str | torch.DeviceObjType = "cpu",
         dtype: torch.dtype | None = None,
         **kwargs,
     ) -> None:
@@ -37,8 +39,6 @@ def __init__(
             particle_transform=particle_transform,
         )
         self.dtype = resolved_dtype
-        self.seed = seed
-        torch.manual_seed(seed)
 
         if isinstance(flow_class, str):
             flow_constructor = getattr(zuko.flows, flow_class)
@@ -221,3 +221,44 @@ def inverse(self, z, context=None):
             x_prime, context=context
         )
         return x, log_j_rescale + log_abs_det_jacobian
+
+
+class ZukoFlowMatching(ZukoFlow):
+    """Flow wrapper for training Zuko flows using flow matching.
+
+    Note that this flow is only compatible with CNF flows, as flow matching is
+    not implemented for discrete flows in Zuko.
+
+    The `eta` parameter controls the amount of noise added to the target
+    velocity field during training.
+    """
+
+    def __init__(
+        self,
+        dims: int,
+        particle_transform: ParticleTransform | None = None,
+        device: str | torch.DeviceObjType = "cpu",
+        eta: float = 1e-3,
+        dtype: torch.dtype | None = None,
+        **kwargs,
+    ):
+        kwargs.setdefault("hidden_features", 4 * [100])
+        super().__init__(
+            dims,
+            device=device,
+            particle_transform=particle_transform,
+            flow_class="CNF",
+            dtype=dtype,
+            **kwargs,
+        )
+        self.eta = eta
+
+    def loss_fn(self, theta: torch.Tensor) -> torch.Tensor:
+        t = torch.rand(
+            theta.shape[:-1], dtype=theta.dtype, device=theta.device
+        )
+        t_ = t[..., None]
+        eps = torch.randn_like(theta)
+        theta_prime = (1 - t_) * theta + (t_ + self.eta) * eps
+        v = eps - theta
+        return (self.flow.transform.f(t, theta_prime) - v).square().mean()

tests/conftest.py

@@ -0,0 +1,7 @@
+import numpy as np
+import pytest
+
+
+@pytest.fixture()
+def rng() -> np.random.Generator:
+    return np.random.default_rng(1234)

tests/test_utils.py

@@ -3,8 +3,6 @@
 import pytest
 from orng import RandomGenerator
 
-torch = pytest.importorskip("torch")
-
 from sequince_density.utils import (  # noqa: E402
     draw_seed,
     torch_random_context,
@@ -26,6 +24,7 @@ def test_draw_seed_rejects_unknown_rng() -> None:
 
 
 def test_torch_random_context_is_reproducible_and_restores_state() -> None:
+    torch = pytest.importorskip("torch")
     initial_state = torch.random.get_rng_state().clone()
 
     with torch_random_context(123):

tests/test_zuko.py

@@ -8,14 +8,16 @@
     FlowTransform,
 )
 
-pytest.importorskip("torch")
+torch = pytest.importorskip("torch")
 pytest.importorskip("zuko")
 
-from sequince_density.flows.zuko import ZukoFlow  # noqa: E402
+from sequince_density.flows.zuko import (  # noqa: E402
+    ZukoFlow,
+    ZukoFlowMatching,
+)
 
 
-def test_zuko_flow_transform_fits_and_round_trips() -> None:
-    rng = np.random.default_rng(1234)
+def test_zuko_flow_transform_fits_and_round_trips(rng) -> None:
     particles = rng.normal(size=(64, 2))
     flow = ZukoFlow(
         dims=2,
@@ -45,3 +47,41 @@ 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)
+
+
+def test_zuko_flow_matching_fits_and_round_trips(rng) -> None:
+    particles = rng.normal(size=(64, 2))
+    flow = ZukoFlowMatching(
+        dims=2,
+        hidden_features=[8, 8],
+        device="cpu",
+        dtype=torch.float64,
+        atol=1e-8,
+        rtol=1e-8,
+        exact=True,
+    )
+
+    transform = CompositeTransform(
+        [
+            AffineTransform(),
+            FlowTransform(
+                flow,
+                fit_kwargs={
+                    "n_epochs": 2,
+                    "batch_size": 16,
+                    "validation_fraction": 0.25,
+                    "rng": 1234,
+                },
+            ),
+        ]
+    )
+    transform = transform.fit(particles)
+    latent, log_det = transform.forward(particles)
+    restored, inverse_log_det = transform.inverse(latent)
+
+    assert isinstance(latent, np.ndarray)
+    assert latent.shape == particles.shape
+    assert log_det.shape == (particles.shape[0],)
+    np.testing.assert_allclose(restored, particles, atol=1e-5, rtol=1e-5)
+    residual = log_det + inverse_log_det
+    assert np.max(np.abs(residual)) < 1e-4