openprotein-python 0.8.12__tar.gz → 0.10.0__tar.gz

{openprotein_python-0.8.12 → openprotein_python-0.10.0}/.gitignore

@@ -1,3 +1,4 @@
+.vscode
 
 # SECRETS FILE
 secrets.config
@@ -363,3 +364,6 @@ __marimo__/
 # Streamlit
 .streamlit/secrets.toml
 /data/
+
+.envrc
+/.direnv/

{openprotein_python-0.8.12 → openprotein_python-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.8.12
+Version: 0.10.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
@@ -8,11 +8,12 @@ License-File: LICENSE.txt
 Classifier: Development Status :: 4 - Beta
 Classifier: Programming Language :: Python :: 3
 Requires-Python: <3.13,>=3.10
-Requires-Dist: gemmi<0.8,>=0.7.0
+Requires-Dist: gemmi<0.8,>=0.6.0
 Requires-Dist: numpy<3,>=1.9
 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
 Description-Content-Type: text/markdown
 

{openprotein_python-0.8.12 → openprotein_python-0.10.0}/openprotein/__init__.py

@@ -6,8 +6,8 @@ A pythonic interface for interacting with our cutting-edge protein engineering p
 isort:skip_file
 """
 
-from typing import TYPE_CHECKING
-import warnings
+import os
+from pathlib import Path
 
 from openprotein._version import __version__
 from openprotein.data import DataAPI
@@ -161,4 +161,48 @@ class OpenProtein(APISession):
         return self._models
 
 
-connect = OpenProtein
+def connect(
+    username: str | None = None,
+    password: str | None = None,
+    backend: str | None = None,
+    timeout: int = 180,
+) -> OpenProtein:
+    """
+    Connect and create a :py:class:`OpenProtein` session.
+
+    Parameters
+    ----------
+    username : str, optional
+        The username of the user. If not provided, taken from the
+        environment variable ``OPENPROTEIN_USERNAME`` or a
+        configuration file at ``~/.openprotein/config.toml``.
+    password : str, optional
+        The password of the user. If not provided, taken from the
+        environment variable ``OPENPROTEIN_PASSWORD`` or a
+        configuration file at ``~/.openprotein/config.toml``.
+
+    Examples
+    --------
+    >>> session = openprotein.connect("username", "password")
+    """
+
+    CREDENTIALS_FILE_PATH = Path.home() / ".openprotein/config.toml"
+    if CREDENTIALS_FILE_PATH.exists():
+        import tomli
+
+        with open(CREDENTIALS_FILE_PATH, "rb") as f:
+            file_config = tomli.load(f)
+    else:
+        file_config = {}
+    USERNAME = os.getenv("OPENPROTEIN_USERNAME", str(file_config.get("username")))
+    PASSWORD = os.getenv("OPENPROTEIN_PASSWORD", str(file_config.get("password")))
+    BACKEND = os.getenv(
+        "OPENPROTEIN_API_BACKEND",
+        file_config.get("backend", "https://api.openprotein.ai/api/"),
+    )
+    return OpenProtein(
+        username=username or USERNAME,
+        password=password or PASSWORD,
+        backend=backend or BACKEND,
+        timeout=timeout,
+    )

{openprotein_python-0.8.12 → openprotein_python-0.10.0}/openprotein/align/align.py

@@ -7,7 +7,7 @@ from typing import BinaryIO, Iterator
 from openprotein.base import APISession
 from openprotein.errors import DeprecationError
 from openprotein.jobs import Job
-from openprotein.protein import Protein
+from openprotein.molecules import Protein
 
 from . import api
 from .msa import MSAFuture
@@ -285,7 +285,7 @@ class AlignAPI:
             session=self.session, job=api.msa_post(self.session, msa_file=msa_file)
         )
 
-    def create_msa(self, seed: bytes) -> MSAFuture:
+    def create_msa(self, seed: bytes | str) -> MSAFuture:
         """
         Construct an MSA via homology search with the seed sequence.
 

{openprotein_python-0.8.12 → openprotein_python-0.10.0}/openprotein/base.py

@@ -1,8 +1,6 @@
-import os
 import sys
 import warnings
-from collections.abc import Container, Mapping
-from typing import Union
+from typing import Mapping, Sequence
 from urllib.parse import urljoin
 
 import requests
@@ -13,10 +11,6 @@ from requests.packages.urllib3.util.retry import Retry  # type: ignore
 import openprotein.config as config
 from openprotein.errors import APIError, AuthError, HTTPError
 
-USERNAME = os.getenv("OPENPROTEIN_USERNAME")
-PASSWORD = os.getenv("OPENPROTEIN_PASSWORD")
-BACKEND = os.getenv("OPENPROTEIN_API_BACKEND", "https://api.openprotein.ai/api/")
-
 
 class BearerAuth(requests.auth.AuthBase):
     """
@@ -34,29 +28,18 @@ class BearerAuth(requests.auth.AuthBase):
 class APISession(requests.Session):
     """
     A class to handle API sessions. This class provides a connection session to the OpenProtein API.
-
-    Parameters
-    ----------
-    username : str
-        The username of the user.
-    password : str
-        The password of the user.
-
-    Examples
-    --------
-    >>> session = APISession("username", "password")
     """
 
     def __init__(
         self,
-        username: str | None = USERNAME,
-        password: str | None = PASSWORD,
-        backend: str = BACKEND,
+        username: str,
+        password: str,
+        backend: str,
         timeout: int = 180,
     ):
         if not username or not password:
             raise AuthError(
-                "Expected username and password. Or use environment variables `OPENPROTEIN_USERNAME` and `OPENPROTEIN_PASSWORD`"
+                "Expected username and password. Or use environment variables `OPENPROTEIN_USERNAME` and `OPENPROTEIN_PASSWORD`. Or provide these variables (`username` and `password`) in ~/.openprotein/config.toml."
             )
         super().__init__()
         self.backend = backend
@@ -132,7 +115,7 @@ class APISession(requests.Session):
         response = super().request(method, full_url, *args, **kwargs)
 
         if (js := kwargs.get("json")) and js is not None:
-            if total_size(js) > 1e6:
+            if _total_size(js) > 1e6:
                 warnings.warn(
                     "The requested payload is >1MB. There might be some delays or issues in processing. If the request fails, please try again with smaller sizes."
                 )
@@ -153,7 +136,7 @@ class APISession(requests.Session):
         return response
 
 
-def total_size(o, seen=None):
+def _total_size(o: Sequence | Mapping, seen=None):
     """Recursively finds size of objects including contents."""
     if seen is None:
         seen = set()
@@ -163,19 +146,27 @@ def total_size(o, seen=None):
     seen.add(obj_id)
     size = sys.getsizeof(o)
     if isinstance(o, dict):
-        size += sum((total_size(k, seen) + total_size(v, seen)) for k, v in o.items())
+        size += sum((_total_size(k, seen) + _total_size(v, seen)) for k, v in o.items())
     elif isinstance(o, (list, tuple, set, frozenset)):
-        size += sum(total_size(i, seen) for i in o)
+        size += sum(_total_size(i, seen) for i in o)
     return size
 
 
-class RestEndpoint:
-    pass
-
-
 class TimeoutError(requests.exceptions.HTTPError):
-    pass
+    """
+    An Exception raised due to timeout, possibly from overly large
+    requests.
+    """
 
 
 class CloudFrontError(requests.exceptions.HTTPError):
-    pass
+    """
+    An Exception raised due to CloudFront.
+
+    This is usually due to the strict timeout from CloudFront.
+    AWS CloudFront limits responses to return within 2 minutes.
+    This can be a bit prohibitive for our system that tends to
+    deal with large data. It is usually safe to just ignore/retry upon
+    hitting this error. Our system will scale up and still handle
+    the job.
+    """

openprotein_python-0.10.0/openprotein/chains.py

@@ -0,0 +1,12 @@
+import warnings
+
+warnings.warn(
+    "openprotein.chains is deprecated and will be removed in v0.11. "
+    "Use `from openprotein.molecules import DNA, RNA, Ligand` instead.",
+    FutureWarning,
+    stacklevel=2,
+)
+
+from openprotein.molecules import DNA, RNA, Ligand
+
+__all__ = ["DNA", "RNA", "Ligand"]

{openprotein_python-0.8.12 → openprotein_python-0.10.0}/openprotein/common/model_metadata.py

@@ -26,8 +26,8 @@ class ModelMetadata(BaseModel):
     id: str = Field(..., alias="model_id")
     description: ModelDescription
     max_sequence_length: int | None = None
-    dimension: int
-    output_types: list[str]
-    input_tokens: list[str] | None
+    dimension: int | None = None
+    output_types: list[str] | None = None
+    input_tokens: list[str] | None = None
     output_tokens: list[str] | None = None
-    token_descriptions: list[list[TokenInfo]]
+    token_descriptions: list[list[TokenInfo]] | None = None

openprotein_python-0.10.0/openprotein/common/reduction.py

@@ -0,0 +1,22 @@
+"""Reduction types used in OpenProtein."""
+
+from enum import Enum
+from typing import Literal
+
+
+class ReductionType(str, Enum):
+    """
+    ReductionType is an enumeration of the possible reduction types available.
+
+    Attributes:
+        MEAN : Mean reduction takes the mean of the embeddings across the sequence length dimension.
+        SUM : Sum reduction takes the sum of the embeddings across the sequence length dimension.
+    """
+
+    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_python-0.10.0/openprotein/common/residue_contants.py

@@ -0,0 +1,75 @@
+# fmt: off
+residue_atoms = {
+    "ALA": ["C", "CA", "CB", "N", "O"],
+    "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"],
+    "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"],
+    "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"],
+    "CYS": ["C", "CA", "CB", "N", "O", "SG"],
+    "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"],
+    "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"],
+    "GLY": ["C", "CA", "N", "O"],
+    "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"],
+    "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"],
+    "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"],
+    "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"],
+    "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"],
+    "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"],
+    "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"],
+    "SER": ["C", "CA", "CB", "N", "O", "OG"],
+    "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"],
+    "TRP": [
+        "C",
+        "CA",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "CE2",
+        "CE3",
+        "CZ2",
+        "CZ3",
+        "CH2",
+        "N",
+        "NE1",
+        "O",
+    ],
+    "TYR": [
+        "C",
+        "CA",
+        "CB",
+        "CG",
+        "CD1",
+        "CD2",
+        "CE1",
+        "CE2",
+        "CZ",
+        "N",
+        "O",
+        "OH",
+    ],
+    "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"],
+}
+
+restype_1to3 = {
+    "A": "ALA",
+    "R": "ARG",
+    "N": "ASN",
+    "D": "ASP",
+    "C": "CYS",
+    "Q": "GLN",
+    "E": "GLU",
+    "G": "GLY",
+    "H": "HIS",
+    "I": "ILE",
+    "L": "LEU",
+    "K": "LYS",
+    "M": "MET",
+    "F": "PHE",
+    "P": "PRO",
+    "S": "SER",
+    "T": "THR",
+    "W": "TRP",
+    "Y": "TYR",
+    "V": "VAL",
+}
+# fmt: on

{openprotein_python-0.8.12 → openprotein_python-0.10.0}/openprotein/embeddings/future.py

@@ -1,7 +1,7 @@
 """Future for embeddings-related jobs."""
 
 from collections import namedtuple
-from typing import Generator
+from typing import Any, Generator, Iterator, TypeVar, Union
 
 import numpy as np
 
@@ -25,7 +25,7 @@ from .schemas import (
 )
 
 
-class EmbeddingsResultFuture(MappedFuture, Future):
+class EmbeddingsResultFuture(MappedFuture[bytes, np.ndarray]):
     """Future for manipulating results for embeddings-related requests."""
 
     job: EmbeddingsJob | AttnJob | LogitsJob
@@ -45,14 +45,17 @@ class EmbeddingsResultFuture(MappedFuture, Future):
             else sequences
         )
 
-    def stream(self):
-        return api.request_get_embeddings_stream(session=self.session, job_id=self.id)
-
-    def get(self, verbose=False) -> list[np.ndarray]:
-        return super().get(verbose=verbose)
+    def stream_sync(self):
+        """
+        Stream the embeddings results synchronously using the streaming endpoint.
+        """
+        for i, array in enumerate(
+            api.request_get_embeddings_stream(session=self.session, job_id=self.id)
+        ):
+            yield self.sequences[i], array
 
     @property
-    def sequences(self) -> list[bytes] | list[str]:
+    def sequences(self):
         if self._sequences is None:
             self._sequences = api.get_request_sequences(
                 session=self.session, job_id=self.job.job_id, job_type=self.job.job_type
@@ -74,12 +77,12 @@ class EmbeddingsResultFuture(MappedFuture, Future):
         """
         return self.sequences
 
-    def get_item(self, sequence: bytes) -> np.ndarray:
+    def get_item(self, sequence: str | bytes) -> np.ndarray:
         """
         Get embedding results for specified sequence.
 
         Args:
-            sequence (bytes): sequence to fetch results for
+            sequence (str | bytes): sequence to fetch results for
 
         Returns:
             np.ndarray: embeddings
@@ -93,10 +96,13 @@ class EmbeddingsResultFuture(MappedFuture, Future):
         return api.result_decode(data)
 
 
-class EmbeddingsScoreFuture(StreamingFuture, Future):
-    """Future for manipulating results for embeddings score-related requests."""
+Score = namedtuple("Score", ["name", "sequence", "score"])
+SingleSiteScore = namedtuple("SingleSiteScore", ["mut_code", "score"])
+S = TypeVar("S", bound=Union[Score, SingleSiteScore])
 
-    job: ScoreJob | ScoreIndelJob | ScoreSingleSiteJob
+
+class BaseScoreFuture(StreamingFuture[S]):
+    """Future for manipulating results for embeddings score-related requests."""
 
     def __init__(
         self,
@@ -105,45 +111,67 @@ class EmbeddingsScoreFuture(StreamingFuture, Future):
         sequences: list[bytes] | list[str] | None = None,
     ):
         super().__init__(session=session, job=job)
-        self._sequences = sequences
+        # ensure all list[bytes]
+        self._sequences = (
+            [seq.encode() if isinstance(seq, str) else seq for seq in sequences]
+            if sequences is not None
+            else sequences
+        )
 
     @property
-    def sequences(self) -> list[bytes] | list[str]:
+    def sequences(self) -> list[bytes]:
         if self._sequences is None:
             self._sequences = api.get_request_sequences(self.session, self.job.job_id)
         return self._sequences
 
-    def stream(self) -> Generator:
-        if (
-            self.job_type == JobType.poet_generate
-            or self.job_type == JobType.embeddings_generate
-        ):
-            stream = api.request_get_generate_result(
-                session=self.session, job_id=self.id
-            )
-        else:
-            stream = api.request_get_score_result(session=self.session, job_id=self.id)
-        # mut_code, ... for ssp
-        # name, sequence, ... for score
-        header = next(stream)
-        score_start_index = 0
-        for i, col_name in enumerate(header):
-            if col_name.startswith("score"):
-                score_start_index = i
-                break
-        Score = namedtuple("Score", header[:score_start_index] + ["score"])
+
+class EmbeddingsScoreFuture(BaseScoreFuture[Score]):
+    """Future for manipulating results for embeddings score-related requests."""
+
+    job: ScoreJob | ScoreIndelJob
+
+    def stream(self) -> Iterator[Score]:
+        stream = api.request_get_score_result(session=self.session, job_id=self.id)
+        # name, sequence, ...
+        next(stream)  # ignore header
+        for line in stream:
+            # combine scores into numpy array
+            scores = np.array([float(s) for s in line[2:]])
+            output = Score(name=line[0], sequence=line[1], score=scores)
+            yield output
+
+
+class EmbeddingsScoreSingleSiteFuture(BaseScoreFuture[SingleSiteScore]):
+    """Future for manipulating results for embeddings score-related requests."""
+
+    job: ScoreSingleSiteJob
+
+    def stream(self) -> Iterator[SingleSiteScore]:
+        stream = api.request_get_score_result(session=self.session, job_id=self.id)
+        # name, sequence, ...
+        next(stream)  # ignore header
         for line in stream:
             # combine scores into numpy array
-            scores = np.array([float(s) for s in line[score_start_index:]])
-            output = Score(*line[:score_start_index], scores)
+            scores = np.array([float(s) for s in line[1:]])
+            output = SingleSiteScore(mut_code=line[0], score=scores)
             yield output
 
 
-class EmbeddingsGenerateFuture(EmbeddingsScoreFuture, StreamingFuture, Future):
+class EmbeddingsGenerateFuture(BaseScoreFuture[Score]):
     """Future for manipulating results for embeddings generate-related requests."""
 
     job: GenerateJob
 
+    def stream(self) -> Iterator[Score]:
+        stream = api.request_get_generate_result(session=self.session, job_id=self.id)
+        # name, sequence, ...
+        next(stream)  # ignore header
+        for line in stream:
+            # combine scores into numpy array
+            scores = np.array([float(s) for s in line[2:]])
+            output = Score(name=line[0], sequence=line[1], score=scores)
+            yield output
+
     @property
     def sequences(self):
         raise Exception("generate job does not support listing sequences")