Release v0.13.1

Author: markgeejw

openprotein/data/data.py

@@ -59,6 +59,7 @@ def create(
         metadata = api.assaydata_post(
             self.session, stream, name, assay_description=description
         )
+        table.columns = [str(column).lower() for column in table.columns]
         metadata.sequence_length = len(table["sequence"].values[0])
         return AssayDataset(self.session, metadata)
 

openprotein/data/schemas.py

@@ -13,6 +13,7 @@ class AssayMetadata(BaseModel):
     num_entries: int
     measurement_names: list[str]
     sequence_length: int | None = None
+    has_name_column: bool = False
 
 
 class AssayDataRow(BaseModel):

openprotein/embeddings/embeddings.py

@@ -3,6 +3,7 @@
 from openprotein.base import APISession
 
 from . import api
+from .ablang import AbLang2Model
 from .esm import ESMModel
 from .future import EmbeddingsResultFuture
 from .models import EmbeddingModel
@@ -55,6 +56,7 @@ class EmbeddingsAPI:
     #: Rotaprot model trained on UniRef90
     rotaprot_large_uniref90_ft: OpenProteinModel
     poet_2: PoET2Model
+    ablang2: AbLang2Model
 
     #: ESM1b model
     esm1b: ESMModel  # alias

openprotein/fold/common.py

@@ -163,14 +163,11 @@ def serialize_input(session: APISession, complexes: list[Complex], needs_msa: bo
     msa_to_seed: dict[str, set[str]] = dict()
     for complex in complexes:
         _complex: list[dict[str, Any]] = []
-        for chain_ids, chain in complex.get_id_groups():
-            id_field: str | list[str] = (
-                chain_ids[0] if len(chain_ids) == 1 else list(chain_ids)
-            )
+        for chain_id, chain in complex.get_chains().items():
             if isinstance(chain, Protein):
                 # add the protein in the unified format
                 p: dict = {
-                    "id": id_field,
+                    "id": chain_id,
                     "sequence": chain.sequence.decode(),
                 }
                 if needs_msa:
@@ -208,17 +205,17 @@ def serialize_input(session: APISession, complexes: list[Complex], needs_msa: bo
                     p["msa_id"] = msa_id
                 _complex.append({"protein": p})
             elif isinstance(chain, Ligand):
-                ligand_payload: dict[str, Any] = {"id": id_field}
+                ligand_payload: dict[str, Any] = {"id": chain_id}
                 if chain.smiles is not None:
                     ligand_payload["smiles"] = chain.smiles
                 if chain.ccd is not None:
                     ligand_payload["ccd"] = chain.ccd
                 _complex.append({"ligand": ligand_payload})
             elif isinstance(chain, DNA):
-                d: dict[str, Any] = {"id": id_field, "sequence": chain.sequence}
+                d: dict[str, Any] = {"id": chain_id, "sequence": chain.sequence}
                 _complex.append({"dna": d})
             elif isinstance(chain, RNA):
-                r: dict[str, Any] = {"id": id_field, "sequence": chain.sequence}
+                r: dict[str, Any] = {"id": chain_id, "sequence": chain.sequence}
                 _complex.append({"rna": r})
             else:
                 raise ValueError(f"Unexpected chain type: {chain}")

openprotein/molecules/complex.py

@@ -24,32 +24,16 @@
 class Complex:
     def __init__(
         self,
-        chains: Mapping[
-            str | tuple[str, ...], Protein | DNA | RNA | Ligand
-        ]
-        | None = None,
+        chains: Mapping[str, Protein | DNA | RNA | Ligand] | None = None,
         name: bytes | str | None = None,
     ):
-        expanded: dict[str, Protein | DNA | RNA | Ligand] = {}
-        groups: list[tuple[str, ...]] = []
+        collected: dict[str, Protein | DNA | RNA | Ligand] = {}
         if chains is not None:
             for key, value in chains.items():
-                ids = (key,) if isinstance(key, str) else key
-                if not isinstance(ids, tuple) or not all(
-                    isinstance(cid, str) for cid in ids
-                ):
-                    raise TypeError(
-                        f"chain id must be str or tuple[str, ...]; got {key!r}"
-                    )
-                if len(ids) == 0:
-                    raise ValueError("tuple chain id must be non-empty")
-                for cid in ids:
-                    if cid in expanded:
-                        raise ValueError(f"duplicate chain id: {cid!r}")
-                    expanded[cid] = value
-                groups.append(tuple(ids))
-        self._chains = dict(sorted(expanded.items()))
-        self._id_groups: list[tuple[str, ...]] = sorted(groups, key=lambda g: g[0])
+                if not isinstance(key, str):
+                    raise TypeError(f"chain id must be str; got {key!r}")
+                collected[key] = value
+        self._chains = dict(sorted(collected.items()))
         self._templates: "Sequence[Protein | Complex | Template]" = ()
         self.name = name
 
@@ -89,17 +73,6 @@ def set_templates(
     def get_chains(self) -> Mapping[str, Protein | DNA | RNA | Ligand]:
         return MappingProxyType(self._chains)
 
-    def get_id_groups(
-        self,
-    ) -> "list[tuple[tuple[str, ...], Protein | DNA | RNA | Ligand]]":
-        """Return ordered (chain_ids, chain) pairs grouped by entity.
-
-        Each group's ``chain_ids`` is the tuple originally passed to the
-        constructor (a 1-tuple for scalar keys), and the chain object is
-        shared across all ids in the group.
-        """
-        return [(ids, self._chains[ids[0]]) for ids in self._id_groups]
-
     def get_proteins(self) -> Mapping[str, Protein]:
         return MappingProxyType(
             {k: v for k, v in self._chains.items() if isinstance(v, Protein)}
@@ -143,18 +116,8 @@ def get_ligand(self, chain_id: str) -> Ligand:
     def set_chain(
         self, chain_id: str, value: Protein | DNA | RNA | Ligand
     ) -> "Complex":
-        new_groups: list[tuple[str, ...]] = []
-        for ids in self._id_groups:
-            if chain_id in ids:
-                remaining = tuple(i for i in ids if i != chain_id)
-                if remaining:
-                    new_groups.append(remaining)
-            else:
-                new_groups.append(ids)
-        new_groups.append((chain_id,))
         self._chains[chain_id] = value
         self._chains = dict(sorted(self._chains.items()))
-        self._id_groups = sorted(new_groups, key=lambda g: g[0])
         return self
 
     def __rand__(self, left: "Complex | Protein | str") -> "Complex":
@@ -184,9 +147,6 @@ def __and__(self, right: "Complex | Protein | str") -> "Complex":
                     f"Trying to combine two sets of chains with overlapping chain ids: {overlapping_chain_ids}"
                 )
             self._chains = dict(sorted((self._chains | right._chains).items()))
-            self._id_groups = sorted(
-                self._id_groups + right._id_groups, key=lambda g: g[0]
-            )
         return self
 
     @overload
@@ -309,12 +269,9 @@ def from_string(
         )
 
     def copy(self) -> "Complex":
-        chains_copy: dict[
-            str | tuple[str, ...], Protein | DNA | RNA | Ligand
-        ] = {}
-        for ids in self._id_groups:
-            value = self._chains[ids[0]].copy()
-            chains_copy[ids if len(ids) > 1 else ids[0]] = value
+        chains_copy: dict[str, Protein | DNA | RNA | Ligand] = {
+            chain_id: chain.copy() for chain_id, chain in self._chains.items()
+        }
         return Complex(chains=chains_copy, name=self._name)
 
     def _assert_valid_templates(self):