Add CirculantNormal distribution. (#1988)

* Add `CirculantNormal` distribution.

* Fix Jacobian of `RealFastFourierTransform`.

* Add transform to pack and unpack Fourier coefficients.

* Implement `CirculantNormal` using transformations.

* Fix typo in skip reason for `test_generated_sample_distribution`.

* Add `CirculantNormal` example notebook.

* Reformat `circulant_gp` example notebook.

* Fix keyword argument for `PackRealFastFourierCoefficientsTransform` test.

* Fix batch shape for mean, covariance_row, covariance_matrix.

* Add analytic KL divergence of diagonal `Normal` from `CirculantNormal`.

* Fix sign bug in `_PositiveDefiniteCirculantVector` validity check.

* Expand docstring and add references for `Circulant` distribution.

* Fix shape promotion for `Circulant` distribution args.

* Skip goodnes of fit test for `CirculantNormal`.

* Expand docstring of `CirculantNormal`.

* Update example and move legend.

---------

Co-authored-by: Du Phan <[email protected]>

Author: tillahoffmann

docs/source/distributions.rst

@@ -136,6 +136,14 @@ Chi2
     :show-inheritance:
     :member-order: bysource
 
+CirculantNormal
+^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.continuous.CirculantNormal
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 Dirichlet
 ^^^^^^^^^
 .. autoclass:: numpyro.distributions.continuous.Dirichlet
@@ -998,6 +1006,14 @@ OrderedTransform
     :show-inheritance:
     :member-order: bysource
 
+PackRealFastFourierCoefficientsTransform
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.transforms.PackRealFastFourierCoefficientsTransform
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 PermuteTransform
 ^^^^^^^^^^^^^^^^
 .. autoclass:: numpyro.distributions.transforms.PermuteTransform

docs/source/index.rst

@@ -39,6 +39,7 @@ NumPyro documentation
    tutorials/censoring
    tutorials/hsgp_example
    tutorials/other_samplers
+   tutorials/circulant_gp
    tutorials/nnx_example
 
 .. nbgallery::

notebooks/source/circulant_gp.ipynb

@@ -19,6 +19,7 @@
     BetaProportion,
     Cauchy,
     Chi2,
+    CirculantNormal,
     Dirichlet,
     EulerMaruyama,
     Exponential,
@@ -132,6 +133,7 @@
     "CategoricalProbs",
     "Cauchy",
     "Chi2",
+    "CirculantNormal",
     "Delta",
     "Dirichlet",
     "DirichletMultinomial",

numpyro/distributions/__init__.py

@@ -47,6 +47,7 @@
     "nonnegative_integer",
     "positive",
     "positive_definite",
+    "positive_definite_circulant_vector",
     "positive_semidefinite",
     "positive_integer",
     "real",
@@ -642,6 +643,19 @@ def feasible_like(self, prototype):
         )
 
 
+class _PositiveDefiniteCirculantVector(_SingletonConstraint):
+    event_dim = 1
+
+    def __call__(self, x):
+        jnp = np if isinstance(x, (np.ndarray, np.generic)) else jax.numpy
+        tol = 10 * jnp.finfo(x.dtype).eps
+        rfft = jnp.fft.rfft(x)
+        return (jnp.abs(rfft.imag) < tol) & (rfft.real > -tol)
+
+    def feasible_like(self, prototype):
+        return jnp.zeros_like(prototype).at[..., 0].set(1.0)
+
+
 class _PositiveSemiDefinite(_SingletonConstraint):
     event_dim = 2
 
@@ -792,6 +806,7 @@ def tree_flatten(self):
 ordered_vector = _OrderedVector()
 positive = _Positive()
 positive_definite = _PositiveDefinite()
+positive_definite_circulant_vector = _PositiveDefiniteCirculantVector()
 positive_semidefinite = _PositiveSemiDefinite()
 positive_integer = _IntegerPositive()
 positive_ordered_vector = _PositiveOrderedVector()

numpyro/distributions/constraints.py

@@ -33,7 +33,7 @@
 import jax.nn as nn
 import jax.numpy as jnp
 import jax.random as random
-from jax.scipy.linalg import cho_solve, solve_triangular
+from jax.scipy.linalg import cho_solve, solve_triangular, toeplitz
 from jax.scipy.special import (
     betaln,
     digamma,
@@ -59,12 +59,15 @@
     CholeskyTransform,
     CorrMatrixCholeskyTransform,
     ExpTransform,
+    PackRealFastFourierCoefficientsTransform,
     PowerTransform,
+    RealFastFourierTransform,
     RecursiveLinearTransform,
     SigmoidTransform,
     ZeroSumTransform,
 )
 from numpyro.distributions.util import (
+    _reshape,
     add_diag,
     assert_one_of,
     betainc,
@@ -3068,3 +3071,144 @@ def entropy(self) -> ArrayLike:
         return jnp.broadcast_to(
             0.5 + 1.5 * jnp.euler_gamma + 0.5 * jnp.log(16 * jnp.pi), self.batch_shape
         ) + jnp.log(self.scale)
+
+
+class CirculantNormal(TransformedDistribution):
+    r"""
+    Multivariate normal distribution with covariance matrix :math:`\mathbf{C}` that is
+    positive-definite and circulant [1], i.e., has periodic boundary conditions. The
+    density of a sample :math:`\mathbf{x}\in\mathbb{R}^n` is the standard multivariate
+    normal density
+
+    .. math::
+
+        p\left(\mathbf{x}\mid\boldsymbol{\mu},\mathbf{C}\right) =
+        \frac{\left(\mathrm{det}\,\mathbf{C}\right)^{-1/2}}{\left(2\pi\right)^{n / 2}}
+        \exp\left(-\frac{1}{2}\left(\mathbf{x}-\boldsymbol{\mu}\right)^\intercal
+        \mathbf{C}^{-1}\left(\mathbf{x}-\boldsymbol{\mu}\right)\right),
+
+    where :math:`\mathrm{det}` denotes the determinant and :math:`^\intercal` the
+    transpose. Circulant matrices can be diagnolized efficiently using the discrete
+    Fourier transform [1], allowing the log likelihood to be evaluated in
+    :math:`n \log n` time for :math:`n` observations [2].
+
+    :param loc: Mean of the distribution :math:`\boldsymbol{\mu}`.
+    :param covariance_row: First row of the circulant covariance matrix
+        :math:`\boldsymbol{C}`. Because of periodic boundary conditions, the covariance
+        matrix is fully determined by its first row (see
+        :func:`jax.scipy.linalg.toeplitz` for further details).
+    :param covariance_rfft: Real part of the real fast Fourier transform of
+        :code:`covariance_row`, the first row of the circulant covariance matrix
+        :math:`\boldsymbol{C}`.
+
+    **References:**
+
+    1. Wikipedia. (n.d.). Circulant matrix. Retrieved March 6, 2025, from
+       https://en.wikipedia.org/wiki/Circulant_matrix
+    2. Wood, A. T. A., & Chan, G. (1994). Simulation of Stationary Gaussian Processes in
+       :math:`\left[0, 1\right]^d`. *Journal of Computational and Graphical Statistics*,
+       3(4), 409--432. https://doi.org/10.1080/10618600.1994.10474655
+    """
+
+    arg_constraints = {
+        "loc": constraints.real_vector,
+        "covariance_row": constraints.positive_definite_circulant_vector,
+        "covariance_rfft": constraints.independent(constraints.positive, 1),
+    }
+    support = constraints.real_vector
+
+    def __init__(
+        self,
+        loc: jnp.ndarray,
+        covariance_row: jnp.ndarray = None,
+        covariance_rfft: jnp.ndarray = None,
+        *,
+        validate_args=None,
+    ) -> None:
+        # We demand a one-dimensional input, because we cannot determine the event shape
+        # if only the `covariance_rfft` is given.
+        assert jnp.ndim(loc) > 0, "Location parameter must have at least one dimension."
+        n = jnp.shape(loc)[-1]
+        n_rfft = n // 2 + 1
+        assert_one_of(covariance_row=covariance_row, covariance_rfft=covariance_rfft)
+
+        if covariance_rfft is None:
+            # Evaluate `covariance_rfft` if not provided and validate.
+            assert covariance_row.shape[-1] == n
+            loc, covariance_row = promote_shapes(loc, covariance_row)
+            covariance_rfft = jnp.fft.rfft(covariance_row).real
+            self.covariance_row = covariance_row
+        else:
+            # The `covariance_rfft` and `loc` are not promotable because the trailing
+            # dimension does not match. We manually retrieve the shapes and then
+            # promote.
+            loc_shape, covariance_rfft_shape = promote_shapes(
+                loc[..., 0], covariance_rfft[..., 0], return_shapes=True
+            )
+            loc = _reshape(loc, loc_shape + (n,))
+            covariance_rfft = _reshape(
+                covariance_rfft, covariance_rfft_shape + (n_rfft,)
+            )
+
+        self.loc = loc
+        self.covariance_rfft = covariance_rfft
+
+        # Construct the base distribution.
+        n_imag = n - n_rfft
+        assert self.covariance_rfft.shape[-1] == n_rfft
+        var_rfft = (n * covariance_rfft / 2).at[..., 0].mul(2)
+        if n % 2 == 0:
+            var_rfft = var_rfft.at[..., -1].mul(2)
+        var_rfft = jnp.concatenate([var_rfft, var_rfft[..., 1 : 1 + n_imag]], axis=-1)
+        assert var_rfft.shape[-1] == n
+        base_distribution = Normal(scale=jnp.sqrt(var_rfft)).to_event(1)
+
+        super().__init__(
+            base_distribution,
+            [
+                PackRealFastFourierCoefficientsTransform((n,)),
+                RealFastFourierTransform((n,)).inv,
+                AffineTransform(loc, scale=1.0),
+            ],
+            validate_args=validate_args,
+        )
+
+    @property
+    def mean(self) -> jnp.ndarray:
+        return jnp.broadcast_to(self.loc, self.shape())
+
+    @lazy_property
+    def covariance_row(self) -> jnp.ndarray:
+        return jnp.fft.irfft(self.covariance_rfft, n=self.event_shape[-1])
+
+    @lazy_property
+    def covariance_matrix(self) -> jnp.ndarray:
+        *leading_shape, n = self.covariance_row.shape
+        if leading_shape:
+            # `toeplitz` flattens the input, and we need to broadcast manually.
+            (n,) = self.event_shape
+            return vmap(toeplitz)(self.covariance_row.reshape((-1, n))).reshape(
+                (*leading_shape, n, n)
+            )
+        else:
+            return toeplitz(self.covariance_row)
+
+    @lazy_property
+    def variance(self) -> jnp.ndarray:
+        return jnp.broadcast_to(self.covariance_row[..., 0, None], self.shape())
+
+    @staticmethod
+    def infer_shapes(
+        loc: tuple = (), covariance_row: tuple = None, covariance_rfft: tuple = None
+    ):
+        assert_one_of(covariance_row=covariance_row, covariance_rfft=covariance_rfft)
+        for cov in [covariance_rfft, covariance_row]:
+            if cov is not None:
+                batch_shape = jnp.broadcast_shapes(loc[:-1], cov[:-1])
+                event_shape = loc[-1:]
+                return batch_shape, event_shape
+
+    def entropy(self):
+        (n,) = self.event_shape
+        log_abs_det_jacobian = 2 * jnp.log(2) * ((n - 1) // 2) - jnp.log(n) * n
+        return self.base_dist.entropy() + log_abs_det_jacobian / 2

numpyro/distributions/continuous.py

@@ -33,6 +33,7 @@
 
 from numpyro.distributions.continuous import (
     Beta,
+    CirculantNormal,
     Dirichlet,
     Gamma,
     Kumaraswamy,
@@ -183,6 +184,27 @@ def _shapes_are_broadcastable(first_shape, second_shape):
     return 0.5 * (tr + t1 - D - log_det_ratio)
 
 
+@dispatch(Independent, CirculantNormal)
+def kl_divergence(p: Independent, q: CirculantNormal):
+    # We can only calculate the KL divergence if the base distribution is normal.
+    if not isinstance(p.base_dist, Normal) or p.reinterpreted_batch_ndims != 1:
+        raise NotImplementedError
+
+    residual = q.mean - p.mean
+    n = residual.shape[-1]
+    log_covariance_rfft = jnp.log(q.covariance_rfft)
+    return (
+        jnp.vecdot(
+            residual, jnp.fft.irfft(jnp.fft.rfft(residual) / q.covariance_rfft, n)
+        )
+        + jnp.fft.irfft(1 / q.covariance_rfft, n)[..., 0] * jnp.sum(p.variance, axis=-1)
+        + log_covariance_rfft.sum(axis=-1)
+        + log_covariance_rfft[..., 1 : (n + 1) // 2].sum(axis=-1)
+        - jnp.log(p.variance).sum(axis=-1)
+        - n
+    ) / 2
+
+
 @dispatch(Beta, Beta)
 def kl_divergence(p, q):
     # From https://en.wikipedia.org/wiki/Beta_distribution#Quantities_of_information_(entropy)

numpyro/distributions/kl.py

@@ -40,6 +40,7 @@
     "LowerCholeskyTransform",
     "ScaledUnitLowerCholeskyTransform",
     "LowerCholeskyAffine",
+    "PackRealFastFourierCoefficientsTransform",
     "PermuteTransform",
     "PowerTransform",
     "RealFastFourierTransform",
@@ -1311,10 +1312,15 @@ def inverse_shape(self, shape: tuple) -> tuple:
     def log_abs_det_jacobian(
         self, x: jnp.ndarray, y: jnp.ndarray, intermediates: None = None
     ) -> jnp.ndarray:
-        shape = jnp.broadcast_shapes(
+        batch_shape = jnp.broadcast_shapes(
             x.shape[: -self.transform_ndims], y.shape[: -self.transform_ndims]
         )
-        return jnp.zeros_like(x, shape=shape)
+        event_shape = x.shape[-self.transform_ndims :]
+        size = math.prod(event_shape)
+        q = math.prod(2 - size % 2 for size in event_shape)
+        return jnp.broadcast_to(
+            (size * jnp.log(size) - jnp.log(2) * (size - q)) / 2, batch_shape
+        )
 
     def tree_flatten(self):
         aux_data = {
@@ -1339,6 +1345,74 @@ def __eq__(self, other):
         )
 
 
+class PackRealFastFourierCoefficientsTransform(Transform):
+    """
+    Transform a real vector to complex coefficients of a real fast Fourier transform.
+
+    :param transform_shape: Shape of the real vector, defaults to the input size.
+    """
+
+    domain = constraints.real_vector
+    codomain = constraints.independent(constraints.complex, 1)
+
+    def __init__(self, transform_shape: tuple = None) -> None:
+        assert transform_shape is None or len(transform_shape) == 1, (
+            "Packing Fourier coefficients is only implemented for vectors."
+        )
+        self.shape = transform_shape
+
+    def tree_flatten(self):
+        return (), ((), {"shape": self.shape})
+
+    def forward_shape(self, shape: tuple) -> tuple:
+        *batch_shape, n = shape
+        assert self.shape is None or self.shape == (n,), (
+            f"`shape` must be `None` or `{self.shape}. Got `{shape}`."
+        )
+        n_rfft = n // 2 + 1
+        return (*batch_shape, n_rfft)
+
+    def inverse_shape(self, shape: tuple) -> tuple:
+        *batch_shape, n_rfft = shape
+        assert self.shape is not None, (
+            "Shape must be specified in `__init__` for inverse transform."
+        )
+        (n,) = self.shape
+        assert n_rfft == n // 2 + 1
+        return (*batch_shape, n)
+
+    def log_abs_det_jacobian(
+        self, x: jnp.ndarray, y: jnp.ndarray, intermediates: None = None
+    ) -> jnp.ndarray:
+        shape = jnp.broadcast_shapes(x.shape[:-1], y.shape[:-1])
+        return jnp.zeros_like(x, shape=shape)
+
+    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
+        assert self.shape is None or self.shape == x.shape[-1:]
+        n = x.shape[-1]
+        n_real = n // 2 + 1
+        n_imag = n - n_real
+        complex_dtype = jnp.result_type(x.dtype, jnp.complex64)
+        return (
+            x[..., :n_real]
+            .astype(complex_dtype)
+            .at[..., 1 : 1 + n_imag]
+            .add(1j * x[..., n_real:])
+        )
+
+    def _inverse(self, y: jnp.ndarray) -> jnp.ndarray:
+        (n,) = self.shape
+        n_real = n // 2 + 1
+        n_imag = n - n_real
+        return jnp.concatenate([y.real, y.imag[..., 1 : n_imag + 1]], axis=-1)
+
+    def __eq__(self, other) -> bool:
+        return (
+            isinstance(other, PackRealFastFourierCoefficientsTransform)
+            and self.shape == other.shape
+        )
+
+
 class RecursiveLinearTransform(Transform):
     """
     Apply a linear transformation recursively such that