protlabel 4.6.0__py3-none-any.whl

protlabel/__init__.py

@@ -0,0 +1,20 @@
+"""protlabel — Embedding Annotation Transfer (EAT) engine.
+
+Nearest-neighbour label transfer in protein-language-model embedding space,
+with the goPredSim reliability index. Pure numpy (plus the standard library);
+no protspace imports.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from protlabel.lookup import Lookup
+from protlabel.transfer import Prediction, eat
+
+try:
+    # Report the installed distribution version rather than a hard-coded literal
+    # that would silently drift across releases.
+    __version__ = version("protlabel")
+except PackageNotFoundError:  # pragma: no cover - source/uninstalled fallback
+    __version__ = "0.0.0"
+
+__all__ = ["Lookup", "Prediction", "eat", "__version__"]

protlabel/backends.py

@@ -0,0 +1,183 @@
+"""Exact (brute-force) k-nearest-neighbour search over reference embeddings.
+
+Distances are computed with a chunked BLAS matrix product (numpy ``@``) plus
+``argpartition`` — the GEMM path, which is roughly an order of magnitude faster
+than a naive per-pair distance loop while staying pure-numpy (no third-party
+distance/ANN library).
+
+The query axis is chunked and, adaptively, bounded against the reference axis so
+the per-chunk distance block is kept at or below ``max_block_bytes`` (default
+256 MiB) regardless of ``n_refs``.  This keeps peak memory close to the
+reference matrix itself even at Swiss-Prot scale (~570 000 references).
+"""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+
+_METRICS = {"euclidean", "cosine"}
+
+# The float32 GEMM distance loses precision to catastrophic cancellation for
+# high-norm pLM embeddings, so its top-k *selection* (argpartition) can drop a
+# true nearest whose float32 distance is noise.  The exact float64 rerank can
+# only reorder the candidates selection kept — it cannot recover a nearest that
+# was never selected.  So over-fetch a wider candidate pool before the rerank;
+# the pad gives small k (e.g. the default k=1) a real margin, the factor scales
+# it for larger k.  Cost is still O(b * k_pool * d) << the O(b * n_refs) GEMM.
+_RERANK_OVERFETCH = 2  # multiplicative widening of the candidate pool
+_RERANK_PAD = 16  # additive floor so small k still over-fetches
+
+# Peak bytes-per-element the float64 rerank holds at once for its
+# (eff_chunk, k_pool, d) tensors: _exact_distances materialises two float64
+# copies simultaneously (the upcast candidates and their difference/normalised
+# form) = 16 B/elem; 24 leaves headroom for the transient f32->f64 upcast.
+_RERANK_PEAK_BYTES = 24
+
+
+def _rerank_pool_size(k: int, n_refs: int) -> int:
+    """Number of float32-selected candidates to rerank in float64 for a top-k."""
+    return min(n_refs, max(k * _RERANK_OVERFETCH, k + _RERANK_PAD))
+
+
+def _exact_distances(block: np.ndarray, sel: np.ndarray, metric: str) -> np.ndarray:
+    """Distances from each query (block[i]) to its pool candidates (sel[i]) in float64.
+
+    The fast GEMM block selects a candidate pool; this recomputes their
+    distances by direct subtraction in float64 to avoid the catastrophic
+    cancellation of ``||q||^2 - 2 q.r + ||r||^2`` for near-identical vectors.
+    Cost is O(b * k_pool * d), not O(b * n_refs), so it is cheap.
+    """
+    blk = block[:, None, :].astype(np.float64)  # (b, 1, d)
+    sel = sel.astype(np.float64)  # (b, k, d)
+    if metric == "euclidean":
+        diff = blk - sel
+        return np.sqrt(np.einsum("bkd,bkd->bk", diff, diff))
+    bn = np.linalg.norm(blk, axis=2, keepdims=True)
+    bn = np.where(bn == 0.0, 1.0, bn)
+    sn = np.linalg.norm(sel, axis=2, keepdims=True)
+    sn = np.where(sn == 0.0, 1.0, sn)
+    cos = np.einsum("bkd,bkd->bk", blk / bn, sel / sn)
+    return np.clip(1.0 - cos, 0.0, 2.0)
+
+
+def nearest(
+    queries: np.ndarray,
+    refs: np.ndarray,
+    k: int,
+    metric: str = "euclidean",
+    chunk: int = 4096,
+    max_block_bytes: int = 256 * 1024 * 1024,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Return (idx, dist) of the k nearest *references* per query.
+
+    idx[i]  -> indices into ``refs`` of the k nearest, ascending by distance.
+    dist[i] -> the corresponding distances.
+    k is capped to the number of references.
+
+    Memory behaviour
+    ----------------
+    The per-chunk distance block has shape ``(query_chunk, n_refs)``.  When
+    ``n_refs`` is large (e.g. Swiss-Prot ~570 000), even a modest ``chunk`` of
+    4096 yields a multi-GiB block.  To bound this, the effective query-chunk
+    size ``eff_chunk`` is computed as ``min(chunk, max_block_bytes // (n_refs *
+    8))`` so each block stays at or below ``max_block_bytes`` (default 256 MiB)
+    independent of ``n_refs``.  The factor 8 budgets for a float64-sized block;
+    each metric holds a single copy of the reference matrix (cosine folds the
+    per-reference norm into the dot product rather than storing a normalized
+    second copy) plus a few float32 ``(query_chunk, n_refs)`` temporaries, so the
+    real peak is comparable to (not far below) that budget.  The exact-distance
+    recompute touches the over-fetched candidate pool (a small multiple of k, not
+    n_refs); ``eff_chunk`` is additionally capped so those ``(eff_chunk, k_pool,
+    d)`` float64 temporaries also stay within ``max_block_bytes`` (they scale with
+    the embedding dim, so a small reference set queried at high dim would
+    otherwise blow the budget).  Peak memory therefore stays bounded — close to
+    the reference matrix at Swiss-Prot scale, laptop-feasible.
+    """
+    if metric not in _METRICS:
+        raise ValueError(f"Unknown metric {metric!r}; expected 'euclidean' or 'cosine'")
+
+    if k < 1:
+        raise ValueError("k must be >= 1")
+
+    queries = np.ascontiguousarray(queries, dtype=np.float32)
+    refs = np.ascontiguousarray(refs, dtype=np.float32)
+    n_refs = refs.shape[0]
+    if n_refs == 0:
+        raise ValueError("refs must contain at least one reference embedding")
+    k = min(k, n_refs)
+    k_pool = _rerank_pool_size(k, n_refs)
+
+    # Adaptively shrink the query chunk so BOTH per-chunk tensors stay within
+    # max_block_bytes:
+    #   * the (eff_chunk, n_refs) distance block (float64-sized = 8 B/row-elem),
+    #   * the (eff_chunk, k_pool, d) float64 rerank temporaries, which scale with
+    #     the embedding dim d (not n_refs), so the block budget alone leaves them
+    #     unbounded for small-reference / high-dim / many-query workloads.
+    dim = refs.shape[1]
+    block_row_bytes = max(1, n_refs * 8)
+    rerank_row_bytes = max(1, k_pool * dim * _RERANK_PEAK_BYTES)
+    eff_chunk = max(
+        1,
+        min(
+            chunk,
+            max_block_bytes // block_row_bytes,
+            max_block_bytes // rerank_row_bytes,
+        ),
+    )
+
+    idx_out = np.empty((queries.shape[0], k), dtype=np.int64)
+    dist_out = np.empty((queries.shape[0], k), dtype=np.float32)
+
+    # Per-reference squared norms (shape (n_refs,), tiny) — used by both metrics:
+    #   euclidean: ||q-r||^2 = ||q||^2 - 2 q.r + ||r||^2
+    #   cosine:    cos = q.r / (||q|| ||r||) — folding ||r|| in here (rather than
+    #              storing a second, full-size normalized copy of the reference
+    #              matrix) keeps cosine at 1x reference memory, like euclidean.
+    refs_sq = np.einsum("ij,ij->i", refs, refs)
+    refs_norm = np.sqrt(refs_sq) if metric == "cosine" else None
+
+    # Row index for the per-chunk fancy-indexing; sliced to each block's length
+    # so it is allocated once rather than per chunk.
+    rows_full = np.arange(eff_chunk)[:, None]
+
+    for start in range(0, queries.shape[0], eff_chunk):
+        block = queries[start : start + eff_chunk]
+        with warnings.catch_warnings():
+            # GEMM on near-overflowing float16-origin values can emit harmless
+            # RuntimeWarnings; the inputs are upcast float32 so values are safe.
+            warnings.simplefilter("ignore", RuntimeWarning)
+            if metric == "euclidean":
+                block_sq = np.einsum("ij,ij->i", block, block)
+                cross = block @ refs.T  # (b, n_refs), BLAS GEMM
+                d2 = block_sq[:, None] - 2.0 * cross + refs_sq[None, :]
+                np.maximum(d2, 0.0, out=d2)  # clip tiny negative fp artifacts
+                d = np.sqrt(d2, dtype=np.float32)
+            else:  # cosine
+                cross = block @ refs.T  # (b, n_refs) raw dot products, BLAS GEMM
+                block_norm = np.sqrt(np.einsum("ij,ij->i", block, block))
+                # Zero-magnitude rows -> safe norm 1; the dot product is 0 there,
+                # so the cosine distance stays a finite 1.0 (matches normalizing a
+                # zero vector to a zero vector).
+                safe_b = np.where(block_norm == 0.0, 1.0, block_norm)
+                safe_r = np.where(refs_norm == 0.0, 1.0, refs_norm)
+                d = (1.0 - cross / (safe_b[:, None] * safe_r[None, :])).astype(
+                    np.float32
+                )
+                np.clip(d, 0.0, 2.0, out=d)  # cosine distance in [0, 2]
+
+        # Over-fetch a candidate pool with the fast (float32) block, recompute
+        # the pool's distances exactly in float64, then take the true top-k from
+        # it — so both the *selection* and the reported distance are robust to
+        # the float32 GEMM's catastrophic cancellation for near-identical (or
+        # high-norm) vectors. Reranking a too-narrow pool cannot recover a true
+        # nearest the float32 selection dropped; see _rerank_pool_size.
+        cand = np.argpartition(d, kth=k_pool - 1, axis=1)[:, :k_pool]  # unsorted pool
+        rows = rows_full[: block.shape[0]]
+        exact = _exact_distances(block, refs[cand], metric)  # (b, k_pool) float64
+        order = np.argsort(exact, axis=1)[:, :k]  # true top-k, ascending by exact dist
+        idx_out[start : start + block.shape[0]] = cand[rows, order]
+        dist_out[start : start + block.shape[0]] = exact[rows, order].astype(np.float32)
+
+    return idx_out, dist_out

protlabel/lookup.py

@@ -0,0 +1,75 @@
+"""A persistable reference lookup: embeddings + ids + labels, plus query().
+
+Serialized as a single .npz sidecar so it can live next to a bundle or in a
+cache dir and be rebuilt on demand from the source HDF5.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+from protlabel.transfer import Prediction, eat
+
+
+@dataclass
+class Lookup:
+    """Reference set for embedding annotation transfer."""
+
+    embeddings: np.ndarray  # (N, D) float32
+    ids: list[str]
+    labels: list[str]
+    metric: str = "euclidean"
+    model: str = field(default="")
+
+    def query(
+        self, query_emb: np.ndarray, query_ids: list[str], *, k: int = 1
+    ) -> list[Prediction]:
+        """Transfer labels to query embeddings from this lookup."""
+        return eat(
+            query_emb,
+            query_ids,
+            self.embeddings,
+            self.ids,
+            self.labels,
+            k=k,
+            metric=self.metric,
+        )
+
+    def save(self, path: Path) -> None:
+        """Serialize to a .npz sidecar.
+
+        Embeddings are stored as float32 (lossless round-trip); ids/labels are
+        stored as unicode arrays rather than pickled object arrays so the sidecar
+        can be loaded with ``allow_pickle=False`` (no arbitrary-code-execution
+        surface on load). ids/labels must not contain trailing NUL bytes, which
+        numpy's fixed-width unicode arrays strip on round-trip.
+        """
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        np.savez(
+            path,
+            # asarray, not astype: no copy of the (potentially multi-GB) matrix
+            # when it is already float32 (the common case).
+            embeddings=np.asarray(self.embeddings, dtype=np.float32),
+            ids=np.asarray(self.ids, dtype=np.str_),
+            labels=np.asarray(self.labels, dtype=np.str_),
+            metric=np.asarray(self.metric, dtype=np.str_),
+            model=np.asarray(self.model, dtype=np.str_),
+        )
+
+    @classmethod
+    def load(cls, path: Path) -> Lookup:
+        """Load a .npz sidecar (with pickling disabled)."""
+        with np.load(path, allow_pickle=False) as data:
+            return cls(
+                # asarray, not astype: the npz array is already float32 (save
+                # writes it as such), so this avoids a redundant full copy.
+                embeddings=np.asarray(data["embeddings"], dtype=np.float32),
+                ids=[str(x) for x in data["ids"]],
+                labels=[str(x) for x in data["labels"]],
+                metric=str(data["metric"]),
+                model=str(data["model"]),
+            )

protlabel/reliability.py

@@ -0,0 +1,45 @@
+"""goPredSim reliability index: map an embedding distance to a [0,1] confidence.
+
+Cosine:     s(d) = 1 - d             (clamped to [0,1]; cosine distance in [0,2])
+Euclidean:  s(d) = 0.5 / (0.5 + d)   (1.0 at d=0, 0.5 at d=0.5, ->0 as d->inf)
+
+Both branches return a value in [0, 1], and a non-finite distance (NaN/inf) maps
+to 0.0 so an invalid neighbour never yields a spuriously high confidence.
+
+The ``protspace transfer`` CLI defaults to cosine — ``1 - d`` is the cosine
+similarity, which is directly interpretable and naturally bounded. The lower-level
+``protlabel`` engine API (``eat``/``Lookup``/``nearest``) keeps euclidean as its
+default to match the published goPredSim reference implementation.
+
+The euclidean transform ``0.5 / (0.5 + d)`` is the published goPredSim constant,
+calibrated against ProtT5 distances. On embedding spaces whose raw distances are
+much larger than ~0.5 (common) it collapses toward 0 even for near neighbours, so
+the *ordering* of confidences stays meaningful but the absolute value is hard to
+interpret.
+
+Reference: Littmann et al., Sci Rep 2021 (Eq. 5); goPredSim calc_reliability_index.
+"""
+
+from __future__ import annotations
+
+import math
+
+
+def similarity(distance: float, metric: str) -> float:
+    """Per-neighbour distance->similarity (the goPredSim reliability transform).
+
+    Always returns a value in [0, 1]. The kNN backend never emits a negative
+    distance (euclidean is a clamped ``sqrt``; cosine distance is in [0, 2]), so
+    the negative-distance branch is purely defensive for direct callers of this
+    function: a negative distance is treated as 0, and a non-finite distance
+    (NaN/inf) maps to 0.0 so an invalid neighbour never yields a spuriously high
+    confidence.
+    """
+    if not math.isfinite(distance):
+        return 0.0
+    if metric == "cosine":
+        return min(1.0, max(0.0, 1.0 - distance))
+    if metric == "euclidean":
+        # d >= 0 makes 0.5 / (0.5 + d) fall in (0, 1] already; no clamp needed.
+        return 0.5 / (0.5 + max(0.0, distance))
+    raise ValueError(f"Unknown metric {metric!r}; expected 'euclidean' or 'cosine'")

protlabel/transfer.py

@@ -0,0 +1,97 @@
+"""Embedding annotation transfer: kNN -> reliability index -> transferred label.
+
+Implements the goPredSim aggregation (Littmann et al. 2021, Eq. 5):
+    RI(p) = (1/eff_k) * sum over neighbours carrying label p of similarity(d),
+where ``eff_k`` is the number of neighbours actually used (``k`` capped to the
+number of references).  The transferred label is argmax RI(p); its source is the
+nearest neighbour carrying that label.  Ties are broken deterministically by
+smallest source distance, then by lexically smallest label, so the result never
+depends on the (arbitrary) ordering of equidistant neighbours.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+from protlabel.backends import nearest
+from protlabel.reliability import similarity
+
+
+@dataclass(frozen=True)
+class Prediction:
+    """One transferred annotation for a query protein."""
+
+    query_id: str
+    label: str
+    source_id: str
+    distance: float
+    reliability: float
+    k: int
+    metric: str
+
+
+def eat(
+    query_emb: np.ndarray,
+    query_ids: list[str],
+    ref_emb: np.ndarray,
+    ref_ids: list[str],
+    ref_labels: list[str],
+    *,
+    k: int = 1,
+    metric: str = "euclidean",
+) -> list[Prediction]:
+    """Transfer the best-guess label to each query from its k nearest references."""
+    if not (len(ref_ids) == len(ref_labels) == ref_emb.shape[0]):
+        raise ValueError("ref_emb, ref_ids and ref_labels must have equal length")
+    if ref_emb.shape[0] == 0:
+        raise ValueError("No reference embeddings to transfer from")
+    if len(query_ids) != query_emb.shape[0]:
+        raise ValueError("query_emb and query_ids must have equal length")
+
+    idx, dist = nearest(query_emb, ref_emb, k=k, metric=metric)
+    eff_k = idx.shape[1]
+    predictions: list[Prediction] = []
+
+    for qi, query_id in enumerate(query_ids):
+        neigh_idx = idx[qi]
+        neigh_dist = dist[qi]
+        # Accumulate RI per label and track the nearest source per label.
+        ri_by_label: dict[str, float] = {}
+        nearest_src: dict[str, tuple[float, str]] = {}
+        for j, ref_i in enumerate(neigh_idx):
+            lab = ref_labels[ref_i]
+            d = float(neigh_dist[j])
+            ri_by_label[lab] = ri_by_label.get(lab, 0.0) + similarity(d, metric)
+            # Break exact-distance ties on the lexically smallest source id so the
+            # reported provenance is independent of the (arbitrary) neighbour order
+            # among equidistant references, matching the label tie-break below.
+            candidate = (d, ref_ids[ref_i])
+            if lab not in nearest_src or candidate < nearest_src[lab]:
+                nearest_src[lab] = candidate
+        # Pick the highest-RI label; break ties deterministically by smallest
+        # source distance, then lexically smallest label. This makes the choice
+        # independent of the order of equidistant neighbours (whose argsort order
+        # is otherwise arbitrary). For distinct distances the nearest neighbour's
+        # label wins, as before.
+        best_label = min(
+            ri_by_label,
+            key=lambda p: (-ri_by_label[p], nearest_src[p][0], p),
+        )
+        # Normalise by eff_k (the goPredSim 1/k term, k capped to n_refs).
+        ri = ri_by_label[best_label] / eff_k
+        src_dist, src_id = nearest_src[best_label]
+        predictions.append(
+            Prediction(
+                query_id=query_id,
+                label=best_label,
+                source_id=src_id,
+                distance=src_dist,
+                reliability=ri,
+                k=eff_k,
+                metric=metric,
+            )
+        )
+
+    return predictions

protlabel-4.6.0.dist-info/METADATA

@@ -0,0 +1,63 @@
+Metadata-Version: 2.4
+Name: protlabel
+Version: 4.6.0
+Summary: Embedding Annotation Transfer (EAT): nearest-neighbour label transfer in protein-language-model embedding space
+Project-URL: Homepage, https://github.com/tsenoner/protspace
+Project-URL: Repository, https://github.com/tsenoner/protspace
+Author-email: Tobias Senoner <tobias.senoner@tum.de>
+License-Expression: GPL-3.0
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.23.0
+Description-Content-Type: text/markdown
+
+# protlabel
+
+**Embedding Annotation Transfer (EAT) engine** — nearest-neighbour label transfer in
+protein-language-model (pLM) embedding space, with the goPredSim reliability index.
+
+`protlabel` is a small, dependency-light library (numpy only). It is
+ProtSpace-agnostic by design — it imports nothing from `protspace` — so it is
+independently testable and reusable from notebooks, `protspace_uniprot`, or any
+other project. The `protspace transfer` CLI is a thin consumer of this engine.
+
+## Install
+
+```bash
+pip install protlabel
+```
+
+## Use
+
+```python
+import numpy as np
+from protlabel import eat, Lookup
+
+preds = eat(
+    query_emb=np.random.rand(3, 1024).astype("float32"),
+    query_ids=["Q1", "Q2", "Q3"],
+    ref_emb=np.random.rand(100, 1024).astype("float32"),
+    ref_ids=[f"R{i}" for i in range(100)],
+    ref_labels=["toxin", "enzyme"] * 50,
+    k=1,
+    metric="cosine",          # or "euclidean" (the goPredSim default)
+)
+for p in preds:
+    print(p.query_id, p.label, round(p.reliability, 3))
+```
+
+`Lookup` builds and serialises a reusable reference set (`.npz` sidecar) so the
+reference matrix can be rebuilt on demand rather than shipped.
+
+## Method
+
+Nearest-neighbour transfer in the *original* pLM embedding space (not a 2-D/3-D
+projection), with the goPredSim reliability index — see Littmann et al.,
+*Sci Rep* 2021 (Eq. 5) and Heinzinger et al., *NAR Genom Bioinform* 2022.
+
+Distances are computed with an exact, chunked brute-force search (numpy BLAS GEMM
++ `argpartition`); queries are processed in batches. No approximate-nearest-neighbour
+index is needed at Swiss-Prot scale.
+
+## License
+
+GPL-3.0 — part of the [ProtSpace](https://github.com/tsenoner/protspace) project.

protlabel-4.6.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+protlabel/__init__.py,sha256=jPPzf3kdVoIBe7iH6ysIKevIQfoYKNcmr6soTZEqduk,726
+protlabel/backends.py,sha256=bhLVnRtEjPD4yEbLU9vEv3GxTe95Sk5eNQGL3ICaFh8,9147
+protlabel/lookup.py,sha256=XoFzeS5_lhHTZQ4BS7QqOQZx2mEfE2rjp_T8yy1hrrs,2689
+protlabel/reliability.py,sha256=B58s27sjecxan2SLZ18GqD34-mvbks_iOwa_z_vSpQw,2118
+protlabel/transfer.py,sha256=4dYGiS0iITZMtqqX1ryqyR8OHD87mQPwq-Zt-TK8ISg,3677
+protlabel-4.6.0.dist-info/METADATA,sha256=Xa4P_x7DgocsIQ8CrGbvFaYBdMR2-_N8T0-8a7O9p7g,2179
+protlabel-4.6.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+protlabel-4.6.0.dist-info/RECORD,,

protlabel-4.6.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any