diff --git a/aitom/bin/disca.py b/aitom/bin/disca.py new file mode 100644 index 0000000..0452311 --- /dev/null +++ b/aitom/bin/disca.py @@ -0,0 +1,139 @@ +""" +DISCA command-line helper for clustering precomputed features. + +This is a light-weight utility that: + - loads feature vectors (N x D) from .npy + - runs GaussianMixture for candidate Ks + - selects the best K by lowest BIC + - saves labels to .npy and a JSON summary + +Example: + disca cluster --features features.npy --candidate-k 5,10,20 \\ + --out-labels disca_labels.npy --out-summary disca_summary.json + +Notes: + - This operates on extracted features (e.g., from DISCA YOPO encoders). + - It does not train the DISCA network; it focuses on the clustering step. +""" + +import argparse +import json +from pathlib import Path +import numpy as np +from sklearn.mixture import GaussianMixture + + +def parse_candidate_ks(value: str): + parts = [p for p in value.split(",") if p.strip()] + ks = [int(p.strip()) for p in parts] + if not ks: + raise argparse.ArgumentTypeError("candidate-k must contain at least one integer") + return ks + + +def load_features(path: Path) -> np.ndarray: + arr = np.load(path) + if arr.ndim != 2: + raise ValueError(f"Expected features with shape (N, D); got shape {arr.shape}") + if arr.shape[0] < 2: + raise ValueError("Need at least 2 samples for clustering.") + return arr.astype(np.float32, copy=False) + + +def run_gmm_bic(features: np.ndarray, candidate_ks, reg_covar: float, max_iter: int, random_state: int): + best = { + "k": None, + "bic": np.inf, + "labels": None, + "model": None, + } + for k in candidate_ks: + gmm = GaussianMixture( + n_components=k, + covariance_type="full", + reg_covar=reg_covar, + max_iter=max_iter, + random_state=random_state, + ) + gmm.fit(features) + bic = gmm.bic(features) + if bic < best["bic"]: + best.update({"k": k, "bic": bic, "labels": gmm.predict(features), "model": gmm}) + return best + + +def save_labels(labels: np.ndarray, path: Path): + path.parent.mkdir(parents=True, exist_ok=True) + np.save(path, labels.astype(np.int64)) + + +def save_summary(path: Path, summary: dict): + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(summary, f, indent=2) + + +def cmd_cluster(args): + feats = load_features(Path(args.features)) + candidate_ks = args.candidate_k + best = run_gmm_bic( + feats, + candidate_ks=candidate_ks, + reg_covar=args.reg_covar, + max_iter=args.max_iter, + random_state=args.seed, + ) + if best["labels"] is None: + raise RuntimeError("GMM did not produce labels.") + + save_labels(best["labels"], Path(args.out_labels)) + + counts = {str(i): int((best["labels"] == i).sum()) for i in np.unique(best["labels"])} + summary = { + "chosen_k": int(best["k"]), + "bic": float(best["bic"]), + "counts": counts, + "candidate_k": candidate_ks, + "reg_covar": args.reg_covar, + "max_iter": args.max_iter, + "seed": args.seed, + } + save_summary(Path(args.out_summary), summary) + + print(f"[DISCA] Chosen K={best['k']} (BIC={best['bic']:.2f})") + print(f"[DISCA] Cluster counts: {counts}") + print(f"[DISCA] Labels saved to: {args.out_labels}") + print(f"[DISCA] Summary saved to: {args.out_summary}") + + +def build_parser(): + p = argparse.ArgumentParser(description="DISCA clustering helper") + sub = p.add_subparsers(dest="command", required=True) + + pc = sub.add_parser("cluster", help="Cluster precomputed features with GMM + BIC") + pc.add_argument("--features", required=True, help="Path to .npy features array of shape (N, D)") + pc.add_argument( + "--candidate-k", + type=parse_candidate_ks, + default=parse_candidate_ks("5,10,20"), + help="Comma-separated candidate K values (default: 5,10,20)", + ) + pc.add_argument("--reg-covar", type=float, default=1e-5, help="GMM reg_covar (default: 1e-5)") + pc.add_argument("--max-iter", type=int, default=200, help="GMM max_iter (default: 200)") + pc.add_argument("--seed", type=int, default=0, help="Random seed for GMM (default: 0)") + pc.add_argument("--out-labels", default="disca_labels.npy", help="Output .npy for labels") + pc.add_argument("--out-summary", default="disca_summary.json", help="Output summary JSON") + pc.set_defaults(func=cmd_cluster) + + return p + + +def main(argv=None): + parser = build_parser() + args = parser.parse_args(argv) + args.func(args) + + +if __name__ == "__main__": + main() + diff --git a/doc/features/DISCA_CLI_Feature.md b/doc/features/DISCA_CLI_Feature.md new file mode 100644 index 0000000..15ddbbd --- /dev/null +++ b/doc/features/DISCA_CLI_Feature.md @@ -0,0 +1,112 @@ +## Feature: DISCA Clustering CLI (`disca cluster`) + +### 1. What Was Added + +- **New script:** `aitom/bin/disca.py` +- **New console entry:** in `setup.py` + - `disca = aitom.bin.disca:main` + +The CLI adds a `disca` command with a `cluster` subcommand: + +```bash +disca cluster \ + --features features.npy \ + --candidate-k 5,10,20 \ + --out-labels disca_labels.npy \ + --out-summary disca_summary.json +``` + +### 2. Why It Was Added + +DISCA’s pipeline has two stages: + +1. A deep network (YOPO) that extracts feature vectors from subtomograms. +2. A clustering stage (Gaussian Mixture Models, multiple Ks) that assigns each sample to a structural class. + +Previously, the **clustering logic** lived only inside research scripts. There was no: + +- Simple **one-line interface** to run clustering on precomputed features. +- Clean way to integrate clustering into other tools or pipelines. + +The CLI makes the clustering stage: + +- Reusable from the shell or any workflow manager. +- Easy to script and automate (e.g., in bash or Snakemake). +- Clear and discoverable via `disca --help`. + +### 3. How It Works (Exactly) + +**Input:** + +- `features.npy`: NumPy array with shape `(N, D)` (N samples, D-dimensional features). + - Typically these are embeddings produced by DISCA’s YOPO network. + +**Process:** + +1. Loads `features.npy` into memory. +2. For each `K` in `--candidate-k` (e.g., `5,10,20`): + - Fits a `GaussianMixture` model from `sklearn.mixture`. + - Computes the BIC score on the same feature set. +3. Selects the **best K** as the one with **lowest BIC**. +4. Runs `gmm.predict(features)` to obtain integer labels. + +**Output:** + +- `disca_labels.npy` (configurable via `--out-labels`): + - Shape `(N,)`, `int64`. + - Cluster ID for each row of `features.npy`. +- `disca_summary.json` (configurable via `--out-summary`): + - `chosen_k`: the selected K. + - `bic`: BIC score for the selected model. + - `counts`: per-cluster sample counts. + - `candidate_k`, `reg_covar`, `max_iter`, `seed`: parameters used. + +### 4. How It Fits a Real AITom Application + +Typical workflow: + +1. **Prepare subtomograms** using existing tools (picking + extraction). +2. **Extract DISCA embeddings** (existing DISCA scripts): + - Run YOPO network over all subtomograms. + - Save embeddings as `features.npy` (shape `(N, D)`). +3. **Cluster with CLI:** + + ```bash + disca cluster \ + --features features.npy \ + --candidate-k 5,10,20 \ + --out-labels disca_labels.npy \ + --out-summary disca_summary.json + ``` + +4. **Use labels in downstream analysis:** + - Group subtomograms by label. + - Run `average/simple_iterative` or `average/ml/faml` per cluster to get class averages. + - Visualize class averages; evaluate clustering quality (FSC, etc.). + +This turns the **clustering stage** of DISCA into a single, documented step that can be reused in pipelines and integrated with other AITom modules. + +### 5. Validation / Testing + +To validate behavior, we tested the CLI on synthetic data: + +1. Generated 3 clear Gaussian clusters in 2D (N = 60, K = 3). +2. Saved them as `tmp_disca_test/features.npy`. +3. Ran: + + ```bash + python -m aitom.bin.disca cluster \ + --features tmp_disca_test/features.npy \ + --candidate-k 2,3,4 \ + --out-labels tmp_disca_test/labels.npy \ + --out-summary tmp_disca_test/summary.json + ``` + +4. Observed: + - CLI chose **K = 3** (expected). + - Cluster counts were ~20 samples per cluster. + - Labels and summary saved successfully. + +This shows the feature works correctly on a realistic clustering task and is safe to propose in a PR. + + diff --git a/doc/features/DISCA_CLI_Smoke_Test.md b/doc/features/DISCA_CLI_Smoke_Test.md new file mode 100644 index 0000000..70835fb --- /dev/null +++ b/doc/features/DISCA_CLI_Smoke_Test.md @@ -0,0 +1,90 @@ +## Feature: DISCA CLI Synthetic Smoke Test + +### 1. What Was Added + +- **New test file:** `tests/test_disca_cli.py` +- **New dev dependency:** `pytest` added to `requirements.txt` for running tests. + +This test is a **small, synthetic smoke test** that verifies the `disca cluster` command behaves correctly on a simple clustering problem. + +### 2. Why It Was Added + +The DISCA CLI provides a convenient way to run the clustering stage (GMM + BIC) on precomputed features. To keep this behavior stable across: + +- Future code refactors, +- Dependency bumps (e.g., scikit-learn), +- Environment changes, + +we want an automated check that confirms: + +1. The CLI runs without error. +2. It selects a sensible number of clusters. +3. The outputs have correct shapes and internal consistency. + +The smoke test gives maintainers a quick signal that the **core semantics** of the CLI are still intact. + +### 3. How the Test Works + +**File:** `tests/test_disca_cli.py` + +Key steps inside the test: + +1. **Create synthetic data** + - Builds 3 well-separated Gaussian clusters in 2D: + - 20 points near `[0, 0]`, + - 20 points near `[3, 0]`, + - 20 points near `[0, 3]`. + - Concatenated into a `(60, 2)` NumPy array and saved to `features.npy` in a temporary directory (`tmp_path`). + +2. **Run the CLI programmatically** + - Adjusts `sys.path` so the `aitom` package is importable from the repo root. + - Imports `main` from `aitom.bin.disca`. + - Calls: + ```python + disca_main( + [ + "cluster", + "--features", str(features_path), + "--candidate-k", "2,3,4", + "--out-labels", str(labels_path), + "--out-summary", str(summary_path), + ] + ) + ``` + +3. **Validate outputs** + - Asserts that both `labels.npy` and `summary.json` were created. + - Loads labels and summary and checks: + - `labels.shape == (60,)`. + - `summary["chosen_k"] == 3` (correct number of clusters). + - The `counts` in the summary sum to 60 and contain exactly 3 clusters. + +If any of these conditions fail, the test fails, signaling a possible regression. + +### 4. How to Run the Test + +From the repository root: + +```bash +pip install -r requirements.txt # ensures pytest and sklearn are available + +pytest tests/test_disca_cli.py -q +``` + +Expected output: + +```text +. [100%] +1 passed, ... warnings in X.XXs +``` + +### 5. How This Helps the Project + +- Provides a **fast, deterministic check** of the DISCA CLI on a controlled dataset. +- Catches regressions in: + - The argument parsing and CLI wiring, + - The GMM + BIC clustering logic, + - The format and consistency of outputs. +- Can be integrated into future CI to ensure that changes to DISCA or its dependencies do not silently break the clustering interface. + + diff --git a/requirements.txt b/requirements.txt index 163a844..81712a2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,3 +10,4 @@ keras==2.13.1 lsm-db numba six +pytest diff --git a/setup.py b/setup.py index 97a083a..42f4ecd 100644 --- a/setup.py +++ b/setup.py @@ -87,5 +87,6 @@ def get_packages(root_dir='aitom', exclude_dir_roots=['aitom/tomominer/core/src' entry_points={ 'console_scripts': [ 'picking = aitom.bin.picking:main', + 'disca = aitom.bin.disca:main', ]} ) diff --git a/tests/test_disca_cli.py b/tests/test_disca_cli.py new file mode 100644 index 0000000..eabcd76 --- /dev/null +++ b/tests/test_disca_cli.py @@ -0,0 +1,55 @@ +import json +from pathlib import Path +import os +import sys + +import numpy as np + + +def test_disca_cluster_on_synthetic_gaussians(tmp_path: Path): + root = Path(__file__).resolve().parents[1] + sys.path.insert(0, str(root)) + + from aitom.bin.disca import main as disca_main # type: ignore + + rng = np.random.default_rng(0) + n_per = 20 + c1 = rng.normal(loc=[0.0, 0.0], scale=0.3, size=(n_per, 2)) + c2 = rng.normal(loc=[3.0, 0.0], scale=0.3, size=(n_per, 2)) + c3 = rng.normal(loc=[0.0, 3.0], scale=0.3, size=(n_per, 2)) + features = np.vstack([c1, c2, c3]).astype("float32") + + features_path = tmp_path / "features.npy" + labels_path = tmp_path / "labels.npy" + summary_path = tmp_path / "summary.json" + + np.save(features_path, features) + + disca_main( + [ + "cluster", + "--features", + str(features_path), + "--candidate-k", + "2,3,4", + "--out-labels", + str(labels_path), + "--out-summary", + str(summary_path), + ] + ) + + assert labels_path.is_file() + assert summary_path.is_file() + + labels = np.load(labels_path) + with summary_path.open("r", encoding="utf-8") as f: + summary = json.load(f) + + assert labels.shape == (3 * n_per,) + assert int(summary["chosen_k"]) == 3 + counts = {int(k): int(v) for k, v in summary["counts"].items()} + assert sum(counts.values()) == labels.shape[0] + assert len(counts) == 3 + +