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

molcraft/__init__.py

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

molcraft/callbacks.py

@@ -1,4 +1,6 @@
 import keras
+import warnings
+import numpy as np
 
 
 class TensorBoard(keras.callbacks.TensorBoard):
@@ -31,3 +33,61 @@ class LearningRateDecay(keras.callbacks.LearningRateScheduler):
             return float(lr * keras.ops.exp(-rate))
         
         super().__init__(schedule=lr_schedule, **kwargs)
+
+
+class Rollback(keras.callbacks.Callback):
+
+    def __init__(
+        self, 
+        frequency: int = None,
+        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')
+
+    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()
+            print("\nRolling back model, found nan or inf loss.\n")
+            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:
+            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/chem.py

@@ -400,7 +400,6 @@ def embed_conformers(
     mol: Mol, 
     num_conformers: int, 
     method: str = 'ETKDGv3',
-    force: bool = True,
     **kwargs
 ) -> None:
     available_embedding_methods = {
@@ -411,27 +410,39 @@ 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)
+        warn(
+            f'Could not embed conformer(s) for {mol.canonical_smiles!r} using the '
+            'speified method. Giving it another try with more permissive methods.'
+        )
+        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 +456,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'):
@@ -469,7 +485,7 @@ def optimize_conformers(
     except RuntimeError as e:
         warn(
             f'{method} force field minimization raised {e}. '
-            '\nProceeding without force field minimization...'
+            '\nProceeding without force field minimization.'
         )
     return mol
 

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/layers.py

@@ -1231,7 +1231,7 @@ class EGConv3D(GraphConv):
     def __init__(
         self, 
         units: int = 128, 
-        activation: keras.layers.Activation | str | None = None, 
+        activation: keras.layers.Activation | str | None = 'silu', 
         use_bias: bool = True,
         normalize: bool = False,
         **kwargs
@@ -1251,31 +1251,52 @@ class EGConv3D(GraphConv):
                 'which is required for Conv3D layers.'
             )
         self._has_edge_feature = 'feature' in spec.edge
-        self.message_fn = self.get_dense(self.units, activation=self._activation)
-        self.dense_position = self.get_dense(1, use_bias=False, kernel_initializer='zeros')
+        self._message_feedforward_intermediate = self.get_dense(
+            self.units, activation=self._activation
+        )
+        self._message_feedforward_final = self.get_dense(
+            self.units, activation=self._activation
+        )
+
+        self._coord_feedforward_intermediate = self.get_dense(
+            self.units, activation=self._activation
+        ) 
+        self._coord_feedforward_final = self.get_dense(
+            1, use_bias=False, activation='tanh'
+        )
 
         has_overridden_update = self.__class__.update != EGConv3D.update 
         if not has_overridden_update:
-            self.update_fn = self.get_dense(self.units, activation=self._activation)
-            self.output_dense = self.get_dense(self.units)
+            self._feedforward_intermediate = self.get_dense(
+                self.units, activation=self._activation
+            )
+            self._feedforward_output = self.get_dense(self.units)
 
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         relative_node_coordinate = keras.ops.subtract(
             tensor.gather('coordinate', 'target'), 
             tensor.gather('coordinate', 'source')
         ) 
-        euclidean_distance = keras.ops.sum(
-            keras.ops.square(
-                relative_node_coordinate
-            ), 
+        squared_distance = keras.ops.sum(
+            keras.ops.square(relative_node_coordinate), 
             axis=-1, 
             keepdims=True
         )
+    
+        # For numerical stability (i.e., to prevent NaN losses), this implementation of `EGConv3D` 
+        # either needs to apply a `tanh` activation to the output of `self._coord_feedforward_final`, 
+        # or normalize `relative_node_cordinate` as follows:
+        #
+        # norm = keras.ops.sqrt(squared_distance) + keras.backend.epsilon()
+        # relative_node_coordinate /= norm
+        #
+        # For now, this implementation does the former.
+
         feature = keras.ops.concatenate(
             [
                 tensor.gather('feature', 'target'), 
                 tensor.gather('feature', 'source'), 
-                euclidean_distance, 
+                squared_distance, 
             ], 
             axis=-1
         )
@@ -1287,10 +1308,15 @@ class EGConv3D(GraphConv):
                 ], 
                 axis=-1
             )
-        message = self.message_fn(feature)
+        message = self._message_feedforward_final(
+            self._message_feedforward_intermediate(feature)
+        )
+
         relative_node_coordinate = keras.ops.multiply(
-            relative_node_coordinate,
-            self.dense_position(message)
+            relative_node_coordinate, 
+            self._coord_feedforward_final(
+                self._coord_feedforward_intermediate(message)
+            )
         )
         return tensor.update(
             {
@@ -1302,26 +1328,26 @@ class EGConv3D(GraphConv):
         )
 
     def aggregate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        # coefficient = keras.ops.bincount(
-        #     tensor.edge['source'], 
-        #     minlength=tensor.num_nodes
-        # )
-        # coefficient = keras.ops.cast(
-        #     coefficient, tensor.node['coordinate'].dtype
-        # )
-        # coefficient = keras.ops.expand_dims(
-        #     keras.ops.divide_no_nan(1, coefficient), axis=1
-        # )  
-
-        updated_coordinate = tensor.aggregate('relative_node_coordinate', mode='mean')# * coefficient
-        updated_coordinate += tensor.node['coordinate']
-
+        coordinate = tensor.node['coordinate']
+        coordinate += tensor.aggregate('relative_node_coordinate', mode='mean')
+
+        # Original implementation seems to apply sum aggregation, which does not
+        # seem work well for this implementation of `EGConv3D`, as it causes 
+        # large output values and large initial losses. The magnitude of the 
+        # aggregated values of a sum aggregation depends on the number of 
+        # neighbors, which may be many and may differ from node to node (or 
+        # graph to graph). Therefore, a mean mean aggregation is performed 
+        # instead:
         aggregate = tensor.aggregate('message', mode='mean')
+
+        # Simply added to silence warning ('no gradients for variables ...')
+        aggregate += (0.0 * keras.ops.sum(coordinate))
+
         return tensor.update(
             {
                 'node': {
                     'aggregate': aggregate, 
-                    'coordinate': updated_coordinate,
+                    'coordinate': coordinate,
                 },
                 'edge': {
                     'message': None,
@@ -1331,16 +1357,16 @@ class EGConv3D(GraphConv):
         ) 
 
     def update(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
-        updated_node_feature = self.update_fn(
-            keras.ops.concatenate(
-                [
-                    tensor.node['aggregate'], 
-                    tensor.node['feature']
-                ], 
-                axis=-1
-            )
+        feature = keras.ops.concatenate(
+            [
+                tensor.node['aggregate'], 
+                tensor.node['feature']
+            ], 
+            axis=-1
+        )
+        updated_node_feature = self._feedforward_output(
+            self._feedforward_intermediate(feature)
         )  
-        updated_node_feature = self.output_dense(updated_node_feature)
         return tensor.update(
             {
                 'node': {
@@ -1694,8 +1720,8 @@ class EdgeEmbedding(GraphLayer):
             mask = keras.ops.expand_dims(mask, -1)
             feature = keras.ops.where(mask, self._mask_feature, feature)
         elif self._allow_masking:
-            # Slience warning of 'no gradients for variables'
-            feature = feature + (self._mask_feature * 0.0)
+            # Simply added to silence warning ('no gradients for variables ...')
+            feature += (0.0 * self._mask_feature)
 
         if self._normalize:
             feature = self._norm(feature)
@@ -1999,14 +2025,11 @@ def Input(spec: tensors.GraphTensor.Spec) -> dict:
     for outer_field, data in spec.__dict__.items():
         inputs[outer_field] = {}
         for inner_field, nested_spec in data.items():
-            if inner_field in ['label', 'weight']:
+            if outer_field == 'context' and inner_field in ['label', 'weight']:
                 # Remove context label and weight from the symbolic input 
                 # as a functional model is strict for what input can be passed.
-                # We want to be able to pass a graph with or without labels 
-                # and sample weights. The __call__ method of the `GraphModel`
-                # temporarily pops label and weight to avoid errors. 
-                if outer_field == 'context':
-                    continue
+                # (We want to train and predict with the model.)
+                continue
             kwargs = {
                 'shape': nested_spec.shape[1:],
                 'dtype': nested_spec.dtype,

molcraft/models.py

@@ -415,6 +415,10 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
     def predict_step(self, tensor: tensors.GraphTensor) -> np.ndarray:
         return self(tensor, training=False)
 
+    def compute_loss(self, x, y, y_pred, sample_weight=None):
+        y, y_pred, sample_weight = _maybe_reshape(y, y_pred, sample_weight)
+        return super().compute_loss(x, y, y_pred, sample_weight)
+        
     def compute_metrics(self, x, y, y_pred, sample_weight=None) -> dict[str, float]:
         loss = self.compute_loss(x, y, y_pred, sample_weight)
         metric_results = {}
@@ -423,7 +427,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
                 metric.update_state(loss)
                 metric_results[metric.name] = metric.result()
             else:
-                metric.update_state(y, y_pred)
+                metric.update_state(y, y_pred, sample_weight=sample_weight)
                 metric_results.update(metric.result())
         return metric_results
     
@@ -593,3 +597,14 @@ def _make_dataset(x: tensors.GraphTensor, batch_size: int):
         .batch(batch_size)
         .prefetch(-1)
     )
+
+def _maybe_reshape(y, y_pred, sample_weight):
+    if (
+        sample_weight is not None and 
+        len(keras.ops.shape(sample_weight)) == 2 and 
+        sample_weight.shape == y_pred.shape
+    ):
+        y = keras.ops.reshape(y, [-1])
+        y_pred = keras.ops.reshape(y_pred, [-1])
+        sample_weight = keras.ops.reshape(sample_weight, [-1])
+    return y, y_pred, sample_weight

molcraft/ops.py

@@ -67,7 +67,7 @@ def edge_softmax(
     edge_target: tf.Tensor
 ) -> tf.Tensor:
     num_segments = keras.ops.cond(
-        keras.ops.shape(edge_target)[0] > 0,
+        keras.ops.greater(keras.ops.shape(edge_target)[0], 0),
         lambda: keras.ops.maximum(keras.ops.max(edge_target) + 1, 1),
         lambda: 0
     )
@@ -100,7 +100,7 @@ def segment_mean(
 ) -> tf.Tensor:
     if num_segments is None:
         num_segments = keras.ops.cond(
-            keras.ops.shape(segment_ids)[0] > 0,
+            keras.ops.greater(keras.ops.shape(segment_ids)[0], 0),
             lambda: keras.ops.max(segment_ids) + 1,
             lambda: 0
         )

{molcraft-0.1.0a6.dist-info → molcraft-0.1.0a7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a6
+Version: 0.1.0a7
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License

molcraft-0.1.0a7.dist-info/RECORD

@@ -0,0 +1,19 @@
+molcraft/__init__.py,sha256=mz_CINPSZxirbRGuhKlGwTh7dXEF4IXDxLrzogvLAyU,463
+molcraft/callbacks.py,sha256=h6Kao830TiPfRGomi7tR__cqRcEtkOwnXJsRpDlqPts,3276
+molcraft/chem.py,sha256=H_7ZiWZcZjYIY6xXkxeDllHAyPaTTArFoHCb781LtEw,21447
+molcraft/conformers.py,sha256=K6ZtiSUNDN_fwqGP9JrPcwALLFFvlMlF_XejEJH3Sr4,4205
+molcraft/datasets.py,sha256=rFgXTC1ZheLhfgQgcCspP_wEE54a33PIneH7OplbS-8,4047
+molcraft/descriptors.py,sha256=gKqlJ3BqJLTeR2ft8isftSEaJDC8cv64eTq5IYhy4XM,3032
+molcraft/features.py,sha256=69oV_GHNdBKPA4sp6Tpo6brvNmaauk_IVIzNjX7VDmg,13648
+molcraft/featurizers.py,sha256=aJJibnHCxvSu3bNbE2xQk34QvFb47Mnm__0MxlRLA0w,27323
+molcraft/layers.py,sha256=zdDckHz4iE4mqyWVHmzfzH-b8WluGzT98CZuLqFY8Bo,74408
+molcraft/losses.py,sha256=JEKZEX2f8vDgky_fUocsF8vZjy9VMzRjZUBa20Uf9Qw,1065
+molcraft/models.py,sha256=FLXpO3OUmRxLmxG3MjBK4ZwcVFlea1gqEgs1ibKly2w,23263
+molcraft/ops.py,sha256=D19aTEdVmgB4utSOAqm-4fLnKwiW7hqTl-6dxjEnsSM,4477
+molcraft/records.py,sha256=0sjOdcr266ZER4F-aTBQ3AVPNAwflKWNiNJVsSc1-PQ,5370
+molcraft/tensors.py,sha256=EOUKx496KUZsjA1zA2ABc7tU_TW3Jv7AXDsug_QsLbA,22407
+molcraft-0.1.0a7.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a7.dist-info/METADATA,sha256=PGq6jEM8RMaCqixZQiIMWqxbjsHY--oWlA2Rj2oZeP4,4063
+molcraft-0.1.0a7.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+molcraft-0.1.0a7.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a7.dist-info/RECORD,,

{molcraft-0.1.0a6.dist-info → molcraft-0.1.0a7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.3.1)
+Generator: setuptools (80.7.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

molcraft-0.1.0a6.dist-info/RECORD

@@ -1,19 +0,0 @@
-molcraft/__init__.py,sha256=eKMk4e5Wki4Ay7_BUuY7B-j3Po2l5FDDefPjkFQw3OM,463
-molcraft/callbacks.py,sha256=mkz4ALjJFPy8nHd2nCAuMbKceKnq4tIpZhUuUOvie2Y,1209
-molcraft/chem.py,sha256=apaECcQSuAMs3Tm12yc6ne4x0BGx5JzfoRhTC1WMhlI,20695
-molcraft/conformers.py,sha256=rojo8OaZrKAesx0JA5kf-JVNEpmsQyLSpcxbWhV9cd4,4324
-molcraft/datasets.py,sha256=rFgXTC1ZheLhfgQgcCspP_wEE54a33PIneH7OplbS-8,4047
-molcraft/descriptors.py,sha256=gKqlJ3BqJLTeR2ft8isftSEaJDC8cv64eTq5IYhy4XM,3032
-molcraft/features.py,sha256=69oV_GHNdBKPA4sp6Tpo6brvNmaauk_IVIzNjX7VDmg,13648
-molcraft/featurizers.py,sha256=aJJibnHCxvSu3bNbE2xQk34QvFb47Mnm__0MxlRLA0w,27323
-molcraft/layers.py,sha256=RyKmdHmHlYJJL15LvHH32daTKsChJ_pHmHUnpUcwS1U,73437
-molcraft/losses.py,sha256=JEKZEX2f8vDgky_fUocsF8vZjy9VMzRjZUBa20Uf9Qw,1065
-molcraft/models.py,sha256=Rl9CkQlOVkj20TLjGlwI8vaQwX07EqqWz22bFYtJlpk,22636
-molcraft/ops.py,sha256=eAi79aawJwxuIVVamjA1kPRHGlUm0PsvN-7d2CYu15I,4441
-molcraft/records.py,sha256=0sjOdcr266ZER4F-aTBQ3AVPNAwflKWNiNJVsSc1-PQ,5370
-molcraft/tensors.py,sha256=EOUKx496KUZsjA1zA2ABc7tU_TW3Jv7AXDsug_QsLbA,22407
-molcraft-0.1.0a6.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a6.dist-info/METADATA,sha256=Zzl1K3WleDp056zbLChy5B1AQ3U26t22oMkIKRUpbMY,4063
-molcraft-0.1.0a6.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-molcraft-0.1.0a6.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a6.dist-info/RECORD,,