molcraft 0.1.0rc9__py3-none-any.whl

molcraft/trainers.py

@@ -0,0 +1,212 @@
+import keras
+import warnings
+
+from molcraft import layers
+from molcraft import models
+from molcraft import tensors 
+
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class Trainer(models.GraphModel):
+
+    '''Base trainer.
+
+    Wraps and (pre)trains a graph neural network for a certain task.
+
+    Args:
+        model: 
+            A `models.GraphModel` to be (pre)trained.
+    '''
+    
+    def __init__(self, model: models.GraphModel, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.model = model 
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config['model'] = keras.saving.serialize_keras_object(self.model)
+        return config
+    
+    @classmethod
+    def from_config(cls, config: dict) -> 'Trainer':
+        config['model'] = keras.saving.deserialize_keras_object(config['model'])
+        return super().from_config(config)
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class NodePredictionTrainer(Trainer):
+
+    '''Node prediction trainer.
+
+    Wraps and (pre)trains a graph neural network to perform node predictions.
+
+    Ignores super nodes and edges, if they exist.
+
+    Args:
+        model: 
+            A `models.GraphModel` to be (pre)trained.
+        decoder:
+            An optional decoder for converting updated node features to node predictions.
+            If None, a `keras.layers.Dense` layer is used with `units` set to `label` dim.
+        select_rate:
+            The rate of which nodes will be selected for prediction. If None, all nodes are predicted.
+        mask_selected:
+            Whether to mask the selected nodes. Only relevant if `select_rate` is specified.
+        edge_masking_rate:
+            The rate of which edges will be masked. If None, edges will not be masked.
+            Only relevant if `select_rate` is specified.
+
+    Example:
+
+    >>> import molcraft
+    >>> import keras
+    >>> 
+    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer(
+    ...     atom_features=[
+    ...         molcraft.features.AtomType(['C', 'N', 'O', 'P', 'S']),
+    ...     ]
+    ... )
+    >>> graph = featurizer(['N[C@@H](C)C(=O)O', 'N[C@@H](CS)C(=O)O'])
+    >>> # Label nodes with the one-hot encoded atom types for illustration
+    >>> graph = graph.update({'node': {'label': graph.node['feature']}})
+    >>> 
+    >>> inputs = molcraft.layers.Input(graph.spec)
+    >>> x = molcraft.layers.NodeEmbedding(128)(inputs)
+    >>> x = molcraft.layers.EdgeEmbedding(128)(x)
+    >>> x = molcraft.layers.GraphConv(128)(x)
+    >>> outputs = molcraft.layers.GraphConv(128)(x)
+    >>> model = molcraft.models.GraphModel(inputs, outputs)
+    >>> 
+    >>> pretrainer = molcraft.trainers.NodePredictionTrainer(
+    ...     model,
+    ...     decoder=None,       # Dense(units=node_label_dim)
+    ...     select_rate=0.5,
+    ...     mask_selected=True,
+    ... )
+    >>> pretrainer.compile(
+    ...     optimizer=keras.optimizers.Adam(1e-4),
+    ...     loss=keras.losses.CategoricalCrossentropy(from_logits=True),
+    ... )
+    >>> pretrainer.fit(graph, epochs=10)
+    >>> # pretrainer.model.save('/tmp/model.keras')
+    '''
+
+    def __init__(
+        self,
+        model: models.GraphModel,
+        decoder: keras.layers.Layer | None = None,
+        select_rate: float | None = None,
+        mask_selected: bool = False,
+        edge_masking_rate: float | None = None,
+        **kwargs
+    ) -> None:
+        super().__init__(model=model, **kwargs)
+        
+        for layer in self.model.layers:
+            if isinstance(layer, layers.NodeEmbedding):
+                break
+        else:
+            raise ValueError('Could not find `NodeEmbedding` layer.')
+        
+        self._embedder = models.GraphModel(
+            self.model.input, layer._symbolic_output
+        )
+        self._model = models.GraphModel(
+            layer._symbolic_output, self.model.output
+        )
+        self._decoder = decoder
+        self._select_rate = select_rate
+        self._mask_selected = mask_selected
+        if edge_masking_rate and not mask_selected:
+            warnings.warn(
+                'Setting `edge_masking_rate` to `None`, '
+                'as `mask_selected` is set to `False`.'
+            )
+            edge_masking_rate = None
+        self._edge_masking_rate = edge_masking_rate
+    
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
+
+        self._has_super = ('super' in spec.node)
+        self._has_edge_feature = ('feature' in spec.edge)
+
+        if self._mask_selected:
+            node_feature_dim = self._embedder._symbolic_output['node']['feature'].shape[-1]
+            self._node_mask_feature = self.get_weight(shape=[node_feature_dim])
+
+        if self._mask_selected and self._has_edge_feature and self._edge_masking_rate:
+            edge_feature_dim = self._embedder._symbolic_output['edge']['feature'].shape[-1]
+            self._edge_mask_feature = self.get_weight(shape=[edge_feature_dim])
+        elif self._edge_masking_rate and not self._has_edge_feature:
+            warnings.warn(
+                'Setting `edge_masking_rate` to `None`, '
+                'as no edge features exist.'
+            )
+            self._edge_masking_rate = None
+
+        if self._decoder is None:
+            label_dim = spec.node['label'].shape[-1]
+            self._decoder = keras.layers.Dense(units=label_dim) 
+
+    def propagate(
+        self, 
+        tensor: tensors.GraphTensor,
+        training: bool | None = None,
+    ) -> tensors.GraphTensor:
+        sample_weight = tensor.node.get('sample_weight')
+        if sample_weight is None:
+            sample_weight = keras.ops.ones([tensor.num_nodes])
+
+        tensor = self._embedder(tensor)
+
+        if self._select_rate is not None and training:
+            # Select nodes to be predicted
+            is_not_super = (
+                True if not self._has_super else keras.ops.logical_not(tensor.node['super'])
+            )
+            r = keras.random.uniform(shape=[tensor.num_nodes])
+            node_mask = keras.ops.logical_and(is_not_super, self._select_rate > r)
+            sample_weight = keras.ops.where(node_mask, sample_weight, 0.0)
+
+            if self._mask_selected:
+                # Mask selected node features
+                node_feature_masked = keras.ops.where(
+                    node_mask[:, None], self._node_mask_feature, tensor.node['feature']
+                )
+                tensor = tensor.update({'node': {'feature': node_feature_masked}})
+
+                if self._edge_masking_rate:
+                    # Mask edge features
+                    is_not_super = (
+                        True if not self._has_super else keras.ops.logical_not(tensor.edge['super'])
+                    )
+                    r = keras.random.uniform(shape=[tensor.num_edges])
+                    edge_mask = keras.ops.logical_and(is_not_super, self._edge_masking_rate > r)
+                    edge_feature_masked = keras.ops.where(
+                        edge_mask[:, None], self._edge_mask_feature, tensor.edge['feature']
+                    )
+                    tensor = tensor.update({'edge': {'feature': edge_feature_masked}})
+
+        node_feature = self._model(tensor).node['feature']
+        node_prediction = self._decoder(node_feature)
+        return tensor.update({
+            'node': {
+                'prediction': node_prediction, 
+                'sample_weight': sample_weight
+            }
+        })
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            'decoder': keras.saving.serialize_keras_object(self._decoder),
+            'select_rate': self._select_rate,
+            'mask_selected': self._mask_selected,
+            'edge_masking_rate': self._edge_masking_rate,
+        })
+        return config
+    
+    @classmethod
+    def from_config(cls, config: dict) -> 'NodePredictionTrainer':
+        config['decoder'] = keras.saving.deserialize_keras_object(config['decoder'])
+        return super().from_config(config)

molcraft-0.1.0rc9.dist-info/METADATA

@@ -0,0 +1,118 @@
+Metadata-Version: 2.4
+Name: molcraft
+Version: 0.1.0rc9
+Summary: Graph Neural Networks for Molecular Machine Learning
+Author-email: Alexander Kensert <alexander.kensert@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2025 Alexander Kensert
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+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
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: tensorflow>=2.16
+Requires-Dist: rdkit>=2023.9.5
+Requires-Dist: pandas>=1.0.3
+Requires-Dist: ipython>=8.12.0
+Provides-Extra: gpu
+Requires-Dist: tensorflow[and-cuda]>=2.16; extra == "gpu"
+Dynamic: license-file
+
+<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo" width="90%">
+
+**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML.
+
+> [!NOTE]  
+> In progress.
+
+## Installation
+
+For CPU users:
+
+```bash
+pip install molcraft
+```
+
+For GPU users:
+```bash
+pip install molcraft[gpu]
+```
+
+## Examples 
+
+```python
+from molcraft import features
+from molcraft import descriptors
+from molcraft import featurizers 
+from molcraft import layers
+from molcraft import models 
+import keras
+
+featurizer = featurizers.MolGraphFeaturizer(
+    atom_features=[
+        features.AtomType(),
+        features.NumHydrogens(),
+        features.Degree(),
+    ],
+    bond_features=[
+        features.BondType(),
+        features.IsRotatable(),
+    ],
+    super_node=True,
+    self_loops=True,
+    include_hydrogens=False,
+)
+
+graph = featurizer([('N[C@@H](C)C(=O)O', 2.5), ('N[C@@H](CS)C(=O)O', 1.5)])
+print(graph)
+
+model = models.GraphModel.from_layers(
+    [
+        layers.Input(graph.spec),
+        layers.NodeEmbedding(dim=128),
+        layers.EdgeEmbedding(dim=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.Readout(),
+        keras.layers.Dense(units=1024, activation='elu'),
+        keras.layers.Dense(units=1024, activation='elu'),
+        keras.layers.Dense(1)
+    ]
+)
+
+pred = model(graph)
+print(pred)
+
+# featurizers.save_featurizer(featurizer, '/tmp/featurizer.json')
+# models.save_model(model, '/tmp/model.keras')
+
+# loaded_featurizer = featurizers.load_featurizer('/tmp/featurizer.json')
+# loaded_model = models.load_model('/tmp/model.keras')
+```
+

molcraft-0.1.0rc9.dist-info/RECORD

@@ -0,0 +1,19 @@
+molcraft/__init__.py,sha256=_1FKaZ23km7d3YPnOGSx3EMeXxjEwSuC30YhWjqKOXI,461
+molcraft/callbacks.py,sha256=B4gGWjVW_1ORrt38jfk1ZFI9c0rOpN5sgjGWVqs3Ess,3571
+molcraft/chem.py,sha256=dPRB4aLk5U6nkzfHCTHosh6f7Cph16UH3Ri4JNCE6fw,22488
+molcraft/datasets.py,sha256=1rHccqra5chIBwo2pz9vduyv0i07uY3CABzmAqWiFBU,4161
+molcraft/descriptors.py,sha256=FI15LYcb0KXqvurCkXUkg_h7rOqRZnRPK5LaWdM7Q8M,4876
+molcraft/features.py,sha256=q-wuRP9YjPu_v5czipsh00VEXEjgFaeuLk6dbgyD_VM,13505
+molcraft/featurizers.py,sha256=_1xq6336s9E2Ho5f3v_dIGLop1dhKf1lscygpeIlVQY,20964
+molcraft/layers.py,sha256=4E8JdI3xLjrI6hvIbHjE23O9PsB7J3H7NIcaaLsUDZM,69879
+molcraft/losses.py,sha256=piu4XYAgjnK7k9LqA4Vkh-SooYZ31sWwRfG1cacCwyA,1081
+molcraft/models.py,sha256=4_658wqfB1rz4QS_68aNr1RWB2cyqwcpNr_k0rvmuj0,23756
+molcraft/ops.py,sha256=cI373Bg51CXXKnn7vhH1GiZ2GrvkgzaG38fhlPmfIs0,6280
+molcraft/records.py,sha256=b9i1azDM_OkY_E4Rl6etwbzOthekVM-Z7oTSxYg5aNM,5888
+molcraft/tensors.py,sha256=KSXAPKhu4kXn6KUmu2ZNE1-CSfmoFDdAPiaxYyIAHUA,22693
+molcraft/trainers.py,sha256=-s3SdBUI2Hfvf3D3JLRAGpPSy0umUlhI5MPIg3eA9Jc,8018
+molcraft-0.1.0rc9.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0rc9.dist-info/METADATA,sha256=3CDa5m0sVaWpnmo5CQYyKPbaCqXuv3RTD9AexLT4rFI,3880
+molcraft-0.1.0rc9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+molcraft-0.1.0rc9.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0rc9.dist-info/RECORD,,

molcraft-0.1.0rc9.dist-info/WHEEL

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

molcraft-0.1.0rc9.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Alexander Kensert
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

molcraft-0.1.0rc9.dist-info/top_level.txt

@@ -0,0 +1 @@
+molcraft