openprotein-python 0.8.8__tar.gz → 0.8.9__tar.gz

{openprotein_python-0.8.8 → openprotein_python-0.8.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openprotein-python
-Version: 0.8.8
+Version: 0.8.9
 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.8.8 → openprotein_python-0.8.9}/openprotein/fold/api.py

@@ -1,7 +1,7 @@
 """Fold REST API interface for making HTTP calls to our fold backend."""
 
 import io
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 from pydantic import TypeAdapter
@@ -12,6 +12,9 @@ from openprotein.errors import HTTPError
 
 from .schemas import FoldJob, FoldMetadata
 
+if TYPE_CHECKING:
+    import pandas as pd
+
 PATH_PREFIX = "v1/fold"
 
 
@@ -160,8 +163,8 @@ def fold_get_complex_result(
 def fold_get_complex_extra_result(
     session: APISession,
     job_id: str,
-    key: Literal["pae", "pde", "plddt", "confidence", "affinity"],
-) -> np.ndarray | list[dict]:
+    key: Literal["pae", "pde", "plddt", "confidence", "affinity", "score", "metrics"],
+) -> "np.ndarray | list[dict] | pd.DataFrame":
     """
     Get extra result for a complex from the request ID.
 
@@ -183,6 +186,10 @@ def fold_get_complex_extra_result(
         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}/complex/{key}"
@@ -194,7 +201,7 @@ def fold_get_complex_extra_result(
         if e.status_code == 400 and key == "affinity":
             raise ValueError("affinity not found for request") from None
         raise e
-    output: np.ndarray | list[dict] = formatter(response)
+    output = formatter(response)
     return output
 
 
@@ -254,34 +261,11 @@ def fold_models_post(
         sequences = kwargs["sequences"]
         # NOTE we are handling the boltz form here too
         sequences = [s.decode() if isinstance(s, bytes) else s for s in sequences]
-        body["sequences"] = sequences
-    if kwargs.get("msa_id"):
-        body["msa_id"] = kwargs["msa_id"]
-    if kwargs.get("num_recycles"):
-        body["num_recycles"] = kwargs["num_recycles"]
-    if kwargs.get("num_models"):
-        body["num_models"] = kwargs["num_models"]
-    if kwargs.get("num_relax"):
-        body["num_relax"] = kwargs["num_relax"]
-    if kwargs.get("use_potentials"):
-        body["use_potentials"] = kwargs["use_potentials"]
-    # boltz
-    if kwargs.get("diffusion_samples"):
-        body["diffusion_samples"] = kwargs["diffusion_samples"]
-    if kwargs.get("recycling_steps"):
-        body["recycling_steps"] = kwargs["recycling_steps"]
-    if kwargs.get("sampling_steps"):
-        body["sampling_steps"] = kwargs["sampling_steps"]
-    if kwargs.get("step_scale"):
-        body["step_scale"] = kwargs["step_scale"]
-    if kwargs.get("constraints"):
-        body["constraints"] = kwargs["constraints"]
-    if kwargs.get("templates"):
-        body["templates"] = kwargs["templates"]
-    if kwargs.get("properties"):
-        body["properties"] = kwargs["properties"]
-    if kwargs.get("method"):
-        body["method"] = kwargs["method"]
+        kwargs["sequences"] = sequences
+    # add non-None args - note this doesnt affect msa_id which is nested
+    for k, v in kwargs.items():
+        if v is not None:
+            body[k] = v
 
     response = session.post(endpoint, json=body)
     return FoldJob.model_validate(response.json())

{openprotein_python-0.8.8 → openprotein_python-0.8.9}/openprotein/fold/future.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
+import pandas as pd
 from pydantic.type_adapter import TypeAdapter
 from typing_extensions import Self
 
@@ -246,6 +247,8 @@ class FoldComplexResultFuture(Future):
         self._pae: np.ndarray | None = None
         self._pde: np.ndarray | None = None
         self._plddt: np.ndarray | None = None
+        self._score: pd.DataFrame | None = None
+        self._metrics: pd.DataFrame | None = None
         self._confidence: list["BoltzConfidence"] | None = None
         self._affinity: "BoltzAffinity | None" = None
 
@@ -439,6 +442,56 @@ class FoldComplexResultFuture(Future):
             self._plddt = plddt
         return self._plddt
 
+    @property
+    def score(self) -> pd.DataFrame:
+        """
+        Get the predicted scores.
+
+        Returns
+        -------
+        pd.DataFrame
+            Structure prediction scores.
+
+        Raises
+        ------
+        AttributeError
+            If score is not supported for the model.
+        """
+        if self.model_id not in {"rosettafold-3"}:
+            raise AttributeError("score not supported for non-RosettaFold model")
+        if self._score is None:
+            score = api.fold_get_complex_extra_result(
+                session=self.session, job_id=self.job.job_id, key="score"
+            )
+            assert isinstance(score, pd.DataFrame)
+            self._score = score
+        return self._score
+
+    @property
+    def metrics(self) -> pd.DataFrame:
+        """
+        Get the predicted metrics.
+
+        Returns
+        -------
+        pd.DataFrame
+            Structure prediction metrics.
+
+        Raises
+        ------
+        AttributeError
+            If metrics is not supported for the model.
+        """
+        if self.model_id not in {"rosettafold-3"}:
+            raise AttributeError("metrics not supported for non-RosettaFold model")
+        if self._metrics is None:
+            metrics = api.fold_get_complex_extra_result(
+                session=self.session, job_id=self.job.job_id, key="metrics"
+            )
+            assert isinstance(metrics, pd.DataFrame)
+            self._metrics = metrics
+        return self._metrics
+
     @property
     def confidence(self) -> list["BoltzConfidence"]:
         """