couplnorm 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Anish Bhat
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,146 @@
1
+ Metadata-Version: 2.4
2
+ Name: couplnorm
3
+ Version: 0.1.0
4
+ Summary: Normalized 4th-order coupling metrics and losses for generative model evaluation.
5
+ Author: Anish Bhat
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/anishbhat28/couplnorm
8
+ Project-URL: Issues, https://github.com/anishbhat28/couplnorm/issues
9
+ Keywords: generative-models,normalizing-flows,pytorch,coupling,spectral-analysis,model-evaluation,non-gaussianity,fourth-order-statistics,lattice-field-theory,scientific-ml
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
15
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
16
+ Classifier: Topic :: Scientific/Engineering :: Physics
17
+ Requires-Python: >=3.10
18
+ Description-Content-Type: text/markdown
19
+ License-File: LICENSE
20
+ Requires-Dist: torch>=2.0
21
+ Provides-Extra: dev
22
+ Requires-Dist: pytest; extra == "dev"
23
+ Requires-Dist: pytest-cov; extra == "dev"
24
+ Requires-Dist: matplotlib; extra == "dev"
25
+ Dynamic: license-file
26
+
27
+ # couplnorm
28
+
29
+ **Normalized 4th-order coupling metrics and losses for generative model evaluation.**
30
+
31
+ `couplnorm` implements the normalized off-diagonal coupling diagnostic **C**, a
32
+ 4th-order spectral statistic that measures *joint* coupling between Fourier modes
33
+ — structure that second-order metrics (power spectrum, two-point function) miss
34
+ by construction.
35
+
36
+ > Marginal Gaussianity does not imply joint independence. **C measures the gap.**
37
+
38
+ ![C across distributions with matched marginals](assets/coupling_hook.png)
39
+
40
+ All three distributions above have (near-)identical per-mode marginals; a power
41
+ spectrum cannot tell them apart. C rises from ~0.04 to ~0.82 as joint 4th-order
42
+ coupling is switched on. (Reproduce with `notebooks/01_basic_usage.ipynb`.)
43
+
44
+ Given field samples, `couplnorm` computes the covariance of the per-mode spectral
45
+ energies `E_k = |φ̃_k|²` and reduces it to a single scale-free number:
46
+
47
+ ```
48
+ C = ‖Σ − diag(Σ)‖_F / ‖Σ‖_F
49
+ ```
50
+
51
+ `C = 0` means the spectral energies are pairwise uncorrelated (a Gaussian /
52
+ free-field signature); `C → 1` means off-diagonal coupling dominates.
53
+
54
+ ## Install
55
+
56
+ ```bash
57
+ pip install -e .[dev] # from a checkout
58
+ ```
59
+
60
+ Requires Python ≥ 3.10 and PyTorch ≥ 2.0.
61
+
62
+ ## Quickstart
63
+
64
+ ```python
65
+ import torch
66
+ from couplnorm import coupling_from_samples, CouplingMetric, CouplingLoss
67
+
68
+ # One-shot number from a batch of (B, N) real fields
69
+ phi = torch.randn(10_000, 32)
70
+ print(coupling_from_samples(phi).item()) # ~0.03 (independent modes)
71
+
72
+ # As a streaming nn.Module metric
73
+ metric = CouplingMetric(mode="running", momentum=0.02, n_modes=32 // 2 + 1)
74
+ for chunk in loader: # doctest: +SKIP
75
+ metric(chunk)
76
+ print(metric.compute().item())
77
+
78
+ # As a training loss that matches a reference distribution's coupling
79
+ loss_fn = CouplingLoss.from_data(reference_data, target_type="matrix") # doctest: +SKIP
80
+ loss = loss_fn(generated_batch) # doctest: +SKIP
81
+ ```
82
+
83
+ ## Three forms
84
+
85
+ | Object | Use |
86
+ | ----------------------- | ------------------------------------------------------------- |
87
+ | `coupling_from_samples` | one-line functional helper — just want the number |
88
+ | `CouplingMetric` | `nn.Module` metric; batch **or** streaming (running EMA) mode |
89
+ | `CouplingLoss` | differentiable loss: match a covariance, match a scalar C, or minimize C (DeCov-style) |
90
+
91
+ ## The FFT convention (important)
92
+
93
+ For a **real** field of length `N`, the DFT is conjugate-symmetric
94
+ (`φ̃_k = φ̃*_{N−k}`), so `|φ̃_k|² = |φ̃_{N−k}|²` *exactly*. Feeding the full FFT into
95
+ C creates `N − 2` off-diagonal entries that equal diagonal entries before any
96
+ physics enters, inflating C by an analytical floor of
97
+
98
+ ```
99
+ C_floor = sqrt((N − 2) / (2N − 2)) ≈ 0.7
100
+ ```
101
+
102
+ `couplnorm` defaults to `real_fft=True` (uses `torch.fft.rfft`, keeping only the
103
+ `N//2 + 1` unique modes), which reflects only genuine coupling. Set
104
+ `real_fft=False` only for intrinsically complex fields. A regression test pins
105
+ the analytical floor so this convention can't silently break.
106
+
107
+ ## What's in the box
108
+
109
+ ```
110
+ src/couplnorm/
111
+ coupling.py # the metric, loss, and helpers (the contribution)
112
+ samplers.py # Phi4Sampler: 1D phi^4 Metropolis-Hastings reference sampler
113
+ models/ # FourierModel, FullGaussian, PCAModel, MAF baselines
114
+ plotting.py # matplotlib helpers for the notebooks
115
+ notebooks/
116
+ 01_basic_usage.ipynb # compare C across distributions
117
+ 02_reproduce_paper.ipynb # phi^4 free vs interacting theory
118
+ 03_coupling_as_loss.ipynb # train a generator to match coupling
119
+ ```
120
+
121
+ Baselines share one interface: construct, `.fit(data)` where applicable, then
122
+ `.sample(n)` returning `(n, N)` real position-space fields. The `FourierModel`
123
+ assumes marginal mode independence, so it collapses C to ≈ 0 regardless of the
124
+ target — the concrete demonstration of why C is worth measuring.
125
+
126
+ ## Case study
127
+
128
+ The φ⁴ lattice study that motivates C is:
129
+
130
+ > A. Bhat, R. Ide, Z. Zhao. *When Independent Gaussian Models Break Down:
131
+ > Characterizing Regime-Dependent Modeling Failures in φ⁴ Theory.*
132
+ > arXiv:[2605.01145](https://arxiv.org/abs/2605.01145).
133
+
134
+ That paper's reported C values (C(N=32) ≈ 0.06 up to C(N=128) ≈ 0.2, strong
135
+ coupling saturating near 0.17–0.18) are computed on the **unique rfft modes** —
136
+ the `real_fft=True` default here. The full FFT would inflate every value to the
137
+ ~0.7 conjugate-symmetry floor (see above), so the published magnitudes are only
138
+ consistent with the unique-mode convention.
139
+
140
+ ## Citation
141
+
142
+ See [`CITATION.cff`](CITATION.cff).
143
+
144
+ ## License
145
+
146
+ MIT — see [`LICENSE`](LICENSE).
@@ -0,0 +1,120 @@
1
+ # couplnorm
2
+
3
+ **Normalized 4th-order coupling metrics and losses for generative model evaluation.**
4
+
5
+ `couplnorm` implements the normalized off-diagonal coupling diagnostic **C**, a
6
+ 4th-order spectral statistic that measures *joint* coupling between Fourier modes
7
+ — structure that second-order metrics (power spectrum, two-point function) miss
8
+ by construction.
9
+
10
+ > Marginal Gaussianity does not imply joint independence. **C measures the gap.**
11
+
12
+ ![C across distributions with matched marginals](assets/coupling_hook.png)
13
+
14
+ All three distributions above have (near-)identical per-mode marginals; a power
15
+ spectrum cannot tell them apart. C rises from ~0.04 to ~0.82 as joint 4th-order
16
+ coupling is switched on. (Reproduce with `notebooks/01_basic_usage.ipynb`.)
17
+
18
+ Given field samples, `couplnorm` computes the covariance of the per-mode spectral
19
+ energies `E_k = |φ̃_k|²` and reduces it to a single scale-free number:
20
+
21
+ ```
22
+ C = ‖Σ − diag(Σ)‖_F / ‖Σ‖_F
23
+ ```
24
+
25
+ `C = 0` means the spectral energies are pairwise uncorrelated (a Gaussian /
26
+ free-field signature); `C → 1` means off-diagonal coupling dominates.
27
+
28
+ ## Install
29
+
30
+ ```bash
31
+ pip install -e .[dev] # from a checkout
32
+ ```
33
+
34
+ Requires Python ≥ 3.10 and PyTorch ≥ 2.0.
35
+
36
+ ## Quickstart
37
+
38
+ ```python
39
+ import torch
40
+ from couplnorm import coupling_from_samples, CouplingMetric, CouplingLoss
41
+
42
+ # One-shot number from a batch of (B, N) real fields
43
+ phi = torch.randn(10_000, 32)
44
+ print(coupling_from_samples(phi).item()) # ~0.03 (independent modes)
45
+
46
+ # As a streaming nn.Module metric
47
+ metric = CouplingMetric(mode="running", momentum=0.02, n_modes=32 // 2 + 1)
48
+ for chunk in loader: # doctest: +SKIP
49
+ metric(chunk)
50
+ print(metric.compute().item())
51
+
52
+ # As a training loss that matches a reference distribution's coupling
53
+ loss_fn = CouplingLoss.from_data(reference_data, target_type="matrix") # doctest: +SKIP
54
+ loss = loss_fn(generated_batch) # doctest: +SKIP
55
+ ```
56
+
57
+ ## Three forms
58
+
59
+ | Object | Use |
60
+ | ----------------------- | ------------------------------------------------------------- |
61
+ | `coupling_from_samples` | one-line functional helper — just want the number |
62
+ | `CouplingMetric` | `nn.Module` metric; batch **or** streaming (running EMA) mode |
63
+ | `CouplingLoss` | differentiable loss: match a covariance, match a scalar C, or minimize C (DeCov-style) |
64
+
65
+ ## The FFT convention (important)
66
+
67
+ For a **real** field of length `N`, the DFT is conjugate-symmetric
68
+ (`φ̃_k = φ̃*_{N−k}`), so `|φ̃_k|² = |φ̃_{N−k}|²` *exactly*. Feeding the full FFT into
69
+ C creates `N − 2` off-diagonal entries that equal diagonal entries before any
70
+ physics enters, inflating C by an analytical floor of
71
+
72
+ ```
73
+ C_floor = sqrt((N − 2) / (2N − 2)) ≈ 0.7
74
+ ```
75
+
76
+ `couplnorm` defaults to `real_fft=True` (uses `torch.fft.rfft`, keeping only the
77
+ `N//2 + 1` unique modes), which reflects only genuine coupling. Set
78
+ `real_fft=False` only for intrinsically complex fields. A regression test pins
79
+ the analytical floor so this convention can't silently break.
80
+
81
+ ## What's in the box
82
+
83
+ ```
84
+ src/couplnorm/
85
+ coupling.py # the metric, loss, and helpers (the contribution)
86
+ samplers.py # Phi4Sampler: 1D phi^4 Metropolis-Hastings reference sampler
87
+ models/ # FourierModel, FullGaussian, PCAModel, MAF baselines
88
+ plotting.py # matplotlib helpers for the notebooks
89
+ notebooks/
90
+ 01_basic_usage.ipynb # compare C across distributions
91
+ 02_reproduce_paper.ipynb # phi^4 free vs interacting theory
92
+ 03_coupling_as_loss.ipynb # train a generator to match coupling
93
+ ```
94
+
95
+ Baselines share one interface: construct, `.fit(data)` where applicable, then
96
+ `.sample(n)` returning `(n, N)` real position-space fields. The `FourierModel`
97
+ assumes marginal mode independence, so it collapses C to ≈ 0 regardless of the
98
+ target — the concrete demonstration of why C is worth measuring.
99
+
100
+ ## Case study
101
+
102
+ The φ⁴ lattice study that motivates C is:
103
+
104
+ > A. Bhat, R. Ide, Z. Zhao. *When Independent Gaussian Models Break Down:
105
+ > Characterizing Regime-Dependent Modeling Failures in φ⁴ Theory.*
106
+ > arXiv:[2605.01145](https://arxiv.org/abs/2605.01145).
107
+
108
+ That paper's reported C values (C(N=32) ≈ 0.06 up to C(N=128) ≈ 0.2, strong
109
+ coupling saturating near 0.17–0.18) are computed on the **unique rfft modes** —
110
+ the `real_fft=True` default here. The full FFT would inflate every value to the
111
+ ~0.7 conjugate-symmetry floor (see above), so the published magnitudes are only
112
+ consistent with the unique-mode convention.
113
+
114
+ ## Citation
115
+
116
+ See [`CITATION.cff`](CITATION.cff).
117
+
118
+ ## License
119
+
120
+ MIT — see [`LICENSE`](LICENSE).
@@ -0,0 +1,41 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "couplnorm"
7
+ version = "0.1.0"
8
+ description = "Normalized 4th-order coupling metrics and losses for generative model evaluation."
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = {text = "MIT"}
12
+ authors = [{name = "Anish Bhat"}]
13
+ dependencies = ["torch>=2.0"]
14
+ keywords = [
15
+ "generative-models", "normalizing-flows", "pytorch",
16
+ "coupling", "spectral-analysis", "model-evaluation",
17
+ "non-gaussianity", "fourth-order-statistics",
18
+ "lattice-field-theory", "scientific-ml",
19
+ ]
20
+ classifiers = [
21
+ "Development Status :: 3 - Alpha",
22
+ "Intended Audience :: Science/Research",
23
+ "License :: OSI Approved :: MIT License",
24
+ "Programming Language :: Python :: 3",
25
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
26
+ "Topic :: Scientific/Engineering :: Mathematics",
27
+ "Topic :: Scientific/Engineering :: Physics",
28
+ ]
29
+
30
+ [project.optional-dependencies]
31
+ dev = ["pytest", "pytest-cov", "matplotlib"]
32
+
33
+ [project.urls]
34
+ Homepage = "https://github.com/anishbhat28/couplnorm"
35
+ Issues = "https://github.com/anishbhat28/couplnorm/issues"
36
+
37
+ [tool.setuptools.packages.find]
38
+ where = ["src"]
39
+
40
+ [tool.pytest.ini_options]
41
+ testpaths = ["tests"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,11 @@
1
+ """couplnorm: normalized 4th-order coupling metrics and losses for
2
+ generative model evaluation on Fourier-decomposable data.
3
+ """
4
+ from couplnorm.coupling import (
5
+ CouplingLoss,
6
+ CouplingMetric,
7
+ coupling_from_samples,
8
+ )
9
+
10
+ __all__ = ["CouplingMetric", "CouplingLoss", "coupling_from_samples"]
11
+ __version__ = "0.1.0"
@@ -0,0 +1,285 @@
1
+ """Normalized off-diagonal coupling diagnostic and loss for evaluating
2
+ generative models on Fourier-decomposable data.
3
+ """
4
+ from __future__ import annotations
5
+
6
+ from typing import Optional, Union
7
+
8
+ import torch
9
+ import torch.nn as nn
10
+ from torch import Tensor
11
+
12
+ __all__ = ["CouplingMetric", "CouplingLoss", "coupling_from_samples"]
13
+
14
+
15
+ def _spectral_energies(
16
+ x: Tensor,
17
+ input_space: str = "position",
18
+ fft_norm: str = "ortho",
19
+ real_fft: bool = True,
20
+ ) -> Tensor:
21
+ if input_space == "position":
22
+ if real_fft:
23
+ x_hat = torch.fft.rfft(x, dim=-1, norm=fft_norm)
24
+ else:
25
+ x_hat = torch.fft.fft(x, dim=-1, norm=fft_norm)
26
+ return x_hat.real.pow(2) + x_hat.imag.pow(2)
27
+ elif input_space == "spectral":
28
+ return x
29
+ else:
30
+ raise ValueError(
31
+ f"input_space must be 'position' or 'spectral', got {input_space!r}"
32
+ )
33
+
34
+
35
+ def _batch_covariance(E: Tensor, unbiased: bool = True) -> Tensor:
36
+ if E.dim() != 2:
37
+ raise ValueError(f"Expected 2D tensor (B, M), got shape {tuple(E.shape)}")
38
+ B = E.shape[0]
39
+ if B < 2 and unbiased:
40
+ raise ValueError(
41
+ f"Need B >= 2 for unbiased covariance; got B={B}. "
42
+ "Set unbiased=False or use mode='running' for streaming."
43
+ )
44
+ mean = E.mean(dim=0, keepdim=True)
45
+ centered = E - mean
46
+ denom = (B - 1) if unbiased else B
47
+ return centered.T @ centered / denom
48
+
49
+
50
+ def _coupling_from_cov(cov: Tensor, eps: float = 1e-12) -> Tensor:
51
+ diag_vec = torch.diagonal(cov)
52
+ off_diag = cov - torch.diag_embed(diag_vec)
53
+ num = torch.linalg.matrix_norm(off_diag, ord="fro")
54
+ den = torch.linalg.matrix_norm(cov, ord="fro")
55
+ return num / (den + eps)
56
+
57
+
58
+ def coupling_from_samples(
59
+ x: Tensor,
60
+ input_space: str = "position",
61
+ fft_norm: str = "ortho",
62
+ real_fft: bool = True,
63
+ eps: float = 1e-12,
64
+ unbiased: bool = True,
65
+ ) -> Tensor:
66
+ if x.dim() == 1:
67
+ raise ValueError(
68
+ "coupling_from_samples requires batched input (B, N). "
69
+ "For streaming, use CouplingMetric(mode='running')."
70
+ )
71
+ E = _spectral_energies(x, input_space=input_space, fft_norm=fft_norm, real_fft=real_fft)
72
+ cov = _batch_covariance(E, unbiased=unbiased)
73
+ return _coupling_from_cov(cov, eps)
74
+
75
+
76
+ class CouplingMetric(nn.Module):
77
+ def __init__(
78
+ self,
79
+ input_space: str = "position",
80
+ mode: str = "batch",
81
+ momentum: float = 0.1,
82
+ eps: float = 1e-12,
83
+ fft_norm: str = "ortho",
84
+ real_fft: bool = True,
85
+ unbiased: bool = True,
86
+ n_modes: Optional[int] = None,
87
+ ):
88
+ super().__init__()
89
+ if input_space not in ("position", "spectral"):
90
+ raise ValueError(
91
+ f"input_space must be 'position' or 'spectral'; got {input_space!r}"
92
+ )
93
+ if mode not in ("batch", "running"):
94
+ raise ValueError(f"mode must be 'batch' or 'running'; got {mode!r}")
95
+ if not (0.0 < momentum <= 1.0):
96
+ raise ValueError(f"momentum must be in (0, 1]; got {momentum}")
97
+
98
+ self.input_space = input_space
99
+ self.mode = mode
100
+ self.momentum = float(momentum)
101
+ self.eps = float(eps)
102
+ self.fft_norm = fft_norm
103
+ self.real_fft = bool(real_fft)
104
+ self.unbiased = bool(unbiased)
105
+
106
+ if mode == "running":
107
+ if n_modes is not None:
108
+ self.register_buffer("running_mean", torch.zeros(n_modes))
109
+ self.register_buffer("running_second", torch.zeros(n_modes, n_modes))
110
+ else:
111
+ self.register_buffer("running_mean", torch.empty(0))
112
+ self.register_buffer("running_second", torch.empty(0))
113
+ self.register_buffer("num_batches_tracked", torch.tensor(0, dtype=torch.long))
114
+
115
+ def _maybe_init_buffers(self, M: int, device, dtype) -> None:
116
+ if self.running_mean.numel() == 0:
117
+ self.running_mean = torch.zeros(M, device=device, dtype=dtype)
118
+ self.running_second = torch.zeros(M, M, device=device, dtype=dtype)
119
+
120
+ def forward(self, x: Tensor) -> Tensor:
121
+ if x.dim() == 1:
122
+ x = x.unsqueeze(0)
123
+ if x.dim() != 2:
124
+ raise ValueError(f"Expected 1D or 2D input, got shape {tuple(x.shape)}")
125
+ E = _spectral_energies(x, self.input_space, self.fft_norm, self.real_fft)
126
+ B, M = E.shape
127
+
128
+ if self.mode == "batch":
129
+ cov = _batch_covariance(E, unbiased=self.unbiased)
130
+ return _coupling_from_cov(cov, self.eps)
131
+
132
+ self._maybe_init_buffers(M, E.device, E.dtype)
133
+ if self.running_mean.shape[0] != M:
134
+ raise ValueError(
135
+ "Number of modes changed: running buffers have "
136
+ f"N={self.running_mean.shape[0]} but got N={M}."
137
+ )
138
+
139
+ with torch.no_grad():
140
+ batch_mean = E.mean(dim=0)
141
+ batch_second = (E.T @ E) / B
142
+ m = self.momentum
143
+ self.running_mean.mul_(1.0 - m).add_(batch_mean, alpha=m)
144
+ self.running_second.mul_(1.0 - m).add_(batch_second, alpha=m)
145
+ self.num_batches_tracked.add_(1)
146
+ return self.compute()
147
+
148
+ def compute(self) -> Tensor:
149
+ if self.mode != "running":
150
+ raise RuntimeError("compute() is only valid in running mode.")
151
+ if self.running_mean.numel() == 0 or self.num_batches_tracked.item() == 0:
152
+ raise RuntimeError("No batches tracked yet; call forward() at least once.")
153
+ mu = self.running_mean
154
+ M = self.running_second
155
+ cov = M - torch.outer(mu, mu)
156
+ return _coupling_from_cov(cov, self.eps)
157
+
158
+ def covariance(self) -> Tensor:
159
+ if self.mode != "running":
160
+ raise RuntimeError(
161
+ "covariance() requires mode='running'. For one-off batches, "
162
+ "use coupling_from_samples or recompute manually."
163
+ )
164
+ mu = self.running_mean
165
+ M = self.running_second
166
+ return M - torch.outer(mu, mu)
167
+
168
+ def reset(self) -> None:
169
+ if self.mode != "running":
170
+ return
171
+ if self.running_mean.numel() > 0:
172
+ self.running_mean.zero_()
173
+ self.running_second.zero_()
174
+ self.num_batches_tracked.zero_()
175
+
176
+ def extra_repr(self) -> str:
177
+ return (
178
+ f"input_space={self.input_space!r}, mode={self.mode!r}, "
179
+ f"momentum={self.momentum}, eps={self.eps}, "
180
+ f"fft_norm={self.fft_norm!r}, real_fft={self.real_fft}"
181
+ )
182
+
183
+
184
+ class CouplingLoss(nn.Module):
185
+ def __init__(
186
+ self,
187
+ target_type: str = "matrix",
188
+ target: Optional[Union[Tensor, float]] = None,
189
+ input_space: str = "position",
190
+ eps: float = 1e-12,
191
+ fft_norm: str = "ortho",
192
+ real_fft: bool = True,
193
+ unbiased: bool = True,
194
+ ):
195
+ super().__init__()
196
+ if target_type not in ("matrix", "scalar", "regularize"):
197
+ raise ValueError(
198
+ "target_type must be 'matrix', 'scalar', or 'regularize'; "
199
+ f"got {target_type!r}"
200
+ )
201
+ if input_space not in ("position", "spectral"):
202
+ raise ValueError(
203
+ f"input_space must be 'position' or 'spectral'; got {input_space!r}"
204
+ )
205
+
206
+ if target_type == "matrix":
207
+ if target is None:
208
+ raise ValueError(
209
+ "target_type='matrix' requires a (M, M) target covariance."
210
+ )
211
+ tgt = torch.as_tensor(target, dtype=torch.float32)
212
+ if tgt.dim() != 2 or tgt.shape[0] != tgt.shape[1]:
213
+ raise ValueError(
214
+ f"matrix target must be (M, M); got shape {tuple(tgt.shape)}"
215
+ )
216
+ self.register_buffer("target_cov", tgt.detach().clone())
217
+ self.register_buffer(
218
+ "target_cov_frob",
219
+ torch.linalg.matrix_norm(tgt, ord="fro").detach().clone(),
220
+ )
221
+ elif target_type == "scalar":
222
+ if target is None:
223
+ raise ValueError("target_type='scalar' requires a target C value.")
224
+ self.register_buffer(
225
+ "target_C", torch.as_tensor(float(target), dtype=torch.float32)
226
+ )
227
+
228
+ self.target_type = target_type
229
+ self.input_space = input_space
230
+ self.eps = float(eps)
231
+ self.fft_norm = fft_norm
232
+ self.real_fft = bool(real_fft)
233
+ self.unbiased = bool(unbiased)
234
+
235
+ @classmethod
236
+ def from_data(
237
+ cls,
238
+ data: Tensor,
239
+ target_type: str = "matrix",
240
+ input_space: str = "position",
241
+ fft_norm: str = "ortho",
242
+ real_fft: bool = True,
243
+ unbiased: bool = True,
244
+ eps: float = 1e-12,
245
+ ) -> "CouplingLoss":
246
+ if target_type not in ("matrix", "scalar"):
247
+ raise ValueError(f"from_data does not support target_type={target_type!r}")
248
+ with torch.no_grad():
249
+ E = _spectral_energies(
250
+ data, input_space=input_space, fft_norm=fft_norm, real_fft=real_fft
251
+ )
252
+ cov = _batch_covariance(E, unbiased=unbiased)
253
+ if target_type == "matrix":
254
+ return cls(
255
+ target_type="matrix", target=cov, input_space=input_space, eps=eps,
256
+ fft_norm=fft_norm, real_fft=real_fft, unbiased=unbiased,
257
+ )
258
+ else:
259
+ C = _coupling_from_cov(cov, eps).item()
260
+ return cls(
261
+ target_type="scalar", target=C, input_space=input_space, eps=eps,
262
+ fft_norm=fft_norm, real_fft=real_fft, unbiased=unbiased,
263
+ )
264
+
265
+ def forward(self, x: Tensor) -> Tensor:
266
+ if x.dim() == 1:
267
+ x = x.unsqueeze(0)
268
+ E = _spectral_energies(x, self.input_space, self.fft_norm, self.real_fft)
269
+ cov = _batch_covariance(E, unbiased=self.unbiased)
270
+ if self.target_type == "matrix":
271
+ diff = cov - self.target_cov
272
+ num = torch.linalg.matrix_norm(diff, ord="fro").pow(2)
273
+ den = self.target_cov_frob.pow(2) + self.eps
274
+ return num / den
275
+ elif self.target_type == "scalar":
276
+ C = _coupling_from_cov(cov, self.eps)
277
+ return (C - self.target_C).pow(2)
278
+ else:
279
+ return _coupling_from_cov(cov, self.eps)
280
+
281
+ def extra_repr(self) -> str:
282
+ return (
283
+ f"target_type={self.target_type!r}, "
284
+ f"input_space={self.input_space!r}, fft_norm={self.fft_norm!r}"
285
+ )
@@ -0,0 +1,25 @@
1
+ """Baseline generative models for couplnorm demos and benchmarks.
2
+
3
+ Every model is an ``nn.Module`` with a common interface:
4
+
5
+ - ``fit(data)`` estimates parameters from a ``(B, N)`` batch of position-space
6
+ field configurations (a no-op / trainer depending on the model).
7
+ - ``sample(n)`` returns ``(n, N)`` real position-space configurations.
8
+
9
+ The models span the spectrum of joint-coupling fidelity:
10
+
11
+ - :class:`FourierModel` assumes marginal independence of Fourier modes, so its
12
+ samples have C near zero regardless of the target. It is the "wrong by
13
+ construction" baseline that motivates the C diagnostic.
14
+ - :class:`FullGaussian` matches the full second-order (covariance) structure of
15
+ the data but, being Gaussian, cannot reproduce genuine 4th-order coupling.
16
+ - :class:`PCAModel` is a low-rank Gaussian; even weaker than FullGaussian.
17
+ - :class:`MAF` is a masked autoregressive flow that can, in principle, capture
18
+ higher-order structure and drive C toward the target.
19
+ """
20
+ from couplnorm.models.fourier import FourierModel
21
+ from couplnorm.models.gaussian import FullGaussian
22
+ from couplnorm.models.maf import MAF
23
+ from couplnorm.models.pca import PCAModel
24
+
25
+ __all__ = ["FourierModel", "FullGaussian", "PCAModel", "MAF"]