couplnorm 0.1.0__py3-none-any.whl

couplnorm/__init__.py

@@ -0,0 +1,11 @@
+"""couplnorm: normalized 4th-order coupling metrics and losses for
+generative model evaluation on Fourier-decomposable data.
+"""
+from couplnorm.coupling import (
+    CouplingLoss,
+    CouplingMetric,
+    coupling_from_samples,
+)
+
+__all__ = ["CouplingMetric", "CouplingLoss", "coupling_from_samples"]
+__version__ = "0.1.0"

couplnorm/coupling.py

@@ -0,0 +1,285 @@
+"""Normalized off-diagonal coupling diagnostic and loss for evaluating
+generative models on Fourier-decomposable data.
+"""
+from __future__ import annotations
+
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+__all__ = ["CouplingMetric", "CouplingLoss", "coupling_from_samples"]
+
+
+def _spectral_energies(
+    x: Tensor,
+    input_space: str = "position",
+    fft_norm: str = "ortho",
+    real_fft: bool = True,
+) -> Tensor:
+    if input_space == "position":
+        if real_fft:
+            x_hat = torch.fft.rfft(x, dim=-1, norm=fft_norm)
+        else:
+            x_hat = torch.fft.fft(x, dim=-1, norm=fft_norm)
+        return x_hat.real.pow(2) + x_hat.imag.pow(2)
+    elif input_space == "spectral":
+        return x
+    else:
+        raise ValueError(
+            f"input_space must be 'position' or 'spectral', got {input_space!r}"
+        )
+
+
+def _batch_covariance(E: Tensor, unbiased: bool = True) -> Tensor:
+    if E.dim() != 2:
+        raise ValueError(f"Expected 2D tensor (B, M), got shape {tuple(E.shape)}")
+    B = E.shape[0]
+    if B < 2 and unbiased:
+        raise ValueError(
+            f"Need B >= 2 for unbiased covariance; got B={B}. "
+            "Set unbiased=False or use mode='running' for streaming."
+        )
+    mean = E.mean(dim=0, keepdim=True)
+    centered = E - mean
+    denom = (B - 1) if unbiased else B
+    return centered.T @ centered / denom
+
+
+def _coupling_from_cov(cov: Tensor, eps: float = 1e-12) -> Tensor:
+    diag_vec = torch.diagonal(cov)
+    off_diag = cov - torch.diag_embed(diag_vec)
+    num = torch.linalg.matrix_norm(off_diag, ord="fro")
+    den = torch.linalg.matrix_norm(cov, ord="fro")
+    return num / (den + eps)
+
+
+def coupling_from_samples(
+    x: Tensor,
+    input_space: str = "position",
+    fft_norm: str = "ortho",
+    real_fft: bool = True,
+    eps: float = 1e-12,
+    unbiased: bool = True,
+) -> Tensor:
+    if x.dim() == 1:
+        raise ValueError(
+            "coupling_from_samples requires batched input (B, N). "
+            "For streaming, use CouplingMetric(mode='running')."
+        )
+    E = _spectral_energies(x, input_space=input_space, fft_norm=fft_norm, real_fft=real_fft)
+    cov = _batch_covariance(E, unbiased=unbiased)
+    return _coupling_from_cov(cov, eps)
+
+
+class CouplingMetric(nn.Module):
+    def __init__(
+        self,
+        input_space: str = "position",
+        mode: str = "batch",
+        momentum: float = 0.1,
+        eps: float = 1e-12,
+        fft_norm: str = "ortho",
+        real_fft: bool = True,
+        unbiased: bool = True,
+        n_modes: Optional[int] = None,
+    ):
+        super().__init__()
+        if input_space not in ("position", "spectral"):
+            raise ValueError(
+                f"input_space must be 'position' or 'spectral'; got {input_space!r}"
+            )
+        if mode not in ("batch", "running"):
+            raise ValueError(f"mode must be 'batch' or 'running'; got {mode!r}")
+        if not (0.0 < momentum <= 1.0):
+            raise ValueError(f"momentum must be in (0, 1]; got {momentum}")
+
+        self.input_space = input_space
+        self.mode = mode
+        self.momentum = float(momentum)
+        self.eps = float(eps)
+        self.fft_norm = fft_norm
+        self.real_fft = bool(real_fft)
+        self.unbiased = bool(unbiased)
+
+        if mode == "running":
+            if n_modes is not None:
+                self.register_buffer("running_mean", torch.zeros(n_modes))
+                self.register_buffer("running_second", torch.zeros(n_modes, n_modes))
+            else:
+                self.register_buffer("running_mean", torch.empty(0))
+                self.register_buffer("running_second", torch.empty(0))
+            self.register_buffer("num_batches_tracked", torch.tensor(0, dtype=torch.long))
+
+    def _maybe_init_buffers(self, M: int, device, dtype) -> None:
+        if self.running_mean.numel() == 0:
+            self.running_mean = torch.zeros(M, device=device, dtype=dtype)
+            self.running_second = torch.zeros(M, M, device=device, dtype=dtype)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.dim() == 1:
+            x = x.unsqueeze(0)
+        if x.dim() != 2:
+            raise ValueError(f"Expected 1D or 2D input, got shape {tuple(x.shape)}")
+        E = _spectral_energies(x, self.input_space, self.fft_norm, self.real_fft)
+        B, M = E.shape
+
+        if self.mode == "batch":
+            cov = _batch_covariance(E, unbiased=self.unbiased)
+            return _coupling_from_cov(cov, self.eps)
+
+        self._maybe_init_buffers(M, E.device, E.dtype)
+        if self.running_mean.shape[0] != M:
+            raise ValueError(
+                "Number of modes changed: running buffers have "
+                f"N={self.running_mean.shape[0]} but got N={M}."
+            )
+
+        with torch.no_grad():
+            batch_mean = E.mean(dim=0)
+            batch_second = (E.T @ E) / B
+            m = self.momentum
+            self.running_mean.mul_(1.0 - m).add_(batch_mean, alpha=m)
+            self.running_second.mul_(1.0 - m).add_(batch_second, alpha=m)
+            self.num_batches_tracked.add_(1)
+        return self.compute()
+
+    def compute(self) -> Tensor:
+        if self.mode != "running":
+            raise RuntimeError("compute() is only valid in running mode.")
+        if self.running_mean.numel() == 0 or self.num_batches_tracked.item() == 0:
+            raise RuntimeError("No batches tracked yet; call forward() at least once.")
+        mu = self.running_mean
+        M = self.running_second
+        cov = M - torch.outer(mu, mu)
+        return _coupling_from_cov(cov, self.eps)
+
+    def covariance(self) -> Tensor:
+        if self.mode != "running":
+            raise RuntimeError(
+                "covariance() requires mode='running'. For one-off batches, "
+                "use coupling_from_samples or recompute manually."
+            )
+        mu = self.running_mean
+        M = self.running_second
+        return M - torch.outer(mu, mu)
+
+    def reset(self) -> None:
+        if self.mode != "running":
+            return
+        if self.running_mean.numel() > 0:
+            self.running_mean.zero_()
+            self.running_second.zero_()
+        self.num_batches_tracked.zero_()
+
+    def extra_repr(self) -> str:
+        return (
+            f"input_space={self.input_space!r}, mode={self.mode!r}, "
+            f"momentum={self.momentum}, eps={self.eps}, "
+            f"fft_norm={self.fft_norm!r}, real_fft={self.real_fft}"
+        )
+
+
+class CouplingLoss(nn.Module):
+    def __init__(
+        self,
+        target_type: str = "matrix",
+        target: Optional[Union[Tensor, float]] = None,
+        input_space: str = "position",
+        eps: float = 1e-12,
+        fft_norm: str = "ortho",
+        real_fft: bool = True,
+        unbiased: bool = True,
+    ):
+        super().__init__()
+        if target_type not in ("matrix", "scalar", "regularize"):
+            raise ValueError(
+                "target_type must be 'matrix', 'scalar', or 'regularize'; "
+                f"got {target_type!r}"
+            )
+        if input_space not in ("position", "spectral"):
+            raise ValueError(
+                f"input_space must be 'position' or 'spectral'; got {input_space!r}"
+            )
+
+        if target_type == "matrix":
+            if target is None:
+                raise ValueError(
+                    "target_type='matrix' requires a (M, M) target covariance."
+                )
+            tgt = torch.as_tensor(target, dtype=torch.float32)
+            if tgt.dim() != 2 or tgt.shape[0] != tgt.shape[1]:
+                raise ValueError(
+                    f"matrix target must be (M, M); got shape {tuple(tgt.shape)}"
+                )
+            self.register_buffer("target_cov", tgt.detach().clone())
+            self.register_buffer(
+                "target_cov_frob",
+                torch.linalg.matrix_norm(tgt, ord="fro").detach().clone(),
+            )
+        elif target_type == "scalar":
+            if target is None:
+                raise ValueError("target_type='scalar' requires a target C value.")
+            self.register_buffer(
+                "target_C", torch.as_tensor(float(target), dtype=torch.float32)
+            )
+
+        self.target_type = target_type
+        self.input_space = input_space
+        self.eps = float(eps)
+        self.fft_norm = fft_norm
+        self.real_fft = bool(real_fft)
+        self.unbiased = bool(unbiased)
+
+    @classmethod
+    def from_data(
+        cls,
+        data: Tensor,
+        target_type: str = "matrix",
+        input_space: str = "position",
+        fft_norm: str = "ortho",
+        real_fft: bool = True,
+        unbiased: bool = True,
+        eps: float = 1e-12,
+    ) -> "CouplingLoss":
+        if target_type not in ("matrix", "scalar"):
+            raise ValueError(f"from_data does not support target_type={target_type!r}")
+        with torch.no_grad():
+            E = _spectral_energies(
+                data, input_space=input_space, fft_norm=fft_norm, real_fft=real_fft
+            )
+            cov = _batch_covariance(E, unbiased=unbiased)
+            if target_type == "matrix":
+                return cls(
+                    target_type="matrix", target=cov, input_space=input_space, eps=eps,
+                    fft_norm=fft_norm, real_fft=real_fft, unbiased=unbiased,
+                )
+            else:
+                C = _coupling_from_cov(cov, eps).item()
+                return cls(
+                    target_type="scalar", target=C, input_space=input_space, eps=eps,
+                    fft_norm=fft_norm, real_fft=real_fft, unbiased=unbiased,
+                )
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.dim() == 1:
+            x = x.unsqueeze(0)
+        E = _spectral_energies(x, self.input_space, self.fft_norm, self.real_fft)
+        cov = _batch_covariance(E, unbiased=self.unbiased)
+        if self.target_type == "matrix":
+            diff = cov - self.target_cov
+            num = torch.linalg.matrix_norm(diff, ord="fro").pow(2)
+            den = self.target_cov_frob.pow(2) + self.eps
+            return num / den
+        elif self.target_type == "scalar":
+            C = _coupling_from_cov(cov, self.eps)
+            return (C - self.target_C).pow(2)
+        else:
+            return _coupling_from_cov(cov, self.eps)
+
+    def extra_repr(self) -> str:
+        return (
+            f"target_type={self.target_type!r}, "
+            f"input_space={self.input_space!r}, fft_norm={self.fft_norm!r}"
+        )

couplnorm/models/__init__.py

@@ -0,0 +1,25 @@
+"""Baseline generative models for couplnorm demos and benchmarks.
+
+Every model is an ``nn.Module`` with a common interface:
+
+- ``fit(data)`` estimates parameters from a ``(B, N)`` batch of position-space
+  field configurations (a no-op / trainer depending on the model).
+- ``sample(n)`` returns ``(n, N)`` real position-space configurations.
+
+The models span the spectrum of joint-coupling fidelity:
+
+- :class:`FourierModel` assumes marginal independence of Fourier modes, so its
+  samples have C near zero regardless of the target. It is the "wrong by
+  construction" baseline that motivates the C diagnostic.
+- :class:`FullGaussian` matches the full second-order (covariance) structure of
+  the data but, being Gaussian, cannot reproduce genuine 4th-order coupling.
+- :class:`PCAModel` is a low-rank Gaussian; even weaker than FullGaussian.
+- :class:`MAF` is a masked autoregressive flow that can, in principle, capture
+  higher-order structure and drive C toward the target.
+"""
+from couplnorm.models.fourier import FourierModel
+from couplnorm.models.gaussian import FullGaussian
+from couplnorm.models.maf import MAF
+from couplnorm.models.pca import PCAModel
+
+__all__ = ["FourierModel", "FullGaussian", "PCAModel", "MAF"]

couplnorm/models/fourier.py

@@ -0,0 +1,73 @@
+"""Independent-mode Fourier baseline (the marginal-independence model)."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+__all__ = ["FourierModel"]
+
+
+class FourierModel(nn.Module):
+    """Models each Fourier mode as an independent Gaussian.
+
+    Fits a per-mode complex Gaussian to ``rfft(data)`` (independent real and
+    imaginary parts, no cross-mode correlation) and samples by drawing each mode
+    independently and inverting the transform. Because the modes are independent
+    by construction, the spectral-energy covariance is diagonal and C is near
+    zero regardless of the target distribution.
+
+    This is the "wrong by construction" baseline: a generative model that
+    assumes marginal independence of modes cannot reproduce joint 4th-order
+    coupling. It is the whole reason the C diagnostic is interesting.
+
+    Parameters
+    ----------
+    N : int
+        Field dimension (length of the real position-space signal).
+    """
+
+    def __init__(self, N: int):
+        super().__init__()
+        self.N = int(N)
+        self.M = N // 2 + 1
+        self.register_buffer("mean_real", torch.zeros(self.M))
+        self.register_buffer("mean_imag", torch.zeros(self.M))
+        self.register_buffer("std_real", torch.ones(self.M))
+        self.register_buffer("std_imag", torch.ones(self.M))
+        self._fitted = False
+
+    @torch.no_grad()
+    def fit(self, data: Tensor) -> "FourierModel":
+        if data.dim() != 2 or data.shape[1] != self.N:
+            raise ValueError(
+                f"Expected data of shape (B, {self.N}); got {tuple(data.shape)}"
+            )
+        data = data.to(self.mean_real)
+        x_hat = torch.fft.rfft(data, dim=-1, norm="ortho")
+        self.mean_real = x_hat.real.mean(dim=0)
+        self.mean_imag = x_hat.imag.mean(dim=0)
+        self.std_real = x_hat.real.std(dim=0).clamp_min(1e-8)
+        self.std_imag = x_hat.imag.std(dim=0).clamp_min(1e-8)
+        # The DC mode (and Nyquist mode when N is even) is purely real for a
+        # real signal; force its imaginary part to zero so irfft stays real.
+        self.mean_imag[0] = 0.0
+        self.std_imag[0] = 0.0
+        if self.N % 2 == 0:
+            self.mean_imag[-1] = 0.0
+            self.std_imag[-1] = 0.0
+        self._fitted = True
+        return self
+
+    @torch.no_grad()
+    def sample(self, n: int) -> Tensor:
+        if not self._fitted:
+            raise RuntimeError("Call fit(data) before sample(n).")
+        dev, dt = self.mean_real.device, self.mean_real.dtype
+        re = self.mean_real + self.std_real * torch.randn(n, self.M, device=dev, dtype=dt)
+        im = self.mean_imag + self.std_imag * torch.randn(n, self.M, device=dev, dtype=dt)
+        x_hat = torch.complex(re, im)
+        return torch.fft.irfft(x_hat, n=self.N, dim=-1, norm="ortho")
+
+    def extra_repr(self) -> str:
+        return f"N={self.N}, M={self.M}, fitted={self._fitted}"

couplnorm/models/gaussian.py

@@ -0,0 +1,61 @@
+"""Full multivariate Gaussian baseline."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+__all__ = ["FullGaussian"]
+
+
+class FullGaussian(nn.Module):
+    """Position-space multivariate Gaussian fit to the data covariance.
+
+    Matches the mean and full covariance of the training fields exactly (up to
+    finite-sample error), then samples via the Cholesky factor. Being Gaussian,
+    it reproduces all second-order structure but no genuine 4th-order coupling:
+    its spectral energies inherit only the coupling implied by the covariance.
+
+    Parameters
+    ----------
+    N : int
+        Number of lattice sites (field dimension).
+    jitter : float
+        Diagonal loading added before the Cholesky factorization for numerical
+        stability with near-singular covariances.
+    """
+
+    def __init__(self, N: int, jitter: float = 1e-5):
+        super().__init__()
+        self.N = int(N)
+        self.jitter = float(jitter)
+        self.register_buffer("mean", torch.zeros(N))
+        self.register_buffer("cov", torch.eye(N))
+        self.register_buffer("chol", torch.eye(N))
+        self._fitted = False
+
+    @torch.no_grad()
+    def fit(self, data: Tensor) -> "FullGaussian":
+        if data.dim() != 2 or data.shape[1] != self.N:
+            raise ValueError(
+                f"Expected data of shape (B, {self.N}); got {tuple(data.shape)}"
+            )
+        data = data.to(self.mean)
+        self.mean = data.mean(dim=0)
+        centered = data - self.mean
+        cov = centered.T @ centered / (data.shape[0] - 1)
+        cov = cov + self.jitter * torch.eye(self.N, device=cov.device, dtype=cov.dtype)
+        self.cov = cov
+        self.chol = torch.linalg.cholesky(cov)
+        self._fitted = True
+        return self
+
+    @torch.no_grad()
+    def sample(self, n: int) -> Tensor:
+        if not self._fitted:
+            raise RuntimeError("Call fit(data) before sample(n).")
+        z = torch.randn(n, self.N, device=self.mean.device, dtype=self.mean.dtype)
+        return self.mean + z @ self.chol.T
+
+    def extra_repr(self) -> str:
+        return f"N={self.N}, fitted={self._fitted}"

couplnorm/models/maf.py

@@ -0,0 +1,171 @@
+"""Masked Autoregressive Flow (MAF) baseline.
+
+A compact MADE-based MAF: a stack of autoregressive Gaussian transforms with a
+fixed permutation between blocks and a standard-normal base density. Unlike the
+Gaussian and Fourier baselines, a MAF can represent higher-order structure, so
+it is the model you would actually train to drive C toward a target.
+
+Kept intentionally small (this is plumbing, not the contribution). Trained by
+maximum likelihood; sampled by inverting each block sequentially.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+__all__ = ["MAF"]
+
+
+class _MaskedLinear(nn.Linear):
+    """Linear layer with a fixed binary mask on the weight matrix."""
+
+    def __init__(self, in_features: int, out_features: int):
+        super().__init__(in_features, out_features)
+        self.register_buffer("mask", torch.ones(out_features, in_features))
+
+    def set_mask(self, mask: Tensor) -> None:
+        self.mask.copy_(mask)
+
+    def forward(self, x: Tensor) -> Tensor:  # type: ignore[override]
+        return F.linear(x, self.weight * self.mask, self.bias)
+
+
+class _MADE(nn.Module):
+    """Masked autoencoder producing (mu, alpha) with the autoregressive property.
+
+    Output dimension ``d`` depends only on inputs ``0..d-1`` (in the current
+    variable order), so ``mu_d`` and ``alpha_d`` are functions of ``x_{<d}``.
+    """
+
+    def __init__(self, D: int, hidden: int = 64, n_hidden: int = 1):
+        super().__init__()
+        self.D = D
+        sizes = [D] + [hidden] * n_hidden + [2 * D]
+        self.net = nn.ModuleList()
+        for i in range(len(sizes) - 1):
+            self.net.append(_MaskedLinear(sizes[i], sizes[i + 1]))
+
+        # Degree assignment (deterministic): inputs 0..D-1, hidden units cycle
+        # through 0..D-2, outputs 0..D-1.
+        m_in = torch.arange(D)
+        degrees = [m_in]
+        for _ in range(n_hidden):
+            degrees.append(torch.arange(hidden) % max(D - 1, 1))
+        m_out = torch.arange(D)
+
+        # Hidden connectivity: non-strict >=.
+        for i, layer in enumerate(self.net[:-1]):
+            prev, cur = degrees[i], degrees[i + 1]
+            mask = (cur.unsqueeze(1) >= prev.unsqueeze(0)).float()
+            layer.set_mask(mask)
+        # Output connectivity: strict > (no self/future leakage), duplicated for
+        # the mu and alpha halves.
+        out_mask = (m_out.unsqueeze(1) > degrees[-1].unsqueeze(0)).float()
+        self.net[-1].set_mask(torch.cat([out_mask, out_mask], dim=0))
+
+    def forward(self, x: Tensor):
+        h = x
+        for layer in self.net[:-1]:
+            h = F.relu(layer(h))
+        out = self.net[-1](h)
+        mu, alpha = out.chunk(2, dim=-1)
+        alpha = torch.tanh(alpha)  # keep the log-scale bounded for stability
+        return mu, alpha
+
+
+class MAF(nn.Module):
+    """Masked autoregressive flow over ``N``-dimensional position-space fields.
+
+    Parameters
+    ----------
+    N : int
+        Field dimension.
+    hidden : int
+        Hidden width of each MADE.
+    n_blocks : int
+        Number of autoregressive blocks (with permutations between them).
+    n_hidden : int
+        Hidden layers per MADE.
+    """
+
+    def __init__(self, N: int, hidden: int = 64, n_blocks: int = 4, n_hidden: int = 1):
+        super().__init__()
+        self.N = int(N)
+        self.blocks = nn.ModuleList(
+            [_MADE(N, hidden, n_hidden) for _ in range(n_blocks)]
+        )
+        # Fixed permutations (reverse ordering is a standard, cheap choice that
+        # guarantees every variable eventually conditions on every other).
+        for i in range(n_blocks):
+            perm = torch.arange(N - 1, -1, -1) if i % 2 == 0 else torch.arange(N)
+            self.register_buffer(f"perm_{i}", perm)
+        self.register_buffer("data_mean", torch.zeros(N))
+        self.register_buffer("data_std", torch.ones(N))
+
+    def _perm(self, i: int) -> Tensor:
+        return getattr(self, f"perm_{i}")
+
+    def log_prob(self, x: Tensor) -> Tensor:
+        """Exact log-density of the flow (up to the constant standardization term)."""
+        y = (x - self.data_mean) / self.data_std
+        log_det = torch.zeros(x.shape[0], device=x.device, dtype=x.dtype)
+        for i, made in enumerate(self.blocks):
+            y = y[:, self._perm(i)]
+            mu, alpha = made(y)
+            y = (y - mu) * torch.exp(-alpha)
+            log_det = log_det - alpha.sum(dim=-1)
+        base = -0.5 * (y.pow(2) + math.log(2 * math.pi)).sum(dim=-1)
+        return base + log_det
+
+    def fit(
+        self,
+        data: Tensor,
+        epochs: int = 200,
+        lr: float = 1e-3,
+        batch_size: int = 256,
+        verbose: bool = False,
+    ) -> "MAF":
+        if data.dim() != 2 or data.shape[1] != self.N:
+            raise ValueError(
+                f"Expected data of shape (B, {self.N}); got {tuple(data.shape)}"
+            )
+        data = data.to(self.data_mean)
+        with torch.no_grad():
+            self.data_mean = data.mean(dim=0)
+            self.data_std = data.std(dim=0).clamp_min(1e-6)
+        opt = torch.optim.Adam(self.parameters(), lr=lr)
+        n = data.shape[0]
+        for epoch in range(epochs):
+            perm = torch.randperm(n, device=data.device)
+            total = 0.0
+            for start in range(0, n, batch_size):
+                idx = perm[start : start + batch_size]
+                loss = -self.log_prob(data[idx]).mean()
+                opt.zero_grad()
+                loss.backward()
+                opt.step()
+                total += loss.item() * idx.numel()
+            if verbose and (epoch % 20 == 0 or epoch == epochs - 1):
+                print(f"[MAF] epoch {epoch:4d}  nll={total / n:.4f}")
+        return self
+
+    @torch.no_grad()
+    def sample(self, n: int) -> Tensor:
+        """Draw ``n`` samples by inverting each block sequentially."""
+        y = torch.randn(n, self.N, device=self.data_mean.device, dtype=self.data_mean.dtype)
+        for i in reversed(range(len(self.blocks))):
+            made = self.blocks[i]
+            t = torch.zeros_like(y)
+            for d in range(self.N):
+                mu, alpha = made(t)
+                t[:, d] = y[:, d] * torch.exp(alpha[:, d]) + mu[:, d]
+            inv = torch.argsort(self._perm(i))
+            y = t[:, inv]
+        return y * self.data_std + self.data_mean
+
+    def extra_repr(self) -> str:
+        return f"N={self.N}, n_blocks={len(self.blocks)}"

couplnorm/models/pca.py

@@ -0,0 +1,71 @@
+"""Low-rank Gaussian (PCA) baseline."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+__all__ = ["PCAModel"]
+
+
+class PCAModel(nn.Module):
+    """Low-rank Gaussian: mean + top-``k`` principal components + isotropic noise.
+
+    Fits a rank-``k`` approximation of the data covariance and adds isotropic
+    residual noise equal to the mean of the discarded eigenvalues, so the model
+    is a proper full-rank Gaussian that under-represents off-principal
+    directions. Weaker than :class:`FullGaussian`; useful for showing that
+    truncating structure lowers coupling fidelity.
+
+    Parameters
+    ----------
+    N : int
+        Field dimension.
+    n_components : int
+        Number of principal components to retain.
+    """
+
+    def __init__(self, N: int, n_components: int = 4):
+        super().__init__()
+        if n_components < 1 or n_components > N:
+            raise ValueError(f"n_components must be in [1, {N}]; got {n_components}")
+        self.N = int(N)
+        self.k = int(n_components)
+        self.register_buffer("mean", torch.zeros(N))
+        self.register_buffer("components", torch.zeros(N, self.k))  # scaled by sqrt(eigval)
+        self.register_buffer("noise_std", torch.zeros(()))
+        self._fitted = False
+
+    @torch.no_grad()
+    def fit(self, data: Tensor) -> "PCAModel":
+        if data.dim() != 2 or data.shape[1] != self.N:
+            raise ValueError(
+                f"Expected data of shape (B, {self.N}); got {tuple(data.shape)}"
+            )
+        data = data.to(self.mean)
+        self.mean = data.mean(dim=0)
+        centered = data - self.mean
+        cov = centered.T @ centered / (data.shape[0] - 1)
+        # eigh returns ascending eigenvalues; take the top k.
+        evals, evecs = torch.linalg.eigh(cov)
+        evals = evals.clamp_min(0.0)
+        top_vals = evals[-self.k:]
+        top_vecs = evecs[:, -self.k:]
+        self.components = top_vecs * top_vals.sqrt().unsqueeze(0)
+        residual = evals[: self.N - self.k]
+        self.noise_std = residual.mean().clamp_min(0.0).sqrt() if residual.numel() else torch.zeros(())
+        self._fitted = True
+        return self
+
+    @torch.no_grad()
+    def sample(self, n: int) -> Tensor:
+        if not self._fitted:
+            raise RuntimeError("Call fit(data) before sample(n).")
+        z = torch.randn(n, self.k, device=self.mean.device, dtype=self.mean.dtype)
+        x = self.mean + z @ self.components.T
+        if self.noise_std > 0:
+            x = x + self.noise_std * torch.randn_like(x)
+        return x
+
+    def extra_repr(self) -> str:
+        return f"N={self.N}, n_components={self.k}, fitted={self._fitted}"

couplnorm/plotting.py

@@ -0,0 +1,116 @@
+"""Matplotlib helpers for the demo notebooks.
+
+Thin convenience wrappers only; matplotlib is an optional dependency (install
+with ``pip install couplnorm[dev]``). Each function accepts an optional ``ax``
+so plots compose into subplot grids.
+"""
+from __future__ import annotations
+
+from typing import Optional, Sequence
+
+import torch
+from torch import Tensor
+
+from couplnorm.coupling import coupling_from_samples
+
+__all__ = [
+    "plot_coupling_bars",
+    "plot_covariance_heatmap",
+    "plot_running_trajectory",
+    "plot_training",
+]
+
+
+def _require_mpl():
+    try:
+        import matplotlib.pyplot as plt  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - trivial guard
+        raise ImportError(
+            "Plotting requires matplotlib. Install with: pip install couplnorm[dev]"
+        ) from exc
+    return plt
+
+
+def plot_coupling_bars(
+    named_samples: dict[str, Tensor],
+    ax=None,
+    real_fft: bool = True,
+    title: str = "Coupling C across distributions",
+):
+    """Bar chart of C for each named ``(B, N)`` sample tensor."""
+    plt = _require_mpl()
+    names, values = [], []
+    for name, phi in named_samples.items():
+        names.append(name)
+        values.append(coupling_from_samples(phi, real_fft=real_fft).item())
+    if ax is None:
+        _, ax = plt.subplots(figsize=(6, 3.5))
+    ax.bar(names, values, color="#3b6ea5")
+    ax.set_ylabel("C")
+    ax.set_ylim(0, max(values + [0.1]) * 1.15)
+    ax.set_title(title)
+    for i, v in enumerate(values):
+        ax.text(i, v, f"{v:.3f}", ha="center", va="bottom", fontsize=9)
+    return ax
+
+
+def plot_covariance_heatmap(cov: Tensor, ax=None, title: str = "Spectral-energy covariance"):
+    """Heatmap of a spectral-energy covariance matrix."""
+    plt = _require_mpl()
+    if ax is None:
+        _, ax = plt.subplots(figsize=(4.5, 4))
+    im = ax.imshow(cov.detach().cpu().numpy(), cmap="magma")
+    ax.set_xlabel("mode k'")
+    ax.set_ylabel("mode k")
+    ax.set_title(title)
+    ax.figure.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+    return ax
+
+
+def plot_running_trajectory(
+    trajectory: Sequence[float],
+    batch_C: Optional[float] = None,
+    ax=None,
+    title: str = "Running C converges to batch C",
+):
+    """Line plot of a running-C trajectory with an optional batch-C reference."""
+    plt = _require_mpl()
+    if ax is None:
+        _, ax = plt.subplots(figsize=(7, 3.5))
+    ax.plot(list(trajectory), label="running C")
+    if batch_C is not None:
+        ax.axhline(batch_C, color="red", linestyle="--", label=f"batch C = {batch_C:.3f}")
+    ax.set_xlabel("chunk index")
+    ax.set_ylabel("running C")
+    ax.set_title(title)
+    ax.legend()
+    return ax
+
+
+def plot_training(
+    losses: Sequence[float],
+    c_steps: Optional[Sequence[int]] = None,
+    c_values: Optional[Sequence[float]] = None,
+    target_C: Optional[float] = None,
+):
+    """Two-panel training summary: loss curve and generator-C trajectory."""
+    plt = _require_mpl()
+    ncols = 2 if c_values is not None else 1
+    fig, axes = plt.subplots(1, ncols, figsize=(5.5 * ncols, 4))
+    ax0 = axes[0] if ncols == 2 else axes
+    ax0.plot(list(losses))
+    ax0.set_yscale("log")
+    ax0.set_xlabel("step")
+    ax0.set_ylabel("loss")
+    ax0.set_title("Training loss")
+    if ncols == 2:
+        ax1 = axes[1]
+        ax1.plot(list(c_steps), list(c_values), marker="o", label="generator C")
+        if target_C is not None:
+            ax1.axhline(target_C, color="red", linestyle="--", label=f"target C = {target_C:.3f}")
+        ax1.set_xlabel("step")
+        ax1.set_ylabel("C")
+        ax1.set_title("Generator C trajectory")
+        ax1.legend()
+    fig.tight_layout()
+    return fig

couplnorm/samplers.py

@@ -0,0 +1,170 @@
+"""Reference sampler for the 1D periodic phi^4 scalar field.
+
+This module is plumbing: it exists so users can install couplnorm and
+reproduce the paper plots without writing their own MCMC. The intellectual
+core of the library is ``coupling.py``; this file just feeds it data.
+
+Lattice action (a = 1, periodic boundary conditions)::
+
+    S(phi) = sum_x [ (1/2) (phi_{x+1} - phi_x)^2
+                     + (1/2) m2 * phi_x^2
+                     + lam * phi_x^4 ]
+
+At ``lam = 0`` the action is quadratic and the Fourier basis diagonalizes it,
+so distinct modes are independent Gaussians. The per-mode variance under the
+unitary ("ortho") FFT is then exactly ``1 / (m2 + 4 sin^2(pi k / N))``; see
+:func:`free_theory_mode_variance`. Turning on ``lam > 0`` couples the modes
+through the quartic vertex, which is exactly the structure the C diagnostic
+picks up.
+"""
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+__all__ = ["Phi4Sampler", "free_theory_mode_variance"]
+
+
+def free_theory_mode_variance(
+    N: int,
+    m2: float = 1.0,
+    real_fft: bool = True,
+    device=None,
+    dtype=torch.float32,
+) -> Tensor:
+    """Analytical spectral-energy expectation for the free (lam=0) theory.
+
+    Returns E[|phi_tilde_k|^2] = 1 / (m2 + 4 sin^2(pi k / N)) under the unitary
+    FFT convention (``norm="ortho"``). With ``real_fft=True`` only the
+    ``N//2 + 1`` unique modes are returned, matching :func:`Phi4Sampler.sample`
+    followed by ``torch.fft.rfft(..., norm="ortho")``.
+
+    Parameters
+    ----------
+    N : int
+        Number of lattice sites.
+    m2 : float
+        Mass-squared parameter of the action.
+    real_fft : bool
+        If True, return the ``N//2 + 1`` rfft modes; else all ``N`` modes.
+    """
+    k = torch.arange(N // 2 + 1 if real_fft else N, device=device, dtype=dtype)
+    lattice_momentum = 4.0 * torch.sin(math.pi * k / N).pow(2)
+    return 1.0 / (m2 + lattice_momentum)
+
+
+class Phi4Sampler:
+    """Checkerboard Metropolis-Hastings sampler for the 1D phi^4 field.
+
+    Each call to :meth:`sample` runs ``n`` independent Markov chains in
+    parallel (fully vectorized over chains and lattice sites), thermalizes
+    them, and returns their final configurations. Because the 1D nearest-
+    neighbor action couples only adjacent sites, even-indexed sites are
+    conditionally independent given odd-indexed sites (and vice versa), so a
+    whole sublattice can be proposed and accepted in one vectorized step.
+
+    Parameters
+    ----------
+    N : int
+        Number of lattice sites per configuration.
+    m2 : float
+        Mass-squared parameter of the action.
+    lam : float
+        Quartic coupling. ``lam = 0`` recovers the free theory.
+    step : float
+        Half-width of the uniform Metropolis proposal. Tune for ~40-60%
+        acceptance; the default is reasonable for ``m2 ~ 1``.
+    n_therm : int
+        Number of thermalization sweeps before a configuration is returned.
+    device, dtype :
+        Passed through to the field tensors.
+    """
+
+    def __init__(
+        self,
+        N: int = 32,
+        m2: float = 1.0,
+        lam: float = 0.0,
+        step: float = 1.0,
+        n_therm: int = 1000,
+        device=None,
+        dtype=torch.float32,
+    ):
+        if N < 2:
+            raise ValueError(f"N must be >= 2; got {N}")
+        if N % 2 != 0:
+            raise ValueError(
+                f"N must be even for checkerboard updates; got {N}"
+            )
+        if step <= 0:
+            raise ValueError(f"step must be > 0; got {step}")
+        self.N = int(N)
+        self.m2 = float(m2)
+        self.lam = float(lam)
+        self.step = float(step)
+        self.n_therm = int(n_therm)
+        self.device = device
+        self.dtype = dtype
+        # Cached acceptance rate from the most recent run, for diagnostics.
+        self.last_acceptance: Optional[float] = None
+        parity = torch.arange(self.N) % 2
+        self._even = (parity == 0)
+        self._odd = (parity == 1)
+
+    def _color_update(self, phi: Tensor, color_mask: Tensor) -> Tensor:
+        """Propose and accept/reject on one sublattice, in place-safe fashion."""
+        left = phi.roll(1, dims=-1)   # phi_{x-1}
+        right = phi.roll(-1, dims=-1)  # phi_{x+1}
+        delta = (torch.rand_like(phi) * 2.0 - 1.0) * self.step
+        phi_new = phi + delta
+
+        d_kin = 0.5 * (
+            (phi_new - right).pow(2) - (phi - right).pow(2)
+            + (phi_new - left).pow(2) - (phi - left).pow(2)
+        )
+        d_mass = 0.5 * self.m2 * (phi_new.pow(2) - phi.pow(2))
+        d_quartic = self.lam * (phi_new.pow(4) - phi.pow(4))
+        dS = d_kin + d_mass + d_quartic
+
+        accept = torch.rand_like(phi) < torch.exp(-dS)
+        do_update = accept & color_mask
+        self._accepted += (do_update & color_mask).sum()
+        self._proposed += color_mask.sum() * phi.shape[0]
+        return torch.where(do_update, phi_new, phi)
+
+    def sweep(self, phi: Tensor) -> Tensor:
+        """One full checkerboard sweep (even sublattice, then odd)."""
+        mask_even = self._even.to(phi.device)
+        mask_odd = self._odd.to(phi.device)
+        phi = self._color_update(phi, mask_even)
+        phi = self._color_update(phi, mask_odd)
+        return phi
+
+    @torch.no_grad()
+    def sample(self, n: int, seed: Optional[int] = None) -> Tensor:
+        """Draw ``n`` thermalized configurations, shape ``(n, N)``.
+
+        Runs ``n`` parallel chains from a random start and returns each chain's
+        configuration after ``n_therm`` sweeps.
+        """
+        if n < 1:
+            raise ValueError(f"n must be >= 1; got {n}")
+        if seed is not None:
+            torch.manual_seed(seed)
+        phi = torch.randn(n, self.N, device=self.device, dtype=self.dtype)
+        self._accepted = torch.zeros((), device=phi.device)
+        self._proposed = torch.zeros((), device=phi.device)
+        for _ in range(self.n_therm):
+            phi = self.sweep(phi)
+        denom = self._proposed.clamp_min(1.0)
+        self.last_acceptance = float((self._accepted / denom).item())
+        return phi
+
+    def __repr__(self) -> str:
+        return (
+            f"Phi4Sampler(N={self.N}, m2={self.m2}, lam={self.lam}, "
+            f"step={self.step}, n_therm={self.n_therm})"
+        )

couplnorm-0.1.0.dist-info/METADATA

@@ -0,0 +1,146 @@
+Metadata-Version: 2.4
+Name: couplnorm
+Version: 0.1.0
+Summary: Normalized 4th-order coupling metrics and losses for generative model evaluation.
+Author: Anish Bhat
+License: MIT
+Project-URL: Homepage, https://github.com/anishbhat28/couplnorm
+Project-URL: Issues, https://github.com/anishbhat28/couplnorm/issues
+Keywords: generative-models,normalizing-flows,pytorch,coupling,spectral-analysis,model-evaluation,non-gaussianity,fourth-order-statistics,lattice-field-theory,scientific-ml
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Dynamic: license-file
+
+# couplnorm
+
+**Normalized 4th-order coupling metrics and losses for generative model evaluation.**
+
+`couplnorm` implements the normalized off-diagonal coupling diagnostic **C**, a
+4th-order spectral statistic that measures *joint* coupling between Fourier modes
+— structure that second-order metrics (power spectrum, two-point function) miss
+by construction.
+
+> Marginal Gaussianity does not imply joint independence. **C measures the gap.**
+
+![C across distributions with matched marginals](assets/coupling_hook.png)
+
+All three distributions above have (near-)identical per-mode marginals; a power
+spectrum cannot tell them apart. C rises from ~0.04 to ~0.82 as joint 4th-order
+coupling is switched on. (Reproduce with `notebooks/01_basic_usage.ipynb`.)
+
+Given field samples, `couplnorm` computes the covariance of the per-mode spectral
+energies `E_k = |φ̃_k|²` and reduces it to a single scale-free number:
+
+```
+C = ‖Σ − diag(Σ)‖_F / ‖Σ‖_F
+```
+
+`C = 0` means the spectral energies are pairwise uncorrelated (a Gaussian /
+free-field signature); `C → 1` means off-diagonal coupling dominates.
+
+## Install
+
+```bash
+pip install -e .[dev]     # from a checkout
+```
+
+Requires Python ≥ 3.10 and PyTorch ≥ 2.0.
+
+## Quickstart
+
+```python
+import torch
+from couplnorm import coupling_from_samples, CouplingMetric, CouplingLoss
+
+# One-shot number from a batch of (B, N) real fields
+phi = torch.randn(10_000, 32)
+print(coupling_from_samples(phi).item())        # ~0.03  (independent modes)
+
+# As a streaming nn.Module metric
+metric = CouplingMetric(mode="running", momentum=0.02, n_modes=32 // 2 + 1)
+for chunk in loader:                            # doctest: +SKIP
+    metric(chunk)
+print(metric.compute().item())
+
+# As a training loss that matches a reference distribution's coupling
+loss_fn = CouplingLoss.from_data(reference_data, target_type="matrix")  # doctest: +SKIP
+loss = loss_fn(generated_batch)                 # doctest: +SKIP
+```
+
+## Three forms
+
+| Object                  | Use                                                            |
+| ----------------------- | ------------------------------------------------------------- |
+| `coupling_from_samples` | one-line functional helper — just want the number             |
+| `CouplingMetric`        | `nn.Module` metric; batch **or** streaming (running EMA) mode  |
+| `CouplingLoss`          | differentiable loss: match a covariance, match a scalar C, or minimize C (DeCov-style) |
+
+## The FFT convention (important)
+
+For a **real** field of length `N`, the DFT is conjugate-symmetric
+(`φ̃_k = φ̃*_{N−k}`), so `|φ̃_k|² = |φ̃_{N−k}|²` *exactly*. Feeding the full FFT into
+C creates `N − 2` off-diagonal entries that equal diagonal entries before any
+physics enters, inflating C by an analytical floor of
+
+```
+C_floor = sqrt((N − 2) / (2N − 2)) ≈ 0.7
+```
+
+`couplnorm` defaults to `real_fft=True` (uses `torch.fft.rfft`, keeping only the
+`N//2 + 1` unique modes), which reflects only genuine coupling. Set
+`real_fft=False` only for intrinsically complex fields. A regression test pins
+the analytical floor so this convention can't silently break.
+
+## What's in the box
+
+```
+src/couplnorm/
+  coupling.py     # the metric, loss, and helpers (the contribution)
+  samplers.py     # Phi4Sampler: 1D phi^4 Metropolis-Hastings reference sampler
+  models/         # FourierModel, FullGaussian, PCAModel, MAF baselines
+  plotting.py     # matplotlib helpers for the notebooks
+notebooks/
+  01_basic_usage.ipynb      # compare C across distributions
+  02_reproduce_paper.ipynb  # phi^4 free vs interacting theory
+  03_coupling_as_loss.ipynb # train a generator to match coupling
+```
+
+Baselines share one interface: construct, `.fit(data)` where applicable, then
+`.sample(n)` returning `(n, N)` real position-space fields. The `FourierModel`
+assumes marginal mode independence, so it collapses C to ≈ 0 regardless of the
+target — the concrete demonstration of why C is worth measuring.
+
+## Case study
+
+The φ⁴ lattice study that motivates C is:
+
+> A. Bhat, R. Ide, Z. Zhao. *When Independent Gaussian Models Break Down:
+> Characterizing Regime-Dependent Modeling Failures in φ⁴ Theory.*
+> arXiv:[2605.01145](https://arxiv.org/abs/2605.01145).
+
+That paper's reported C values (C(N=32) ≈ 0.06 up to C(N=128) ≈ 0.2, strong
+coupling saturating near 0.17–0.18) are computed on the **unique rfft modes** —
+the `real_fft=True` default here. The full FFT would inflate every value to the
+~0.7 conjugate-symmetry floor (see above), so the published magnitudes are only
+consistent with the unique-mode convention.
+
+## Citation
+
+See [`CITATION.cff`](CITATION.cff).
+
+## License
+
+MIT — see [`LICENSE`](LICENSE).

couplnorm-0.1.0.dist-info/RECORD

@@ -0,0 +1,14 @@
+couplnorm/__init__.py,sha256=a3VlmkTPH0tPyYWQsETWbfvkvQ6hEz-OJJsUkieOYns,322
+couplnorm/coupling.py,sha256=-E6WaSOmG8VyUL757kYOpKMDLBtsZgtJdxH1nm6B_jE,10508
+couplnorm/plotting.py,sha256=Vo221PiUPsOTK-JfMT17Gf6e5RIYVD8JlGEc-gETydM,3670
+couplnorm/samplers.py,sha256=k9J3G4pCSBx9SxfSpP7iwSrN4RjFAQLyimGc_D9fppg,6296
+couplnorm/models/__init__.py,sha256=7lWgWvbc1oJ2vi73WYDeF1wwVZTcdsQf4hC2CW3Y6AY,1225
+couplnorm/models/fourier.py,sha256=GMLZoUckhQIWl8rLf6VIKhRCsimnDCJOCC8VvLn8FK8,2932
+couplnorm/models/gaussian.py,sha256=pTVz12ueYEsJ3uOxDkwOEtZ8phHROO2r8lzTRDZFNeQ,2160
+couplnorm/models/maf.py,sha256=ckYFU32Ui1HzvJbjoAGkGO9TwYUmUxEgNNPhlXVwodo,6476
+couplnorm/models/pca.py,sha256=nZpT8sleIdie2Vp6lWerbsmSvV6zZohGrdsLIO16Mww,2683
+couplnorm-0.1.0.dist-info/licenses/LICENSE,sha256=aAGv9YvWMyE7dbE_qCI-swi8JzEqG7IYb5gkokxGeRg,1067
+couplnorm-0.1.0.dist-info/METADATA,sha256=2GazRukFWRQ0N7iwnT6sYQh8YI7hHLo98ItXTzVLtK0,5843
+couplnorm-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+couplnorm-0.1.0.dist-info/top_level.txt,sha256=9KoE72M7B4GVWSK3K7Yh81R3HeKYAepSL9J3WjATMzk,10
+couplnorm-0.1.0.dist-info/RECORD,,

couplnorm-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (83.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

couplnorm-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Anish Bhat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

couplnorm-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+couplnorm