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

openprotein/fold/esmfold.py

@@ -0,0 +1,54 @@
+"""Community-based ESMFold model."""
+
+from collections.abc import Sequence
+
+from openprotein.base import APISession
+from openprotein.common import ModelMetadata
+
+from . import api
+from .future import FoldResultFuture
+from .models import FoldModel
+
+
+class ESMFoldModel(FoldModel):
+    """
+    Class providing inference endpoints for Facebook's ESMFold structure prediction models.
+    """
+
+    model_id: str = "esmfold"
+
+    def __init__(
+        self,
+        session: APISession,
+        model_id: str,
+        metadata: ModelMetadata | None = None,
+    ):
+        super().__init__(session=session, model_id=model_id, metadata=metadata)
+
+    def fold(
+        self, sequences: Sequence[bytes | str], num_recycles: int | None = None
+    ) -> FoldResultFuture:
+        """
+        Fold sequences using this model.
+
+        Parameters
+        ----------
+        sequences : Sequence[bytes | str]
+            sequences to fold
+        num_recycles : int | None
+            number of times to recycle models
+        Returns
+        -------
+            FoldResultFuture
+        """
+        result = FoldResultFuture.create(
+            session=self.session,
+            job=api.fold_models_post(
+                session=self.session,
+                model_id=self.model_id,
+                sequences=sequences,
+                num_recycles=num_recycles,
+            ),
+        )
+        assert isinstance(result, FoldResultFuture)
+        return result

openprotein/fold/fold.py

@@ -0,0 +1,107 @@
+"""Fold API interface for making structure prediction requests."""
+
+from openprotein.base import APISession
+
+from . import api
+from .alphafold2 import AlphaFold2Model
+from .boltz import Boltz1Model, Boltz1xModel, Boltz2Model
+from .esmfold import ESMFoldModel
+from .future import FoldComplexResultFuture, FoldResultFuture
+from .models import (
+    FoldModel,
+)
+
+
+class FoldAPI:
+    """
+    Fold API provides a high level interface for making protein structure predictions.
+    """
+
+    #: Boltz-2 model
+    boltz2: Boltz2Model
+    boltz_2: Boltz2Model
+    #: Boltz-1x model
+    boltz1x: Boltz1xModel
+    boltz_1x: Boltz1xModel
+    #: Boltz-1 model
+    boltz1: Boltz1Model
+    boltz_1: Boltz1Model
+    af2: AlphaFold2Model
+    #: AlphaFold-2 model
+    alphafold2: AlphaFold2Model
+    #: ESMFold model
+    esmfold: ESMFoldModel
+
+    def __init__(self, session: APISession):
+        self.session = session
+        self._load_models()
+
+    def _load_models(self):
+        # Dynamically add model instances as attributes - precludes any drift
+        models = self.list_models()
+        for model in models:
+            model_name = model.id.replace("-", "_")  # hyphens out
+            setattr(self, model_name, model)
+        # Setup aliases safely
+        if getattr(self, "alphafold2", None):
+            self.af2 = self.alphafold2
+        if getattr(self, "boltz_1", None):
+            self.boltz1 = self.boltz_1
+        if getattr(self, "boltz_1x", None):
+            self.boltz1x = self.boltz_1x
+        if getattr(self, "boltz_2", None):
+            self.boltz2 = self.boltz_2
+
+    def list_models(self) -> list[FoldModel]:
+        """list models available for creating folds of your sequences"""
+        models = []
+        for model_id in api.fold_models_list_get(self.session):
+            models.append(
+                FoldModel.create(
+                    session=self.session, model_id=model_id, default=FoldModel
+                )
+            )
+        return models
+
+    def get_model(self, model_id: str) -> FoldModel:
+        """
+        Get model by model_id. 
+
+        FoldModel allows all the usual job manipulation: \
+            e.g. making POST and GET requests for this model specifically. 
+
+
+        Parameters
+        ----------
+        model_id : str
+            the model identifier
+
+        Returns
+        -------
+        FoldModel
+            The model
+
+        Raises
+        ------
+        HTTPError
+            If the GET request does not succeed.
+        """
+        return FoldModel.create(
+            session=self.session, model_id=model_id, default=FoldModel
+        )
+
+    def get_results(self, job) -> FoldResultFuture | FoldComplexResultFuture:
+        """
+        Retrieves the results of a fold job.
+
+        Parameters
+        ----------
+        job : Job
+            The fold job whose results are to be retrieved.
+
+        Returns
+        -------
+        FoldResultFuture
+            An instance of FoldResultFuture
+        """
+        return FoldResultFuture.create(job=job, session=self.session)

openprotein/fold/future.py

@@ -0,0 +1,509 @@
+"""Fold prediction results represented as futures."""
+
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+from pydantic.type_adapter import TypeAdapter
+from typing_extensions import Self
+
+from openprotein import config
+from openprotein.base import APISession
+from openprotein.chains import DNA, RNA, Ligand
+from openprotein.jobs import Future, MappedFuture
+from openprotein.protein import Protein
+
+from . import api
+from .schemas import FoldJob
+
+if TYPE_CHECKING:
+    from .boltz import BoltzAffinity, BoltzConfidence
+
+
+class FoldResultFuture(MappedFuture, Future):
+    """
+    Fold results represented as a future.
+
+    Attributes
+    ----------
+    job : FoldJob
+        The fold job associated with this future.
+    """
+
+    job: FoldJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: FoldJob,
+        sequences: list[bytes] | None = None,
+        max_workers: int = config.MAX_CONCURRENT_WORKERS,
+    ):
+        """
+        Initialize a FoldResultFuture instance.
+
+        Parameters
+        ----------
+        session : APISession
+            The API session to use for requests.
+        job : FoldJob
+            The fold job associated with this future.
+        sequences : list[bytes], optional
+            List of sequences submitted for the fold request. If None, sequences will be fetched.
+        max_workers : int, optional
+            Maximum number of concurrent workers. Default is config.MAX_CONCURRENT_WORKERS.
+        """
+        super().__init__(session, job, max_workers)
+        if sequences is None:
+            sequences = api.fold_get_sequences(self.session, job_id=job.job_id)
+        self._sequences = sequences
+
+    @classmethod
+    def create(
+        cls: type[Self],
+        session: APISession,
+        job: FoldJob,
+        **kwargs,
+    ) -> "Self | FoldComplexResultFuture":
+        """
+        Factory method to create a FoldResultFuture or FoldComplexResultFuture.
+
+        Parameters
+        ----------
+        session : APISession
+            The API session to use for requests.
+        job : FoldJob
+            The fold job associated with this future.
+        **kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        FoldResultFuture or FoldComplexResultFuture
+            An instance of FoldResultFuture or FoldComplexResultFuture depending on the model.
+        """
+        model_id = api.fold_get(session=session, job_id=job.job_id).model_id
+        if model_id.startswith("boltz") or model_id.startswith("alphafold"):
+            return FoldComplexResultFuture(session=session, job=job, **kwargs)
+        else:
+            return cls(session=session, job=job, **kwargs)
+
+    @property
+    def sequences(self) -> list[bytes]:
+        """
+        Get the sequences submitted for the fold request.
+
+        Returns
+        -------
+        list[bytes]
+            List of sequences.
+        """
+        if self._sequences is None:
+            self._sequences = api.fold_get_sequences(self.session, self.job.job_id)
+        return self._sequences
+
+    @property
+    def model_id(self) -> str:
+        """
+        Get the model ID used for the fold request.
+
+        Returns
+        -------
+        str
+            Model ID.
+        """
+        if self._model_id is None:
+            self._model_id = api.fold_get(
+                session=self.session, job_id=self.job.job_id
+            ).model_id
+        return self._model_id
+
+    @property
+    def id(self):
+        """
+        Get the ID of the fold request.
+
+        Returns
+        -------
+        str
+            Fold job ID.
+        """
+        return self.job.job_id
+
+    def __keys__(self):
+        """
+        Get the list of sequences submitted for the fold request.
+
+        Returns
+        -------
+        list of bytes
+            List of sequences.
+        """
+        return self.sequences
+
+    def get(self, verbose=False) -> list[tuple[str, bytes]]:
+        """
+        Retrieve the fold results as a list of tuples mapping sequence to PDB-encoded string.
+
+        Parameters
+        ----------
+        verbose : bool, optional
+            If True, print verbose output. Default is False.
+
+        Returns
+        -------
+        list[tuple[str, str]]
+            List of tuples mapping sequence to PDB-encoded string.
+        """
+        return super().get(verbose=verbose)
+
+    def get_item(self, sequence: bytes) -> bytes:
+        """
+        Get fold results for a specified sequence.
+
+        Parameters
+        ----------
+        sequence : bytes
+            Sequence to fetch results for.
+
+        Returns
+        -------
+        bytes
+            Fold result for the specified sequence.
+        """
+        data = api.fold_get_sequence_result(self.session, self.job.job_id, sequence)
+        return data
+
+
+class FoldComplexResultFuture(Future):
+    """
+    Future for manipulating results of a fold complex request.
+
+    Attributes
+    ----------
+    job : FoldJob
+        The fold job associated with this future.
+    """
+
+    job: FoldJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: FoldJob,
+        model_id: str | None = None,
+        proteins: list[Protein] | None = None,
+        ligands: list[Ligand] | None = None,
+        dnas: list[DNA] | None = None,
+        rnas: list[RNA] | None = None,
+    ):
+        """
+        Initialize a FoldComplexResultFuture instance.
+
+        Parameters
+        ----------
+        session : APISession
+            The API session to use for requests.
+        job : FoldJob
+            The fold job associated with this future.
+        model_id : str, optional
+            Model ID used for the fold request.
+        proteins : list[Protein], optional
+            List of proteins submitted for fold request.
+        ligands : list[Ligand], optional
+            List of ligands submitted for fold request.
+        dnas : list[DNA], optional
+            List of DNAs submitted for fold request.
+        rnas : list[RNA], optional
+            List of RNAs submitted for fold request.
+        """
+        super().__init__(session, job)
+        self._model_id = model_id
+        self._proteins = proteins
+        self._ligands = ligands
+        self._dnas = dnas
+        self._rnas = rnas
+        self._initialized = not (proteins == ligands == dnas == rnas == None)
+        self._pae: np.ndarray | None = None
+        self._pde: np.ndarray | None = None
+        self._plddt: np.ndarray | None = None
+        self._confidence: list["BoltzConfidence"] | None = None
+        self._affinity: "BoltzAffinity | None" = None
+
+    @property
+    def model_id(self) -> str:
+        """
+        Get the model ID used for the fold request.
+
+        Returns
+        -------
+        str
+            Model ID.
+        """
+        if self._model_id is None:
+            self._model_id = api.fold_get(
+                session=self.session, job_id=self.job.job_id
+            ).model_id
+        return self._model_id
+
+    def __get_chains(self):
+        """
+        Internal method to initialize chain objects (proteins, dnas, rnas, ligands)
+        from the fold job arguments.
+        """
+        args = api.fold_get(session=self.session, job_id=self.job.job_id).args
+        assert args is not None and "sequences" in args
+        for chain in args["sequences"]:
+            assert isinstance(chain, dict)
+            for chain_type, chain_info in chain:
+                if chain_type == "protein":
+                    self._proteins = self._proteins or []
+                    protein = Protein(sequence=chain_info["sequence"])
+                    protein.chain_id = chain_info.get("id")
+                    protein.msa = chain_info.get("msa_id")
+                    self._proteins.append(protein)
+                elif chain_type == "dna":
+                    self._dnas = self._dnas or []
+                    dna = DNA(sequence=chain_info["sequence"])
+                    dna.chain_id = chain_info.get("id")
+                    self._dnas.append(dna)
+                elif chain_type == "rna":
+                    self._rnas = self._rnas or []
+                    rna = RNA(sequence=chain_info["sequence"])
+                    rna.chain_id = chain_info.get("id")
+                    self._rnas.append(rna)
+                elif chain_type == "ligand":
+                    self._ligands = self._ligands or []
+                    ligand = Ligand(
+                        chain_id=chain_info.get("id"),
+                        ccd=chain_info.get("ccd"),
+                        smiles=chain_info.get("smiles"),
+                    )
+                    self._ligands.append(ligand)
+                else:
+                    pass
+        self._initialized = True
+
+    @property
+    def proteins(self) -> list[Protein] | None:
+        """
+        Get the proteins submitted for the fold request.
+
+        Returns
+        -------
+        list[Protein] or None
+            List of Protein objects or None.
+        """
+        if not self._initialized:
+            self.__get_chains()
+        return self._proteins
+
+    @property
+    def dnas(self) -> list[DNA] | None:
+        """
+        Get the DNAs submitted for the fold request.
+
+        Returns
+        -------
+        list[DNA] or None
+            List of DNA objects or None.
+        """
+        if not self._initialized:
+            self.__get_chains()
+        return self._dnas
+
+    @property
+    def rnas(self) -> list[RNA] | None:
+        """
+        Get the RNAs submitted for the fold request.
+
+        Returns
+        -------
+        list[RNA] or None
+            List of RNA objects or None.
+        """
+        if not self._initialized:
+            self.__get_chains()
+        return self._rnas
+
+    @property
+    def ligands(self) -> list[Ligand] | None:
+        """
+        Get the ligands submitted for the fold request.
+
+        Returns
+        -------
+        list[Ligand] or None
+            List of Ligand objects or None.
+        """
+        if not self._initialized:
+            self.__get_chains()
+        return self._ligands
+
+    @property
+    def pae(self) -> np.ndarray:
+        """
+        Get the Predicted Aligned Error (PAE) matrix.
+
+        Returns
+        -------
+        np.ndarray
+            PAE matrix.
+
+        Raises
+        ------
+        AttributeError
+            If PAE is not supported for the model.
+        """
+        if self.model_id not in {"boltz-1", "boltz-1x", "boltz-2"}:
+            raise AttributeError("pae not supported for non-Boltz model")
+        if self._pae is None:
+            pae = api.fold_get_complex_extra_result(
+                session=self.session, job_id=self.job.job_id, key="pae"
+            )
+            assert isinstance(pae, np.ndarray)
+            self._pae = pae
+        return self._pae
+
+    @property
+    def pde(self) -> np.ndarray:
+        """
+        Get the Predicted Distance Error (PDE) matrix.
+
+        Returns
+        -------
+        np.ndarray
+            PDE matrix.
+
+        Raises
+        ------
+        AttributeError
+            If PDE is not supported for the model.
+        """
+        if self.model_id not in {"boltz-1", "boltz-1x", "boltz-2"}:
+            raise AttributeError("pde not supported for non-Boltz model")
+        if self._pde is None:
+            pde = api.fold_get_complex_extra_result(
+                session=self.session, job_id=self.job.job_id, key="pde"
+            )
+            assert isinstance(pde, np.ndarray)
+            self._pde = pde
+        return self._pde
+
+    @property
+    def plddt(self) -> np.ndarray:
+        """
+        Get the Predicted Local Distance Difference Test (pLDDT) scores.
+
+        Returns
+        -------
+        np.ndarray
+            pLDDT scores.
+
+        Raises
+        ------
+        AttributeError
+            If pLDDT is not supported for the model.
+        """
+        if self.model_id not in {"boltz-1", "boltz-1x", "boltz-2"}:
+            raise AttributeError("plddt not supported for non-Boltz model")
+        if self._plddt is None:
+            plddt = api.fold_get_complex_extra_result(
+                session=self.session, job_id=self.job.job_id, key="plddt"
+            )
+            assert isinstance(plddt, np.ndarray)
+            self._plddt = plddt
+        return self._plddt
+
+    @property
+    def confidence(self) -> list["BoltzConfidence"]:
+        """
+        Retrieve the confidences of the structure prediction.
+
+        Note
+        ----
+        This is only currently supported for Boltz models.
+
+        Returns
+        -------
+        list[BoltzConfidence]
+            List of BoltzConfidence objects.
+
+        Raises
+        ------
+        AttributeError
+            If confidence is not supported for the model.
+        """
+        if self.model_id not in {"boltz-1", "boltz-1x", "boltz-2"}:
+            raise AttributeError("confidence not supported for non-Boltz model")
+        if self._confidence is None:
+            confidence = api.fold_get_complex_extra_result(
+                session=self.session, job_id=self.job.job_id, key="confidence"
+            )
+            assert isinstance(confidence, list)
+            self._confidence = TypeAdapter(list[BoltzConfidence]).validate_python(
+                confidence
+            )
+        return self._confidence
+
+    @property
+    def affinity(self) -> "BoltzAffinity":
+        """
+        Retrieve the predicted binding affinities.
+
+        Note
+        ----
+        This is only currently supported for Boltz models.
+
+        Returns
+        -------
+        BoltzAffinity
+            BoltzAffinity object containing the predicted affinities.
+
+        Raises
+        ------
+        AttributeError
+            If affinity is not supported for the model.
+        """
+        if self.model_id not in {"boltz-1", "boltz-1x", "boltz-2"}:
+            raise AttributeError("affinity not supported for non-Boltz model")
+        if self._affinity is None:
+            affinity = api.fold_get_complex_extra_result(
+                session=self.session, job_id=self.job.job_id, key="affinity"
+            )
+            assert isinstance(affinity, dict)
+            self._affinity = BoltzAffinity.parse_obj_with_models(affinity)
+        return self._affinity
+
+    @property
+    def id(self):
+        """
+        Get the ID of the fold request.
+
+        Returns
+        -------
+        str
+            Fold job ID.
+        """
+        return self.job.job_id
+
+    def get(self, format: Literal["pdb", "mmcif"] = "mmcif", verbose=False) -> bytes:
+        """
+        Retrieve the fold results as a single bytestring.
+
+        Defaults to mmCIF for complexes. Additional predicted properties like plddt and pae should be accessed from their respective properties, i.e. `.plddt` and `.pae`.
+
+        Parameters
+        ----------
+        format : {'pdb', 'mmcif'}, optional
+            Output format. Default is 'mmcif'.
+        verbose : bool, optional
+            If True, print verbose output. Default is False.
+
+        Returns
+        -------
+        bytes
+            Fold result as a bytestring.
+        """
+        return api.fold_get_complex_result(
+            session=self.session, job_id=self.id, format=format
+        )