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

molcraft/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0a8'
+__version__ = '0.1.0a10'
 
 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

@@ -350,7 +350,7 @@ class GraphConv(GraphLayer):
         )
         if self._project_residual:
             warnings.warn(
-                '`skip_connect` is set to `True`, but found incompatible dim ' 
+                '`skip_connect` is set to `True`, but found incompatible dim '
                 'between input (node feature dim) and output (`self.units`). '
                 'Automatically applying a projection layer to residual to '
                 'match input and output. ',
@@ -369,7 +369,7 @@ class GraphConv(GraphLayer):
             self._message_intermediate_activation = self.activation
             self._message_final_dense = self.get_dense(self.units)
 
-        has_overridden_aggregate = self.__class__.message != GraphConv.aggregate 
+        has_overridden_aggregate = self.__class__.message != GraphConv.aggregate
         if not has_overridden_aggregate:
             pass
 
@@ -401,13 +401,15 @@ class GraphConv(GraphLayer):
                 residual = self._residual_dense(residual)
 
         message = self.message(tensor)
-        if not isinstance(message, tensors.GraphTensor):
+        add_message = not isinstance(message, tensors.GraphTensor)
+        if add_message:
             message = tensor.update({'edge': {'message': message}})
         elif not 'message' in message.edge:
             raise ValueError('Could not find `message` in `edge` output.')
-            
+        
         aggregate = self.aggregate(message)
-        if not isinstance(aggregate, tensors.GraphTensor):
+        add_aggregate = not isinstance(aggregate, tensors.GraphTensor)
+        if add_aggregate:
             aggregate = tensor.update({'node': {'aggregate': aggregate}})
         elif not 'aggregate' in aggregate.node:
             raise ValueError('Could not find `aggregate` in `node` output.')
@@ -421,6 +423,16 @@ class GraphConv(GraphLayer):
         if update.node['feature'].shape[-1] != self.units:
             raise ValueError('Updated node `feature` is not equal to `self.units`.')
 
+        if add_message and add_aggregate:
+            update = update.update({'node': {'aggregate': None}, 'edge': {'message': None}})
+        elif add_message:
+            update = update.update({'edge': {'message': None}})
+        elif add_aggregate:
+            update = update.update({'node': {'aggregate': None}})
+
+        if not self._skip_connect and not self._normalize:
+            return update
+        
         feature = update.node['feature']
 
         if self._skip_connect:
@@ -1303,17 +1315,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 +1334,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 +1347,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 +1366,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 +1382,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 +1414,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)
+        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)
 
-        if self._normalize:
-            feature = self._norm(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/ops.py

@@ -105,7 +105,11 @@ def segment_mean(
             lambda: 0
         )
     if backend.backend() == 'tensorflow':
-        return tf.math.unsorted_segment_mean(
+        segment_mean_fn = (
+            tf.math.unsorted_segment_mean if not sorted else
+            tf.math.segment_mean
+        )
+        return segment_mean_fn(
             data=data,
             segment_ids=segment_ids,
             num_segments=num_segments

molcraft/records.py

@@ -51,19 +51,24 @@ def write(
         if num_files is None:
             num_files = min(len(inputs), max(1, math.ceil(len(inputs) / 1_000)))
             
-        chunk_size = math.ceil(len(inputs) / num_files)
-        num_files = math.ceil(len(inputs) / chunk_size)
+        num_examples = len(inputs)
+        chunk_sizes = [0] * num_files
+        for i in range(num_examples):
+            chunk_sizes[i % num_files] += 1
+        
+        input_chunks = []
+        current_index = 0
+        for size in chunk_sizes:
+            input_chunks.append(inputs[current_index: current_index + size])
+            current_index += size 
+        
+        assert current_index == num_examples
         
         paths = [
             os.path.join(path, f'tfrecord-{i:04d}.tfrecord')
             for i in range(num_files)
         ]
         
-        input_chunks = [
-            inputs[i * chunk_size: (i + 1) * chunk_size]
-            for i in range(num_files)
-        ]
-        
         if not multiprocessing:
             for path, input_chunk in zip(paths, input_chunks):
                 _write_tfrecord(input_chunk, path, featurizer)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a8
+Version: 0.1.0a10
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -25,7 +25,7 @@ License: MIT License
         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
         SOFTWARE.
         
-Project-URL: Homepage, https://github.com/akensert/molcraft
+Project-URL: Homepage, https://github.com/compomics/molcraft
 Keywords: python,machine-learning,deep-learning,graph-neural-networks,molecular-machine-learning,molecular-graphs,computational-chemistry,computational-biology
 Classifier: Programming Language :: Python :: 3
 Classifier: Intended Audience :: Science/Research

molcraft-0.1.0a10.dist-info/RECORD

@@ -0,0 +1,19 @@
+molcraft/__init__.py,sha256=1te1sOK-k4LT9l-mlxRmOhY3_Za-7jPezx_B3gdToiQ,464
+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=ybJ1djH747cgsftztWHxAX2iTq6k03MYr17btQ2Gtcs,27063
+molcraft/layers.py,sha256=G-ZFhnyiSny0YHGXg5tBYrvmhZsurBEJj_0mHD1zmlw,60135
+molcraft/losses.py,sha256=JEKZEX2f8vDgky_fUocsF8vZjy9VMzRjZUBa20Uf9Qw,1065
+molcraft/models.py,sha256=FLXpO3OUmRxLmxG3MjBK4ZwcVFlea1gqEgs1ibKly2w,23263
+molcraft/ops.py,sha256=PVxKfY_XbWCyntiSnmpyeBb-coFGT_VNNP9QzmeUwC0,4870
+molcraft/records.py,sha256=MbvYkcCunbAmpy_MWXmQ9WBGi2WvwxFUlwQSPKPvSSk,5534
+molcraft/tensors.py,sha256=EOUKx496KUZsjA1zA2ABc7tU_TW3Jv7AXDsug_QsLbA,22407
+molcraft-0.1.0a10.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a10.dist-info/METADATA,sha256=Tmh4KckmdKr20q8RVPOKlogt343qTdOMzci6zgT6CfQ,4064
+molcraft-0.1.0a10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+molcraft-0.1.0a10.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a10.dist-info/RECORD,,

{molcraft-0.1.0a8.dist-info → molcraft-0.1.0a10.dist-info}/WHEEL

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

molcraft-0.1.0a8.dist-info/RECORD

@@ -1,19 +0,0 @@
-molcraft/__init__.py,sha256=s8dUh6Fjq34j2aNgF13Y2NUkDwBWmsOAuIJVgY3gwCE,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/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,,