rcsb-embedding-model 0.0.7__tar.gz → 0.0.9__tar.gz

{rcsb_embedding_model-0.0.7 → rcsb_embedding_model-0.0.9}/.gitignore

@@ -2,3 +2,4 @@
 /rcsb-embedding-model.iml
 /dist/
 /.pypi.rc
+__pycache__

{rcsb_embedding_model-0.0.7 → rcsb_embedding_model-0.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.7
+Version: 0.0.9
 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.7
+**Version** 0.0.9
 
 
 ## Overview
@@ -60,14 +60,17 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
 
 ### CLI
 
-    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3.
-    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --src-format mmcif --batch-size 8 --devices auto
+    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3. Predictions are stored as torch tensor files.
+    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --structure-format mmcif --batch-size 8 --devices auto
 
-    # 2. Compute structure embeddings: Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas data-frame file.
-    inference structure-embedding --src-file results/residue_embeddings.csv --output-path results/structure_embeddings --out-df-id embeddings.pkl --batch-size 4 --devices 0 --devives 1
+    # 2. Compute structure embeddings: Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas DataFrame file.
+    inference structure-embedding --src-file data/structures.csv --output-path results/residue_embeddings --out-df-name df-res-embeddings --batch-size 4 --devices 0 --devives 1
 
-    # 3. Compute chain embeddings: Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files.
-    inference chain-embedding --src-file results/residue_embeddings.csv --output-path results/chain_embeddings --batch-size 4
+    # 3. Compute chain embeddings: Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files. Predictions a re stored as csv files.
+    inference chain-embedding --src-file data/structures.csv --output-path results/chain_embeddings --batch-size 4
+
+    # 4. Compute assembly embeddings: Calculate assembly embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files.
+    inference assembly-embedding --src-file data/structures.csv --res-embedding-location results/residue_embeddings --output-path results/assembly_embeddings
 
 ### Python API
 
@@ -85,7 +88,7 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
     # Aggregate to structure-level embedding
     struct_emb = model.aggregator_embedding(res_emb)
 
-See the examples directory for complete scripts.
+See the examples and tests directories for more use cases.
 
 ---
 
@@ -117,7 +120,7 @@ The aggregation component consists of six transformer encoder layers, each with
 
 ## Citation
 
-Segura, J., Bittrich, S., et al. (2024). *Multi-scale structural similarity embedding search across entire proteomes*. bioRxiv. (Preprint: https://www.biorxiv.org/content/10.1101/2024.03.07.XXXXX)
+Segura, J., Bittrich, S., et al. (2024). *Multi-scale structural similarity embedding search across entire proteomes*. bioRxiv. (Preprint: https://www.biorxiv.org/content/10.1101/2025.02.28.640875v1)
 
 ---
 

{rcsb_embedding_model-0.0.7 → rcsb_embedding_model-0.0.9}/README.md

@@ -1,6 +1,6 @@
 # RCSB Embedding Model
 
-**Version** 0.0.7
+**Version** 0.0.9
 
 
 ## Overview
@@ -42,14 +42,17 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
 
 ### CLI
 
-    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3.
-    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --src-format mmcif --batch-size 8 --devices auto
+    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3. Predictions are stored as torch tensor files.
+    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --structure-format mmcif --batch-size 8 --devices auto
 
-    # 2. Compute structure embeddings: Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas data-frame file.
-    inference structure-embedding --src-file results/residue_embeddings.csv --output-path results/structure_embeddings --out-df-id embeddings.pkl --batch-size 4 --devices 0 --devives 1
+    # 2. Compute structure embeddings: Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas DataFrame file.
+    inference structure-embedding --src-file data/structures.csv --output-path results/residue_embeddings --out-df-name df-res-embeddings --batch-size 4 --devices 0 --devives 1
 
-    # 3. Compute chain embeddings: Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files.
-    inference chain-embedding --src-file results/residue_embeddings.csv --output-path results/chain_embeddings --batch-size 4
+    # 3. Compute chain embeddings: Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files. Predictions a re stored as csv files.
+    inference chain-embedding --src-file data/structures.csv --output-path results/chain_embeddings --batch-size 4
+
+    # 4. Compute assembly embeddings: Calculate assembly embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files.
+    inference assembly-embedding --src-file data/structures.csv --res-embedding-location results/residue_embeddings --output-path results/assembly_embeddings
 
 ### Python API
 
@@ -67,7 +70,7 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
     # Aggregate to structure-level embedding
     struct_emb = model.aggregator_embedding(res_emb)
 
-See the examples directory for complete scripts.
+See the examples and tests directories for more use cases.
 
 ---
 
@@ -99,7 +102,7 @@ The aggregation component consists of six transformer encoder layers, each with
 
 ## Citation
 
-Segura, J., Bittrich, S., et al. (2024). *Multi-scale structural similarity embedding search across entire proteomes*. bioRxiv. (Preprint: https://www.biorxiv.org/content/10.1101/2024.03.07.XXXXX)
+Segura, J., Bittrich, S., et al. (2024). *Multi-scale structural similarity embedding search across entire proteomes*. bioRxiv. (Preprint: https://www.biorxiv.org/content/10.1101/2025.02.28.640875v1)
 
 ---
 

{rcsb_embedding_model-0.0.7 → rcsb_embedding_model-0.0.9}/pyproject.toml

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

rcsb_embedding_model-0.0.9/src/rcsb_embedding_model/cli/inference.py

@@ -0,0 +1,271 @@
+import sys
+from typing import Annotated, List
+
+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
+
+app = typer.Typer(
+    add_completion=False
+)
+
+
+@app.command(
+    name="residue-embedding",
+    help="Calculate residue level embeddings of protein structures using ESM3. Predictions are stored as torch tensor files."
+)
+def residue_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            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.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions. Embeddings are stored as torch tensor files.'
+        )],
+        src_from: Annotated[SrcProteinFrom, typer.Option(
+            help='Use specific chains or all chains in a structure.'
+        )] = SrcProteinFrom.chain,
+        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,
+        num_workers: Annotated[int, typer.Option(
+            help='Number of subprocesses to use for data loading.'
+        )] = 0,
+        num_nodes: Annotated[int, typer.Option(
+            help='Number of nodes to use for inference.'
+        )] = 1,
+        accelerator: Annotated[Accelerator, typer.Option(
+            help='Device used for inference.'
+        )] = Accelerator.auto,
+        devices: Annotated[List[str], typer.Option(
+            help='The devices to use. Can be set to a positive number or "auto". Repeat this argument to indicate multiple indices of devices. "auto" for automatic selection based on the chosen accelerator.'
+        )] = tuple(['auto'])
+):
+    from rcsb_embedding_model.inference.esm_inference import predict
+    predict(
+        src_stream=src_file,
+        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,
+        accelerator=accelerator,
+        devices=arg_devices(devices),
+        out_path=output_path
+    )
+
+
+@app.command(
+    name="structure-embedding",
+    help="Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas DataFrame file."
+)
+def structure_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            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.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions. Embeddings are stored as a single DataFrame file (see out-df-name).'
+        )],
+        out_df_name: Annotated[str, typer.Option(
+            help='File name (without extension) for storing embeddings as a pandas DataFrame pickle (.pkl). The DataFrame contains 2 columns: Id | Embedding'
+        )],
+        src_from: Annotated[SrcProteinFrom, typer.Option(
+            help='Use specific chains or all chains in a structure.'
+        )] = SrcProteinFrom.chain,
+        structure_location: Annotated[StructureLocation, typer.Option(
+            help='Source input 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,
+        num_workers: Annotated[int, typer.Option(
+            help='Number of subprocesses to use for data loading.'
+        )] = 0,
+        num_nodes: Annotated[int, typer.Option(
+            help='Number of nodes to use for inference.'
+        )] = 1,
+        accelerator: Annotated[Accelerator, typer.Option(
+            help='Device used for inference.'
+        )] = Accelerator.auto,
+        devices: Annotated[List[str], typer.Option(
+            help='The devices to use. Can be set to a positive number or "auto". Repeat this argument to indicate multiple indices of devices. "auto" for automatic selection based on the chosen accelerator.'
+        )] = tuple(['auto'])
+):
+    from rcsb_embedding_model.inference.structure_inference import predict
+    predict(
+        src_stream=src_file,
+        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,
+        accelerator=accelerator,
+        devices=arg_devices(devices),
+        out_path=output_path,
+        out_df_name=out_df_name
+    )
+
+
+@app.command(
+    name="chain-embedding",
+    help="Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files."
+)
+def chain_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+            help='CSV file 2 columns: Residue embedding torch tensor file | Output embedding name.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions. Embeddings are stored as csv files.'
+        )],
+        batch_size: Annotated[int, typer.Option(
+            help='Number of samples processed together in one iteration.'
+        )] = 1,
+        num_workers: Annotated[int, typer.Option(
+            help='Number of subprocesses to use for data loading.'
+        )] = 0,
+        num_nodes: Annotated[int, typer.Option(
+            help='Number of nodes to use for inference.'
+        )] = 1,
+        accelerator: Annotated[Accelerator, typer.Option(
+            help='Device used for inference.'
+        )] = Accelerator.auto,
+        devices: Annotated[List[str], typer.Option(
+            help='The devices to use. Can be set to a positive number or "auto". Repeat this argument to indicate multiple indices of devices. "auto" for automatic selection based on the chosen accelerator.'
+        )] = tuple(['auto'])
+):
+    from rcsb_embedding_model.inference.chain_inference import predict
+    predict(
+        src_stream=src_file,
+        src_location=SrcLocation.local,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        num_nodes=num_nodes,
+        accelerator=accelerator,
+        devices=arg_devices(devices),
+        out_path=output_path
+    )
+
+@app.command(
+    name="assembly-embedding",
+    help="Calculate assembly embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files."
+)
+def assembly_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+            help='CSV file 4 columns: Structure Name | Structure File Path | Assembly Id | Output embedding name.'
+        )],
+        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 for single chains are located.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions. Embeddings are stored as csv files.'
+        )],
+        src_from: Annotated[SrcAssemblyFrom, typer.Option(
+            help='Use specific assembly or all assemblies in a structure.'
+        )] = SrcAssemblyFrom.assembly,
+        structure_location: Annotated[StructureLocation, typer.Option(
+            help='Source input location.'
+        )] = StructureLocation.local,
+        structure_format: Annotated[StructureFormat, typer.Option(
+            help='Structure file format.'
+        )] = StructureFormat.mmcif,
+        min_res_n: Annotated[int, typer.Option(
+            help='Consider only assembly chains with more than <min_res_n> residues.'
+        )] = 0,
+        max_res_n: Annotated[int, typer.Option(
+            help='Stop adding assembly chains when number of residues is greater than <max_res_n> residues.'
+        )] = sys.maxsize,
+        batch_size: Annotated[int, typer.Option(
+            help='Number of samples processed together in one iteration.'
+        )] = 1,
+        num_workers: Annotated[int, typer.Option(
+            help='Number of subprocesses to use for data loading.'
+        )] = 0,
+        num_nodes: Annotated[int, typer.Option(
+            help='Number of nodes to use for inference.'
+        )] = 1,
+        accelerator: Annotated[Accelerator, typer.Option(
+            help='Device used for inference.'
+        )] = Accelerator.auto,
+        devices: Annotated[List[str], typer.Option(
+            help='The devices to use. Can be set to a positive number or "auto". Repeat this argument to indicate multiple indices of devices. "auto" for automatic selection based on the chosen accelerator.'
+        )] = tuple(['auto'])
+):
+    from rcsb_embedding_model.inference.assembly_inferece 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,
+        max_res_n=max_res_n,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        num_nodes=num_nodes,
+        accelerator=accelerator,
+        devices=arg_devices(devices),
+        out_path=output_path
+    )
+
+
+if __name__ == "__main__":
+    app()

rcsb_embedding_model-0.0.9/src/rcsb_embedding_model/dataset/esm_prot_from_chain.py

@@ -0,0 +1,102 @@
+import argparse
+
+import torch
+from biotite.structure import chain_iter
+from esm.models.esm3 import ESM3
+from esm.sdk.api import ESMProtein, SamplingConfig
+from esm.utils.constants.models import ESM3_OPEN_SMALL
+from esm.utils.structure.protein_chain import ProteinChain
+from torch.utils.data import Dataset, DataLoader
+import pandas as pd
+
+from rcsb_embedding_model.types.api_types import StructureFormat, StructureLocation, SrcLocation
+from rcsb_embedding_model.utils.data import stringio_from_url
+from rcsb_embedding_model.utils.structure_parser import rename_atom_ch
+from rcsb_embedding_model.utils.structure_provider import StructureProvider
+
+
+class EsmProtFromChain(Dataset):
+
+    STREAM_NAME_ATTR = 'stream_name'
+    STREAM_ATTR = 'stream'
+    CH_ATTR = 'chain_id'
+    ITEM_NAME_ATTR = 'item_name'
+
+    COLUMNS = [STREAM_NAME_ATTR, STREAM_ATTR, CH_ATTR, ITEM_NAME_ATTR]
+
+    def __init__(
+        self,
+        src_stream,
+        src_location=SrcLocation.local,
+        structure_location=StructureLocation.local,
+        structure_format=StructureFormat.mmcif,
+        structure_provider=StructureProvider()
+    ):
+        super().__init__()
+        self.__structure_provider = structure_provider
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.data = pd.DataFrame()
+        self.__load_stream(src_stream)
+
+    def __load_stream(self, src_stream):
+        self.data = pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=EsmProtFromChain.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=EsmProtFromChain.COLUMNS
+        )
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        src_name = self.data.loc[idx, EsmProtFromChain.STREAM_NAME_ATTR]
+        src_structure = self.data.loc[idx, EsmProtFromChain.STREAM_ATTR]
+        chain_id = self.data.loc[idx, EsmProtFromChain.CH_ATTR]
+        item_name = self.data.loc[idx, EsmProtFromChain.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,
+            chain_id=chain_id
+        )
+        for atom_ch in chain_iter(structure):
+            protein_chain = ProteinChain.from_atomarray(rename_atom_ch(atom_ch))
+            return ESMProtein.from_protein_chain(protein_chain), item_name
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--file_list', type=argparse.FileType('r'), required=True)
+    args = parser.parse_args()
+
+    dataset = EsmProtFromChain(
+        args.file_list
+    )
+
+    esm3 = ESM3.from_pretrained(
+        ESM3_OPEN_SMALL,
+        torch.device("cpu")
+    )
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=2,
+        collate_fn=lambda _: _
+    )
+
+    for _batch in dataloader:
+        for esm_prot, prot_name in _batch:
+            protein_tensor = esm3.encode(esm_prot)
+            embeddings = esm3.forward_and_sample(
+                protein_tensor, SamplingConfig(return_per_residue_embeddings=True)
+            ).per_residue_embedding
+            print(prot_name, embeddings.shape)

rcsb_embedding_model-0.0.9/src/rcsb_embedding_model/dataset/esm_prot_from_structure.py

@@ -0,0 +1,64 @@
+
+import pandas as pd
+
+from rcsb_embedding_model.dataset.esm_prot_from_chain import EsmProtFromChain
+from rcsb_embedding_model.types.api_types import StructureLocation, StructureFormat, SrcLocation
+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 EsmProtFromStructure(EsmProtFromChain):
+
+    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,
+            src_location=SrcLocation.local,
+            structure_location=StructureLocation.local,
+            structure_format=StructureFormat.mmcif,
+            min_res_n=0,
+            structure_provider=StructureProvider()
+    ):
+        self.min_res_n = min_res_n
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.__structure_provider = structure_provider
+        super().__init__(
+            src_stream=self.__get_chains(src_stream),
+            src_location=SrcLocation.stream,
+            structure_location=structure_location,
+            structure_format=structure_format,
+            structure_provider=structure_provider
+        )
+
+    def __get_chains(self, src_stream):
+        chains = []
+        for idx, row in (pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=self.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=EsmProtFromStructure.COLUMNS
+        )).iterrows():
+            src_name = row[EsmProtFromStructure.STREAM_NAME_ATTR]
+            src_structure = row[EsmProtFromStructure.STREAM_ATTR]
+            item_name = row[EsmProtFromStructure.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((src_name, src_structure, ch, f"{item_name}.{ch}"))
+        return tuple(chains)

rcsb_embedding_model-0.0.9/src/rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py

@@ -0,0 +1,68 @@
+import sys
+
+import pandas as pd
+
+from rcsb_embedding_model.dataset.residue_assembly_embedding_from_tensor_file import ResidueAssemblyEmbeddingFromTensorFile
+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_assemblies
+from rcsb_embedding_model.utils.structure_provider import StructureProvider
+
+
+class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile):
+
+    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,
+            max_res_n=sys.maxsize,
+            structure_provider=StructureProvider()
+    ):
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.min_res_n = min_res_n
+        self.max_res_n = max_res_n
+        self.__structure_provider = structure_provider
+        super().__init__(
+            src_stream=self.__get_assemblies(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,
+            max_res_n=max_res_n,
+            structure_provider=structure_provider
+        )
+
+    def __get_assemblies(self, src_stream):
+        assemblies = []
+        for idx, row in (pd.DataFrame(
+                src_stream,
+                dtype=str,
+                columns=self.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=ResidueAssemblyDatasetFromStructure.COLUMNS
+        )).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
+            item_name = row[ResidueAssemblyDatasetFromStructure.ITEM_NAME_ATTR]
+            for assembly_id in get_assemblies(src_structure=src_structure, structure_format=self.structure_format):
+                assemblies.append((src_name, src_structure, str(assembly_id), f"{item_name}.{assembly_id}"))
+
+        return tuple(assemblies)