molcraft 0.1.0a14__tar.gz → 0.1.0a16__tar.gz

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a14
+Version: 0.1.0a16
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -35,6 +35,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: tensorflow>=2.16
+Requires-Dist: tensorflow-text>=2.16
 Requires-Dist: rdkit>=2023.9.5
 Requires-Dist: pandas>=1.0.3
 Requires-Dist: ipython>=8.12.0

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft/__init__.py

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

molcraft-0.1.0a16/molcraft/applications/proteomics.py

@@ -0,0 +1,239 @@
+import re
+import keras
+import numpy as np
+import tensorflow as tf
+import tensorflow_text as tf_text
+import json
+
+from molcraft import featurizers
+from molcraft import tensors
+from molcraft import layers
+from molcraft import models 
+from molcraft import chem
+
+
+# TODO: Add regex pattern for residue (C-term mod + N-term mod)?
+# TODO: Add regex pattern for residue (C-term mod + N-term mod + mod)?
+residue_pattern: str = "|".join([
+    r'(\[[A-Za-z0-9]+\]-[A-Z]\[[A-Za-z0-9]+\])', # residue (N-term mod + mod)
+    r'([A-Z]\[[A-Za-z0-9]+\]-\[[A-Za-z0-9]+\])', # residue (C-term mod + mod)
+    r'([A-Z]-\[[A-Za-z0-9]+\])', # residue (C-term mod)
+    r'(\[[A-Za-z0-9]+\]-[A-Z])', # residue (N-term mod)
+    r'([A-Z]\[[A-Za-z0-9]+\])', # residue (mod)
+    r'([A-Z])', # residue (no mod)
+])
+
+default_residues: dict[str, str] = {
+    "A": "N[C@@H](C)C(=O)O",
+    "C": "N[C@@H](CS)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",
+    "L": "N[C@@H](CC(C)C)C(=O)O",
+    "M": "N[C@@H](CCSC)C(=O)O",
+    "N": "N[C@@H](CC(=O)N)C(=O)O",
+    "P": "N1[C@@H](CCC1)C(=O)O",
+    "Q": "N[C@@H](CCC(=O)N)C(=O)O",
+    "R": "N[C@@H](CCCNC(=N)N)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",
+}
+
+
+class Peptide(chem.Mol):
+
+    @classmethod
+    def from_sequence(cls, sequence: str, **kwargs) -> 'Peptide':
+        sequence = [
+            match.group(0) for match in re.finditer(residue_pattern, sequence)
+        ]
+        peptide_smiles = []
+        for i, residue in enumerate(sequence):
+            if i < len(sequence) - 1:
+                residue_smiles = registered_residues[residue + '*']
+            else:
+                residue_smiles = registered_residues[residue]
+            peptide_smiles.append(residue_smiles)
+        peptide_smiles = ''.join(peptide_smiles)
+        return super().from_encoding(peptide_smiles, **kwargs)
+
+
+@keras.saving.register_keras_serializable(package='proteomics')
+class ResidueEmbedding(keras.layers.Layer):
+
+    def __init__(
+        self, 
+        featurizer: featurizers.MolGraphFeaturizer,
+        embedder: models.GraphModel, 
+        **kwargs
+    ) -> None:
+        residues = kwargs.pop('_residues', None)
+        super().__init__(**kwargs)
+        if residues is None:
+            residues = registered_residues.copy()
+        self._residues = residues
+        self.embedder = embedder
+        self.featurizer = featurizer
+        self.ragged_split = SequenceSplitter(pad=False)
+        self.split = SequenceSplitter(pad=True)
+        self.supports_masking = True
+        
+    def build(self, input_shape) -> None:
+        embedding_dim = self.embedder.output.shape[-1]
+        residues = sorted(self._residues.keys())
+        smiles = [self._residues[residue] for residue in residues]
+        num_residues = len(residues)
+        self.oov_index = np.where(np.array(residues) == "G")[0][0]
+        self.mapping = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                keys=residues, 
+                values=range(num_residues)
+            ),
+            default_value=-1,
+        )
+        self.graph = tf.stack([self.featurizer(s) for s in smiles], axis=0)
+        self.cached_embeddings = tf.Variable(
+            initial_value=tf.zeros((num_residues, embedding_dim))
+        )
+        self.use_cached_embeddings = tf.Variable(False)
+        super().build(input_shape)
+
+    def call(self, sequences, training=None) -> tensors.GraphTensor:
+        if training is False:
+            self.use_cached_embeddings.assign(True)
+        else:
+            self.use_cached_embeddings.assign(False)
+        embeddings = tf.cond(
+            pred=self.use_cached_embeddings,
+            true_fn=lambda: self.cached_embeddings,
+            false_fn=lambda: self.embeddings(),
+        )
+        sequences = self.ragged_split(sequences)
+        sequences = keras.ops.concatenate([
+            tf.strings.join([sequences[:, :-1], '*']), sequences[:, -1:]
+        ], axis=1)
+        indices = self.mapping.lookup(sequences)
+        indices = keras.ops.where(indices == -1, self.oov_index, indices)
+        return tf.gather(embeddings, indices).to_tensor()
+    
+    def embeddings(self) -> tf.Tensor:
+        embeddings = self.embedder(self.graph)
+        self.cached_embeddings.assign(embeddings)
+        return embeddings
+    
+    def compute_mask(
+        self, 
+        inputs: tensors.GraphTensor, 
+        mask: bool | None = None
+    ) -> tf.Tensor | None:
+        sequences = self.split(inputs)
+        return keras.ops.not_equal(sequences, '')
+                
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            '_residues': self._residues,
+            'featurizer': keras.saving.serialize_keras_object(self.featurizer),
+            'embedder': keras.saving.serialize_keras_object(self.embedder)
+        })
+        return config
+    
+    @classmethod
+    def from_config(cls, config: dict) -> 'ResidueEmbedding':
+        config['featurizer'] = keras.saving.deserialize_keras_object(config['featurizer'])
+        config['embedder'] = keras.saving.deserialize_keras_object(config['embedder'])
+        return super().from_config(config)
+    
+
+@keras.saving.register_keras_serializable(package='proteomics')
+class SequenceSplitter(keras.layers.Layer):
+
+    def __init__(self, pad: bool, **kwargs):
+        super().__init__(**kwargs)
+        self.pad = pad 
+
+    def call(self, inputs):
+        inputs = tf_text.regex_split(inputs, residue_pattern, residue_pattern)
+        if self.pad:
+            inputs = inputs.to_tensor()
+        return inputs
+    
+
+def interpret(model: keras.models.Model, sequence: list[str]) -> tensors.GraphTensor:
+    
+    if not tf.is_tensor(sequence):
+        sequence = keras.ops.convert_to_tensor(sequence)
+
+    # Find embedding layer
+    for layer in model.layers:
+        if isinstance(layer, ResidueEmbedding):
+            break
+    
+    # Use embedding layer to convert the sequence to a graph
+    residues = layer.ragged_split(sequence)
+    residues = keras.ops.concatenate([
+        tf.strings.join([residues[:, :-1], '*']), residues[:, -1:]
+    ], axis=1)
+    indices = layer.mapping.lookup(residues)
+    graph = tf.concat([
+        layer.graph[residue_ids] for residue_ids in indices
+    ], axis=0)
+
+    # Define layer which reshapes data into sequences of residue embeddings
+    num_residues = indices.row_lengths()
+    to_sequence = (
+        lambda x: tf.RaggedTensor.from_row_lengths(x, num_residues).to_tensor()
+    )
+    reshape = keras.layers.Lambda(to_sequence)
+    
+    # Obtain the embedder part of the original model
+    embedder = layer.embedder
+    # Obtain the remaining part of the original model
+    predictor = keras.models.Model(embedder.output, model.output)
+    # Obtain an 'interpretable model', based on the original model
+    inputs = layers.Input(graph.spec)
+    x = inputs 
+    for layer in embedder.layers: # Loop over layers to expose them
+        x = layer(x)
+    x = reshape(x)
+    outputs = predictor(x)
+    interpretable_model = models.GraphModel(inputs, outputs)
+    
+    # Interpret original model through the 'interpretable model'
+    graph = models.interpret(interpretable_model, graph)
+    del interpretable_model
+
+    # Update 'size' field with new sizes corresponding to peptides for convenience
+    # Allows the user to obtain n:th peptide graph using indexing: nth_peptide = graph[n]
+    peptide_indices = range(len(num_residues))
+    peptide_indicator = keras.ops.repeat(peptide_indices, num_residues)
+    residue_sizes = graph.context['size']
+    peptide_sizes = keras.ops.segment_sum(residue_sizes, peptide_indicator)
+    return graph.update({'context': {'size': peptide_sizes, 'sequence': sequence}})
+
+
+def register_residues(residues: dict[str, str]) -> None:
+    # TODO: Implement functions that check if residue has N- or C-terminal mod
+    #       if C-terminal mod, no need to enforce concatenatable perm.
+    #       if N-terminal mod, enforce only 'C(=O)O'
+    #       if normal mod, enforce concatenateable perm ('N[C@@H]' and 'C(=O)O)).
+    for residue, smiles in residues.items():
+        if residue.startswith('P'):
+            smiles.startswith('N'), f'Incorrect SMILES permutation for {residue}.'
+        elif not residue.startswith('['):
+            smiles.startswith('N[C@@H]'), f'Incorrect SMILES permutation for {residue}.'
+        if len(residue) > 1 and not residue[1] == "-":
+            assert smiles.endswith('C(=O)O'), f'Incorrect SMILES permutation for {residue}.'
+        registered_residues[residue] = smiles
+        registered_residues[residue + '*'] = smiles.strip('O')
+    
+
+registered_residues: dict[str, str] = {}
+register_residues(default_residues)

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft/featurizers.py

@@ -169,7 +169,7 @@ class MolGraphFeaturizer(Featurizer):
         if default_atom_features:
             atom_features = [features.AtomType()]
             if not self.include_hs:
-                atom_features.append(features.TotalNumHs())
+                atom_features.append(features.NumHydrogens())
             atom_features.append(features.Degree())
         if not isinstance(self, MolGraphFeaturizer3D):
             default_bond_features = (

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft/layers.py

@@ -1430,6 +1430,56 @@ class EdgeEmbedding(GraphLayer):
         return config
     
 
+@keras.saving.register_keras_serializable(package='molcraft')
+class AddContext(GraphLayer):
+
+    """Context adding layer.
+
+    Adds context to super nodes.
+    """
+
+    def __init__(
+        self, 
+        field: str = 'feature',
+        drop: bool = True,
+        normalize: bool = False,
+        **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+        self.field = field
+        self.drop = drop
+        self._normalize = normalize
+        
+    def build(self, spec: tensors.GraphTensor.Spec) -> None:
+        feature_dim = spec.node['feature'].shape[-1]
+        self._context_dense = self.get_dense(feature_dim)
+        if not self._normalize:
+            self._norm = keras.layers.Identity()
+        elif str(self._normalize).lower().startswith('layer'):
+            self._norm = keras.layers.LayerNormalization()
+        else:
+            self._norm = keras.layers.BatchNormalization()
+
+    def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
+        context = tensor.context[self.field]
+        context = self._context_dense(context)
+        context = self._norm(context)
+        node_feature = ops.scatter_add(
+            tensor.node['feature'], tensor.node['super'], context
+        )
+        data = {'node': {'feature': node_feature}}
+        if self.drop:
+            data['context'] = {self.field: None}
+        return tensor.update(data)
+
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config['field'] = self.field 
+        config['drop'] = self.drop
+        config['normalize'] = self._normalize
+        return config
+
+
 @keras.saving.register_keras_serializable(package='molcraft')
 class GraphNetwork(GraphLayer):
 

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft/models.py

@@ -397,7 +397,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
             raise ValueError(
                 'Could not extract output. `Readout` layer not found.'
             )
-        return self.__class__(inputs, outputs, name=f'{self.name}_head')
+        return self.__class__(inputs, outputs, name=f'{self.name}_backbone')
 
     def head(self) -> functional.Functional:
         if not isinstance(self, FunctionalGraphModel):

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft/ops.py

@@ -4,6 +4,7 @@ import tensorflow as tf
 from keras import backend
 
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def gather(
     node_feature: tf.Tensor, 
     edge: tf.Tensor
@@ -16,6 +17,7 @@ def gather(
         edge = keras.ops.expand_dims(edge, axis=-1)
     return keras.ops.take_along_axis(node_feature, edge, axis=0)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def aggregate(
     node_feature: tf.Tensor, 
     edge: tf.Tensor, 
@@ -30,6 +32,7 @@ def aggregate(
         node_feature, edge, num_nodes, sorted=False
     )
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def propagate(
     node_feature: tf.Tensor,
     edge_source: tf.Tensor,
@@ -49,6 +52,7 @@ def propagate(
         
     return aggregate(node_feature, edge_target, num_nodes)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def scatter_update(
     inputs: tf.Tensor,
     indices: tf.Tensor,
@@ -62,6 +66,7 @@ def scatter_update(
         indices = keras.ops.expand_dims(indices, axis=-1)
     return keras.ops.scatter_update(inputs, indices, updates)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def scatter_add(
     inputs: tf.Tensor,
     indices: tf.Tensor,
@@ -78,6 +83,7 @@ def scatter_add(
     updates = scatter_update(keras.ops.zeros_like(inputs), indices, updates)
     return inputs + updates
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def edge_softmax(
     score: tf.Tensor, 
     edge_target: tf.Tensor
@@ -98,6 +104,7 @@ def edge_softmax(
     denominator = gather(denominator, edge_target)
     return numerator / denominator
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def edge_weight(
     edge: tf.Tensor,
     edge_weight: tf.Tensor,
@@ -108,6 +115,7 @@ def edge_weight(
         edge_weight = keras.ops.expand_dims(edge_weight, axis=-1)
     return edge * edge_weight
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def segment_mean(
     data: tf.Tensor,
     segment_ids: tf.Tensor,
@@ -142,6 +150,7 @@ def segment_mean(
     )
     return x / sizes[:, None]
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def gaussian(
     x: tf.Tensor, 
     mean: tf.Tensor, 
@@ -155,6 +164,7 @@ def gaussian(
     a = (2 * np.pi) ** 0.5
     return keras.ops.exp(-0.5 * (((x - mean) / std) ** 2)) / (a * std)
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def euclidean_distance(
     x1: tf.Tensor, 
     x2: tf.Tensor, 
@@ -169,6 +179,7 @@ def euclidean_distance(
         )
     )
 
+@keras.saving.register_keras_serializable(package='molcraft')
 def displacement(
     x1: tf.Tensor,
     x2: tf.Tensor,

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a14
+Version: 0.1.0a16
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -35,6 +35,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: tensorflow>=2.16
+Requires-Dist: tensorflow-text>=2.16
 Requires-Dist: rdkit>=2023.9.5
 Requires-Dist: pandas>=1.0.3
 Requires-Dist: ipython>=8.12.0

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft.egg-info/SOURCES.txt

@@ -20,8 +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
+molcraft/applications/__init__.py
+molcraft/applications/proteomics.py
 tests/test_chem.py
 tests/test_featurizers.py
 tests/test_layers.py

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/molcraft.egg-info/requires.txt

@@ -1,4 +1,5 @@
 tensorflow>=2.16
+tensorflow-text>=2.16
 rdkit>=2023.9.5
 pandas>=1.0.3
 ipython>=8.12.0

{molcraft-0.1.0a14 → molcraft-0.1.0a16}/pyproject.toml

@@ -26,6 +26,7 @@ classifiers = [
 requires-python = ">=3.10"
 dependencies = [
     "tensorflow>=2.16",
+    "tensorflow-text>=2.16",
     "rdkit>=2023.9.5",
     "pandas>=1.0.3",
     "ipython>=8.12.0"

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

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