point-affinity-transformer 0.1.0__tar.gz

point_affinity_transformer-0.1.0/.gitignore

@@ -0,0 +1,19 @@
+wandb/*
+output/*
+outputs/*
+logs/*
+*.pyc
+util/__pycache__/*
+util/__pycache__/datasets.cpython-39.pyc
+*.png
+output.log
+*.log
+misc/*
+data/*
+site/
+*.egg-info/
+.pytest_cache/
+ckpt/
+dist/
+.ipynb_checkpoints/
+.cache/

point_affinity_transformer-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jakob Troidl
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

point_affinity_transformer-0.1.0/PKG-INFO

@@ -0,0 +1,132 @@
+Metadata-Version: 2.5
+Name: point-affinity-transformer
+Version: 0.1.0
+Summary: Point Affinity Transformers for global neuron shape reasoning
+Project-URL: Paper, https://www.biorxiv.org/content/10.1101/2024.11.24.625067v1
+Project-URL: Source, https://github.com/jakobtroidl/neuron-shape-reasoning
+Author: Jakob Troidl
+License-Expression: MIT
+License-File: LICENSE
+Keywords: connectomics,neuron,point cloud,segmentation,transformer
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.9
+Requires-Dist: navis>=1.4
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: pyyaml
+Requires-Dist: safetensors>=0.4
+Requires-Dist: scikit-learn>=1.2
+Requires-Dist: torch>=2.1
+Requires-Dist: tqdm
+Provides-Extra: dev
+Requires-Dist: einops; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: python-voi==0.0.2; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Requires-Dist: timm; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: mkdocs-material>=9; extra == 'docs'
+Requires-Dist: mkdocstrings[python]>=0.24; extra == 'docs'
+Provides-Extra: pretrained
+Requires-Dist: gdown>=5; extra == 'pretrained'
+Provides-Extra: wandb
+Requires-Dist: wandb; extra == 'wandb'
+Description-Content-Type: text/markdown
+
+[![Paper](https://img.shields.io/badge/paper-arxiv-blue.svg?colorB=4AC8F4)](https://www.biorxiv.org/content/10.1101/2024.11.24.625067v1)
+[![Data](https://img.shields.io/badge/data-gdrive-red.svg?colorB=f25100)](https://drive.google.com/drive/folders/1vgPSYsqDJyhv1s9aD09GgXqAeGm0Rb1V?usp=sharing)
+[![Models](https://img.shields.io/badge/models-gdrive-purple.svg?colorB=C46CFD)](https://drive.google.com/drive/folders/1vgPSYsqDJyhv1s9aD09GgXqAeGm0Rb1V?usp=sharing)
+
+## Point Affinity Transformers
+
+
+### Installation
+```bash
+pip install point-affinity-transformer
+```
+The package is imported as `pat`. Install a PyTorch build matching your CUDA version first ([pytorch.org](https://pytorch.org/get-started/locally/)).
+
+### Getting Started
+
+```python
+import pat
+from pat.io import read_swcs
+
+points, gt = read_swcs(["a.swc", "b.swc"])          # (N, 3) coordinates in µm
+points = points / 659.88                            # your normalization: the same as for the training data
+model = pat.load_pretrained("my_model", device="cuda")  # directory from pat train or pat.save_pretrained
+
+seg = pat.segment(model, points, threshold=0.8)     # one label per point, -1 = noise
+print(seg.num_neurons, pat.metrics.voi(seg.labels, gt))
+
+enc = pat.encode(model, points)                     # encode once ...
+A = enc.affinity(queries, others)                   # ... then query affinities between any points
+```
+
+### Point Clouds and Features
+
+Point clouds of shape `(B, N, 3)` with `D` extra features are passed to the model as `(B, N, 3 + D)`. Training only needs a neuron label per point.
+
+```python
+import torch
+import torch.nn.functional as F
+import pat
+
+B, N, D = 8, 1024, 4                           # random toy data: replace with your own
+xyz = torch.rand(B, N, 3) * 2 - 1              # coordinates, already normalized to roughly [-1, 1]
+features = torch.randn(B, N, D)                # D features per point
+labels = torch.randint(0, 3, (B, N))           # neuron id of every point (training targets)
+pc = torch.cat([xyz, features], dim=-1)        # (B, N, 3 + D)
+
+model = pat.create_model("ae_d256_m64", depth=4, feature_dim=D)
+optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+
+# Training: sample random point pairs, target = 1 if both lie on the same neuron
+def gather(x, idx):                            # x: (B, N, C), idx: (B, Q) -> (B, Q, C)
+    return x.gather(1, idx[..., None].expand(-1, -1, x.shape[-1]))
+
+model.train()
+for step in range(100):
+    i, j = torch.randint(0, N, (2, B, 2048))   # 2048 query pairs per cloud
+    target = (labels.gather(1, i) == labels.gather(1, j)).float()
+    logits = model(pc, None, gather(pc, i), gather(pc, j), logits=True)  # mask=None: no padding
+    loss = F.binary_cross_entropy_with_logits(logits, target)
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+# Inference
+model.eval()
+with torch.no_grad():
+    latents = model.encode(pc)                             # (B, num_latents, dim), once per cloud
+    affinity = model.decode(latents, pc[:, :5], pc[:, 5:10])  # (B, 5): any query pairs, with features
+
+seg = pat.segment(model, pc[0])                # cluster one cloud: labels (N,), -1 = noise
+pat.save_pretrained(model, "my_model")         # reload with pat.load_pretrained("my_model")
+```
+
+Coordinates should lie within roughly [-1, 1]. Clouds of different sizes can be padded with `pat.preprocess.pad_batch`, which returns the `mask` to pass instead of
+`None`. A trained 3D model can be extended to `D` features with `pat.expand_features(model, D)` and then fine-tuned
+the same way. To train from SWC files, run `pat train config.yaml` (examples in [`research/configs`](research/configs)),
+optionally with `data.features: [radius]` or, from Python, `pat.train.train(cfg, feature_fn=...)`.
+
+See the [quickstart notebook](examples/quickstart.ipynb) for an end-to-end example, and the documentation
+(`pip install -e ".[docs]" && mkdocs serve`) for the API reference.
+
+### Citation
+```bibtex
+@techreport{troidlgnsr2024,
+  title = {Global Neuron Shape Reasoning with Point Affinity Transformers},
+  author = {Troidl, Jakob and Knittel, Johannes and Li, Wanhua and Zhan, Fengnang and Pfister*, Hanspeter and Turaga*, Srinivas},
+  journal = {bioRxiv},
+  year = {2024},
+  publisher = {Cold Spring Harbor Laboratory},
+  keywords = {preprint}
+}
+```
+
+### Acknowledgements
+We acknowledge NSF grants CRCNS-2309041, NCS-FO-2124179, and NIH grant R01HD104969. We also thank the HHMI Janelia Visiting Scientist Program and the Harvard Data Science Initiative Postdoctoral Fellowship for their support. The code is partially based on [3DShape2VecSet](https://arxiv.org/abs/2301.11445) by Zhang et. al.
+

point_affinity_transformer-0.1.0/README.md

@@ -0,0 +1,95 @@
+[![Paper](https://img.shields.io/badge/paper-arxiv-blue.svg?colorB=4AC8F4)](https://www.biorxiv.org/content/10.1101/2024.11.24.625067v1)
+[![Data](https://img.shields.io/badge/data-gdrive-red.svg?colorB=f25100)](https://drive.google.com/drive/folders/1vgPSYsqDJyhv1s9aD09GgXqAeGm0Rb1V?usp=sharing)
+[![Models](https://img.shields.io/badge/models-gdrive-purple.svg?colorB=C46CFD)](https://drive.google.com/drive/folders/1vgPSYsqDJyhv1s9aD09GgXqAeGm0Rb1V?usp=sharing)
+
+## Point Affinity Transformers
+
+
+### Installation
+```bash
+pip install point-affinity-transformer
+```
+The package is imported as `pat`. Install a PyTorch build matching your CUDA version first ([pytorch.org](https://pytorch.org/get-started/locally/)).
+
+### Getting Started
+
+```python
+import pat
+from pat.io import read_swcs
+
+points, gt = read_swcs(["a.swc", "b.swc"])          # (N, 3) coordinates in µm
+points = points / 659.88                            # your normalization: the same as for the training data
+model = pat.load_pretrained("my_model", device="cuda")  # directory from pat train or pat.save_pretrained
+
+seg = pat.segment(model, points, threshold=0.8)     # one label per point, -1 = noise
+print(seg.num_neurons, pat.metrics.voi(seg.labels, gt))
+
+enc = pat.encode(model, points)                     # encode once ...
+A = enc.affinity(queries, others)                   # ... then query affinities between any points
+```
+
+### Point Clouds and Features
+
+Point clouds of shape `(B, N, 3)` with `D` extra features are passed to the model as `(B, N, 3 + D)`. Training only needs a neuron label per point.
+
+```python
+import torch
+import torch.nn.functional as F
+import pat
+
+B, N, D = 8, 1024, 4                           # random toy data: replace with your own
+xyz = torch.rand(B, N, 3) * 2 - 1              # coordinates, already normalized to roughly [-1, 1]
+features = torch.randn(B, N, D)                # D features per point
+labels = torch.randint(0, 3, (B, N))           # neuron id of every point (training targets)
+pc = torch.cat([xyz, features], dim=-1)        # (B, N, 3 + D)
+
+model = pat.create_model("ae_d256_m64", depth=4, feature_dim=D)
+optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+
+# Training: sample random point pairs, target = 1 if both lie on the same neuron
+def gather(x, idx):                            # x: (B, N, C), idx: (B, Q) -> (B, Q, C)
+    return x.gather(1, idx[..., None].expand(-1, -1, x.shape[-1]))
+
+model.train()
+for step in range(100):
+    i, j = torch.randint(0, N, (2, B, 2048))   # 2048 query pairs per cloud
+    target = (labels.gather(1, i) == labels.gather(1, j)).float()
+    logits = model(pc, None, gather(pc, i), gather(pc, j), logits=True)  # mask=None: no padding
+    loss = F.binary_cross_entropy_with_logits(logits, target)
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+# Inference
+model.eval()
+with torch.no_grad():
+    latents = model.encode(pc)                             # (B, num_latents, dim), once per cloud
+    affinity = model.decode(latents, pc[:, :5], pc[:, 5:10])  # (B, 5): any query pairs, with features
+
+seg = pat.segment(model, pc[0])                # cluster one cloud: labels (N,), -1 = noise
+pat.save_pretrained(model, "my_model")         # reload with pat.load_pretrained("my_model")
+```
+
+Coordinates should lie within roughly [-1, 1]. Clouds of different sizes can be padded with `pat.preprocess.pad_batch`, which returns the `mask` to pass instead of
+`None`. A trained 3D model can be extended to `D` features with `pat.expand_features(model, D)` and then fine-tuned
+the same way. To train from SWC files, run `pat train config.yaml` (examples in [`research/configs`](research/configs)),
+optionally with `data.features: [radius]` or, from Python, `pat.train.train(cfg, feature_fn=...)`.
+
+See the [quickstart notebook](examples/quickstart.ipynb) for an end-to-end example, and the documentation
+(`pip install -e ".[docs]" && mkdocs serve`) for the API reference.
+
+### Citation
+```bibtex
+@techreport{troidlgnsr2024,
+  title = {Global Neuron Shape Reasoning with Point Affinity Transformers},
+  author = {Troidl, Jakob and Knittel, Johannes and Li, Wanhua and Zhan, Fengnang and Pfister*, Hanspeter and Turaga*, Srinivas},
+  journal = {bioRxiv},
+  year = {2024},
+  publisher = {Cold Spring Harbor Laboratory},
+  keywords = {preprint}
+}
+```
+
+### Acknowledgements
+We acknowledge NSF grants CRCNS-2309041, NCS-FO-2124179, and NIH grant R01HD104969. We also thank the HHMI Janelia Visiting Scientist Program and the Harvard Data Science Initiative Postdoctoral Fellowship for their support. The code is partially based on [3DShape2VecSet](https://arxiv.org/abs/2301.11445) by Zhang et. al.
+

point_affinity_transformer-0.1.0/pyproject.toml

@@ -0,0 +1,61 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "point-affinity-transformer"
+version = "0.1.0"
+description = "Point Affinity Transformers for global neuron shape reasoning"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.9"
+authors = [{ name = "Jakob Troidl" }]
+keywords = ["connectomics", "neuron", "point cloud", "transformer", "segmentation"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+dependencies = [
+    "torch>=2.1",
+    "numpy",
+    "scikit-learn>=1.2",
+    "safetensors>=0.4",
+    "navis>=1.4",
+    "pandas",
+    "pyyaml",
+    "tqdm",
+]
+
+[project.optional-dependencies]
+pretrained = ["gdown>=5"]
+wandb = ["wandb"]
+docs = ["mkdocs-material>=9", "mkdocstrings[python]>=0.24"]
+# timm, einops and python-voi are for comparing against research/legacy;
+# python-voi 0.0.3 ships an empty wheel, so pin 0.0.2.
+dev = ["pytest", "ruff", "timm", "einops", "python-voi==0.0.2"]
+
+[project.scripts]
+pat = "pat.cli:main"
+
+[project.urls]
+Paper = "https://www.biorxiv.org/content/10.1101/2024.11.24.625067v1"
+Source = "https://github.com/jakobtroidl/neuron-shape-reasoning"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/pat"]
+
+[tool.hatch.build.targets.sdist]
+include = ["/src/pat", "/tests", "/README.md", "/LICENSE"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 120
+src = ["src"]
+target-version = "py39"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "B"]

point_affinity_transformer-0.1.0/src/pat/__init__.py

@@ -0,0 +1,30 @@
+"""Point Affinity Transformer: global neuron shape reasoning on point clouds."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from . import cluster, metrics
+from .inference import EncodedCloud, Segmentation, encode, segment
+from .model import PATConfig, PointAffinityTransformer, expand_features
+from .presets import PRESETS, create_model
+from .pretrained import load_pretrained, save_pretrained
+
+try:
+    __version__ = version("point-affinity-transformer")  # defined in pyproject.toml
+except PackageNotFoundError:  # running from a source tree without installing
+    __version__ = "unknown"
+
+__all__ = [
+    "PATConfig",
+    "PointAffinityTransformer",
+    "EncodedCloud",
+    "Segmentation",
+    "encode",
+    "segment",
+    "create_model",
+    "PRESETS",
+    "load_pretrained",
+    "save_pretrained",
+    "expand_features",
+    "cluster",
+    "metrics",
+]

point_affinity_transformer-0.1.0/src/pat/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+main()

point_affinity_transformer-0.1.0/src/pat/cli.py

@@ -0,0 +1,85 @@
+"""Command line interface: ``pat train | eval | convert``."""
+
+from __future__ import annotations
+
+import argparse
+
+
+def _train(args):
+    from .train import load_config, train
+
+    train(load_config(args.config, args.set))
+
+
+def _eval(args):
+    import json
+
+    import torch
+
+    from .pretrained import load_pretrained
+    from .train import MergedNeuronDataset, SkeletonStore, TensorDirDataset, evaluate, format_results, load_config
+    from .train.data import read_root_ids
+
+    cfg = load_config(args.config, args.set)
+    torch.manual_seed(cfg.seed)
+    model = load_pretrained(args.model, device=args.device or cfg.device)
+    if args.items:
+        dataset = TensorDirDataset(args.items)
+    else:
+        store = SkeletonStore(cfg.data.swc_dir, cfg.data.units, cfg.data.dust_size, cfg.data.cache_dir,
+                              features=cfg.data.features)
+        ids = read_root_ids(cfg.data.eval_neuron_ids or cfg.data.neuron_ids)
+        dataset = MergedNeuronDataset(store, ids, cfg.data, train=False)
+    thresholds = [float(t) for t in args.thresholds.split(",")]
+    results = evaluate(model, dataset, thresholds, min_cluster_size=args.min_cluster_size,
+                       symmetrize=args.symmetrize, limit=args.limit)
+    print(format_results(results))
+    if args.out:
+        with open(args.out, "w") as f:
+            json.dump(results, f, indent=2)
+
+
+def _convert(args):
+    import torch
+
+    from .pretrained import load_legacy_checkpoint, save_pretrained
+
+    model = load_legacy_checkpoint(args.checkpoint, preset=args.preset)
+    print(save_pretrained(model, args.out, dtype=torch.float16 if args.fp16 else None))
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="pat", description="Point Affinity Transformer")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p = sub.add_parser("train", help="train or fine-tune a model")
+    p.add_argument("config", help="YAML training config")
+    p.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE", help="override config entries")
+    p.set_defaults(func=_train)
+
+    p = sub.add_parser("eval", help="evaluate segmentation on merged clouds")
+    p.add_argument("config", nargs="?", help="YAML config describing the data (not needed with --items)")
+    p.add_argument("--model", required=True, help="model directory (from pat train or save_pretrained)")
+    p.add_argument("--items", help="directory of pre-generated items (pc_*.pt, labels_*.pt, mask_*.pt)")
+    p.add_argument("--thresholds", default="0.8")
+    p.add_argument("--min-cluster-size", type=int, default=30)
+    p.add_argument("--symmetrize", action="store_true", help="symmetrize affinities before clustering")
+    p.add_argument("--limit", type=int, help="evaluate only the first N items")
+    p.add_argument("--device")
+    p.add_argument("--out", help="write results as JSON")
+    p.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE")
+    p.set_defaults(func=_eval)
+
+    p = sub.add_parser("convert", help="convert a research .pth checkpoint into a model directory")
+    p.add_argument("checkpoint")
+    p.add_argument("out")
+    p.add_argument("--preset", help="model preset; overrides the one stored in the checkpoint")
+    p.add_argument("--fp16", action="store_true", help="store weights in float16")
+    p.set_defaults(func=_convert)
+
+    args = parser.parse_args(argv)
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

point_affinity_transformer-0.1.0/src/pat/cluster.py

@@ -0,0 +1,97 @@
+"""Clustering of affinity matrices into neuron labels.
+
+The label ``-1`` marks noise: DBSCAN outliers and, after
+[`remove_small_clusters`][pat.cluster.remove_small_clusters], points of clusters
+smaller than ``min_cluster_size``.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from sklearn import cluster as _skcluster
+
+NOISE = -1
+
+
+def to_numpy(x) -> np.ndarray:
+    """Convert a tensor or array-like to a NumPy array."""
+    return x.detach().cpu().numpy() if torch.is_tensor(x) else np.asarray(x)
+
+
+def num_clusters(labels) -> int:
+    """Number of clusters in ``labels`` numbered ``0..K-1``, ignoring noise (``-1``)."""
+    labels = torch.as_tensor(labels)
+    return int(labels.max().item()) + 1 if (labels >= 0).any() else 0
+
+
+def agglomerative(distance, threshold: float = 0.8, linkage: str = "average") -> torch.Tensor:
+    """Agglomerative clustering of a precomputed ``(N, N)`` distance matrix.
+
+    Clusters are merged while their linkage distance is below ``threshold``.
+    """
+    d = to_numpy(distance)
+    model = _skcluster.AgglomerativeClustering(
+        n_clusters=None, metric="precomputed", distance_threshold=threshold, linkage=linkage
+    )
+    return torch.as_tensor(model.fit_predict(d), dtype=torch.int64)
+
+
+def dbscan(distance, eps: float = 0.2, min_samples: int = 5) -> torch.Tensor:
+    """DBSCAN on a precomputed distance matrix. Noise points get label ``-1``."""
+    d = to_numpy(distance)
+    labels = _skcluster.DBSCAN(eps=eps, min_samples=min_samples, metric="precomputed").fit_predict(d)
+    return torch.as_tensor(labels, dtype=torch.int64)
+
+
+def remove_small_clusters(labels: torch.Tensor, min_size: int) -> torch.Tensor:
+    """Mark clusters with fewer than ``min_size`` points as noise and relabel.
+
+    The remaining clusters are numbered ``0..K-1``, largest first. Existing
+    ``-1`` labels stay noise.
+    """
+    labels = torch.as_tensor(labels, dtype=torch.int64)
+    out = torch.full_like(labels, NOISE)
+    valid = labels >= 0
+    if not valid.any():
+        return out
+    ids, inverse, counts = torch.unique(labels[valid], return_inverse=True, return_counts=True)
+    keep = counts >= min_size
+    order = torch.argsort(counts.masked_fill(~keep, -1), descending=True, stable=True)
+    new_id = torch.full_like(ids, NOISE)
+    new_id[order[: int(keep.sum())]] = torch.arange(int(keep.sum()))
+    out[valid] = new_id[inverse]
+    return out
+
+
+def cluster_affinity(
+    affinity,
+    threshold: float = 0.8,
+    min_cluster_size: int = 30,
+    method: str = "agglomerative",
+    symmetrize: bool = True,
+) -> torch.Tensor:
+    """Turn an ``(N, N)`` affinity matrix into one label per point.
+
+    Args:
+        affinity: Pairwise same-neuron probabilities.
+        threshold: Distance threshold on ``1 - affinity``. For average linkage,
+            two groups merge while their mean affinity exceeds ``1 - threshold``.
+            For DBSCAN this is ``eps``.
+        min_cluster_size: Smaller clusters become noise (``-1``). Use ``0`` to keep all.
+        method: ``"agglomerative"`` or ``"dbscan"``.
+        symmetrize: Average ``A`` and ``A.T`` first. The model is not exactly
+            symmetric in its two inputs; the paper evaluation used ``False``.
+    """
+    a = to_numpy(affinity).astype(np.float32)
+    if symmetrize:
+        a = 0.5 * (a + a.T)
+    dist = np.clip(1.0 - a, 0.0, None)
+    np.fill_diagonal(dist, 0.0)
+    if method == "agglomerative":
+        labels = agglomerative(dist, threshold=threshold)
+    elif method == "dbscan":
+        labels = dbscan(dist, eps=threshold)
+    else:
+        raise ValueError(f"Unknown method {method!r}")
+    return remove_small_clusters(labels, min_cluster_size) if min_cluster_size > 0 else labels