molcraft 0.1.0a22__py3-none-any.whl → 0.1.0a23__py3-none-any.whl

molcraft/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0a22'
+__version__ = '0.1.0a23'
 
 import os
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -14,4 +14,4 @@ from molcraft import records
 from molcraft import tensors
 from molcraft import callbacks
 from molcraft import datasets
-from molcraft import losses
+from molcraft import losses

molcraft/callbacks.py

@@ -1,5 +1,5 @@
-import keras
 import warnings
+import keras
 import numpy as np
 
 

molcraft/chem.py

@@ -3,6 +3,7 @@ import collections
 import numpy as np
 
 from rdkit import Chem
+from rdkit.Chem import AllChem
 from rdkit.Chem import Lipinski
 from rdkit.Chem import rdDistGeom
 from rdkit.Chem import rdDepictor
@@ -31,10 +32,8 @@ class Mol(Chem.Mol):
     
     @property 
     def encoding(self):
-        if hasattr(self, '_encoding'):
-            return self._encoding 
-        return None 
-    
+        return getattr(self, '_encoding', None)
+        
     @property
     def bonds(self) -> list['Bond']:
         if not hasattr(self, '_bonds'):
@@ -67,7 +66,7 @@ class Mol(Chem.Mol):
             atom = atom.GetIdx()
         return Atom.cast(self.GetAtomWithIdx(int(atom)))
     
-    def get_path_between_atoms(
+    def get_shortest_path_between_atoms(
         self, 
         atom_i: int | Chem.Atom, 
         atom_j: int | Chem.Atom
@@ -107,13 +106,13 @@ class Mol(Chem.Mol):
 
     def get_conformer(self, index: int = 0) -> 'Conformer':
         if self.num_conformers == 0:
-            warnings.warn('Molecule has no conformer.')
+            warnings.warn(f'{self} has no conformer. Returning None.')
             return None
         return Conformer.cast(self.GetConformer(index))
     
     def get_conformers(self) -> list['Conformer']:
         if self.num_conformers == 0:
-            warnings.warn('Molecule has no conformer.')
+            warnings.warn(f'{self} has no conformers. Returning an empty list.')
             return []
         return [Conformer.cast(x) for x in self.GetConformers()]
      
@@ -124,7 +123,8 @@ class Mol(Chem.Mol):
         return None
     
     def __repr__(self) -> str:
-        return f'<{self.__class__.__name__} {self.canonical_smiles} at {hex(id(self))}>'
+        encoding = self.encoding or self.canonical_smiles
+        return f'<{self.__class__.__name__} {encoding} at {hex(id(self))}>'
 
 
 class Conformer(Chem.Conformer):
@@ -251,7 +251,10 @@ def sanitize_mol(
     flag = Chem.SanitizeMol(mol, catchErrors=True)
     if flag != Chem.SanitizeFlags.SANITIZE_NONE:
         if strict:
-            return None
+            raise ValueError(f'Could not sanitize {mol}.')
+        warnings.warn(
+            f'Could not sanitize {mol}. Proceeding with partial sanitization.'
+        )
         # Sanitize mol, excluding the steps causing the error previously
         Chem.SanitizeMol(mol, sanitizeOps=Chem.SanitizeFlags.SANITIZE_ALL^flag)
     if assign_stereo_chemistry:
@@ -411,13 +414,12 @@ def embed_conformers(
         'KDG': rdDistGeom.KDG()
     }
     mol = Mol(mol)
-    encoding = mol.encoding or mol.canonical_smiles
     embedding_method = available_embedding_methods.get(method)
     if embedding_method is None:
-        raise ValueError(
-            f'Could not find `method` {method!r}. Specify either of: '
-            '`ETDG`, `ETKDG`, `ETKDGv2`, `ETKDGv3`, `srETKDGv3` or `KDG`.'
+        warnings.warn(
+            f'{method} is not available. Proceeding with ETKDGv3.'
         )
+        embedding_method = available_embedding_methods['ETKDGv3']
 
     for key, value in kwargs.items():
         setattr(embedding_method, key, value)
@@ -438,8 +440,8 @@ def embed_conformers(
     if num_successes < num_conformers:
         warnings.warn(
             f'Could only embed {num_successes} out of {num_conformers} conformer(s) for '
-            f'{encoding!r} using the specified method ({method!r}) and parameters. Attempting '
-            f'to embed the remaining {num_conformers-num_successes} using fallback methods.',
+            f'{mol} using the specified method ({method}) and parameters. Attempting to '
+            f'embed the remaining {num_conformers-num_successes} using fallback methods.',
         )
         max_iters = 20 * mol.num_atoms # Doubling the number of iterations
         for fallback_method in [method, 'ETDG', 'KDG']:
@@ -457,10 +459,13 @@ def embed_conformers(
                 break
         else:
             raise RuntimeError(
-                f'Could not embed {num_conformers} conformer(s) for {encoding!r}. '
+                f'Could not embed {num_conformers} conformer(s) for {mol}. '
             )
     return mol
 
+
+import warnings 
+
 def optimize_conformers(
     mol: Mol,
     method: str = 'UFF',
@@ -469,14 +474,17 @@ def optimize_conformers(
     ignore_interfragment_interactions: bool = True,
     vdw_threshold: float = 10.0,
 ) -> Mol:
-    available_force_field_methods = [
-        'MMFF', 'MMFF94', 'MMFF94s', 'UFF'
-    ]
+    if mol.num_conformers == 0:
+        warnings.warn(
+            f'{mol} has no conformers to optimize. Proceeding without it.'
+        )
+        return Mol(mol)
+    available_force_field_methods = ['MMFF', 'MMFF94', 'MMFF94s', 'UFF']
     if method not in available_force_field_methods:
-        raise ValueError(
-            f'Could not find `method` {method!r}. Specify either of: '
-            '`UFF`, `MMFF`, `MMFF94` or `MMFF94s`.'
+        warnings.warn(
+            f'{method} is not available. Proceeding with universal force field (UFF).'
         )
+        method = 'UFF'
     mol_optimized = Mol(mol)
     try:
         if method.startswith('MMFF'):
@@ -500,7 +508,7 @@ def optimize_conformers(
             )
     except RuntimeError as e:
         warnings.warn(
-            f'{method} force field minimization did not succeed. Proceeding without it.',
+            f'Unsuccessful {method} force field minimization for {mol}. Proceeding without it.',
         )
         return Mol(mol)
     return mol_optimized
@@ -513,10 +521,9 @@ def prune_conformers(
 ) -> Mol:
     if mol.num_conformers == 0:
         warnings.warn(
-            'Molecule has no conformers. To embed conformers, invoke the `embed` method, '
-            'and optionally followed by `minimize()` to perform force field minimization.',
+            f'{mol} has no conformers to prune. Proceeding without it.'
         )
-        return mol
+        return Chem.Mol(mol)
     
     threshold = threshold or 0.0
     deviations = conformer_deviations(mol)

molcraft/datasets.py

@@ -1,3 +1,4 @@
+import warnings
 import numpy as np
 import pandas as pd
 import typing

molcraft/descriptors.py

@@ -1,5 +1,7 @@
+import warnings
 import keras
 import numpy as np
+
 from rdkit.Chem import rdMolDescriptors
 
 from molcraft import chem
@@ -12,9 +14,7 @@ class Descriptor(features.Feature):
     def __call__(self, mol: chem.Mol) -> np.ndarray:
         if not isinstance(mol, chem.Mol):
             raise ValueError(
-                f'Input to {self.name} needs to be a `chem.Mol`, which '
-                'implements two properties that should be iterated over '
-                'to compute features: `atoms` and `bonds`.'
+                f'Input to {self.name} must be a `chem.Mol` object.'
             )
         descriptor = self.call(mol)
         func = (
@@ -30,6 +30,23 @@ class Descriptor(features.Feature):
         return np.concatenate(descriptors)
     
 
+@keras.saving.register_keras_serializable(package='molcraft')
+class Descriptor3D(Descriptor):
+    
+    def __call__(self, mol: chem.Mol) -> np.ndarray:
+        if not isinstance(mol, chem.Mol):
+            raise ValueError(
+                f'Input to {self.name} must be a `chem.Mol` object.'
+            ) 
+        if mol.num_conformers == 0:
+            raise ValueError(
+                f'The inputted `chem.Mol` to {self.name} must embed a conformer. '
+                f'It is recommended that {self.name} is used as a molecule feature '
+                'for `MolGraphFeaturizer3D`, which by default embeds a conformer.'
+            )
+        return super().__call__(mol)
+
+
 @keras.saving.register_keras_serializable(package='molcraft')
 class MolWeight(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
@@ -77,7 +94,7 @@ class NumHydrogenDonors(Descriptor):
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumHydrogenAcceptors(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return rdMolDescriptors.CalcNumHBA(mol) 
+        return rdMolDescriptors.CalcNumHBA(mol)
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
@@ -89,7 +106,7 @@ class NumRotatableBonds(Descriptor):
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumRings(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return rdMolDescriptors.CalcNumRings(mol) 
+        return rdMolDescriptors.CalcNumRings(mol)
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
@@ -105,3 +122,18 @@ class AtomCount(Descriptor):
             if atom.GetSymbol() == self.atom_type:
                 count += 1
         return count
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config['atom_type'] = self.atom_type
+        return config
+    
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class ForceFieldEnergy(Descriptor3D):
+    """Universal Force Field (UFF) Energy."""
+    def call(self, mol: chem.Mol) -> np.ndarray:
+        mol_copy = chem.Mol(mol)
+        mol_copy = chem.add_hs(mol_copy)
+        return chem.conformer_energies(mol_copy, method="UFF")
+        

molcraft/diffusion.py

@@ -0,0 +1,241 @@
+import warnings
+import keras 
+import tensorflow as tf
+import numpy as np
+
+from molcraft import ops
+from molcraft import tensors 
+from molcraft import layers 
+from molcraft import models 
+
+
+# smiles = pd.read_csv('../../data/rt/RIKEN.csv')['smiles'].values
+
+
+# graph = featurizers.MolGraphFeaturizer3D(super_node=False)(smiles)
+# graph.node['coordinate']
+
+
+# encoder = molcraft.models.GraphModel.from_layers(
+#     [
+#         diffusion.CoordinateNoise(),
+#         molcraft.layers.NodeEmbedding(128),
+#         molcraft.layers.EdgeEmbedding(128),
+#         molcraft.layers.AddContext('position'),
+#         molcraft.layers.MPConv(128),
+#         molcraft.layers.AddContext('position'),
+#         molcraft.layers.MPConv(128),
+#         molcraft.layers.AddContext('position'),
+#     ]
+# )
+
+# decoder = keras.Sequential([
+#     keras.layers.Dense(128, activation='relu'),
+#     keras.layers.Dense(3),
+# ])
+
+# model = diffusion.CoordinateNoisePredictor(encoder, decoder)
+
+# model(graph)
+
+# model.save('/tmp/model.keras')
+# model = molcraft.models.load_model('/tmp/model.keras')
+
+# model.compile(keras.optimizers.Adam(1e-3), 'mse')
+# model.fit(graph, epochs=100)
+
+
+# from rdkit.Geometry import Point3D
+
+# def energy(smiles, coordinate):
+#     m = chem.Mol.from_encoding(smiles)
+#     m = chem.embed_conformers(m, 1)
+#     conf = m.GetConformer()
+#     for i in range(m.GetNumAtoms()):
+#         x, y, z = coordinate[i]
+#         conf.SetAtomPosition(i, Point3D(float(x), float(y), float(z)))
+#     return m, chem.conformer_energies(m)[0]
+
+# def denoise(
+#     graph: tensors.GraphTensor,
+#     model,
+# ):
+
+#     print("----")
+#     print(energy(smiles[0], graph[0].node['coordinate'])[-1])
+#     print("----")
+    
+#     beta = keras.ops.linspace(1e-4, 1e-2, 100)
+#     alpha = 1 - beta
+#     alpha_bar = keras.ops.cumprod(alpha)
+#     sigma = keras.ops.sqrt(beta[1:] * (1.0 - alpha_bar[:-1]) / (1.0 - alpha_bar[1:]))
+    
+#     graph = graph.update(
+#         {
+#             'context': {
+#                 'position': keras.ops.ones_like(graph.context['size']) * 99
+#             },
+#             'node': {
+#                 'coordinate': keras.random.normal(graph.node['coordinate'].shape)
+#             }
+#         }
+#     )
+
+#     for t in reversed(range(100)):
+#         alpha_t = alpha[t]
+#         alpha_bar_t = alpha_bar[t]
+
+#         a = 1 / keras.ops.sqrt(alpha_t)
+
+#         b = (1 - alpha_t) / keras.ops.sqrt(1 - alpha_bar_t)
+
+#         if t > 0:
+#             z = keras.random.normal(()) * sigma[t-1]
+#         else:
+#             z = 0.0
+
+#         graph = graph.update({
+#             'node': {
+#                 'coordinate': (
+#                     a * (graph.node['coordinate'] - b * model(graph)) + z
+#                 ) 
+#             } 
+#         }) 
+        
+        
+#         print(energy(smiles[0], graph[0].node['coordinate'])[-1])
+
+#     return graph
+        
+# graph_updated = denoise(graph[:1], model)x
+# mol, e = energy(smiles[0], graph_updated[0].node['coordinate'])
+# print(e)
+# Chem.Mol(mol)
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class CoordinateNoisePredictor(models.GraphModel):
+
+    def __init__(self, encoder, decoder, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.encoder = encoder 
+        self.decoder = decoder 
+
+    def propagate(self, tensor):
+        return self.decoder(self.encoder(tensor).node['feature'])
+    
+    def train_step(self, tensor: tensors.GraphTensor) -> dict[str, float]:
+        with tf.GradientTape() as tape:
+            tensor = self.encoder(tensor)
+            feature = tensor.node['feature']
+            noise_true = tensor.node['label']
+            noise_pred = self.decoder(feature)
+            loss = self.compute_loss(tensor, noise_true, noise_pred)
+            loss = self.optimizer.scale_loss(loss)
+        trainable_weights = self.trainable_weights 
+        gradients = tape.gradient(loss, trainable_weights)
+        self.optimizer.apply_gradients(zip(gradients, trainable_weights))
+        return self.compute_metrics(tensor, noise_true, noise_pred)
+    
+    def test_step(self, tensor: tensors.GraphTensor) -> dict[str, float]:
+        tensor = self.encoder(tensor)
+        feature = tensor.node['feature']
+        noise_true = tensor.node['label']
+        noise_pred = self.decoder(feature)
+        return self.compute_metrics(tensor, noise_true, noise_pred)
+    
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config['encoder'] = keras.saving.serialize_keras_object(self.encoder)
+        config['decoder'] = keras.saving.serialize_keras_object(self.decoder)
+        return config 
+    
+    @classmethod
+    def from_config(cls, config: dict):
+        config['encoder'] = keras.saving.deserialize_keras_object(config['encoder'])
+        config['decoder'] = keras.saving.deserialize_keras_object(config['decoder'])
+        return super().from_config(config)
+
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class CoordinateNoise(layers.GraphLayer):
+
+    def __init__(
+        self, 
+        beta: tuple[float, float] = (1e-4, 1e-2),
+        position_dim: int = 128,
+        max_timesteps: int = 100,
+        **kwargs
+    ) -> None:
+        super().__init__(**kwargs)
+        self._beta = beta 
+        self._max_timesteps = max_timesteps
+        beta = keras.ops.linspace(*self._beta, self._max_timesteps)
+        alpha = 1 - beta
+        alpha_cumprod = keras.ops.cumprod(alpha)
+        alpha_cumprod = keras.ops.expand_dims(alpha_cumprod, -1)
+        self._alpha_cumprod = alpha_cumprod
+        self._timestep_embedding = TimestepEmbedding(dim=position_dim)
+
+    def propagate(self, graph: tensors.GraphTensor) -> tensors.GraphTensor:
+        if 'position' in graph.context:
+            return graph.update({'context': {'position': self._timestep_embedding(graph.context['position'])}})
+        
+        timestep = keras.random.randint(
+            shape=(graph.num_subgraphs,), minval=0, maxval=self._max_timesteps
+        )
+        alpha_cumprod = ops.gather(
+            ops.gather(self._alpha_cumprod, timestep), graph.graph_indicator
+        )
+        epsilon = keras.random.normal(
+            shape=keras.ops.shape(graph.node['coordinate']), mean=0, stddev=1
+        )
+        noisy_coordinate = (
+            keras.ops.sqrt(alpha_cumprod) * graph.node['coordinate'] + 
+            keras.ops.sqrt(1 - alpha_cumprod) * epsilon
+        )
+        timestep = self._timestep_embedding(timestep)
+        return graph.update(
+            {
+                'context': {
+                    'position': timestep,
+                },
+                'node': {
+                    'coordinate': noisy_coordinate, 
+                    'label': epsilon
+                },
+            }
+        )
+
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config['beta'] = self._beta 
+        config['max_timesteps'] = self._max_timesteps
+        return config 
+    
+
+class TimestepEmbedding(keras.layers.Layer):
+    
+    def __init__(self, dim: int, max_wavelength: int = 10000, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._dim = dim
+        self._max_wavelength = max_wavelength
+
+    def call(self, inputs: tf.Tensor) -> tf.Tensor:
+        timestep = keras.ops.cast(inputs, 'float32')
+        embedding = keras.ops.log(self._max_wavelength) / (self._dim // 2 - 1)
+        embedding = keras.ops.exp(
+            -embedding * keras.ops.arange(self._dim // 2, dtype='float32')
+        )
+        embedding = timestep[:, None] * embedding[None, :]
+        embedding = keras.ops.concatenate(
+            [keras.ops.sin(embedding), keras.ops.cos(embedding)], axis=-1
+        )
+        return embedding
+
+
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config['dim'] = self._dim
+        config['max_wavelength'] = self._max_wavelength
+        return config 
+    

molcraft/features.py

@@ -1,7 +1,7 @@
+import warnings
 import abc
 import math
 import keras
-import warnings
 import numpy as np
 
 from molcraft import chem
@@ -41,14 +41,14 @@ class Feature(abc.ABC):
     
     def __call__(self, mol: chem.Mol) -> np.ndarray:
         if not isinstance(mol, chem.Mol):
-            raise TypeError(f'Input to {self.name} must be a `chem.Mol` instance.')
+            raise TypeError(f'Input to {self.name} must be a `chem.Mol` object.')
         features = self.call(mol)
         if len(features) != mol.num_atoms and len(features) != mol.num_bonds:
             raise ValueError(
                 f'The number of features computed by {self.name} does not '
                 'match the number of atoms or bonds of the `chem.Mol` object. '
-                'Make sure to iterate over `atoms` or `bonds` of `chem.Mol` '
-                'when computing features.'
+                'Make sure to iterate over `atoms` or `bonds` of the `chem.Mol` '
+                'object when computing features.'
             )
         if len(features) == 0:
             # Edge case: no atoms or bonds in the molecule.
@@ -109,7 +109,6 @@ class Feature(abc.ABC):
             warnings.warn(
                 f'Found value of {self.name} to be non-finite. '
                 f'Value received: {value}. Converting it to a value of 0.',
-                stacklevel=2
             )
             value = 0.0
         return np.asarray([value], dtype=self.dtype)

molcraft/featurizers.py

@@ -1,8 +1,8 @@
+import warnings
 import keras 
 import json
 import abc
 import typing 
-import os
 import numpy as np
 import pandas as pd
 import tensorflow as tf

molcraft/layers.py

@@ -1,6 +1,6 @@
+import warnings
 import keras 
 import tensorflow as tf
-import warnings
 import functools
 from keras.src.models import functional
 
@@ -350,11 +350,8 @@ class GraphConv(GraphLayer):
         )
         if self._project_residual:
             warnings.warn(
-                '`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. ',
-                stacklevel=2,
+                'Found incompatible dim between input and output. Applying '
+                'a projection layer to residual to match input and output dim.',
             )
             self._residual_dense = self.get_dense(
                 self.units, name='residual_dense'
@@ -613,10 +610,8 @@ class GIConv(GraphConv):
             if not self._update_edge_feature:
                 if (edge_feature_dim != node_feature_dim):
                     warnings.warn(
-                        'Found edge feature dim to be incompatible with node feature dim. '
-                        'Automatically adding a edge feature projection layer to match '
-                        'the dim of node features.',
-                        stacklevel=2,
+                        'Found edge and node feature dim to be incompatible. Applying a '
+                        'projection layer to edge features to match the dim of the node features.',
                     )
                     self._update_edge_feature = True 
 
@@ -870,10 +865,10 @@ class MPConv(GraphConv):
         self._project_previous_node_feature = node_feature_dim != self.units
         if self._project_previous_node_feature:
             warnings.warn(
-                'Input node feature dim does not match updated node feature dim. '
-                'To make sure input node feature can be passed as `states` to the '
-                'GRU cell, it will automatically be projected prior to it.',
-                stacklevel=2
+                'Inputted node feature dim does not match updated node feature dim, '
+                'which is required for the GRU update. Applying a projection layer to '
+                'the inputted node features prior to the GRU update, to match dim '
+                'of the updated node feature dim.'
             )
             self._previous_node_dense = self.get_dense(self.units)
 
@@ -1497,6 +1492,7 @@ class AddContext(GraphLayer):
         
     def build(self, spec: tensors.GraphTensor.Spec) -> None:
         feature_dim = spec.node['feature'].shape[-1]
+        self._has_super_node = 'super' in spec.node
         if self._intermediate_dim is None:
             self._intermediate_dim = feature_dim * 2
         self._intermediate_dense = self.get_dense(
@@ -1515,9 +1511,14 @@ class AddContext(GraphLayer):
         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
-        )
+        if self._has_super_node:
+            node_feature = ops.scatter_add(
+                tensor.node['feature'], tensor.node['super'], context
+            )
+        else:
+            node_feature = (
+                tensor.node['feature'] + ops.gather(context, tensor.graph_indicator)
+            )
         data = {'node': {'feature': node_feature}}
         if self._drop:
             data['context'] = {self._field: None}
@@ -1561,8 +1562,7 @@ class GraphNetwork(GraphLayer):
         if self._update_node_feature:
             warnings.warn(
                 'Node feature dim does not match `units` of the first layer. '
-                'Automatically adding a node projection layer to match `units`.',
-                stacklevel=2
+                'Applying a projection layer to node features to match `units`.',
             )
             self._node_dense = self.get_dense(units)
         self._has_edge_feature = 'feature' in spec.edge 
@@ -1572,8 +1572,7 @@ class GraphNetwork(GraphLayer):
             if self._update_edge_feature:
                 warnings.warn(
                     'Edge feature dim does not match `units` of the first layer. '
-                    'Automatically adding a edge projection layer to match `units`.',
-                    stacklevel=2
+                    'Applying projection layer to edge features to match `units`.'
                 )
                 self._edge_dense = self.get_dense(units)
 

molcraft/losses.py

@@ -1,3 +1,4 @@
+import warnings
 import keras 
 import numpy as np
 

molcraft/models.py

@@ -1,3 +1,4 @@
+import warnings
 import typing
 import keras
 import numpy as np
@@ -111,7 +112,7 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
     def __new__(cls, *args, **kwargs):
         if _functional_init_arguments(args, kwargs) and cls == GraphModel:
             return FunctionalGraphModel(*args, **kwargs)
-        return typing.cast(GraphModel, super().__new__(cls))
+        return super().__new__(cls)
     
     def __init__(self, *args, **kwargs):
         self._model_layers = kwargs.pop('model_layers', None)
@@ -137,6 +138,8 @@ class GraphModel(layers.GraphLayer, keras.models.Model):
         """
         if not tensors.is_graph(graph_layers[0]):
             return cls(model_layers=graph_layers)
+        elif cls != GraphModel:
+            return cls(model_layers=graph_layers[1:])
         inputs: dict = graph_layers.pop(0)
         x = inputs
         for layer in graph_layers:

molcraft/ops.py

@@ -1,3 +1,4 @@
+import warnings
 import keras
 import numpy as np
 import tensorflow as tf

molcraft/records.py

@@ -1,9 +1,9 @@
+import warnings
 import os
 import math
 import glob
 import time
 import typing 
-import warnings
 import tensorflow as tf
 import numpy as np
 import pandas as pd
@@ -164,8 +164,8 @@ def _write_tfrecord(
                 _write_example(tensor)
             except Exception as e:
                 warnings.warn(
-                    f"Could not write record for index {i + start_index}, proceeding without it."
-                    f"Exception raised:\n{e}"
+                    f'Could not write record for index {i + start_index}, proceeding without it.'
+                    f'Exception raised:\n{e}'
                 )
 
 def _serialize_example(

molcraft/tensors.py

@@ -1,3 +1,4 @@
+import warnings
 import tensorflow as tf 
 import keras
 import typing 

{molcraft-0.1.0a22.dist-info → molcraft-0.1.0a23.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a22
+Version: 0.1.0a23
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -35,7 +35,6 @@ 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
@@ -43,7 +42,7 @@ Provides-Extra: gpu
 Requires-Dist: tensorflow[and-cuda]>=2.16; extra == "gpu"
 Dynamic: license-file
 
-<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo", width="90%">
+<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo" width="90%">
 
 **Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML.
 

molcraft-0.1.0a23.dist-info/RECORD

@@ -0,0 +1,22 @@
+molcraft/__init__.py,sha256=QAJQSS_jOBzLxGRL7ciskMY_kn2ARLCg7FVTWeF-D_I,432
+molcraft/callbacks.py,sha256=B4gGWjVW_1ORrt38jfk1ZFI9c0rOpN5sgjGWVqs3Ess,3571
+molcraft/chem.py,sha256=0Zni91J4fQJW16R6g3jlOX9Vm8FH0Z5NOx0s7_X-xQw,22232
+molcraft/datasets.py,sha256=1rHccqra5chIBwo2pz9vduyv0i07uY3CABzmAqWiFBU,4161
+molcraft/descriptors.py,sha256=uqMPeIKqfkHC04FgztxS1FsfC3zsFJhvniZO70D22l0,4553
+molcraft/diffusion.py,sha256=HR1kp2MuCWyUtGoGXvEA6kXTdWMGD2w5EZEoKLI1ilM,7902
+molcraft/features.py,sha256=q-wuRP9YjPu_v5czipsh00VEXEjgFaeuLk6dbgyD_VM,13505
+molcraft/featurizers.py,sha256=nGdV9G-aO43-vgKPNFfEOESW2hVvIvixHu3EHjIRrgU,18097
+molcraft/layers.py,sha256=ba0WdQC2IUNsLy9pV0mIm5BBO3MCvR2lYWWuq1-8M4M,64522
+molcraft/losses.py,sha256=piu4XYAgjnK7k9LqA4Vkh-SooYZ31sWwRfG1cacCwyA,1081
+molcraft/models.py,sha256=-at-yFWj8mIkGchVY39m9-HtTnKqAUDQrF6wDrQXNuQ,22040
+molcraft/ops.py,sha256=Qf9l1oOg20HNi9L9nLgf_c_5v09GtXDCc-1fkEQqn54,6194
+molcraft/records.py,sha256=dAaq5tr3B1jXNgd3tkvVxWgUW2Qa9gtXWE3TToeiKmQ,6283
+molcraft/tensors.py,sha256=JILwU9l6kUsrtJJ9YSzmT90_G5kbZZC6LAShvsnZOrk,22493
+molcraft/applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+molcraft/applications/chromatography.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+molcraft/applications/proteomics.py,sha256=BL3EtW-q-0j79pLYO7npC67mA2ApRhH-XI4rOaP8_wc,8407
+molcraft-0.1.0a23.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
+molcraft-0.1.0a23.dist-info/METADATA,sha256=RomzSM8GmqbILgyZMmzR4GjrBLWDPHGZcedb1VDmqxs,3892
+molcraft-0.1.0a23.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+molcraft-0.1.0a23.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
+molcraft-0.1.0a23.dist-info/RECORD,,

molcraft-0.1.0a22.dist-info/RECORD

@@ -1,21 +0,0 @@
-molcraft/__init__.py,sha256=O88EmicQAD8oz9oFMXk_IzFChQEbbU-BCs3IE-c9Dkk,431
-molcraft/callbacks.py,sha256=x5HnkZhqcFRrW6xdApt_jZ4X08A-0fxcnFKfdmRKa0c,3571
-molcraft/chem.py,sha256=ynrEpWZL2D370p7CqH2kE1KhBByq7IiuQbUNoKQt96I,22028
-molcraft/datasets.py,sha256=Nd2lw5USUZE52vvAiNr-q-n03Y3--NlZlK0NzqHgp-E,4145
-molcraft/descriptors.py,sha256=Cl3KnBPsTST7XLgRLktkX5LwY9MV0P_lUlrt8iPV5no,3508
-molcraft/features.py,sha256=s0WeV8eZcDEypPgC1m37f4s9QkvWIlVgn-L43Cdsa14,13525
-molcraft/featurizers.py,sha256=1yBz5-JA7IhNm0dGivvVm1nJ5QGck8VQXtwHPWFbTuQ,18091
-molcraft/layers.py,sha256=H7XZru4XGJA6gbRO9V1BsGqh1mIrMdhzNCKS5o6oNok,64544
-molcraft/losses.py,sha256=qnS2yC5g-O3n_zVea9MR6TNiFraW2yqRgePOisoUP4A,1065
-molcraft/models.py,sha256=2Pc1htT9fCukGd8ZxrvE0rzEHsPBm0pluHw4FZXaUE4,21963
-molcraft/ops.py,sha256=bQbdFDt9waxVCzF5-dkTB6vlpj9eoSt8I4Qg7ZGXbsU,6178
-molcraft/records.py,sha256=sopYElKWC3A9QE5I8_957v3faLb2Wt5WILHZv_FLLds,6283
-molcraft/tensors.py,sha256=vk-W8zZu-re1g18YevDEEoVQRxT4AdIiMdI-4EvtJI4,22477
-molcraft/applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-molcraft/applications/chromatography.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-molcraft/applications/proteomics.py,sha256=BL3EtW-q-0j79pLYO7npC67mA2ApRhH-XI4rOaP8_wc,8407
-molcraft-0.1.0a22.dist-info/licenses/LICENSE,sha256=sbVeqlrtZ0V63uYhZGL5dCxUm8rBAOqe2avyA1zIQNk,1074
-molcraft-0.1.0a22.dist-info/METADATA,sha256=1OHx3-Q94fFEi21l0p3bnMjU-Q0EHaZLm4PU1A6QbkU,3930
-molcraft-0.1.0a22.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-molcraft-0.1.0a22.dist-info/top_level.txt,sha256=dENV6MfOceshM6MQCgJlcN1ojZkiCL9B4F7XyUge3QM,9
-molcraft-0.1.0a22.dist-info/RECORD,,