rcsb-embedding-model 0.0.10__tar.gz → 0.0.12__tar.gz

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.10
+Version: 0.0.12
 Summary: Protein Embedding Model for Structure Search
 Project-URL: Homepage, https://github.com/rcsb/rcsb-embedding-model
 Project-URL: Issues, https://github.com/rcsb/rcsb-embedding-model/issues
@@ -18,7 +18,7 @@ Description-Content-Type: text/markdown
 
 # RCSB Embedding Model
 
-**Version** 0.0.10
+**Version** 0.0.12
 
 
 ## Overview

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/README.md

@@ -1,6 +1,6 @@
 # RCSB Embedding Model
 
-**Version** 0.0.10
+**Version** 0.0.12
 
 
 ## Overview

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.10"
+version = "0.0.12"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/cli/inference.py

@@ -5,7 +5,7 @@ import typer
 
 from rcsb_embedding_model.cli.args_utils import arg_devices
 from rcsb_embedding_model.types.api_types import StructureFormat, Accelerator, SrcLocation, SrcProteinFrom, \
-    StructureLocation, SrcAssemblyFrom
+    StructureLocation, SrcAssemblyFrom, SrcTensorFrom
 
 app = typer.Typer(
     add_completion=False
@@ -22,7 +22,7 @@ def residue_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding Name.'
+            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path or URL (switch structure-location) | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding Name.'
         )],
         output_path: Annotated[typer.FileText, typer.Option(
             exists=True,
@@ -86,7 +86,7 @@ def structure_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding Name.'
+            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path or URL (switch structure-location) | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding Name.'
         )],
         output_path: Annotated[typer.FileText, typer.Option(
             exists=True,
@@ -102,7 +102,7 @@ def structure_embedding(
             help='Use specific chains or all chains in a structure.'
         )] = SrcProteinFrom.chain,
         structure_location: Annotated[StructureLocation, typer.Option(
-            help='Source input location.'
+            help='Structure file location.'
         )] = StructureLocation.local,
         structure_format: Annotated[StructureFormat, typer.Option(
             help='Structure file format.'
@@ -154,7 +154,7 @@ def chain_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='CSV file 2 columns: Residue embedding torch tensor file | Output embedding name.'
+            help='Option 1 (src-from=file) - CSV file 2 columns: Residue Embedding Torch Tensor File | Output Embedding Name. Option 2 (src-from=structure) - CSV file 3 columns: Structure Name | Structure File Path or URL (switch structure-location) | Output Embedding Name.'
         )],
         output_path: Annotated[typer.FileText, typer.Option(
             exists=True,
@@ -163,6 +163,25 @@ def chain_embedding(
             resolve_path=True,
             help='Output path to store predictions. Embeddings are stored as csv files.'
         )],
+        res_embedding_location: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Path where residue level embeddings are located. This argument is required if src-from=structure.'
+        )] = None,
+        src_from: Annotated[SrcTensorFrom, typer.Option(
+            help='Use file names or all chains in a structure.'
+        )] = SrcTensorFrom.file,
+        structure_location: Annotated[StructureLocation, typer.Option(
+            help='Structure file location.'
+        )] = StructureLocation.local,
+        structure_format: Annotated[StructureFormat, typer.Option(
+            help='Structure file format.'
+        )] = StructureFormat.mmcif,
+        min_res_n: Annotated[int, typer.Option(
+            help='When using all chains in a structure, consider only chains with more than <min_res_n> residues.'
+        )] = 0,
         batch_size: Annotated[int, typer.Option(
             help='Number of samples processed together in one iteration.'
         )] = 1,
@@ -182,7 +201,12 @@ def chain_embedding(
     from rcsb_embedding_model.inference.chain_inference import predict
     predict(
         src_stream=src_file,
+        res_embedding_location=res_embedding_location,
         src_location=SrcLocation.local,
+        src_from=src_from,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n,
         batch_size=batch_size,
         num_workers=num_workers,
         num_nodes=num_nodes,
@@ -201,7 +225,7 @@ def assembly_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='CSV file 4 columns: Structure Name | Structure File Path | Assembly Id | Output embedding name.'
+            help='CSV file 4 columns: Structure Name | Structure File Path or URL (switch structure-location) | Assembly Id | Output embedding name.'
         )],
         res_embedding_location: Annotated[typer.FileText, typer.Option(
             exists=True,
@@ -221,7 +245,7 @@ def assembly_embedding(
             help='Use specific assembly or all assemblies in a structure.'
         )] = SrcAssemblyFrom.assembly,
         structure_location: Annotated[StructureLocation, typer.Option(
-            help='Source input location.'
+            help='Structure file location.'
         )] = StructureLocation.local,
         structure_format: Annotated[StructureFormat, typer.Option(
             help='Structure file format.'

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/dataset/esm_prot_from_structure.py

@@ -43,7 +43,7 @@ class EsmProtFromStructure(EsmProtFromChain):
         for idx, row in (pd.DataFrame(
             src_stream,
             dtype=str,
-            columns=self.COLUMNS
+            columns=EsmProtFromStructure.COLUMNS
         ) if self.src_location == SrcLocation.stream else pd.read_csv(
             src_stream,
             header=None,

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py

@@ -50,7 +50,7 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
         for idx, row in (pd.DataFrame(
                 src_stream,
                 dtype=str,
-                columns=self.COLUMNS
+                columns=ResidueAssemblyDatasetFromStructure.COLUMNS
         ) if self.src_location == SrcLocation.stream else pd.read_csv(
             src_stream,
             header=None,
@@ -60,9 +60,9 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
         )).iterrows():
             src_name = row[ResidueAssemblyDatasetFromStructure.STREAM_NAME_ATTR]
             src_structure = row[ResidueAssemblyDatasetFromStructure.STREAM_ATTR]
-            src_structure = stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure
+            structure = stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure
             item_name = row[ResidueAssemblyDatasetFromStructure.ITEM_NAME_ATTR]
-            for assembly_id in get_assemblies(src_structure=src_structure, structure_format=self.structure_format):
+            for assembly_id in get_assemblies(structure=structure, structure_format=self.structure_format):
                 assemblies.append((src_name, src_structure, str(assembly_id), f"{item_name}.{assembly_id}"))
 
         return tuple(assemblies)

rcsb_embedding_model-0.0.12/src/rcsb_embedding_model/dataset/residue_embedding_from_structure.py

@@ -0,0 +1,66 @@
+import os
+
+import pandas as pd
+
+from rcsb_embedding_model.dataset.residue_embedding_from_tensor_file import ResidueEmbeddingFromTensorFile
+from rcsb_embedding_model.types.api_types import SrcLocation, StructureLocation, StructureFormat
+from rcsb_embedding_model.utils.data import stringio_from_url
+from rcsb_embedding_model.utils.structure_parser import get_protein_chains
+from rcsb_embedding_model.utils.structure_provider import StructureProvider
+
+
+class ResidueEmbeddingFromStructure(ResidueEmbeddingFromTensorFile):
+
+    STREAM_NAME_ATTR = 'stream_name'
+    STREAM_ATTR = 'stream'
+    ITEM_NAME_ATTR = 'item_name'
+
+    COLUMNS = [STREAM_NAME_ATTR, STREAM_ATTR, ITEM_NAME_ATTR]
+
+    def __init__(
+            self,
+            src_stream,
+            res_embedding_location,
+            src_location=SrcLocation.local,
+            structure_location=StructureLocation.local,
+            structure_format=StructureFormat.mmcif,
+            min_res_n=0,
+            structure_provider=StructureProvider()
+    ):
+        if not os.path.isdir(res_embedding_location):
+            raise FileNotFoundError(f"Folder {res_embedding_location} does not exist")
+        self.res_embedding_location = res_embedding_location
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.min_res_n = min_res_n
+        self.__structure_provider = structure_provider
+        super().__init__(
+            src_stream=self.__get_chains(src_stream),
+            src_location=SrcLocation.stream
+        )
+
+    def __get_chains(self, src_stream):
+        chains = []
+        for idx, row in (pd.DataFrame(
+                src_stream,
+                dtype=str,
+                columns=ResidueEmbeddingFromStructure.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=ResidueEmbeddingFromStructure.COLUMNS
+        )).iterrows():
+            src_name = row[ResidueEmbeddingFromStructure.STREAM_NAME_ATTR]
+            src_structure = row[ResidueEmbeddingFromStructure.STREAM_ATTR]
+            item_name = row[ResidueEmbeddingFromStructure.ITEM_NAME_ATTR]
+            structure = self.__structure_provider.get_structure(
+                src_name=src_name,
+                src_structure=stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure,
+                structure_format=self.structure_format
+            )
+            for ch in get_protein_chains(structure, self.min_res_n):
+                chains.append((os.path.join(self.res_embedding_location, f"{src_name}.{ch}.pt"), f"{item_name}.{ch}"))
+        return tuple(chains)

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/dataset/residue_embedding_from_tensor_file.py

@@ -2,7 +2,7 @@ import pandas as pd
 import torch
 from torch.utils.data import Dataset
 
-from rcsb_embedding_model.types.api_types import StructureLocation, SrcLocation
+from rcsb_embedding_model.types.api_types import SrcLocation
 
 
 class ResidueEmbeddingFromTensorFile(Dataset):
@@ -26,7 +26,7 @@ class ResidueEmbeddingFromTensorFile(Dataset):
         self.data = pd.DataFrame(
             src_stream,
             dtype=str,
-            columns=self.COLUMNS
+            columns=ResidueEmbeddingFromTensorFile.COLUMNS
         ) if self.src_location == SrcLocation.stream else pd.read_csv(
             src_stream,
             header=None,

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/inference/chain_inference.py

@@ -1,16 +1,23 @@
 from torch.utils.data import DataLoader
 from lightning import Trainer
 
+from rcsb_embedding_model.dataset.residue_embedding_from_structure import ResidueEmbeddingFromStructure
 from rcsb_embedding_model.dataset.residue_embedding_from_tensor_file import ResidueEmbeddingFromTensorFile
 from rcsb_embedding_model.modules.chain_module import ChainModule
-from rcsb_embedding_model.types.api_types import Accelerator, Devices, OptionalPath, FileOrStreamTuple, SrcLocation
+from rcsb_embedding_model.types.api_types import Accelerator, Devices, OptionalPath, FileOrStreamTuple, SrcLocation, \
+    SrcTensorFrom, StructureLocation, StructureFormat
 from rcsb_embedding_model.utils.data import collate_seq_embeddings
 from rcsb_embedding_model.writer.batch_writer import CsvBatchWriter
 
 
 def predict(
         src_stream: FileOrStreamTuple,
+        res_embedding_location: OptionalPath = None,
         src_location: SrcLocation = SrcLocation.local,
+        src_from: SrcTensorFrom = SrcTensorFrom.file,
+        structure_location: StructureLocation = StructureLocation.local,
+        structure_format: StructureFormat = StructureFormat.mmcif,
+        min_res_n: int = 0,
         batch_size: int = 1,
         num_workers: int = 0,
         num_nodes: int = 1,
@@ -24,6 +31,13 @@ def predict(
         inference_set = ResidueEmbeddingFromTensorFile(
             src_stream=src_stream,
             src_location=src_location
+        ) if src_from == SrcTensorFrom.file else ResidueEmbeddingFromStructure(
+            src_stream=src_stream,
+            res_embedding_location=res_embedding_location,
+            src_location=src_location,
+            structure_location=structure_location,
+            structure_format=structure_format,
+            min_res_n=min_res_n
         )
 
     inference_dataloader = DataLoader(

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/types/api_types.py

@@ -49,3 +49,8 @@ class SrcProteinFrom(str, Enum):
 class SrcAssemblyFrom(str, Enum):
     assembly = "assembly"
     structure = "structure"
+
+
+class SrcTensorFrom(str, Enum):
+    file = "file"
+    structure = "structure"

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/utils/data.py

@@ -44,6 +44,7 @@ def stringio_from_url(url):
         print(f"Error fetching URL: {e}")
         return None
 
+
 def concatenate_tensors(file_list, max_residues, dim=0):
     """
     Concatenates a list of tensors stored in individual files along a specified dimension.

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/src/rcsb_embedding_model/utils/structure_parser.py

@@ -37,13 +37,13 @@ def get_protein_chains(structure, min_res_n=0):
     return tuple(chain_ids)
 
 
-def get_assemblies(src_structure, structure_format="mmcif"):
+def get_assemblies(structure, structure_format="mmcif"):
     if structure_format == "pdb":
-        return tuple(list_pdb_assemblies(PDBFile.read(src_structure)))
+        return tuple(list_pdb_assemblies(PDBFile.read(structure)))
     elif structure_format == "mmcif":
-        return tuple(list_assemblies(CIFFile.read(src_structure)).keys())
+        return tuple(list_assemblies(CIFFile.read(structure)).keys())
     elif structure_format == "binarycif":
-        return tuple(list_assemblies(BinaryCIFFile.read(src_structure)))
+        return tuple(list_assemblies(BinaryCIFFile.read(structure)))
     else:
         raise RuntimeError(f"Unknown file format {structure_format}")
 

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/tests/test_inference.py

@@ -2,7 +2,7 @@ import os
 import unittest
 
 from rcsb_embedding_model.types.api_types import StructureLocation, Accelerator, SrcProteinFrom, SrcLocation, \
-    StructureFormat, SrcAssemblyFrom
+    StructureFormat, SrcAssemblyFrom, SrcTensorFrom
 
 
 class TestInference(unittest.TestCase):
@@ -72,6 +72,29 @@ class TestInference(unittest.TestCase):
         self.assertEqual(tuple(chain_embeddings[3][0][0].shape), (1536,))
         self.assertEqual(tuple(chain_embeddings[4][0][0].shape), (1536,))
 
+    def test_chain_inference_from_structure(self):
+        from rcsb_embedding_model.inference.chain_inference import predict
+        chain_embeddings = predict(
+            src_stream=[
+                ("1acb", f"{self.__test_path}/resources/pdb/1acb.cif", "1acb"),
+                ("2uzi", f"{self.__test_path}/resources/pdb/2uzi.cif", "2uzi"),
+            ],
+            res_embedding_location=f"{self.__test_path}/resources/embeddings",
+            src_location=SrcLocation.stream,
+            src_from=SrcTensorFrom.structure,
+            structure_location=StructureLocation.local,
+            structure_format=StructureFormat.mmcif,
+            min_res_n=0,
+            accelerator=Accelerator.cpu
+        )
+
+        self.assertEqual(len(chain_embeddings), 5)
+        self.assertEqual(tuple(chain_embeddings[0][0][0].shape), (1536,))
+        self.assertEqual(tuple(chain_embeddings[1][0][0].shape), (1536,))
+        self.assertEqual(tuple(chain_embeddings[2][0][0].shape), (1536,))
+        self.assertEqual(tuple(chain_embeddings[3][0][0].shape), (1536,))
+        self.assertEqual(tuple(chain_embeddings[4][0][0].shape), (1536,))
+
     def test_structure_inference_from_chain(self):
         from rcsb_embedding_model.inference.structure_inference import predict
 

{rcsb_embedding_model-0.0.10 → rcsb_embedding_model-0.0.12}/tests/test_remote_inference.py

@@ -1,11 +1,14 @@
+import os.path
 import unittest
 
 from rcsb_embedding_model.types.api_types import SrcLocation, SrcProteinFrom, StructureLocation, StructureFormat, \
-    Accelerator
+    Accelerator, SrcAssemblyFrom
 
 
 class TestRemoteInference(unittest.TestCase):
 
+    __test_path = os.path.dirname(__file__)
+
     def test_esm_inference_from_structure(self):
         from rcsb_embedding_model.inference.esm_inference import predict
 
@@ -26,3 +29,20 @@ class TestRemoteInference(unittest.TestCase):
         for idx, shape in enumerate(shapes):
             self.assertEqual(tuple(esm_embeddings[idx][0][0].shape), shape)
 
+    def test_assembly_inference_from_structure(self):
+        from rcsb_embedding_model.inference.assembly_inferece import predict
+
+        assembly_embeddings = predict(
+            src_stream=[
+                ("1acb", "https://files.rcsb.org/download/1acb.cif", "1acb"),
+                ("2uzi", "https://files.rcsb.org/download/2uzi.cif", "2uzi")
+            ],
+            res_embedding_location=f"{self.__test_path}/resources/embeddings",
+            src_location=SrcLocation.stream,
+            src_from=SrcAssemblyFrom.structure,
+            structure_location=StructureLocation.remote,
+            structure_format=StructureFormat.mmcif,
+            accelerator=Accelerator.cpu
+        )
+
+        self.assertEqual(len(assembly_embeddings), 2)