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

openprotein/umap/api.py

@@ -0,0 +1,259 @@
+"""UMAP REST API for making HTTP calls to our UMAP backend."""
+
+import io
+
+import numpy as np
+import pandas as pd
+from pydantic import TypeAdapter
+
+from openprotein.base import APISession
+from openprotein.errors import APIError, InvalidParameterError
+
+from .schemas import FeatureType, UMAPEmbeddingsJob, UMAPFitJob, UMAPMetadata
+
+PATH_PREFIX = "v1/umap"
+
+
+def umap_list_get(session: APISession) -> list[UMAPMetadata]:
+    """Get UMAP job metadata for all UMAPs. Including UMAP dimension and sequence lengths."""
+    endpoint = PATH_PREFIX
+    response = session.get(endpoint)
+    return TypeAdapter(list[UMAPMetadata]).validate_python(response.json())
+
+
+def umap_get(session: APISession, umap_id: str) -> UMAPMetadata:
+    """Get UMAP job metadata. Including UMAP dimension and sequence lengths."""
+    endpoint = PATH_PREFIX + f"/{umap_id}"
+    response = session.get(endpoint)
+    return UMAPMetadata.model_validate(response.json())
+
+
+def umap_get_sequences(session: APISession, umap_id: str) -> list[bytes]:
+    """
+    Get sequences used to fit an UMAP.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    umap_id : str
+        UMAP ID whose sequences to fetch
+
+    Returns
+    -------
+    sequences : List[bytes]
+    """
+    endpoint = PATH_PREFIX + f"/{umap_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 umap 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_get_batch_result(session: APISession, job_id: str) -> bytes:
+    """
+    Get encoded umap embeddings batched result from the request ID.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    job_id : str
+        Job ID to retrieve results from
+
+    Returns
+    -------
+    result : bytes
+    """
+    endpoint = PATH_PREFIX + f"/embed/{job_id}/csv"
+    response = session.get(endpoint)
+    return response.content
+
+
+def embed_decode(data: bytes) -> np.ndarray:
+    """
+    Decode embedding as numpy array.
+
+    Parameters
+    ----------
+        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 embed_batch_decode(data: bytes) -> np.ndarray:
+    """
+    Decode prediction scores.
+
+    Args:
+        data (bytes): raw bytes encoding the array received over the API
+        batched (bool): whether or not the result was batched. affects the retrieved csv format whether they contain additional columns and header rows.
+
+    Returns:
+        mus (np.ndarray): decoded array of means
+        vars (np.ndarray): decoded array of variances
+    """
+    s = io.BytesIO(data)
+    # should contain header and sequence column
+    df = pd.read_csv(s)
+    umaps = df.iloc[:, 1:].values
+    return umaps
+
+
+def umap_delete(session: APISession, umap_id: str) -> bool:
+    """
+    Delete and UMAP model.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    umap_id : str
+        UMAP model to delete
+
+    Returns
+    -------
+    bool
+    """
+
+    endpoint = PATH_PREFIX + f"/{umap_id}"
+    response = session.delete(endpoint)
+    if 200 <= response.status_code < 300:
+        return True
+    else:
+        raise APIError(response.text)
+
+
+def umap_fit_post(
+    session: APISession,
+    model_id: str,
+    feature_type: str,
+    sequences: list[bytes] | list[str] | None = None,
+    assay_id: str | None = None,
+    n_components: int = 2,
+    n_neighbors: int = 15,
+    min_dist: float = 0.1,
+    reduction: str | None = None,
+    **kwargs,
+) -> UMAPFitJob:
+    """
+    Create UMAP fit job.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    model_id : str
+        Model to use. Can be either svd_id or id of a foundational model.
+    feature_type: str
+        Type of feature to use for fitting UMAP. Either PLM or SVD.
+    sequences : list[bytes] | None, optional
+        Optional sequences to fit UMAP with. Either use sequences or
+        assay_id. sequences is preferred.
+    assay_id: str | None, optional
+        Optional ID of assay containing sequences to fit UMAP with.
+        Either use sequences or assay_id. Ignored if sequences are
+        provided.
+    n_components: int
+        Number of UMAP components to fit. Defaults to 2.
+    n_neighbors: int
+        Number of neighbors to use for fitting. Defaults to 15.
+    min_dist: float
+        Minimum distance in UMAP fitting. Defaults to 0.1.
+    reduction : str | None
+        Embedding reduction to use for fitting the UMAP. Defaults to None.
+    kwargs:
+        Additional keyword arguments to be passed to foundational models, e.g. prompt_id for PoET models.
+
+    Returns
+    -------
+    UMAPFitJob
+    """
+
+    endpoint = PATH_PREFIX
+
+    body = {
+        "model_id": model_id,
+        "feature_type": feature_type,
+        "n_components": n_components,
+        "n_neighbors": n_neighbors,
+        "min_dist": min_dist,
+    }
+    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 UMAPFitJob.model_validate(response.json())
+
+
+def umap_embed_post(
+    session: APISession, umap_id: str, sequences: list[bytes] | list[str]
+) -> UMAPEmbeddingsJob:
+    """
+    POST a request for embeddings from the given UMAP model.
+
+    Parameters
+    ----------
+    session : APISession
+        Session object for API communication.
+    umap_id : str
+        UMAP model to use
+    sequences : List[bytes]
+        sequences to UMAP
+
+    Returns
+    -------
+    UMAPEmbeddingsJob
+    """
+    endpoint = PATH_PREFIX + f"/{umap_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 UMAPEmbeddingsJob.model_validate(response.json())

openprotein/umap/models.py

@@ -0,0 +1,211 @@
+"""UMAP models on the OpenProtein system which can be used directly to create projected embeddings useful for visualization."""
+
+import numpy as np
+
+from openprotein import config
+from openprotein.base import APISession
+from openprotein.embeddings import EmbeddingModel, EmbeddingsResultFuture
+from openprotein.jobs import Future, JobsAPI
+
+from . import api
+from .schemas import UMAPEmbeddingsJob, UMAPFitJob, UMAPMetadata
+
+
+class UMAPModel(Future):
+    """
+    UMAP model that can be used to create projected embeddings.
+
+    The model is also implemented as a `Future` to allow waiting for a fit job.
+    The projected embeddings of the sequences used to fit the UMAP can be
+    accessed using `embeddings`.
+    """
+
+    job: UMAPFitJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: UMAPFitJob | None = None,
+        metadata: UMAPMetadata | None = None,
+    ):
+        # Initializes with either job get or umap metadata get.
+        if metadata is None:
+            # use job to fetch metadata
+            if job is None:
+                raise ValueError("Expected umap metadata or job")
+            metadata = api.umap_get(session, job.job_id)
+        self._metadata = metadata
+        if job is None:
+            jobs_api = getattr(session, "jobs", None)
+            assert isinstance(jobs_api, JobsAPI)
+            job = UMAPFitJob.create(jobs_api.get_job(job_id=metadata.id))
+        self._sequences = None
+        self._embeddings = None
+        # getter initializes job if not provided
+        super().__init__(session, job)
+
+    def __str__(self) -> str:
+        return str(self.metadata)
+
+    def __repr__(self) -> str:
+        return repr(self.metadata)
+
+    @property
+    def id(self):
+        """UMAP unique identifier."""
+
+        return self._metadata.id
+
+    @property
+    def n_components(self):
+        """Number of components specified for the UMAP."""
+
+        return self._metadata.n_components
+
+    @property
+    def n_neighbors(self):
+        """Number of neighbors specified for the UMAP."""
+
+        return self._metadata.n_neighbors
+
+    @property
+    def min_dist(self):
+        """Minimum distance specified for the UMAP."""
+
+        return self._metadata.min_dist
+
+    @property
+    def sequence_length(self):
+        """Sequence length constraint of the UMAP."""
+
+        return self._metadata.sequence_length
+
+    @property
+    def reduction(self):
+        """Reduction used to fit the UMAP."""
+
+        return self._metadata.reduction
+
+    @property
+    def metadata(self):
+        """Metadata of the UMAP."""
+
+        self._refresh_metadata()
+        return self._metadata
+
+    @property
+    def sequences(self):
+        """The sequences used to fit the UMAP."""
+
+        if self._sequences is not None:
+            return self._sequences
+        self._sequences = self.get_inputs()
+        return self._sequences
+
+    @property
+    def embeddings(self):
+        """The projected embeddings of the sequences used to fit the UMAP."""
+
+        if self._embeddings is not None:
+            return self._embeddings
+        data = api.embed_get_batch_result(session=self.session, job_id=self.id)
+        embeddings = [
+            (seq, umap)
+            for seq, umap in zip(self.sequences, api.embed_batch_decode(data))
+        ]
+        self._embeddings = embeddings
+        return self._embeddings
+
+    def _refresh_metadata(self):
+        if not self._metadata.is_done():
+            self._metadata = api.umap_get(self.session, 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 UMAP."""
+        return self.get_model()
+
+    def delete(self) -> bool:
+        """
+        Delete this UMAP model.
+        """
+        return api.umap_delete(self.session, self.id)
+
+    def get(self, verbose: bool = False):
+        """Retrieve this UMAP model itself."""
+        return self
+
+    def get_inputs(self) -> list[bytes]:
+        """
+        Get sequences used for umap job.
+
+        Returns
+        -------
+        list[bytes]
+            list of sequences
+        """
+        return api.umap_get_sequences(session=self.session, umap_id=self.id)
+
+    def embed(
+        self, sequences: list[bytes] | list[str], **kwargs
+    ) -> "UMAPEmbeddingsResultFuture":
+        """
+        Use this UMAP model to get projected embeddings from input sequences.
+
+        Parameters
+        ----------
+        sequences : List[bytes]
+            List of protein sequences.
+
+        Returns
+        -------
+        UMAPEmbeddingsResultFuture
+            Future result containing the projected embeddings.
+        """
+        return UMAPEmbeddingsResultFuture.create(
+            session=self.session,
+            job=api.umap_embed_post(
+                session=self.session, umap_id=self.id, sequences=sequences, **kwargs
+            ),
+            sequences=sequences,
+        )
+
+
+class UMAPEmbeddingsResultFuture(EmbeddingsResultFuture, Future):
+    """UMAP embeddings results represented as a future."""
+
+    job: UMAPEmbeddingsJob
+
+    def wait(
+        self,
+        interval: int = config.POLLING_INTERVAL,
+        timeout: int | None = None,
+        verbose: bool = False,
+    ) -> list[np.ndarray]:
+        """Wait for the UMAP embeddings job and retrieve the embeddings."""
+        return super().wait(interval, timeout, verbose)
+
+    def get(self, verbose=False) -> list[np.ndarray]:
+        """Get all the UMAP projected embeddings from the job."""
+        return super().get(verbose)
+
+    def get_item(self, sequence: bytes) -> np.ndarray:
+        """
+        Get UMAP embeddings for specified sequence.
+
+        Parameters
+        ----------
+        sequence: bytes
+            Sequence to fetch UMAP embeddings for.
+
+        Returns
+        -------
+        np.ndarray
+            UMAP embeddings represented a numpy array.
+        """
+        data = api.embed_get_sequence_result(self.session, self.job.job_id, sequence)
+        return api.embed_decode(data)

openprotein/umap/schemas.py

@@ -0,0 +1,35 @@
+"""Schemas for OpenProtein UMAP system."""
+
+from datetime import datetime
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict
+
+from openprotein.common import FeatureType
+from openprotein.jobs import BatchJob, Job, JobStatus, JobType
+
+
+class UMAPMetadata(BaseModel):
+    id: str
+    status: JobStatus
+    created_date: datetime | None = None
+    model_id: str
+    feature_type: FeatureType
+    n_components: int = 2
+    n_neighbors: int = 15
+    min_dist: float = 0.1
+    reduction: str | None = None
+    sequence_length: int | None = None
+
+    def is_done(self):
+        return self.status.done()
+
+    model_config = ConfigDict(protected_namespaces=())
+
+
+class UMAPFitJob(Job):
+    job_type: Literal[JobType.umap_fit]
+
+
+class UMAPEmbeddingsJob(Job, BatchJob):
+    job_type: Literal[JobType.umap_embed]

openprotein/umap/umap.py

@@ -0,0 +1,175 @@
+"""UMAP API providing the interface to fit and run UMAP visualizations."""
+
+from openprotein.base import APISession
+from openprotein.common import FeatureType, ReductionType
+from openprotein.data import AssayDataset, AssayMetadata
+from openprotein.embeddings import EmbeddingModel, EmbeddingsAPI
+from openprotein.errors import InvalidParameterError
+from openprotein.jobs import JobsAPI
+from openprotein.svd import SVDAPI, SVDModel
+
+from . import api
+from .models import UMAPModel
+
+
+class UMAPAPI:
+    """UMAP API providing the interface to fit and run UMAP visualizations."""
+
+    def __init__(
+        self,
+        session: APISession,
+    ):
+        self.session = session
+
+    def fit_umap(
+        self,
+        model: EmbeddingModel | SVDModel | str,
+        feature_type: FeatureType | None = None,
+        sequences: list[bytes] | list[str] | None = None,
+        assay: AssayMetadata | AssayDataset | str | None = None,
+        n_components: int = 2,
+        n_neighbors: int = 15,
+        min_dist: float = 0.1,
+        reduction: ReductionType | None = None,
+        **kwargs,
+    ) -> UMAPModel:
+        """
+        Fit an UMAP on the sequences with the specified model_id and hyperparameters (n_components).
+
+        Parameters
+        ----------
+        sequences: list of bytes or None, optional
+            Optional sequences to fit UMAP 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.
+        model : EmbeddingModel or SVDModel or str
+            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, optional
+            Type of features to use for encoding sequences. "SVD" or "PLM".
+            None would require model to be EmbeddingModel or SVDModel.
+        n_components : int, optional
+            Number of UMAP components to fit. Defaults to 2.
+        n_neighbors : int, optional
+            Number of neighbors to use for fitting. Defaults to 15.
+        min_dist : float, optional
+            Minimum distance in UMAP fitting. Defaults to 0.1.
+        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
+        -------
+        UMAPModel
+            The UMAP model being fit.
+        """
+        # extract feature 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"
+            )
+        # get model if model_id
+        if feature_type == FeatureType.PLM:
+            if reduction is None:
+                raise InvalidParameterError(
+                    "Expected reduction if using EmbeddingModel"
+                )
+            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 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
+        # get assay_id
+        assay_id = (
+            assay.assay_id
+            if isinstance(assay, AssayMetadata)
+            else assay.id if isinstance(assay, AssayDataset) else assay
+        )
+        return UMAPModel(
+            session=self.session,
+            job=api.umap_fit_post(
+                session=self.session,
+                model_id=model_id,
+                feature_type=feature_type,
+                sequences=sequences,
+                assay_id=assay_id,
+                n_components=n_components,
+                n_neighbors=n_neighbors,
+                min_dist=min_dist,
+                reduction=reduction,
+                **kwargs,
+            ),
+        )
+
+    def get_umap(self, umap_id: str) -> UMAPModel:
+        """
+        Get UMAP job results. Including UMAP dimension and sequence lengths.
+
+        Requires a successful UMAP job from fit_umap.
+
+        Parameters
+        ----------
+        umap_id : str
+            The ID of the UMAP  job.
+        Returns
+        -------
+        UMAPModel
+            The model with the UMAP fit.
+        """
+        metadata = api.umap_get(self.session, umap_id)
+        return UMAPModel(session=self.session, metadata=metadata)
+
+    def __delete_umap(self, umap_id: str) -> bool:
+        """
+        Delete UMAP model.
+
+        Parameters
+        ----------
+        umap_id : str
+            The ID of the UMAP  job.
+        Returns
+        -------
+        bool
+            True: successful deletion
+
+        """
+        return api.umap_delete(self.session, umap_id)
+
+    def list_umap(self) -> list[UMAPModel]:
+        """
+        List UMAP models made by user.
+
+        Takes no args.
+
+        Returns
+        -------
+        list[UMAPModel]
+            UMAPModels
+
+        """
+        jobs_api = getattr(self.session, "jobs", None)
+        assert isinstance(jobs_api, JobsAPI)
+        return [
+            UMAPModel(session=self.session, metadata=metadata)
+            for metadata in api.umap_list_get(self.session)
+        ]

openprotein/utils/uuid.py

@@ -0,0 +1,29 @@
+from uuid import UUID
+
+
+def is_valid_uuid(uuid_to_test: str, version=4):
+    """
+    Check if uuid_to_test is a valid UUID.
+
+     Parameters
+    ----------
+    uuid_to_test : str
+    version : {1, 2, 3, 4}
+
+     Returns
+    -------
+    `True` if uuid_to_test is a valid UUID, otherwise `False`.
+
+     Examples
+    --------
+    >>> is_valid_uuid('c9bf9e57-1685-4c89-bafb-ff5af830be8a')
+    True
+    >>> is_valid_uuid('c9bf9e58')
+    False
+    """
+
+    try:
+        uuid_obj = UUID(uuid_to_test, version=version)
+    except ValueError:
+        return False
+    return str(uuid_obj) == uuid_to_test