molcraft 0.1.0a14__tar.gz → 0.1.0a15__tar.gz

{molcraft-0.1.0a14 → molcraft-0.1.0a15}/PKG-INFO

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

{molcraft-0.1.0a14 → molcraft-0.1.0a15}/molcraft/__init__.py

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

molcraft-0.1.0a15/molcraft/apps/qsrr.py

@@ -0,0 +1,47 @@
+import molcraft
+import keras
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class AuxiliaryFeatureInjection(molcraft.layers.GraphLayer):
+
+    def __init__(
+        self,
+        field: str = 'auxiliary_feature',
+        depth: int = 2,
+        drop: bool = True,
+        activation: str | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.field = field 
+        self.depth = depth 
+        self.drop = drop
+        self.activation = keras.activations.get(activation)
+
+    def build(self, spec: molcraft.tensors.GraphTensor.Spec) -> None:
+        units = spec.node['feature'].shape[1]
+        for i in range(self.depth):
+            setattr(
+                self, f'dense_{i}', self.get_dense(units, activation=self.activation)
+            )
+
+    def propagate(self, tensor: molcraft.tensors.GraphTensor) -> None:
+        x = tensor.context[self.field] 
+        if self.drop:
+            tensor = tensor.update({'context': {self.field: None}})
+        for i in range(self.depth):
+            x = getattr(self, f'dense_{i}')(x)
+        node_feature = molcraft.ops.scatter_add(
+            tensor.node['feature'], tensor.node['super'], x
+        )
+        return tensor.update({'node': {'feature': node_feature}})
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            'field': self.field,
+            'depth': self.depth,
+            'drop': self.drop,
+            'activation': keras.activations.serialize(self.activation)
+        })
+        return config

{molcraft-0.1.0a14 → molcraft-0.1.0a15}/molcraft/featurizers.py

@@ -169,7 +169,7 @@ class MolGraphFeaturizer(Featurizer):
         if default_atom_features:
             atom_features = [features.AtomType()]
             if not self.include_hs:
-                atom_features.append(features.TotalNumHs())
+                atom_features.append(features.NumHydrogens())
             atom_features.append(features.Degree())
         if not isinstance(self, MolGraphFeaturizer3D):
             default_bond_features = (

{molcraft-0.1.0a14 → molcraft-0.1.0a15}/molcraft/models.py

@@ -397,7 +397,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
             raise ValueError(
                 'Could not extract output. `Readout` layer not found.'
             )
-        return self.__class__(inputs, outputs, name=f'{self.name}_head')
+        return self.__class__(inputs, outputs, name=f'{self.name}_backbone')
 
     def head(self) -> functional.Functional:
         if not isinstance(self, FunctionalGraphModel):

{molcraft-0.1.0a14 → molcraft-0.1.0a15}/molcraft/ops.py

@@ -4,6 +4,7 @@ import tensorflow as tf
 from keras import backend
 
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def gather(
     node_feature: tf.Tensor, 
     edge: tf.Tensor
@@ -16,6 +17,7 @@ def gather(
         edge = keras.ops.expand_dims(edge, axis=-1)
     return keras.ops.take_along_axis(node_feature, edge, axis=0)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def aggregate(
     node_feature: tf.Tensor, 
     edge: tf.Tensor, 
@@ -30,6 +32,7 @@ def aggregate(
         node_feature, edge, num_nodes, sorted=False
     )
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def propagate(
     node_feature: tf.Tensor,
     edge_source: tf.Tensor,
@@ -49,6 +52,7 @@ def propagate(
         
     return aggregate(node_feature, edge_target, num_nodes)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def scatter_update(
     inputs: tf.Tensor,
     indices: tf.Tensor,
@@ -62,6 +66,7 @@ def scatter_update(
         indices = keras.ops.expand_dims(indices, axis=-1)
     return keras.ops.scatter_update(inputs, indices, updates)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def scatter_add(
     inputs: tf.Tensor,
     indices: tf.Tensor,
@@ -78,6 +83,7 @@ def scatter_add(
     updates = scatter_update(keras.ops.zeros_like(inputs), indices, updates)
     return inputs + updates
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def edge_softmax(
     score: tf.Tensor, 
     edge_target: tf.Tensor
@@ -98,6 +104,7 @@ def edge_softmax(
     denominator = gather(denominator, edge_target)
     return numerator / denominator
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def edge_weight(
     edge: tf.Tensor,
     edge_weight: tf.Tensor,
@@ -108,6 +115,7 @@ def edge_weight(
         edge_weight = keras.ops.expand_dims(edge_weight, axis=-1)
     return edge * edge_weight
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def segment_mean(
     data: tf.Tensor,
     segment_ids: tf.Tensor,
@@ -142,6 +150,7 @@ def segment_mean(
     )
     return x / sizes[:, None]
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def gaussian(
     x: tf.Tensor, 
     mean: tf.Tensor, 
@@ -155,6 +164,7 @@ def gaussian(
     a = (2 * np.pi) ** 0.5
     return keras.ops.exp(-0.5 * (((x - mean) / std) ** 2)) / (a * std)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def euclidean_distance(
     x1: tf.Tensor, 
     x2: tf.Tensor, 
@@ -169,6 +179,7 @@ def euclidean_distance(
         )
     )
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def displacement(
     x1: tf.Tensor,
     x2: tf.Tensor,

{molcraft-0.1.0a14 → molcraft-0.1.0a15}/molcraft.egg-info/PKG-INFO

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

{molcraft-0.1.0a14 → molcraft-0.1.0a15}/molcraft.egg-info/SOURCES.txt

@@ -22,6 +22,7 @@ molcraft.egg-info/requires.txt
 molcraft.egg-info/top_level.txt
 molcraft/apps/__init__.py
 molcraft/apps/peptides.py
+molcraft/apps/qsrr.py
 tests/test_chem.py
 tests/test_featurizers.py
 tests/test_layers.py