molcraft 0.1.0a2__py3-none-any.whl → 0.1.0a4__py3-none-any.whl

molcraft/models.py

@@ -270,6 +270,19 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         """
         super().load_weights(filepath, *args, **kwargs)
 
+    def embedding(self) -> 'FunctionalGraphModel':
+        model = self
+        if not isinstance(model, FunctionalGraphModel):
+            raise ValueError(
+                'Currently, to extract the embedding part of the model, '
+                'it needs to be a `FunctionalGraphModel`. '
+            )
+        inputs = model.input 
+        for layer in model.layers:
+            if isinstance(layer, layers.Readout):
+                outputs = layer.output 
+        return self.__class__(inputs, outputs, name=f'{self.name}_embedding')
+    
     def train_step(self, tensor: tensors.GraphTensor) -> dict[str, float]:
         y = tensor.context.get('label')
         sample_weight = tensor.context.get('weight')
@@ -315,10 +328,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(
@@ -333,7 +352,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()
@@ -373,6 +392,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/ops.py

@@ -19,9 +19,16 @@ def gather(
 def aggregate(
     node_feature: tf.Tensor, 
     edge: tf.Tensor, 
-    num_nodes: tf.Tensor
+    num_nodes: tf.Tensor,
+    mode: str = 'sum',
 ) -> tf.Tensor:
-    return keras.ops.segment_sum(node_feature, edge, num_nodes)
+    if mode == 'mean':
+        return segment_mean(
+            node_feature, edge, num_nodes, sorted=False
+        )
+    return keras.ops.segment_sum(
+        node_feature, edge, num_nodes, sorted=False
+    )
 
 def propagate(
     node_feature: tf.Tensor,
@@ -82,7 +89,11 @@ def segment_mean(
     sorted: bool = False,
 ) -> tf.Tensor:
     if num_segments is None:
-        num_segments = keras.ops.max(segment_ids) + 1
+        num_segments = keras.ops.cond(
+            keras.ops.shape(segment_ids)[0] > 0,
+            lambda: keras.ops.max(segment_ids) + 1,
+            lambda: 0
+        )
     if backend.backend() == 'tensorflow':
         return tf.math.unsorted_segment_mean(
             data=data,

molcraft/tensors.py

@@ -219,13 +219,13 @@ class GraphTensor(tf.experimental.BatchableExtensionType):
             raise ValueError
         return ops.gather(self.node[node_attr], self.edge[edge_type])
     
-    def aggregate(self, edge_attr: str, edge_type: str = 'target') -> tf.Tensor:
+    def aggregate(self, edge_attr: str, edge_type: str = 'target', mode: str = 'sum') -> tf.Tensor:
         if edge_type != 'source' and edge_type != 'target':
-            raise ValueError
+            raise ValueError('`edge_attr` needs to be `source` or `target`.')
         edge_attr = self.edge[edge_attr]
         if 'weight' in self.edge:
             edge_attr = edge_attr * self.edge['weight']
-        return ops.aggregate(edge_attr, self.edge[edge_type], self.num_nodes)
+        return ops.aggregate(edge_attr, self.edge[edge_type], self.num_nodes, mode=mode)
     
     def propagate(self, add_edge_feature: bool = False):
         updated_feature = ops.propagate(

{molcraft-0.1.0a2.dist-info → molcraft-0.1.0a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a2
+Version: 0.1.0a4
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -51,11 +51,11 @@ Dynamic: license-file
 
 ## Highlights
 - Compatible with **Keras 3**
-- Simplified API
-- Fast featurization
-- Modular graph **layers**
-- Serializable graph **featurizers** and **models**
-- Flexible **GraphTensor**
+- Customizable and serializable **featurizers**
+- Customizable and serializable **layers** and **models**
+- Customizable **GraphTensor**
+- Fast and efficient featurization of molecular graphs
+- Efficient and easy-to-use input pipelines using TF **records**
 
 ## Examples 
 

molcraft-0.1.0a4.dist-info/RECORD

@@ -0,0 +1,20 @@
+molcraft/__init__.py,sha256=FQyasgy1kEz2v9sKdr3am6ap7Cm1oHEuCKhHwH-CQpM,435
+molcraft/callbacks.py,sha256=mkz4ALjJFPy8nHd2nCAuMbKceKnq4tIpZhUuUOvie2Y,1209
+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=gKqlJ3BqJLTeR2ft8isftSEaJDC8cv64eTq5IYhy4XM,3032
+molcraft/features.py,sha256=69oV_GHNdBKPA4sp6Tpo6brvNmaauk_IVIzNjX7VDmg,13648
+molcraft/featurizers.py,sha256=kV5RN_Z2pELjDcwE65KYy_JagbDUueXoClpsIOFsI9I,27073
+molcraft/layers.py,sha256=y-sBLXWttr-fkGZ-acL1srMB8QqeXnHotYK9KCcyJNU,70581
+molcraft/models.py,sha256=0MN4PAlsacni7RfIcYm_imxuzBVL2K8w3MnaUM24DeI,18021
+molcraft/ops.py,sha256=uSnBYQwxYJ1ATdDpr290bxiyQZkrSCVxlB7btlh_n2I,4112
+molcraft/records.py,sha256=w4-bcWZEC0oVInrE1e0kQBroIaSCA0PN1JBPOtO6VUY,5251
+molcraft/tensors.py,sha256=8hwlad000wQ5pNLSdzd3rCXVbaUHBxUq2MbBx27dKzU,22391
+molcraft/experimental/__init__.py,sha256=x5h6LOO8bo3NPjkKKM9M1H-Kz6R3yxYhRSePoxHCdRE,42
+molcraft/experimental/peptides.py,sha256=82Bzw9FEnlymOUgTIIKha-ELNbqEFkv9T4hspDGRetw,9266
+molcraft-0.1.0a4.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a4.dist-info/METADATA,sha256=bhsytRfa6BIbfmph0Cm2NfubmZJPumsMQt4lbch33kQ,4201
+molcraft-0.1.0a4.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+molcraft-0.1.0a4.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a4.dist-info/RECORD,,

{molcraft-0.1.0a2.dist-info → molcraft-0.1.0a4.dist-info}/WHEEL

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

molcraft-0.1.0a2.dist-info/RECORD

@@ -1,20 +0,0 @@
-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,,