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

openprotein/jobs/jobs.py

@@ -0,0 +1,69 @@
+from datetime import datetime
+
+from openprotein import config
+from openprotein.base import APISession
+
+from . import api
+from .futures import Future
+from .schemas import Job, JobStatus, JobType
+
+
+class JobsAPI:
+    """API interface to get jobs."""
+
+    def __init__(self, session: APISession):
+        self.session = session
+
+    def list(
+        self,
+        status: JobStatus | None = None,
+        job_type: JobType | None = None,
+        assay_id: str | None = None,
+        more_recent_than: datetime | str | None = None,
+        limit: int = 100,
+    ) -> list[Job]:
+        """List jobs."""
+        more_recent_than_str = (
+            more_recent_than.isoformat()
+            if isinstance(more_recent_than, datetime)
+            else more_recent_than
+        )
+        return [
+            Job.create(j)
+            for j in api.jobs_list(
+                self.session,
+                status=status,
+                job_type=job_type,
+                assay_id=assay_id,
+                more_recent_than=more_recent_than_str,
+                limit=limit,
+            )
+        ]
+
+    def get_job(self, job_id: str) -> Job:
+        return api.job_get(session=self.session, job_id=job_id)
+
+    def get(self, job_id: str, verbose: bool = False) -> Future:  # Job:
+        """
+        Get job by ID.
+
+        Notes
+        -----
+        This retrieves the job and loads it as a future so you can do `wait` and `get`.
+        """
+        return self.__load(job_id=job_id)
+        # return Job.create(job.job_get(session=self.session, job_id=job_id))
+
+    def __load(self, job_id: str) -> Future:
+        """Loads a job by ID and returns the future."""
+        return Future.create(session=self.session, job_id=job_id)
+
+    def wait(
+        self,
+        future: Future,
+        interval=config.POLLING_INTERVAL,
+        timeout: int | None = None,
+        verbose: bool = False,
+    ):
+        """Waits on a job result."""
+        return future.wait(interval=interval, timeout=timeout, verbose=verbose)

openprotein/jobs/schemas.py

@@ -0,0 +1,135 @@
+import logging
+from datetime import datetime
+from enum import Enum
+from typing import Union
+
+from pydantic import BaseModel, ConfigDict, TypeAdapter
+from requests import Response
+from typing_extensions import Self
+
+logger = logging.getLogger(__name__)
+
+
+class JobType(str, Enum):
+    """
+    Type of job.
+
+    Describes the types of jobs that can be done.
+    """
+
+    stub = "stub"
+
+    workflow_preprocess = "/workflow/preprocess"
+    workflow_train = "/workflow/train"
+    workflow_embed_umap = "/workflow/embed/umap"
+    workflow_predict = "/workflow/predict"
+    workflow_predict_single_site = "/workflow/predict/single_site"
+    workflow_crossvalidate = "/workflow/crossvalidate"
+    workflow_evaluate = "/workflow/evaluate"
+    workflow_design = "/workflow/design"
+
+    align_align = "/align/align"
+    align_prompt = "/align/prompt"
+
+    clustalo = "/align/clustalo"
+    mafft = "/align/mafft"
+    abnumber = "/align/abnumber"
+
+    poet = "/poet"
+    poet_score = "/poet/score"
+    poet_single_site = "/poet/single_site"
+    poet_generate = "/poet/generate"
+    poet_score_indel = "/poet/score/indel"
+
+    embeddings_embed = "/embeddings/embed"
+    embeddings_svd = "/embeddings/svd"
+    embeddings_attn = "/embeddings/attn"
+    embeddings_logits = "/embeddings/logits"
+    embeddings_embed_reduced = "/embeddings/embed_reduced"
+
+    svd_fit = "/svd/fit"
+    svd_embed = "/svd/embed"
+
+    umap_fit = "/umap/fit"
+    umap_embed = "/umap/embed"
+
+    embeddings_fold = "/embeddings/fold"
+
+    # predictor jobs
+    predictor_train = "/predictor/train"
+    predictor_predict = "/predictor/predict"
+    predictor_crossvalidate = "/predictor/crossvalidate"
+    predictor_predict_single_site = "/predictor/predict_single_site"
+    predictor_predict_multi = "/predictor/predict_multi"
+    predictor_predict_multi_single_site = "/predictor/predict_multi_single_site"
+
+    # designer
+    designer = "/design"
+
+
+class JobStatus(str, Enum):
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+    SUCCESS = "SUCCESS"
+    FAILURE = "FAILURE"
+    RETRYING = "RETRYING"
+    CANCELED = "CANCELED"
+
+    def done(self):
+        return (
+            (self is self.SUCCESS) or (self is self.FAILURE) or (self is self.CANCELED)
+        )  # noqa: E501
+
+    def cancelled(self):
+        return self is self.CANCELED
+
+
+class Job(BaseModel):
+    job_id: str
+    # new emb service get doesnt have job_type
+    job_type: str
+    status: JobStatus
+    created_date: datetime
+    start_date: datetime | None = None
+    end_date: datetime | None = None
+    prerequisite_job_id: str | None = None
+    progress_message: str | None = None
+    progress_counter: int | None = None
+    sequence_length: int | None = None
+
+    @classmethod
+    def create(cls, obj: "Job | Response | dict", **kwargs) -> Self:
+        # parse specific child Job from base Job or Response
+        try:
+            # try to parse as subclass job
+            # get dict form
+            d = (
+                obj.json()
+                if isinstance(obj, Response)
+                else obj.model_dump() if isinstance(obj, Job) else obj
+            )
+            job_classes = Job.__subclasses__()
+            job = TypeAdapter(Union[tuple(job_classes)]).validate_python(d | kwargs)  # type: ignore
+        except Exception as e:
+            job = Job.model_validate(d | kwargs)
+        return job  # type: ignore - static checker cannot know runtime type
+
+    # hide extra allowed fields
+    def __repr_args__(self):
+        for k, v in self.__dict__.items():
+            field = self.model_fields.get(k)
+            if field and field.repr:
+                yield k, v
+
+        yield from (
+            (k, getattr(self, k))
+            for k, v in self.model_computed_fields.items()
+            if v.repr
+        )
+
+    # allows to carry over subclassed job fields when factory creating
+    model_config = ConfigDict(extra="allow")
+
+
+class BatchJob(BaseModel):
+    num_records: int | None = None

openprotein/models/__init__.py

@@ -0,0 +1,4 @@
+"""Make the ModelsAPI class available on the package."""
+
+from .foundation.rfdiffusion import RFdiffusionFuture, RFdiffusionJob
+from .models import ModelsAPI

openprotein/models/base.py

@@ -0,0 +1,63 @@
+"""Base protein models for working with proteins."""
+
+from openprotein.base import APISession
+from openprotein.common import ModelMetadata
+from openprotein.jobs import Future
+
+
+class ProteinModel:
+    def __init__(
+        self,
+        session: APISession,
+        model_id: str,
+        metadata: ModelMetadata | None = None,
+    ):
+        self.session = session
+        self.id = model_id
+        self._metadata = metadata
+        self.__doc__ = self.__fmt_doc()
+
+    def __fmt_doc(self):
+        summary = str(self.metadata.description.summary)
+        return f"""\t{summary}
+        \t max_sequence_length = {self.metadata.max_sequence_length}
+        \t supported outputs = {self.metadata.output_types}
+        \t supported tokens = {self.metadata.input_tokens} 
+        """
+
+    def __str__(self) -> str:
+        return self.id
+
+    def __repr__(self) -> str:
+        return self.id
+
+    @property
+    def metadata(self):
+        """
+        ModelMetadata for this model.
+
+        Returns
+        -------
+        ModelMetadata
+            The metadata associated with this model.
+        """
+        if self._metadata is None:
+            self._metadata = self.get_metadata()
+        return self._metadata
+
+    def get_metadata(self) -> ModelMetadata:
+        """
+        Get model metadata for this model.
+
+        Returns
+        -------
+        ModelMetadata
+            The metadata associated with this model.
+        """
+        raise NotImplementedError("`get_metadata` not implemented for this model")
+
+    def predict(self, *args, **kwargs) -> Future:
+        """
+        Alias for the `design` method to conform to the base ProteinModel.
+        """
+        raise NotImplementedError("`predict` not implemented for this model")

openprotein/models/foundation/rfdiffusion.py

@@ -0,0 +1,283 @@
+"""RFdiffusion model for protein structure and sequence design."""
+
+from typing import BinaryIO, Literal
+
+from pydantic import BaseModel, Field
+
+from openprotein.base import APISession
+from openprotein.common import ModelMetadata
+from openprotein.common.model_metadata import ModelDescription
+from openprotein.jobs import Future, Job
+from openprotein.models.base import ProteinModel
+from openprotein.protein import Protein
+
+
+class Contig(BaseModel):
+    """Defines a contig segment for protein design."""
+
+    length: str = Field(..., description="Length range, e.g., '10-20' or '100'")
+    chain: str | None = Field(None, description="Chain to sample from")
+
+
+class Hotspot(BaseModel):
+    """Specifies a hotspot residue constraint."""
+
+    res_id: str = Field(
+        ..., description="Residue identifier, e.g., 'A100' for chain A, residue 100"
+    )
+
+
+class RFdiffusionRequest(BaseModel):
+    "Specification for an RFdiffusion request."
+
+    n: int = 1
+    # protein: Protein
+    structure_text: str | None = None
+    # contigs: list[Contig]
+    contigs: str | None = None
+    inpaint_seq: str | None = None
+    provide_seq: str | None = None
+    # hotspots: list[Hotspot]
+    hotspot: str | None = None
+    T: int | None = None
+    partial_T: int | None = None
+    use_active_site_model: bool | None = None
+    use_beta_model: bool | None = None
+
+    # Simplified symmetry options
+    symmetry: Literal["cyclic", "dihedral", "tetrahedral"] | None = None
+    order: int | None = None
+    add_potential: bool | None = None
+
+    # Fold conditioning
+    scaffold_target_structure_text: str | None = None
+    scaffold_target_use_struct: bool = False
+
+
+class RFdiffusionJob(Job):
+    """Job schema for an RFdiffusion request."""
+
+    job_type: Literal["/models/rfdiffusion"]
+
+
+class RFdiffusionFuture(Future):
+    """Future for handling the results of an RFdiffusion job."""
+
+    job: RFdiffusionJob
+
+    def get_pdb(self, replicate: int = 0) -> str:
+        """
+        Retrieve the PDB file for a specific design.
+
+        Args:
+            design_index (int): The 0-based index of the design to retrieve.
+
+        Returns:
+            str: The content of the PDB file as a string.
+        """
+        return _rfdiffusion_api_result_get(
+            session=self.session, job_id=self.id, replicate=replicate
+        )
+
+    def get(self, replicate: int = 0):
+        """Default result accessor, returns the first PDB."""
+        # TODO handle different design index
+        return self.get_pdb(replicate=replicate)
+
+
+def _rfdiffusion_api_post(
+    session: APISession, request: RFdiffusionRequest, **kwargs
+) -> RFdiffusionJob:
+    """
+    POST a request for RFdiffusion design.
+
+    Returns a Job object that can be used to retrieve results later.
+    """
+    endpoint = "v1/design/models/rfdiffusion"
+    body = request.model_dump(exclude_none=True)
+    body.update(kwargs)
+    response = session.post(endpoint, json=body)
+    return RFdiffusionJob.model_validate(response.json())
+
+
+def _rfdiffusion_api_get_metadata(session: APISession) -> ModelMetadata:
+    """
+    POST a request for RFdiffusion design.
+
+    Returns a Job object that can be used to retrieve results later.
+    """
+    endpoint = f"v1/design/models/rfdiffusion"
+    response = session.get(endpoint)
+    return ModelMetadata.model_validate(response.json())
+
+
+def _rfdiffusion_api_result_get(
+    session: APISession, job_id: str, replicate: int = 0
+) -> str:
+    """
+    POST a request for RFdiffusion design.
+
+    # Returns a Job object that can be used to retrieve results later.
+    """
+    endpoint = f"v1/design/{job_id}/results"
+    response = session.get(endpoint, params={"replicate": replicate})
+    return response.text
+
+
+class RFdiffusionModel(ProteinModel):
+    """
+    RFdiffusion model for generating de novo protein structures.
+
+    This model supports functionalities like unconditional design, scaffolding,
+    and binder design.
+    """
+
+    model_id: str = "rfdiffusion"
+
+    def __init__(self, session: APISession, model_id: str = "rfdiffusion"):
+        # The model_id from the API might be more specific, e.g., "rfdiffusion-v1.1"
+        super().__init__(session, model_id)
+
+    def get_metadata(self) -> ModelMetadata:
+        return ModelMetadata(
+            model_id="rfdiffusion",
+            description=ModelDescription(summary="RFdiffusion"),
+            dimension=0,
+            output_types=["pdb"],
+            input_tokens=[],
+            token_descriptions=[[]],
+        )
+
+    def generate(
+        self,
+        n: int = 1,
+        structure_file: str | bytes | BinaryIO | None = None,
+        contigs: int | str | None = None,
+        inpaint_seq: str | None = None,
+        provide_seq: str | None = None,
+        hotspot: str | None = None,
+        T: int | None = None,
+        partial_T: int | None = None,
+        use_active_site_model: bool | None = None,
+        use_beta_model: bool | None = None,
+        # Symmetry options
+        symmetry: Literal["cyclic", "dihedral", "tetrahedral"] | None = None,
+        order: int | None = None,
+        add_potential: bool | None = None,
+        # Fold conditioning
+        scaffold_target_structure_file: str | bytes | BinaryIO | None = None,
+        scaffold_target_use_struct: bool = False,
+        **kwargs,
+    ) -> RFdiffusionFuture:
+        """
+        Run a protein structure generate job using RFdiffusion.
+
+        Parameters
+        ----------
+        n : int, optional
+            The number of unique design trajectories to run (default is 1).
+        structure_file : BinaryIO, optional
+            An input PDB file (as a file-like object) used for inpainting or other
+            guided design tasks where parts of an existing structure are provided.
+        contigs : int, str, optional
+            Defines the lengths and connectivity of chain segments for the desired
+            structure, specified in RFdiffusion's contig string format.
+            Required for most design tasks. Example: 150, '10-20/A100-110/10-20' for a
+            binder design.
+        inpaint_seq : str, optional
+            A string specifying the regions in the input structure to mask for
+            in-painting. Example: 'A1-A10/A30-40'.
+        provide_seq : str, optional
+            A string specifying which segments of the contig have a provided
+            sequence. Example: 'A1-A10/A30-40'.
+        hotspot : str, optional
+            A string specifying hotspot residues to constrain during design,
+            typically for functional sites. Example: 'A10,A12,A14'.
+        T : int, optional
+            The number of timesteps for the diffusion process.
+        partial_T : int, optional
+            The number of timesteps for partial diffusion.
+        use_active_site_model : bool, optional
+            If True, uses the active site model checkpoint, which has been finetuned to
+            better keep very small motifs in place in the output for motif scaffolding
+            (default is False).
+        use_beta_model : bool, optional
+            If True, uses the complex beta model checkpoint, which generates a
+            greater diversity of topologies but has not been extensively
+            experimentally validated (default is False).
+        symmetry : {"cyclic", "dihedral", "tetrahedral"}, optional
+            The type of symmetry to apply to the design.
+        order : int, optional
+            The order of the symmetry (e.g., 3 for C3 or D3 symmetry).
+            Must be provided if `symmetry` is set.
+        add_potential : bool, optional
+            A flag to toggle an additional potential to guide the design.
+            This defaults to true in the case of symmetric design.
+        scaffold_target_structure_file : str, bytes, BinaryIO, optional
+            A PDB file (which can be the text string or bytes or the file-like
+            object) containing a scaffold structure to be used as a structural
+            guide. It could also be used as a target when doing scaffold guided
+            binder design with `scaffold_target_use_struct`.
+        scaffold_target_use_struct : bool, optional
+            Whether or not to use the provided scaffold structure as a target.
+            Otherwise, it is used only as a topology guide.
+
+        Other Parameters
+        ----------------
+        **kwargs : dict
+            Additional keyword args that are passed directly to the rfdiffusion
+            inference script. Overwrites any preceding options.
+
+        Returns
+        -------
+        RFdiffusionFuture
+            A future object that can be used to retrieve the results of the design
+            job upon completion.
+        """
+        if isinstance(contigs, int):
+            contigs = f"{contigs}-{contigs}"
+        request = RFdiffusionRequest(
+            n=n,
+            contigs=contigs,
+            inpaint_seq=inpaint_seq,
+            provide_seq=provide_seq,
+            hotspot=hotspot,
+            T=T,
+            partial_T=partial_T,
+            use_active_site_model=use_active_site_model,
+            use_beta_model=use_beta_model,
+            symmetry=symmetry,
+            order=order,
+            add_potential=add_potential,
+            scaffold_target_use_struct=scaffold_target_use_struct,
+        )
+        if structure_file is not None:
+            if isinstance(structure_file, bytes):
+                structure_text = structure_file.decode()
+            elif isinstance(structure_file, str):
+                structure_text = structure_file
+            else:
+                structure_text = structure_file.read().decode()
+            request.structure_text = structure_text
+        if scaffold_target_structure_file is not None:
+            if isinstance(scaffold_target_structure_file, bytes):
+                scaffold_target_structure_text = scaffold_target_structure_file.decode()
+            elif isinstance(scaffold_target_structure_file, str):
+                scaffold_target_structure_text = scaffold_target_structure_file
+            else:
+                scaffold_target_structure_text = (
+                    scaffold_target_structure_file.read().decode()
+                )
+            request.scaffold_target_structure_text = scaffold_target_structure_text
+
+        # Submit the job via the private API function
+        job = _rfdiffusion_api_post(
+            session=self.session,
+            request=request,
+            **kwargs,
+        )
+
+        # Return the future object
+        return RFdiffusionFuture(session=self.session, job=job)
+
+    predict = generate

openprotein/models/models.py

@@ -0,0 +1,33 @@
+"""The ModelsAPI class, providing access to all protein models."""
+
+from openprotein.base import APISession
+
+from .foundation.rfdiffusion import RFdiffusionModel
+
+# In the future, we would import other models here:
+# from .foundation.esm import ESMModel
+# from .foundation.alphafold import AlphaFoldModel
+# from .custom.gp import GaussianProcessModel
+
+
+class ModelsAPI:
+    """
+    API-like accessor that groups all available protein models.
+
+    This class is attached to the main APISession and provides a single,
+    consistent entry point for accessing various models.
+    """
+
+    def __init__(self, session: APISession):
+        """
+        Initializes the ModelsAPI and attaches instances of all available models.
+
+        Args:
+            session: The active APISession to be used by the models for API calls.
+        """
+        self.rfdiffusion = RFdiffusionModel(session)
+
+        # To add new models, you would simply instantiate them here:
+        # self.esm = ESMModel(session)
+        # self.alphafold = AlphaFoldModel(session)
+        # self.gp = GaussianProcessModel(session)

openprotein/predictor/__init__.py

@@ -0,0 +1,25 @@
+"""
+Predictor module for training predictors on OpenProtein.
+
+isort:skip_file
+"""
+
+from .schemas import (
+    Kernel,
+    Constraints,
+    Features,
+    Dataset,
+    PredictorMetadata,
+    PredictorType,
+    PredictorArgs,
+    PredictJob,
+    PredictMultiJob,
+    PredictMultiSingleSiteJob,
+    PredictSingleSiteJob,
+    PredictorTrainJob,
+    PredictorEnsembleJob,
+    PredictorCVJob,
+)
+from .models import PredictorModel
+from .prediction import PredictionResultFuture
+from .predictor import PredictorAPI