mockcraft 0.1.1

Reusable Python package for multimodal astronomical source generation using AION

mockcraft

mockcraft is a Python package for generating synthetic astronomical catalogues using the 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 is a multimodal foundation model trained on one of the largest astronomical datasets ever assembled (see the Multimodal Universe paper). 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.

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

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:

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

From source:

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

---

Quick Start

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

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.

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.

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

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

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)

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

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:

# 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]).

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.

---

Citation

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

@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}
}