openprotein-python 0.15.0__tar.gz → 0.16.0__tar.gz

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.15.0
+Version: 0.16.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.15.0 → openprotein_python-0.16.0}/openprotein/embeddings/models.py

@@ -17,6 +17,7 @@ from . import api
 from .future import EmbeddingsResultFuture
 
 if TYPE_CHECKING:
+    from openprotein.predictor.schemas import KernelType
     from openprotein.predictor import PredictorModel
     from openprotein.svd import SVDModel
     from openprotein.umap import UMAPModel
@@ -352,6 +353,9 @@ class EmbeddingModel:
         reduction: ReductionType,
         name: str | None = None,
         description: str | None = None,
+        kernel: "KernelType | str" = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> "PredictorModel":
         """
@@ -363,6 +367,13 @@ class EmbeddingModel:
             Assay to fit GP on.
         properties : list of str
             Properties in the assay to fit the GP on.
+        kernel : KernelType or str, optional
+            Kernel for the GP, e.g. "rbf", "matern52", "periodic",
+            "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length; only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter; only valid for "rational_quadratic" (must be > 0).
         reduction : ReductionType
             Type of embedding reduction to use for computing features. PLM must use reduction.
         name : str or None, optional
@@ -398,6 +409,9 @@ class EmbeddingModel:
             reduction=reduction,
             name=name,
             description=description,
+            kernel=kernel,
+            period=period,
+            alpha=alpha,
             **kwargs,
         )
 

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/embeddings/poet.py

@@ -17,6 +17,7 @@ from .future import (
 from .models import EmbeddingModel
 
 if TYPE_CHECKING:
+    from openprotein.predictor.schemas import KernelType
     from openprotein.predictor import PredictorModel
     from openprotein.svd import SVDModel
     from openprotein.umap import UMAPModel
@@ -432,6 +433,9 @@ class PoETModel(EmbeddingModel):
         name: str | None = None,
         description: str | None = None,
         prompt: str | Prompt | None = None,
+        kernel: "KernelType | str" = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> "PredictorModel":
         """
@@ -443,6 +447,13 @@ class PoETModel(EmbeddingModel):
             Assay to fit GP on.
         properties : list of str
             Properties in the assay to fit the GP on.
+        kernel : KernelType or str, optional
+            Kernel for the GP, e.g. "rbf", "matern52", "periodic",
+            "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length; only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter; only valid for "rational_quadratic" (must be > 0).
         prompt : str or Prompt or None, optional
             Prompt from an align workflow to condition the PoET model.
         **kwargs
@@ -464,5 +475,8 @@ class PoETModel(EmbeddingModel):
             name=name,
             description=description,
             prompt_id=prompt_id,
+            kernel=kernel,
+            period=period,
+            alpha=alpha,
             **kwargs,
         )

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/embeddings/poet2.py

@@ -21,6 +21,7 @@ from .models import EmbeddingModel
 from .poet import PoETModel
 
 if TYPE_CHECKING:
+    from openprotein.predictor.schemas import KernelType
     from openprotein.models.structure_generation import StructureGenerationFuture
     from openprotein.predictor import PredictorModel
     from openprotein.svd import SVDModel
@@ -533,6 +534,9 @@ class PoET2Model(PoETModel, EmbeddingModel):
         query: str | bytes | Protein | Complex | Query | None = None,
         use_query_structure_in_decoder: bool = True,
         decoder_type: Literal["mlm", "clm"] | None = None,
+        kernel: "KernelType | str" = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> "PredictorModel":
         """
@@ -544,6 +548,13 @@ class PoET2Model(PoETModel, EmbeddingModel):
             Assay to fit GP on.
         properties : list of str
             Properties in the assay to fit the GP on.
+        kernel : KernelType or str, optional
+            Kernel for the GP, e.g. "rbf", "matern52", "periodic",
+            "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length; only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter; only valid for "rational_quadratic" (must be > 0).
         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
@@ -570,5 +581,8 @@ class PoET2Model(PoETModel, EmbeddingModel):
             query_id=query_id,
             use_query_structure_in_decoder=use_query_structure_in_decoder,
             decoder_type=decoder_type,
+            kernel=kernel,
+            period=period,
+            alpha=alpha,
             **kwargs,
         )

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/fold/common.py

@@ -13,40 +13,28 @@ from openprotein.prompt import PromptAPI
 def normalize_inputs(
     proteins: Sequence[Complex | Protein | str | bytes],
 ):
-    # collate the id's used
-    used_ids = []
     normalized_complexes: list[Complex] = []
-    remaining_proteins: list[Protein] = []
-    remaining_protein_strings: list[str] = []
     for protein in proteins:
-        if isinstance(protein, Protein):
-            # collate these to init with id_gen
-            remaining_proteins.append(protein)
-        elif isinstance(protein, Complex):
-            used_ids.extend(list(protein.get_chains().keys()))
+        if isinstance(protein, Complex):
             normalized_complexes.append(protein)
+        elif isinstance(protein, Protein):
+            # a lone protein becomes a single-chain complex
+            id_gen = id_generator()
+            normalized_complexes.append(Complex(chains={next(id_gen): protein}))
         else:
             if isinstance(protein, bytes):
                 protein = protein.decode()
-            # handle ':'-delimited
+            # a ':'-delimited string is one complex, one chain per split sequence.
+            # chain ids only need to be unique within the complex, so use a fresh
+            # generator per complex.
+            id_gen = id_generator()
+            chains: dict[str, Protein] = {}
             for seq in protein.split(":"):
-                # collate these to init with id_gen
-                remaining_protein_strings.append(seq)
-
-    # auto generate the chain ids
-    id_gen = id_generator(used_ids)
-
-    # add the remaining proteins with id gen
-    for protein in remaining_proteins:
-        id = next(id_gen)
-        normalized_complexes.append(Complex(chains={id: protein}))
-
-    for protein_str in remaining_protein_strings:
-        id = next(id_gen)
-        protein = Protein(sequence=protein_str)
-        # protein strings default to null msa
-        protein.msa = Protein.NullMSA
-        normalized_complexes.append(Complex(chains={id: protein}))
+                chain = Protein(sequence=seq)
+                # protein strings default to null msa
+                chain.msa = Protein.NullMSA
+                chains[next(id_gen)] = chain
+            normalized_complexes.append(Complex(chains=chains))
 
     return normalized_complexes
 

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/models/__init__.py

@@ -19,6 +19,7 @@ from .foundation.boltzgen import BoltzGenFuture, BoltzGenModel
 from .foundation.esmif1 import ESMIF1Model
 from .foundation.proteinmpnn import ProteinMPNNModel
 from .foundation.rfdiffusion import RFdiffusionFuture, RFdiffusionModel
+from .foundation.solublempnn import SolubleMPNNModel
 from .models import ModelsAPI
 from .structure_generation import (
     StructureGenerationFuture,

openprotein_python-0.16.0/openprotein/models/foundation/solublempnn.py

@@ -0,0 +1,25 @@
+"""SolubleMPNN model: ProteinMPNN variant tuned for soluble protein targets."""
+
+from .proteinmpnn import ProteinMPNNModel
+
+
+class SolubleMPNNModel(ProteinMPNNModel):
+    """
+    Class for SolubleMPNN model.
+
+    Soluble-protein variant of ProteinMPNN.
+
+    Examples
+    --------
+    View specific model details (including supported tokens) with the `?` operator.
+
+    Examples
+    --------
+    .. code-block:: ipython3
+
+        >>> import openprotein
+        >>> session = openprotein.connect(username="user", password="password")
+        >>> session.models.solublempnn?
+    """
+
+    model_id = "solublempnn"

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/models/models.py

@@ -6,6 +6,7 @@ from .foundation.boltzgen import BoltzGenModel
 from .foundation.esmif1 import ESMIF1Model
 from .foundation.proteinmpnn import ProteinMPNNModel
 from .foundation.rfdiffusion import RFdiffusionModel
+from .foundation.solublempnn import SolubleMPNNModel
 
 # In the future, we would import other models here:
 # from .foundation.esm import ESMModel
@@ -30,6 +31,7 @@ class ModelsAPI:
         """
         self.rfdiffusion = RFdiffusionModel(session)
         self.proteinmpnn = ProteinMPNNModel(session)
+        self.solublempnn = SolubleMPNNModel(session)
         self.boltzgen = BoltzGenModel(session)
         self.esmif1 = ESMIF1Model(session)
         self.esm_if1 = self.esmif1  # alias mirroring the `esm-if1` model_id

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/molecules/protein.py

@@ -738,6 +738,7 @@ class Protein:
         # extract coordinates and plddt
         coordinates = np.full((len(chain_seq), _N_ATOM, 3), np.nan, dtype=np.float32)
         plddt = np.full(len(chain_seq), np.nan, dtype=np.float32)
+        max_atom_b_iso = -np.inf  # max B-factor over all atoms (for pLDDT check below)
         for residue_idx, residue in enumerate(residues):
             i = (
                 residue.label_seq - label_seq_offset
@@ -764,6 +765,8 @@ class Protein:
                         )
             for atom in residue.first_conformer():
                 atom_name = atom.name
+                if atom.b_iso != _NAN_BFACTOR_VALUE:
+                    max_atom_b_iso = max(max_atom_b_iso, atom.b_iso)
                 if residue.name == "MSE" and atom_name == "SE":
                     atom_name = "SD"
                 if (j := _ATOM_TYPE_TO_IDX.get(atom.name)) is None:
@@ -779,12 +782,12 @@ class Protein:
                 and np.isfinite(coordinates[i, _ATOM_TYPE_TO_IDX["CA"]]).all()
             ):
                 plddt[i] = 100.0
-        if (
-            maybe_use_bfactor_as_plddt
-            and not np.isnan(plddt).all()
-            and np.nanmax(plddt) <= 10
-        ):
-            plddt[~np.isnan(plddt)] = 100.0  # these were almost surely not plddts
+        if maybe_use_bfactor_as_plddt and not np.isnan(plddt).all():
+            # B-factors that are implausibly low (all <= 10) or impossible for pLDDT
+            # (any atom > 100) are almost surely not pLDDTs, so treat the structure
+            # as experimental and assign full confidence.
+            if np.nanmax(plddt) <= 10 or max_atom_b_iso > 100:
+                plddt[~np.isnan(plddt)] = 100.0
         assert np.isnan(plddt).all() or (
             (np.nanmin(plddt) >= 0) and (np.nanmax(plddt) <= 100)
         )

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/predictor/__init__.py

@@ -6,6 +6,7 @@ isort:skip_file
 
 from .schemas import (
     Kernel,
+    KernelType,
     Constraints,
     Features,
     Dataset,

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/predictor/api.py

@@ -109,6 +109,9 @@ def predictor_fit_gp_post(
     reduction: str | None = None,
     name: str | None = None,
     description: str | None = None,
+    kernel: str = "rbf",
+    period: float | None = None,
+    alpha: float | None = None,
     **kwargs,
 ) -> PredictorTrainJob:
     """
@@ -132,6 +135,13 @@ def predictor_fit_gp_post(
         Optional name of predictor model. Randomly generated if not provided.
     description: str | None
         Optional description to attach to the model.
+    kernel : str
+        Kernel to train the GP with (e.g. "rbf", "matern52", "periodic",
+        "rational_quadratic"). default = "rbf".
+    period : float | None
+        Period length for the periodic kernel. default = None.
+    alpha : float | None
+        Scale-mixture parameter for the rational_quadratic kernel. default = None.
     kwargs:
         Additional keyword arguments to be passed to foundational models, e.g. prompt_id for PoET models.
 
@@ -142,6 +152,12 @@ def predictor_fit_gp_post(
     """
     endpoint = PATH_PREFIX + "/gp"
 
+    kernel_spec: dict = {"type": kernel}
+    if period is not None:
+        kernel_spec["period"] = period
+    if alpha is not None:
+        kernel_spec["alpha"] = alpha
+
     body = {
         "dataset": {
             "assay_id": assay_id,
@@ -151,10 +167,7 @@ def predictor_fit_gp_post(
             "type": feature_type,
             "model_id": model_id,
         },
-        "kernel": {
-            "type": "rbf",
-            # "multitask": True
-        },
+        "kernel": kernel_spec,
     }
     if reduction is not None:
         body["features"]["reduction"] = reduction

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/predictor/predictor.py

@@ -13,6 +13,7 @@ from openprotein.svd import SVDAPI, SVDModel
 
 from . import api
 from .models import PredictorModel
+from .schemas import KernelType
 
 
 class PredictorAPI:
@@ -125,6 +126,9 @@ class PredictorAPI:
         reduction: Reduction | ReductionType | None = None,
         name: str | None = None,
         description: str | None = None,
+        kernel: KernelType | str = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> PredictorModel:
         """
@@ -147,6 +151,14 @@ class PredictorAPI:
             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.
+        kernel : KernelType or str, optional
+            Kernel for the GP. One of "linear", "rbf", "matern12", "matern32",
+            "matern52", "periodic", "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length, only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter, only valid for the "rational_quadratic"
+            kernel (must be > 0).
         kwargs :
             Additional keyword arguments to be passed to foundational models, e.g. prompt_id for PoET models.
 
@@ -177,6 +189,24 @@ class PredictorAPI:
             raise InvalidParameterError(
                 "Training a multitask GP is not yet supported (i.e. number of properties should only be 1 for now)"
             )
+        # kernel selection + hyperparameter validation (fail fast, client-side)
+        if not isinstance(kernel, KernelType):
+            try:
+                kernel = KernelType(kernel)
+            except ValueError as e:
+                raise InvalidParameterError(
+                    f"Unknown kernel {kernel!r}; valid options: {[k.value for k in KernelType]}"
+                ) from e
+        if period is not None:
+            if kernel != KernelType.PERIODIC:
+                raise InvalidParameterError("`period` is only valid for the 'periodic' kernel")
+            if period <= 0:
+                raise InvalidParameterError("`period` must be > 0")
+        if alpha is not None:
+            if kernel != KernelType.RATIONAL_QUADRATIC:
+                raise InvalidParameterError("`alpha` is only valid for the 'rational_quadratic' kernel")
+            if alpha <= 0:
+                raise InvalidParameterError("`alpha` must be > 0")
         # 2. Check features input
         # extract feature type
         feature_type = (
@@ -228,6 +258,9 @@ class PredictorAPI:
                 reduction=reduction,
                 name=name,
                 description=description,
+                kernel=kernel.value,
+                period=period,
+                alpha=alpha,
                 **kwargs,
             ),
         )

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/predictor/schemas.py

@@ -10,9 +10,24 @@ from openprotein.common import FeatureType
 from openprotein.jobs import Job, JobStatus, JobType
 
 
+class KernelType(str, Enum):
+    """Kernel options for a GP predictor."""
+
+    LINEAR = "linear"
+    RBF = "rbf"
+    MATERN12 = "matern12"
+    MATERN32 = "matern32"
+    MATERN52 = "matern52"
+    PERIODIC = "periodic"
+    RATIONAL_QUADRATIC = "rational_quadratic"
+
+
 class Kernel(BaseModel):
     type: str
     multitask: bool = False
+    # period applies only to the periodic kernel; alpha only to rational_quadratic
+    period: float | None = None
+    alpha: float | None = None
 
 
 class Constraints(BaseModel):

{openprotein_python-0.15.0 → openprotein_python-0.16.0}/openprotein/svd/models.py

@@ -15,6 +15,7 @@ from . import api
 from .schemas import SVDEmbeddingsJob, SVDFitJob, SVDMetadata
 
 if TYPE_CHECKING:
+    from openprotein.predictor.schemas import KernelType
     from openprotein.predictor import PredictorModel
     from openprotein.umap import UMAPModel
 
@@ -201,6 +202,9 @@ class SVDModel(Future):
         properties: list[str],
         name: str | None = None,
         description: str | None = None,
+        kernel: "KernelType | str" = "rbf",
+        period: float | None = None,
+        alpha: float | None = None,
         **kwargs,
     ) -> "PredictorModel":
         """
@@ -212,6 +216,13 @@ class SVDModel(Future):
             Assay to fit GP on. Or its assay_id.
         properties: list of str
             Properties in the assay to fit the gp on.
+        kernel : KernelType or str, optional
+            Kernel for the GP, e.g. "rbf", "matern52", "periodic",
+            "rational_quadratic". Defaults to "rbf".
+        period : float or None, optional
+            Period length; only valid for the "periodic" kernel (must be > 0).
+        alpha : float or None, optional
+            Scale-mixture parameter; only valid for "rational_quadratic" (must be > 0).
 
         Returns
         -------
@@ -254,6 +265,9 @@ class SVDModel(Future):
             model=self,
             name=name,
             description=description,
+            kernel=kernel,
+            period=period,
+            alpha=alpha,
             **kwargs,
         )