rbonnetguerrini/ml4transients

1. What this repo is

ML4transients is an end-to-end Python package for transient-vs-bogus classification, centered on LSST data products. It is not a single model repo; it spans data generation, cutout/lightcurve preparation, model training, inference, and post-hoc evaluation/interpretability.

2. Core pipeline boundary

The architecture is organized as a pipeline over astronomy data artifacts:

raw LSST difference-image/visit data → injected catalogs and cutouts → indexed lightcurve datasets and cross-match annotations → PyTorch training datasets → trained CNN/ensemble models → batch inference and evaluation outputs.

Most steps are driven by YAML configs and shell submit scripts.

3. Package layout in src/ML4transients

The library is split by subsystem:

- data_preparation: crossmatch, cutouts, lightcurves

- data_access: dataset loading and lazy access

- training: models, losses, trainers, PyTorch dataset wrappers

- evaluation: inference, metrics, visualizations, interpretability, lightcurve labeling

- injection: synthetic transient catalog generation

- utils.py: shared helpers

This mirrors the operational workflow and makes the main architecture easy to diagram as separate stages.

4. Data preparation subsystem

This is the heaviest external-boundary layer and depends on LSST-distributed tooling. It includes:

- run_cutout.py and batch submit scripts to produce image cutouts/features from selected visits

- run_lightcurves.py to extract and index time-series data into patch-based HDF5 files

- run_crossmatch.py to annotate/filter against external catalogs such as Gaia/MilliQuas

- noise perturbation and batch orchestration utilities for splitting large collections across jobs

The output is structured for downstream ML consumption, not for direct human use.

5. Injection subsystem

scripts/injection and src/ML4transients/injection generate synthetic source catalogs for training. The repo supports both host-associated and hostless transients, with band-specific catalogue generation and batch submission. This stage feeds the LSST pipeline and is a distinct source of supervised labels in an otherwise weakly/unsupervised data regime.

6. Dataset access boundary

DatasetLoader in data_access is the main abstraction for consuming prepared data. It lazily loads cutouts, lightcurves, features, and inference outputs from disk, and builds object/source ID indexing for lookup. This is the boundary between storage formats (HDF5, indices, per-visit artifacts) and model-facing code.

7. Training subsystem

Training is PyTorch-based (torch 2.0 per README) and includes multiple strategies:

- standard CNN training

- deep ensemble training

- co-teaching and asymmetric co-teaching

- stochastic co-teaching variants

- multichannel variants that use difference images plus coadded images

The models.py, losses.py, trainers.py, and pytorch_dataset.py modules separate network definition, objective functions, training loop logic, and data feeding. Hyperparameter search is supported through Bayesian optimization configs.

8. Inference and evaluation subsystem

evaluation/inference.py and scripts/evaluation/run_inference.py run model prediction over prepared data, often in batch mode. Evaluation extends beyond classification metrics into:

- UMAP/embedding-based interpretability

- uncertainty analysis

- visualizations for cutouts and lightcurves

- object-subset evaluation via ID lists

Results are written to disk and can be re-ingested for further analysis.

9. Lightcurve-specific workflow and SuperNNova integration

A separate lightcurve inference path exists around SuperNNova (SNN). The repo converts HDF5 patch-based lightcurves to CSV, runs SNN in a different environment, then writes ensemble predictions back into the original HDF5 files. This is a clear cross-system boundary: ML4transients manages the data staging and result reintegration, while SNN performs the actual sequence-model inference.

10. Operational model: scripts, configs, and batch jobs

The repo is designed for HPC/cluster execution rather than a single interactive app. scripts/ contains job submission wrappers for SLURM-style batch execution, and configs/ defines stage-specific YAMLs for data preparation, injection, training, and evaluation. The README explicitly distinguishes LSST-distrib environments for pipeline-related steps from a normal conda/PyTorch environment for training and evaluation.

11. Artifacts and persistence

The main persistent outputs are HDF5 datasets, CSV exports for SNN, config-derived intermediate files, model checkpoints/weights, TensorBoard logs, and HTML/lightcurve visualizations. This means the architecture is file- and job-oriented, with disk-based handoffs between subsystems rather than a central service or database.

12. How to diagram it

A good architecture diagram would show four main lanes:

data generation/preparation, ML training, inference/evaluation, and external catalog/SNN integration.

The key interfaces are the on-disk dataset formats, the DatasetLoader abstraction, and the config-driven batch scripts that connect each stage.