Add DISCA clustering CLI and documentation

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()
+

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

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

requirements.txt

@@ -10,3 +10,4 @@ keras==2.13.1
 lsm-db
 numba
 six
+pytest

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',
           ]}
       )