Release v0.8.5

- Fixed boltz affinity
- Fixed predictor train with poet prompt
- Add paging to listing assays and jobs
- bfactor for plddt
- Some typing improvements

Author: markgeejw

anaconda_build/meta.yaml

@@ -1,6 +1,6 @@
 package:
   name: openprotein-python
-  version: "0.8.4"
+  version: "0.8.5"
 
 source:
   path: ../

openprotein/common/__init__.py

@@ -1,5 +1,5 @@
 """Common classes and utilities for OpenProtein."""
 
-from .features import FeatureType
+from .features import Feature, FeatureType
 from .model_metadata import ModelDescription, ModelMetadata, TokenInfo
-from .reduction import ReductionType
+from .reduction import Reduction, ReductionType

openprotein/common/features.py

@@ -1,7 +1,15 @@
+"""Feature types used in OpenProtein."""
+
 from enum import Enum
+from typing import Literal
 
 
 class FeatureType(str, Enum):
 
     PLM = "PLM"
     SVD = "SVD"
+
+
+# NOTE: only works with python 3.12+
+# Feature = Literal[*tuple([r.value for r in FeatureType])]
+Feature = Literal["PLM", "SVD"]

openprotein/common/reduction.py

@@ -1,8 +1,14 @@
 """Reduction types used in OpenProtein."""
 
 from enum import Enum
+from typing import Literal
 
 
 class ReductionType(str, Enum):
     MEAN = "MEAN"
     SUM = "SUM"
+
+
+# NOTE: only works with python 3.12+
+# Reduction = Literal[*tuple([r.value for r in ReductionType])]
+Reduction = Literal["MEAN", "SUM"]

openprotein/data/api.py

@@ -64,14 +64,20 @@ def assaydata_post(
         raise APIError(f"Unable to post assay data: {response.text}")
 
 
-def assaydata_list(session: APISession) -> list[AssayMetadata]:
+def assaydata_list(
+    session: APISession, limit: int | None = None, offset: int | None = None
+) -> list[AssayMetadata]:
     """
     Get a list of all assay metadata.
 
     Parameters
     ----------
     session : APISession
         Session object for API communication.
+    limit : int, optional
+        Limit the number of assays to return.
+    offset : int, optional
+        Offset of assays to retrieve. Useful with limit.
 
     Returns
     -------
@@ -84,7 +90,12 @@ def assaydata_list(session: APISession) -> list[AssayMetadata]:
         If an error occurs during the API request.
     """
     endpoint = "v1/assaydata"
-    response = session.get(endpoint)
+    params = {}
+    if limit is not None:
+        params["limit"] = limit
+    if offset is not None:
+        params["offset"] = offset
+    response = session.get(endpoint, params=params)
     if response.status_code == 200:
         return TypeAdapter(list[AssayMetadata]).validate_python(response.json())
     else:

openprotein/data/data.py

@@ -14,16 +14,23 @@ class DataAPI:
     def __init__(self, session: APISession):
         self.session = session
 
-    def list(self) -> list[AssayDataset]:
+    def list(
+        self, limit: int | None = None, offset: int | None = None
+    ) -> list[AssayDataset]:
         """
         List all assay datasets.
 
+        limit : int, optional
+            Limit the number of assays to return.
+        offset : int, optional
+            Offset of assays to retrieve. Useful with limit.
+
         Returns
         -------
         List[AssayDataset]
             List of all assay datasets.
         """
-        metadata = api.assaydata_list(self.session)
+        metadata = api.assaydata_list(session=self.session, limit=limit, offset=offset)
         return [AssayDataset(self.session, x) for x in metadata]
 
     def create(

openprotein/embeddings/models.py

@@ -3,7 +3,13 @@
 from typing import TYPE_CHECKING
 
 from openprotein.base import APISession
-from openprotein.common import FeatureType, ModelMetadata, ReductionType
+from openprotein.common import (
+    Feature,
+    FeatureType,
+    ModelMetadata,
+    Reduction,
+    ReductionType,
+)
 from openprotein.data import AssayDataset, AssayMetadata, DataAPI
 from openprotein.errors import InvalidParameterError
 
@@ -199,9 +205,9 @@ def logits(
     def fit_svd(
         self,
         sequences: list[bytes] | list[str] | None = None,
-        assay: AssayDataset | None = None,
+        assay: AssayDataset | AssayMetadata | None = None,
         n_components: int = 1024,
-        reduction: ReductionType | None = None,
+        reduction: Reduction | ReductionType | None = None,
         **kwargs,
     ) -> "SVDModel":
         """
@@ -236,6 +242,11 @@ def fit_svd(
         # local import for cyclic dep
         from openprotein.svd import SVDAPI
 
+        # runtime check on value
+        if isinstance(reduction, str):
+            reduction = ReductionType(reduction)
+            reduction = reduction.value
+
         svd_api = getattr(self.session, "svd", None)
         assert isinstance(svd_api, SVDAPI)
 
@@ -246,9 +257,8 @@ def fit_svd(
             raise InvalidParameterError(
                 "Expected either assay or sequences to fit SVD on!"
             )
-        model_id = self.id
         return svd_api.fit_svd(
-            model_id=model_id,
+            model=self,
             sequences=sequences,
             assay=assay,
             n_components=n_components,
@@ -259,9 +269,9 @@ def fit_svd(
     def fit_umap(
         self,
         sequences: list[bytes] | list[str] | None = None,
-        assay: AssayDataset | None = None,
+        assay: AssayDataset | AssayMetadata | None = None,
         n_components: int = 2,
-        reduction: ReductionType | None = ReductionType.MEAN,
+        reduction: Reduction | ReductionType = "MEAN",
         **kwargs,
     ) -> "UMAPModel":
         """
@@ -274,11 +284,11 @@ def fit_umap(
         ----------
         sequences : list of bytes or list of str or None, optional
             Optional sequences to fit UMAP with. Either use sequences or assay. Sequences is preferred.
-        assay : AssayDataset or None, optional
+        assay : AssayDataset or AssayMetadata or None, optional
             Optional assay containing sequences to fit UMAP with. Either use sequences or assay. Ignored if sequences are provided.
         n_components : int, optional
             Number of components in UMAP fit. Determines output shapes. Default is 2.
-        reduction : ReductionType or None, optional
+        reduction : Reduction or ReductionType or None, optional
             Embeddings reduction to use (e.g. mean). Defaults to MEAN.
         kwargs :
             Additional keyword arguments to be used from foundational models, e.g. prompt_id for PoET models.
@@ -296,6 +306,16 @@ def fit_umap(
         # local import for cyclic dep
         from openprotein.umap import UMAPAPI
 
+        if reduction is None:
+            raise InvalidParameterError(
+                "Expected reduction if using EmbeddingModel to fit UMAP"
+            )
+
+        # runtime check on value
+        if isinstance(reduction, str):
+            reduction = ReductionType(reduction)
+            reduction = reduction.value
+
         umap_api = getattr(self.session, "umap", None)
         assert isinstance(umap_api, UMAPAPI)
 
@@ -306,20 +326,26 @@ def fit_umap(
             raise InvalidParameterError(
                 "Expected either assay or sequences to fit UMAP on!"
             )
+        # get assay_id
+        assay_id = (
+            assay.assay_id
+            if isinstance(assay, AssayMetadata)
+            else assay.id if isinstance(assay, AssayDataset) else assay
+        )
         model_id = self.id
         return umap_api.fit_umap(
             model_id=model_id,
             feature_type=FeatureType.PLM,
             sequences=sequences,
-            assay_id=assay.id if assay is not None else None,
+            assay_id=assay_id,
             n_components=n_components,
             reduction=reduction,
             **kwargs,
         )
 
     def fit_gp(
         self,
-        assay: AssayMetadata | AssayDataset | str,
+        assay: AssayDataset | AssayMetadata | str,
         properties: list[str],
         reduction: ReductionType,
         name: str | None = None,
@@ -358,26 +384,9 @@ def fit_gp(
         # local import to resolve cyclic
         from openprotein.predictor import PredictorAPI
 
-        data_api = getattr(self.session, "data", None)
-        assert isinstance(data_api, DataAPI)
         predictor_api = getattr(self.session, "predictor", None)
         assert isinstance(predictor_api, PredictorAPI)
 
-        # get assay if str
-        assay = data_api.get(assay_id=assay) if isinstance(assay, str) else assay
-        # extract assay_id
-        if len(properties) == 0:
-            raise InvalidParameterError("Expected (at-least) 1 property to train")
-        if not set(properties) <= set(assay.measurement_names):
-            raise InvalidParameterError(
-                f"Expected all provided properties to be a subset of assay's measurements: {assay.measurement_names}"
-            )
-        # TODO - support multitask
-        if len(properties) > 1:
-            raise InvalidParameterError(
-                "Training a multitask GP is not yet supported (i.e. number of properties should only be 1 for now)"
-            )
-
         # inject into predictor api
         return predictor_api.fit_gp(
             assay=assay,

openprotein/fold/future.py

@@ -464,6 +464,8 @@ def affinity(self) -> "BoltzAffinity":
         AttributeError
             If affinity is not supported for the model.
         """
+        from .boltz import BoltzAffinity
+
         if self.model_id not in {"boltz-1", "boltz-1x", "boltz-2"}:
             raise AttributeError("affinity not supported for non-Boltz model")
         if self._affinity is None:

openprotein/predictor/api.py

@@ -162,8 +162,8 @@ def predictor_fit_gp_post(
         body["name"] = name
     if description is not None:
         body["description"] = description
-    # add kwargs for embeddings kwargs
-    body.update(kwargs)
+    # add kwargs for embeddings kwargs to features
+    body["features"].update(kwargs)
 
     response = session.post(endpoint, json=body)
     return PredictorTrainJob.model_validate(response.json())

openprotein/predictor/predictor.py

@@ -1,10 +1,11 @@
 """Predictor API providing the interface to train and predict predictors."""
 
 from openprotein.base import APISession
-from openprotein.common import FeatureType, ReductionType
+from openprotein.common import Feature, FeatureType, Reduction, ReductionType
 from openprotein.data import (
     AssayDataset,
     AssayMetadata,
+    DataAPI,
 )
 from openprotein.embeddings import EmbeddingModel, EmbeddingsAPI
 from openprotein.errors import InvalidParameterError
@@ -120,8 +121,8 @@ def fit_gp(
         assay: AssayDataset | AssayMetadata | str,
         properties: list[str],
         model: EmbeddingModel | SVDModel | str,
-        feature_type: FeatureType | None = None,
-        reduction: ReductionType | None = None,
+        feature_type: Feature | FeatureType | None = None,
+        reduction: Reduction | ReductionType | None = None,
         name: str | None = None,
         description: str | None = None,
         **kwargs,
@@ -139,10 +140,10 @@ def fit_gp(
             Instance of either EmbeddingModel or SVDModel to use depending
             on feature type. Can also be a str specifying the model id,
             but then feature_type would have to be specified.
-        feature_type : FeatureType or None
+        feature_type : Feature or FeatureType or None
             Type of features to use for encoding sequences. "SVD" or "PLM".
             None would require model to be EmbeddingModel or SVDModel.
-        reduction  : str or None, optional
+        reduction  : Reduction or ReductionType or None, optional
             Type of embedding reduction to use for computing features.
             E.g. "MEAN" or "SUM". Used only if using EmbeddingModel, and
             must be non-nil if using an EmbeddingModel. Defaults to None.
@@ -154,6 +155,29 @@ def fit_gp(
         PredictorModel
             The GP model being fit.
         """
+        data_api = getattr(self.session, "data", None)
+        assert isinstance(data_api, DataAPI)
+        # 1. Check assay data input
+        # get assay if str
+        assay = data_api.get(assay_id=assay) if isinstance(assay, str) else assay
+        # extract assay_id
+        assay_id = (
+            assay.assay_id
+            if isinstance(assay, AssayMetadata)
+            else assay.id if isinstance(assay, AssayDataset) else assay
+        )
+        if len(properties) == 0:
+            raise InvalidParameterError("Expected (at-least) 1 property to train")
+        if not set(properties) <= set(assay.measurement_names):
+            raise InvalidParameterError(
+                f"Expected all provided properties to be a subset of assay's measurements: {assay.measurement_names}"
+            )
+        # TODO - support multitask
+        if len(properties) > 1:
+            raise InvalidParameterError(
+                "Training a multitask GP is not yet supported (i.e. number of properties should only be 1 for now)"
+            )
+        # 2. Check features input
         # extract feature type
         feature_type = (
             FeatureType.PLM
@@ -164,6 +188,15 @@ def fit_gp(
             raise InvalidParameterError(
                 "Expected feature_type to be provided if passing str model_id as model"
             )
+        # runtime check on value
+        if isinstance(feature_type, str):
+            feature_type = FeatureType(feature_type)
+
+        # 3. Check reduction
+        if isinstance(reduction, str):
+            reduction = ReductionType(reduction)
+            reduction = reduction.value
+
         # get model if model_id
         if feature_type == FeatureType.PLM:
             if reduction is None:
@@ -183,19 +216,14 @@ def fit_gp(
                 model = svd_api.get_svd(model)
             assert isinstance(model, SVDModel), "Expected SVDModel"
             model_id = model.id
-        # get assay_id
-        assay_id = (
-            assay.assay_id
-            if isinstance(assay, AssayMetadata)
-            else assay.id if isinstance(assay, AssayDataset) else assay
-        )
+
         return PredictorModel(
             session=self.session,
             job=api.predictor_fit_gp_post(
                 session=self.session,
                 assay_id=assay_id,
                 properties=properties,
-                feature_type=feature_type,
+                feature_type=feature_type.value,
                 model_id=model_id,
                 reduction=reduction,
                 name=name,