molcraft 0.1.0a10__tar.gz → 0.1.0a12__tar.gz

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a10
+Version: 0.1.0a12
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -47,15 +47,20 @@ Dynamic: license-file
 **Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML. 
 
 > [!NOTE]  
-> In progress/Unfinished.
+> In progress.
 
-## Highlights
-- Compatible with **Keras 3**
-- Customizable and serializable **featurizers**
-- Customizable and serializable **layers** and **models**
-- Customizable **GraphTensor**
-- Fast and efficient featurization of molecular graphs
-- Fast and efficient input pipelines using TF **records**
+## Installation
+
+For CPU users:
+
+```bash
+pip install --pre molcraft
+```
+
+For GPU users:
+```bash
+pip install --pre molcraft[gpu]
+```
 
 ## Examples 
 
@@ -70,7 +75,7 @@ import keras
 featurizer = featurizers.MolGraphFeaturizer(
     atom_features=[
         features.AtomType(),
-        features.TotalNumHs(),
+        features.NumHydrogens(),
         features.Degree(),
     ],
     bond_features=[

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/README.md

@@ -3,15 +3,20 @@
 **Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML. 
 
 > [!NOTE]  
-> In progress/Unfinished.
+> In progress.
 
-## Highlights
-- Compatible with **Keras 3**
-- Customizable and serializable **featurizers**
-- Customizable and serializable **layers** and **models**
-- Customizable **GraphTensor**
-- Fast and efficient featurization of molecular graphs
-- Fast and efficient input pipelines using TF **records**
+## Installation
+
+For CPU users:
+
+```bash
+pip install --pre molcraft
+```
+
+For GPU users:
+```bash
+pip install --pre molcraft[gpu]
+```
 
 ## Examples 
 
@@ -26,7 +31,7 @@ import keras
 featurizer = featurizers.MolGraphFeaturizer(
     atom_features=[
         features.AtomType(),
-        features.TotalNumHs(),
+        features.NumHydrogens(),
         features.Degree(),
     ],
     bond_features=[

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/__init__.py

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

molcraft-0.1.0a12/molcraft/apps/peptides.py

@@ -0,0 +1,429 @@
+import re
+import keras
+import numpy as np
+import tensorflow as tf
+import tensorflow_text as tf_text
+from rdkit import Chem
+
+from molcraft import ops
+from molcraft import chem 
+from molcraft import features 
+from molcraft import featurizers
+from molcraft import tensors
+from molcraft import descriptors
+from molcraft import layers
+from molcraft import models 
+
+
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class SequenceSplitter(keras.layers.Layer):
+
+    _pattern = "|".join([
+        r'(\[[A-Za-z0-9]+\]-[A-Z]\[[A-Za-z0-9]+\])', # N-term mod + mod
+        r'([A-Z]\[[A-Za-z0-9]+\]-\[[A-Za-z0-9]+\])', # C-term mod + mod
+        r'([A-Z]-\[[A-Za-z0-9]+\])', # C-term mod
+        r'(\[[A-Za-z0-9]+\]-[A-Z])', # N-term mod
+        r'([A-Z]\[[A-Za-z0-9]+\])', # Mod
+        r'([A-Z])', # No mod
+    ])
+
+    def call(self, inputs):
+        inputs = tf_text.regex_split(inputs, self._pattern, self._pattern)
+        inputs = keras.ops.concatenate([
+            tf.strings.join([inputs[:, :-1], '-[X]']),
+            inputs[:, -1:]
+        ], axis=1)
+        return inputs.to_tensor()
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class Gather(keras.layers.Layer):
+
+    def __init__(
+        self, 
+        padding: list[tuple[int]] | tuple[int] | int = 1, 
+        mask_value: int = 0,
+        **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+        self._splitter = SequenceSplitter()
+        self.padding = padding
+        self.mask_value = mask_value
+        self.supports_masking = True
+
+        self._tags = list(sorted(residues.keys()))
+        self._mapping = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                keys=self._tags,
+                values=range(len(self._tags)),
+            ),
+            default_value=-1,
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config['mask_value'] = self.mask_value 
+        config['padding'] = self.padding
+        return config 
+
+    def call(self, inputs) -> tf.Tensor:
+        embedding, sequence = inputs
+        sequence = self._splitter(sequence)
+        sequence = self._mapping.lookup(sequence)
+        readout = ops.gather(embedding, keras.ops.where(sequence == -1, 0, sequence)) 
+        readout = keras.ops.where(sequence[..., None] == -1, 0.0, readout)
+        return readout
+
+    def compute_mask(
+        self, 
+        inputs: tensors.GraphTensor, 
+        mask: bool | None = None
+    ) -> tf.Tensor | None:
+        # if self.mask_value is None:
+        #     return None
+        _, sequence = inputs
+        sequence = self._splitter(sequence)
+        return keras.ops.not_equal(sequence, '')
+
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class Embedding(keras.layers.Layer):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        tags = list(sorted(residues.keys()))
+        self.mapping = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                keys=tags,
+                values=range(len(tags)),
+            ),
+            default_value=-1,
+        )
+        self.splitting = SequenceSplitter()
+        featurizer = featurizers.MolGraphFeaturizer(super_atom=True)
+        tensor_list = [featurizer(residues[tag]) for tag in tags]
+        graph = tf.stack(tensor_list, axis=0)
+        self._build_on_init(graph)
+        self.embedder = models.GraphModel.from_layers(
+            [
+                layers.Input(graph.spec),
+                layers.NodeEmbedding(128),
+                layers.EdgeEmbedding(128),
+                layers.GraphTransformer(128),
+                layers.Readout()
+            ]
+        )
+        self.embedding = tf.Variable(
+            initial_value=tf.zeros((114, 128)), trainable=True
+        )
+        self.new_state = tf.Variable(True, dtype=tf.bool, trainable=False)
+        self.gather = Gather()
+        self.update_state()
+
+        # Keep AA as is (most simple?), add positional embedding to distingusih N-, C- and non-terminal
+
+    def update_state(self, inputs=None):
+        graph = self._graph_tensor
+        graph = tensors.to_dict(graph)
+        embedding = self.embedder(graph)
+        self.embedding.assign(embedding)
+        tf.print("STATE UPDATED")
+        return embedding
+
+    def call(self, inputs=None, training=None) -> tensors.GraphTensor:
+        if training:
+            embedding = self.update_state()
+            self.new_state.assign(True)
+            return self.gather([embedding, inputs])
+        else:
+            embedding = tf.cond(
+                pred=self.new_state,
+                true_fn=lambda: self.update_state(),
+                false_fn=lambda: self.embedding
+            )   
+            self.new_state.assign(False)
+            return self.gather([embedding, inputs])
+    
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def _build_on_init(self, x):
+
+        if isinstance(x, tensors.GraphTensor):
+            tensor = tensors.to_dict(x)
+            self._spec = tf.nest.map_structure(
+                tf.type_spec_from_value, tensor
+            )
+        else:
+            self._spec = x
+
+        self._graph = tf.nest.map_structure(
+            lambda s: self.add_weight(
+                shape=s.shape, 
+                dtype=s.dtype, 
+                trainable=False,
+                initializer='zeros'
+            ),
+            self._spec
+        )
+
+        if isinstance(x, tensors.GraphTensor):
+            tf.nest.map_structure(
+                lambda v, x: v.assign(x),
+                self._graph, tensor
+            )
+
+        graph = tf.nest.map_structure(
+            keras.ops.convert_to_tensor, self._graph
+        )
+        self._graph_tensor = tensors.from_dict(graph)
+
+    # def get_config(self) -> dict:
+    #     config = super().get_config()
+    #     spec = keras.saving.serialize_keras_object(self._spec)
+    #     config['spec'] = spec
+    #     #config['layers'] = keras.saving.serialize_keras_object(self.embedding.layers)
+    #     return config 
+    
+    # @classmethod
+    # def from_config(cls, config: dict) -> 'SequenceToGraph':
+    #     spec = config.pop('spec')
+    #     spec = keras.saving.deserialize_keras_object(spec)
+    #    # config['layers'] = keras.saving.deserialize_keras_object(config['layers'])
+    #     layer = cls(**config)
+    #     layer._build_on_init(spec)
+    #     return layer
+
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class SequenceToGraph(keras.layers.Layer):
+
+    def __init__(
+        self, 
+        atom_features: list[features.Feature] | str | None = 'auto',
+        bond_features: list[features.Feature] | str | None = 'auto',
+        molecule_features: list[descriptors.Descriptor] | str | None = 'auto',
+        super_atom: bool = True,
+        radius: int | float | None = None,
+        self_loops: bool = False,
+        include_hs: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._splitter = SequenceSplitter()
+        featurizer = featurizers.MolGraphFeaturizer(
+            atom_features=atom_features,
+            bond_features=bond_features,
+            molecule_features=molecule_features,
+            super_atom=super_atom,
+            radius=radius,
+            self_loops=self_loops,
+            include_hs=include_hs,
+            **kwargs,
+        )
+        tensor_list: list[tensors.GraphTensor] = [
+            featurizer(residues[tag]).update({'context': {'tag': tag}}) for tag in residues
+        ]
+        graph = tf.stack(tensor_list, axis=0)
+        self._build_on_init(graph)
+
+    def call(self, sequence: tf.Tensor) -> tensors.GraphTensor:
+        sequence = self._splitter(sequence)
+        indices = self._tag_to_index.lookup(sequence)
+        indices = tf.sort(tf.unique(tf.reshape(indices, [-1]))[0])[1:]
+        graph = self._graph_tensor[indices]
+        return tensors.to_dict(graph)
+
+    def _build_on_init(self, x):
+
+        if isinstance(x, tensors.GraphTensor):
+            tensor = tensors.to_dict(x)
+            self._spec = tf.nest.map_structure(
+                tf.type_spec_from_value, tensor
+            )
+        else:
+            self._spec = x
+
+        self._graph = tf.nest.map_structure(
+            lambda s: self.add_weight(
+                shape=s.shape, 
+                dtype=s.dtype, 
+                trainable=False,
+                initializer='zeros'
+            ),
+            self._spec
+        )
+
+        if isinstance(x, tensors.GraphTensor):
+            tf.nest.map_structure(
+                lambda v, x: v.assign(x),
+                self._graph, tensor
+            )
+
+        graph = tf.nest.map_structure(
+            keras.ops.convert_to_tensor, self._graph
+        )
+        self._graph_tensor = tensors.from_dict(graph)
+
+        tags = self._graph_tensor.context['tag']
+
+        self._tag_to_index = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                keys=tags,
+                values=range(len(tags)),
+            ),
+            default_value=-1,
+        )
+
+    def get_config(self) -> dict:
+        config = super().get_config()
+        spec = keras.saving.serialize_keras_object(self._spec)
+        config['spec'] = spec
+        return config 
+    
+    @classmethod
+    def from_config(cls, config: dict) -> 'SequenceToGraph':
+        spec = config.pop('spec')
+        spec = keras.saving.deserialize_keras_object(spec)
+        layer = cls(**config)
+        layer._build_on_init(spec)
+        return layer
+
+    # @property 
+    # def graph(self) -> tensors.GraphTensor:
+    #     return self._graph_tensor
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class GraphToSequence(keras.layers.Layer):
+
+    def __init__(
+        self, 
+        padding: list[tuple[int]] | tuple[int] | int = 1, 
+        mask_value: int = 0,
+        **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+        self._splitter = SequenceSplitter()
+        self.padding = padding
+        self.mask_value = mask_value
+        self._readout_layer = layers.Readout(mode='mean')
+        self.supports_masking = True
+
+    def get_config(self):
+        config = super().get_config()
+        config['mask_value'] = self.mask_value 
+        config['padding'] = self.padding
+        return config 
+
+    def call(self, inputs) -> tf.Tensor:
+
+        graph, sequence = inputs
+        sequence = self._splitter(sequence)
+        tag = graph['context']['tag']
+        data = self._readout_layer(graph)
+
+        table = tf.lookup.experimental.MutableHashTable(
+            key_dtype=tf.string,
+            value_dtype=tf.int32,
+            default_value=-1
+        )
+
+        table.insert(tag, tf.range(tf.shape(tag)[0]))
+        sequence = table.lookup(sequence)
+
+        readout = ops.gather(data, keras.ops.where(sequence == -1, 0, sequence)) 
+        readout = keras.ops.where(sequence[..., None] == -1, 0.0, readout)
+        return readout
+
+    def compute_mask(
+        self, 
+        inputs: tensors.GraphTensor, 
+        mask: bool | None = None
+    ) -> tf.Tensor | None:
+        # if self.mask_value is None:
+        #     return None
+        _, sequence = inputs
+        sequence = self._splitter(sequence)
+        return keras.ops.not_equal(sequence, '')
+
+
+residues = {
+    "A": "N[C@@H](C)C(=O)O",
+    "C": "N[C@@H](CS)C(=O)O",
+    "C[Carbamidomethyl]": "N[C@@H](CSCC(=O)N)C(=O)O",
+    "D": "N[C@@H](CC(=O)O)C(=O)O",
+    "E": "N[C@@H](CCC(=O)O)C(=O)O",
+    "F": "N[C@@H](Cc1ccccc1)C(=O)O",
+    "G": "NCC(=O)O",
+    "H": "N[C@@H](CC1=CN=C-N1)C(=O)O",
+    "I": "N[C@@H](C(CC)C)C(=O)O",
+    "K": "N[C@@H](CCCCN)C(=O)O",
+    "K[Acetyl]": "N[C@@H](CCCCNC(=O)C)C(=O)O",
+    "K[Crotonyl]": "N[C@@H](CCCCNC(C=CC)=O)C(=O)O",
+    "K[Dimethyl]": "N[C@@H](CCCCN(C)C)C(=O)O",
+    "K[Formyl]": "N[C@@H](CCCCNC=O)C(=O)O",
+    "K[Malonyl]": "N[C@@H](CCCCNC(=O)CC(O)=O)C(=O)O",
+    "K[Methyl]": "N[C@@H](CCCCNC)C(=O)O",
+    "K[Propionyl]": "N[C@@H](CCCCNC(=O)CC)C(=O)O",
+    "K[Succinyl]": "N[C@@H](CCCCNC(CCC(O)=O)=O)C(=O)O",
+    "K[Trimethyl]": "N[C@@H](CCCC[N+](C)(C)C)C(=O)O",
+    "L": "N[C@@H](CC(C)C)C(=O)O",
+    "M": "N[C@@H](CCSC)C(=O)O",
+    "M[Oxidation]": "N[C@@H](CCS(=O)C)C(=O)O",
+    "N": "N[C@@H](CC(=O)N)C(=O)O",
+    "P": "N1[C@@H](CCC1)C(=O)O",
+    "P[Oxidation]": "N1CC(O)C[C@H]1C(=O)O",
+    "Q": "N[C@@H](CCC(=O)N)C(=O)O",
+    "R": "N[C@@H](CCCNC(=N)N)C(=O)O",
+    "R[Deamidated]": "N[C@@H](CCCNC(N)=O)C(=O)O",
+    "R[Dimethyl]": "N[C@@H](CCCNC(N(C)C)=N)C(=O)O",
+    "R[Methyl]": "N[C@@H](CCCNC(=N)NC)C(=O)O",
+    "S": "N[C@@H](CO)C(=O)O",
+    "T": "N[C@@H](C(O)C)C(=O)O",
+    "V": "N[C@@H](C(C)C)C(=O)O",
+    "W": "N[C@@H](CC(=CN2)C1=C2C=CC=C1)C(=O)O",
+    "Y": "N[C@@H](Cc1ccc(O)cc1)C(=O)O",
+    "Y[Nitro]": "N[C@@H](Cc1ccc(O)c(N(=O)=O)c1)C(=O)O",
+    "Y[Phospho]": "N[C@@H](Cc1ccc(OP(O)(=O)O)cc1)C(=O)O",
+    "[Acetyl]-A": "N(C(C)=O)[C@@H](C)C(=O)O",
+    "[Acetyl]-C": "N(C(C)=O)[C@@H](CS)C(=O)O",
+    "[Acetyl]-D": "N(C(=O)C)[C@H](C(=O)O)CC(=O)O",
+    "[Acetyl]-E": "N(C(=O)C)[C@@H](CCC(O)=O)C(=O)O",
+    "[Acetyl]-F": "N(C(C)=O)[C@@H](Cc1ccccc1)C(=O)O",
+    "[Acetyl]-G": "N(C(=O)C)CC(=O)O",
+    "[Acetyl]-H": "N(C(=O)C)[C@@H](Cc1[nH]cnc1)C(=O)O",
+    "[Acetyl]-I": "N(C(=O)C)[C@@H]([C@H](CC)C)C(=O)O",
+    "[Acetyl]-K": "N(C(C)=O)[C@@H](CCCCN)C(=O)O",
+    "[Acetyl]-L": "N(C(=O)C)[C@@H](CC(C)C)C(=O)O",
+    "[Acetyl]-M": "N(C(=O)C)[C@@H](CCSC)C(=O)O",
+    "[Acetyl]-N": "N(C(C)=O)[C@@H](CC(=O)N)C(=O)O",
+    "[Acetyl]-P": "N1(C(=O)C)CCC[C@H]1C(=O)O",
+    "[Acetyl]-Q": "N(C(=O)C)[C@@H](CCC(=O)N)C(=O)O",
+    "[Acetyl]-R": "N(C(C)=O)[C@@H](CCCN=C(N)N)C(=O)O",
+    "[Acetyl]-S": "N(C(C)=O)[C@@H](CO)C(=O)O",
+    "[Acetyl]-T": "N(C(=O)C)[C@@H]([C@H](O)C)C(=O)O",
+    "[Acetyl]-V": "N(C(=O)C)[C@@H](C(C)C)C(=O)O",
+    "[Acetyl]-W": "N(C(C)=O)[C@@H](Cc1c2ccccc2[nH]c1)C(=O)O",
+    "[Acetyl]-Y": "N(C(C)=O)[C@@H](Cc1ccc(O)cc1)C(=O)O"
+}
+
+residues_reverse = {}
+def register_peptide_residues(residues_: dict[str, str], canonicalize=True):
+    for residue, smiles in residues_.items():
+        if canonicalize:
+            smiles = Chem.MolToSmiles(Chem.MolFromSmiles(smiles))
+        residues[residue] = smiles
+        residues_reverse[residues[residue]] = residue
+    
+register_peptide_residues(residues, canonicalize=False)
+
+def _extract_residue_type(residue_tag: str) -> str:
+    pattern = r"(?<!\[)[A-Z](?![^\[]*\])"
+    return [match.group(0) for match in re.finditer(pattern, residue_tag)][0]
+
+special_residues = {}
+for key, value in residues.items():
+    special_residues[key + '-[X]'] = value.rstrip('O')
+
+register_peptide_residues(special_residues, canonicalize=False)

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/chem.py

@@ -426,10 +426,12 @@ def embed_conformers(
     success = rdDistGeom.EmbedMultipleConfs(
         mol, numConfs=num_conformers, params=embedding_method
     )
-    if not len(success):
+    num_successes = len(success)
+    if num_successes < num_conformers:
         warnings.warn(
-            f'Could not embed conformer(s) for {mol.canonical_smiles!r} using the '
-            'speified method. Giving it another try with more permissive methods.',
+            f'Could only embed {num_successes} out of {num_conformers} conformer(s) '
+            f'for {mol.canonical_smiles!r} using {method}. Embedding the remaining '
+            f'{num_conformers - num_successes} conformer(s) using different embedding methods.',
             stacklevel=2
         )
         max_attempts = (20 * mol.num_atoms) # increasing it from 10xN to 20xN
@@ -437,14 +439,16 @@ def embed_conformers(
             fallback_embedding_method = available_embedding_methods[fallback_method]
             fallback_embedding_method.useRandomCoords = True
             fallback_embedding_method.maxAttempts = max_attempts
+            fallback_embedding_method.clearConfs = False
             success = rdDistGeom.EmbedMultipleConfs(
-                mol, numConfs=num_conformers, params=fallback_embedding_method
+                mol, numConfs=(num_conformers - num_successes), params=fallback_embedding_method
             )
-            if len(success):
+            num_successes += len(success)
+            if num_successes == num_conformers:
                 break
         else:
             raise RuntimeError(
-                f'Could not embed conformer(s) for {mol.canonical_smiles!r}. '
+                f'Could not embed {num_conformers} conformer(s) for {mol.canonical_smiles!r}. '
             )
     return mol
 

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/descriptors.py

@@ -61,7 +61,7 @@ class NumHeavyAtoms(Descriptor):
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class NumHeteroAtoms(Descriptor):
+class NumHeteroatoms(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
         return rdMolDescriptors.CalcNumHeteroatoms(mol) 
     

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/features.py

@@ -185,13 +185,13 @@ class Degree(Feature):
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class TotalNumHs(Feature):
+class NumHydrogens(Feature):
     def call(self, mol: chem.Mol) -> list[int, float, str]: 
         return [atom.GetTotalNumHs() for atom in mol.atoms]
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class TotalValence(Feature):
+class Valence(Feature):
     def call(self, mol: chem.Mol) -> list[int, float, str]: 
         return [atom.GetTotalValence() for atom in mol.atoms]
     
@@ -218,10 +218,17 @@ class CIPCode(Feature):
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class IsChiralityPossible(Feature):
+class RingSize(Feature):
     def call(self, mol: chem.Mol) -> list[int, float, str]: 
-        return [atom.HasProp("_ChiralityPossible") for atom in mol.atoms]
-
+        def ring_size(atom):
+            if not atom.IsInRing():
+                return -1
+            size = 3
+            while not atom.IsInRingSize(size):
+                size += 1 
+            return size
+        return [ring_size(atom) for atom in mol.atoms]
+    
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class FormalCharge(Feature):
@@ -229,6 +236,12 @@ class FormalCharge(Feature):
         return [atom.GetFormalCharge() for atom in mol.atoms]
 
 
+@keras.saving.register_keras_serializable(package='molcraft')
+class IsChiralityPossible(Feature):
+    def call(self, mol: chem.Mol) -> list[int, float, str]: 
+        return [atom.HasProp("_ChiralityPossible") for atom in mol.atoms]
+
+
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumRadicalElectrons(Feature):
     def call(self, mol: chem.Mol) -> list[int, float, str]: 
@@ -242,7 +255,7 @@ class IsAromatic(Feature):
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class IsHetero(Feature):
+class IsHeteroatom(Feature):
     def call(self, mol: chem.Mol) -> list[int, float, str]: 
         return chem.hetero_atoms(mol)
 
@@ -259,19 +272,6 @@ class IsHydrogenAcceptor(Feature):
         return chem.hydrogen_acceptors(mol)
 
 
-@keras.saving.register_keras_serializable(package='molcraft')
-class RingSize(Feature):
-    def call(self, mol: chem.Mol) -> list[int, float, str]: 
-        def ring_size(atom):
-            if not atom.IsInRing():
-                return -1
-            size = 3
-            while not atom.IsInRingSize(size):
-                size += 1 
-            return size
-        return [ring_size(atom) for atom in mol.atoms]
-
-
 @keras.saving.register_keras_serializable(package='molcraft')
 class IsInRing(Feature):
     def call(self, mol: chem.Mol) -> list[int, float, str]:         

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/featurizers.py

@@ -196,7 +196,7 @@ class MolGraphFeaturizer(Featurizer):
                 descriptors.CrippenLogP(),
                 descriptors.CrippenMolarRefractivity(),
                 descriptors.NumHeavyAtoms(),
-                descriptors.NumHeteroAtoms(),
+                descriptors.NumHeteroatoms(),
                 descriptors.NumHydrogenDonors(),
                 descriptors.NumHydrogenAcceptors(),
                 descriptors.NumRotatableBonds(),
@@ -556,6 +556,7 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
 
         molecule_feature = self.molecule_feature(mol)
         molecule_size = self.num_atoms(mol) + int(self.super_atom)
+        molecule_size = molecule_size.astype(self.index_dtype)
 
         if isinstance(context, dict):
             if 'x' in context:

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/layers.py

@@ -661,7 +661,7 @@ class GIConv(GraphConv):
         return config
 
 
-@keras.saving.register_keras_serializable(package='molgraphx')
+@keras.saving.register_keras_serializable(package='molcraft')
 class GAConv(GraphConv):
 
     """Graph attention network layer.

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/losses.py

@@ -2,7 +2,7 @@ import keras
 import numpy as np
 
 
-@keras.saving.register_keras_serializable(package='molgraph')
+@keras.saving.register_keras_serializable(package='molcraft')
 class GaussianNegativeLogLikelihood(keras.losses.Loss):
 
     def __init__(

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft/models.py

@@ -114,6 +114,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         return typing.cast(GraphModel, super().__new__(cls))
     
     def __init__(self, *args, **kwargs):
+        self._model_layers = kwargs.pop('model_layers', None)
         super().__init__(*args, **kwargs)
         self.jit_compile = False 
 
@@ -135,10 +136,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
                 `molcraft.layers.Input(spec)`.
         """
         if not tensors.is_graph(graph_layers[0]):
-            # TODO: Allow this. E.g.: return cls(layers=graph_layers)
-            raise ValueError(
-                'Graph input not found. Make sure to add `Input`.'
-            )
+            return cls(model_layers=graph_layers)
         inputs: dict = graph_layers.pop(0)
         x = inputs
         for layer in graph_layers:
@@ -148,6 +146,31 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         outputs = x
         return cls(inputs=inputs, outputs=outputs, **kwargs)
     
+    def propagate(self, graph: tensors.GraphTensor) -> tensors.GraphTensor:
+        if self._model_layers is None:
+            return super().propagate(graph)
+        for layer in self._model_layers:
+            graph = layer(graph)
+        return graph
+    
+    def get_config(self):
+        config = super().get_config()
+        if hasattr(self, '_model_layers') and self._model_layers is not None:
+            config['model_layers'] = [
+                keras.saving.serialize_keras_object(l) 
+                for l in self._model_layers
+            ]
+        return config 
+    
+    @classmethod
+    def from_config(cls, config: dict):
+        if 'model_layers' in config:
+            config['model_layers'] = [
+                keras.saving.deserialize_keras_object(l) 
+                for l in config['model_layers']
+            ]
+        return super().from_config(config)
+
     def compile(
         self, 
         optimizer: keras.optimizers.Optimizer | str | None = 'rmsprop',
@@ -416,7 +439,6 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         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]:
@@ -531,58 +553,6 @@ def saliency(
             }
         }
     ) 
-
-def predict(
-    model: GraphModel,
-    x: tensors.GraphTensor | tf.data.Dataset,
-    repeats: int | None = 16, 
-    batch_size: int = 256,
-    verbose: int = 0,
-    **kwargs,
-) -> tuple[tf.Tensor | np.ndarray, tf.Tensor | np.ndarray]:
-    """Predict with model.
-
-    By default performs monte-carlo predictions. Namely, it performs
-    `repeats` number of predictions for each example with `training = True`,
-    and subsequently computes mean and standard deviations of the predictions.
-
-    Args:
-        x: 
-            A `GraphTensor` instance.
-        repeats:
-            Number of predictions per example.
-        batch_size:
-            Number of samples per batch of computation.
-        kwargs:
-            See `Model.predict` in Keras documentation. 
-            May or may not apply here. 
-    """
-    if not repeats:
-        return model.predict(
-            x, batch_size=batch_size, verbose=verbose, **kwargs
-        )
-    if isinstance(x, tensors.GraphTensor):
-        ds = tf.data.Dataset.from_tensor_slices(x)
-        ds = ds.repeat(repeats)
-        ds = ds.batch(batch_size)
-    elif isinstance(x, tf.data.Dataset):
-        ds = x.repeat(repeats)
-    else:
-        raise ValueError(
-            'Input `x` needs to be a `tensors.GraphTensor` instance '
-            'or a `tf.data.Dataset` instance constructed from `tensors.GraphTensor`.'
-        )
-    ds = ds.prefetch(-1)
-    y_pred = keras.ops.concatenate([
-        model(x, training=True) for x in ds])
-    global_batch_size = len(y_pred) // repeats
-    y_pred = np.reshape(y_pred, (repeats, global_batch_size, -1))
-    y_pred_loc = keras.ops.mean(y_pred, axis=0)
-    y_pred_scale = keras.ops.std(y_pred, axis=0)
-    if tf.executing_eagerly():
-        y_pred_loc = y_pred_loc.numpy()
-        y_pred_scale = y_pred_scale.numpy()
-    return (y_pred_loc, y_pred_scale)
     
 def _functional_init_arguments(args, kwargs):
     return (
@@ -597,14 +567,3 @@ 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-0.1.0a10 → molcraft-0.1.0a12}/molcraft.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a10
+Version: 0.1.0a12
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -47,15 +47,20 @@ Dynamic: license-file
 **Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML. 
 
 > [!NOTE]  
-> In progress/Unfinished.
+> In progress.
 
-## Highlights
-- Compatible with **Keras 3**
-- Customizable and serializable **featurizers**
-- Customizable and serializable **layers** and **models**
-- Customizable **GraphTensor**
-- Fast and efficient featurization of molecular graphs
-- Fast and efficient input pipelines using TF **records**
+## Installation
+
+For CPU users:
+
+```bash
+pip install --pre molcraft
+```
+
+For GPU users:
+```bash
+pip install --pre molcraft[gpu]
+```
 
 ## Examples 
 
@@ -70,7 +75,7 @@ import keras
 featurizer = featurizers.MolGraphFeaturizer(
     atom_features=[
         features.AtomType(),
-        features.TotalNumHs(),
+        features.NumHydrogens(),
         features.Degree(),
     ],
     bond_features=[

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/molcraft.egg-info/SOURCES.txt

@@ -20,6 +20,8 @@ molcraft.egg-info/SOURCES.txt
 molcraft.egg-info/dependency_links.txt
 molcraft.egg-info/requires.txt
 molcraft.egg-info/top_level.txt
+molcraft/apps/__init__.py
+molcraft/apps/peptides.py
 tests/test_chem.py
 tests/test_featurizers.py
 tests/test_layers.py

{molcraft-0.1.0a10 → molcraft-0.1.0a12}/tests/test_featurizers.py

@@ -31,7 +31,7 @@ class TestFeaturizer(unittest.TestCase):
         featurizer = featurizers.MolFeaturizer(
             atom_features=[
                 features.AtomType({'C', 'N', 'O', 'H'}),
-                features.TotalNumHs({0, 1, 2, 3, 4})
+                features.NumHydrogens({0, 1, 2, 3, 4})
             ],
             bond_features=[
                 features.BondType({'single', 'double', 'aromatic'}),
@@ -119,7 +119,7 @@ class TestFeaturizer(unittest.TestCase):
         featurizer = featurizers.MolFeaturizer3D(
             atom_features=[
                 features.AtomType({'C', 'N', 'O', 'H'}, encode_oov=True),
-                features.TotalNumHs({0, 1, 2, 3, 4})
+                features.NumHydrogens({0, 1, 2, 3, 4})
             ],
             bond_features=[
                 features.Distance(max_distance=20)