molcraft 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

molcraft/models.py

@@ -194,8 +194,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         Args:
             x: 
                 A `GraphTensor` instance or a `tf.data.Dataset` constructed from
-                a `GraphTensor` instance. Context `label`s and/or `weight`s may
-                be encoded and will be ignored.
+                a `GraphTensor` instance.
             batch_size:
                 Number of samples per batch of computation.
             kwargs:
@@ -316,10 +315,16 @@ class FunctionalGraphModel(functional.Functional, GraphModel):
         ]
 
 
-def save_model(model: keras.Model, filepath: str | Path, *args, **kwargs) -> None:
+def save_model(model: GraphModel, filepath: str | Path, *args, **kwargs) -> None:
+    if not model.built:
+        raise ValueError(
+            'Model and its layers have not yet been (fully) built. '
+            'Build the model before saving it: `model.build(graph_spec)` '
+            'or `model(graph)`.'
+        )
     keras.models.save_model(model, filepath, *args, **kwargs)
 
-def load_model(filepath: str | Path, inputs=None, *args, **kwargs) -> None:
+def load_model(filepath: str | Path, inputs=None, *args, **kwargs) -> GraphModel:
     return keras.models.load_model(filepath, *args, **kwargs)
 
 def create(
@@ -334,7 +339,7 @@ def create(
 def interpret(
     model: GraphModel,
     graph_tensor: tensors.GraphTensor,
-) -> tuple[tf.Tensor | tf.RaggedTensor | np.ndarray, tf.Tensor | np.ndarray]:
+) -> tensors.GraphTensor:
     x = graph_tensor
     if tensors.is_ragged(x):
         x = x.flatten()
@@ -374,6 +379,31 @@ def interpret(
         }
     )
 
+def saliency(
+    model: GraphModel,
+    graph_tensor: tensors.GraphTensor
+) -> tensors.GraphTensor:
+    x = graph_tensor
+    if tensors.is_ragged(x):
+        x = x.flatten()
+    y_true = x.context.get('label')
+    with tf.GradientTape(watch_accessed_variables=False) as tape:
+        tape.watch(x.node['feature'])
+        y_pred = model(x, training=False)
+        if y_true is not None and len(y_true.shape) > 1: 
+            target = tf.gather_nd(y_pred, tf.where(y_true != 0))
+        else:
+            target = y_pred
+    gradients = tape.gradient(target, x.node['feature'])
+    gradients = keras.ops.absolute(gradients)
+    return graph_tensor.update(
+        {
+            'node': {
+                'feature_saliency': gradients
+            }
+        }
+    ) 
+
 def predict(
     model: GraphModel,
     x: tensors.GraphTensor | tf.data.Dataset,

molcraft/tensors.py

@@ -9,6 +9,9 @@ from molcraft import ops
 class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
     
     def batch(self, spec: 'GraphTensor.Spec', batch_size: int | None):
+        """Batches spec.
+        """
+
         def batch_field(f):
             if isinstance(f, tf.TensorSpec):
                 return tf.TensorSpec(
@@ -36,6 +39,9 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         return batched_spec
     
     def unbatch(self, spec: 'GraphTensor.Spec'):
+        """Unbatches spec.
+        """
+
         def unbatch_field(f):   
             if isinstance(f, tf.TensorSpec):
                 return tf.TensorSpec(
@@ -63,6 +69,8 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         return unbatched_spec
         
     def encode(self, spec: 'GraphTensor.Spec', value: 'GraphTensor', minimum_rank: int = 0):
+        """Encodes value.
+        """
         unflatten = False if (is_ragged(spec) or is_scalar(spec)) else True 
         if unflatten:
             value = value.unflatten()
@@ -74,6 +82,8 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         return value_components
     
     def encoding_specs(self, spec: 'GraphTensor.Spec'):
+        """Matches spec and encoded value of `encode(spec, value)`.
+        """
         def encode_fields(f):
             if isinstance(f, tf.TensorSpec):
                 scalar = is_scalar(spec)
@@ -95,6 +105,8 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         return spec_components
     
     def decode(self, spec, encoded_value):
+        """Decodes encoded value.
+        """
         spec_tuple = tuple(spec.__dict__.values())
         encoded_value = iter(encoded_value)
         value_tuple = [
@@ -122,18 +134,27 @@ class GraphTensor(tf.experimental.BatchableExtensionType):
     __name__ = 'GraphTensor'
 
     def __validate__(self):
-        assert 'size' in self.context, "graph.context['size'] is required."
-        assert self.context['size'].dtype == tf.int32, (
-            "dtype of graph.context['size'] needs to be int32.")
-        assert 'feature' in self.node, "graph.node['feature'] is required."
-        assert 'source' in self.edge, "graph.edge['source'] is required."
-        assert 'target' in self.edge, "graph.edge['target'] is required."
-        assert self.edge['source'].dtype == tf.int32, (
-            "dtype of graph.edge['source'] needs to be int32.")
-        assert self.edge['target'].dtype == tf.int32, (
-            "dtype of graph.edge['target'] needs to be int32.")
-        # TODO: Assert node sizes (based on context['size'])
-        #       Assert edge sizes (based on edge['source'])
+        if tf.executing_eagerly():
+            assert 'size' in self.context, "graph.context['size'] is required."
+            assert self.context['size'].dtype == tf.int32, (
+                "dtype of graph.context['size'] needs to be int32."
+            )
+            assert 'feature' in self.node, "graph.node['feature'] is required."
+            assert 'source' in self.edge, "graph.edge['source'] is required."
+            assert 'target' in self.edge, "graph.edge['target'] is required."
+            assert self.edge['source'].dtype == tf.int32, (
+                "dtype of graph.edge['source'] needs to be int32."
+            )
+            assert self.edge['target'].dtype == tf.int32, (
+                "dtype of graph.edge['target'] needs to be int32."
+            )
+            if isinstance(self.node['feature'], tf.Tensor):
+                num_nodes = keras.ops.shape(self.node['feature'])[0]
+            else:
+                num_nodes = keras.ops.sum(self.node['feature'].row_lengths())
+            assert keras.ops.sum(self.context['size']) == num_nodes, (
+                "graph.node['feature'] tensor is incompatible with graph.context['size']"
+            )
         
     @property 
     def spec(self):

{molcraft-0.1.0a1.dist-info → molcraft-0.1.0a3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a1
+Version: 0.1.0a3
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -56,3 +56,70 @@ Dynamic: license-file
 - Modular graph **layers**
 - Serializable graph **featurizers** and **models**
 - Flexible **GraphTensor**
+
+## Examples 
+
+```python
+from molcraft import features
+from molcraft import descriptors
+from molcraft import featurizers 
+from molcraft import layers
+from molcraft import models 
+import keras
+
+featurizer = featurizers.MolGraphFeaturizer(
+    atom_features=[
+        features.AtomType(),
+        features.TotalNumHs(),
+        features.Degree(),
+    ],
+    bond_features=[
+        features.BondType(),
+        features.IsRotatable(),
+    ],
+    super_atom=True,
+    self_loops=False,
+)
+
+graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])
+print(graph)
+
+model = models.GraphModel.from_layers(
+    [
+        layers.Input(graph.spec),
+        layers.NodeEmbedding(dim=128),
+        layers.EdgeEmbedding(dim=128),
+        layers.GraphTransformer(units=128),
+        layers.GraphTransformer(units=128),
+        layers.GraphTransformer(units=128),
+        layers.GraphTransformer(units=128),
+        layers.Readout(mode='mean'),
+        keras.layers.Dense(units=1024, activation='relu'),
+        keras.layers.Dense(units=1024, activation='relu'),
+        keras.layers.Dense(1)
+    ]
+)
+
+pred = model(graph)
+print(pred)
+
+# featurizers.save_featurizer(featurizer, '/tmp/featurizer.json')
+# models.save_model(model, '/tmp/model.keras')
+
+# featurizers.load_featurizer('/tmp/featurizer.json')
+# models.load_model('/tmp/model.keras')
+```
+
+## Installation
+
+Install the pre-release of molcraft via pip:
+
+```bash
+pip install molcraft --pre
+```
+
+with GPU support:
+
+```bash
+pip install molcraft[gpu] --pre
+```

molcraft-0.1.0a3.dist-info/RECORD

@@ -0,0 +1,20 @@
+molcraft/__init__.py,sha256=2ZNfWBjGl8DscOwjdDiRkgIsuPnKit29Q3MhZyP336Q,435
+molcraft/callbacks.py,sha256=6gwCwdsHGb-fVB4m1QGmtBwQwZ9mFq9QUkmPKSMn05U,849
+molcraft/chem.py,sha256=_UO5O-I7KUtGf3vRrFEYoAUGlW5xi2x8ylu5f-Ybumo,18696
+molcraft/conformers.py,sha256=p09gOQOdxLSj3yohZOMkxxLriHsZ1ZqOoiWLi73OpIg,4325
+molcraft/datasets.py,sha256=rFgXTC1ZheLhfgQgcCspP_wEE54a33PIneH7OplbS-8,4047
+molcraft/descriptors.py,sha256=x6RfZ-gK7D_WSvmK6sh6yHyEjQqovPnRU0xwC3dAKfg,2880
+molcraft/features.py,sha256=69oV_GHNdBKPA4sp6Tpo6brvNmaauk_IVIzNjX7VDmg,13648
+molcraft/featurizers.py,sha256=Yu8I6I_zkzB__WYSiqz-FDGjvKFOmyWFxojRBr39Aw8,26236
+molcraft/layers.py,sha256=HjnAtqhuP0uZ5yP4L33k3xT4IUdLavWBrjd3wO9_Rmw,64915
+molcraft/models.py,sha256=DXqWR_XnMVXQseVR91XnDLXvmHa1hv-6_Y_wvpQZBFI,17476
+molcraft/ops.py,sha256=iiE6zgA2P7cmjKO1RHmL9GE_Tv7Tyuo_xDoxB_ELZQM,3824
+molcraft/records.py,sha256=w4-bcWZEC0oVInrE1e0kQBroIaSCA0PN1JBPOtO6VUY,5251
+molcraft/tensors.py,sha256=b7PO-YOvV72s9g057ILJACKS2n2fn10VkO35gHXpssI,22312
+molcraft/experimental/__init__.py,sha256=x5h6LOO8bo3NPjkKKM9M1H-Kz6R3yxYhRSePoxHCdRE,42
+molcraft/experimental/peptides.py,sha256=RCuOTSwoYHGSdeYi6TWHdPIv2WC3avCZjKLdhEZQeXw,8997
+molcraft-0.1.0a3.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a3.dist-info/METADATA,sha256=f_5sBinpFcGSqKLaSqGkZJ83gGQtZw1Pb3fkgq9aCBM,4088
+molcraft-0.1.0a3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+molcraft-0.1.0a3.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a3.dist-info/RECORD,,

molcraft-0.1.0a1.dist-info/RECORD

@@ -1,19 +0,0 @@
-molcraft/__init__.py,sha256=lx_frHVGfTmGUmzLhCt5LJVy47W_d_mTrKtRsw9RTQ8,406
-molcraft/callbacks.py,sha256=6gwCwdsHGb-fVB4m1QGmtBwQwZ9mFq9QUkmPKSMn05U,849
-molcraft/chem.py,sha256=_UO5O-I7KUtGf3vRrFEYoAUGlW5xi2x8ylu5f-Ybumo,18696
-molcraft/conformers.py,sha256=p09gOQOdxLSj3yohZOMkxxLriHsZ1ZqOoiWLi73OpIg,4325
-molcraft/descriptors.py,sha256=x6RfZ-gK7D_WSvmK6sh6yHyEjQqovPnRU0xwC3dAKfg,2880
-molcraft/features.py,sha256=nZDfX9fsWWjhUbUbrWSUI0ny1QIDbxb4MO8umjcdQqw,13572
-molcraft/featurizers.py,sha256=epv-K5ah9MVFEDZ7c2UzT7f9Vxglr28sbQoTYk1fev8,24583
-molcraft/layers.py,sha256=XRxAUXnYOQ_fD8OegkuJTb1zMHHKm_4lAjxnIr2fdl4,43119
-molcraft/models.py,sha256=z2V9_I_cnJr3VNqfY_CWkZKojlLoD_8r9MHkv3pKOh8,16605
-molcraft/ops.py,sha256=iiE6zgA2P7cmjKO1RHmL9GE_Tv7Tyuo_xDoxB_ELZQM,3824
-molcraft/records.py,sha256=w4-bcWZEC0oVInrE1e0kQBroIaSCA0PN1JBPOtO6VUY,5251
-molcraft/tensors.py,sha256=_zlcTHH47icUe_S_wYzRp7J-4M2U_4lLkAZOIInea0w,21677
-molcraft/experimental/__init__.py,sha256=x5h6LOO8bo3NPjkKKM9M1H-Kz6R3yxYhRSePoxHCdRE,42
-molcraft/experimental/peptides.py,sha256=AcjuilPUPzn6l2ui5YPYE00VpXnGj8H1osTjz58biGw,10442
-molcraft-0.1.0a1.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a1.dist-info/METADATA,sha256=0fczlpAEjq5hLHqGCQPdocQXeoCRfF7L4cyoz8WzYeM,2573
-molcraft-0.1.0a1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-molcraft-0.1.0a1.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a1.dist-info/RECORD,,