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

molcraft/__init__.py

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

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/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/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/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/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(),

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:
@@ -649,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/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/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/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.0a11.dist-info}/METADATA

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

@@ -0,0 +1,21 @@
+molcraft/__init__.py,sha256=Huk8xSj59YLku1q0poDWWsKArf7_HULYSFbA9Jpn8u0,464
+molcraft/callbacks.py,sha256=x5HnkZhqcFRrW6xdApt_jZ4X08A-0fxcnFKfdmRKa0c,3571
+molcraft/chem.py,sha256=JARpv4IgFBtuNia0FLW_VF_DdmaA6e-_eZgH9dFAykA,21796
+molcraft/conformers.py,sha256=K6ZtiSUNDN_fwqGP9JrPcwALLFFvlMlF_XejEJH3Sr4,4205
+molcraft/datasets.py,sha256=rFgXTC1ZheLhfgQgcCspP_wEE54a33PIneH7OplbS-8,4047
+molcraft/descriptors.py,sha256=W8GLuDpc38RtwmreNsPOcn-PRvMjTfVng9ksJwcrVyM,3032
+molcraft/features.py,sha256=FpvT_9zk9EiOhvrk6OA5eEvUAYalquF7V6IvpiEJCns,13559
+molcraft/featurizers.py,sha256=A_0wJfvz9JuPEZINi2iZoFNhhHgid608XJTTuVO_jwo,27063
+molcraft/layers.py,sha256=cUpo9dqqNEnc7rNf-Dze8adFhOkTV5F9IhHOKs13OUI,60134
+molcraft/losses.py,sha256=qnS2yC5g-O3n_zVea9MR6TNiFraW2yqRgePOisoUP4A,1065
+molcraft/models.py,sha256=0x74B4WsaZgmUrHmpX9YNr9QXqd1rR3QF_ygyegHoXU,21770
+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/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+molcraft/apps/peptides.py,sha256=N5wJDGDIDRbmOmxin_dTY-odLqb0avAX9FU22U6x6c0,14576
+molcraft-0.1.0a11.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a11.dist-info/METADATA,sha256=jIcab-EvRqLHqM13ftx_eWNz5WjPZTkdmdNM8VttMYA,3893
+molcraft-0.1.0a11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+molcraft-0.1.0a11.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a11.dist-info/RECORD,,

{molcraft-0.1.0a9.dist-info → molcraft-0.1.0a11.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
 

molcraft-0.1.0a9.dist-info/RECORD

@@ -1,19 +0,0 @@
-molcraft/__init__.py,sha256=8f1z8Lhuhh8TxB-QGHI5w4a3M_ZZNH8EWGD4Y6pB578,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=ybJ1djH747cgsftztWHxAX2iTq6k03MYr17btQ2Gtcs,27063
-molcraft/layers.py,sha256=r6hEAyJxO_Yrw5hD1r2v8yb_UxLRK9S4FMjDCUQedH8,59655
-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.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,,