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

molcraft/models.py

@@ -415,6 +415,10 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
     def predict_step(self, tensor: tensors.GraphTensor) -> np.ndarray:
         return self(tensor, training=False)
 
+    def compute_loss(self, x, y, y_pred, sample_weight=None):
+        y, y_pred, sample_weight = _maybe_reshape(y, y_pred, sample_weight)
+        return super().compute_loss(x, y, y_pred, sample_weight)
+        
     def compute_metrics(self, x, y, y_pred, sample_weight=None) -> dict[str, float]:
         loss = self.compute_loss(x, y, y_pred, sample_weight)
         metric_results = {}
@@ -423,7 +427,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
                 metric.update_state(loss)
                 metric_results[metric.name] = metric.result()
             else:
-                metric.update_state(y, y_pred)
+                metric.update_state(y, y_pred, sample_weight=sample_weight)
                 metric_results.update(metric.result())
         return metric_results
     
@@ -593,3 +597,14 @@ def _make_dataset(x: tensors.GraphTensor, batch_size: int):
         .batch(batch_size)
         .prefetch(-1)
     )
+
+def _maybe_reshape(y, y_pred, sample_weight):
+    if (
+        sample_weight is not None and 
+        len(keras.ops.shape(sample_weight)) == 2 and 
+        sample_weight.shape == y_pred.shape
+    ):
+        y = keras.ops.reshape(y, [-1])
+        y_pred = keras.ops.reshape(y_pred, [-1])
+        sample_weight = keras.ops.reshape(sample_weight, [-1])
+    return y, y_pred, sample_weight

molcraft/ops.py

@@ -67,7 +67,7 @@ def edge_softmax(
     edge_target: tf.Tensor
 ) -> tf.Tensor:
     num_segments = keras.ops.cond(
-        keras.ops.shape(edge_target)[0] > 0,
+        keras.ops.greater(keras.ops.shape(edge_target)[0], 0),
         lambda: keras.ops.maximum(keras.ops.max(edge_target) + 1, 1),
         lambda: 0
     )
@@ -100,7 +100,7 @@ def segment_mean(
 ) -> tf.Tensor:
     if num_segments is None:
         num_segments = keras.ops.cond(
-            keras.ops.shape(segment_ids)[0] > 0,
+            keras.ops.greater(keras.ops.shape(segment_ids)[0], 0),
             lambda: keras.ops.max(segment_ids) + 1,
             lambda: 0
         )
@@ -147,4 +147,15 @@ def euclidean_distance(
             axis=axis, 
             keepdims=True
         )
-    )
+    )
+
+def displacement(
+    x1: tf.Tensor,
+    x2: tf.Tensor,
+    normalize: bool = False,
+    axis=-1,
+) -> tf.Tensor:
+    displacement = keras.ops.subtract(x1, x2)
+    if not normalize:
+        return displacement
+    return displacement / euclidean_distance(x1, x2, axis=axis)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a6
+Version: 0.1.0a8
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -78,7 +78,7 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.IsRotatable(),
     ],
     super_atom=True,
-    self_loops=False,
+    self_loops=True,
 )
 
 graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])

molcraft-0.1.0a8.dist-info/RECORD

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

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

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

molcraft-0.1.0a6.dist-info/RECORD

@@ -1,19 +0,0 @@
-molcraft/__init__.py,sha256=eKMk4e5Wki4Ay7_BUuY7B-j3Po2l5FDDefPjkFQw3OM,463
-molcraft/callbacks.py,sha256=mkz4ALjJFPy8nHd2nCAuMbKceKnq4tIpZhUuUOvie2Y,1209
-molcraft/chem.py,sha256=apaECcQSuAMs3Tm12yc6ne4x0BGx5JzfoRhTC1WMhlI,20695
-molcraft/conformers.py,sha256=rojo8OaZrKAesx0JA5kf-JVNEpmsQyLSpcxbWhV9cd4,4324
-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=aJJibnHCxvSu3bNbE2xQk34QvFb47Mnm__0MxlRLA0w,27323
-molcraft/layers.py,sha256=RyKmdHmHlYJJL15LvHH32daTKsChJ_pHmHUnpUcwS1U,73437
-molcraft/losses.py,sha256=JEKZEX2f8vDgky_fUocsF8vZjy9VMzRjZUBa20Uf9Qw,1065
-molcraft/models.py,sha256=Rl9CkQlOVkj20TLjGlwI8vaQwX07EqqWz22bFYtJlpk,22636
-molcraft/ops.py,sha256=eAi79aawJwxuIVVamjA1kPRHGlUm0PsvN-7d2CYu15I,4441
-molcraft/records.py,sha256=0sjOdcr266ZER4F-aTBQ3AVPNAwflKWNiNJVsSc1-PQ,5370
-molcraft/tensors.py,sha256=EOUKx496KUZsjA1zA2ABc7tU_TW3Jv7AXDsug_QsLbA,22407
-molcraft-0.1.0a6.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a6.dist-info/METADATA,sha256=Zzl1K3WleDp056zbLChy5B1AQ3U26t22oMkIKRUpbMY,4063
-molcraft-0.1.0a6.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-molcraft-0.1.0a6.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a6.dist-info/RECORD,,