molcraft 0.1.0a8__py3-none-any.whl → 0.1.0a9__py3-none-any.whl

molcraft/__init__.py

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

molcraft/featurizers.py

@@ -280,8 +280,6 @@ class MolGraphFeaturizer(Featurizer):
                             mol.get_bond_between_atoms(atom_i, atom_j).index
                         )
                 edge['feature'] = bond_feature[bond_indices]
-                if self.self_loops:
-                    edge['self_loop'] = (edge['source'] == edge['target'])
         else:
             paths = chem.get_shortest_paths(
                 mol, radius=self.radius, self_loops=self.self_loops
@@ -376,7 +374,7 @@ class MolGraphFeaturizer(Featurizer):
         num_nodes = node['feature'].shape[0]
         node = _add_super_nodes(node, num_super_nodes)
         edge = _add_super_edges(
-            edge, num_nodes, num_super_nodes, self.feature_dtype, self.index_dtype
+            edge, num_nodes, num_super_nodes, self.feature_dtype, self.index_dtype, self.self_loops
         )
         return node, edge
 
@@ -708,11 +706,15 @@ def _add_super_edges(
     num_super_nodes: int,
     feature_dtype: str,
     index_dtype: str,
+    self_loops: bool,
 ) -> dict[str, np.ndarray]:
     edge = copy.deepcopy(edge)
-    super_node_indices = (
-        np.repeat(np.arange(num_super_nodes), [num_nodes]) + num_nodes
-    )
+
+    super_node_indices = np.arange(num_super_nodes) + num_nodes
+    if self_loops:
+        edge['source'] = np.concatenate([edge['source'], super_node_indices])
+        edge['target'] = np.concatenate([edge['target'], super_node_indices])
+    super_node_indices = np.repeat(super_node_indices, [num_nodes])
     node_indices = (
         np.tile(np.arange(num_nodes), [num_super_nodes])
     )
@@ -727,6 +729,8 @@ def _add_super_edges(
     if 'feature' in edge:
         num_edges = int(edge['feature'].shape[0])
         num_super_edges = int(num_super_nodes * num_nodes * 2)
+        if self_loops:
+            num_super_edges += num_super_nodes
         edge['super'] = np.asarray(
             ([False] * num_edges + [True] * num_super_edges),
             dtype=bool
@@ -741,12 +745,6 @@ def _add_super_edges(
             ]
         )
 
-    if 'self_loop' in edge:
-        edge['self_loop'] = np.pad(
-            edge['self_loop'], [(0, num_nodes * num_super_nodes * 2)],
-            constant_values=False,
-        )
-
     return edge
 
 

molcraft/layers.py

@@ -1303,17 +1303,12 @@ class NodeEmbedding(GraphLayer):
         dim: int = None, 
         normalize: bool = False,
         embed_context: bool = False,
-        allow_reconstruction: bool = False,
-        allow_masking: bool = False, 
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
         self._normalize = normalize
         self._embed_context = embed_context
-        self._masking_rate = None
-        self._allow_masking = allow_masking
-        self._allow_reconstruction = allow_reconstruction
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
         feature_dim = spec.node['feature'].shape[-1]
@@ -1327,8 +1322,6 @@ class NodeEmbedding(GraphLayer):
             self._embed_context = False 
         if self._has_super and not self._embed_context:
             self._super_feature = self.get_weight(shape=[self.dim], name='super_node_feature')
-        if self._allow_masking:
-            self._mask_feature = self.get_weight(shape=[self.dim], name='mask_node_feature')
         if self._embed_context:
             self._context_dense = self.get_dense(self.dim)
         
@@ -1342,28 +1335,18 @@ class NodeEmbedding(GraphLayer):
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         feature = self._node_dense(tensor.node['feature'])
 
-        if self._has_super:
-            super_feature = (0 if self._embed_context else self._super_feature)
+        if self._has_super and not self._embed_context:
             super_mask = keras.ops.expand_dims(tensor.node['super'], 1)
-            feature = keras.ops.where(super_mask, super_feature, feature)
+            feature = keras.ops.where(super_mask, self._super_feature, feature)
 
         if self._embed_context:
             context_feature = self._context_dense(tensor.context['feature'])
             feature = ops.scatter_update(feature, tensor.node['super'], context_feature)
             tensor = tensor.update({'context': {'feature': None}})
 
-        apply_mask = (self._allow_masking and 'mask' in tensor.node)
-        if apply_mask:
-            mask = keras.ops.expand_dims(tensor.node['mask'], -1)
-            feature = keras.ops.where(mask, self._mask_feature, feature)
-        elif self._allow_masking:
-            feature = feature + (self._mask_feature * 0.0)
-
         feature = self._norm(feature)
 
-        if not self._allow_reconstruction:
-            return tensor.update({'node': {'feature': feature}})
-        return tensor.update({'node': {'feature': feature, 'target_feature': feature}})
+        return tensor.update({'node': {'feature': feature}})
 
     def get_config(self) -> dict:
         config = super().get_config()
@@ -1371,8 +1354,6 @@ class NodeEmbedding(GraphLayer):
             'dim': self.dim,
             'normalize': self._normalize,
             'embed_context': self._embed_context,
-            'allow_masking': self._allow_masking,
-            'allow_reconstruction': self._allow_reconstruction,
         })
         return config
     
@@ -1389,39 +1370,30 @@ class EdgeEmbedding(GraphLayer):
         self, 
         dim: int = None, 
         normalize: bool = False,
-        allow_masking: bool = True, 
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
         self._normalize = normalize
-        self._masking_rate = None
-        self._allow_masking = allow_masking
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
         feature_dim = spec.edge['feature'].shape[-1]
         if not self.dim:
             self.dim = feature_dim
-        self._edge_dense = self.get_dense(self.dim)
+        self._edge_dense = self.get_dense(self.dim) 
+
+        self._self_loop_feature = self.get_weight(shape=[self.dim], name='self_loop_edge_feature')
 
         self._has_super = 'super' in spec.edge
-        self._has_self_loop = 'self_loop' in spec.edge
         if self._has_super:
             self._super_feature = self.get_weight(shape=[self.dim], name='super_edge_feature')
-        if self._has_self_loop:
-            self._self_loop_feature = self.get_weight(shape=[self.dim], name='self_loop_edge_feature')
-        if self._allow_masking:
-            self._mask_feature = self.get_weight(shape=[self.dim], name='mask_edge_feature')
-
-        if self._normalize:
-            if str(self._normalize).lower().startswith('batch'):
-                self._norm = keras.layers.BatchNormalization(
-                    name='output_batch_norm'
-                )
-            else:
-                self._norm = keras.layers.LayerNormalization(
-                    name='output_layer_norm'
-                )
+        
+        if not self._normalize:
+            self._norm = keras.layers.Identity()
+        elif str(self._normalize).lower().startswith('layer'):
+            self._norm = keras.layers.LayerNormalization()
+        else:
+            self._norm = keras.layers.BatchNormalization()
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         feature = self._edge_dense(tensor.edge['feature'])
@@ -1430,51 +1402,18 @@ class EdgeEmbedding(GraphLayer):
             super_mask = keras.ops.expand_dims(tensor.edge['super'], 1)
             feature = keras.ops.where(super_mask, self._super_feature, feature)
 
-        if self._has_self_loop:
-            self_loop_mask = keras.ops.expand_dims(tensor.edge['self_loop'], 1)
-            feature = keras.ops.where(self_loop_mask, self._self_loop_feature, feature)
-            
-        if (
-            self._allow_masking and 
-            self._masking_rate is not None and 
-            self._masking_rate > 0
-        ):
-            random = keras.random.uniform(shape=[tensor.num_edges])
-            mask = random <= self._masking_rate
-            if self._has_super:
-                mask = keras.ops.logical_and(
-                    mask, keras.ops.logical_not(tensor.edge['super'])
-                )
-            mask = keras.ops.expand_dims(mask, -1)
-            feature = keras.ops.where(mask, self._mask_feature, feature)
-        elif self._allow_masking:
-            # Simply added to silence warning ('no gradients for variables ...')
-            feature += (0.0 * self._mask_feature)
-
-        if self._normalize:
-            feature = self._norm(feature)
+        self_loop_mask = keras.ops.expand_dims(tensor.edge['source'] == tensor.edge['target'], 1)
+        feature = keras.ops.where(self_loop_mask, self._self_loop_feature, feature)
 
-        return tensor.update({'edge': {'feature': feature, 'embedding': feature}})
+        feature = self._norm(feature)
 
-    @property 
-    def masking_rate(self):
-        return self._masking_rate 
-    
-    @masking_rate.setter
-    def masking_rate(self, rate: float):
-        if not self._allow_masking and rate is not None:
-            raise ValueError(
-                f'Cannot set `masking_rate` for layer {self} '
-                'as `allow_masking` was set to `False`.'
-            )
-        self._masking_rate = float(rate)
+        return tensor.update({'edge': {'feature': feature}})
 
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
             'dim': self.dim,
             'normalize': self._normalize,
-            'allow_masking': self._allow_masking
         })
         return config
     

{molcraft-0.1.0a8.dist-info → molcraft-0.1.0a9.dist-info}/METADATA

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

{molcraft-0.1.0a8.dist-info → molcraft-0.1.0a9.dist-info}/RECORD

@@ -1,19 +1,19 @@
-molcraft/__init__.py,sha256=s8dUh6Fjq34j2aNgF13Y2NUkDwBWmsOAuIJVgY3gwCE,463
+molcraft/__init__.py,sha256=8f1z8Lhuhh8TxB-QGHI5w4a3M_ZZNH8EWGD4Y6pB578,463
 molcraft/callbacks.py,sha256=x5HnkZhqcFRrW6xdApt_jZ4X08A-0fxcnFKfdmRKa0c,3571
 molcraft/chem.py,sha256=zHH7iX0ZJ7QmP-YqR_IXCpylTwCXHXptWf1DsblnZR4,21496
 molcraft/conformers.py,sha256=K6ZtiSUNDN_fwqGP9JrPcwALLFFvlMlF_XejEJH3Sr4,4205
 molcraft/datasets.py,sha256=rFgXTC1ZheLhfgQgcCspP_wEE54a33PIneH7OplbS-8,4047
 molcraft/descriptors.py,sha256=gKqlJ3BqJLTeR2ft8isftSEaJDC8cv64eTq5IYhy4XM,3032
 molcraft/features.py,sha256=aBYxDfQqQsVuyjKaPUlwEgvCjbNZ-FJhuKo2Cg5ajrA,13554
-molcraft/featurizers.py,sha256=qNmXSOAeplICN3j-nzvWACVuKoJ_ZBzhYP9LterKVH8,27042
-molcraft/layers.py,sha256=KKaH58zuov5aARj72BS_xK3ZQEwSFJrIPkoXQAAcqz8,62285
+molcraft/featurizers.py,sha256=ybJ1djH747cgsftztWHxAX2iTq6k03MYr17btQ2Gtcs,27063
+molcraft/layers.py,sha256=r6hEAyJxO_Yrw5hD1r2v8yb_UxLRK9S4FMjDCUQedH8,59655
 molcraft/losses.py,sha256=JEKZEX2f8vDgky_fUocsF8vZjy9VMzRjZUBa20Uf9Qw,1065
 molcraft/models.py,sha256=FLXpO3OUmRxLmxG3MjBK4ZwcVFlea1gqEgs1ibKly2w,23263
 molcraft/ops.py,sha256=dLIUq-KG8nOzEcphJqNbF_f82VZRDNrB1UKrcPt5JNM,4752
 molcraft/records.py,sha256=0sjOdcr266ZER4F-aTBQ3AVPNAwflKWNiNJVsSc1-PQ,5370
 molcraft/tensors.py,sha256=EOUKx496KUZsjA1zA2ABc7tU_TW3Jv7AXDsug_QsLbA,22407
-molcraft-0.1.0a8.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a8.dist-info/METADATA,sha256=CtHK0DVlQECWUdlhg0KzvvpPyUD150BSyfzkdNF3fT8,4062
-molcraft-0.1.0a8.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-molcraft-0.1.0a8.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a8.dist-info/RECORD,,
+molcraft-0.1.0a9.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a9.dist-info/METADATA,sha256=HiwS2wmntCA7m_YpgSWKiJTP0BFpl4GWWz4a77w1XBw,4062
+molcraft-0.1.0a9.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+molcraft-0.1.0a9.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a9.dist-info/RECORD,,