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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.12.1
+Version: 0.13.1
 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

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/align/align.py

@@ -189,6 +189,7 @@ class AlignAPI:
         sequences: Sequence[bytes | str],
         names: Sequence[str] | None = None,
         scheme: AbNumberScheme = AbNumberScheme.CHOTHIA,
+        drop_minority_chains: bool = False,
     ) -> MSAFuture:
         """
         Align antibody sequences using `AbNumber`.
@@ -205,6 +206,10 @@ class AlignAPI:
             Optional list of sequence names, must be the same length as sequences if provided.
         scheme : AbNumberScheme, default=AbNumberScheme.CHOTHIA
             Antibody numbering scheme.
+        drop_minority_chains : bool, default=False
+            If True, drop sequences belonging to chain types that are in the
+            minority (e.g. heavy vs light) so the resulting alignment contains
+            only the dominant chain type.
 
         Returns
         -------
@@ -235,10 +240,15 @@ class AlignAPI:
                 lines.append(sequence)
         content = b"\n".join(lines)
         stream = BytesIO(content)
-        return self.abnumber_file(stream, scheme=scheme)
+        return self.abnumber_file(
+            stream, scheme=scheme, drop_minority_chains=drop_minority_chains
+        )
 
     def abnumber_file(
-        self, file, scheme: AbNumberScheme = AbNumberScheme.CHOTHIA
+        self,
+        file,
+        scheme: AbNumberScheme = AbNumberScheme.CHOTHIA,
+        drop_minority_chains: bool = False,
     ) -> MSAFuture:
         """
         Align antibody sequences using `AbNumber`.
@@ -253,13 +263,22 @@ class AlignAPI:
             Sequences to align in FASTA or CSV format.
         scheme : AbNumberScheme, default=AbNumberScheme.CHOTHIA
             Antibody numbering scheme.
+        drop_minority_chains : bool, default=False
+            If True, drop sequences belonging to chain types that are in the
+            minority (e.g. heavy vs light) so the resulting alignment contains
+            only the dominant chain type.
 
         Returns
         -------
         MSAFuture
             Future object awaiting the contents of the MSA upload.
         """
-        job = api.abnumber_post(self.session, file, scheme=scheme)
+        job = api.abnumber_post(
+            self.session,
+            file,
+            scheme=scheme,
+            drop_minority_chains=drop_minority_chains,
+        )
         return MSAFuture.create(session=self.session, job=job)
 
     def upload_msa(self, msa_file: BinaryIO) -> MSAFuture:

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/align/api.py

@@ -269,6 +269,7 @@ def abnumber_post(
     session: APISession,
     sequence_file: BinaryIO,
     scheme: AbNumberScheme | str = AbNumberScheme.IMGT,
+    drop_minority_chains: bool = False,
 ) -> Job:
     """
     Align antibody sequences using AbNumber.
@@ -284,6 +285,10 @@ def abnumber_post(
         Sequences to align in FASTA or CSV format.
     scheme : AbNumberScheme, optional
         Antibody numbering scheme. Default is IMGT.
+    drop_minority_chains : bool, optional
+        If True, drop sequences belonging to chain types that are in the
+        minority (e.g. heavy vs light) so the resulting alignment contains
+        only the dominant chain type. Default is False.
 
     Returns
     -------
@@ -297,7 +302,10 @@ def abnumber_post(
             raise InvalidParameterError(f"Antibody numbering {scheme} not recognized")
 
     files = {"file": sequence_file}
-    params = {"scheme": scheme if isinstance(scheme, str) else scheme.value}
+    params = {
+        "scheme": scheme if isinstance(scheme, str) else scheme.value,
+        "drop_minority_chains": drop_minority_chains,
+    }
 
     response = session.post(endpoint, files=files, params=params)
     return Job.model_validate(response.json())

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/align/msa.py

@@ -45,7 +45,7 @@ class MSAFuture(AlignFuture, Future):
         self.page_size = page_size
         self.msa_id = self.job.job_id
 
-    def get(self, verbose: bool = False) -> Iterator[tuple[str, str]]:
+    def _get(self, verbose: bool = False) -> Iterator[tuple[str, str]]:
         """
         Retrieve the MSA of the job.
 

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/align/schemas.py

@@ -157,7 +157,24 @@ class AbNumberJob(MSAJob, Job):
 
 
 class AbNumberScheme(str, Enum):
-    """Antibody numbering scheme."""
+    """
+    Antibody numbering scheme.
+
+    Attributes
+    ----------
+    IMGT : str
+        IMGT numbering scheme, a standardized system unifying numbering across
+        all immunoglobulin and T-cell receptor variable domains.
+    CHOTHIA : str
+        Chothia numbering scheme, based on the structural location of
+        canonical loops in antibody variable domains.
+    KABAT : str
+        Kabat numbering scheme, based on sequence variability across
+        antibody variable domains.
+    AHO : str
+        AHo numbering scheme, a structure-based unified numbering scheme
+        for antibody variable domains.
+    """
 
     IMGT = "imgt"
     CHOTHIA = "chothia"

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/base.py

@@ -8,8 +8,8 @@ import requests.auth
 from requests.adapters import HTTPAdapter
 from requests.packages.urllib3.util.retry import Retry  # type: ignore
 
-import openprotein.config as config
-from openprotein.errors import APIError, AuthError, HTTPError
+from openprotein._version import __version__
+from openprotein.errors import AuthError, HTTPError
 
 
 class BearerAuth(requests.auth.AuthBase):
@@ -17,11 +17,19 @@ class BearerAuth(requests.auth.AuthBase):
     See https://stackoverflow.com/a/58055668
     """
 
-    def __init__(self, token):
-        self.token = token
+    def __init__(self, username, token):
+        self._username = username
+        self._token = token
+
+    @property
+    def token(self):
+        warnings.warn(
+            "DeprecationWarning: Accessing session.auth.token is deprecated and will throw an error in the future. Use session.auth._token if you need it."
+        )
+        return self._token
 
     def __call__(self, r):
-        r.headers["Authorization"] = "Bearer " + self.token
+        r.headers["Authorization"] = "Bearer " + self._token
         return r
 
 
@@ -56,6 +64,7 @@ class APISession(requests.Session):
         )
         adapter = HTTPAdapter(max_retries=retry)
         self.mount("https://", adapter)
+        self.headers["User-Agent"] = f"openprotein-python/{__version__}"
         self.login(username, password)
 
     def post(self, url, data=None, json=None, **kwargs):
@@ -101,14 +110,14 @@ class APISession(requests.Session):
         except HTTPError as e:
             # if an error occured during auth, we raise an AuthError with reference to the HTTPError
             raise AuthError(
-                f"Authentication failed. Please check your credentials and connection."
+                "Authentication failed. Please check your credentials and connection."
             ) from e
 
         result = response.json()
         token = result.get("access_token")
         if token is None:
             raise AuthError("Unable to authenticate with given credentials.")
-        return BearerAuth(token)
+        return BearerAuth(username=username, token=token)
 
     def request(self, method: str, url: str, *args, **kwargs):
         full_url = urljoin(self.backend, url)
@@ -136,7 +145,7 @@ class APISession(requests.Session):
         return response
 
 
-def _total_size(o: Sequence | Mapping, seen=None):
+def _total_size(o: Sequence | Mapping | object, seen=None):
     """Recursively finds size of objects including contents."""
     if seen is None:
         seen = set()

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

@@ -59,6 +59,7 @@ class DataAPI:
         metadata = api.assaydata_post(
             self.session, stream, name, assay_description=description
         )
+        table.columns = [str(column).lower() for column in table.columns]
         metadata.sequence_length = len(table["sequence"].values[0])
         return AssayDataset(self.session, metadata)
 

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/data/schemas.py

@@ -13,6 +13,7 @@ class AssayMetadata(BaseModel):
     num_entries: int
     measurement_names: list[str]
     sequence_length: int | None = None
+    has_name_column: bool = False
 
 
 class AssayDataRow(BaseModel):

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/design/schemas.py

@@ -67,7 +67,14 @@ class Subcriterion(BaseModel):
             raise ValueError(
                 f"Expected to chain only with criterion or subcriterion, got {type(other)}"
             )
-        return Criterion([self] + others)  # type: ignore - doesnt like Self
+        all_subcriteria = [self] + others
+        types = {sc.criterion_type for sc in all_subcriteria}
+        if len(types) > 1:
+            raise ValueError(
+                "Cannot AND model and n_mutations criteria in the same group. "
+                "Use | (OR) to put them in separate groups instead."
+            )
+        return Criterion(all_subcriteria)  # type: ignore - doesnt like Self
 
     def __or__(self, other: "Subcriterion | Criterion | Any") -> "Criteria":
         """
@@ -332,7 +339,14 @@ class Criterion(RootModel):
         elif isinstance(other, Criterion):
             others = other.root
 
-        return Criterion(self.root + others)
+        all_subcriteria = self.root + others
+        types = {sc.criterion_type for sc in all_subcriteria}
+        if len(types) > 1:
+            raise ValueError(
+                "Cannot AND model and n_mutations criteria in the same group. "
+                "Use | (OR) to put them in separate groups instead."
+            )
+        return Criterion(all_subcriteria)
 
     def __or__(self, other: "Criterion | Subcriterion") -> "Criteria":
         """

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/embeddings/api.py

@@ -28,20 +28,28 @@ from .schemas import (
 PATH_PREFIX = "v1/embeddings"
 
 
-def list_models(session: APISession) -> list[str]:
+def list_models(
+    session: APISession, verbose: bool = False
+) -> list[str] | list[ModelMetadata]:
     """
     List available embeddings models.
 
     Args:
         session (APISession): API session
+        verbose (bool): If True, return full ModelMetadata objects instead of model ID strings.
 
     Returns:
-        list[str]: list of model names.
+        list[str] | list[ModelMetadata]: list of model names or metadata objects.
     """
 
     endpoint = PATH_PREFIX + "/models"
-    response = session.get(endpoint)
+    params = {}
+    if verbose:
+        params["verbose"] = "true"
+    response = session.get(endpoint, params=params)
     result = response.json()
+    if verbose:
+        return TypeAdapter(list[ModelMetadata]).validate_python(result)
     return result
 
 

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/embeddings/embeddings.py

@@ -3,6 +3,7 @@
 from openprotein.base import APISession
 
 from . import api
+from .ablang import AbLang2Model
 from .esm import ESMModel
 from .future import EmbeddingsResultFuture
 from .models import EmbeddingModel
@@ -55,6 +56,7 @@ class EmbeddingsAPI:
     #: Rotaprot model trained on UniRef90
     rotaprot_large_uniref90_ft: OpenProteinModel
     poet_2: PoET2Model
+    ablang2: AbLang2Model
 
     #: ESM1b model
     esm1b: ESMModel  # alias
@@ -100,10 +102,13 @@ class EmbeddingsAPI:
     def list_models(self) -> list[EmbeddingModel]:
         """list models available for creating embeddings of your sequences"""
         models = []
-        for model_id in api.list_models(self.session):
+        for metadata in api.list_models(self.session, verbose=True):
             models.append(
                 EmbeddingModel.create(
-                    session=self.session, model_id=model_id, default=EmbeddingModel
+                    session=self.session,
+                    model_id=metadata.id,
+                    default=EmbeddingModel,
+                    metadata=metadata,
                 )
             )
         return models

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/embeddings/models.py

@@ -74,6 +74,7 @@ class EmbeddingModel:
         session: APISession,
         model_id: str,
         default: type["EmbeddingModel"] | None = None,
+        metadata: ModelMetadata | None = None,
         **kwargs,
     ):
         """
@@ -87,6 +88,8 @@ class EmbeddingModel:
             The model identifier.
         default : type variable of EmbeddingModel or None, optional
             Default EmbeddingModel subclass to use if no match is found.
+        metadata : ModelMetadata or None, optional
+            Pre-fetched metadata to avoid an extra API call.
         kwargs :
             Additional keyword arguments to pass to the model constructor.
 
@@ -106,11 +109,11 @@ class EmbeddingModel:
         # Find the EmbeddingModel class that matches the model_id
         for model_class in model_classes:
             if model_id in model_class.get_model():
-                return model_class(session=session, model_id=model_id, **kwargs)
+                return model_class(session=session, model_id=model_id, metadata=metadata, **kwargs)
         # default to ProtembedModel
         if default is not None:
             try:
-                return default(session=session, model_id=model_id, **kwargs)
+                return default(session=session, model_id=model_id, metadata=metadata, **kwargs)
             except:
                 # continue to throw error as unsupported
                 pass

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/errors.py

@@ -72,6 +72,23 @@ class TimeoutException(Exception):
         super().__init__(self.message)
 
 
+class JobFailedException(Exception):
+    """Raised when retrieving results from a job that has failed."""
+
+    def __init__(
+        self,
+        job_id: str,
+        failure_message: str | None = None,
+    ):
+        self.job_id = job_id
+        self.failure_message = failure_message
+        message = f"Job {job_id} failed"
+        if failure_message:
+            message = f"{message}: {failure_message}"
+        self.message = message
+        super().__init__(self.message)
+
+
 class DeprecationError(Exception):
     """DeprecationError used for flagging to the user to not use this interface anymore."""
 

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/fold/api.py

@@ -140,7 +140,7 @@ def fold_get_extra_result(
     session: APISession,
     job_id: str,
     sequence_or_index: bytes | str | int,
-    key: Literal["pae", "pde", "plddt", "ptm"],
+    key: Literal["pae", "pde", "plddt", "ptm", "ipae"],
 ) -> np.ndarray: ...
 
 
@@ -176,7 +176,15 @@ def fold_get_extra_result(
     job_id: str,
     sequence_or_index: bytes | str | int,
     key: Literal[
-        "pae", "pde", "plddt", "ptm", "confidence", "affinity", "score", "metrics"
+        "pae",
+        "pde",
+        "plddt",
+        "ptm",
+        "ipae",
+        "confidence",
+        "affinity",
+        "score",
+        "metrics",
     ],
 ) -> "np.ndarray | list[dict] | dict | pd.DataFrame":
     """
@@ -190,15 +198,16 @@ def fold_get_extra_result(
         Job ID to retrieve results from.
     sequence_or_index : bytes or str or int
         Sequence to retrieve results for or its index in the job.
-    key : {'pae', 'pde', 'plddt', 'ptm', 'confidence', 'affinity', 'score', 'metrics'}
-        The type of result to retrieve.
+    key : {'pae', 'pde', 'plddt', 'ptm', 'ipae', 'confidence', 'affinity', 'score', 'metrics'}
+        The type of result to retrieve. ``ipae`` is a synthetic scalar per
+        unit (shape ``(1,)``) derived from ``pae`` and the chain layout.
 
     Returns
     -------
     numpy.ndarray or list of dict
-        The result as a numpy array (for "pae", "pde", "plddt") or a list of dictionaries (for "confidence", "affinity").
+        The result as a numpy array (for "pae", "pde", "plddt", "ptm", "ipae") or a list of dictionaries (for "confidence", "affinity").
     """
-    if key in {"pae", "pde", "plddt", "ptm"}:
+    if key in {"pae", "pde", "plddt", "ptm", "ipae"}:
 
         def formatter(response):
             return np.load(io.BytesIO(response.content))
@@ -226,6 +235,82 @@ def fold_get_extra_result(
     return output
 
 
+@typing.overload
+def fold_get_batch_extra_result(
+    session: APISession,
+    job_id: str,
+    key: Literal["pae", "pde", "plddt", "ptm", "ipae"],
+) -> np.ndarray: ...
+
+
+@typing.overload
+def fold_get_batch_extra_result(
+    session: APISession,
+    job_id: str,
+    key: Literal["confidence"],
+) -> list[list[dict] | None]: ...
+
+
+@typing.overload
+def fold_get_batch_extra_result(
+    session: APISession,
+    job_id: str,
+    key: Literal["affinity"],
+) -> list[dict | None]: ...
+
+
+def fold_get_batch_extra_result(
+    session: APISession,
+    job_id: str,
+    key: Literal["pae", "pde", "plddt", "ptm", "ipae", "confidence", "affinity"],
+) -> "np.ndarray | list":
+    """
+    Fetch an extra result key stacked across every unit in a fold job.
+
+    Backed by a single server call to ``/v1/fold/{job_id}/results/{key}``
+    (not N per-unit calls). For ``.npy`` keys the server pads per-unit
+    arrays with NaN to the per-axis max shape and stacks along a new
+    leading dim; for ``.json`` keys it returns a length-``N`` array where
+    element ``i`` is unit ``i``'s parsed JSON (``None`` if that unit's
+    result was missing or failed to fetch).
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    job_id : str
+        Job ID to retrieve results from.
+    key : {'pae', 'pde', 'plddt', 'ptm', 'ipae', 'confidence', 'affinity'}
+        The type of result to retrieve. ``ipae`` returns shape ``[N]``.
+
+    Returns
+    -------
+    numpy.ndarray or list
+        ``np.ndarray`` of shape ``[N, ...]`` for npy keys; ``list`` of
+        length ``N`` for json keys (``confidence`` / ``affinity``), with
+        ``None`` entries for units whose per-unit result could not be
+        fetched.
+    """
+    if key in {"pae", "pde", "plddt", "ptm", "ipae"}:
+
+        def formatter(response):
+            return np.load(io.BytesIO(response.content))
+    elif key in {"confidence", "affinity"}:
+
+        def formatter(response):
+            return response.json()
+    else:
+        raise ValueError(f"Unexpected key: {key}")
+    endpoint = PATH_PREFIX + f"/{job_id}/results/{key}"
+    try:
+        response = session.get(endpoint)
+    except HTTPError as e:
+        if e.status_code == 400 and key == "affinity":
+            raise ValueError("affinity not found for request") from None
+        raise e
+    return formatter(response)
+
+
 def fold_models_post(
     session: APISession,
     model_id: str,

{openprotein_python-0.12.1 → openprotein_python-0.13.1}/openprotein/fold/common.py

@@ -139,13 +139,19 @@ def resolve_templates(session: APISession, templates: Sequence[Template]) -> lis
                 query=template.template
             )
 
-        template_dict = {"query_id": struct_id_to_query_id[struct_id]}
+        template_dict: dict = {"query_id": struct_id_to_query_id[struct_id]}
         if template.mapping is not None:
             if isinstance(template.mapping, str):
                 template_dict["chain_id"] = template.mapping
             else:
                 template_dict["chain_id"] = list(template.mapping.values())
                 template_dict["template_id"] = list(template.mapping.keys())
+        if template.index is not None:
+            template_dict["index"] = list(template.index)
+        if template.index_intervals is not None:
+            template_dict["index_intervals"] = [
+                [int(start), int(end)] for start, end in template.index_intervals
+            ]
         template_dicts.append(template_dict)
 
     return template_dicts
@@ -199,19 +205,17 @@ def serialize_input(session: APISession, complexes: list[Complex], needs_msa: bo
                     p["msa_id"] = msa_id
                 _complex.append({"protein": p})
             elif isinstance(chain, Ligand):
-                ligand_payload = {
-                    "id": chain_id,
-                }
+                ligand_payload: dict[str, Any] = {"id": chain_id}
                 if chain.smiles is not None:
                     ligand_payload["smiles"] = chain.smiles
                 if chain.ccd is not None:
                     ligand_payload["ccd"] = chain.ccd
                 _complex.append({"ligand": ligand_payload})
             elif isinstance(chain, DNA):
-                d = {"id": chain_id, "sequence": chain.sequence}
+                d: dict[str, Any] = {"id": chain_id, "sequence": chain.sequence}
                 _complex.append({"dna": d})
             elif isinstance(chain, RNA):
-                r = {"id": chain_id, "sequence": chain.sequence}
+                r: dict[str, Any] = {"id": chain_id, "sequence": chain.sequence}
                 _complex.append({"rna": r})
             else:
                 raise ValueError(f"Unexpected chain type: {chain}")