molcraft 0.1.0a1__py3-none-any.whl → 0.1.0a2__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:

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.0a2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a1
+Version: 0.1.0a2
 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.0a2.dist-info/RECORD

@@ -0,0 +1,20 @@
+molcraft/__init__.py,sha256=lE7_mCo7lLcP1AopGZtGyWqzAN1qgjZnH5juymdjrJc,406
+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=nZDfX9fsWWjhUbUbrWSUI0ny1QIDbxb4MO8umjcdQqw,13572
+molcraft/featurizers.py,sha256=gAUe7Ui8gF32aotuiDAUoRUuw8bTbkMgB2C2BO1VWDM,26176
+molcraft/layers.py,sha256=zs6Ae6p7ASeAy3eF113f35d55yQmyk2Z7vUUfkfJUmY,49677
+molcraft/models.py,sha256=Nvm5LKCtH-xj395f1OvIEmYVTTrnutoSthL2DxGicnY,16519
+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.0a2.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a2.dist-info/METADATA,sha256=TYf32YHTSrK9OaaGKCCk89uPlf_REWsK-LKf93c6V4M,4088
+molcraft-0.1.0a2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+molcraft-0.1.0a2.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a2.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,,