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

openprotein/predictor/models.py

@@ -0,0 +1,374 @@
+"""Predictor models for making predictions on new sequences."""
+
+from typing import TYPE_CHECKING
+
+from openprotein.base import APISession
+from openprotein.data import AssayDataset, DataAPI
+from openprotein.embeddings import EmbeddingModel, EmbeddingsAPI
+from openprotein.errors import InvalidParameterError
+from openprotein.jobs import Future, JobsAPI, JobType
+from openprotein.svd import SVDAPI, SVDModel
+
+from . import api
+from .prediction import PredictionResultFuture
+from .schemas import (
+    PredictorEnsembleJob,
+    PredictorMetadata,
+    PredictorTrainJob,
+    PredictorType,
+)
+from .validate import CVResultFuture
+
+if TYPE_CHECKING:
+    from openprotein.design import ModelCriterion
+
+
+class PredictorModel(Future):
+    """
+    Class providing predict endpoint for fitted predictor models.
+
+    Also implements a Future that waits for train job.
+    """
+
+    job: PredictorTrainJob | None
+
+    def __init__(
+        self,
+        session: APISession,
+        job: PredictorTrainJob | PredictorEnsembleJob | None = None,
+        metadata: PredictorMetadata | None = None,
+    ):
+        """
+        Construct a predictor model.
+
+        Takes in either a train job, or the predictor metadata.
+
+        :meta private:
+        """
+        self._training_assay = None
+
+        # initialize the predictor metadata
+        if metadata is None:
+            if job is None or job.job_id is None:
+                raise ValueError("Expected predictor metadata or job")
+            metadata = api.predictor_get(session, job.job_id)
+        self._metadata = metadata
+        if job is None:
+            if metadata.model_spec.type != PredictorType.ENSEMBLE:
+                jobs_api = getattr(session, "jobs", None)
+                assert isinstance(jobs_api, JobsAPI)
+                job = PredictorTrainJob.create(jobs_api.get_job(job_id=metadata.id))
+            else:
+                job = PredictorEnsembleJob(
+                    created_date=self._metadata.created_date,
+                    status=self._metadata.status,
+                    job_type=JobType.predictor_train,
+                )
+        super().__init__(session, job)
+
+    def __str__(self) -> str:
+        return str(self.metadata)
+
+    def __repr__(self) -> str:
+        return repr(self.metadata)
+
+    def __or__(self, model: "PredictorModel") -> "PredictorModelGroup":
+        if self.sequence_length is not None:
+            if model.sequence_length != self.sequence_length:
+                raise ValueError(
+                    "Expected sequence lengths to either match or be None."
+                )
+        return PredictorModelGroup(
+            session=self.session,
+            models=[self, model],
+            sequence_length=self.sequence_length or model.sequence_length,
+            check_sequence_length=False,
+        )
+
+    def __lt__(self, target: float) -> "ModelCriterion":
+        from openprotein.design import ModelCriterion
+
+        if len(self.training_properties) == 1:
+            return ModelCriterion(
+                model_id=self.id,
+                measurement_name=self.training_properties[0],
+                criterion=ModelCriterion.Criterion(
+                    target=target, direction=ModelCriterion.Criterion.DirectionEnum.lt
+                ),
+            )
+        raise self.InvalidMultitaskModelToCriterion()
+
+    def __gt__(self, target: float) -> "ModelCriterion":
+        from openprotein.design import ModelCriterion
+
+        if len(self.training_properties) == 1:
+            return ModelCriterion(
+                model_id=self.id,
+                measurement_name=self.training_properties[0],
+                criterion=ModelCriterion.Criterion(
+                    target=target, direction=ModelCriterion.Criterion.DirectionEnum.gt
+                ),
+            )
+        raise self.InvalidMultitaskModelToCriterion()
+
+    def __eq__(self, target: float) -> "ModelCriterion":
+        from openprotein.design import ModelCriterion
+
+        if len(self.training_properties) == 1:
+            return ModelCriterion(
+                model_id=self.id,
+                measurement_name=self.training_properties[0],
+                criterion=ModelCriterion.Criterion(
+                    target=target, direction=ModelCriterion.Criterion.DirectionEnum.eq
+                ),
+            )
+        raise self.InvalidMultitaskModelToCriterion()
+
+    class InvalidMultitaskModelToCriterion(Exception):
+        """
+        Exception raised when trying to create model criterion from multitask predictor.
+
+        :meta private:
+        """
+
+    @property
+    def id(self):
+        """ID of predictor."""
+        return self._metadata.id
+
+    @property
+    def reduction(self):
+        """The reduction of th embeddings used to train the predictor, if any."""
+        return (
+            self._metadata.model_spec.features.reduction
+            if self._metadata.model_spec.features is not None
+            else None
+        )
+
+    @property
+    def sequence_length(self):
+        """The sequence length constraint on the predictor, if any."""
+        if (constraints := self._metadata.model_spec.constraints) is not None:
+            return constraints.sequence_length
+        return None
+
+    @property
+    def training_assay(self) -> AssayDataset:
+        """The assay the predictor was trained on."""
+        if self._training_assay is None:
+            self._training_assay = self.get_assay()
+        return self._training_assay
+
+    @property
+    def training_properties(self) -> list[str]:
+        """The list of properties the predictor was trained on."""
+        return self._metadata.training_dataset.properties
+
+    @property
+    def metadata(self):
+        """The predictor metadata."""
+        self._refresh_metadata()
+        return self._metadata
+
+    def _refresh_metadata(self):
+        if not self._metadata.is_done():
+            self._metadata = api.predictor_get(self.session, self._metadata.id)
+
+    def get_model(self) -> EmbeddingModel | SVDModel | None:
+        """Retrieve the embeddings or SVD model used to create embeddings to train on."""
+        if (
+            (features := self._metadata.model_spec.features)
+            and (model_id := features.model_id) is None
+            or features is None
+        ):
+            return None
+        elif features.type.upper() == "PLM":
+            model = EmbeddingModel.create(session=self.session, model_id=model_id)
+        elif features.type.upper() == "SVD":
+            svd_api = getattr(self.session, "svd", None)
+            assert isinstance(svd_api, SVDAPI)
+            model = svd_api.get_svd(svd_id=model_id)
+        else:
+            raise ValueError(f"Unexpected feature type {features.type}")
+        return model
+
+    @property
+    def model(self) -> EmbeddingModel | SVDModel | None:
+        """The embeddings or SVD model used to create embeddings to train on."""
+        return self.get_model()
+
+    def delete(self) -> bool:
+        """
+        Delete this predictor model.
+        """
+        return api.predictor_delete(self.session, self.id)
+
+    def get(self, verbose: bool = False):
+        """
+        Returns the train loss curves.
+        """
+        return self.metadata.traingraphs
+
+    def get_assay(self) -> AssayDataset:
+        """
+        Get assay used for train job.
+
+        Returns
+        -------
+            AssayDataset: Assay dataset used for train job.
+        """
+        data_api = getattr(self.session, "data", None)
+        assert isinstance(data_api, DataAPI)
+        return data_api.get(assay_id=self._metadata.training_dataset.assay_id)
+
+    def crossvalidate(self, n_splits: int | None = None) -> CVResultFuture:
+        """
+        Run a crossvalidation on the trained predictor.
+        """
+        return CVResultFuture.create(
+            session=self.session,
+            job=api.predictor_crossvalidate_post(
+                session=self.session,
+                predictor_id=self.id,
+                n_splits=n_splits,
+            ),
+        )
+
+    def predict(self, sequences: list[bytes] | list[str]) -> PredictionResultFuture:
+        """
+        Make predictions about the trained properties for a list of sequences.
+        """
+        if self.sequence_length is not None:
+            for sequence in sequences:
+                # convert to string to check token length
+                sequence = sequence if isinstance(sequence, str) else sequence.decode()
+                if len(sequence) != self.sequence_length:
+                    raise InvalidParameterError(
+                        f"Expected sequences to predict to be of length {self.sequence_length}"
+                    )
+        return PredictionResultFuture.create(
+            session=self.session,
+            job=api.predictor_predict_post(
+                session=self.session, predictor_id=self.id, sequences=sequences
+            ),
+        )
+
+    def single_site(self, sequence: bytes | str) -> PredictionResultFuture:
+        """
+        Compute the single-site mutated predictions of a base sequence.
+        """
+        if self.sequence_length is not None:
+            # convert to string to check token length
+            seq = sequence if isinstance(sequence, str) else sequence.decode()
+            if len(seq) != self.sequence_length:
+                raise InvalidParameterError(
+                    f"Expected sequence to predict to be of length {self.sequence_length}"
+                )
+        return PredictionResultFuture.create(
+            session=self.session,
+            job=api.predictor_predict_single_site_post(
+                session=self.session, predictor_id=self.id, base_sequence=sequence
+            ),
+        )
+
+
+class PredictorModelGroup(Future):
+    """
+    Class providing predict endpoint for fitted predictor models.
+
+    Also implements a Future that waits for train job.
+    """
+
+    __models__: list[PredictorModel]
+
+    def __init__(
+        self,
+        session: APISession,
+        models: list[PredictorModel],
+        sequence_length: int | None = None,
+        check_sequence_length: bool = True,  # turn off checking - prevent n^2 operation when chaining many
+    ):
+        if len(models) == 0:
+            raise ValueError("Expected at least one model to group")
+        # calculate and check sequence length compatibility
+        if check_sequence_length:
+            for m in models:
+                if m.sequence_length is not None:
+                    if sequence_length is None:
+                        sequence_length = m.sequence_length
+                    elif sequence_length != m.sequence_length:
+                        raise ValueError(
+                            "Expected sequence lengths of all models to either match or be None."
+                        )
+        self.sequence_length = sequence_length
+        self.session = session
+        self.__models__ = models
+
+    def __str__(self) -> str:
+        return repr(self.__models__)
+
+    def __repr__(self) -> str:
+        return repr(self.__models__)
+
+    def __or__(self, model: PredictorModel) -> "PredictorModelGroup":
+        if self.sequence_length is not None:
+            if model.sequence_length != self.sequence_length:
+                raise ValueError(
+                    "Expected sequence lengths to either match or be None."
+                )
+        return PredictorModelGroup(
+            session=self.session,
+            models=self.__models__ + [model],
+            sequence_length=self.sequence_length or model.sequence_length,
+            check_sequence_length=False,
+        )
+
+    def predict(self, sequences: list[bytes] | list[str]) -> PredictionResultFuture:
+        """
+        Make predictions about the trained properties for a list of sequences.
+        """
+        if self.sequence_length is not None:
+            for sequence in sequences:
+                # convert to string to check token length
+                sequence = sequence if isinstance(sequence, str) else sequence.decode()
+                if len(sequence) != self.sequence_length:
+                    raise InvalidParameterError(
+                        f"Expected sequences to predict to be of length {self.sequence_length}"
+                    )
+        return PredictionResultFuture.create(
+            session=self.session,
+            job=api.predictor_predict_multi_post(
+                session=self.session,
+                predictor_ids=[m.id for m in self.__models__],
+                sequences=sequences,
+            ),
+        )
+
+    def single_site(self, sequence: bytes | str) -> PredictionResultFuture:
+        """
+        Compute the single-site mutated predictions of a base sequence.
+        """
+        if self.sequence_length is not None:
+            # convert to string to check token length
+            seq = sequence if isinstance(sequence, str) else sequence.decode()
+            if len(seq) != self.sequence_length:
+                raise InvalidParameterError(
+                    f"Expected sequence to predict to be of length {self.sequence_length}"
+                )
+        return PredictionResultFuture.create(
+            session=self.session,
+            job=api.predictor_predict_single_site_post(
+                session=self.session, predictor_id=self.id, base_sequence=sequence
+            ),
+        )
+
+    def get(self, verbose: bool = False):
+        """
+        Returns the predictor model.
+
+        :meta private:
+        """
+        return self
+
+    def delete(self):
+        return api.predictor_delete(session=self.session, predictor_id=self.id)

openprotein/predictor/prediction.py

@@ -0,0 +1,79 @@
+"""Prediction results represented as futures."""
+
+import numpy as np
+
+from openprotein.base import APISession
+from openprotein.jobs import Future
+
+from . import api
+from .schemas import (
+    PredictJob,
+    PredictMultiJob,
+    PredictMultiSingleSiteJob,
+    PredictSingleSiteJob,
+)
+
+
+class PredictionResultFuture(Future):
+    """Prediction results represented as a future."""
+
+    job: PredictJob | PredictSingleSiteJob | PredictMultiJob | PredictMultiSingleSiteJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: (
+            PredictJob
+            | PredictSingleSiteJob
+            | PredictMultiJob
+            | PredictMultiSingleSiteJob
+        ),
+        sequences: list[bytes] | None = None,
+    ):
+        super().__init__(session, job)
+        self._sequences = sequences
+
+    @property
+    def sequences(self):
+        if self._sequences is None:
+            self._sequences = api.predictor_predict_get_sequences(
+                self.session, self.job.job_id
+            )
+        return self._sequences
+
+    @property
+    def id(self):
+        return self.job.job_id
+
+    def __keys__(self):
+        return self.sequences
+
+    def get_item(self, sequence: bytes) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Get embedding results for specified sequence.
+
+        Args:
+            sequence (bytes): sequence to fetch results for
+
+        Returns:
+            mu (np.ndarray): means of sequence prediction
+            var (np.ndarray): variances of sequence prediction
+        """
+        data = api.predictor_predict_get_sequence_result(
+            self.session, self.job.job_id, sequence
+        )
+        return api.decode_predict(data)
+
+    def get(self, verbose: bool = False) -> tuple[np.ndarray, np.ndarray]:
+        """
+        Get embedding results for specified sequence.
+
+        Args:
+            sequence (bytes): sequence to fetch results for
+
+        Returns:
+            mu (np.ndarray): means of predictions
+            var (np.ndarray): variances of predictions
+        """
+        data = api.predictor_predict_get_batched_result(self.session, self.job.job_id)
+        return api.decode_predict(data, batched=True)

openprotein/predictor/predictor.py

@@ -0,0 +1,242 @@
+"""Predictor API providing the interface to train and predict predictors."""
+
+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.svd import SVDAPI, SVDModel
+
+from . import api
+from .models import PredictorModel
+
+
+class PredictorAPI:
+    """Predictor API providing the interface to train and predict predictors."""
+
+    def __init__(
+        self,
+        session: APISession,
+    ):
+        self.session = session
+
+    def get_predictor(
+        self,
+        predictor_id: str,
+        include_stats: bool = False,
+        include_calibration_curves: bool = False,
+    ) -> PredictorModel:
+        """
+        Get predictor by model_id.
+
+        PredictorModel allows all the usual prediction job manipulation:
+        e.g. making POST and GET requests for this predictor specifically.
+
+        Parameters
+        ----------
+        predictor_id : str
+            The model identifier.
+        include_stats : bool
+            Whether to include stats of the predictor from the latest evaluation
+            (pearson, spearman, ece). Default is False.
+        include_calibration_curves : bool
+            Whether to include calibration curves of the predictor from the latest
+            evaluation. Default is False.
+
+        Returns
+        -------
+        PredictorModel
+            The predictor model to inspect and make predictions with.
+
+        Raises
+        ------
+        HTTPError
+            If the GET request does not succeed.
+        """
+        metadata = api.predictor_get(
+            session=self.session,
+            predictor_id=predictor_id,
+            include_stats=include_stats,
+            include_calibration_curves=include_calibration_curves,
+        )
+        return PredictorModel(
+            session=self.session,
+            metadata=metadata,
+        )
+
+    def list_predictors(
+        self,
+        limit: int = 100,
+        offset: int = 0,
+        include_stats: bool = False,
+        include_calibration_curves: bool = False,
+    ) -> list[PredictorModel]:
+        """
+        List predictors available.
+
+        Parameters
+        ----------
+        limit : int
+            Limit of the number of predictors to return in list. Default is 100.
+        offset : int
+            Offset to the predictors to query for paged queries. Default is 0.
+        include_stats : bool
+            Whether to include stats of each predictor from the latest evaluation
+            (pearson, spearman, ece). Default is False.
+        include_calibration_curves : bool
+            Whether to include calibration curves of each predictor from the latest
+            evaluation. Default is False.
+
+        Returns
+        -------
+        list[PredictorModel]
+            List of predictor models to inspect and make predictions with.
+
+        Raises
+        ------
+        HTTPError
+            If the GET request does not succeed.
+        """
+        metadatas = api.predictor_list(
+            session=self.session,
+            limit=limit,
+            offset=offset,
+            include_stats=include_stats,
+            include_calibration_curves=include_calibration_curves,
+        )
+        return [
+            PredictorModel(
+                session=self.session,
+                metadata=m,
+            )
+            for m in metadatas
+        ]
+
+    def fit_gp(
+        self,
+        assay: AssayDataset | AssayMetadata | str,
+        properties: list[str],
+        model: EmbeddingModel | SVDModel | str,
+        feature_type: FeatureType | None = None,
+        reduction: ReductionType | None = None,
+        name: str | None = None,
+        description: str | None = None,
+        **kwargs,
+    ) -> PredictorModel:
+        """
+        Fit a GP on an assay with the specified feature model and hyperparameters.
+
+        Parameters
+        ----------
+        assay : AssayMetadata or AssayDataset or str
+            Assay to fit GP on.
+        properties : list of str
+            Properties in the assay to fit the gp on.
+        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
+            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
+            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.
+        kwargs :
+            Additional keyword arguments to be passed to foundational models, e.g. prompt_id for PoET models.
+
+        Returns
+        -------
+        PredictorModel
+            The GP 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 PredictorModel(
+            session=self.session,
+            job=api.predictor_fit_gp_post(
+                session=self.session,
+                assay_id=assay_id,
+                properties=properties,
+                feature_type=feature_type,
+                model_id=model_id,
+                reduction=reduction,
+                name=name,
+                description=description,
+                **kwargs,
+            ),
+        )
+
+    def delete_predictor(self, predictor_id: str) -> bool:
+        """
+        Delete predictor model.
+
+        Parameters
+        ----------
+        predictor_id : str
+            The ID of the predictor.
+        Returns
+        -------
+        bool
+            True: successful deletion
+
+        """
+        return api.predictor_delete(session=self.session, predictor_id=predictor_id)
+
+    def ensemble(self, predictors: list[PredictorModel]) -> PredictorModel:
+        """
+        Ensemble predictor models together.
+
+        Parameters
+        __________
+        predictors: list[PredictorModel]
+            List of predictors to ensemble together.
+        Returns
+        -------
+        PredictorModel
+            Ensembled predictor model
+        """
+        return PredictorModel(
+            session=self.session,
+            metadata=api.predictor_ensemble(
+                session=self.session,
+                predictor_ids=[predictor.id for predictor in predictors],
+            ),
+        )