desisky 0.4.0

No project description provided

desisky

---

About

desisky provides machine learning models and tools for DESI sky modeling:

1. Predictive broadband model - Predicts surface brightness in V, g, r, and z photometric bands from observational metadata (moon position, transparency, eclipse fraction)
2. Variational Autoencoder (VAE) - Compresses sky spectra (7,781 wavelength points → 8-dimensional latent space) for analysis, anomaly detection, and dimensionality reduction
3. Latent Diffusion Models (LDM) - Generates realistic night-sky emission spectra conditioned on observational parameters:
   • LDM Dark - Dark-time spectra conditioned on sun position, transparency, galactic/ecliptic coordinates, and solar flux
   • LDM Moon - Moon-contaminated spectra conditioned on moon position, separation, and illumination fraction
   • LDM Twilight - Twilight spectra conditioned on observation altitude, transparency, sun altitude, and sun separation
4. Data utilities - Download and load DESI DR1 Sky Spectra Value-Added Catalog (VAC) with automatic integrity verification and subset filtering

Built with JAX/Equinox for high-performance model inference and designed to integrate with SpecSim and survey forecasting workflows. This repository hosts the code and notebooks supporting the forthcoming paper by Dowicz et al. (20XX).

Table of Contents

• Installation
• Quick Start
  • Broadband Model
  • VAE
  • Latent Diffusion Model
  • Data Loading
• Data Subsets
• Loading Pre-trained Models
• Data Download
• Examples
• Development
• License

Installation

Basic installation (model inference only)

pip install desisky[cpu]

With data utilities (includes FITS file reading)

pip install desisky[cpu,data]

For GPU support

pip install desisky[cuda12,data]

Note: CUDA wheels require manual installation. See JAX installation guide for details.

Quick Start

Load pre-trained broadband model and run inference

import desisky
import jax.numpy as jnp

# Load the pre-trained broadband model
model, meta = desisky.io.load_model("broadband")

# Example input: [placeholder for actual feature names]
x = jnp.array([...])  # Shape: (6,)

# Predict surface brightness in V, g, r, z bands
y = model(x)  # Shape: (4,)
print(f"Predicted magnitudes: {y}")

Load pre-trained VAE and encode sky spectra

from desisky.io import load_builtin
from desisky.data import SkySpecVAC
import jax.random as jr

# Load DESI sky spectra
vac = SkySpecVAC(version="v1.0", download=True)
wavelength, flux, metadata = vac.load()

# Load pre-trained VAE
vae, meta = load_builtin("vae")

# Encode a sky spectrum to latent representation
spectrum = flux[0].squeeze()
mean, logvar = vae.encode(spectrum)
print(f"Latent mean: {mean}")  # Shape: (8,)

# Sample and decode
latent = vae.sample(mean, logvar, jr.PRNGKey(0))
reconstructed = vae.decode(latent)
print(f"Reconstructed shape: {reconstructed.shape}")  # Shape: (7781,)

# Batch processing with vmap
import jax
batch_means, batch_logvars = jax.vmap(vae.encode)(flux.squeeze())
print(f"Batch latents shape: {batch_means.shape}")  # Shape: (9176, 8)

Generate sky spectra with Latent Diffusion Model

from desisky.io import load_builtin
from desisky.inference import LatentDiffusionSampler
import jax.random as jr
import jax.numpy as jnp

# Load pre-trained VAE and LDM
vae, _ = load_builtin("vae")
ldm, _ = load_builtin("ldm_dark")

# Create sampler (Heun method recommended for quality)
sampler = LatentDiffusionSampler(
    ldm_model=ldm,
    vae_model=vae,
    method="heun",
    num_steps=1000
)

# Define conditioning: [OBSALT, TRANSP, SUNALT, SOLFLUX, ECLLON, ECLLAT, GALLON, GALLAT]
conditioning = jnp.array([
    [2100.0, 0.9, -30.0, 150.0, 45.0, 10.0, 120.0, 5.0],  # Dark sky conditions
])

# Generate spectrum
generated_spectra = sampler.sample(
    key=jr.PRNGKey(42),
    conditioning=conditioning,
    guidance_scale=2.0
)

print(f"Generated spectrum shape: {generated_spectra.shape}")  # (1, 7781)

Download and load DESI sky spectra data

from desisky.data import SkySpecVAC

# Download DR1 VAC (274 MB, with SHA-256 verification)
vac = SkySpecVAC(version="v1.0", download=True)

# Load wavelength, flux, and metadata
wavelength, flux, metadata = vac.load()

print(f"Wavelength shape: {wavelength.shape}")  # (7781,)
print(f"Flux shape: {flux.shape}")              # (9176, 7781)
print(f"Metadata columns: {list(metadata.columns)}")
# ['NIGHT', 'EXPID', 'TILEID', 'AIRMASS', 'EBV', 'MOONFRAC', 'MOONALT', ...]

# Load with enrichment (adds V-band magnitudes and eclipse fraction)
wavelength, flux, metadata = vac.load(enrich=True)
print('SKY_MAG_V_SPEC' in metadata.columns)  # True
print('ECLIPSE_FRAC' in metadata.columns)    # True

Data Subsets

The VAC provides three subset methods for filtering observations by sky conditions:

Dark Time (Non-contaminated)

# Load observations with minimal sun/moon contamination
wave, flux, meta = vac.load_dark_time()

# Filtering criteria:
# - SUNALT < -20° (Sun well below horizon)
# - MOONALT < -5° (Moon below horizon)
# - TRANSPARENCY_GFA > 0 (valid measurements)

Sun Contaminated (Twilight)

# Load twilight observations
wave, flux, meta = vac.load_sun_contaminated()

# Filtering criteria:
# - SUNALT > -20° (Sun near or above horizon)
# - MOONALT <= -5° (Moon below horizon)
# - MOONSEP <= 110° (Sun-Moon separation)
# - TRANSPARENCY_GFA > 0

Moon Contaminated

# Load moon-bright observations
wave, flux, meta = vac.load_moon_contaminated()

# Filtering criteria:
# - SUNALT < -20° (nighttime)
# - MOONALT > 5° (Moon above horizon)
# - MOONFRAC > 0.5 (Moon >50% illuminated)
# - MOONSEP <= 90° (Moon within 90°)
# - TRANSPARENCY_GFA > 0

All subset methods include enrichment by default (enrich=True), adding computed columns for V-band magnitude and lunar eclipse fraction.

Loading Pre-trained Models

The desisky.io.load_model() function provides a unified interface for loading models:

import desisky

# Load packaged pre-trained weights
model, meta = desisky.io.load_model("broadband")

# Load from a custom checkpoint
model, meta = desisky.io.load_model("broadband", path="path/to/checkpoint.eqx")

# Save your own trained model
desisky.io.save(
    "my_model.eqx",
    model,
    meta={
        "schema": 1,
        "arch": {"in_size": 6, "out_size": 4, "width_size": 128, "depth": 5},
        "training": {"date": "2025-01-15", "commit": "abc123"},
    }
)

Available models:

• "broadband" - Multi-layer perceptron (6 inputs → 4 outputs) for V, g, r, z magnitude prediction from moon/transparency conditions
• "vae" - Variational autoencoder (7781 → 8 → 7781) for sky spectra compression, reconstruction, and latent space analysis
• "ldm_dark" - Latent diffusion model (1D U-Net) for generating dark-time sky spectra conditioned on 8 observational parameters:
  • Conditioning: [OBSALT, TRANSPARENCY_GFA, SUNALT, SOLFLUX, ECLLON, ECLLAT, GALLON, GALLAT]
• "ldm_moon" - Latent diffusion model (1D U-Net) for generating moon-contaminated sky spectra conditioned on 6 observational parameters:
  • Conditioning: [OBSALT, TRANSPARENCY_GFA, SUNALT, MOONALT, MOONSEP, MOONFRAC]
• "ldm_twilight" - Latent diffusion model (1D U-Net) for generating twilight sky spectra conditioned on 4 observational parameters:
  • Conditioning: [OBSALT, TRANSPARENCY_GFA, SUNALT, SUNSEP]

Data Download

Python API

from desisky.data import SkySpecVAC

# Download to default location (~/.desisky/data)
vac = SkySpecVAC(download=True)

# Download to custom location
vac = SkySpecVAC(root="/path/to/data", download=True)

# Skip SHA-256 verification (not recommended)
vac = SkySpecVAC(download=True, verify=False)

# Get path to downloaded file
print(vac.filepath())

Command-line interface

# Show default data directory
desisky-data dir

# Download DESI DR1 sky spectra VAC
desisky-data fetch --version v1.0

# Download to custom location
desisky-data fetch --root /path/to/data

# Skip checksum verification
desisky-data fetch --no-verify

Environment variable

Override the default data directory:

export DESISKY_DATA_DIR=/path/to/data
desisky-data dir  # Shows /path/to/data

Examples

See examples/ directory for Jupyter notebooks demonstrating:

• 00_quickstart.ipynb - Quick introduction: loading models, data subsets, and running inference
• 01_broadband_training.ipynb - Train the broadband model on moon-contaminated subset
• 02_vae_inference.ipynb - VAE inference: encoding/decoding sky spectra and latent space visualization
• 03_vae_analysis.ipynb - Advanced VAE analysis: latent space interpolation and anomaly detection
• 04_vae_training.ipynb - Train a VAE from scratch with InfoVAE-MMD objective
• 05_ldm_inference.ipynb - Generate dark-time and moon-contaminated sky spectra using the latent diffusion models with custom conditioning

Development

Setting up development environment

git clone https://github.com/MatthewDowicz/desisky.git
cd desisky
pip install -e ".[cpu,data]"
pip install pytest pytest-cov

Running tests

# Run all tests
pytest

# Run with coverage
pytest --cov=desisky --cov-report=html

# Run specific test file
pytest tests/test_model_io.py -v

Project Structure

desisky/
├── src/desisky/
│   ├── io/              # Model I/O (save/load checkpoints with metadata)
│   ├── models/          # Model architectures
│   │   ├── broadband.py # Broadband MLP for magnitude prediction
│   │   ├── vae.py       # Variational autoencoder (InfoVAE-MMD)
│   │   └── ldm.py       # Latent diffusion model (1D U-Net)
│   ├── data/            # Data downloading, loading, and enrichment
│   │   ├── skyspec.py   # SkySpecVAC class with subset filtering
│   │   ├── _enrich.py   # V-band, eclipse, solar flux, coordinates
│   │   └── _core.py     # Download utilities with SHA-256 verification
│   ├── training/        # Training infrastructure
│   │   ├── dataset.py   # PyTorch Dataset wrappers
│   │   ├── vae_trainer.py    # VAE training loop
│   │   ├── losses.py         # Loss functions
│   │   └── vae_losses.py     # InfoVAE-MMD loss
│   ├── inference/       # Sampling algorithms
│   │   └── sampling.py  # DDPM, DDIM, Heun samplers for LDM
│   ├── visualization/   # Plotting utilities
│   ├── scripts/         # CLI tools (desisky-data)
│   └── weights/         # Pre-trained model weights (small models)
├── tests/               # Comprehensive test suite (123+ tests)
│   ├── test_vae.py           # VAE unit tests
│   ├── test_model_io.py      # Model I/O tests
│   ├── test_enrichment.py    # Data enrichment tests
│   ├── test_ldm_sampling.py  # LDM sampling tests
│   └── ...                   # Other test modules
├── examples/            # Jupyter notebook tutorials
│   ├── 00_quickstart.ipynb
│   ├── 01_broadband_training.ipynb
│   ├── 02_vae_inference.ipynb
│   ├── 03_vae_analysis.ipynb
│   ├── 04_vae_training.ipynb
│   └── 05_ldm_inference.ipynb
└── pyproject.toml       # Package configuration

Key Features

• JAX/Equinox models: High-performance, functional ML models with automatic differentiation
• Production-ready I/O: Checkpoint format with JSON metadata + binary weights
• Automatic caching: Downloaded data and models cached locally for fast re-use
• Integrity verification: SHA-256 checksums for all downloaded files
• Subset filtering: Easy access to dark-time, twilight, and moon-contaminated observations
• Data enrichment: Automatic computation of V-band magnitudes, eclipse fractions, solar flux, and coordinate transformations
• Multiple sampling methods: DDPM, DDIM, and Heun (probability-flow ODE) for LDM inference
• Comprehensive tests: 123+ unit tests ensuring reliability

License

desisky is distributed under the terms of the MIT license.