How guyardon/umap-clustering works
1. Repository shape
- This is a single Python package, umap_clustering, split into one shared kernel and two optional consumers: a Lightning callback and an Optuna sweep.
- The architectural boundary is explicit: kernel code should not import Lightning or Optuna; callback/ and sweep/ are sibling subsystems that depend on the kernel, but not on each other.
2. Shared kernel
- The core package lives in src/umap_clustering/ and centers on three modules: config.py, fitting.py, and plotting/.
- This layer owns the UMAP + HDBSCAN domain logic, plus plotting and MLflow artifact rendering helpers. It is the reusable computation core for both higher-level flows.
3. Kernel data flow
- Inputs are embedding arrays plus optional labels.
- fitting.py fits UMAP, optionally clusters with HDBSCAN, and can also transform new points and assign cluster labels from an existing fit.
- plotting/scatter.py and plotting/heatmap.py turn results into Plotly figures, while plotting/mlflow_sink.py writes those figures to MLflow as HTML artifacts.
4. Lightning callback subsystem
- src/umap_clustering/callback/ is the training-time integration for PyTorch Lightning.
- callback.py collects embeddings from step outputs, ddp.py gathers them across distributed ranks, buffers.py manages per-stage storage, cache.py holds reusable train-fit state, and strategies.py selects the fit mode.
- lightning_mlflow_sink.py adapts the kernel’s Plotly output to MLflow logging in a Lightning-friendly way.
5. Callback runtime behavior
- During training, the callback periodically pulls embeddings from the model, moves them off-device, aggregates them across DDP workers, and runs the kernel fit.
- It then emits UMAP scatter plots and class-vs-cluster heatmaps to MLflow.
- A special reuse path lets validation/test/predict project into the train epoch’s learned manifold instead of refitting, which keeps cluster IDs stable across stages.
6. Sweep subsystem
- src/umap_clustering/sweep/ is a separate orchestration layer for offline hyperparameter search.
- orchestrator.py is the entry point, objective.py defines the Optuna objective and metrics, embeddings.py materializes frozen-encoder embeddings, parallel.py handles multi-process execution, and storage.py chooses Optuna backend storage.
- mlflow_logging.py coordinates parent/child run logging, and plotting/ generates sweep-specific summary figures.
7. Sweep data flow
- The sweep path starts with a frozen encoder checkpoint and a Hydra-instantiated datamodule.
- Embeddings are computed once, cached on disk, and then reused across many Optuna trials.
- Each trial fits UMAP + HDBSCAN on a fixed training subsample and evaluates on a fixed validation subsample, producing metrics such as ARI, NMI, V-measure, silhouette, cluster count, and noise ratio.
8. Parallelism and shared state
- Multi-worker sweeps use multiprocessing with a shared Optuna storage backend, so trial history is coordinated across workers.
- The embedding cache is SHA1-keyed by the encoder/data/extractor identity, which makes repeated sweeps deterministic and avoids recomputing the same frozen embeddings.
- For large worker counts, memory-mapped cached arrays reduce duplicate RAM usage.
9. Configuration boundary
- Configuration is strongly typed and Hydra-driven in the sweep subsystem under src/umap_clustering/sweep/configs/.
- The root sweep config composes dataclasses for data, encoder, Optuna, MLflow, and search space, while the callback side uses simpler frozen config objects in config.py.
- This means the repo has two config styles: lightweight dataclass configs for the kernel/callback, and Hydra composition for the sweep.
10. External technologies that shape the architecture
- The main runtime dependencies are PyTorch Lightning, UMAP, HDBSCAN, Plotly, MLflow, Hydra, Optuna, and joblib.
- Lightning is only relevant to the callback path, Optuna/Hydra only to the sweep path, and MLflow is the common artifact sink across both consumers.
- uv and extras-based installs enforce the package boundaries at install time.
11. Scripts and demos
- scripts/demo_mlflow.py exercises the callback subsystem on synthetic data.
- scripts/demo_sweep.py exercises the sweep subsystem end-to-end, including MLflow hierarchy and best-trial artifacts.
- The scripts are demonstration and validation entry points rather than core architecture, but they show the two primary execution modes of the package.
12. Tests as architecture coverage
- The tests are organized along the same boundaries as the code: kernel tests, callback tests, and sweep tests.
- This suggests the repo’s main contract is the stability of the shared kernel plus the correctness of the two adapters around it.
Main entry points: callback.py Callback · orchestrator.py Orchestrator · __init__.py umap_clustering · fitting.py Fitting
Components 22
Shared kernel
Reusable UMAP/HDBSCAN computation and plotting core.
- umap_clustering · Top-level package that exposes the shared kernel and the two optional consumers.
src/umap_clustering/__init__.py - Config · Lightweight frozen configuration objects for the shared kernel and callback-facing APIs.
src/umap_clustering/config.py - Fitting · Fits UMAP, optionally clusters with HDBSCAN, and projects new points through an existing fit.
src/umap_clustering/fitting.py - Scatter · Builds UMAP scatter plots from fitted embeddings and labels.
src/umap_clustering/plotting/scatter.py - Heatmap · Builds class-vs-cluster heatmaps for evaluating clustering quality.
src/umap_clustering/plotting/heatmap.py - MLflow sink · Writes Plotly figures to MLflow as HTML artifacts.
src/umap_clustering/plotting/mlflow_sink.py
Lightning callback
Training-time integration built on the shared kernel.
- Callback · Coordinates training-time embedding collection, fitting, and artifact emission.
src/umap_clustering/callback/callback.py - Buffers · Manages per-stage embedding storage and lifecycle inside the callback.
src/umap_clustering/callback/buffers.py - DDP · Aggregates embeddings and state across distributed Lightning ranks.
src/umap_clustering/callback/ddp.py - Cache · Holds reusable train-fit state for validation, test, and predict reuse paths.
src/umap_clustering/callback/cache.py - Strategies · Chooses whether the callback refits, reuses, or projects with the learned manifold.
src/umap_clustering/callback/strategies.py - Lightning sink · Adapts kernel Plotly outputs to MLflow logging in the Lightning runtime.
src/umap_clustering/callback/lightning_mlflow_sink.py
Optuna sweep
Offline hyperparameter search orchestration built on the shared kernel.
- Orchestrator · Entry point that assembles and runs the Optuna sweep workflow.
src/umap_clustering/sweep/orchestrator.py - Objective · Defines trial evaluation, metrics, and the train/validation subsampling logic.
src/umap_clustering/sweep/objective.py - Embeddings · Materializes and reuses frozen-encoder embeddings keyed by data and extractor identity.
src/umap_clustering/sweep/embeddings.py - Parallel · Runs trials across workers while sharing Optuna state through common storage.
src/umap_clustering/sweep/parallel.py - Storage · Selects and configures the Optuna storage backend used by workers.
src/umap_clustering/sweep/storage.py - MLflow logging · Coordinates parent and child run logging for sweep execution and artifacts.
src/umap_clustering/sweep/mlflow_logging.py - Best trial · Renders best-trial summary visuals for sweep results.
src/umap_clustering/sweep/plotting/best_trial.py - Parcoords · Renders parallel-coordinate summaries over sweep search dimensions.
src/umap_clustering/sweep/plotting/parcoords.py
Demos
Executable examples for the two supported runtime modes.
- Demo MLflow · Synthetic-data demo for the callback subsystem and MLflow logging.
scripts/demo_mlflow.py - Demo sweep · End-to-end demo for the Optuna sweep workflow and best-trial artifacts.
scripts/demo_sweep.py
Connections 28
- umap_clustering → Config: exposes · Package surface includes lightweight shared config.
- umap_clustering → Fitting: exposes · Package surface includes the fitting core.
- umap_clustering → Scatter: exposes · Package surface includes scatter plotting helpers.
- umap_clustering → Heatmap: exposes · Package surface includes heatmap plotting helpers.
- umap_clustering → MLflow sink: exposes · Package surface includes MLflow artifact writing helpers.
- Fitting → Scatter: figures · Fitted embeddings and labels feed scatter plot generation.
- Fitting → Heatmap: figures · Cluster assignments feed heatmap generation.
- Scatter → MLflow sink: logs · Scatter figures are written to MLflow as HTML artifacts.
- Heatmap → MLflow sink: logs · Heatmap figures are written to MLflow as HTML artifacts.
- Callback → Buffers: stores · Callback collects embeddings into staged buffers.
- Callback → DDP: gathers · Callback aggregates embeddings across distributed workers.
- Callback → Cache: reuses · Callback reuses learned train-fit state across stages.
- Callback → Strategies: selects · Callback chooses fit or reuse behavior through strategy logic.
- Callback → Fitting: fits · Callback hands collected embeddings to the shared fitting core.
- Callback → Lightning sink: logs · Callback emits MLflow artifacts through the Lightning adapter.
- Lightning sink → MLflow sink: adapts · Lightning-specific logging delegates to the kernel MLflow sink.
- Orchestrator → Objective: runs · Orchestrator executes the Optuna objective across trials.
- Orchestrator → Parallel: distributes · Orchestrator can fan out execution to multiple workers.
- Objective → Embeddings: loads · Objective consumes cached frozen embeddings.
- Objective → Fitting: fits · Each trial fits UMAP and HDBSCAN using the shared core.
- Objective → MLflow logging: logs · Objective reports trial metrics and artifacts through MLflow logging.
- Parallel → Storage: shares · Workers coordinate trial history via shared Optuna storage.
- Orchestrator → Storage: configures · Orchestrator selects the storage backend used by workers.
- Orchestrator → MLflow logging: hierarchy · Orchestrator manages parent/child run structure for sweeps.
- Orchestrator → Best trial: summarizes · Sweep results are rendered into best-trial summary figures.
- Orchestrator → Parcoords: summarizes · Sweep results are rendered into search-space summary figures.
- Demo MLflow → Callback: exercises · Demo validates the Lightning callback path end to end.
- Demo sweep → Orchestrator: exercises · Demo validates the sweep orchestration path end to end.