diff --git a/src/sequince_density/flows/zuko.py b/src/sequince_density/flows/zuko.py index 458704d..a1eb272 100644 --- a/src/sequince_density/flows/zuko.py +++ b/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() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..e3f43ec --- /dev/null +++ b/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) diff --git a/tests/test_utils.py b/tests/test_utils.py index a63a653..7fbd94a 100644 --- a/tests/test_utils.py +++ b/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): diff --git a/tests/test_zuko.py b/tests/test_zuko.py index e5162cf..168ddf3 100644 --- a/tests/test_zuko.py +++ b/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