cosmufr 0.7.3

CosmUFR: Cosmological Parameter Inference via Unified Field Representations. 125M-parameter belief-settling neural emulator. Recovers 8 cosmological parameters from matter power spectrum P(k).

CosmUFR

Cosmological Parameter Inference via Unified Field Representations

125M-parameter belief-settling neural emulator. Recovers 8 cosmological parameters from matter power spectrum P(k) in < 40ms. Uncertainty-quantified. Sequentially updatable.

Quick Start · Architecture · Training Data · Results · API Reference · Sprint Plan

---

What It Does

CosmUFR takes a matter power spectrum P(k) as input and recovers:

• 8 cosmological parameters — Ω_m, σ_8, h, n_s, Ω_b, w₀, mν, w_a — with calibrated 1-sigma uncertainties
• Reconstructed P(k) at 200 k-bins across [0.1, 4.5] h/Mpc
• Halo mass function n(M) — 20 bins from 10¹¹ to 10¹⁵ M☉/h
• Constraint energy score — flags anomalous or out-of-distribution inputs

What takes a classical MCMC pipeline hours, CosmUFR does in one forward pass (~36ms on A100, ~400ms on CPU).

The key architectural difference: CosmUFR doesn't regress parameters directly. It forms a belief state about the universe, settles that belief via energy minimisation over 16 steps, then reads off parameters from the settled state. This allows sequential updating — feeding in Planck data then DESI data produces a belief that reflects both, with uncertainty narrowing as evidence accumulates.

---

Quick Start

pip install cosmufr

Recover parameters from a P(k) observation:

import cosmufr
import numpy as np

# Load model (requires checkpoint file — download from HuggingFace below)
model = cosmufr.load("checkpoints/best.pt")

# pk_z0 and pk_z047: log10 P(k) at z=0 and z≈0.47, shape [200]
# k-grid: 200 bins, log-spaced from 0.1 to 4.5 h/Mpc
result = model.predict(pk_z0, pk_z047)

print(result.summary())
# CosmUFR Parameter Recovery:
# ────────────────────────────────────────
#   Omega_m    =  +0.3012  ±  0.0082
#   sigma_8    =  +0.8134  ±  0.0124
#   h          =  +0.6731  ±  0.0091
#   ...
#   E_con      =  -0.42σ   (typical)

With real telescope data:

# Interpolate any observed P(k) onto the model's k-grid
from cosmufr.utils import interpolate_to_model_grid

# k_obs, pk_obs can be from DESI, Planck, Euclid, or any survey
pk_model = interpolate_to_model_grid(k_obs, pk_obs)

result = model.predict(pk_model, pk_model_z047)

Sequential update (Hubble tension demo):

session = cosmufr.CosmUFRSession(model)

r1 = session.add_observation(*cosmufr.data.load_planck2018(), label="Planck 2018")
r2 = session.add_observation(*cosmufr.data.load_desi_dr2(),   label="DESI DR2")

evol = session.param_evolution()
print(evol["h"])          # [0.674, 0.680] — Hubble tension visible
print(evol["Omega_m"])    # parameter shift across surveys

# Uncertainty narrows as evidence accumulates
unc = session.uncertainty_evolution()
print(unc["sigma_8"])     # [0.012, 0.009] — tighter after DESI

Model comparison (ΛCDM vs w₀w_aCDM):

from cosmufr import compare

result = compare(
    model,
    lcdm_pk_z0, lcdm_pk_z047,    # Model A: ΛCDM
    wcdm_pk_z0, wcdm_pk_z047,    # Model B: w₀w_aCDM (DESI-hinted)
)
print(result.preferred_model)       # "B"
print(result.sigma_significance)    # 2.3  (σ units, requires E_con calibration)

Cosmic web visualisation:

from cosmufr.viz import render_cosmic_web
import matplotlib.pyplot as plt

field = render_cosmic_web(result.pk_reconstructed, model.k_grid)  # [512, 512]
plt.imshow(field, cmap='inferno', origin='lower')
plt.title("Inferred Cosmic Web — DESI DR2")
plt.show()

---

Download Checkpoint

from huggingface_hub import hf_hub_download

path = hf_hub_download(repo_id="arajgor1/cosmufr-v07", filename="best.pt")
model = cosmufr.load(path)

Or manually from huggingface.co/arajgor1/cosmufr-v07.

---

Architecture

CosmUFR is a 125M-parameter belief-settling emulator. No attention. No transformers. No MoE.

Input: log10 P(k) at z=0 + z≈0.47  →  [400-dim observation]
          │
          ▼
  ObsEncoder (22M params)         — 400d → 1024d belief encoding
          │
          ▼
  BeliefProposal (21M params)     — (z_t, b_prev) → b_hat [1024d]
          │
          ▼
  SettlingCore × 16 steps (3M)    — energy gradient descent on b_hat
  ┌──────────────────────────────┐
  │  E(b) = w_obs·E_obs          │  ← ObsEnergyHead (6M)
  │       + w_con·E_con          │  ← ConstraintHead via AttractorBank (4.2M)
  │       + w_dyn·E_dyn          │  ← DynEnergyHead (6M)
  └──────────────────────────────┘
          │
          ▼
     b_star [1024d]              — settled belief
    ┌──────┼──────────┬──────────┐
    ▼      ▼          ▼          ▼
 Params  Uncert.  GenerativeHead  HaloMassHead
  [8]    [8]      P(k) at any k    n(M) [20]
                  (k-continuous)

Belief state decomposition: b = [s | c | u | p] — 512d world-state, 256d context, 128d uncertainty, 128d prototype proximity.

AttractorBank: 4096 prototype belief vectors, top-16 cosine similarity, EMA updates (α=0.99). Provides the constraint energy signal and anomaly detection.

GenerativeHead: k-continuous implicit field — queries P(k) at any k value, not just the training grid. 2048-wide, 8-block residual MLP, 25M params.

---

Training Data

57,528,155 total samples from five simulation suites:

Source	Type	Samples	Params varied
BACCO	Emulator	~30M	Ω_m, σ_8, h, n_s, Ω_b
BCemu	Baryon emulator	~10M	Ω_m, σ_8 + baryonic
DarkEmulator	N-body	~1M	Ω_m, σ_8, w₀
SP(k)	Semi-analytic	~10M	Ω_m, σ_8
CAMB NL	Boltzmann	~6.5M	Ω_m, σ_8, h, n_s, Ω_b, mν

Known gap: Only DarkEmulator (1.7% of data) varies w₀. No suite independently varies w_a. The w₀w_aCDM comparison in the demo is directional but not quantitatively reliable until Run 3 (planned: 10× upsampling of dark energy simulations).

---

Results (Run 1)

Trained on A100 40GB — Phase 3, Epoch 12 best checkpoint.

Parameter	R²	Notes
Ω_m	0.901	Demo-ready
σ_8	~0.85	Demo-ready
h	TBD	Full eval in Sprint 0
n_s	TBD	Full eval in Sprint 0
w₀	Low	Limited training coverage
w_a	Low	No training coverage — Run 3 needed
mν	~0.6	Improving with 3× upsampling in Run 2

Run 2 (in preparation): BF16, batch 2048, 16 epochs per phase, ECE bug fixed, 3× neutrino upsampling, 2× learning rate. Expected improvement: R²(Ω_m) → 0.95+.

---

API Reference

cosmufr.load(path, device="auto") → CosmUFR

Load a checkpoint. device is "cuda", "cpu", or "auto".

CosmUFR.predict(pk_z0, pk_z047, b_prev=None) → CosmUFRResult

Run inference. Both inputs: np.ndarray shape [200], log10 P(k) on model k-grid.

CosmUFRResult

Field	Type	Description
params	ndarray [8]	[Ω_m, σ_8, h, n_s, Ω_b, w₀, mν, w_a] in physical units
uncertainties	ndarray [8]	1-sigma per parameter, physical units
pk_reconstructed	ndarray [200]	Reconstructed log10 P(k) at z=0
belief	ndarray [1024]	Settled belief vector b_star
energy_log	list[float]	Energy at each of 16 settling steps + final
e_con	float	Raw constraint energy (AttractorBank prototype loss)
e_con_zscore	float | None	Anomaly score in σ — None until calibration run
nmass	ndarray [20]	log10 n(M) halo mass function

CosmUFRSession

session = CosmUFRSession(model)
result  = session.add_observation(pk_z0, pk_z047, label="DESI DR2")
session.reset()
evol    = session.param_evolution()        # dict[str, ndarray]
unc     = session.uncertainty_evolution()  # dict[str, ndarray]

compare(model, pk_z0_a, pk_z047_a, pk_z0_b, pk_z047_b) → CompareResult

Compare two cosmological models. Returns delta_E, delta_params, sigma_significance, preferred_model.

render_cosmic_web(pk, k_grid, N=512, smoothing_sigma=1.5, seed=42) → ndarray [N,N]

Generate 2D lognormal density field from P(k). No GPU required.

interpolate_to_model_grid(k_obs, pk_obs, log10_output=True) → ndarray [200]

Map any observed P(k) onto model k-grid. Cubic spline in log-log space. Issues warning if k-coverage < 90%.

---

Repository Structure

calybre-ufr/
├── cosmufr/                  # PyPI package — user-facing inference API
│   ├── __init__.py
│   ├── inference.py          # CosmUFRResult, CosmUFR, load()
│   ├── session.py            # CosmUFRSession — sequential inference
│   ├── compare.py            # compare(), CompareResult
│   ├── viz/
│   │   └── cosmic_web.py     # render_cosmic_web()
│   ├── utils/
│   │   └── interpolation.py  # interpolate_to_model_grid()
│   └── data/
│       └── loaders.py        # load_desi_dr2(), load_planck2018(), load_euclid_q1()
├── model/                    # Model architecture + training
│   ├── cosmufr_arch.py       # CosmUFRLite — full model, all components
│   ├── cosmufr_config.py     # CosmUFRConfig — all hyperparameters
│   ├── cosmufr_losses.py     # param_loss, nll_loss, pk_loss, energy_margin_loss
│   └── train_vertex.py       # 5-phase training loop for Vertex AI
├── scripts/                  # Data generation + infrastructure
│   ├── submit_vertex_job.py  # Submit training job to Vertex AI
│   ├── calibrate_econ.py     # (Sprint 1) E_con baseline calibration
│   └── eval_all_params.py    # (Sprint 0) Full 8-param evaluation
├── experiments/              # Benchmark results
├── configs/                  # Training configs
├── checkpoints/              # Local checkpoint storage (not in git — use HF)
├── TRAINING_RUN2.md          # Run 2 specification + product vision gaps
└── SPRINT_PLAN.md            # Full 5-sprint product build plan

---

Sprint Plan

Sprint	Days	Deliverable
0	1–2	Full 8-param eval on Run 1, DESI real-data sanity check
1 ✅	3–7	cosmufr package v0.7.1 on PyPI, Run 2 submitted
2	8–14	Real data loaders (DESI/Planck/Euclid), CosmUFRSession
3	15–21	render_cosmic_web(), compare(), notebook smoke test
4	22–30	Gradio Space live at huggingface.co/spaces/arajgor1/cosmufr-demo
5	Post-Run 2	arXiv preprint (astro-ph.CO), researcher outreach

Full detail: SPRINT_PLAN.md

---

Known Limitations

• w₀/w_a recovery: Only 1.7% of training data varies w₀; none varies w_a. Model comparison for dark energy is directional only — not quantitatively reliable until Run 3.
• E_con calibration: e_con_zscore is None until scripts/calibrate_econ.py is run on a GPU against the validation set. Raw e_con values are the AttractorBank prototype loss and require a reference distribution to be interpretable.
• k-range: Model is trained on k ∈ [0.1, 4.5] h/Mpc. Extrapolation beyond this range is flat (constant P(k)) and unreliable.
• Cross-sim generalisation: Trained on BACCO/BCemu/DarkEmulator/SP(k)/CAMB NL. Not yet tested on IllustrisTNG or SIMBA directly.
• Data loaders: load_desi_dr2(), load_planck2018(), load_euclid_q1() return synthetic placeholders in v0.7.1. Real implementations in Sprint 2.

---

Contact

Aaditya Rajgor — research inquiries and collaboration via GitHub issues.

HuggingFace: arajgor1/cosmufr-v07 · PyPI: cosmufr