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

molcraft/losses.py

@@ -0,0 +1,36 @@
+import keras 
+import numpy as np
+
+
+@keras.saving.register_keras_serializable(package='molgraph')
+class GaussianNegativeLogLikelihood(keras.losses.Loss):
+
+    def __init__(
+        self, 
+        events: int = 1, 
+        name="gaussian_nll", 
+        **kwargs
+    ):
+        super().__init__(name=name, **kwargs)
+        self.events = events
+    
+    def call(self, y_true, y_pred):
+        mean = y_pred[..., :self.events]
+        scale = y_pred[..., self.events:]
+        variance = keras.ops.square(scale)
+        expected_rank = len(keras.ops.shape(mean))
+        current_rank = len(keras.ops.shape(y_true))
+        for _ in range(expected_rank - current_rank):
+            y_true = keras.ops.expand_dims(y_true, axis=-1)
+        return keras.ops.mean(
+            0.5 * keras.ops.log(2.0 * np.pi * variance) + 
+            0.5 * keras.ops.square(y_true - mean) / variance 
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config['events'] = self.events 
+        return config 
+    
+
+GaussianNLL = GaussianNegativeLogLikelihood

molcraft/models.py

@@ -17,7 +17,70 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
 
     Currently, the `GraphModel` only supports `GraphTensor` input.
 
-    Example (using `from_layers`):
+    Build a subclassed GraphModel:
+
+    >>> import molcraft
+    >>> import keras
+    >>>
+    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer()
+    >>> graph = featurizer([('N[C@@H](C)C(=O)O', 1.0), ('N[C@@H](CS)C(=O)O', 2.0)])
+    >>>
+    >>> @keras.saving.register_keras_serializable()
+    >>> class GraphNeuralNetwork(molcraft.models.GraphModel):
+    ...     def __init__(self, units, **kwargs):
+    ...         super().__init__(**kwargs)
+    ...         self.units = units
+    ...         self.node_embedding = molcraft.layers.NodeEmbedding(self.units)
+    ...         self.edge_embedding = molcraft.layers.EdgeEmbedding(self.units)
+    ...         self.conv_1 = molcraft.layers.GraphTransformer(self.units)
+    ...         self.conv_2 = molcraft.layers.GraphTransformer(self.units)
+    ...         self.readout = molcraft.layers.Readout('mean')
+    ...         self.dense = keras.layers.Dense(1)
+    ...     def propagate(self, graph):
+    ...         x = self.edge_embedding(self.node_embedding(graph))
+    ...         x = self.conv_2(self.conv_1(x))
+    ...         return self.dense(self.readout(x))
+    ...     def get_config(self):
+    ...         config = super().get_config()
+    ...         config['units'] = self.units
+    ...         return config
+    >>>
+    >>> model = GraphNeuralNetwork(128)
+    >>> model.compile(
+    ...     optimizer=keras.optimizers.Adam(1e-3),
+    ...     loss=keras.losses.MeanSquaredError(),
+    ...     metrics=[keras.metrics.MeanAbsolutePercentageError(name='mape')]
+    ... )
+    >>> model.fit(graph, epochs=10)
+    >>> mse, mape = model.evaluate(graph)
+    >>> preds = model.predict(graph)
+
+    Build a functional GraphModel:
+
+    >>> import molcraft
+    >>> import keras
+    >>>
+    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer()
+    >>> graph = featurizer([('N[C@@H](C)C(=O)O', 1.0), ('N[C@@H](CS)C(=O)O', 2.0)])
+    >>>
+    >>> inputs = molcraft.layers.Input(graph.spec)
+    >>> x = molcraft.layers.NodeEmbedding(128)(inputs)
+    >>> x = molcraft.layers.EdgeEmbedding(128)(x)
+    >>> x = molcraft.layers.GraphTransformer(128)(x)
+    >>> x = molcraft.layers.GraphTransformer(128)(x)
+    >>> x = molcraft.layers.Readout('mean')(x)
+    >>> outputs = keras.layers.Dense(1)(x)
+    >>> model = molcraft.models.GraphModel(inputs, outputs)
+    >>> model.compile(
+    ...     optimizer=keras.optimizers.Adam(1e-3),
+    ...     loss=keras.losses.MeanSquaredError(),
+    ...     metrics=[keras.metrics.MeanAbsolutePercentageError(name='mape')]
+    ... )
+    >>> model.fit(graph, epochs=10)
+    >>> mse, mape = model.evaluate(graph)
+    >>> preds = model.predict(graph)
+
+    Build a GraphModel using `from_layers`:
 
     >>> import molcraft 
     >>> import keras
@@ -25,14 +88,14 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
     >>> featurizer = molcraft.featurizers.MolGraphFeaturizer()
     >>> graph = featurizer([('N[C@@H](C)C(=O)O', 1.0), ('N[C@@H](CS)C(=O)O', 2.0)])
     >>> 
-    >>> model = molcraft.models.GraphModel.from_layers(
+    >>> model = molcraft.models.GraphModel.from_layers([
     ...     molcraft.layers.Input(graph.spec),
     ...     molcraft.layers.NodeEmbedding(128),
     ...     molcraft.layers.EdgeEmbedding(128),
     ...     molcraft.layers.GraphTransformer(128),
     ...     molcraft.layers.GraphTransformer(128),
     ...     molcraft.layers.Readout('mean'),
-    ...     molcraft.layers.Dense(1)
+    ...     keras.layers.Dense(1)
     ... ])
     >>> model.compile(
     ...     optimizer=keras.optimizers.Adam(1e-3),
@@ -42,6 +105,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
     >>> model.fit(graph, epochs=10)
     >>> mse, mape = model.evaluate(graph)
     >>> preds = model.predict(graph)
+
     """
 
     def __new__(cls, *args, **kwargs):
@@ -270,7 +334,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         """
         super().load_weights(filepath, *args, **kwargs)
 
-    def embedding(self) -> 'FunctionalGraphModel':
+    def embedding(self, layer_name: str = None) -> 'FunctionalGraphModel':
         model = self
         if not isinstance(model, FunctionalGraphModel):
             raise ValueError(
@@ -278,11 +342,58 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
                 'it needs to be a `FunctionalGraphModel`. '
             )
         inputs = model.input 
-        for layer in model.layers:
-            if isinstance(layer, layers.Readout):
-                outputs = layer.output 
+        if not layer_name:
+            for layer in model.layers:
+                if isinstance(layer, layers.Readout):
+                    outputs = layer.output 
+        else:
+            layer = model.get_layer(layer_name)
+            outputs = (
+                layer.output if isinstance(layer, keras.layers.Layer) else None
+            )
+            if outputs is None:
+                raise ValueError(
+                    f'Could not find `{layer_name}` or '
+                    f'`{layer_name} is not a `keras.layers.Layer`.'
+                )
         return self.__class__(inputs, outputs, name=f'{self.name}_embedding')
-    
+
+    def backbone(self) -> 'FunctionalGraphModel':
+        if not isinstance(self, FunctionalGraphModel):
+            raise ValueError(
+                'Currently, to extract the backbone part of the model, '
+                'it needs to be a `FunctionalGraphModel`, with a `Readout` '
+                'layer dividing the backbone and the head part of the model.'
+            )
+        inputs = self.input
+        outputs = None
+        for layer in self.layers:
+            if isinstance(layer, layers.Readout):
+                outputs = layer.output
+        if outputs is None:
+            raise ValueError(
+                'Could not extract output. `Readout` layer not found.'
+            )
+        return self.__class__(inputs, outputs, name=f'{self.name}_head')
+
+    def head(self) -> functional.Functional:
+        if not isinstance(self, FunctionalGraphModel):
+            raise ValueError(
+                'Currently, to extract the head part of the model, '
+                'it needs to be a `FunctionalGraphModel`, with a `Readout` '
+                'layer dividing the backbone and the head part of the model.'
+            )
+        inputs = None
+        for layer in self.layers:
+            if isinstance(layer, layers.Readout):
+                inputs = layer.output
+        if inputs is None:
+            raise ValueError(
+                'Could not extract input. `Readout` layer not found.'
+            )
+        outputs = layer.output
+        return keras.models.Model(inputs, outputs, name=f'{self.name}_head')
+
     def train_step(self, tensor: tensors.GraphTensor) -> dict[str, float]:
         y = tensor.context.get('label')
         sample_weight = tensor.context.get('weight')

molcraft/ops.py

@@ -82,6 +82,16 @@ def edge_softmax(
     denominator = gather(denominator, edge_target)
     return numerator / denominator
 
+def edge_weight(
+    edge: tf.Tensor,
+    edge_weight: tf.Tensor,
+) -> tf.Tensor:
+    expected_rank = len(keras.ops.shape(edge))
+    current_rank = len(keras.ops.shape(edge_weight))
+    for _ in range(expected_rank - current_rank):
+        edge_weight = keras.ops.expand_dims(edge_weight, axis=-1)
+    return edge * edge_weight
+
 def segment_mean(
     data: tf.Tensor,
     segment_ids: tf.Tensor,

molcraft/records.py

@@ -36,7 +36,10 @@ def write(
         if not isinstance(inputs, list):
             inputs = list(inputs)
 
-        example = _featurize_input(inputs[0], featurizer)
+        example = inputs[0]
+        if isinstance(example, (list, np.ndarray)):
+            example = tuple(example)
+        example = featurizer(example)
         if not isinstance(example, tensors.GraphTensor):
             example = example[0]
 
@@ -46,7 +49,7 @@ def write(
             num_processes = mp.cpu_count()
 
         if num_files is None:
-            num_files = min(len(inputs), num_processes)
+            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)
@@ -88,7 +91,7 @@ def write(
         for process in processes:
             process.join()         
     
-def load(
+def read(
     path: str, 
     shuffle_files: bool = False
 ) -> tf.data.Dataset:
@@ -107,13 +110,28 @@ def load(
         ds = ds.unbatch()
     return ds
 
+def save_spec(path: str, spec: tensors.GraphTensor.Spec) -> None:
+    proto = spec.experimental_as_proto()
+    with open(path, 'wb') as fh:
+        fh.write(proto.SerializeToString())
+
+def load_spec(path: str) -> tensors.GraphTensor.Spec:
+    with open(path, 'rb') as fh:
+        serialized_proto = fh.read()
+    spec = tensors.GraphTensor.Spec.experimental_from_proto(
+        tensors.GraphTensor.Spec
+        .experimental_type_proto()
+        .FromString(serialized_proto)
+    )
+    return spec
+    
 def _write_tfrecord(
     inputs, 
     path: str,
     featurizer: featurizers.Featurizer, 
 ) -> None:
     
-    def write_example(tensor):
+    def _write_example(tensor):
         flat_values = tf.nest.flatten(tensor, expand_composites=True)
         flat_values = [tf.io.serialize_tensor(value).numpy() for value in flat_values]
         feature = tf.train.Feature(bytes_list=tf.train.BytesList(value=flat_values))
@@ -122,17 +140,15 @@ def _write_tfrecord(
 
     with tf.io.TFRecordWriter(path) as writer:
         for x in inputs:
-            tensor = _featurize_input(x, featurizer)
-            if isinstance(tensor, tensors.GraphTensor):
-                write_example(tensor)
-            else:
-                for t in tensor:
-                    write_example(t)
-
-def _featurize_input(x, featurizer):
-    if isinstance(x, (list, np.ndarray)):
-        x = tuple(x)
-    return featurizer(x)
+            if isinstance(x, (list, np.ndarray)):
+                x = tuple(x)
+            tensor = featurizer(x)
+            if tensor is not None:
+                if isinstance(tensor, tensors.GraphTensor):
+                    _write_example(tensor)
+                else:
+                    for t in tensor:
+                        _write_example(t)
 
 def _serialize_example(
     feature: dict[str, tf.train.Feature]
@@ -151,19 +167,4 @@ def _parse_example(
         tf.split(out, len(tf.nest.flatten(spec, expand_composites=True))), 
         tf.nest.flatten(spec, expand_composites=True))]
     out = tf.nest.pack_sequence_as(spec, tf.nest.flatten(out), expand_composites=True)
-    return out
-
-def save_spec(path: str, spec: tensors.GraphTensor.Spec) -> None:
-    proto = spec.experimental_as_proto()
-    with open(path, 'wb') as fh:
-        fh.write(proto.SerializeToString())
-
-def load_spec(path: str) -> tensors.GraphTensor.Spec:
-    with open(path, 'rb') as fh:
-        serialized_proto = fh.read()
-    spec = tensors.GraphTensor.Spec.experimental_from_proto(
-        tensors.GraphTensor.Spec
-        .experimental_type_proto()
-        .FromString(serialized_proto))
-    return spec
-    
+    return out

molcraft/tensors.py

@@ -224,7 +224,7 @@ class GraphTensor(tf.experimental.BatchableExtensionType):
             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']
+            edge_attr = ops.edge_weight(edge_attr, self.edge['weight'])
         return ops.aggregate(edge_attr, self.edge[edge_type], self.num_nodes, mode=mode)
     
     def propagate(self, add_edge_feature: bool = False):

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a4
+Version: 0.1.0a6
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -55,7 +55,7 @@ Dynamic: license-file
 - 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**
+- Fast and efficient input pipelines using TF **records**
 
 ## Examples 
 
@@ -106,20 +106,7 @@ print(pred)
 # featurizers.save_featurizer(featurizer, '/tmp/featurizer.json')
 # models.save_model(model, '/tmp/model.keras')
 
-# featurizers.load_featurizer('/tmp/featurizer.json')
-# models.load_model('/tmp/model.keras')
+# loaded_featurizer = featurizers.load_featurizer('/tmp/featurizer.json')
+# loaded_model = models.load_model('/tmp/model.keras')
 ```
 
-## Installation
-
-Install the pre-release of molcraft via pip:
-
-```bash
-pip install molcraft --pre
-```
-
-with GPU support:
-
-```bash
-pip install molcraft[gpu] --pre
-```

molcraft-0.1.0a6.dist-info/RECORD

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

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

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

molcraft/experimental/__init__.py

@@ -1 +0,0 @@
-from molcraft.experimental import peptides