Metadata-Version: 2.4
Name: cosmufr
Version: 0.7.3
Summary: CosmUFR: Cosmological Parameter Inference via Unified Field Representations. 125M-parameter belief-settling neural emulator. Recovers 8 cosmological parameters from matter power spectrum P(k).
Home-page: https://github.com/arajgor1/calybre-ufr
Author: Aaditya Rajgor
Author-email: arajgor1@gmail.com
License: MIT
Project-URL: GitHub, https://github.com/arajgor1/calybre-ufr
Project-URL: HuggingFace, https://huggingface.co/arajgor1/cosmufr-v07
Keywords: cosmology,machine learning,power spectrum,neural network,parameter inference,MCMC,emulator,DESI,Euclid,Planck
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Astronomy
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.24.0
Requires-Dist: scipy>=1.10.0
Requires-Dist: torch>=2.0.0
Requires-Dist: astropy>=5.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0; extra == "dev"
Requires-Dist: twine>=4.0; extra == "dev"
Requires-Dist: build>=0.10; extra == "dev"
Provides-Extra: viz
Requires-Dist: matplotlib>=3.7; extra == "viz"
Provides-Extra: demo
Requires-Dist: gradio>=4.0; extra == "demo"
Requires-Dist: matplotlib>=3.7; extra == "demo"
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: keywords
Dynamic: license
Dynamic: project-url
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

<div align="center">

<img src="https://img.shields.io/badge/CosmUFR-v0.7.1-blueviolet?style=for-the-badge" />
<img src="https://img.shields.io/badge/PyPI-cosmufr-blue?style=for-the-badge&logo=pypi&logoColor=white" />
<img src="https://img.shields.io/badge/PyTorch-2.x-EE4C2C?style=for-the-badge&logo=pytorch&logoColor=white" />
<img src="https://img.shields.io/badge/Status-Active%20Research-brightgreen?style=for-the-badge" />

<br/><br/>

# 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.**

<br/>

[Quick Start](#quick-start) · [Architecture](#architecture) · [Training Data](#training-data) · [Results](#results) · [API Reference](#api-reference) · [Sprint Plan](#sprint-plan)

</div>

---

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

```bash
pip install cosmufr
```

**Recover parameters from a P(k) observation:**
```python
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:**
```python
# 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):**
```python
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):**
```python
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:**
```python
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

```python
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](https://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`

```python
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](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](https://huggingface.co/arajgor1/cosmufr-v07) · PyPI: [cosmufr](https://pypi.org/project/cosmufr/)
