docstrings

clipzyme/lightning/clipzyme.py

@@ -14,7 +14,7 @@
 from clipzyme.utils.protein_utils import create_protein_graph
 from clipzyme.utils.loading import default_collate
 
-CHECKPOINT_URL = "https://github.com/pgmikhael/clipzyme/releases/download/v1.0.0/clipzyme_checkpoints.zip"  # TODO: Update this
+CHECKPOINT_URL = "https://github.com/pgmikhael/clipzyme/releases/download/v0.0.5/clipzyme_files.zip"  # TODO: Update this
 
 
 def download_and_extract(remote_model_url: str, local_model_dir) -> List[str]:
@@ -286,6 +286,10 @@ def extract_protein_features(
         ----------
         batch : dict
             Batch of data.
+        cif_path : Union[str, List[str]], optional
+            Path to CIF file, by default None
+        esm_dir : str, optional
+            Path to ESM model (esm2_t33_650M_UR50D), by default None
 
         Returns
         -------
@@ -322,6 +326,8 @@ def extract_reaction_features(
         ----------
         batch : dict
             Batch of data.
+        reaction : Union[str, List[str]], optional
+            Mapped raction string, by default None
 
         Returns
         -------
@@ -347,6 +353,21 @@ def extract_reaction_features(
         return model_output["hidden"]
 
     def store_in_predictions(self, preds: dict, storage_dict: dict) -> dict:
+        """
+        Store values in predictions.
+
+        Parameters
+        ----------
+        preds : dict
+            prediction dictionary
+        storage_dict : dict
+            dictionary from which to get values
+
+        Returns
+        -------
+        dict
+            updated predictions dictionary
+        """
         for key, val in storage_dict.items():
             if torch.is_tensor(val) and val.requires_grad:
                 preds[key] = val.detach()

clipzyme/utils/protein_utils.py

@@ -293,30 +293,16 @@ def get_sequences_from_structure(protein_parser, file_path):
     return sequence
 
 
-# Process SDF files and SMILES
-# def read_ligand_files():
-#     pass
-# # add ligand graphs or other protein graphs to existing graph
-# def combine_graphs():
-#     pass
-# - Parallel processing
-# #     - Complex generation
-# #     - ESM embeddings
-# #     - AF downloading
-# - Can compute intelligent splits
-#     - mmseqs split
-#     - random split
-#     - easy to add additional splitting
-
-
 def create_protein_graph(cif_path: str, esm_path: str) -> Union[pygData, None]:
     """
     Create pyg protein graph from CIF file
 
     Parameters
     ----------
-    sample : dict
-        dataset sample
+    cif_path : str
+        Path to CIF file
+    esm_path : str
+        Path to ESM model (esm2_t33_650M_UR50D.pt)
 
     Returns
     -------

clipzyme/utils/screening.py

@@ -112,6 +112,23 @@ def collect_screening_results(config_path: str) -> ScreeningResults:
 def process_mapped_reaction(
     reaction: str, bond_changes=None, use_one_hot_mol_features: bool = False
 ):
+    """
+    Process a mapped reaction string into a PyG data object
+
+    Parameters
+    ----------
+    reaction : str
+        mapped reaction string
+    bond_changes : list, optional
+        list of bond change t (i,j,t) associated with edges (i,j), by default None
+    use_one_hot_mol_features : bool, optional
+        whether to use one-hot features for molecules , by default False
+
+    Returns
+    -------
+    tuple of PyG data objects
+        reactants and products as PyG data objects
+    """
     reactants, products = reaction.split(">>")
 
     reactants, atom_map2new_index = from_mapped_smiles(

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "clipzyme"
-version = "0.0.4"
+version = "0.0.5"
 description = "Reaction-Conditioned Virtual Screening of Enzymes"
 authors = ["Peter G Mikhael <[email protected]>"]
 license = "Apache 2.0"