mockcraft 0.1.1__tar.gz

mockcraft-0.1.1/LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2026 Bibiana Terres Stumpf, Joao Victor Motta da Silva,
+Pedro Jann Luna, Anne Laure Mealier, Julien Zoubian
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mockcraft-0.1.1/PKG-INFO

@@ -0,0 +1,392 @@
+Metadata-Version: 2.4
+Name: mockcraft
+Version: 0.1.1
+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: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: <3.13,>=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy<2.0,>=1.26
+Requires-Dist: pandas<3.0,>=2.2
+Requires-Dist: scipy<1.16,>=1.11
+Requires-Dist: scikit-learn<1.8,>=1.5
+Requires-Dist: astropy<7.0,>=6.1
+Requires-Dist: astroquery<0.5,>=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<0.6,>=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}
+}
+```