Metadata-Version: 2.4
Name: reionemu
Version: 0.3.0
Summary: Machine-learning emulator pipeline for the kinetic Sunyaev-Zel'dovich angular power spectrum, with uncertainty quantification.
Project-URL: Homepage, https://robertxpearce.github.io/reionization-emulator/
Project-URL: Repository, https://github.com/RobertxPearce/reionization-emulator
Project-URL: Issues, https://github.com/RobertxPearce/reionization-emulator/issues
Author-email: Robert Pearce <robertbdpearce@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: astrophysics,cosmology,kSZ,machine learning,reionization
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Astronomy
Requires-Python: >=3.10
Requires-Dist: h5py
Requires-Dist: numpy
Requires-Dist: ray[tune]>=2.0
Requires-Dist: torch
Provides-Extra: dev
Requires-Dist: build; extra == 'dev'
Requires-Dist: ipykernel; extra == 'dev'
Requires-Dist: jupyter; extra == 'dev'
Requires-Dist: matplotlib; extra == 'dev'
Requires-Dist: mkdocs-material>=9.5; extra == 'dev'
Requires-Dist: mkdocs>=1.6; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: scipy; extra == 'dev'
Requires-Dist: twine; extra == 'dev'
Provides-Extra: docs
Requires-Dist: mkdocs-material>=9.5; extra == 'docs'
Requires-Dist: mkdocs>=1.6; extra == 'docs'
Description-Content-Type: text/markdown

[![CI](https://github.com/RobertxPearce/reionization-emulator/actions/workflows/ci.yml/badge.svg)](https://github.com/RobertxPearce/reionization-emulator/actions/workflows/ci.yml)
[![Docs](https://img.shields.io/badge/docs-online-blue)](https://robertxpearce.github.io/reionization-emulator/)
[![PyPI](https://img.shields.io/pypi/v/reionemu)](https://pypi.org/project/reionemu/)
[![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13-blue)](https://www.python.org/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
<!-- [![DOI](https://zenodo.org/badge/1051925247.svg)](https://zenodo.org/badge/latestdoi/1051925247) -->


<p align="center">
    <img src="https://raw.githubusercontent.com/RobertxPearce/reionization-emulator/main/docs/assets/reionemu-logo.png" alt="reionemu logo" width="300">
</p>

# reionemu

A modular Python package for building machine-learning emulators of the kinetic Sunyaev-Zel'dovich (kSZ) angular power spectrum from kSZ 2LPT reionization simulations. It includes tools to condense simulation outputs, compute flat-sky power spectra, assemble training datasets, train neural networks that predict binned rescaled kSZ power spectra from reionization parameters, and save lightweight experiment artifacts for reproducibility.

The goal is to learn a fast surrogate model that maps reionization parameters → binned kSZ power spectrum, enabling rapid exploration of cosmological parameter space without re-running expensive simulations.

---

## Installation

```bash
pip install reionemu
```

Or from source (editable):

```bash
git clone https://github.com/RobertxPearce/reionization-emulator.git
cd reionization-emulator
python -m pip install -e .
```

**Requirements:** Python 3.10+, NumPy, HDF5, PyTorch, and Ray Tune.

---

## Quick start

After installing, you can load a processed HDF5 training dataset, create dataloaders, and train the baseline deterministic 4-parameter emulator:

```python
from pathlib import Path
import torch
import reionemu

# Path to a condensed HDF5 that already has /training (X, Y, ell)
h5_path = Path("path/to/condensed.h5")

# Dataloaders with train/val split and optional normalization
loaders, normalizers, ell = reionemu.make_dataloaders(
    h5_path,
    split={"train": 0.8, "val": 0.2},
    config=reionemu.DataLoaderConfig(batch_size=32, seed=42),
)

# Baseline 4-parameter model, optimizer, loss
model = reionemu.FourParamEmulator()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = torch.nn.MSELoss()

# Train for a few epochs
history = reionemu.fit(
    model,
    loaders["train"],
    loaders["val"],
    optimizer,
    loss_fn,
    config=reionemu.FitConfig(epochs=10, device="cpu"),
)

# Validation loss per epoch
print(history["val_loss"])

# Save a lightweight experiment artifact
artifact_dir = reionemu.save_artifact(
    "baseline_four_param",
    Path("artifacts"),
    dataset_path=h5_path,
    dataloader_config=reionemu.DataLoaderConfig(batch_size=32, seed=42),
    fit_config=reionemu.FitConfig(epochs=10, device="cpu"),
    model_config={
        "class_name": "FourParamEmulator",
        "input_dim": 4,
        "output_dim": 5,
    },
    optimizer_config={"name": "Adam", "lr": 1e-3},
    history=history,
    normalizers=normalizers,
    checkpoint=model.state_dict(),
)
```

For MC-dropout experiments, use `MCDropoutEmulator` with the MC evaluation path:

```python
model = reionemu.MCDropoutEmulator(dropout_rate=0.2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

history = reionemu.fit(
    model,
    loaders["train"],
    loaders["val"],
    optimizer,
    torch.nn.MSELoss(),
    config=reionemu.FitConfig(epochs=10, device="cpu"),
    evaluation="evaluate_mc_metrics",
    n_mc_samples=50,
)

print(history["val_mean_predictive_std"])
```

Once trained, use `predict_mc` to get predictions and uncertainties in physical units. Dropout is redrawn on every call, so seed torch first when you need a reproducible number:

```python
import numpy as np
import torch

theta = np.array([[8.0, 0.5, 1.0, 0.45]], dtype=np.float32)

torch.manual_seed(42)
pred_mean, pred_std, samples_dl, samples_log = reionemu.predict_mc(
    theta,
    model,
    X_mean=normalizers["X"].mean,
    X_std=normalizers["X"].std,
    normalize_X=True,
    n_mc_samples=200,
)

print(pred_mean, pred_std)
```

If you want to tune the four-parameter architecture with Ray Tune before training a final model, you can work directly with the loaded arrays:

```python
from pathlib import Path

import reionemu
from ray import tune

h5_path = Path("path/to/condensed.h5")
X, Y, ell = reionemu.load_training_arrays(h5_path)

split_idx = int(0.8 * len(X))
X_train, X_val = X[:split_idx], X[split_idx:]
Y_train, Y_val = Y[:split_idx], Y[split_idx:]

param_space = {
    "hidden_dim": tune.choice([20, 32, 64]),
    "num_hidden_layers": tune.choice([1, 2, 3]),
    "activation": tune.choice(["relu", "silu", "tanh"]),
    "optimizer": tune.choice(["adam", "adamw"]),
    "lr": tune.loguniform(3e-4, 2e-3),
    "weight_decay": tune.loguniform(1e-8, 1e-4),
    "batch_size": tune.choice([16, 32, 64]),
    "epochs": 150,
    "early_stopping_patience": tune.choice([10, 15]),
    "gradient_clipping": tune.choice([None, 0.5, 1.0]),
    "normalize_X": True,
    "normalize_Y": False,
}

results = reionemu.run_tune_four_param(
    X_train=X_train,
    Y_train=Y_train,
    X_val=X_val,
    Y_val=Y_val,
    param_space=param_space,
    num_samples=20,
    max_concurrent_trials=2,
    device="cpu",
    storage_path="ray_results",
    experiment_name="four_param_search",
)

best = results.get_best_result(metric="val_loss", mode="min")
print(best.config)
print(best.metrics["best_val_loss"])
```

For a full pipeline example (condense → compute power spectra → build training data → tune/train/evaluate), scientific context, and complete usage examples, see the full documentation: [Homepage](https://robertxpearce.github.io/reionization-emulator/)

---

## Scientific context

The kinetic Sunyaev-Zel'dovich (kSZ) effect arises from the scattering of CMB photons by free electrons with bulk motion, generating secondary temperature anisotropies. The kSZ angular power spectrum carries information about the timing, duration, and structure of reionization. This emulator provides a fast surrogate that maps reionization parameters (zmean_zre, alpha_zre, kb_zre, b0_zre) to binned, rescaled kSZ power spectra, making parameter-space exploration much faster than rerunning the full simulations.

---

## Repository structure

| Path                     | Description                                                                            |
|--------------------------|----------------------------------------------------------------------------------------|
| **`src/reionemu/`**      | Core library (pip-installable package)                                                 |
| `src/reionemu/simio/`    | Simulation I/O, power spectrum computation, training-array building                    |
| `src/reionemu/data/`     | Dataloaders, normalization                                                             |
| `src/reionemu/artifact/` | JSON experiment manifests, config/results saving, normalizer and checkpoint sidecars   |
| `src/reionemu/models/`   | Baseline and experimental emulator architectures                                       |
| `src/reionemu/training/` | Training loop, K-fold cross-validation, metrics, and model builders                    |
| `src/reionemu/tuning/`   | Ray Tune integration for hyperparameter search                                         |
| **`scripts/`**           | Dataset builder, HPC runners, sampling (environment-specific)                          |
| **`notebooks/`**         | Analysis and training examples                                                         |
| **`docs/`**              | Documentation source code                                                              |
| `datasets/`              | Raw and processed datasets (not tracked)                                               |
| `results/`               | Visualizations for simulation checks, parameter-space validation, and model evaluation |

The **core API** is in `src/reionemu/`. Scripts under `scripts/hpc/` and `scripts/sampling/` are for cluster and sampling workflows and may use machine-specific paths; the library itself is portable.

---

## Main public API

Import from the top-level package after `pip install reionemu`:

- **Simulation I/O:** `condense_sim_root`, `CondenseConfig`, `add_cl_to_condensed_h5`, `ClConfig`, `build_and_write_training`, `build_training_arrays`, `BuildXYConfig`, `BuildStats`, `CondenseStats`
- **Data:** `make_dataloaders`, `load_training_arrays`, `DataLoaderConfig`, `Normalizer`
- **Artifacts:** `create_artifact_dir`, `save_artifact`, `save_configs`, `save_results`, `save_info`, `save_normalizers`, `load_normalizers`, `save_model_checkpoint`, `dataset_summary`, `file_fingerprint`, `read_json`
- **Models:** `FourParamEmulator`, `MCDropoutEmulator`, `predict_mc`, `enable_dropout_only` (experimental variants live in `reionemu.models.experimental`)
- **Training:** `fit`, `FitConfig`, `train_one_epoch`, `evaluate`, `evaluate_metrics`, `evaluate_mc_metrics`, `kfold_cross_validate`, `KFoldConfig`
- **Training helpers:** `build_four_param_model`, `build_mc_dropout_model`, `build_optimizer`, `mse`, `rmse`, `mean_relative_error`, `physical_mean_relative_error`
- **Tuning:** `train_four_param_tune`, `default_param_space`, `run_tune_four_param`

For full API reference, module documentation, and usage guides, visit: [Homepage](https://robertxpearce.github.io/reionization-emulator/)

---

## Typical workflow

1. **Parameter sampling** - Latin Hypercube Sampling over the 4D reionization parameter space.
2. **Simulation (HPC)** - Run Zreion (or compatible) simulations; outputs per sim in HDF5.
3. **Dataset construction** - Use `condense_sim_root` → `add_cl_to_condensed_h5` → `build_and_write_training` to produce a single condensed HDF5 with `/sims` and `/training`.
4. **Hyperparameter search (optional)** - Use `load_training_arrays` and `run_tune_four_param` to search over model and optimizer settings with Ray Tune.
5. **Training and evaluation** - Use `make_dataloaders` and `fit` (or `kfold_cross_validate`) to train and evaluate the selected emulator configuration.
6. **Artifact saving** - Use `save_artifact` to record JSON configs/results plus optional `.npz` normalizers and `.pt` model checkpoints.

---

## Acknowledgments

This research is conducted in the LEADS Lab at the University of Nevada, Las Vegas, under [Dr. Paul La Plante](https://plaplant.github.io/), with computing resources from the Pittsburgh Supercomputing Center (Bridges-2).
