molcraft 0.1.0a6__tar.gz → 0.1.0a8__tar.gz

{molcraft-0.1.0a6 → molcraft-0.1.0a8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a6
+Version: 0.1.0a8
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -78,7 +78,7 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.IsRotatable(),
     ],
     super_atom=True,
-    self_loops=False,
+    self_loops=True,
 )
 
 graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])

{molcraft-0.1.0a6 → molcraft-0.1.0a8}/README.md

@@ -34,7 +34,7 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.IsRotatable(),
     ],
     super_atom=True,
-    self_loops=False,
+    self_loops=True,
 )
 
 graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])

{molcraft-0.1.0a6 → molcraft-0.1.0a8}/molcraft/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0a6'
+__version__ = '0.1.0a8'
 
 import os
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

molcraft-0.1.0a8/molcraft/callbacks.py

@@ -0,0 +1,100 @@
+import keras
+import warnings
+import numpy as np
+
+
+class TensorBoard(keras.callbacks.TensorBoard):
+
+    def _log_weights(self, epoch):
+        with self._train_writer.as_default():
+            for layer in self.model.layers:
+                for weight in layer.weights:
+                    # Use weight.path istead of weight.name to distinguish
+                    # weights of different layers.
+                    histogram_weight_name = weight.path + "/histogram"
+                    self.summary.histogram(
+                        histogram_weight_name, weight, step=epoch
+                    )
+                    if self.write_images:
+                        image_weight_name = weight.path + "/image"
+                        self._log_weight_as_image(
+                            weight, image_weight_name, epoch
+                        )
+            self._train_writer.flush()
+
+
+class LearningRateDecay(keras.callbacks.LearningRateScheduler):
+
+    def __init__(self, rate: float, delay: int = 0, **kwargs):
+
+        def lr_schedule(epoch: int, lr: float):
+            if epoch < delay:
+                return float(lr)
+            return float(lr * keras.ops.exp(-rate))
+        
+        super().__init__(schedule=lr_schedule, **kwargs)
+
+
+class Rollback(keras.callbacks.Callback):
+    """Rollback callback.
+
+    Currently, this callback simply restores the model and (optionally) the optimizer 
+    variables if current loss deviates too much from the best observed loss.
+
+    This callback might be useful in situations where the loss tend to spike and put 
+    the model in an undesired/problematic high-loss parameter space.
+
+    Args:
+        tolerance (float):
+            The threshold for when the restoration is triggered. The devaiation is 
+            calculated as follows: (current_loss - best_loss) / best_loss.
+    """
+
+    def __init__(
+        self,
+        tolerance: float = 0.5,
+        rollback_optimizer: bool = True,
+    ):
+        super().__init__()
+        self.tolerance = tolerance
+        self.rollback_optimizer = rollback_optimizer
+
+    def on_train_begin(self, logs=None):
+        self._rollback_weights = self._get_model_vars()
+        if self.rollback_optimizer:
+            self._rollback_optimizer_vars = self._get_optimizer_vars()
+        self._rollback_loss = float('inf')
+
+    def on_epoch_end(self, epoch: int, logs: dict = None):
+        current_loss = logs.get('val_loss', logs.get('loss'))
+        deviation = (current_loss - self._rollback_loss) / self._rollback_loss
+
+        if np.isnan(current_loss) or np.isinf(current_loss):
+            self._rollback()
+            # Rolling back model because of nan or inf loss
+            return
+
+        if deviation > self.tolerance:
+            self._rollback()
+            # Rolling back model because of large loss deviation.
+            return
+
+        if current_loss < self._rollback_loss:
+            self._save_state(current_loss)
+
+    def _save_state(self, current_loss: float) -> None:
+        self._rollback_loss = current_loss
+        self._rollback_weights = self._get_model_vars()
+        if self.rollback_optimizer:
+            self._rollback_optimizer_vars = self._get_optimizer_vars()
+
+    def _rollback(self) -> None:
+        self.model.set_weights(self._rollback_weights)
+        if self.rollback_optimizer:
+            self.model.optimizer.set_weights(self._rollback_optimizer_vars)
+
+    def _get_optimizer_vars(self):
+        return [v.numpy() for v in self.model.optimizer.variables]
+
+    def _get_model_vars(self):
+        return self.model.get_weights()

{molcraft-0.1.0a6 → molcraft-0.1.0a8}/molcraft/chem.py

@@ -102,18 +102,20 @@ class Mol(Chem.Mol):
 
     def get_conformer(self, index: int = 0) -> 'Conformer':
         if self.num_conformers == 0:
-            warn(
+            warnings.warn(
                 'Molecule has no conformer. To embed conformer(s), invoke the `embed` method, '
-                'and optionally followed by `minimize()` to perform force field minimization.'
+                'and optionally followed by `minimize()` to perform force field minimization.',
+                stacklevel=2
             )
             return None
         return Conformer.cast(self.GetConformer(index))
     
     def get_conformers(self) -> list['Conformer']:
         if self.num_conformers == 0:
-            warn(
+            warnings.warn(
                 'Molecule has no conformers. To embed conformers, invoke the `embed` method, '
-                'and optionally followed by `minimize()` to perform force field minimization.'
+                'and optionally followed by `minimize()` to perform force field minimization.',
+                stacklevel=2
             )
             return []
         return [Conformer.cast(x) for x in self.GetConformers()]
@@ -400,7 +402,6 @@ def embed_conformers(
     mol: Mol, 
     num_conformers: int, 
     method: str = 'ETKDGv3',
-    force: bool = True,
     **kwargs
 ) -> None:
     available_embedding_methods = {
@@ -411,27 +412,40 @@ def embed_conformers(
         'srETKDGv3': rdDistGeom.srETKDGv3(),
         'KDG': rdDistGeom.KDG()
     }
-    default_embedding_method = 'ETKDGv3'
     mol = Mol(mol)
-    params = available_embedding_methods.get(method)
-    if params is None:
-        warn(
-            f"Could not find `method` {method}. "
-            f"Automatically setting method to {default_embedding_method}."
+    embedding_method = available_embedding_methods.get(method)
+    if embedding_method is None:
+        raise ValueError(
+            f'Could not find `method` {method!r}. Specify either of: '
+            '`ETDG`, `ETKDG`, `ETKDGv2`, `ETKDGv3`, `srETKDGv3` or `KDG`.'
         )
-        params = available_embedding_methods[default_embedding_method]
+
     for key, value in kwargs.items():
-        setattr(params, key, value)
+        setattr(embedding_method, key, value)
 
-    success = rdDistGeom.EmbedMultipleConfs(mol, numConfs=num_conformers, params=params)
+    success = rdDistGeom.EmbedMultipleConfs(
+        mol, numConfs=num_conformers, params=embedding_method
+    )
     if not len(success):
-        warning = 'Could not embed conformer(s).'
-        if not force:
-            warn(warning)
+        warnings.warn(
+            f'Could not embed conformer(s) for {mol.canonical_smiles!r} using the '
+            'speified method. Giving it another try with more permissive methods.',
+            stacklevel=2
+        )
+        max_attempts = (20 * mol.num_atoms) # increasing it from 10xN to 20xN
+        for fallback_method in [method, 'ETDG', 'KDG']:
+            fallback_embedding_method = available_embedding_methods[fallback_method]
+            fallback_embedding_method.useRandomCoords = True
+            fallback_embedding_method.maxAttempts = max_attempts
+            success = rdDistGeom.EmbedMultipleConfs(
+                mol, numConfs=num_conformers, params=fallback_embedding_method
+            )
+            if len(success):
+                break
         else:
-            solution = ' Embedding a conformer (in 3D space) using (x, y) coordinates.'
-            warn(warning + solution)
-            rdDepictor.Compute2DCoords(mol)
+            raise RuntimeError(
+                f'Could not embed conformer(s) for {mol.canonical_smiles!r}. '
+            )
     return mol
 
 def optimize_conformers(
@@ -445,6 +459,11 @@ def optimize_conformers(
     available_force_field_methods = [
         'MMFF', 'MMFF94', 'MMFF94s', 'UFF'
     ]
+    if method not in available_force_field_methods:
+        raise ValueError(
+            f'Could not find `method` {method!r}. Specify either of: '
+            '`UFF`, `MMFF`, `MMFF94` or `MMFF94s`.'
+        )
     mol = Mol(mol)
     try:
         if method.startswith('MMFF'):
@@ -467,9 +486,10 @@ def optimize_conformers(
                 ignore_interfragment_interactions=ignore_interfragment_interactions,
             )
     except RuntimeError as e:
-        warn(
+        warnings.warn(
             f'{method} force field minimization raised {e}. '
-            '\nProceeding without force field minimization...'
+            '\nProceeding without force field minimization.',
+            stacklevel=2
         )
     return mol
 
@@ -480,9 +500,10 @@ def prune_conformers(
     energy_force_field: str = 'UFF',
 ):
     if mol.num_conformers == 0:
-        warn(
+        warnings.warn(
             'Molecule has no conformers. To embed conformers, invoke the `embed` method, '
-            'and optionally followed by `minimize()` to perform force field minimization.'
+            'and optionally followed by `minimize()` to perform force field minimization.',
+            stacklevel=2
         )
         return mol
     
@@ -658,9 +679,3 @@ def _atom_pair_fingerprint(
     fp_param = {'fpSize': size}
     return _get_fingerprint(mol, 'atom_pair', binary, dtype, **fp_param)
 
-def warn(message: str) -> None:
-    warnings.warn(
-        message=message,
-        category=UserWarning,
-        stacklevel=1,
-    )

{molcraft-0.1.0a6 → molcraft-0.1.0a8}/molcraft/conformers.py

@@ -24,19 +24,16 @@ class ConformerEmbedder(ConformerProcessor):
         self, 
         method: str = 'ETKDGv3',
         num_conformers: int = 5, 
-        force: bool = True,
         **kwargs,
     ) -> None:
         self.method = method 
         self.num_conformers = num_conformers 
-        self.force = force
         self.kwargs = kwargs 
 
     def get_config(self) -> dict:
         config = {
             'method': self.method, 
             'num_conformers': self.num_conformers, 
-            'force': self.force, 
         }
         config.update({
             k: v for (k, v) in self.kwargs.items()
@@ -48,7 +45,6 @@ class ConformerEmbedder(ConformerProcessor):
             mol, 
             method=self.method,
             num_conformers=self.num_conformers,
-            force=self.force,
             **self.kwargs,
         )
 

{molcraft-0.1.0a6 → molcraft-0.1.0a8}/molcraft/features.py

@@ -110,9 +110,10 @@ class Feature(abc.ABC):
                 'type `float`, `int`, `bool` or `None`.'
             )
         if not math.isfinite(value):
-            warn(
+            warnings.warn(
                 f'Found value of {self.name} to be non-finite. '
-                f'Value received: {value}. Converting it to a value of 0.'
+                f'Value received: {value}. Converting it to a value of 0.',
+                stacklevel=2
             )
             value = 0.0
         return np.asarray([value], dtype=self.dtype)
@@ -380,10 +381,3 @@ default_vocabulary = {
     ],
 }
 
-
-def warn(message: str) -> None:
-    warnings.warn(
-        message=message,
-        category=UserWarning,
-        stacklevel=1
-    )

{molcraft-0.1.0a6 → molcraft-0.1.0a8}/molcraft/featurizers.py

@@ -180,6 +180,12 @@ class MolGraphFeaturizer(Featurizer):
                 bond_features = [
                     features.BondType(vocab)
                 ]
+                if not default_bond_features and self.radius > 1:
+                    warnings.warn(
+                        'Replacing user-specified bond features with default bond features, '
+                        'as `radius`>1. When `radius`>1, only bond types are considered.',
+                        stacklevel=2
+                    )
         default_molecule_features = (
             molecule_features == 'auto' or molecule_features == 'default'
         )
@@ -213,9 +219,10 @@ class MolGraphFeaturizer(Featurizer):
         mol = chem.Mol.from_encoding(x, explicit_hs=self.include_hs)
 
         if mol is None:
-            warn(
+            warnings.warn(
                 f'Could not obtain `chem.Mol` from {x}. '
-                'Returning `None` (proceeding without it).'
+                'Returning `None` (proceeding without it).',
+                stacklevel=2
             )
             return None
         
@@ -245,10 +252,11 @@ class MolGraphFeaturizer(Featurizer):
 
         if molecule_feature is not None:
             if 'feature' in context:
-                warn(
+                warnings.warn(
                     'Found both inputted and computed context feature. '
                     'Overwriting inputted context feature with computed '
-                    'context feature (based on `molecule_features`).'
+                    'context feature (based on `molecule_features`).',
+                    stacklevel=2
                 )
             context['feature'] = molecule_feature
 
@@ -284,9 +292,6 @@ class MolGraphFeaturizer(Featurizer):
             edge['target'] = np.asarray(
                 [path[-1] for path in paths], dtype=self.index_dtype
             )
-            edge['length'] = np.asarray(
-                [len(path) - 1 for path in paths], dtype=self.index_dtype
-            )
             if bond_feature is not None:
                 zero_bond_feature = np.array(
                     [[1., 0., 0., 0., 0.]], dtype=bond_feature.dtype
@@ -297,7 +302,6 @@ class MolGraphFeaturizer(Featurizer):
                 edge['feature'] = self._expand_bond_features(
                     mol, paths, bond_feature,
                 )
-            edge['length'] = np.eye(self.radius + 1, dtype=self.feature_dtype)[edge['length']]
 
         if self.super_atom:
             node, edge = self._add_super_atom(node, edge)
@@ -533,9 +537,10 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
         mol = chem.Mol.from_encoding(x, explicit_hs=explicit_hs)
         
         if mol is None:
-            warn(
+            warnings.warn(
                 f'Could not obtain `chem.Mol` from {x}. '
-                'Proceeding without it.'
+                'Proceeding without it.',
+                stacklevel=2
             )
             return None
 
@@ -575,10 +580,11 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
 
         if molecule_feature is not None:
             if 'feature' in context:
-                warn(
+                warnings.warn(
                     'Found both inputted and computed context feature. '
                     'Overwriting inputted context feature with computed '
-                    'context feature (based on `molecule_features`).'
+                    'context feature (based on `molecule_features`).',
+                    stacklevel=2
                 )
             context['feature'] = molecule_feature
             
@@ -740,23 +746,9 @@ def _add_super_edges(
             edge['self_loop'], [(0, num_nodes * num_super_nodes * 2)],
             constant_values=False,
         )
-    if 'length' in edge:
-        edge['length'] = np.pad(edge['length'], [(0, 0), (1, 0)])
-        zero_array = np.zeros([num_nodes * num_super_nodes * 2], dtype='int32')
-        edge_length_dim = edge['length'].shape[1]
-        virtual_edge_length = np.eye(edge_length_dim)[zero_array]
-        edge['length'] = np.concatenate([edge['length'], virtual_edge_length])
-        edge['length'] = edge['length'].astype(feature_dtype)
 
     return edge
 
-    
-def warn(message: str) -> None:
-    warnings.warn(
-        message=message,
-        category=UserWarning,
-        stacklevel=1
-    )
 
 MolFeaturizer = MolGraphFeaturizer
 MolFeaturizer3D = MolGraphFeaturizer3D