molcraft 0.1.0a16__tar.gz → 0.1.0a18__tar.gz

{molcraft-0.1.0a16 → molcraft-0.1.0a18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a16
+Version: 0.1.0a18
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -43,9 +43,9 @@ Provides-Extra: gpu
 Requires-Dist: tensorflow[and-cuda]>=2.16; extra == "gpu"
 Dynamic: license-file
 
-<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo">
+<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo", width="90%">
 
-**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML. 
+**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML.
 
 > [!NOTE]  
 > In progress.
@@ -83,11 +83,12 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.BondType(),
         features.IsRotatable(),
     ],
-    super_atom=True,
+    super_node=True,
     self_loops=True,
+    include_hydrogens=False,
 )
 
-graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])
+graph = featurizer([('N[C@@H](C)C(=O)O', 2.5), ('N[C@@H](CS)C(=O)O', 1.5)])
 print(graph)
 
 model = models.GraphModel.from_layers(
@@ -95,13 +96,13 @@ model = models.GraphModel.from_layers(
         layers.Input(graph.spec),
         layers.NodeEmbedding(dim=128),
         layers.EdgeEmbedding(dim=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.Readout(mode='mean'),
-        keras.layers.Dense(units=1024, activation='relu'),
-        keras.layers.Dense(units=1024, activation='relu'),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.Readout(),
+        keras.layers.Dense(units=1024, activation='elu'),
+        keras.layers.Dense(units=1024, activation='elu'),
         keras.layers.Dense(1)
     ]
 )

{molcraft-0.1.0a16 → molcraft-0.1.0a18}/README.md

@@ -1,6 +1,6 @@
-<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo">
+<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo", width="90%">
 
-**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML. 
+**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML.
 
 > [!NOTE]  
 > In progress.
@@ -38,11 +38,12 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.BondType(),
         features.IsRotatable(),
     ],
-    super_atom=True,
+    super_node=True,
     self_loops=True,
+    include_hydrogens=False,
 )
 
-graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])
+graph = featurizer([('N[C@@H](C)C(=O)O', 2.5), ('N[C@@H](CS)C(=O)O', 1.5)])
 print(graph)
 
 model = models.GraphModel.from_layers(
@@ -50,13 +51,13 @@ model = models.GraphModel.from_layers(
         layers.Input(graph.spec),
         layers.NodeEmbedding(dim=128),
         layers.EdgeEmbedding(dim=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.Readout(mode='mean'),
-        keras.layers.Dense(units=1024, activation='relu'),
-        keras.layers.Dense(units=1024, activation='relu'),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.Readout(),
+        keras.layers.Dense(units=1024, activation='elu'),
+        keras.layers.Dense(units=1024, activation='elu'),
         keras.layers.Dense(1)
     ]
 )

{molcraft-0.1.0a16 → molcraft-0.1.0a18}/molcraft/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0a16'
+__version__ = '0.1.0a18'
 
 import os
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -6,7 +6,6 @@ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 from molcraft import chem
 from molcraft import features
 from molcraft import descriptors
-from molcraft import conformers 
 from molcraft import featurizers
 from molcraft import layers 
 from molcraft import models 
@@ -15,4 +14,6 @@ from molcraft import records
 from molcraft import tensors
 from molcraft import callbacks
 from molcraft import datasets
-from molcraft import losses
+from molcraft import losses
+
+from molcraft.applications import proteomics

molcraft-0.1.0a18/molcraft/applications/proteomics.py

@@ -0,0 +1,274 @@
+import re
+import keras
+import numpy as np
+import tensorflow as tf
+import tensorflow_text as tf_text
+from rdkit import Chem
+
+from molcraft import featurizers
+from molcraft import tensors
+from molcraft import layers
+from molcraft import models 
+from molcraft import chem
+
+"""
+
+
+
+
+
+
+No need to correct smiles for modeling, only for interpretation.
+
+Use added smiles data to rearrange list of saliency values. 
+
+
+
+
+
+
+
+
+
+
+
+
+"""
+
+# TODO: Add regex pattern for residue (C-term mod + N-term mod)?
+# TODO: Add regex pattern for residue (C-term mod + N-term mod + mod)?
+
+no_mod_pattern = r'([A-Z])'
+side_chain_mod_pattern = r'([A-Z]\[[A-Za-z0-9]+\])'
+n_term_mod_pattern = r'(\[[A-Za-z0-9]+\]-[A-Z])'
+c_term_mod_pattern = r'([A-Z]-\[[A-Za-z0-9]+\])'
+side_chain_and_n_term_mod_pattern = r'(\[[A-Za-z0-9]+\]-[A-Z]\[[A-Za-z0-9]+\])'
+side_chain_and_c_term_mod_pattern = r'([A-Z]\[[A-Za-z0-9]+\]-\[[A-Za-z0-9]+\])'
+
+residue_pattern: str = "|".join([
+    side_chain_and_n_term_mod_pattern,
+    side_chain_and_c_term_mod_pattern,
+    n_term_mod_pattern,
+    c_term_mod_pattern,
+    side_chain_mod_pattern,
+    no_mod_pattern
+])
+
+default_residues: dict[str, str] = {
+    "A": "N[C@@H](C)C(=O)O",
+    "C": "N[C@@H](CS)C(=O)O",
+    "D": "N[C@@H](CC(=O)O)C(=O)O",
+    "E": "N[C@@H](CCC(=O)O)C(=O)O",
+    "F": "N[C@@H](Cc1ccccc1)C(=O)O",
+    "G": "NCC(=O)O",
+    "H": "N[C@@H](CC1=CN=C-N1)C(=O)O",
+    "I": "N[C@@H](C(CC)C)C(=O)O",
+    "K": "N[C@@H](CCCCN)C(=O)O",
+    "L": "N[C@@H](CC(C)C)C(=O)O",
+    "M": "N[C@@H](CCSC)C(=O)O",
+    "N": "N[C@@H](CC(=O)N)C(=O)O",
+    "P": "N1[C@@H](CCC1)C(=O)O",
+    "Q": "N[C@@H](CCC(=O)N)C(=O)O",
+    "R": "N[C@@H](CCCNC(=N)N)C(=O)O",
+    "S": "N[C@@H](CO)C(=O)O",
+    "T": "N[C@@H](C(O)C)C(=O)O",
+    "V": "N[C@@H](C(C)C)C(=O)O",
+    "W": "N[C@@H](CC(=CN2)C1=C2C=CC=C1)C(=O)O",
+    "Y": "N[C@@H](Cc1ccc(O)cc1)C(=O)O",
+}
+
+def has_c_terminal_mod(residue: str):
+    if re.search(c_term_mod_pattern, residue):
+        return True 
+    return False 
+
+def has_n_terminal_mod(residue: str):
+    if re.search(n_term_mod_pattern, residue):
+        return True 
+    return False 
+
+# def register_residues(residues: dict[str, str]) -> None:
+#     for residue, smiles in residues.items():
+#         if residue.startswith('P'):
+#             smiles.startswith('N'), f'Incorrect SMILES permutation for {residue}.'
+#         elif not residue.startswith('['):
+#             smiles.startswith('N[C@@H]'), f'Incorrect SMILES permutation for {residue}.'
+#         if len(residue) > 1 and not residue[1] == "-":
+#             assert smiles.endswith('C(=O)O'), f'Incorrect SMILES permutation for {residue}.'
+#         registered_residues[residue] = smiles
+#         registered_residues[residue + '*'] = smiles.strip('O')
+    
+    
+class Peptide(chem.Mol):
+
+    @classmethod
+    def from_sequence(cls, sequence: str, **kwargs) -> 'Peptide':
+        sequence = [
+            match.group(0) for match in re.finditer(residue_pattern, sequence)
+        ]
+        peptide_smiles = []
+        for i, residue in enumerate(sequence):
+            if i < len(sequence) - 1:
+                residue_smiles = registered_residues[residue + '*']
+            else:
+                residue_smiles = registered_residues[residue]
+            peptide_smiles.append(residue_smiles)
+        peptide_smiles = ''.join(peptide_smiles)
+        return super().from_encoding(peptide_smiles, **kwargs)
+
+
+def permute_residue_smiles(smiles: str) -> str:
+    glycine = chem.Mol.from_encoding("NCC(=O)O")
+    mol = chem.Mol.from_encoding(smiles)
+    nitrogen_index = mol.GetSubstructMatch(glycine)[0]
+    permuted_smiles = Chem.MolToSmiles(
+        mol, rootedAtAtom=nitrogen_index
+    )
+    return permuted_smiles
+
+def check_peptide_residue_smiles(smiles: list[str]) -> bool:
+    backbone = 'NCC(=O)' * (len(smiles) - 1) + 'NC'
+    backbone = chem.Mol.from_encoding(backbone)
+    mol = chem.Mol.from_encoding(''.join(smiles))
+    is_valid = mol.HasSubstructMatch(backbone)
+    return is_valid
+
+@keras.saving.register_keras_serializable(package='proteomics')
+class ResidueEmbedding(keras.layers.Layer):
+
+    def __init__(
+        self, 
+        featurizer: featurizers.MolGraphFeaturizer,
+        embedder: models.GraphModel, 
+        residues: dict[str, str] | None = None,
+        **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+        if residues is None:
+            residues = {}
+        self.embedder = embedder
+        self.featurizer = featurizer
+        self.embedding_dim = int(self.embedder.output.shape[-1])
+        self.ragged_split = SequenceSplitter(pad=False)
+        self.split = SequenceSplitter(pad=True)
+        self.use_cached_embeddings = tf.Variable(False)
+        self.residues = residues
+        self.supports_masking = True
+        
+    @property
+    def residues(self) -> dict[str, str]:
+        return self._residues
+
+    @residues.setter
+    def residues(self, residues: dict[str, str]) -> None:
+        
+        residues = {**default_residues, **residues}
+        self._residues = {}
+        for residue, smiles in residues.items():
+            permuted_smiles = permute_residue_smiles(smiles)
+            # Returned smiles should begin with the amino group.
+            # It seems that the returned smiles ends with carboxyl group,
+            # though we do another check just in case. 
+            if not has_c_terminal_mod(residue):
+                carboxyl_group = 'C(=O)O'
+                if not permuted_smiles.endswith(carboxyl_group):
+                    raise ValueError(
+                        f'Unsupported permutation of {residue!r} smiles: {permuted_smiles!r}.'
+                    )
+            self._residues[residue] = permuted_smiles
+            self._residues[residue + '*'] = permuted_smiles.rstrip('O')
+
+        residue_keys = sorted(self._residues.keys())
+        residue_values = range(len(residue_keys))
+        residue_oov_value = np.where(np.array(residue_keys) == "G")[0][0]
+
+        self.mapping = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                keys=residue_keys, 
+                values=residue_values
+            ),
+            default_value=residue_oov_value,
+        )
+
+        self.graph = tf.stack([
+            self.featurizer(self._residues[r]) for r in residue_keys
+        ], axis=0)
+
+        zeros = tf.zeros((residue_values[-1] + 1, self.embedding_dim))
+        self.cached_embeddings = tf.Variable(initial_value=zeros)
+        _ = self.cache_and_get_embeddings()
+
+    def build(self, input_shape) -> None:
+        self.residues = self._residues
+        super().build(input_shape)
+
+    def call(self, sequences: tf.Tensor, training: bool = None) -> tf.Tensor:
+        if training is False:
+            self.use_cached_embeddings.assign(True)
+        else:
+            self.use_cached_embeddings.assign(False)
+        embeddings = tf.cond(
+            pred=self.use_cached_embeddings,
+            true_fn=lambda: self.cached_embeddings,
+            false_fn=lambda: self.cache_and_get_embeddings(),
+        )
+        sequences = self.ragged_split(sequences)
+        sequences = keras.ops.concatenate([
+            tf.strings.join([sequences[:, :-1], '*']), sequences[:, -1:]
+        ], axis=1)
+        indices = self.mapping.lookup(sequences)
+        return tf.gather(embeddings, indices).to_tensor()
+    
+    def cache_and_get_embeddings(self) -> tf.Tensor:
+        embeddings = self.embedder(self.graph)
+        self.cached_embeddings.assign(embeddings)
+        return embeddings
+    
+    def compute_mask(
+        self, 
+        inputs: tensors.GraphTensor, 
+        mask: bool | None = None
+    ) -> tf.Tensor | None:
+        sequences = self.split(inputs)
+        return keras.ops.not_equal(sequences, '')
+                
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            'featurizer': keras.saving.serialize_keras_object(
+                self.featurizer
+            ),
+            'embedder': keras.saving.serialize_keras_object(
+                self.embedder
+            ),
+            'residues': self._residues,
+        })
+        return config
+    
+    @classmethod
+    def from_config(cls, config: dict) -> 'ResidueEmbedding':
+        config['featurizer'] = keras.saving.deserialize_keras_object(
+            config['featurizer']
+        )
+        config['embedder'] = keras.saving.deserialize_keras_object(
+            config['embedder']
+        )
+        return super().from_config(config)
+    
+
+@keras.saving.register_keras_serializable(package='proteomics')
+class SequenceSplitter(keras.layers.Layer): 
+
+    def __init__(self, pad: bool, **kwargs):
+        super().__init__(**kwargs)
+        self.pad = pad 
+
+    def call(self, inputs: tf.Tensor) -> tf.Tensor | tf.RaggedTensor:
+        inputs = tf_text.regex_split(inputs, residue_pattern, residue_pattern)
+        if self.pad:
+            inputs = inputs.to_tensor()
+        return inputs
+    
+
+# registered_residues: dict[str, str] = {}
+# register_residues(default_residues)

{molcraft-0.1.0a16 → molcraft-0.1.0a18}/molcraft/chem.py

@@ -19,8 +19,6 @@ class Mol(Chem.Mol):
     @classmethod
     def from_encoding(cls, encoding: str, explicit_hs: bool = False, **kwargs) -> 'Mol':
         rdkit_mol = get_mol(encoding, **kwargs)
-        if not rdkit_mol:
-            return None
         if explicit_hs:
             rdkit_mol = Chem.AddHs(rdkit_mol) 
         rdkit_mol.__class__ = cls 
@@ -102,21 +100,13 @@ class Mol(Chem.Mol):
 
     def get_conformer(self, index: int = 0) -> 'Conformer':
         if self.num_conformers == 0:
-            warnings.warn(
-                'Molecule has no conformer. To embed conformer(s), invoke the `embed` method, '
-                'and optionally followed by `minimize()` to perform force field minimization.',
-                stacklevel=2
-            )
+            warnings.warn('Molecule has no conformer.')
             return None
         return Conformer.cast(self.GetConformer(index))
     
     def get_conformers(self) -> list['Conformer']:
         if self.num_conformers == 0:
-            warnings.warn(
-                'Molecule has no conformers. To embed conformers, invoke the `embed` method, '
-                'and optionally followed by `minimize()` to perform force field minimization.',
-                stacklevel=2
-            )
+            warnings.warn('Molecule has no conformer.')
             return []
         return [Conformer.cast(x) for x in self.GetConformers()]
      
@@ -222,11 +212,10 @@ def get_mol(
     else:
         mol = Chem.MolFromSmiles(encoding, sanitize=False)
     if mol is not None:
-        return sanitize_mol(mol, strict, assign_stereo_chemistry)
-    raise ValueError(
-        f"{encoding} is invalid; "
-        f"make sure {encoding} is a valid SMILES or InChI string."
-    )
+        mol = sanitize_mol(mol, strict, assign_stereo_chemistry)
+    if mol is not None:
+        return mol
+    raise ValueError(f'Could not obtain `chem.Mol` from {encoding}.')
 
 def get_adjacency_matrix(
     mol: Chem.Mol,
@@ -402,8 +391,9 @@ def embed_conformers(
     mol: Mol, 
     num_conformers: int, 
     method: str = 'ETKDGv3',
+    random_seed: int | None = None, 
     **kwargs
-) -> None:
+) -> Mol:
     available_embedding_methods = {
         'ETDG': rdDistGeom.ETDG(),
         'ETKDG': rdDistGeom.ETKDG(),
@@ -423,6 +413,9 @@ def embed_conformers(
     for key, value in kwargs.items():
         setattr(embedding_method, key, value)
 
+    if random_seed is not None:
+        embedding_method.randomSeed = random_seed
+
     success = rdDistGeom.EmbedMultipleConfs(
         mol, numConfs=num_conformers, params=embedding_method
     )
@@ -440,6 +433,8 @@ def embed_conformers(
             fallback_embedding_method.useRandomCoords = True
             fallback_embedding_method.maxAttempts = max_attempts
             fallback_embedding_method.clearConfs = False
+            if random_seed is not None:
+                fallback_embedding_method.randomSeed = random_seed
             success = rdDistGeom.EmbedMultipleConfs(
                 mol, numConfs=(num_conformers - num_successes), params=fallback_embedding_method
             )
@@ -459,7 +454,7 @@ def optimize_conformers(
     num_threads: bool = 1,
     ignore_interfragment_interactions: bool = True,
     vdw_threshold: float = 10.0,
-):
+) -> Mol:
     available_force_field_methods = [
         'MMFF', 'MMFF94', 'MMFF94s', 'UFF'
     ]
@@ -502,7 +497,7 @@ def prune_conformers(
     keep: int = 1,
     threshold: float = 0.0,
     energy_force_field: str = 'UFF',
-):
+) -> Mol:
     if mol.num_conformers == 0:
         warnings.warn(
             'Molecule has no conformers. To embed conformers, invoke the `embed` method, '
@@ -539,7 +534,7 @@ def _uff_optimize_conformers(
     vdw_threshold: float = 10.0,
     ignore_interfragment_interactions: bool = True,
     **kwargs,
-) -> Mol:
+) -> tuple[list[float], list[bool]]:
     """Universal Force Field Minimization.
     """
     results = rdForceFieldHelpers.UFFOptimizeMoleculeConfs(
@@ -560,7 +555,7 @@ def _mmff_optimize_conformers(
     variant: str = 'MMFF94',
     ignore_interfragment_interactions: bool = True,
     **kwargs,
-) -> Mol:
+) -> tuple[list[float], list[bool]]:
     """Merck Molecular Force Field Minimization.
     """
     if not rdForceFieldHelpers.MMFFHasAllMoleculeParams(mol):

{molcraft-0.1.0a16 → molcraft-0.1.0a18}/molcraft/datasets.py

@@ -11,7 +11,7 @@ def split(
     test_size: float | None = None,
     groups: str | np.ndarray = None,
     shuffle: bool = False, 
-    random_state: int | None = None,
+    random_seed: int | None = None,
 ) -> tuple[np.ndarray | pd.DataFrame, ...]:
     """Splits the dataset into subsets.
 
@@ -28,7 +28,7 @@ def split(
             The groups to perform the splitting on.
         shuffle:
             Whether the dataset should be shuffled prior to splitting.
-        random_state:
+        random_seed:
             The random state/seed. Only applicable if shuffling.
     """
     if not isinstance(data, (pd.DataFrame, np.ndarray)):
@@ -69,7 +69,7 @@ def split(
     train_size += remainder
 
     if shuffle:
-        np.random.seed(random_state)
+        np.random.seed(random_seed)
         np.random.shuffle(indices)
 
     train_mask = np.isin(groups, indices[:train_size])
@@ -84,7 +84,7 @@ def cv_split(
     num_splits: int = 10,
     groups: str | np.ndarray = None,
     shuffle: bool = False, 
-    random_state: int | None = None,
+    random_seed: int | None = None,
 ) -> typing.Iterator[
         tuple[np.ndarray | pd.DataFrame, np.ndarray | pd.DataFrame]
     ]:
@@ -99,7 +99,7 @@ def cv_split(
             The groups to perform the splitting on.
         shuffle:
             Whether the dataset should be shuffled prior to splitting.
-        random_state:
+        random_seed:
             The random state/seed. Only applicable if shuffling.
     """
     if not isinstance(data, (pd.DataFrame, np.ndarray)):
@@ -119,7 +119,7 @@ def cv_split(
             f'the data size or the number of groups ({size}).'
         )
     if shuffle:
-        np.random.seed(random_state)
+        np.random.seed(random_seed)
         np.random.shuffle(indices)
 
     indices_splits = np.array_split(indices, num_splits)

{molcraft-0.1.0a16 → molcraft-0.1.0a18}/molcraft/descriptors.py

@@ -91,3 +91,17 @@ class NumRings(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
         return rdMolDescriptors.CalcNumRings(mol) 
 
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class AtomCount(Descriptor):
+
+    def __init__(self, atom_type: str, **kwargs):
+        super().__init__(**kwargs)
+        self.atom_type = atom_type
+
+    def call(self, mol: chem.Mol) -> np.ndarray:
+        count = 0
+        for atom in mol.atoms:
+            if atom.GetSymbol() == self.atom_type:
+                count += 1
+        return count