# maniacurgency42/ml4em architecture

> Architecture diagram of the GitHub repository maniacurgency42/ml4em, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-16). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/maniacurgency42/ml4em
- Repository: https://github.com/maniacurgency42/ml4em

## Overview

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.

Main entry points: [`src/ml4em/features/pipeline.py`](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/pipeline.py) (Pipeline), [`src/ml4em/types.py`](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/types.py) (Types), [`src/ml4em/config`](https://github.com/maniacurgency42/ml4em/tree/main/src/ml4em/config) (Config), [`src/ml4em/features/base.py`](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/base.py) (Extractor API)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_foundation["Foundation"]
  node_types["Types<br/>contracts<br/>[types.py]"]
  node_constants["Constants<br/>[constants.py]"]
  node_config["Config"]
end

subgraph group_data["Data"]
  node_data_base["Source API<br/>ingestion protocol<br/>[base.py]"]
  node_data_ztf["ZTF<br/>source adapter<br/>[ztf.py]"]
  node_data_rubin["Rubin<br/>stub source<br/>[rubin.py]"]
  node_data_sim["Simulation<br/>stub source<br/>[simulation.py]"]
end

subgraph group_features["Features"]
  node_feat_base["Extractor API<br/>feature protocol<br/>[base.py]"]
  node_feat_pipe["Pipeline<br/>feature orchestrator<br/>[pipeline.py]"]
  node_feat_stats["Statistics<br/>scalar features<br/>[statistics.py]"]
  node_feat_period["Period<br/>[period.py]"]
  node_feat_dmdt["dm/dt<br/>image features<br/>[dmdt.py]"]
  node_feat_catalog["Catalog<br/>[catalog.py]"]
end

subgraph group_models["Models"]
  node_model_base["Model API<br/>model protocol<br/>[base.py]"]
  node_xgboost["XGBoost<br/>classifier<br/>[xgboost.py]"]
end

subgraph group_training["Training"]
  node_train_dataset["Dataset<br/>label joiner<br/>[dataset.py]"]
  node_train_trainer["Trainer<br/>training shell<br/>[trainer.py]"]
end

subgraph group_inference["Inference"]
  node_infer_loader["Loader<br/>model dispatcher<br/>[loader.py]"]
  node_infer_predictor["Predictor<br/>batch inference<br/>[predictor.py]"]
  node_infer_post["Postprocess<br/>candidate builder<br/>[postprocess.py]"]
end

subgraph group_external["External"]
  node_periodfind{{"periodfind<br/>rust/cuda backend"}}
end

node_config -->|"params"| node_data_base
node_config -->|"params"| node_feat_pipe
node_config -->|"params"| node_xgboost
node_config -->|"manifest"| node_infer_loader
node_data_base -->|"implements"| node_data_ztf
node_data_base -->|"implements"| node_data_rubin
node_data_base -->|"implements"| node_data_sim
node_data_ztf -->|"LightCurve"| node_types
node_data_rubin -.->|"LightCurve"| node_types
node_data_sim -.->|"LightCurve"| node_types
node_feat_base -->|"implements"| node_feat_stats
node_feat_base -->|"implements"| node_feat_period
node_feat_base -->|"implements"| node_feat_dmdt
node_feat_base -->|"implements"| node_feat_catalog
node_feat_pipe -->|"composes"| node_feat_stats
node_feat_pipe -->|"composes"| node_feat_period
node_feat_pipe -->|"composes"| node_feat_dmdt
node_feat_pipe -->|"composes"| node_feat_catalog
node_feat_stats -->|"backend"| node_periodfind
node_feat_period -->|"backend"| node_periodfind
node_feat_dmdt -->|"backend"| node_periodfind
node_feat_pipe -->|"FeatureVector"| node_types
node_model_base -->|"implements"| node_xgboost
node_xgboost -->|"SCALAR_FIELDS"| node_constants
node_xgboost -->|"FeatureVector"| node_types
node_train_dataset -->|"LabeledSample"| node_types
node_train_dataset -->|"features"| node_feat_pipe
node_train_trainer -->|"uses"| node_train_dataset
node_train_trainer -->|"trains"| node_model_base
node_infer_loader -->|"loads"| node_model_base
node_infer_predictor -->|"predicts"| node_model_base
node_infer_predictor -->|"features"| node_feat_pipe
node_infer_post -->|"Candidate"| node_types
node_infer_post -->|"consumes"| node_infer_predictor

click node_types "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/types.py"
click node_constants "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/constants.py"
click node_config "https://github.com/maniacurgency42/ml4em/tree/main/src/ml4em/config"
click node_data_base "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/base.py"
click node_data_ztf "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/ztf.py"
click node_data_rubin "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/rubin.py"
click node_data_sim "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/simulation.py"
click node_feat_base "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/base.py"
click node_feat_pipe "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/pipeline.py"
click node_feat_stats "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/statistics.py"
click node_feat_period "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/period.py"
click node_feat_dmdt "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/dmdt.py"
click node_feat_catalog "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/catalog.py"
click node_periodfind "https://github.com/maniacurgency42/ml4em/tree/main/external/periodfind"
click node_model_base "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/models/base.py"
click node_xgboost "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/models/xgboost.py"
click node_train_dataset "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/training/dataset.py"
click node_train_trainer "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/training/trainer.py"
click node_infer_loader "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/inference/loader.py"
click node_infer_predictor "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/inference/predictor.py"
click node_infer_post "https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/inference/postprocess.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_types,node_constants,node_config,node_periodfind toneBlue
class node_data_base,node_data_ztf,node_data_rubin,node_data_sim toneAmber
class node_feat_base,node_feat_pipe,node_feat_stats,node_feat_period,node_feat_dmdt,node_feat_catalog toneMint
class node_model_base,node_xgboost toneRose
class node_train_dataset,node_train_trainer toneIndigo
class node_infer_loader,node_infer_predictor,node_infer_post toneTeal
```

## Components

### Foundation

Shared contracts, runtime configuration, and core data types.

- [Types](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/types.py): contracts · Core dataclasses for light curves, features, labeled samples, and candidates. · `src/ml4em/types.py`
- [Constants](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/constants.py): shared constants · System-wide constants and stable field ordering contracts. · `src/ml4em/constants.py`
- [Config](https://github.com/maniacurgency42/ml4em/tree/main/src/ml4em/config): runtime config · Validated configuration models and loaders for YAML and environment-based secrets. · `src/ml4em/config`

### Data

Survey ingestion and normalization into light curves.

- [Source API](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/base.py): ingestion protocol · LightCurveSource interface for fetching single objects or batches. · `src/ml4em/data/base.py`
- [ZTF](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/ztf.py): source adapter · Implemented survey ingestion for ZTF light curves. · `src/ml4em/data/ztf.py`
- [Rubin](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/rubin.py): stub source · Placeholder survey adapter for Rubin ingestion. · `src/ml4em/data/rubin.py`
- [Simulation](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/data/simulation.py): stub source · Placeholder adapter for simulated light-curve generation. · `src/ml4em/data/simulation.py`

### Features

Batch feature extraction and heavy compute backends.

- [Extractor API](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/base.py): feature protocol · FeatureExtractor interface implemented by all batch extractors. · `src/ml4em/features/base.py`
- [Pipeline](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/pipeline.py): feature orchestrator · Composes extractors, tolerates partial failures, and returns feature vectors. · `src/ml4em/features/pipeline.py`
- [Statistics](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/statistics.py): scalar features · Statistical feature extraction for light-curve summaries. · `src/ml4em/features/statistics.py`
- [Period](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/period.py): period features · Period-finding features backed by the external compute backend. · `src/ml4em/features/period.py`
- [dm/dt](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/dmdt.py): image features · Delta-magnitude over delta-time histogram features. · `src/ml4em/features/dmdt.py`
- [Catalog](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/features/catalog.py): catalog features · Gaia and catalog-derived feature fields. · `src/ml4em/features/catalog.py`

### Models

Model interfaces and concrete classifiers.

- [Model API](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/models/base.py): model protocol · MLModel interface with prediction and persistence contracts. · `src/ml4em/models/base.py`
- [XGBoost](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/models/xgboost.py): classifier · Reference scalar-only model implementation. · `src/ml4em/models/xgboost.py`

### Training

Label joining, dataset shaping, and training shell.

- [Dataset](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/training/dataset.py): label joiner · Builds LabeledSample objects and computes splits and class stats. · `src/ml4em/training/dataset.py`
- [Trainer](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/training/trainer.py): training shell · Trainer scaffold for saving models; full training loop is pending. · `src/ml4em/training/trainer.py`

### Inference

Model loading, prediction batching, and candidate postprocessing.

- [Loader](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/inference/loader.py): model dispatcher · Loads saved model manifests and dispatches to concrete model classes. · `src/ml4em/inference/loader.py`
- [Predictor](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/inference/predictor.py): batch inference · Batches predict_proba calls over feature vectors. · `src/ml4em/inference/predictor.py`
- [Postprocess](https://github.com/maniacurgency42/ml4em/blob/main/src/ml4em/inference/postprocess.py): candidate builder · Converts probabilities into immutable Candidate outputs with confidence bands. · `src/ml4em/inference/postprocess.py`

### External

Bundled period-finding backend used by feature extraction.

- [periodfind](https://github.com/maniacurgency42/ml4em/tree/main/external/periodfind): rust/cuda backend · External batched implementation for statistics, period finding, Fourier decomposition, and dm/dt histograms. · `external/periodfind`

## Connections

- Config → Source API: params · Runtime settings configure ingestion behavior.
- Config → Pipeline: params · Extraction pipeline uses config to choose algorithms and thresholds.
- Config → XGBoost: params · Model hyperparameters and paths come from config.
- Config → Loader: manifest · Inference uses config and manifest-derived paths.
- Source API → ZTF: implements · ZTF is a concrete source adapter.
- Source API → Rubin: implements · Rubin is a stub adapter.
- Source API → Simulation: implements · Simulation is a stub adapter.
- ZTF → Types: LightCurve · Ingestion normalizes survey rows into light-curve objects.
- Rubin → Types: LightCurve · Planned ingestion target is the shared light-curve type.
- Simulation → Types: LightCurve · Planned simulation output is the shared light-curve type.
- Extractor API → Statistics: implements · Statistics extractor conforms to the protocol.
- Extractor API → Period: implements · Period extractor conforms to the protocol.
- Extractor API → dm/dt: implements · dm/dt extractor conforms to the protocol.
- Extractor API → Catalog: implements · Catalog extractor conforms to the protocol.
- Pipeline → Statistics: composes · Pipeline runs statistics extraction as one stage.
- Pipeline → Period: composes · Pipeline runs period extraction as one stage.
- Pipeline → dm/dt: composes · Pipeline runs dm/dt extraction as one stage.
- Pipeline → Catalog: composes · Pipeline runs catalog feature extraction as one stage.
- Statistics → periodfind: backend · Statistics stage delegates heavy batched compute to the backend.
- Period → periodfind: backend · Period stage delegates period search to the backend.
- dm/dt → periodfind: backend · dm/dt histograms are built by the backend.
- Pipeline → Types: FeatureVector · Pipeline outputs the shared feature-vector contract.
- Model API → XGBoost: implements · XGBoost is the concrete model implementation.
- XGBoost → Constants: SCALAR_FIELDS · Scalar feature ordering is a stable model contract.
- XGBoost → Types: FeatureVector · Model consumes feature vectors as input.
- Dataset → Types: LabeledSample · Training dataset builds labeled samples from features and labels.
- Dataset → Pipeline: features · Training data starts from extracted features.
- Trainer → Dataset: uses · Trainer consumes prepared datasets.
- Trainer → Model API: trains · Trainer operates on model interface contracts.
- Loader → Model API: loads · Loader resolves concrete model implementations through the shared interface.
- Predictor → Model API: predicts · Predictor batches model probability calls.
- Predictor → Pipeline: features · Inference consumes the same feature vectors produced upstream.
- Postprocess → Types: Candidate · Postprocessing emits immutable candidate objects.
- Postprocess → Predictor: consumes · Postprocessing transforms predicted probabilities into candidates.