openprotein-python 0.8.2__1-py3-none-any.whl

openprotein/svd/api.py

@@ -0,0 +1,206 @@
+"""SVD REST API for making HTTP calls to our SVD backend."""
+
+import io
+
+import numpy as np
+from pydantic import TypeAdapter
+
+from openprotein.base import APISession
+from openprotein.errors import APIError, InvalidParameterError
+
+from .schemas import SVDEmbeddingsJob, SVDFitJob, SVDMetadata
+
+PATH_PREFIX = "v1/embeddings/svd"
+
+
+def svd_list_get(session: APISession) -> list[SVDMetadata]:
+    """Get SVD job metadata for all SVDs. Including SVD dimension and sequence lengths."""
+    endpoint = PATH_PREFIX
+    response = session.get(endpoint)
+    return TypeAdapter(list[SVDMetadata]).validate_python(response.json())
+
+
+def svd_get(session: APISession, svd_id: str) -> SVDMetadata:
+    """Get SVD job metadata. Including SVD dimension and sequence lengths."""
+    endpoint = PATH_PREFIX + f"/{svd_id}"
+    response = session.get(endpoint)
+    return SVDMetadata.model_validate(response.json())
+
+
+def svd_get_sequences(session: APISession, svd_id: str) -> list[bytes]:
+    """
+    Get sequences used to fit an SVD.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    svd_id : str
+        SVD ID whose sequences to fetch
+
+    Returns
+    -------
+    sequences : List[bytes]
+    """
+    endpoint = PATH_PREFIX + f"/{svd_id}/sequences"
+    response = session.get(endpoint)
+    return TypeAdapter(list[bytes]).validate_python(response.json())
+
+
+def embed_get_sequence_result(
+    session: APISession, job_id: str, sequence: str | bytes
+) -> bytes:
+    """
+    Get encoded svd embeddings result for a sequence from the request ID.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    job_id : str
+        job ID to retrieve results from
+    sequence : bytes
+        sequence to retrieve results for
+
+    Returns
+    -------
+    result : bytes
+    """
+    if isinstance(sequence, bytes):
+        sequence = sequence.decode()
+    endpoint = PATH_PREFIX + f"/embed/{job_id}/{sequence}"
+    response = session.get(endpoint)
+    return response.content
+
+
+def embed_decode(data: bytes) -> np.ndarray:
+    """
+    Decode embedding as numpy array.
+
+    Args:
+        data (bytes): raw bytes encoding the array received over the API
+
+    Returns:
+        np.ndarray: decoded array
+    """
+    s = io.BytesIO(data)
+    return np.load(s, allow_pickle=False)
+
+
+def svd_delete(session: APISession, svd_id: str):
+    """
+    Delete and SVD model.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    svd_id : str
+        SVD model to delete
+
+    Returns
+    -------
+    bool
+    """
+
+    endpoint = PATH_PREFIX + f"/{svd_id}"
+    response = session.delete(endpoint)
+    if 200 <= response.status_code < 300:
+        return True
+    else:
+        raise APIError(response.text)
+
+
+def svd_fit_post(
+    session: APISession,
+    model_id: str,
+    sequences: list[bytes] | list[str] | None = None,
+    assay_id: str | None = None,
+    n_components: int = 1024,
+    reduction: str | None = None,
+    **kwargs,
+) -> SVDFitJob:
+    """
+    Create SVD fit job.
+
+    Parameters
+    ----------
+    session: APISession
+        Session object for API communication.
+    model_id: str
+        ID of embeddings model to use.
+    sequences: list of bytes or None, optional
+        Optional sequences to fit SVD with. Either use sequences or
+        assay_id. sequences is preferred.
+    assay_id: str | None, optional
+        Optional ID of assay containing sequences to fit SVD with. Either
+        use sequences or assay_id. Ignored if sequences are provided.
+    n_components: int
+        Number of SVD components to fit. Defaults to 1024
+    reduction: str | None
+        Type of embedding reduction to use for computing features.
+        E.g. "MEAN" or "SUM". Useful when dealing with variable length
+        sequence. Defaults to None.
+    kwargs:
+        Additional keyword arguments to be passed to foundational models, e.g. prompt_id for PoET models.
+
+    Returns
+    -------
+    Job
+    """
+
+    endpoint = PATH_PREFIX
+
+    body = {
+        "model_id": model_id,
+        "n_components": n_components,
+    }
+    if reduction is not None:
+        body["reduction"] = reduction
+    if sequences is not None:
+        # both provided
+        if assay_id is not None:
+            raise InvalidParameterError("Expected only either sequences or assay_id")
+        sequences = [(s if isinstance(s, str) else s.decode()) for s in sequences]
+        body["sequences"] = sequences
+    else:
+        # both are none
+        if assay_id is None:
+            raise InvalidParameterError("Expected either sequences or assay_id")
+        body["assay_id"] = assay_id
+    # add kwargs for embeddings kwargs
+    body.update(**kwargs)
+
+    response = session.post(endpoint, json=body)
+    # return job for metadata
+    return SVDFitJob.model_validate(response.json())
+
+
+def svd_embed_post(
+    session: APISession, svd_id: str, sequences: list[bytes] | list[str]
+) -> SVDEmbeddingsJob:
+    """
+    POST a request for embeddings from the given SVD model.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    svd_id : str
+        SVD model to use
+    sequences : List[bytes]
+        sequences to SVD
+
+    Returns
+    -------
+    Job
+    """
+    endpoint = PATH_PREFIX + f"/{svd_id}/embed"
+
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
+    body = {
+        "sequences": sequences_unicode,
+    }
+    response = session.post(endpoint, json=body)
+
+    return SVDEmbeddingsJob.model_validate(response.json())

openprotein/svd/models.py

@@ -0,0 +1,288 @@
+"""SVD model representations which can be used for creating reduced embeddings."""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from openprotein import config
+from openprotein.base import APISession
+from openprotein.common import FeatureType
+from openprotein.data import AssayDataset, AssayMetadata, DataAPI
+from openprotein.embeddings import EmbeddingModel, EmbeddingsResultFuture
+from openprotein.errors import InvalidParameterError
+from openprotein.jobs import Future, JobsAPI
+
+from . import api
+from .schemas import SVDEmbeddingsJob, SVDFitJob, SVDMetadata
+
+if TYPE_CHECKING:
+    from openprotein.predictor import PredictorModel
+    from openprotein.umap import UMAPModel
+
+
+class SVDModel(Future):
+    """
+    SVD model that can be used to create reduced embeddings.
+
+    The model is also implemented as a `Future` to allow waiting for a fit job.
+    """
+
+    job: SVDFitJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: SVDFitJob | None = None,
+        metadata: SVDMetadata | None = None,
+    ):
+        """Construct the SVD model using either job get or svd metadata get."""
+        # initialize the metadata
+        if metadata is None:
+            # use job to fetch metadata
+            if job is None:
+                raise ValueError("Expected svd metadata or job")
+            metadata = api.svd_get(session=session, svd_id=job.job_id)
+        self._metadata = metadata
+        if job is None:
+            jobs_api = getattr(session, "jobs", None)
+            assert isinstance(jobs_api, JobsAPI)
+            job = SVDFitJob.create(jobs_api.get_job(job_id=metadata.id))
+        # getter initializes job if not provided
+        super().__init__(session=session, job=job)
+
+    def __str__(self) -> str:
+        return str(self.metadata)
+
+    def __repr__(self) -> str:
+        return repr(self.metadata)
+
+    @property
+    def id(self):
+        return self._metadata.id
+
+    @property
+    def n_components(self):
+        """Number of components of the SVD."""
+        return self._metadata.n_components
+
+    @property
+    def sequence_length(self):
+        """Sequence length constraint of the SVD."""
+        return self._metadata.sequence_length
+
+    @property
+    def reduction(self):
+        """Reduction of embeddings used to fit the SVD."""
+        return self._metadata.reduction
+
+    @property
+    def metadata(self):
+        """Metadata of the SVD."""
+        self._refresh_metadata()
+        return self._metadata
+
+    def _refresh_metadata(self):
+        if not self._metadata.is_done():
+            self._metadata = api.svd_get(session=self.session, svd_id=self._metadata.id)
+
+    def get_model(self) -> EmbeddingModel:
+        model = EmbeddingModel.create(session=self.session, model_id=self._metadata.id)
+        return model
+
+    @property
+    def model(self) -> EmbeddingModel:
+        """Base embeddings model used for the SVD."""
+        return self.get_model()
+
+    def delete(self) -> bool:
+        """
+        Delete this SVD model.
+        """
+        return api.svd_delete(self.session, self.id)
+
+    def get(self, verbose: bool = False):
+        """Retrieve this SVD model itself."""
+        return self
+
+    def get_inputs(self) -> list[bytes]:
+        """
+        Get sequences used for svd job.
+
+        Returns
+        -------
+        list[bytes]
+            List of sequences
+        """
+        return api.svd_get_sequences(session=self.session, svd_id=self.id)
+
+    def embed(
+        self, sequences: list[bytes] | list[str], **kwargs
+    ) -> "SVDEmbeddingsResultFuture":
+        """
+        Use this SVD model to get reduced embeddings from input sequences.
+
+        Parameters
+        ----------
+        sequences : List[bytes]
+            List of protein sequences.
+
+        Returns
+        -------
+        SVDEmbeddingsResultFuture
+            Future result containing the reduced embeddings.
+        """
+        return SVDEmbeddingsResultFuture.create(
+            session=self.session,
+            job=api.svd_embed_post(
+                session=self.session, svd_id=self.id, sequences=sequences, **kwargs
+            ),
+            sequences=sequences,
+        )
+
+    def fit_umap(
+        self,
+        sequences: list[bytes] | list[str] | None = None,
+        assay: AssayDataset | None = None,
+        n_components: int = 2,
+        **kwargs,
+    ) -> "UMAPModel":
+        """
+        Fit an UMAP on the embedding results of this model. 
+
+        This function will create an UMAPModel based on the embeddings from this model \
+            as well as the hyperparameters specified in the args.  
+
+        Parameters
+        ----------
+        sequences : List[bytes] 
+            sequences to UMAP
+        n_components: int 
+            number of components in UMAP. Will determine output shapes
+        reduction: ReductionType | None
+            embeddings reduction to use (e.g. mean)
+
+        Returns
+        -------
+        UMAPModel
+            UMAP model fitted on the reduced embeddings from provided sequences or assay.    
+        """
+        # local import for cyclic dep
+        from openprotein.umap import UMAPAPI
+
+        umap_api = getattr(self.session, "umap", None)
+        assert isinstance(umap_api, UMAPAPI)
+
+        # Ensure either or
+        if (assay is None and sequences is None) or (
+            assay is not None and sequences is not None
+        ):
+            raise InvalidParameterError(
+                "Expected either assay or sequences to fit UMAP on!"
+            )
+        model_id = self.id
+        return umap_api.fit_umap(
+            model=model_id,
+            feature_type=FeatureType.SVD,
+            sequences=sequences,
+            assay_id=assay.id if assay is not None else None,
+            n_components=n_components,
+            **kwargs,
+        )
+
+    def fit_gp(
+        self,
+        assay: AssayMetadata | AssayDataset | str,
+        properties: list[str],
+        name: str | None = None,
+        description: str | None = None,
+        **kwargs,
+    ) -> "PredictorModel":
+        """
+        Fit a GP on assay using this embedding model and hyperparameters.
+
+        Parameters
+        ----------
+        assay : AssayMetadata or AssayDataset or str
+            Assay to fit GP on. Or its assay_id.
+        properties: list of str
+            Properties in the assay to fit the gp on.
+
+        Returns
+        -------
+        PredictorModel
+            Property predictor model trained using the reduced embeddings with provided assay and properties.
+        """
+        # 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
+        if (
+            self.sequence_length is not None
+            and assay.sequence_length != self.sequence_length
+        ):
+            raise InvalidParameterError(
+                f"Expected dataset to be of sequence length {self.sequence_length} due to svd fitted constraints"
+            )
+        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)"
+            )
+        return predictor_api.fit_gp(
+            assay=assay,
+            properties=properties,
+            feature_type=FeatureType.SVD,
+            model=self,
+            name=name,
+            description=description,
+            **kwargs,
+        )
+
+
+class SVDEmbeddingsResultFuture(EmbeddingsResultFuture, Future):
+    """SVD embeddings results represented as a future."""
+
+    job: SVDEmbeddingsJob
+
+    def wait(
+        self,
+        interval: int = config.POLLING_INTERVAL,
+        timeout: int | None = None,
+        verbose: bool = False,
+    ) -> list[np.ndarray]:
+        """Wait for the SVD embeddings job and retrieve the embeddings."""
+        return super().wait(interval, timeout, verbose)
+
+    def get(self, verbose=False) -> list[np.ndarray]:
+        """Get all the SVD reduced embeddings from the job."""
+        return super().get(verbose)
+
+    def get_item(self, sequence: bytes) -> np.ndarray:
+        """
+        Get SVD embeddings for specified sequence.
+
+        Parameters
+        ----------
+        sequence: bytes
+            Sequence to fetch SVD embeddings for.
+
+        Returns
+        -------
+        np.ndarray
+            SVD embeddings represented a numpy array.
+        """
+        data = api.embed_get_sequence_result(self.session, self.job.job_id, sequence)
+        return api.embed_decode(data)

openprotein/svd/schemas.py

@@ -0,0 +1,31 @@
+"""Schemas for OpenProtein SVD system."""
+
+from datetime import datetime
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+from openprotein.jobs import BatchJob, Job, JobStatus, JobType
+
+
+class SVDMetadata(BaseModel):
+    id: str
+    status: JobStatus
+    created_date: datetime | None = None
+    model_id: str
+    n_components: int
+    reduction: str | None = None
+    sequence_length: int | None = None
+
+    def is_done(self):
+        return self.status.done()
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+class SVDFitJob(Job):
+    job_type: Literal[JobType.svd_fit]
+
+
+class SVDEmbeddingsJob(Job, BatchJob):
+    job_type: Literal[JobType.svd_embed]

openprotein/svd/svd.py

@@ -0,0 +1,134 @@
+"""SVD API providing the interface for creating and using SVD models."""
+
+from openprotein.base import APISession
+from openprotein.common import ReductionType
+from openprotein.data import AssayDataset, AssayMetadata
+from openprotein.embeddings import EmbeddingModel, EmbeddingsAPI
+
+from . import api
+from .models import SVDModel
+
+
+class SVDAPI:
+    """SVD API providing the interface for creating and using SVD models."""
+
+    def __init__(
+        self,
+        session: APISession,
+    ):
+        self.session = session
+
+    def fit_svd(
+        self,
+        model_id: str,
+        sequences: list[bytes] | list[str] | None = None,
+        assay: AssayMetadata | AssayDataset | str | None = None,
+        n_components: int = 1024,
+        reduction: ReductionType | None = None,
+        **kwargs,
+    ) -> SVDModel:
+        """
+        Fit an SVD on the sequences with the specified model_id and hyperparameters (n_components).
+
+        Parameters
+        ----------
+        model_id : str
+            ID of embeddings model to use.
+        sequences : list of bytes or None, optional
+            Optional sequences to fit SVD with. Either use sequences or
+            assay_id. sequences is preferred.
+        assay : AssayMetadata or AssayDataset or str or None, optional
+            Optional assay containing sequences to fit SVD with.
+            Or its assay_id. Either use sequences or assay.
+            Ignored if sequences are provided.
+        n_components : int, optional
+            The number of components for the SVD. Defaults to 1024.
+        reduction : str or None, optional
+            Type of embedding reduction to use for computing features.
+            E.g. "MEAN" or "SUM". Useful when dealing with variable length
+            sequence. Defaults to None.
+        kwargs :
+            Additional keyword arguments to be passed to foundational models, e.g. prompt_id for PoET models.
+
+        Returns
+        -------
+        SVDModel
+            The SVD model being fit.
+        """
+        embeddings_api = getattr(self.session, "embedding", None)
+        assert isinstance(embeddings_api, EmbeddingsAPI)
+        model = embeddings_api.get_model(model_id)
+        assert isinstance(model, EmbeddingModel), "Expected EmbeddingModel"
+        # get assay_id
+        assay_id = (
+            assay.assay_id
+            if isinstance(assay, AssayMetadata)
+            else assay.id if isinstance(assay, AssayDataset) else assay
+        )
+        return SVDModel(
+            session=self.session,
+            job=api.svd_fit_post(
+                session=self.session,
+                model_id=model.id,
+                sequences=sequences,
+                assay_id=assay_id,
+                n_components=n_components,
+                reduction=reduction,
+                **kwargs,
+            ),
+        )
+
+    def get_svd(self, svd_id: str) -> SVDModel:
+        """
+        Get SVD job results. Including SVD dimension and sequence lengths.
+
+        Requires a successful SVD job from fit_svd
+
+        Parameters
+        ----------
+        svd_id : str
+            The ID of the SVD  job.
+        Returns
+        -------
+        SVDModel
+            The model with the SVD fit.
+        """
+        metadata = api.svd_get(self.session, svd_id)
+        return SVDModel(
+            session=self.session,
+            metadata=metadata,
+        )
+
+    def __delete_svd(self, svd_id: str) -> bool:
+        """
+        Delete SVD model.
+
+        Parameters
+        ----------
+        svd_id : str
+            The ID of the SVD job.
+        Returns
+        -------
+        bool
+            Whether or not the SVD was successfully deleted.
+
+        """
+        return api.svd_delete(self.session, svd_id)
+
+    def list_svd(self) -> list[SVDModel]:
+        """
+        List SVD models made by user.
+
+        Returns
+        -------
+        list of SVDModel
+            List of SVDs that the user has access to.
+
+        """
+        return [
+            SVDModel(
+                session=self.session,
+                metadata=metadata,
+            )
+            for metadata in api.svd_list_get(self.session)
+        ]

openprotein/umap/__init__.py

@@ -0,0 +1,9 @@
+"""
+UMAP module for OpenProtein for visualizing embeddings.
+
+isort:skip_file
+"""
+
+from .schemas import UMAPMetadata, UMAPEmbeddingsJob, UMAPFitJob
+from .models import UMAPModel, UMAPEmbeddingsResultFuture
+from .umap import UMAPAPI