- paper/pnas -> paper/manuscript (venue-neutral)
- configs/layer1 -> configs/inheritance, src/knowledge -> src/inheritance
(imported as `inheritance`), make layer1 -> make inheritance; layer2 alias dropped
- inheritance and trained-network bundles named after the manuscript figure
they feed (fig2_grounding_sweep, figS3_rebaselining, ...), or descriptively
where they feed none; configs keep their `experiment:` value so parquet
hashes are unchanged, only output.dir moves
- figure scripts, SI figure sources, notebooks, REPRODUCING.md, README and the
SI Methods/tables updated; make clean no longer deletes tracked manifests;
reproduce.sh hashes the s{seed}/ layouts too
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Y64o8FKP7rCuXzC48pxpMm
106 lines
4.6 KiB
Python
106 lines
4.6 KiB
Python
"""MNIST-tier tests (skipped when torchvision is absent).
|
|
|
|
Gate the real-image confirmation tier: the oracle measures modes accurately, thickness binning
|
|
and mode assignment are well-formed, the sampler draws the requested modes, the VAE clears a
|
|
floor-aware gen-0 fidelity check, and a short dry lineage collapses more than a grounded one.
|
|
These are sign checks (the tier is confirmation-only), and they reuse the cached classifier so
|
|
they stay fast.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
pytest.importorskip("torchvision")
|
|
|
|
from inheritance.metrics import forward_kl # noqa: E402
|
|
from neural.config import MnistCfg, ModelCfg, OracleCfg # noqa: E402
|
|
from neural.mnist_data import ( # noqa: E402
|
|
MnistSampler, assign_modes, fit_thickness_thresholds, make_mnist_truth, thickness_bin,
|
|
)
|
|
from neural.mnist_loop import run_mnist_lineage # noqa: E402
|
|
from neural.mnist_oracle import build_oracle, confusion_summary # noqa: E402
|
|
from neural.mnist_vae import ConvVAEGenerator # noqa: E402
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def world():
|
|
"""Load MNIST + build the oracle/sampler/truth once (cached CNN -> fast)."""
|
|
from neural.mnist_data import load_mnist
|
|
|
|
cfg = MnistCfg()
|
|
data = load_mnist(cfg.data_root)
|
|
td = make_mnist_truth(cfg)
|
|
oracle, cuts, _ = build_oracle(cfg, OracleCfg(), data, seed=1)
|
|
modes = assign_modes(data.train_x, data.train_y, cuts, cfg)
|
|
sampler = MnistSampler(data.train_x, modes, cfg.K)
|
|
return cfg, data, td, oracle, cuts, sampler
|
|
|
|
|
|
def test_thickness_bins_and_mode_bijection(world):
|
|
cfg, data, td, oracle, cuts, sampler = world
|
|
sub = data.train_x[:2000]
|
|
lab = data.train_y[:2000]
|
|
bins = thickness_bin(sub, lab, cuts)
|
|
assert bins.min() >= 0 and bins.max() < cfg.style_bins # bins in range
|
|
modes = assign_modes(sub, lab, cuts, cfg)
|
|
assert np.array_equal(modes, lab * cfg.style_bins + bins) # exact bijection
|
|
assert modes.min() >= 0 and modes.max() < cfg.K
|
|
|
|
|
|
def test_sampler_draws_requested_modes(world):
|
|
cfg, data, td, oracle, cuts, sampler = world
|
|
counts = np.zeros(cfg.K, dtype=int)
|
|
counts[5] = 30 # mode 5 -> class 5//style_bins
|
|
counts[17] = 10 # mode 17 -> class 17//style_bins
|
|
imgs = sampler.draw(counts, np.random.default_rng(0))
|
|
assert imgs.shape == (40, 1, 28, 28)
|
|
# The sampler draws from the true per-mode pools, so the oracle should recover the two
|
|
# requested classes for the large majority of the draw (allowing the ~2% oracle error).
|
|
pred_classes = oracle.classify(imgs) // cfg.style_bins
|
|
want = {5 // cfg.style_bins, 17 // cfg.style_bins}
|
|
assert np.mean(np.isin(pred_classes, list(want))) > 0.9
|
|
|
|
|
|
def test_oracle_mode_accuracy_high(world):
|
|
cfg, data, td, oracle, cuts, sampler = world
|
|
conf = confusion_summary(oracle, data, cuts, cfg)
|
|
assert conf["mode_accuracy"] > 0.95 # solid measurement-noise floor
|
|
assert conf["class_accuracy"] > 0.97
|
|
|
|
|
|
def test_vae_gen0_recovers_full_support(world):
|
|
# gen-0 fidelity gate (floor-aware): the VAE must represent every mode (incl. the tail),
|
|
# else later tail loss would be underfitting, not collapse. Frequencies over-smooth (KL~0.5).
|
|
cfg, data, td, oracle, cuts, sampler = world
|
|
rng = np.random.default_rng(0)
|
|
mcfg = ModelCfg(kind="convvae", latent=32, epochs=15, n_eval=8000)
|
|
X0 = sampler.draw(rng.multinomial(6000, td.p_star), rng)
|
|
vae = ConvVAEGenerator(cfg, mcfg, oracle)
|
|
vae.fit(X0, rng)
|
|
p_hat = vae.mode_distribution(rng)
|
|
assert float(np.mean(p_hat[td.tail_mask] > 1e-9)) > 0.8 # tail represented
|
|
assert forward_kl(td.p_star, p_hat, 1e-9) < 1.2 # loose, floor-aware
|
|
|
|
|
|
def test_dry_collapses_more_than_grounded(world):
|
|
# The N1+N2 signs on real images: dry end-KL > grounded end-KL. Tiny config for speed.
|
|
cfg, data, td, oracle, cuts, sampler = world
|
|
base = {
|
|
"mnist": cfg.__dict__,
|
|
"model": ModelCfg(kind="convvae", latent=32, epochs=20, n_eval=6000).__dict__,
|
|
"generations": 8,
|
|
"metrics": {"kl_floor": 1e-9, "support_eps": 1e-9},
|
|
}
|
|
|
|
def run(g):
|
|
n = 6000
|
|
m = 0 if g == 0 else round(n * g / (1 - g))
|
|
c = {**base, "dynamics": {"n": n, "grounding": {"m": m, "policy": "proportional"}}}
|
|
return run_mnist_lineage(c, seed=0, oracle=oracle, sampler=sampler, td=td)
|
|
|
|
dry, grd = run(0.0), run(0.1)
|
|
assert dry["forward_kl"].iloc[-1] > dry["forward_kl"].iloc[0] + 1.0 # dry collapses
|
|
assert grd["forward_kl"].iloc[-1] < dry["forward_kl"].iloc[-1] # grounding arrests it
|
|
assert grd["support_size"].iloc[-1] > dry["support_size"].iloc[-1] # keeps more modes
|