molcraft 0.1.0a6__tar.gz → 0.1.0a7__tar.gz

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

@@ -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.0a6 → molcraft-0.1.0a7}/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-0.1.0a7/molcraft/callbacks.py

@@ -0,0 +1,93 @@
+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):
+
+    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-0.1.0a6 → molcraft-0.1.0a7}/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-0.1.0a6 → molcraft-0.1.0a7}/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.0a7}/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-0.1.0a6 → molcraft-0.1.0a7}/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-0.1.0a6 → molcraft-0.1.0a7}/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 → molcraft-0.1.0a7}/molcraft.egg-info/PKG-INFO

@@ -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.0a6 → molcraft-0.1.0a7}/tests/test_models.py

@@ -1,6 +1,8 @@
 import unittest 
-
+import tempfile
+import shutil
 import keras
+import numpy as np
 
 from molcraft import tensors 
 from molcraft import layers 
@@ -138,9 +140,64 @@ class TestModel(unittest.TestCase):
                 self.assertTrue(isinstance(metrics, list))
                 del model
 
-    # TODO: Write test for saving and loading model: model(tensor) == loaded_model(tensor)
     def test_saved_model(self):
-        pass
+
+        def get_model(tensor):
+            inputs = layers.Input(tensor.spec)
+            x = layers.NodeEmbedding(32)(inputs)
+            x = layers.EdgeEmbedding(32)(x)
+            x = layers.GTConv(32)(x)
+            x = layers.GTConv(32)(x)
+            x = layers.Readout('sum')(x)
+            outputs = keras.layers.Dense(1)(x)
+            return models.GraphModel(inputs, outputs)
+
+        @keras.saving.register_keras_serializable()
+        class Model(models.GraphModel):
+            def __init__(self, **kwargs):
+                super().__init__(**kwargs)
+                self.e1 = layers.NodeEmbedding(32)
+                self.e2 = layers.EdgeEmbedding(32)
+                self.c1 = layers.GTConv(32)
+                self.c2 = layers.GTConv(32)
+                self.r = layers.Readout('sum')
+                self.d = keras.layers.Dense(1)
+            def propagate(self, tensor):
+                return self.d(self.r(self.c2(self.c1(self.e2(self.e1(tensor))))))
+            
+        example = self.tensors[-1]
+
+        tmp_dir = tempfile.mkdtemp()
+        tmp_file = tmp_dir + '/model.keras'
+        with self.subTest(functional_model=True):
+            model = get_model(example)
+            model.compile('adam', 'mse')
+            model.fit(example, verbose=0)
+            model.save(tmp_file)
+            loaded_model = models.load_model(tmp_file)
+            pred_1 = model.predict(example, verbose=0)
+            pred_2 = loaded_model.predict(example, verbose=0)
+            test_preds = np.all(pred_1.round(4) == pred_2.round(4))
+            self.assertTrue(test_preds)
+            test_vars = np.all([np.all((v1 == v2).numpy()) for (v1, v2) in zip(model.variables, loaded_model.variables)])
+            self.assertTrue(test_vars)
+        shutil.rmtree(tmp_dir)
+
+        tmp_dir = tempfile.mkdtemp()
+        tmp_file = tmp_dir + '/model.keras'
+        with self.subTest(functional_model=False):
+            model = Model()
+            model.compile('adam', 'mse')
+            model.fit(example, verbose=0)
+            model.save(tmp_file)
+            loaded_model = models.load_model(tmp_file)
+            pred_1 = model.predict(example, verbose=0)
+            pred_2 = loaded_model.predict(example, verbose=0)
+            test_preds = np.all(pred_1.round(4) == pred_2.round(4))
+            self.assertTrue(test_preds)
+            test_vars = np.all([np.all((v1 == v2).numpy()) for (v1, v2) in zip(model.variables, loaded_model.variables)])
+            self.assertTrue(test_vars)
+        shutil.rmtree(tmp_dir)
 
     def test_subclassed_model(self):
 

molcraft-0.1.0a6/molcraft/callbacks.py

@@ -1,33 +0,0 @@
-import keras
-
-
-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)