maniacurgency42/ml4em

1. What this repo is

ml4em is a modular Python library for turning astronomical photometric time series into ML-ready features, then using those features for training or inference. The architecture is intentionally science-case agnostic: the target class is defined by labels and the chosen model, not by the library.

2. Core architectural rule

The code is split into six layers with strict boundaries: foundation, data, features, models, training, and inference. Dependencies flow downward only through shared contracts; training and inference are parallel consumers of the same feature/model interfaces and do not depend on each other.

3. Foundation layer

types.py, constants.py, and config/ define the shared language of the system. The key dataclasses are LightCurve, FeatureVector, LabeledSample, and Candidate. config/schema.py holds validated Pydantic config models, while config/loader.py loads YAML and environment-based secrets. This layer controls runtime parameters, not model architecture.

4. Data layer: survey-specific ingestion

data/base.py defines a LightCurveSource protocol with fetch() and fetch_batch(). Concrete sources live in data/ztf.py, data/rubin.py, and data/simulation.py. Only ZTF is implemented; Rubin and simulation are stubs. This layer normalizes survey data into lists of LightCurve objects, one per source/band combination.

5. Feature layer: batch extraction pipeline

features/base.py defines a FeatureExtractor protocol, and features/pipeline.py composes multiple extractors into a FeatureVector. The main extractors are statistics, period, dm/dt, and catalog. The pipeline is batch-first, tolerates partial failures, and returns all-NaN feature vectors for sources with insufficient observations instead of aborting the job.

6. Heavy compute backend

Most feature computation delegates to the external periodfind submodule under external/periodfind. That dependency is central to the architecture because it provides Rust/CUDA-backed batched implementations for statistics, period finding, Fourier decomposition, and dm/dt histograms. The feature layer prebuilds algorithm objects and runs on CPU or GPU depending on config and container build.

7. Feature outputs and shape contracts

FeatureVector is the main interchange object between feature extraction, training, and inference. It contains scalar fields, a period result, a 26×26 dm/dt image, and Gaia fields. models/base.py defines the scalar feature ordering (SCALAR_FIELDS) used by scalar-only models, which is a stable contract; changing field order would invalidate saved models.

8. Models layer

models/base.py defines the MLModel protocol with predict_proba() and save(). models/xgboost.py provides the reference implementation, XGBoostClassifier, which uses only scalar features and ignores the dm/dt image. Concrete model loading is dispatched through inference/loader.py, so the rest of the system stays model-agnostic.

9. Training flow

training/dataset.py joins feature files with researcher-provided labels to produce LabeledSample objects, then supports dataset splitting and class statistics. training/trainer.py defines the trainer shell: it can save models, but the full training loop is still pending. Labels are external to the library and are expected to come from upstream catalog cross-matching or similar preparation.

10. Inference flow

inference/loader.py loads a saved model by reading a manifest and dispatching to the right concrete class. inference/predictor.py is the predictor shell that batches predict_proba() calls. inference/postprocess.py converts probabilities into immutable Candidate outputs, attaching source identity, period metadata, and a confidence band derived from config thresholds.

11. End-to-end data path

The primary runtime path is: survey source → LightCurve objects → feature extraction → FeatureVector records → model probability prediction → Candidate outputs. Training follows a similar path but joins feature vectors with labels and learns model parameters instead of producing candidates.

12. Deployment and external infrastructure

The repo is built around Docker because periodfind needs a Rust/CUDA build stack that is hard to manage directly in local or HPC environments. The Dockerfile supports CPU and GPU targets, and MSI deployment uses Apptainer containers pulled from GHCR. This is not just packaging detail: container choice affects whether GPU-accelerated period finding is available.

13. Extension points

New data sources, feature extractors, and models are meant to be added as single files implementing the relevant protocol, with minimal registration. The only central registry-like coupling is model loading in inference/loader.py. That makes the architecture easy to diagram as a set of pluggable adapters around the shared feature and model contracts.

14. Implementation maturity

The repository is partially implemented, not fully production-complete. ZTF ingestion, the feature pipeline, statistics/period/dm/dt extraction, model base utilities, and inference postprocessing/loading are complete. Rubin ingestion, simulation, Gaia catalog lookup, and the actual training/inference loops remain stubs or shells.