Metadata-Version: 2.4
Name: mockcraft
Version: 0.1.2
Summary: Reusable Python package for multimodal astronomical source generation using AION
Author: Bibiana Terres Stumpf, Joao Victor Motta da Silva, Pedro Jann Luna, Anne Laure Mealier, Julien Zoubian
License-Expression: MIT
Project-URL: Repository, https://github.com/CentraleDigitaleLab/mockcraft
Project-URL: Issues, https://github.com/CentraleDigitaleLab/mockcraft/issues
Project-URL: Documentation, https://github.com/CentraleDigitaleLab/mockcraft/tree/main/docs
Keywords: astronomy,astrophysics,machine-learning,generative-ai,multimodal,synthetic-catalogues,AION
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: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Astronomy
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: <3.15,>=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy<3,>=1.26
Requires-Dist: pandas<3.0,>=2.2
Requires-Dist: scipy<2,>=1.11
Requires-Dist: scikit-learn<2,>=1.5
Requires-Dist: astropy<8,>=6.1
Requires-Dist: astroquery<1,>=0.4.7
Requires-Dist: matplotlib<4,>=3.8
Requires-Dist: torch<3,>=2.4
Requires-Dist: einops<1,>=0.8
Requires-Dist: jaxtyping<0.4,>=0.2.28
Requires-Dist: huggingface-hub<1.0,>=0.23
Requires-Dist: tokenizers<1,>=0.19
Requires-Dist: transformers<5,>=4.40
Requires-Dist: datasets<4,>=2.18
Requires-Dist: safetensors<1,>=0.4.3
Requires-Dist: polymathic-aion==0.0.2
Requires-Dist: filelock<4,>=3.13
Requires-Dist: tqdm<5,>=4.66
Requires-Dist: typing-extensions<5,>=4.9
Provides-Extra: viz
Requires-Dist: matplotlib<4,>=3.8; extra == "viz"
Requires-Dist: umap-learn<0.6,>=0.5.5; extra == "viz"
Provides-Extra: notebooks
Requires-Dist: ipython<9,>=8.20; extra == "notebooks"
Requires-Dist: ipykernel<7,>=6.29; extra == "notebooks"
Requires-Dist: jupyter<2,>=1.0; extra == "notebooks"
Requires-Dist: ipywidgets<9,>=8.1; extra == "notebooks"
Requires-Dist: matplotlib<4,>=3.8; extra == "notebooks"
Requires-Dist: umap-learn<0.6,>=0.5.5; extra == "notebooks"
Requires-Dist: tqdm<5,>=4.66; extra == "notebooks"
Provides-Extra: data
Requires-Dist: pyarrow<20,>=15.0; extra == "data"
Requires-Dist: datasets<4,>=2.18; extra == "data"
Provides-Extra: docs
Requires-Dist: Sphinx<9,>=7.2; extra == "docs"
Requires-Dist: docutils<0.22,>=0.20; extra == "docs"
Provides-Extra: network
Requires-Dist: redis<8,>=5.0; extra == "network"
Requires-Dist: urllib3<3,>=2.0; extra == "network"
Requires-Dist: requests<3,>=2.31; extra == "network"
Provides-Extra: security
Requires-Dist: cryptography<47,>=42.0; extra == "security"
Requires-Dist: pyOpenSSL<26,>=24.0; extra == "security"
Dynamic: license-file

# mockcraft

`mockcraft` is a Python package for generating synthetic astronomical catalogues using the [AION](https://github.com/PolymathicAI/aion) foundation model. It provides a simple interface for cross-modal generation across astronomical data types and includes a catalogue utility for fetching real sources from Gaia DR3, DESI DR1, and Legacy Survey DR8.

---

## What is MockCraft?

MockCraft is a pipeline for generating synthetic astronomical mock catalogues using foundation models. Given a set of real astronomical observations as input — redshifts, photometric fluxes, or images — MockCraft can generate realistic synthetic counterparts: spectra, multi-band images, and morphological parameters.

[AION](https://github.com/PolymathicAI/aion) is a multimodal foundation model trained on one of the largest astronomical datasets ever assembled (see the [Multimodal Universe paper](https://arxiv.org/pdf/2412.02527)). It learns joint representations across spectra, images, and physical parameters, enabling cross-modal generation: given any subset of observables, it can generate any other. MockCraft uses AION as its generative backbone.

---

## Installation

Requires Python 3.10–3.12.

```bash
pip install mockcraft
```

**Optional extras:**

| Extra | When to use |
|-------|-------------|
| `viz` | Visualization helpers (`plot_xp_spectrum`, `plot_image`) and UMAP projections |
| `notebooks` | Jupyter notebook workflows |
| `data` | Local Parquet files or HuggingFace datasets |

```bash
pip install mockcraft[viz]
pip install mockcraft[notebooks]
pip install mockcraft[viz,notebooks]   # both
```

**NVIDIA GPU (CUDA 12.4, Linux) — install PyTorch before the package:**

```bash
pip install torch --index-url https://download.pytorch.org/whl/cu124
pip install mockcraft
```

**From source:**

```bash
git clone https://github.com/CentraleDigitaleLab/mockcraft.git
cd mockcraft
pip install -e ".[viz]"
```

---

## Quick Start

```python
from mockcraft import SourceGenerator

gen = SourceGenerator(model="aion", seed=42)

# Redshift → spectrum + image
result = gen.generate(
    inputs={"redshift": 0.3},
    outputs=["spectrum", "image"],
)

print(result.spectrum.shape)  # (8704,)
print(result.image.shape)     # (3, 128, 128)
```

---

## API Reference

### `SourceGenerator`

```python
from mockcraft import SourceGenerator

gen = SourceGenerator(model="aion", seed=42)
```

| Parameter | Type | Description |
|-----------|------|-------------|
| `model` | `str` | Model identifier. Only `"aion"` is currently supported. |
| `seed` | `int` or `None` | Fixed random seed for reproducibility. |
| `device` | `str` or `None` | PyTorch device: `"cuda"`, `"mps"`, or `"cpu"`. Auto-detected if not set. |

---

### `generate(inputs, outputs, cfg=None, type=None, compute_embeddings=False)`

The primary generation method. Accepts any combination of input and output modalities.

```python
result = gen.generate(
    inputs={"redshift": 0.3},
    outputs=["spectrum", "image"],
)
```

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `inputs` | `dict` | — | Modality name → value. Floats for scalars, numpy arrays for spectra/images. |
| `outputs` | `list[str]` | — | List of modality names to generate. |
| `cfg` | `float` or `None` | `None` | Classifier-free guidance override. Uses per-modality defaults if not set. |
| `type` | `str` or `None` | `None` | Morphological type prior: `"elliptical"`, `"spiral"`, or `"irregular"`. |
| `compute_embeddings` | `bool` | `False` | If `True`, computes AION latent embeddings `(768,)` for generated spectra and images. |

Returns a `GeneratedSource` object (see below).

---

### `star(temperature, logg=None, metallicity=None)`

Generate a synthetic stellar XP spectrum conditioned on effective temperature.

```python
result = gen.star(temperature=5778.0)

print(result.xp_bp.shape)  # (55,)
print(result.xp_rp.shape)  # (55,)
```

| Parameter | Type | Description |
|-----------|------|-------------|
| `temperature` | `float` | Effective temperature in Kelvin. |
| `logg` | `float` or `None` | Surface gravity (accepted for API compatibility, not yet used as conditioning). |
| `metallicity` | `float` or `None` | Metallicity [Fe/H] (accepted for API compatibility, not yet used as conditioning). |

Temperature is converted to approximate Gaia G/BP/RP fluxes via bolometric scaling relative to a solar reference (T☉ = 5778 K). Returns a `GeneratedSource` with `xp_bp` and `xp_rp` outputs.

---

## Supported Modality Keys

| Key | Type | Description |
|-----|------|-------------|
| `redshift` | scalar | Spectroscopic redshift |
| `parallax` | scalar | Gaia parallax (mas) |
| `flux_g`, `flux_r`, `flux_i`, `flux_z` | scalar | Legacy Survey photometric fluxes (nanomaggies) |
| `gaia_flux_g`, `gaia_flux_bp`, `gaia_flux_rp` | scalar | Gaia G / BP / RP fluxes |
| `shape_r`, `shape_e1`, `shape_e2` | scalar | Legacy Survey morphology parameters |
| `spectrum` | array `(8704,)` | DESI spectrum (flux) |
| `xp_bp`, `xp_rp` | array `(55,)` | Gaia XP coefficient arrays |
| `image` | array `(3, 128, 128)` | Legacy Survey 3-band image (g, r, z) |

Any of the above can be used as inputs or outputs in `generate()`.

---

## Return Type

`generate()` and `star()` return a `GeneratedSource` object.

| Field | Description |
|-------|-------------|
| `.outputs` | Dictionary of generated modalities: key → numpy array |
| `.<key>` | Attribute-style access, e.g. `.spectrum`, `.image`, `.redshift` |
| `.embedding_spectrum` | AION latent embedding of the generated spectrum `(768,)` — `None` if `compute_embeddings=False` |
| `.embedding_image` | AION latent embedding of the generated image `(768,)` — `None` if `compute_embeddings=False` |

---

## Generation Examples

### Redshift → spectrum + image

```python
result = gen.generate(
    inputs={"redshift": 0.3},
    outputs=["spectrum", "image"],
)
```

Runs a two-step chained pipeline: redshift → DESI spectrum (CFG=1.0), then spectrum → Legacy Survey image (CFG=2.0).

### Morphological type conditioning

```python
result = gen.generate(
    inputs={"redshift": 0.3},
    outputs=["spectrum", "image"],
    type="elliptical",  # or "spiral", "irregular"
)
```

Internally injects median `shape_r`, `shape_e1`, `shape_e2` values from DESI DR1 as additional conditioning inputs.

### Real sources from catalogue → spectrum (with embeddings)

```python
from mockcraft.catalogue import fetch_sources

sources = fetch_sources(
    surveys=["desi"],
    region="cosmos",
    columns=["redshift", "flux_g", "flux_r", "flux_z"],
    max_sources=10,
)

for _, row in sources.iterrows():
    result = gen.generate(
        inputs={"redshift": float(row["redshift"]), "flux_g": float(row["flux_g"]),
                "flux_r": float(row["flux_r"]), "flux_z": float(row["flux_z"])},
        outputs=["spectrum"],
        compute_embeddings=True,
    )
    print(result.spectrum.shape)            # (8704,)
    print(result.embedding_spectrum.shape)  # (768,)
```

---

## Catalogue Utility

```python
from mockcraft.catalogue import fetch_sources

sources = fetch_sources(
    surveys=["gaia", "desi"],
    region="cosmos",
    columns=["ra", "dec", "magnitude", "redshift"],
    max_sources=100,
)
```

`region` accepts either a named field or an explicit `(RA, Dec, radius_deg)` tuple:

```python
# Named region
fetch_sources(surveys=["desi"], region="cosmos", ...)

# Explicit coordinates
fetch_sources(surveys=["gaia"], region=(150.1, 2.18, 0.18), ...)
```

### Supported surveys and columns

| Survey | Supported columns |
|--------|-------------------|
| `"gaia"` | `ra`, `dec`, `magnitude`, `gaia_flux_g`, `gaia_flux_bp`, `gaia_flux_rp`, `gaia_parallax` |
| `"desi"` | `ra`, `dec`, `redshift`, `flux_g`, `flux_r`, `flux_z`, `targetid`, `otype` |
| `"legacy"` | `ra`, `dec`, `redshift`, `type` |

When combining multiple surveys, columns not available in a given survey are filled with `NaN`. DESI results are automatically filtered to `ZWARN == 0` (good redshift quality only).

---

## Visualization

Requires the `viz` extra (`pip install mockcraft[viz]`).

```python
from mockcraft import plot_xp_spectrum, plot_image

# Plot a predicted Gaia XP spectrum
result = gen.star(temperature=5778.0)
plot_xp_spectrum(result)

# Plot a predicted Legacy Survey image
result = gen.generate(inputs={"redshift": 0.3}, outputs=["image"])
plot_image(result)
```

---

## Model and Hyperparameters

The package loads `polymathic-ai/aion-base` automatically on first use.

| Parameter | Value | Role |
|-----------|-------|------|
| `CFG_SPEC` | `1.0` | Guidance scale for anything → spectrum |
| `CFG_GALAXY` | `2.0` | Guidance scale for spectrum → image |
| `MASKGIT_STEPS` | `8` | Number of iterative decoding steps |
| `TEMPERATURE` | `0.8` | Sampling temperature |
| `N_ROAR_DRAWS` | `50` | Posterior samples for redshift estimation (higher = more accurate, slower) |

---

## Embeddings and Validation

Setting `compute_embeddings=True` re-encodes generated spectra and images through AION's encoder to produce latent vectors of shape `(768,)`. These can be used for:

- comparing generated vs real source distributions in embedding space (cosine similarity, MMD)
- UMAP projection to inspect coverage of the latent space
- detecting mode collapse or out-of-distribution generation

Embeddings are disabled by default because they add a second forward pass per modality.

---

## Repository Structure

```
mockcraft/
├── mockcraft/
│   ├── __init__.py           # Public API exports
│   ├── source_generator.py   # SourceGenerator, GeneratedSource, plot helpers
│   └── catalogue.py          # fetch_sources — query Gaia, DESI, Legacy via VizieR
├── pyproject.toml
├── README.md
└── LICENSE
```

---

## License

MIT — see [LICENSE](LICENSE).

---

## Citation

If you use MockCraft in your research, please cite the AION paper:

```bibtex
@article{multimodal_universe_2024,
  title   = {Multimodal Universe: Enabling Large-Scale Machine Learning with 100TB of Astronomical Scientific Data},
  year    = {2024},
  url     = {https://arxiv.org/pdf/2412.02527}
}
```
