molcraft 0.1.0a3__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')
@@ -381,7 +394,7 @@ def interpret(
 
 def saliency(
     model: GraphModel,
-    graph_tensor: tensors.GraphTensor
+    graph_tensor: tensors.GraphTensor,
 ) -> tensors.GraphTensor:
     x = graph_tensor
     if tensors.is_ragged(x):

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.0a3.dist-info → molcraft-0.1.0a4.dist-info}/METADATA

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

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