molcraft 0.1.0a17__tar.gz → 0.1.0a19__tar.gz

{molcraft-0.1.0a17 → molcraft-0.1.0a19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a17
+Version: 0.1.0a19
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License

{molcraft-0.1.0a17 → molcraft-0.1.0a19}/molcraft/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0a17'
+__version__ = '0.1.0a19'
 
 import os
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -14,4 +14,6 @@ from molcraft import records
 from molcraft import tensors
 from molcraft import callbacks
 from molcraft import datasets
-from molcraft import losses
+from molcraft import losses
+
+from molcraft.applications import proteomics

{molcraft-0.1.0a17 → molcraft-0.1.0a19}/molcraft/applications/proteomics.py

@@ -3,6 +3,7 @@ import keras
 import numpy as np
 import tensorflow as tf
 import tensorflow_text as tf_text
+from rdkit import Chem
 
 from molcraft import featurizers
 from molcraft import tensors
@@ -10,16 +11,47 @@ from molcraft import layers
 from molcraft import models 
 from molcraft import chem
 
+"""
+
+
+
+
+
+
+No need to correct smiles for modeling, only for interpretation.
+
+Use added smiles data to rearrange list of saliency values. 
+
+
+
+
+
+
+
+
+
+
+
+
+"""
 
 # 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)?
+
+no_mod_pattern = r'([A-Z])'
+side_chain_mod_pattern = r'([A-Z]\[[A-Za-z0-9]+\])'
+n_term_mod_pattern = r'(\[[A-Za-z0-9]+\]-[A-Z])'
+c_term_mod_pattern = r'([A-Z]-\[[A-Za-z0-9]+\])'
+side_chain_and_n_term_mod_pattern = r'(\[[A-Za-z0-9]+\]-[A-Z]\[[A-Za-z0-9]+\])'
+side_chain_and_c_term_mod_pattern = r'([A-Z]\[[A-Za-z0-9]+\]-\[[A-Za-z0-9]+\])'
+
 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)
+    side_chain_and_n_term_mod_pattern,
+    side_chain_and_c_term_mod_pattern,
+    n_term_mod_pattern,
+    c_term_mod_pattern,
+    side_chain_mod_pattern,
+    no_mod_pattern
 ])
 
 default_residues: dict[str, str] = {
@@ -45,20 +77,26 @@ default_residues: dict[str, str] = {
     "Y": "N[C@@H](Cc1ccc(O)cc1)C(=O)O",
 }
 
-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')
+def has_c_terminal_mod(residue: str):
+    if re.search(c_term_mod_pattern, residue):
+        return True 
+    return False 
+
+def has_n_terminal_mod(residue: str):
+    if re.search(n_term_mod_pattern, residue):
+        return True 
+    return False 
+
+# def register_residues(residues: dict[str, str]) -> None:
+#     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')
     
     
 class Peptide(chem.Mol):
@@ -79,6 +117,22 @@ class Peptide(chem.Mol):
         return super().from_encoding(peptide_smiles, **kwargs)
 
 
+def permute_residue_smiles(smiles: str) -> str:
+    glycine = chem.Mol.from_encoding("NCC(=O)O")
+    mol = chem.Mol.from_encoding(smiles)
+    nitrogen_index = mol.GetSubstructMatch(glycine)[0]
+    permuted_smiles = Chem.MolToSmiles(
+        mol, rootedAtAtom=nitrogen_index
+    )
+    return permuted_smiles
+
+def check_peptide_residue_smiles(smiles: list[str]) -> bool:
+    backbone = 'NCC(=O)' * (len(smiles) - 1) + 'NC'
+    backbone = chem.Mol.from_encoding(backbone)
+    mol = chem.Mol.from_encoding(''.join(smiles))
+    is_valid = mol.HasSubstructMatch(backbone)
+    return is_valid
+
 @keras.saving.register_keras_serializable(package='proteomics')
 class ResidueEmbedding(keras.layers.Layer):
 
@@ -92,42 +146,50 @@ class ResidueEmbedding(keras.layers.Layer):
         super().__init__(**kwargs)
         if residues is None:
             residues = {}
-        self._residue_dict = {**default_residues, **residues}
         self.embedder = embedder
         self.featurizer = featurizer
-        self.embedding_dim = self.embedder.output.shape[-1]
+        self.embedding_dim = int(self.embedder.output.shape[-1])
         self.ragged_split = SequenceSplitter(pad=False)
         self.split = SequenceSplitter(pad=True)
         self.use_cached_embeddings = tf.Variable(False)
+        self.residues = residues
         self.supports_masking = True
         
     @property
     def residues(self) -> dict[str, str]:
-        return self._residue_dict
+        return self._residues
 
     @residues.setter
     def residues(self, residues: dict[str, str]) -> None:
-        self._residue_dict = residues 
-        num_residues = len(residues)
-        residue_keys = sorted(residues.keys())
-        oov_value = np.where(np.array(residue_keys) == "G")[0][0]
+        
+        residues = {**default_residues, **residues}
+        self._residues = {}
+        for residue, smiles in residues.items():
+            self._residues[residue] = smiles
+            self._residues[residue + '*'] = smiles.rstrip('O')
+
+        residue_keys = sorted(self._residues.keys())
+        residue_values = range(len(residue_keys))
+        residue_oov_value = np.where(np.array(residue_keys) == "G")[0][0]
+
         self.mapping = tf.lookup.StaticHashTable(
             tf.lookup.KeyValueTensorInitializer(
                 keys=residue_keys, 
-                values=range(num_residues)
+                values=residue_values
             ),
-            default_value=oov_value,
+            default_value=residue_oov_value,
         )
+
         self.graph = tf.stack([
-            self.featurizer(residues[residue]) for residue in residue_keys
+            self.featurizer(self._residues[r]) for r in residue_keys
         ], axis=0)
-        self.cached_embeddings = tf.Variable(
-            initial_value=tf.zeros((num_residues, self.embedding_dim))
-        )
+
+        zeros = tf.zeros((residue_values[-1] + 1, self.embedding_dim))
+        self.cached_embeddings = tf.Variable(initial_value=zeros)
         _ = self.cache_and_get_embeddings()
 
     def build(self, input_shape) -> None:
-        self.residues = self._residue_dict
+        self.residues = self._residues
         super().build(input_shape)
 
     def call(self, sequences: tf.Tensor, training: bool = None) -> tf.Tensor:
@@ -163,16 +225,24 @@ class ResidueEmbedding(keras.layers.Layer):
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
-            'featurizer': keras.saving.serialize_keras_object(self.featurizer),
-            'embedder': keras.saving.serialize_keras_object(self.embedder),
-            'residues': self._residue_dict,
+            'featurizer': keras.saving.serialize_keras_object(
+                self.featurizer
+            ),
+            'embedder': keras.saving.serialize_keras_object(
+                self.embedder
+            ),
+            'residues': self._residues,
         })
         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'])
+        config['featurizer'] = keras.saving.deserialize_keras_object(
+            config['featurizer']
+        )
+        config['embedder'] = keras.saving.deserialize_keras_object(
+            config['embedder']
+        )
         return super().from_config(config)
     
 
@@ -190,5 +260,5 @@ class SequenceSplitter(keras.layers.Layer):
         return inputs
     
 
-registered_residues: dict[str, str] = {}
-register_residues(default_residues)
+# registered_residues: dict[str, str] = {}
+# register_residues(default_residues)

{molcraft-0.1.0a17 → molcraft-0.1.0a19}/molcraft/featurizers.py

@@ -413,6 +413,7 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
 
     def get_config(self):
         config = super().get_config()
+        config.pop('bond_features', None)
         config['radius'] = self._radius
         config['pair_features'] = keras.saving.serialize_keras_object(
             self._pair_features

{molcraft-0.1.0a17 → molcraft-0.1.0a19}/molcraft/layers.py

@@ -279,7 +279,7 @@ class GraphConv(GraphLayer):
         use_bias (bool):
             Whether bias should be used in the dense layers. Default to `True`.
         normalize (bool, str):
-            Whether normalization should be applied to the final output. Default to `False`.
+            Whether a normalization layer should be obtain by `get_norm()`. Default to `False`.
         skip_connect (bool):
             Whether node feature input should be added to the node feature output. Default to `True`.
         kernel_initializer (keras.initializers.Initializer, str):
@@ -366,6 +366,7 @@ class GraphConv(GraphLayer):
         has_overridden_message = self.__class__.message != GraphConv.message 
         if not has_overridden_message:
             self._message_intermediate_dense = self.get_dense(self.units)
+            self._message_norm = self.get_norm()
             self._message_intermediate_activation = self.activation
             self._message_final_dense = self.get_dense(self.units)
 
@@ -376,16 +377,10 @@ class GraphConv(GraphLayer):
         has_overridden_update = self.__class__.update != GraphConv.update 
         if not has_overridden_update:
             self._update_intermediate_dense = self.get_dense(self.units)
+            self._update_norm = self.get_norm()
             self._update_intermediate_activation = self.activation
             self._update_final_dense = self.get_dense(self.units)
 
-        if not self._normalize:
-            self._normalization = keras.layers.Identity()
-        elif str(self._normalize).lower().startswith('layer'):
-            self._normalization = keras.layers.LayerNormalization()
-        else:
-            self._normalization = keras.layers.BatchNormalization()
-
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Forward pass.
 
@@ -430,7 +425,7 @@ class GraphConv(GraphLayer):
         elif add_aggregate:
             update = update.update({'node': {'aggregate': None}})
 
-        if not self._skip_connect and not self._normalize:
+        if not self._skip_connect:
             return update
         
         feature = update.node['feature']
@@ -438,8 +433,6 @@ class GraphConv(GraphLayer):
         if self._skip_connect:
             feature += residual 
 
-        feature = self._normalization(feature)
-
         return update.update({'node': {'feature': feature}})
 
     def message(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
@@ -480,6 +473,7 @@ class GraphConv(GraphLayer):
                 axis=-1
             )
         message = self._message_intermediate_dense(message)
+        message = self._message_norm(message)
         message = self._message_intermediate_activation(message)
         message = self._message_final_dense(message)
         return tensor.update({'edge': {'message': message}})
@@ -519,6 +513,7 @@ class GraphConv(GraphLayer):
         """
         aggregate = tensor.node['aggregate']
         node_feature = self._update_intermediate_dense(aggregate)
+        node_feature = self._update_norm(node_feature)
         node_feature = self._update_intermediate_activation(node_feature)
         node_feature = self._update_final_dense(node_feature)
         return tensor.update(
@@ -530,6 +525,14 @@ class GraphConv(GraphLayer):
             }
         )
 
+    def get_norm(self, **kwargs):
+        if not self._normalize:
+            return keras.layers.Identity()
+        elif str(self._normalize).lower().startswith('layer'):
+            return keras.layers.LayerNormalization(**kwargs)
+        else:
+            return keras.layers.BatchNormalization(**kwargs)
+        
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
@@ -1312,13 +1315,19 @@ class NodeEmbedding(GraphLayer):
 
     def __init__(
         self, 
-        dim: int = None, 
+        dim: int | None = None, 
+        intermediate_dim: int | None = None,
+        intermediate_activation: str | keras.layers.Activation | None = 'relu',
         normalize: bool = False,
         embed_context: bool = False,
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
+        self._intermediate_dim = intermediate_dim
+        self._intermediate_activation = keras.activations.get(
+            intermediate_activation
+        )
         self._normalize = normalize
         self._embed_context = embed_context
 
@@ -1326,30 +1335,38 @@ class NodeEmbedding(GraphLayer):
         feature_dim = spec.node['feature'].shape[-1]
         if not self.dim:
             self.dim = feature_dim
-        self._node_dense = self.get_dense(self.dim)
-
+        if not self._intermediate_dim:
+            self._intermediate_dim = self.dim * 2 
+        self._node_dense = self.get_dense(
+            self._intermediate_dim, activation=self._intermediate_activation
+        )
         self._has_super = 'super' in spec.node
         has_context_feature = 'feature' in spec.context
         if not has_context_feature:
             self._embed_context = False 
         if self._has_super and not self._embed_context:
-            self._super_feature = self.get_weight(shape=[self.dim], name='super_node_feature')
+            self._super_feature = self.get_weight(
+                shape=[self._intermediate_dim], name='super_node_feature'
+            )
         if self._embed_context:
-            self._context_dense = self.get_dense(self.dim)
-        
+            self._context_dense = self.get_dense(
+                self._intermediate_dim, activation=self._intermediate_activation
+            )
         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()
+        self._dense = self.get_dense(self.dim)
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         feature = self._node_dense(tensor.node['feature'])
 
         if self._has_super and not self._embed_context:
             super_mask = keras.ops.expand_dims(tensor.node['super'], 1)
-            feature = keras.ops.where(super_mask, self._super_feature, feature)
+            super_feature = self._intermediate_activation(self._super_feature)
+            feature = keras.ops.where(super_mask, super_feature, feature)
 
         if self._embed_context:
             context_feature = self._context_dense(tensor.context['feature'])
@@ -1357,6 +1374,7 @@ class NodeEmbedding(GraphLayer):
             tensor = tensor.update({'context': {'feature': None}})
 
         feature = self._norm(feature)
+        feature = self._dense(feature)
 
         return tensor.update({'node': {'feature': feature}})
 
@@ -1364,6 +1382,10 @@ class NodeEmbedding(GraphLayer):
         config = super().get_config()
         config.update({
             'dim': self.dim,
+            'intermediate_dim': self._intermediate_dim,
+            'intermediate_activation': keras.activations.serialize(
+                self._intermediate_activation
+            ),
             'normalize': self._normalize,
             'embed_context': self._embed_context,
         })
@@ -1381,50 +1403,67 @@ class EdgeEmbedding(GraphLayer):
     def __init__(
         self, 
         dim: int = None, 
+        intermediate_dim: int | None = None,
+        intermediate_activation: str | keras.layers.Activation | None = 'relu',
         normalize: bool = False,
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
         self.dim = dim
+        self._intermediate_dim = intermediate_dim
+        self._intermediate_activation = keras.activations.get(
+            intermediate_activation
+        )
         self._normalize = normalize
 
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
         feature_dim = spec.edge['feature'].shape[-1]
         if not self.dim:
             self.dim = feature_dim
-        self._edge_dense = self.get_dense(self.dim) 
-
-        self._self_loop_feature = self.get_weight(shape=[self.dim], name='self_loop_edge_feature')
-
+        if not self._intermediate_dim:
+            self._intermediate_dim = self.dim * 2 
+        self._edge_dense = self.get_dense(
+            self._intermediate_dim, activation=self._intermediate_activation
+        ) 
+        self._self_loop_feature = self.get_weight(
+            shape=[self._intermediate_dim], name='self_loop_edge_feature'
+        )
         self._has_super = 'super' in spec.edge
         if self._has_super:
-            self._super_feature = self.get_weight(shape=[self.dim], name='super_edge_feature')
-        
+            self._super_feature = self.get_weight(
+                shape=[self._intermediate_dim], name='super_edge_feature'
+            )
         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()
+        self._dense = self.get_dense(self.dim)
 
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         feature = self._edge_dense(tensor.edge['feature'])
 
         if self._has_super:
             super_mask = keras.ops.expand_dims(tensor.edge['super'], 1)
-            feature = keras.ops.where(super_mask, self._super_feature, feature)
+            super_feature = self._intermediate_activation(self._super_feature)
+            feature = keras.ops.where(super_mask, super_feature, feature)
 
         self_loop_mask = keras.ops.expand_dims(tensor.edge['source'] == tensor.edge['target'], 1)
-        feature = keras.ops.where(self_loop_mask, self._self_loop_feature, feature)
-
+        self_loop_feature = self._intermediate_activation(self._self_loop_feature)
+        feature = keras.ops.where(self_loop_mask, self_loop_feature, feature)
         feature = self._norm(feature)
-
+        feature = self._dense(feature)
         return tensor.update({'edge': {'feature': feature}})
 
     def get_config(self) -> dict:
         config = super().get_config()
         config.update({
             'dim': self.dim,
+            'intermediate_dim': self._intermediate_dim,
+            'intermediate_activation': keras.activations.serialize(
+                self._intermediate_activation
+            ),
             'normalize': self._normalize,
         })
         return config
@@ -1441,42 +1480,60 @@ class AddContext(GraphLayer):
     def __init__(
         self, 
         field: str = 'feature',
+        intermediate_dim: int | None = None,
+        intermediate_activation: str | keras.layers.Activation | None = 'relu',
         drop: bool = False,
         normalize: bool = False,
         **kwargs
     ) -> None:
         super().__init__(**kwargs)
-        self.field = field
-        self.drop = drop
+        self._field = field
+        self._drop = drop
+        self._intermediate_dim = intermediate_dim
+        self._intermediate_activation = keras.activations.get(
+            intermediate_activation
+        )
         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 self._intermediate_dim is None:
+            self._intermediate_dim = feature_dim * 2
+        self._intermediate_dense = self.get_dense(
+            self._intermediate_dim, activation=self._intermediate_activation
+        )
+        self._final_dense = self.get_dense(feature_dim)
         if not self._normalize:
-            self._norm = keras.layers.Identity()
+            self._intermediate_norm = keras.layers.Identity()
         elif str(self._normalize).lower().startswith('layer'):
-            self._norm = keras.layers.LayerNormalization()
+            self._intermediate_norm = keras.layers.LayerNormalization()
         else:
-            self._norm = keras.layers.BatchNormalization()
+            self._intermediate_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)
+        context = tensor.context[self._field]
+        context = self._intermediate_dense(context)
+        context = self._intermediate_norm(context)
+        context = self._final_dense(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}
+        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
+        config.update({
+            'field': self._field,
+            'intermediate_dim': self._intermediate_dim,
+            'intermediate_activation': keras.activations.serialize(
+                self._intermediate_activation
+            ),
+            'drop': self._drop,
+            'normalize': self._normalize,
+        })
         return config
 
 
@@ -1738,5 +1795,3 @@ def _spec_from_inputs(inputs):
         return spec
     return tensors.GraphTensor.Spec(**nested_specs)
 
-
-GraphTransformer = GTConv

{molcraft-0.1.0a17 → molcraft-0.1.0a19}/molcraft/tensors.py

@@ -16,13 +16,15 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
             if isinstance(f, tf.TensorSpec):
                 return tf.TensorSpec(
                     shape=[None] + f.shape[1:],
-                    dtype=f.dtype)
+                    dtype=f.dtype
+                )
             elif isinstance(f, tf.RaggedTensorSpec):
                 return tf.RaggedTensorSpec(
                     shape=[batch_size, None] + f.shape[1:],
                     dtype=f.dtype,
                     ragged_rank=1,
-                    row_splits_dtype=f.row_splits_dtype)
+                    row_splits_dtype=f.row_splits_dtype
+                )
             elif isinstance(f, tf.TypeSpec):
                 return f.__batch_encoder__.batch(f, batch_size)
             return f
@@ -33,7 +35,8 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         batched_spec = object.__new__(type(spec))
         batched_context_fields = tf.nest.map_structure(
             lambda spec: tf.TensorSpec([batch_size] + spec.shape, spec.dtype), 
-            context_fields)
+            context_fields
+        )
         batched_spec.__dict__.update({'context': batched_context_fields})
         batched_spec.__dict__.update(batched_fields)
         return batched_spec
@@ -46,13 +49,15 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
             if isinstance(f, tf.TensorSpec):
                 return tf.TensorSpec(
                     shape=[None] + f.shape[1:],
-                    dtype=f.dtype)
+                    dtype=f.dtype
+                )
             elif isinstance(f, tf.RaggedTensorSpec):
                 return tf.RaggedTensorSpec(
                     shape=[None] + f.shape[2:],
                     dtype=f.dtype,
                     ragged_rank=0,
-                    row_splits_dtype=f.row_splits_dtype)
+                    row_splits_dtype=f.row_splits_dtype
+                )
             elif isinstance(f, tf.TypeSpec):
                 return f.__batch_encoder__.unbatch(f)
             return f
@@ -62,7 +67,8 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         unbatched_fields = tf.nest.map_structure(unbatch_field, fields)
         unbatched_context_fields = tf.nest.map_structure(
             lambda spec: tf.TensorSpec(spec.shape[1:], spec.dtype), 
-            context_fields)
+            context_fields
+        )
         unbatched_spec = object.__new__(type(spec))
         unbatched_spec.__dict__.update({'context': unbatched_context_fields})
         unbatched_spec.__dict__.update(unbatched_fields)
@@ -91,7 +97,8 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
                     shape=([None] if scalar else [None, None]) + f.shape[1:], 
                     dtype=f.dtype, 
                     ragged_rank=(0 if scalar else 1),
-                    row_splits_dtype=spec.context['size'].dtype)
+                    row_splits_dtype=spec.context['size'].dtype
+                )
             return f
         fields = dict(spec.__dict__)
         context_fields = fields.pop('context')
@@ -99,7 +106,7 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         encoded_fields = {**{'context': context_fields}, **encoded_fields}
         spec_components = tuple(encoded_fields.values())
         spec_components = tuple(
-            x for x in tf.nest.flatten(spec_components) 
+            x for x in tf.nest.flatten(spec_components)
             if isinstance(x, tf.TypeSpec)
         )
         return spec_components
@@ -117,7 +124,6 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
         fields = dict(zip(spec.__dict__.keys(), value_tuple))
         value = object.__new__(spec.value_type)
         value.__dict__.update(fields)
-
         flatten = is_ragged(value) and not is_ragged(spec)
         if flatten:
             value = value.flatten()
@@ -125,7 +131,7 @@ class GraphTensorBatchEncoder(tf.experimental.ExtensionTypeBatchEncoder):
     
 
 class GraphTensor(tf.experimental.BatchableExtensionType):
-    context: typing.Mapping[str, typing.Union[tf.Tensor, tf.RaggedTensor]]
+    context: typing.Mapping[str, tf.Tensor]
     node: typing.Mapping[str, typing.Union[tf.Tensor, tf.RaggedTensor]]
     edge: typing.Mapping[str, typing.Union[tf.Tensor, tf.RaggedTensor]]
 

{molcraft-0.1.0a17 → molcraft-0.1.0a19}/molcraft.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a17
+Version: 0.1.0a19
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License