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.
- couplnorm-0.1.0/LICENSE +21 -0
- couplnorm-0.1.0/PKG-INFO +146 -0
- couplnorm-0.1.0/README.md +120 -0
- couplnorm-0.1.0/pyproject.toml +41 -0
- couplnorm-0.1.0/setup.cfg +4 -0
- couplnorm-0.1.0/src/couplnorm/__init__.py +11 -0
- couplnorm-0.1.0/src/couplnorm/coupling.py +285 -0
- couplnorm-0.1.0/src/couplnorm/models/__init__.py +25 -0
- couplnorm-0.1.0/src/couplnorm/models/fourier.py +73 -0
- couplnorm-0.1.0/src/couplnorm/models/gaussian.py +61 -0
- couplnorm-0.1.0/src/couplnorm/models/maf.py +171 -0
- couplnorm-0.1.0/src/couplnorm/models/pca.py +71 -0
- couplnorm-0.1.0/src/couplnorm/plotting.py +116 -0
- couplnorm-0.1.0/src/couplnorm/samplers.py +170 -0
- couplnorm-0.1.0/src/couplnorm.egg-info/PKG-INFO +146 -0
- couplnorm-0.1.0/src/couplnorm.egg-info/SOURCES.txt +21 -0
- couplnorm-0.1.0/src/couplnorm.egg-info/dependency_links.txt +1 -0
- couplnorm-0.1.0/src/couplnorm.egg-info/requires.txt +6 -0
- couplnorm-0.1.0/src/couplnorm.egg-info/top_level.txt +1 -0
- couplnorm-0.1.0/tests/test_coupling.py +239 -0
- couplnorm-0.1.0/tests/test_models.py +94 -0
- couplnorm-0.1.0/tests/test_reproducibility.py +63 -0
- couplnorm-0.1.0/tests/test_samplers.py +70 -0
couplnorm-0.1.0/LICENSE
ADDED
|
@@ -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.
|
couplnorm-0.1.0/PKG-INFO
ADDED
|
@@ -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
|
+

|
|
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
|
+

|
|
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,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"]
|