molcraft 0.1.0a18__tar.gz → 0.1.0a20__tar.gz

{molcraft-0.1.0a18 → molcraft-0.1.0a20}/PKG-INFO

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

{molcraft-0.1.0a18 → molcraft-0.1.0a20}/molcraft/__init__.py

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

{molcraft-0.1.0a18 → molcraft-0.1.0a20}/molcraft/applications/proteomics.py

@@ -165,18 +165,8 @@ class ResidueEmbedding(keras.layers.Layer):
         residues = {**default_residues, **residues}
         self._residues = {}
         for residue, smiles in residues.items():
-            permuted_smiles = permute_residue_smiles(smiles)
-            # Returned smiles should begin with the amino group.
-            # It seems that the returned smiles ends with carboxyl group,
-            # though we do another check just in case. 
-            if not has_c_terminal_mod(residue):
-                carboxyl_group = 'C(=O)O'
-                if not permuted_smiles.endswith(carboxyl_group):
-                    raise ValueError(
-                        f'Unsupported permutation of {residue!r} smiles: {permuted_smiles!r}.'
-                    )
-            self._residues[residue] = permuted_smiles
-            self._residues[residue + '*'] = permuted_smiles.rstrip('O')
+            self._residues[residue] = smiles
+            self._residues[residue + '*'] = smiles.rstrip('O')
 
         residue_keys = sorted(self._residues.keys())
         residue_values = range(len(residue_keys))

{molcraft-0.1.0a18 → molcraft-0.1.0a20}/molcraft/featurizers.py

@@ -323,7 +323,9 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
         include_hydrogens: bool = False,
         radius: int | float | None = 6.0,
         random_seed: int | None = None,
+        **kwargs,
     ) -> None:
+        kwargs.pop('bond_features', None)
         super().__init__(
             atom_features=atom_features,
             bond_features=None,

{molcraft-0.1.0a18 → molcraft-0.1.0a20}/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,19 +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._message_norm = keras.layers.Identity()
-            self._update_norm = keras.layers.Identity()
-        elif str(self._normalize).lower().startswith('layer'):
-            self._message_norm = keras.layers.LayerNormalization()
-            self._update_norm = keras.layers.LayerNormalization()
-        else:
-            self._message_norm = keras.layers.BatchNormalization()
-            self._update_norm = keras.layers.BatchNormalization()
-
     def propagate(self, tensor: tensors.GraphTensor) -> tensors.GraphTensor:
         """Forward pass.
 
@@ -533,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({

{molcraft-0.1.0a18 → molcraft-0.1.0a20}/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.0a18 → molcraft-0.1.0a20}/molcraft.egg-info/PKG-INFO

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

{molcraft-0.1.0a18 → molcraft-0.1.0a20}/tests/test_featurizers.py

@@ -1,4 +1,6 @@
 import unittest 
+import tempfile
+import shutil
 
 from molcraft import features
 from molcraft import featurizers
@@ -42,6 +44,12 @@ class TestFeaturizer(unittest.TestCase):
             include_hydrogens=False, 
         ) 
 
+        tmp_dir = tempfile.mkdtemp()
+        tmp_file = tmp_dir + '/featurizer.json'
+        featurizers.save_featurizer(featurizer, tmp_file)
+        _ = featurizers.load_featurizer(tmp_file)
+        shutil.rmtree(tmp_dir)
+        
         node_dim = 9
         edge_dim = 4
 
@@ -127,6 +135,12 @@ class TestFeaturizer(unittest.TestCase):
             radius=5.0, 
         ) 
 
+        tmp_dir = tempfile.mkdtemp()
+        tmp_file = tmp_dir + '/featurizer.json'
+        featurizers.save_featurizer(featurizer, tmp_file)
+        _ = featurizers.load_featurizer(tmp_file)
+        shutil.rmtree(tmp_dir)
+
         node_dim = 10
         edge_dim = 22