molcraft 0.1.0a7__py3-none-any.whl → 0.1.0a9__py3-none-any.whl

molcraft/__init__.py

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

molcraft/callbacks.py

@@ -36,58 +36,65 @@ class LearningRateDecay(keras.callbacks.LearningRateScheduler):
 
 
 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, 
-        frequency: int = None,
-        tolerance: float = 0.5, 
+        self,
+        tolerance: float = 0.5,
         rollback_optimizer: bool = True,
     ):
         super().__init__()
-        self.frequency = frequency or 1_000_000_000
         self.tolerance = tolerance
         self.rollback_optimizer = rollback_optimizer
 
     def on_train_begin(self, logs=None):
-        self.rollback_weights = self._get_model_vars()
-        self.rollback_optimizer_vars = self._get_optimizer_vars()
-        self.rollback_loss = float('inf')
+        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
+        deviation = (current_loss - self._rollback_loss) / self._rollback_loss
 
         if np.isnan(current_loss) or np.isinf(current_loss):
             self._rollback()
-            print("\nRolling back model, found nan or inf loss.\n")
-            return 
-        
+            # Rolling back model because of nan or inf loss
+            return
+
         if deviation > self.tolerance:
             self._rollback()
-            print(f"\nRolling back model, found too large deviation: {deviation:.3f}\n")
-        
-        if epoch and epoch % self.frequency == 0:
-            self._rollback()
-            print(f"\nRolling back model, {epoch} % {self.frequency} == 0\n")
-            return 
-        
-        if current_loss < self.rollback_loss:
+            # 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()
+        self._rollback_loss = current_loss
+        self._rollback_weights = self._get_model_vars()
         if self.rollback_optimizer:
-            self.rollback_optimizer_vars = self._get_optimizer_vars()
+            self._rollback_optimizer_vars = self._get_optimizer_vars()
 
     def _rollback(self) -> None:
-        self.model.set_weights(self.rollback_weights)
+        self.model.set_weights(self._rollback_weights)
         if self.rollback_optimizer:
-            self.model.optimizer.set_weights(self.rollback_optimizer_vars)
+            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/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()]
@@ -425,9 +427,10 @@ def embed_conformers(
         mol, numConfs=num_conformers, params=embedding_method
     )
     if not len(success):
-        warn(
+        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.'
+            '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']:
@@ -483,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
 
@@ -496,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
     
@@ -674,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/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/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
 
@@ -272,8 +280,6 @@ class MolGraphFeaturizer(Featurizer):
                             mol.get_bond_between_atoms(atom_i, atom_j).index
                         )
                 edge['feature'] = bond_feature[bond_indices]
-                if self.self_loops:
-                    edge['self_loop'] = (edge['source'] == edge['target'])
         else:
             paths = chem.get_shortest_paths(
                 mol, radius=self.radius, self_loops=self.self_loops
@@ -284,9 +290,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 +300,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)
@@ -372,7 +374,7 @@ class MolGraphFeaturizer(Featurizer):
         num_nodes = node['feature'].shape[0]
         node = _add_super_nodes(node, num_super_nodes)
         edge = _add_super_edges(
-            edge, num_nodes, num_super_nodes, self.feature_dtype, self.index_dtype
+            edge, num_nodes, num_super_nodes, self.feature_dtype, self.index_dtype, self.self_loops
         )
         return node, edge
 
@@ -533,9 +535,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 +578,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
             
@@ -702,11 +706,15 @@ def _add_super_edges(
     num_super_nodes: int,
     feature_dtype: str,
     index_dtype: str,
+    self_loops: bool,
 ) -> dict[str, np.ndarray]:
     edge = copy.deepcopy(edge)
-    super_node_indices = (
-        np.repeat(np.arange(num_super_nodes), [num_nodes]) + num_nodes
-    )
+
+    super_node_indices = np.arange(num_super_nodes) + num_nodes
+    if self_loops:
+        edge['source'] = np.concatenate([edge['source'], super_node_indices])
+        edge['target'] = np.concatenate([edge['target'], super_node_indices])
+    super_node_indices = np.repeat(super_node_indices, [num_nodes])
     node_indices = (
         np.tile(np.arange(num_nodes), [num_super_nodes])
     )
@@ -721,6 +729,8 @@ def _add_super_edges(
     if 'feature' in edge:
         num_edges = int(edge['feature'].shape[0])
         num_super_edges = int(num_super_nodes * num_nodes * 2)
+        if self_loops:
+            num_super_edges += num_super_nodes
         edge['super'] = np.asarray(
             ([False] * num_edges + [True] * num_super_edges),
             dtype=bool
@@ -735,28 +745,8 @@ def _add_super_edges(
             ]
         )
 
-    if 'self_loop' in edge:
-        edge['self_loop'] = np.pad(
-            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