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

molcraft/__init__.py

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

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:

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

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

@@ -1,4 +1,4 @@
-molcraft/__init__.py,sha256=8f1z8Lhuhh8TxB-QGHI5w4a3M_ZZNH8EWGD4Y6pB578,463
+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
@@ -6,14 +6,14 @@ 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=r6hEAyJxO_Yrw5hD1r2v8yb_UxLRK9S4FMjDCUQedH8,59655
+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=dLIUq-KG8nOzEcphJqNbF_f82VZRDNrB1UKrcPt5JNM,4752
-molcraft/records.py,sha256=0sjOdcr266ZER4F-aTBQ3AVPNAwflKWNiNJVsSc1-PQ,5370
+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.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,,
+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.0a9.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