openprotein-python 0.13.1__tar.gz → 0.14.0__tar.gz

{openprotein_python-0.13.1 → openprotein_python-0.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.13.1
+Version: 0.14.0
 Summary: OpenProtein Python interface.
 Author-email: Mark Gee <markgee@ne47.bio>, "Timothy Truong Jr." <ttruong@ne47.bio>, Tristan Bepler <tbepler@ne47.bio>
 License-Expression: MIT
@@ -14,7 +14,7 @@ Requires-Dist: pandas<3,>=2.2.2
 Requires-Dist: pydantic<3,>=2.5
 Requires-Dist: requests<3,>=2.32.3
 Requires-Dist: tomli<3,>=2.3.0
-Requires-Dist: tqdm<5,>=4.66.5
+Requires-Dist: tqdm<5,>=4.66.4
 Description-Content-Type: text/markdown
 
 [![PyPI version](https://badge.fury.io/py/openprotein-python.svg)](https://pypi.org/project/openprotein-python/)

{openprotein_python-0.13.1 → openprotein_python-0.14.0}/openprotein/__init__.py

@@ -20,6 +20,7 @@ from openprotein.fold import FoldAPI
 from openprotein.models import ModelsAPI
 from openprotein.svd import SVDAPI
 from openprotein.umap import UMAPAPI
+from openprotein.clustering import ClusteringAPI
 from openprotein.predictor import PredictorAPI
 from openprotein.design import DesignAPI
 from openprotein.jobs import Future
@@ -38,6 +39,7 @@ class OpenProtein(APISession):
     _embeddings = None
     _svd = None
     _umap = None
+    _clustering = None
     _fold = None
     _predictor = None
     _design = None
@@ -81,7 +83,7 @@ class OpenProtein(APISession):
     @property
     def prompt(self) -> PromptAPI:
         """
-        The Align submodule gives access to the sequence alignment capabilities by building MSAs and prompts that can be used with PoET.
+        The prompt submodule gives access to creating and retrieving prompts and queries used to condition PoET models.
         """
         if self._prompt is None:
             self._prompt = PromptAPI(self)
@@ -120,6 +122,15 @@ class OpenProtein(APISession):
             )
         return self._umap
 
+    @property
+    def clustering(self) -> ClusteringAPI:
+        """The clustering API for fitting hierarchical clustering jobs."""
+        if self._clustering is None:
+            self._clustering = ClusteringAPI(
+                session=self,
+            )
+        return self._clustering
+
     @property
     def predictor(self) -> PredictorAPI:
         """

openprotein_python-0.14.0/openprotein/clustering/__init__.py

@@ -0,0 +1,14 @@
+"""Clustering module for OpenProtein.
+
+isort:skip_file
+"""
+
+from .schemas import (
+    ClusteringMetadata,
+    HierarchicalClusteringResult,
+    HierarchicalFitJob,
+    LinkageMethod,
+    Metric,
+)
+from .models import HierarchicalClusteringFuture
+from .clustering import ClusteringAPI

openprotein_python-0.14.0/openprotein/clustering/api.py

@@ -0,0 +1,117 @@
+"""Clustering REST API — HTTP calls to the backend."""
+
+from pydantic import TypeAdapter
+
+from openprotein.base import APISession
+from openprotein.errors import APIError, InvalidParameterError
+
+from .schemas import (
+    ClusteringMetadata,
+    HierarchicalClusteringResult,
+    HierarchicalFitJob,
+)
+
+PATH_PREFIX = "v1/clustering"
+
+
+def clustering_list_get(
+    session: APISession,
+    method: str | None = None,
+    page_size: int | None = None,
+    page_offset: int | None = None,
+) -> list[ClusteringMetadata]:
+    """List clustering jobs, optionally filtered by method."""
+    params: dict = {}
+    if method is not None:
+        params["method"] = method
+    if page_size is not None:
+        params["page_size"] = page_size
+    if page_offset is not None:
+        params["page_offset"] = page_offset
+    response = session.get(PATH_PREFIX, params=params or None)
+    return TypeAdapter(list[ClusteringMetadata]).validate_python(response.json())
+
+
+def clustering_get(session: APISession, clustering_id: str) -> ClusteringMetadata:
+    """Fetch clustering job metadata."""
+    response = session.get(f"{PATH_PREFIX}/{clustering_id}")
+    return ClusteringMetadata.model_validate(response.json())
+
+
+def clustering_get_result(
+    session: APISession, clustering_id: str
+) -> HierarchicalClusteringResult:
+    """Fetch the clustering result (linkage + leaf_order). Sequences are NOT
+    filled by this function — callers that need them should call
+    `clustering_get_sequences` and assign to `.sequences`."""
+    response = session.get(f"{PATH_PREFIX}/{clustering_id}/result")
+    return HierarchicalClusteringResult.model_validate(response.json())
+
+
+def clustering_get_sequences(session: APISession, clustering_id: str) -> list[bytes]:
+    """Fetch the input sequences used for the clustering job."""
+    response = session.get(f"{PATH_PREFIX}/{clustering_id}/sequences")
+    return TypeAdapter(list[bytes]).validate_python(response.json())
+
+
+def clustering_delete(session: APISession, clustering_id: str) -> bool:
+    """Delete a clustering job."""
+    response = session.delete(f"{PATH_PREFIX}/{clustering_id}")
+    if 200 <= response.status_code < 300:
+        return True
+    raise APIError(response.text)
+
+
+def clustering_redispatch_post(
+    session: APISession, clustering_id: str
+) -> HierarchicalFitJob:
+    """Redispatch a clustering job."""
+    response = session.post(f"{PATH_PREFIX}/{clustering_id}/redispatch")
+    return HierarchicalFitJob.model_validate(response.json())
+
+
+def clustering_hierarchical_post(
+    session: APISession,
+    model_id: str,
+    feature_type: str,
+    linkage_method: str,
+    metric: str,
+    sequences: list[bytes] | list[str] | None = None,
+    assay_id: str | None = None,
+    reduction: str | None = None,
+    svd_id: str | None = None,
+    force_recompute: bool = False,
+    **kwargs,
+) -> HierarchicalFitJob:
+    """POST to create a hierarchical clustering fit job.
+
+    Parameters
+    ----------
+    force_recompute : bool
+        If True, send ?force=true so the backend bypasses the result cache and recomputes. Default: False.
+    """
+    body: dict = {
+        "model_id": model_id,
+        "feature_type": feature_type,
+        "linkage_method": linkage_method,
+        "metric": metric,
+    }
+    if reduction is not None:
+        body["reduction"] = reduction
+    if svd_id is not None:
+        body["svd_id"] = svd_id
+    if sequences is not None:
+        if assay_id is not None:
+            raise InvalidParameterError("Expected only either sequences or assay_id")
+        body["sequences"] = [
+            (s if isinstance(s, str) else s.decode()) for s in sequences
+        ]
+    else:
+        if assay_id is None:
+            raise InvalidParameterError("Expected either sequences or assay_id")
+        body["assay_id"] = assay_id
+    body.update(**kwargs)
+
+    params = {"force": "true"} if force_recompute else None
+    response = session.post(f"{PATH_PREFIX}/hierarchical", json=body, params=params)
+    return HierarchicalFitJob.model_validate(response.json())

openprotein_python-0.14.0/openprotein/clustering/clustering.py

@@ -0,0 +1,194 @@
+"""ClusteringAPI — user-facing entry point for clustering jobs."""
+
+import typing
+
+from openprotein.base import APISession
+from openprotein.common import Feature, FeatureType, Reduction, ReductionType
+from openprotein.data import AssayDataset, AssayMetadata
+from openprotein.embeddings import EmbeddingModel, EmbeddingsAPI
+from openprotein.errors import InvalidParameterError
+from openprotein.svd import SVDAPI, SVDModel
+
+from . import api
+from .models import HierarchicalClusteringFuture
+from .schemas import LinkageMethod, Metric
+
+
+class ClusteringAPI:
+    """Top-level clustering API. Use `session.clustering.hierarchical(...)` to
+    fit a hierarchical clustering job on a sequence set."""
+
+    def __init__(self, session: APISession):
+        self.session = session
+
+    @typing.overload
+    def hierarchical(
+        self,
+        model: EmbeddingModel,
+        reduction: Reduction | ReductionType,
+        feature_type: FeatureType = FeatureType.PLM,
+        linkage_method: LinkageMethod | str = LinkageMethod.WARD,
+        metric: Metric | str = Metric.EUCLIDEAN,
+        sequences: list[bytes] | list[str] | None = None,
+        assay: AssayDataset | AssayMetadata | str | None = None,
+        **kwargs,
+    ) -> HierarchicalClusteringFuture: ...
+
+    @typing.overload
+    def hierarchical(
+        self,
+        model: SVDModel,
+        reduction: None = None,
+        feature_type: FeatureType = FeatureType.SVD,
+        linkage_method: LinkageMethod | str = LinkageMethod.WARD,
+        metric: Metric | str = Metric.EUCLIDEAN,
+        sequences: list[bytes] | list[str] | None = None,
+        assay: AssayDataset | AssayMetadata | str | None = None,
+        **kwargs,
+    ) -> HierarchicalClusteringFuture: ...
+
+    def hierarchical(
+        self,
+        model: EmbeddingModel | SVDModel | str,
+        reduction: Reduction | ReductionType | None = None,
+        feature_type: Feature | FeatureType | None = None,
+        linkage_method: LinkageMethod | str = LinkageMethod.WARD,
+        metric: Metric | str = Metric.EUCLIDEAN,
+        sequences: list[bytes] | list[str] | None = None,
+        assay: AssayDataset | AssayMetadata | str | None = None,
+        **kwargs,
+    ) -> HierarchicalClusteringFuture:
+        """Fit a hierarchical clustering on sequences, returning a future
+        that resolves to a HierarchicalClusteringResult (scipy linkage + leaf order)."""
+        # resolve assay_id
+        assay_id = (
+            assay.assay_id
+            if isinstance(assay, AssayMetadata)
+            else assay.id
+            if isinstance(assay, AssayDataset)
+            else assay
+        )
+        if sequences is not None and assay_id is not None:
+            raise InvalidParameterError(
+                "Expected only either sequences or assay, not both"
+            )
+        if sequences is not None:
+            n = len(sequences)
+            if n < 2:
+                raise InvalidParameterError(
+                    f"clustering requires at least 2 sequences, got {n}"
+                )
+            if n > 10000:
+                raise InvalidParameterError(
+                    f"clustering size cap exceeded: N={n} > 10000"
+                )
+        # infer feature_type from model type
+        feature_type = (
+            FeatureType.PLM
+            if isinstance(model, EmbeddingModel)
+            else FeatureType.SVD
+            if isinstance(model, SVDModel)
+            else feature_type
+        )
+        if feature_type is None:
+            raise InvalidParameterError(
+                "Expected feature_type to be provided if passing str model_id as model"
+            )
+        if isinstance(feature_type, str):
+            feature_type = FeatureType(feature_type)
+        if isinstance(reduction, str):
+            reduction = ReductionType(reduction)
+        # combo validation — mirrors the Go server rule
+        _WARD_ONLY_METRICS = {"ward", "centroid", "median"}
+        lm_val = (
+            linkage_method.value
+            if isinstance(linkage_method, LinkageMethod)
+            else str(linkage_method)
+        )
+        m_val = metric.value if isinstance(metric, Metric) else str(metric)
+        if lm_val in _WARD_ONLY_METRICS and m_val != "euclidean":
+            raise InvalidParameterError(
+                f"linkage_method={lm_val!r} requires metric='euclidean', got {m_val!r}"
+            )
+        # resolve model_id and svd_id
+        svd_id: str | None = None
+        if feature_type == FeatureType.PLM:
+            if reduction is None:
+                raise InvalidParameterError(
+                    "Expected reduction when using PLM feature_type"
+                )
+            if isinstance(model, str):
+                embeddings_api = getattr(self.session, "embedding", None)
+                assert isinstance(embeddings_api, EmbeddingsAPI)
+                model = embeddings_api.get_model(model)
+            assert isinstance(model, EmbeddingModel), "Expected EmbeddingModel"
+            model_id = model.id
+        elif feature_type == FeatureType.SVD:
+            if reduction is not None:
+                raise InvalidParameterError(
+                    "Unexpected reduction when using SVD feature_type"
+                )
+            if isinstance(model, str):
+                svd_api = getattr(self.session, "svd", None)
+                assert isinstance(svd_api, SVDAPI)
+                model = svd_api.get_svd(model)
+            assert isinstance(model, SVDModel), "Expected SVDModel"
+            model_id = model.id
+            svd_id = model.id
+        else:
+            raise InvalidParameterError(f"Unsupported feature_type: {feature_type}")
+
+        linkage_method = (
+            linkage_method.value
+            if isinstance(linkage_method, LinkageMethod)
+            else str(linkage_method)
+        )
+        metric = metric.value if isinstance(metric, Metric) else str(metric)
+        reduction_str = (
+            reduction.value if isinstance(reduction, ReductionType) else reduction
+        )
+
+        # Advanced flags such as `force_recompute` are accepted via **kwargs
+        # only (intentionally kept out of the typed signature / autocomplete):
+        # they bypass the backend result cache and are easy to misuse.
+        job = api.clustering_hierarchical_post(
+            session=self.session,
+            model_id=model_id,
+            feature_type=feature_type.value,
+            linkage_method=linkage_method,
+            metric=metric,
+            sequences=sequences,
+            assay_id=assay_id,
+            reduction=reduction_str,
+            svd_id=svd_id,
+            **kwargs,
+        )
+        return HierarchicalClusteringFuture(session=self.session, job=job)
+
+    def get(self, clustering_id: str) -> HierarchicalClusteringFuture:
+        """Fetch a clustering job by ID."""
+        # Single-method dispatch today; when new methods are added,
+        # branch on metadata.method to pick the concrete future type.
+        metadata = api.clustering_get(self.session, clustering_id)
+        return HierarchicalClusteringFuture(session=self.session, metadata=metadata)
+
+    def list(
+        self,
+        method: str | None = None,
+        page_size: int | None = None,
+        page_offset: int | None = None,
+    ) -> list[HierarchicalClusteringFuture]:
+        """List clustering jobs, optionally filtering by method.
+
+        Pass `page_size` / `page_offset` to paginate."""
+        # Single-method dispatch today; when new methods are added,
+        # branch on metadata.method to pick the concrete future type.
+        return [
+            HierarchicalClusteringFuture(session=self.session, metadata=md)
+            for md in api.clustering_list_get(
+                self.session,
+                method=method,
+                page_size=page_size,
+                page_offset=page_offset,
+            )
+        ]

openprotein_python-0.14.0/openprotein/clustering/models.py

@@ -0,0 +1,75 @@
+"""HierarchicalClusteringFuture — the future type returned by ClusteringAPI calls."""
+
+from openprotein.base import APISession
+from openprotein.jobs import Future, JobsAPI
+
+from . import api
+from .schemas import ClusteringMetadata, HierarchicalClusteringResult, HierarchicalFitJob
+
+
+class HierarchicalClusteringFuture(Future["HierarchicalClusteringResult"]):
+    """Future for a hierarchical clustering job. Waits for the fit job to complete, then
+    returns a HierarchicalClusteringResult (linkage + leaf_order + sequences)."""
+
+    job: HierarchicalFitJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: HierarchicalFitJob | None = None,
+        metadata: ClusteringMetadata | None = None,
+    ):
+        if metadata is None:
+            if job is None:
+                raise ValueError("Expected clustering metadata or job")
+            metadata = api.clustering_get(session, job.job_id)
+        self._metadata = metadata
+        if job is None:
+            jobs_api = getattr(session, "jobs", None)
+            assert isinstance(jobs_api, JobsAPI)
+            job = HierarchicalFitJob.create(
+                jobs_api.get_job(job_id=metadata.id)
+            )
+        self._sequences: list[bytes] | None = None
+        super().__init__(session, job)
+
+    def __str__(self) -> str:
+        return str(self.metadata)
+
+    def __repr__(self) -> str:
+        return repr(self.metadata)
+
+    @property
+    def id(self) -> str:
+        return self._metadata.id
+
+    @property
+    def metadata(self) -> ClusteringMetadata:
+        """Metadata — lazily refreshes until the job reaches a terminal state."""
+        if not self._metadata.is_done():
+            self._metadata = api.clustering_get(self.session, self._metadata.id)
+        return self._metadata
+
+    @property
+    def sequences(self) -> list[bytes]:
+        """Input sequences used to fit the clustering job (cached)."""
+        if self._sequences is None:
+            self._sequences = api.clustering_get_sequences(
+                session=self.session, clustering_id=self.id
+            )
+        return self._sequences
+
+    def _get(self, verbose: bool = False) -> HierarchicalClusteringResult:
+        """Fetch the clustering result. Called by Future.get after SUCCESS."""
+        result = api.clustering_get_result(self.session, self.id)
+        result.sequences = self.sequences
+        return result
+
+    def delete(self) -> bool:
+        """Delete this clustering job."""
+        return api.clustering_delete(self.session, self.id)
+
+    def redispatch(self) -> "HierarchicalClusteringFuture":
+        """Redispatch this clustering job."""
+        job = api.clustering_redispatch_post(self.session, self.id)
+        return HierarchicalClusteringFuture(session=self.session, job=job)

openprotein_python-0.14.0/openprotein/clustering/schemas.py

@@ -0,0 +1,85 @@
+"""Schemas for OpenProtein clustering."""
+
+from datetime import datetime
+from enum import Enum
+from typing import Literal
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from openprotein.common import FeatureType
+from openprotein.jobs import Job, JobStatus, JobType
+
+
+class LinkageMethod(str, Enum):
+    """Hierarchical clustering linkage methods (scipy.cluster.hierarchy)."""
+
+    WARD = "ward"
+    SINGLE = "single"
+    COMPLETE = "complete"
+    AVERAGE = "average"
+    WEIGHTED = "weighted"
+    CENTROID = "centroid"
+    MEDIAN = "median"
+
+
+class Metric(str, Enum):
+    """Pairwise distance metrics supported by scipy.spatial.distance.pdist."""
+
+    EUCLIDEAN = "euclidean"
+    COSINE = "cosine"
+    CORRELATION = "correlation"
+    HAMMING = "hamming"
+    CHEBYSHEV = "chebyshev"
+    CITYBLOCK = "cityblock"
+    SQEUCLIDEAN = "sqeuclidean"
+    CANBERRA = "canberra"
+    BRAYCURTIS = "braycurtis"
+
+
+class ClusteringMetadata(BaseModel):
+    id: str
+    status: JobStatus
+    created_date: datetime | None = None
+    method: str
+    linkage_method: LinkageMethod | None = None
+    metric: Metric | None = None
+    model_id: str
+    feature_type: FeatureType
+    reduction: str | None = None
+    svd_id: str | None = None
+
+    def is_done(self) -> bool:
+        return self.status.done()
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+class HierarchicalFitJob(Job):
+    job_type: Literal[JobType.clustering_hierarchical]
+
+
+class HierarchicalClusteringResult(BaseModel):
+    """Result of a hierarchical clustering job.
+
+    `linkage` is the scipy linkage matrix with shape (N-1, 4); pass it
+    directly to `scipy.cluster.hierarchy.dendrogram` or `fcluster`.
+    `sequences` is filled by `HierarchicalClusteringFuture._get` after the API fetch.
+
+    Note: `linkage` is a numpy array, so this model is not JSON-serializable via
+    `model_dump_json()` (use `linkage.tolist()` first if you need to serialize).
+    """
+
+    n_leaves: int
+    linkage: np.ndarray
+    leaf_order: list[int]
+    sequences: list[bytes] = Field(default_factory=list)
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @field_validator("linkage", mode="before")
+    @classmethod
+    def _to_ndarray(cls, v):
+        if isinstance(v, np.ndarray):
+            return v
+        return np.asarray(v, dtype=float)

{openprotein_python-0.13.1 → openprotein_python-0.14.0}/openprotein/data/data.py

@@ -14,9 +14,11 @@ class DataAPI:
     def __init__(self, session: APISession):
         self.session = session
 
+    # `list` shadows the builtin, so the return annotation is quoted to keep it
+    # resolving to the builtin (not this method) under PEP 649 deferred evaluation.
     def list(
         self, limit: int | None = None, offset: int | None = None
-    ) -> list[AssayDataset]:
+    ) -> "list[AssayDataset]":
         """
         List all assay datasets.
 

{openprotein_python-0.13.1 → openprotein_python-0.14.0}/openprotein/embeddings/ablang.py

@@ -11,7 +11,7 @@ class AbLang2Model(EmbeddingModel):
     --------
     View specific model details (inc supported tokens) with the `?` operator.
 
-    .. code-block:: python
+    .. code-block:: ipython3
 
         >>> import openprotein
         >>> session = openprotein.connect(username="user", password="password")