openprotein-python 0.11.1__tar.gz → 0.12.0__tar.gz

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.11.1
+Version: 0.12.0
 Summary: OpenProtein Python interface.
 Author-email: Mark Gee <markgee@ne47.bio>, "Timothy Truong Jr." <ttruong@ne47.bio>, Tristan Bepler <tbepler@ne47.bio>
 License-Expression: MIT

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/embeddings/__init__.py

@@ -10,6 +10,7 @@ from .openprotein import OpenProteinModel
 from .esm import ESMModel
 from .poet import PoETModel
 from .poet2 import PoET2Model
+from .ablang import AbLang2Model
 from .schemas import (
     EmbeddedSequence,
     EmbeddingsJob,

openprotein_python-0.12.0/openprotein/embeddings/ablang.py

@@ -0,0 +1,21 @@
+"""AbLang model."""
+
+from .models import EmbeddingModel
+
+
+class AbLang2Model(EmbeddingModel):
+    """
+    Community AbLang2 model that targets antibodies.
+
+    Examples
+    --------
+    View specific model details (inc supported tokens) with the `?` operator.
+
+    .. code-block:: python
+
+        >>> import openprotein
+        >>> session = openprotein.connect(username="user", password="password")
+        >>> session.embedding.ablang2?
+    """
+
+    model_id = ["ablang2"]

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/embeddings/api.py

@@ -598,9 +598,16 @@ def request_generate_post(
         body["seed"] = random_seed
     if kwargs.get("prompt_id"):
         body["prompt_id"] = kwargs["prompt_id"]
-    if kwargs.get("query_id"):
+    if kwargs.get("design_id"):
+        body["design_id"] = kwargs["design_id"]
+    query_id = kwargs.get("query_id")
+    if query_id is not None:
         assert model_id != "poet", f"Model with id {model_id} does not support query"
-        body["query_id"] = kwargs["query_id"]
+        body["query_id"] = (
+            list(query_id)
+            if isinstance(query_id, list)
+            else query_id
+        )
         if "use_query_structure_in_decoder" in kwargs:
             body["use_query_structure_in_decoder"] = kwargs[
                 "use_query_structure_in_decoder"

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/embeddings/future.py

@@ -165,10 +165,14 @@ class EmbeddingsGenerateFuture(BaseScoreFuture[Score]):
     def stream(self) -> Iterator[Score]:
         stream = api.request_get_generate_result(session=self.session, job_id=self.id)
         # name, sequence, ...
-        next(stream)  # ignore header
+        header = next(stream)
+        has_query_id = (
+            len(header) > 2 and header[-1].strip().lower() == "query_id"
+        )
         for line in stream:
             # combine scores into numpy array
-            scores = np.array([float(s) for s in line[2:]])
+            score_values = line[2:-1] if has_query_id else line[2:]
+            scores = np.array([float(s) for s in score_values])
             output = Score(name=line[0], sequence=line[1], score=scores)
             yield output
 

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/embeddings/poet2.py

@@ -20,6 +20,7 @@ from .models import EmbeddingModel
 from .poet import PoETModel
 
 if TYPE_CHECKING:
+    from openprotein.models.structure_generation import StructureGenerationFuture
     from openprotein.predictor import PredictorModel
     from openprotein.svd import SVDModel
     from openprotein.umap import UMAPModel
@@ -290,7 +291,16 @@ class PoET2Model(PoETModel, EmbeddingModel):
     def generate(
         self,
         prompt: str | Prompt | None,
-        query: str | bytes | Protein | Complex | Query | None = None,
+        query: (
+            str
+            | bytes
+            | Protein
+            | Complex
+            | Query
+            | list[str | bytes | Protein | Complex | Query]
+            | None
+        ) = None,
+        design: "str | StructureGenerationFuture | None" = None,
         use_query_structure_in_decoder: bool = True,
         num_samples: int = 100,
         temperature: float = 1.0,
@@ -308,7 +318,7 @@ class PoET2Model(PoETModel, EmbeddingModel):
         ----------
         prompt : str or Prompt or None, optional
             Prompt from an align workflow to condition PoET model.
-        query : str or bytes or Protein or Complex or Query or None, optional
+        query : str or bytes or Protein or Complex or Query or list of these or None, optional
             Query to use with prompt.
         use_query_structure_in_decoder : bool, optional
             Whether to use query structure in decoder. Default is True.
@@ -340,9 +350,14 @@ class PoET2Model(PoETModel, EmbeddingModel):
         EmbeddingsGenerateFuture
             A future object representing the status and information about the generation job.
         """
+        from openprotein.models.structure_generation import StructureGenerationFuture
+
         prompt_api = getattr(self.session, "prompt", None)
         assert isinstance(prompt_api, PromptAPI)
-        query_id = prompt_api._resolve_query(query=query)
+        query_id = prompt_api._resolve_query(query=query) if query is not None else None
+        design_id = (
+            design.job_id if isinstance(design, StructureGenerationFuture) else design
+        )
         if ensemble_weights is not None:
             # NB: for now, ensemble_method is None -> ensemble_method == "arithmetic"
             if ensemble_method is None or (ensemble_method == "arithmetic"):
@@ -364,6 +379,7 @@ class PoET2Model(PoETModel, EmbeddingModel):
             max_length=max_length,
             seed=seed,
             query_id=query_id,
+            design_id=design_id,
             use_query_structure_in_decoder=use_query_structure_in_decoder,
             ensemble_weights=ensemble_weights,
             ensemble_method=ensemble_method,

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/fold/__init__.py

@@ -1,13 +1,7 @@
 """
 Fold module for predicting structures on OpenProtein.
-
-isort:skip_file
 """
 
-from .schemas import FoldJob, FoldMetadata
-from .models import FoldModel
-from .esmfold import ESMFoldModel
-from .minifold import MiniFoldModel
 from .alphafold2 import AlphaFold2Model
 from .boltz import (
     Boltz1Model,
@@ -18,6 +12,31 @@ from .boltz import (
     BoltzConstraint,
     BoltzProperty,
 )
-from .rosettafold3 import RosettaFold3Model
-from .future import FoldResultFuture
+from .esmfold import ESMFoldModel
 from .fold import FoldAPI
+from .future import FoldResultFuture
+from .minifold import MiniFoldModel
+from .models import FoldModel
+from .protenix import ProtenixModel
+from .rosettafold3 import RosettaFold3Model
+from .schemas import FoldJob, FoldMetadata
+
+__all__ = [
+    "FoldJob",
+    "FoldMetadata",
+    "FoldModel",
+    "ESMFoldModel",
+    "MiniFoldModel",
+    "AlphaFold2Model",
+    "ProtenixModel",
+    "Boltz1Model",
+    "Boltz1xModel",
+    "Boltz2Model",
+    "BoltzAffinity",
+    "BoltzConfidence",
+    "BoltzConstraint",
+    "BoltzProperty",
+    "RosettaFold3Model",
+    "FoldResultFuture",
+    "FoldAPI",
+]

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/fold/alphafold2.py

@@ -1,15 +1,17 @@
 """Community-based AlphaFold 2 model running using ColabFold."""
 
-import io
 import warnings
-from typing import Any, Sequence
+from typing import Sequence
 
-from openprotein.align import AlignAPI, MSAFuture
+from openprotein.align import MSAFuture
 from openprotein.base import APISession
 from openprotein.common import ModelMetadata
-from openprotein.fold.common import normalize_inputs, serialize_input
-from openprotein.fold.complex import id_generator
-from openprotein.molecules import Protein, DNA, RNA, Ligand, Complex
+from openprotein.fold.common import (
+    msa_future_to_complex,
+    normalize_inputs,
+    serialize_input,
+)
+from openprotein.molecules import DNA, RNA, Complex, Ligand, Protein
 
 from . import api
 from .future import FoldResultFuture
@@ -33,7 +35,7 @@ class AlphaFold2Model(FoldModel):
 
     def fold(
         self,
-        sequences: Sequence[Complex | Protein | str] | MSAFuture | None = None,
+        sequences: Sequence[Complex | Protein | str | bytes] | MSAFuture,
         num_recycles: int | None = None,
         num_models: int = 1,
         num_relax: int = 0,
@@ -44,7 +46,7 @@ class AlphaFold2Model(FoldModel):
 
         Parameters
         ----------
-        sequences : List[Complex | Protein | str] | MSAFuture
+        sequences : Sequence[Complex | Protein | str | bytes] | MSAFuture
             List of protein sequences to include in folded output. `Protein` objects must be tagged with an `msa`, which can be a `Protein.single_sequence_mode` for single sequence mode. Alternatively, supply an `MSAFuture` to use all query sequences as a multimer.
         num_recycles : int
             number of times to recycle models
@@ -57,7 +59,6 @@ class AlphaFold2Model(FoldModel):
         -------
         job : Job
         """
-        from openprotein.align import AlignAPI
 
         if "msa" in kwargs:
             warnings.warn(
@@ -71,18 +72,7 @@ class AlphaFold2Model(FoldModel):
 
         # build the normalized_models from msa
         if isinstance(sequences, MSAFuture):
-            id_gen = id_generator()
-            align_api = getattr(self.session, "align", None)
-            assert isinstance(align_api, AlignAPI)
-            msa = sequences  # rename
-            seed = align_api.get_seed(job_id=msa.job.job_id)
-            _proteins: dict[str, Protein] = {}
-            for seq in seed.split(":"):
-                protein = Protein(sequence=seq)
-                id = next(id_gen)
-                protein.msa = msa.id
-                _proteins[id] = protein
-            normalized_complexes = [Complex(chains=_proteins)]
+            normalized_complexes = [msa_future_to_complex(self.session, sequences)]
 
         else:
             normalized_complexes = normalize_inputs(sequences)

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/fold/api.py

@@ -199,93 +199,21 @@ def fold_get_extra_result(
         The result as a numpy array (for "pae", "pde", "plddt") or a list of dictionaries (for "confidence", "affinity").
     """
     if key in {"pae", "pde", "plddt", "ptm"}:
-        formatter = lambda response: np.load(io.BytesIO(response.content))
-    elif key in {"confidence", "affinity"}:
-        formatter = lambda response: response.json()
-    elif key in {"score", "metrics"}:
-        import pandas as pd
-
-        formatter = lambda response: pd.read_csv(io.StringIO(response.content.decode()))
-    else:
-        raise ValueError(f"Unexpected key: {key}")
-    endpoint = PATH_PREFIX + f"/{job_id}/{sequence_or_index}/{key}"
-    try:
-        response = session.get(
-            endpoint,
-        )
-    except HTTPError as e:
-        if e.status_code == 400 and key == "affinity":
-            raise ValueError("affinity not found for request") from None
-        raise e
-    output = formatter(response)
-    return output
-
-
-def fold_get_complex_result(
-    session: APISession, job_id: str, format: Literal["pdb", "mmcif"]
-) -> bytes:
-    """
-    Get encoded result for a complex from the request ID.
 
-    Parameters
-    ----------
-    session : APISession
-        Session object for API communication.
-    job_id : str
-        Job ID to retrieve results from.
-    format : {'pdb', 'mmcif'}
-        Format of the result.
-
-    Returns
-    -------
-    bytes
-        Encoded result for the complex.
-    """
-    endpoint = PATH_PREFIX + f"/{job_id}/complex"
-    response = session.get(
-        endpoint,
-        params={
-            "format": format,
-        },
-    )
-    return response.content
-
-
-def fold_get_complex_extra_result(
-    session: APISession,
-    job_id: str,
-    key: Literal[
-        "pae", "pde", "plddt", "ptm", "confidence", "affinity", "score", "metrics"
-    ],
-) -> "np.ndarray | list[dict] | pd.DataFrame":
-    """
-    Get extra result for a complex from the request ID.
-
-    Parameters
-    ----------
-    session : APISession
-        Session object for API communication.
-    job_id : str
-        Job ID to retrieve results from.
-    key : {'pae', 'pde', 'plddt', 'ptm', 'confidence', 'affinity', 'score', 'metrics'}
-        The type of result to retrieve.
-
-    Returns
-    -------
-    numpy.ndarray or list of dict
-        The result as a numpy array (for "pae", "pde", "plddt") or a list of dictionaries (for "confidence", "affinity").
-    """
-    if key in {"pae", "pde", "plddt", "ptm"}:
-        formatter = lambda response: np.load(io.BytesIO(response.content))
+        def formatter(response):
+            return np.load(io.BytesIO(response.content))
     elif key in {"confidence", "affinity"}:
-        formatter = lambda response: response.json()
+
+        def formatter(response):
+            return response.json()
     elif key in {"score", "metrics"}:
         import pandas as pd
 
-        formatter = lambda response: pd.read_csv(io.StringIO(response.content.decode()))
+        def formatter(response):
+            return pd.read_csv(io.StringIO(response.content.decode()))
     else:
         raise ValueError(f"Unexpected key: {key}")
-    endpoint = PATH_PREFIX + f"/{job_id}/complex/{key}"
+    endpoint = PATH_PREFIX + f"/{job_id}/{sequence_or_index}/{key}"
     try:
         response = session.get(
             endpoint,
@@ -321,28 +249,8 @@ def fold_models_post(
         The outer list represents the batch of requests, and the inner
         list represents the complex, with each item in the list being
         an entity in that complex. A monomer would thus be a single item.
-    num_recycles : int, optional
-        Number of recycles for structure prediction.
-    num_models : int, optional
-        Number of models to generate.
-    num_relax : int, optional
-        Number of relaxation steps.
-    use_potentials : bool, optional
-        Whether to use potentials.
-    diffusion_samples : int, optional
-        Number of diffusion samples (boltz).
-    recycling_steps : int, optional
-        Number of recycling steps (boltz).
-    sampling_steps : int, optional
-        Number of sampling steps (boltz).
-    step_scale : float, optional
-        Step scale (boltz).
-    constraints : dict, optional
-        Constraints to apply.
-    templates : list, optional
-        Templates to use.
-    properties : dict, optional
-        Additional properties.
+    **kwargs
+        Additional keyword arguments to be sent with POST body.
 
     Returns
     -------

{openprotein_python-0.11.1 → openprotein_python-0.12.0}/openprotein/fold/boltz.py

@@ -1,20 +1,24 @@
 """Community-based Boltz models for complex structure prediction with ligands/dna/rna."""
 
 import warnings
-from typing import Mapping, Sequence, cast
+from typing import Mapping, Sequence
 
-from pydantic import BaseModel, Field, TypeAdapter, model_validator
+from pydantic import BaseModel, TypeAdapter, model_validator
 
-from openprotein.align import AlignAPI, MSAFuture
+from openprotein.align import MSAFuture
 from openprotein.base import APISession
 from openprotein.common import ModelMetadata
-from openprotein.fold.common import normalize_inputs, serialize_input
+from openprotein.fold.common import (
+    msa_future_to_complex,
+    normalize_inputs,
+    normalize_templates,
+    resolve_templates,
+    serialize_input,
+)
 from openprotein.molecules import Complex, Ligand, Protein
 from openprotein.molecules.template import Template
-from openprotein.prompt import PromptAPI
 
 from . import api
-from .complex import id_generator
 from .future import FoldResultFuture
 from .models import FoldModel
 
@@ -51,7 +55,7 @@ class BoltzModel(FoldModel):
         Parameters
         ----------
         sequences : Sequence[Complex | Protein | str | bytes] | MSAFuture
-            List of protein sequences to include in folded output. `Protein` objects must be tagged with an `msa`, which can be a `Protein.single_sequence_mode` for single sequence mode. Alternatively, supply an `MSAFuture` to use all query sequences as a multimer.
+            List of protein complexes to include in folded output. `Protein` objects must be tagged with an `msa`, which can be a `Protein.single_sequence_mode` for single sequence mode. Alternatively, supply an `MSAFuture` to use all query sequences as a multimer.
         diffusion_samples: int
             Number of diffusion samples to use
         num_recycles : int
@@ -85,9 +89,7 @@ class BoltzModel(FoldModel):
 
         # build the normalized_models from msa
         if isinstance(sequences, MSAFuture):
-            normalized_complexes = [
-                _msa_future_to_complex(session=self.session, msa=sequences)
-            ]
+            normalized_complexes = [msa_future_to_complex(self.session, sequences)]
         else:
             normalized_complexes = normalize_inputs(sequences)
 
@@ -135,6 +137,7 @@ class Boltz2Model(BoltzModel, FoldModel):
         templates: Sequence[Protein | Complex | Template] | None = None,
         properties: Sequence[Mapping] | None = None,
         method: str | None = None,
+        **_,
     ) -> FoldResultFuture:
         """
         Request structure prediction with Boltz-2 model.
@@ -174,83 +177,14 @@ class Boltz2Model(BoltzModel, FoldModel):
         FoldResultFuture
              Future for the folding result.
         """
-        prompt_api = getattr(self.session, "prompt", None)
-        assert isinstance(prompt_api, PromptAPI)
-
-        # validate templates
-        # mapping chain_id (to predict) to template
-        # needs to be consistent
-        templates_: list[Template] = []
-        if not isinstance(sequences, MSAFuture):
-            first_chain_id_to_template = {}
-            for batch_idx, seq in enumerate(sequences):
-                # validate templates and normalize to complex
-                if isinstance(seq, str) or isinstance(seq, bytes):
-                    seq = Protein(seq)
-                seq._assert_valid_templates()
-                if isinstance(seq, Protein):
-                    complex = Complex({"A": seq})
-                else:
-                    complex = seq
-                # resolve chain-level templates
-                for chain_id, protein in complex.get_proteins().items():
-                    # Verify same chain_id should have same templates
-                    if batch_idx == 0:
-                        first_chain_id_to_template[chain_id] = protein.templates
-                        for template in protein.templates:
-                            templates_.append(_to_template(template, chain_id=chain_id))
-                    elif first_chain_id_to_template[chain_id] != protein.templates:
-                        raise ValueError(
-                            "Expected same chain across batches to have the same templates"
-                        )
-                # resolve complex-level templates
-                if batch_idx == 0:
-                    first_templates = complex.templates
-                    for template in complex.templates:
-                        templates_.append(_to_template(template))
-                elif first_templates != complex.templates:
-                    raise ValueError(
-                        "Expected templates across complexes in batch to be the same"
-                    )
-        # method level argument
-        if templates is not None:
-            if isinstance(sequences, MSAFuture):
-                # need to convert to complex for template validation
-                sequences = [
-                    _msa_future_to_complex(session=self.session, msa=sequences)
-                ]
-            for template in templates:
-                template = _to_template(template)
-                # validate the template for all sequences before accepting it
-                for seq in sequences:
-                    if isinstance(seq, str) or isinstance(seq, bytes):
-                        seq = Protein(seq)
-                    template.validate_for_target(seq)
-                templates_.append(template)
-
-        # resolve list of Templates into expected dict arg
-        template_dicts: list[dict] = []
-        # track resolved queries to reduce network calls - use id() for identity-based caching
-        struct_id_to_query_id = {}
-
-        for template in templates_:
-            # Use id() for caching - only resolve each unique structure once
-            struct_id = id(template.template)
-            if struct_id not in struct_id_to_query_id:
-                struct_id_to_query_id[struct_id] = prompt_api._resolve_query(
-                    query=template.template
-                )
-
-            template_dict = {"query_id": struct_id_to_query_id[struct_id]}
-
-            if template.mapping is not None:
-                if isinstance(template.mapping, str):
-                    template_dict["chain_id"] = template.mapping
-                else:
-                    template_dict["chain_id"] = list(template.mapping.values())
-                    template_dict["template_id"] = list(template.mapping.keys())
-
-            template_dicts.append(template_dict)
+        template_dicts = resolve_templates(
+            session=self.session,
+            templates=normalize_templates(
+                session=self.session,
+                sequences=sequences,
+                templates=templates,
+            ),
+        )
 
         # validate properties
         if properties is not None:
@@ -302,6 +236,7 @@ class Boltz1Model(BoltzModel, FoldModel):
         step_scale: float = 1.638,
         use_potentials: bool = False,
         constraints: Sequence[Mapping] | None = None,
+        **_,
     ) -> FoldResultFuture:
         """
         Request structure prediction with Boltz-1 model.
@@ -337,12 +272,12 @@ class Boltz1Model(BoltzModel, FoldModel):
                 if "pocket" in constraint:
                     pocket_constraint = constraint["pocket"]
                     if len(pocket_constraints) > 0:
-                        msg = f"Only one pocket binders is supported in Boltz-1!"
+                        msg = "Only one pocket binders is supported in Boltz-1!"
                         raise ValueError(msg)
 
                     max_distance = constraint["pocket"].get("max_distance", 6.0)
                     if max_distance != 6.0:
-                        msg = f"Max distance != 6.0 is not supported in Boltz-1!"
+                        msg = "Max distance != 6.0 is not supported in Boltz-1!"
                         raise ValueError(msg)
                     pocket_constraints.append(pocket_constraint)
 
@@ -371,7 +306,9 @@ class Boltz1xModel(Boltz1Model, BoltzModel, FoldModel):
         num_recycles: int = 3,
         num_steps: int = 200,
         step_scale: float = 1.638,
+        use_potentials: bool = True,
         constraints: Sequence[Mapping] | None = None,
+        **_,
     ) -> FoldResultFuture:
         """
         Request structure prediction with Boltz-1x model. Uses potentials with Boltz-1 model.
@@ -396,6 +333,8 @@ class Boltz1xModel(Boltz1Model, BoltzModel, FoldModel):
         FoldResultFuture
             Future for the folding complex result.
         """
+        if not use_potentials:
+            raise ValueError("Boltz-1x uses potentials!")
 
         return super().fold(
             sequences=sequences,
@@ -582,21 +521,3 @@ class BoltzAffinity(BaseModel):
 
     class Config:
         extra = "allow"  # Allow extra fields
-
-
-def _msa_future_to_complex(session: APISession, msa: MSAFuture) -> Complex:
-    align_api = getattr(session, "align", None)
-    assert isinstance(align_api, AlignAPI)
-    seed = align_api.get_seed(job_id=msa.job.job_id)
-    proteins: dict[str, Protein] = {}
-    for chain_id, seq in zip(id_generator(), seed.split(":")):
-        protein = Protein(sequence=seq)
-        protein.msa = msa.id
-        proteins[chain_id] = protein
-    return Complex(chains=proteins)
-
-
-def _to_template(obj, chain_id: str | None = None):
-    if not isinstance(obj, Template):
-        obj = Template(template=obj, mapping=chain_id)
-    return obj