# guyardon/umap-clustering architecture

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

- Interactive diagram: https://gitdiagram.com/guyardon/umap-clustering
- Repository: https://github.com/guyardon/umap-clustering

## Overview

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: [`src/umap_clustering/callback/callback.py`](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/callback.py) (Callback), [`src/umap_clustering/sweep/orchestrator.py`](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/orchestrator.py) (Orchestrator), [`src/umap_clustering/__init__.py`](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/__init__.py) (umap_clustering), [`src/umap_clustering/fitting.py`](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/fitting.py) (Fitting)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_kernel["Shared kernel"]
  node_pkg["umap_clustering<br/>python package<br/>[__init__.py]"]
  node_kernel_config["Config<br/>dataclasses<br/>[config.py]"]
  node_kernel_fitting["Fitting<br/>umap+hdbscan core<br/>[fitting.py]"]
  node_plot_scatter["Scatter<br/>plotly figure<br/>[scatter.py]"]
  node_plot_heatmap["Heatmap<br/>plotly figure<br/>[heatmap.py]"]
  node_plot_sink["MLflow sink<br/>artifact writer<br/>[mlflow_sink.py]"]
end

subgraph group_callback["Lightning callback"]
  node_cb_callback["Callback<br/>[callback.py]"]
  node_cb_buffers["Buffers<br/>stage storage<br/>[buffers.py]"]
  node_cb_ddp["DDP<br/>distributed gather<br/>[ddp.py]"]
  node_cb_cache["Cache<br/>fit reuse state<br/>[cache.py]"]
  node_cb_strategies["Strategies<br/>fit mode selector<br/>[strategies.py]"]
  node_cb_sink["Lightning sink<br/>mlflow adapter"]
end

subgraph group_sweep["Optuna sweep"]
  node_sw_orchestrator["Orchestrator<br/>sweep entrypoint<br/>[orchestrator.py]"]
  node_sw_objective["Objective<br/>[objective.py]"]
  node_sw_embeddings["Embeddings<br/>frozen cache<br/>[embeddings.py]"]
  node_sw_parallel["Parallel<br/>multiprocessing<br/>[parallel.py]"]
  node_sw_storage["Storage<br/>optuna backend<br/>[storage.py]"]
  node_sw_mlflow["MLflow logging<br/>run hierarchy<br/>[mlflow_logging.py]"]
  node_sw_plot_best["Best trial<br/>summary plot<br/>[best_trial.py]"]
  node_sw_plot_parcoords["Parcoords<br/>summary plot<br/>[parcoords.py]"]
end

subgraph group_scripts["Demos"]
  node_demo_mlflow["Demo MLflow<br/>example script<br/>[demo_mlflow.py]"]
  node_demo_sweep["Demo sweep<br/>example script<br/>[demo_sweep.py]"]
end

node_pkg -->|"exposes"| node_kernel_config
node_pkg -->|"exposes"| node_kernel_fitting
node_pkg -->|"exposes"| node_plot_scatter
node_pkg -->|"exposes"| node_plot_heatmap
node_pkg -->|"exposes"| node_plot_sink
node_kernel_fitting -->|"figures"| node_plot_scatter
node_kernel_fitting -->|"figures"| node_plot_heatmap
node_plot_scatter -->|"logs"| node_plot_sink
node_plot_heatmap -->|"logs"| node_plot_sink
node_cb_callback -->|"stores"| node_cb_buffers
node_cb_callback -->|"gathers"| node_cb_ddp
node_cb_callback -->|"reuses"| node_cb_cache
node_cb_callback -->|"selects"| node_cb_strategies
node_cb_callback -->|"fits"| node_kernel_fitting
node_cb_callback -->|"logs"| node_cb_sink
node_cb_sink -->|"adapts"| node_plot_sink
node_sw_orchestrator -->|"runs"| node_sw_objective
node_sw_orchestrator -->|"distributes"| node_sw_parallel
node_sw_objective -->|"loads"| node_sw_embeddings
node_sw_objective -->|"fits"| node_kernel_fitting
node_sw_objective -->|"logs"| node_sw_mlflow
node_sw_parallel -->|"shares"| node_sw_storage
node_sw_orchestrator -->|"configures"| node_sw_storage
node_sw_orchestrator -->|"hierarchy"| node_sw_mlflow
node_sw_orchestrator -->|"summarizes"| node_sw_plot_best
node_sw_orchestrator -->|"summarizes"| node_sw_plot_parcoords
node_demo_mlflow -->|"exercises"| node_cb_callback
node_demo_sweep -->|"exercises"| node_sw_orchestrator

click node_pkg "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/__init__.py"
click node_kernel_config "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/config.py"
click node_kernel_fitting "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/fitting.py"
click node_plot_scatter "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/plotting/scatter.py"
click node_plot_heatmap "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/plotting/heatmap.py"
click node_plot_sink "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/plotting/mlflow_sink.py"
click node_cb_callback "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/callback.py"
click node_cb_buffers "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/buffers.py"
click node_cb_ddp "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/ddp.py"
click node_cb_cache "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/cache.py"
click node_cb_strategies "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/strategies.py"
click node_cb_sink "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/lightning_mlflow_sink.py"
click node_sw_orchestrator "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/orchestrator.py"
click node_sw_objective "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/objective.py"
click node_sw_embeddings "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/embeddings.py"
click node_sw_parallel "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/parallel.py"
click node_sw_storage "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/storage.py"
click node_sw_mlflow "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/mlflow_logging.py"
click node_sw_plot_best "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/plotting/best_trial.py"
click node_sw_plot_parcoords "https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/plotting/parcoords.py"
click node_demo_mlflow "https://github.com/guyardon/umap-clustering/blob/main/scripts/demo_mlflow.py"
click node_demo_sweep "https://github.com/guyardon/umap-clustering/blob/main/scripts/demo_sweep.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_pkg,node_kernel_config,node_kernel_fitting,node_plot_scatter,node_plot_heatmap,node_plot_sink toneBlue
class node_cb_callback,node_cb_buffers,node_cb_ddp,node_cb_cache,node_cb_strategies,node_cb_sink toneAmber
class node_sw_orchestrator,node_sw_objective,node_sw_embeddings,node_sw_parallel,node_sw_storage,node_sw_mlflow,node_sw_plot_best,node_sw_plot_parcoords toneMint
class node_demo_mlflow,node_demo_sweep toneRose
```

## Components

### Shared kernel

Reusable UMAP/HDBSCAN computation and plotting core.

- [umap_clustering](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/__init__.py): python package · Top-level package that exposes the shared kernel and the two optional consumers. · `src/umap_clustering/__init__.py`
- [Config](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/config.py): dataclasses · Lightweight frozen configuration objects for the shared kernel and callback-facing APIs. · `src/umap_clustering/config.py`
- [Fitting](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/fitting.py): umap+hdbscan core · Fits UMAP, optionally clusters with HDBSCAN, and projects new points through an existing fit. · `src/umap_clustering/fitting.py`
- [Scatter](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/plotting/scatter.py): plotly figure · Builds UMAP scatter plots from fitted embeddings and labels. · `src/umap_clustering/plotting/scatter.py`
- [Heatmap](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/plotting/heatmap.py): plotly figure · Builds class-vs-cluster heatmaps for evaluating clustering quality. · `src/umap_clustering/plotting/heatmap.py`
- [MLflow sink](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/plotting/mlflow_sink.py): artifact writer · 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](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/callback.py): lightning callback · Coordinates training-time embedding collection, fitting, and artifact emission. · `src/umap_clustering/callback/callback.py`
- [Buffers](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/buffers.py): stage storage · Manages per-stage embedding storage and lifecycle inside the callback. · `src/umap_clustering/callback/buffers.py`
- [DDP](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/ddp.py): distributed gather · Aggregates embeddings and state across distributed Lightning ranks. · `src/umap_clustering/callback/ddp.py`
- [Cache](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/cache.py): fit reuse state · Holds reusable train-fit state for validation, test, and predict reuse paths. · `src/umap_clustering/callback/cache.py`
- [Strategies](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/strategies.py): fit mode selector · Chooses whether the callback refits, reuses, or projects with the learned manifold. · `src/umap_clustering/callback/strategies.py`
- [Lightning sink](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/callback/lightning_mlflow_sink.py): mlflow adapter · 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](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/orchestrator.py): sweep entrypoint · Entry point that assembles and runs the Optuna sweep workflow. · `src/umap_clustering/sweep/orchestrator.py`
- [Objective](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/objective.py): optuna objective · Defines trial evaluation, metrics, and the train/validation subsampling logic. · `src/umap_clustering/sweep/objective.py`
- [Embeddings](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/embeddings.py): frozen cache · Materializes and reuses frozen-encoder embeddings keyed by data and extractor identity. · `src/umap_clustering/sweep/embeddings.py`
- [Parallel](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/parallel.py): multiprocessing · Runs trials across workers while sharing Optuna state through common storage. · `src/umap_clustering/sweep/parallel.py`
- [Storage](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/storage.py): optuna backend · Selects and configures the Optuna storage backend used by workers. · `src/umap_clustering/sweep/storage.py`
- [MLflow logging](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/mlflow_logging.py): run hierarchy · Coordinates parent and child run logging for sweep execution and artifacts. · `src/umap_clustering/sweep/mlflow_logging.py`
- [Best trial](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/plotting/best_trial.py): summary plot · Renders best-trial summary visuals for sweep results. · `src/umap_clustering/sweep/plotting/best_trial.py`
- [Parcoords](https://github.com/guyardon/umap-clustering/blob/main/src/umap_clustering/sweep/plotting/parcoords.py): summary plot · 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](https://github.com/guyardon/umap-clustering/blob/main/scripts/demo_mlflow.py): example script · Synthetic-data demo for the callback subsystem and MLflow logging. · `scripts/demo_mlflow.py`
- [Demo sweep](https://github.com/guyardon/umap-clustering/blob/main/scripts/demo_sweep.py): example script · End-to-end demo for the Optuna sweep workflow and best-trial artifacts. · `scripts/demo_sweep.py`

## Connections

- 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.