molcraft 0.1.0a2__tar.gz → 0.1.0a4__tar.gz

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a2
+Version: 0.1.0a4
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -51,11 +51,11 @@ Dynamic: license-file
 
 ## Highlights
 - Compatible with **Keras 3**
-- Simplified API
-- Fast featurization
-- Modular graph **layers**
-- Serializable graph **featurizers** and **models**
-- Flexible **GraphTensor**
+- Customizable and serializable **featurizers**
+- Customizable and serializable **layers** and **models**
+- Customizable **GraphTensor**
+- Fast and efficient featurization of molecular graphs
+- Efficient and easy-to-use input pipelines using TF **records**
 
 ## Examples 
 

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/README.md

@@ -7,11 +7,11 @@
 
 ## Highlights
 - Compatible with **Keras 3**
-- Simplified API
-- Fast featurization
-- Modular graph **layers**
-- Serializable graph **featurizers** and **models**
-- Flexible **GraphTensor**
+- Customizable and serializable **featurizers**
+- Customizable and serializable **layers** and **models**
+- Customizable **GraphTensor**
+- Fast and efficient featurization of molecular graphs
+- Efficient and easy-to-use input pipelines using TF **records**
 
 ## Examples 
 

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/molcraft/__init__.py

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

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/molcraft/callbacks.py

@@ -19,3 +19,15 @@ class TensorBoard(keras.callbacks.TensorBoard):
                             weight, image_weight_name, epoch
                         )
             self._train_writer.flush()
+
+
+class LearningRateDecay(keras.callbacks.LearningRateScheduler):
+
+    def __init__(self, rate: float, delay: int = 0, **kwargs):
+
+        def lr_schedule(epoch: int, lr: float):
+            if epoch < delay:
+                return float(lr)
+            return float(lr * keras.ops.exp(-rate))
+        
+        super().__init__(schedule=lr_schedule, **kwargs)

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/molcraft/descriptors.py

@@ -1,6 +1,6 @@
 import keras
 import numpy as np
-from rdkit.Chem import Descriptors
+from rdkit.Chem import rdMolDescriptors
 
 from molcraft import chem
 from molcraft import features
@@ -8,9 +8,6 @@ from molcraft import features
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class Descriptor(features.Feature):
-    def __init__(self, scale: float | None = None, **kwargs):
-        super().__init__(**kwargs)
-        self.scale = scale 
 
     def __call__(self, mol: chem.Mol) -> np.ndarray:
         if not isinstance(mol, chem.Mol):
@@ -24,67 +21,71 @@ class Descriptor(features.Feature):
             self._featurize_categorical if self.vocab else 
             self._featurize_floating
         )
-        scale_value = self.scale and not self.vocab
         if not isinstance(descriptor, (tuple, list, np.ndarray)):
             descriptor = [descriptor]
         
         descriptors = []
         for value in descriptor:
-            if scale_value:
-                value /= self.scale 
             descriptors.append(func(value))
         return np.concatenate(descriptors)
     
-    def get_config(self):
-        config = super().get_config()
-        config['scale'] = self.scale 
-        return config 
-    
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class MolWeight(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.MolWt(mol) 
+        return rdMolDescriptors.CalcExactMolWt(mol) 
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class MolTPSA(Descriptor):
+class TPSA(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.TPSA(mol)
+        return rdMolDescriptors.CalcTPSA(mol)
+
 
+@keras.saving.register_keras_serializable(package='molcraft')
+class CrippenLogP(Descriptor):
+    def call(self, mol: chem.Mol) -> np.ndarray:
+        return rdMolDescriptors.CalcCrippenDescriptors(mol)[0]
+    
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class MolLogP(Descriptor):
+class CrippenMolarRefractivity(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.MolLogP(mol)
+        return rdMolDescriptors.CalcCrippenDescriptors(mol)[1]
     
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumHeavyAtoms(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.HeavyAtomCount(mol)
+        return rdMolDescriptors.CalcNumHeavyAtoms(mol)
 
 
+@keras.saving.register_keras_serializable(package='molcraft')
+class NumHeteroAtoms(Descriptor):
+    def call(self, mol: chem.Mol) -> np.ndarray:
+        return rdMolDescriptors.CalcNumHeteroatoms(mol) 
+    
+    
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumHydrogenDonors(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.NumHDonors(mol) 
+        return rdMolDescriptors.CalcNumHBD(mol) 
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumHydrogenAcceptors(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.NumHAcceptors(mol) 
-
+        return rdMolDescriptors.CalcNumHBA(mol) 
+    
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumRotatableBonds(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.NumRotatableBonds(mol) 
+        return rdMolDescriptors.CalcNumRotatableBonds(mol) 
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
 class NumRings(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
-        return Descriptors.RingCount(mol) 
+        return rdMolDescriptors.CalcNumRings(mol) 
 

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/molcraft/experimental/peptides.py

@@ -2,6 +2,7 @@ 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
@@ -10,12 +11,14 @@ from molcraft import features
 from molcraft import featurizers
 from molcraft import tensors
 from molcraft import descriptors
+from molcraft import layers 
 
     
 def Graph(
     inputs,
     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,
@@ -25,7 +28,7 @@ def Graph(
     featurizer = featurizers.MolGraphFeaturizer(
         atom_features=atom_features,
         bond_features=bond_features,
-        molecule_features=[AminoAcidType()],
+        molecule_features=molecule_features,
         super_atom=super_atom,
         radius=radius,
         self_loops=self_loops,
@@ -33,41 +36,55 @@ def Graph(
         **kwargs,
     )
 
-    inputs = [
-        residues[x] for x in ['G'] + inputs
+    tensor_list: list[tensors.GraphTensor] = [
+        featurizer(residues[tag]).update({'context': {'tag': tag}}) for tag in inputs
     ]
-    tensor_list = [featurizer(x) for x in inputs]
     return tf.stack(tensor_list, axis=0)
+        
 
-
-def GraphLookup(graph: tensors.GraphTensor) -> 'GraphLookupLayer':
-    lookup = GraphLookupLayer()
+def Lookup(graph: tensors.GraphTensor) -> 'LookupLayer':
+    lookup = LookupLayer()
     lookup._build(graph)
     return lookup
 
 
 @keras.saving.register_keras_serializable(package='molcraft')
-class GraphLookupLayer(keras.layers.Layer):
+class SequenceSplit(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()
+
+
+
 
-    def call(self, indices: tf.Tensor) -> tensors.GraphTensor:
-        indices = tf.sort(tf.unique(tf.reshape(indices, [-1]))[0])
+@keras.saving.register_keras_serializable(package='molcraft')
+class LookupLayer(keras.layers.Layer):
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._sequence_splitter = SequenceSplit()
+
+    def call(self, sequence: tf.Tensor) -> tensors.GraphTensor:
+        sequence = self._sequence_splitter(sequence)
+        indices = self._tag_to_index.lookup(sequence)
+        indices = tf.sort(tf.unique(tf.reshape(indices, [-1]))[0])[1:]
         graph = self.graph[indices]
-        sizes = graph.context['size']
-        max_index = keras.ops.max(indices)
-        sizes = tf.tensor_scatter_nd_update(
-            tensor=tf.zeros([max_index + 1], dtype=indices.dtype),
-            indices=indices[:, None],
-            updates=sizes
-        )
-        graph = graph.update(
-            {
-                'context': {
-                    'size': sizes
-                }
-            },
-        )
         return tensors.to_dict(graph)
-    
+
     def _build(self, x):
 
         if isinstance(x, tensors.GraphTensor):
@@ -98,7 +115,17 @@ class GraphLookupLayer(keras.layers.Layer):
             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):
         config = super().get_config()
         spec = keras.saving.serialize_keras_object(self._spec)
@@ -106,7 +133,7 @@ class GraphLookupLayer(keras.layers.Layer):
         return config 
     
     @classmethod
-    def from_config(cls, config: dict) -> 'GraphLookupLayer':
+    def from_config(cls, config: dict) -> 'LookupLayer':
         spec = config.pop('spec')
         spec = keras.saving.deserialize_keras_object(spec)
         layer = cls(**config)
@@ -130,66 +157,48 @@ class Gather(keras.layers.Layer):
         super().__init__(**kwargs)
         self.padding = padding
         self.mask_value = mask_value
+        self._readout_layer = layers.Readout(mode='mean')
         self.supports_masking = True
+        self._sequence_splitter = SequenceSplit()
 
     def get_config(self):
         config = super().get_config()
         config['mask_value'] = self.mask_value 
         config['padding'] = self.padding
         return config 
-    
-    def call(self, inputs: tuple[tf.Tensor, tf.Tensor]) -> tf.Tensor:
-        data, indices = inputs 
-        # if self.padding:
-        #     padding = self.padding
-        #     if isinstance(self.padding, int):
-        #         padding = [(self.padding, 0)]
-        #     if isinstance(self.padding, tuple):
-        #         padding = [self.padding] 
-        #     data_rank = len(keras.ops.shape(data))
-        #     for _ in range(data_rank - len(padding)):
-        #         padding.append((0, 0))
-        #     data = keras.ops.pad(data, padding)
-        return ops.gather(data, indices) 
-    
+
+    def call(self, inputs) -> tf.Tensor:
+
+        graph, sequence = inputs
+
+        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 = self._sequence_splitter(sequence)
+        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: tuple[tf.Tensor, tf.Tensor], 
+        inputs: tensors.GraphTensor, 
         mask: bool | None = None
     ) -> tf.Tensor | None:
         # if self.mask_value is None:
         #     return None
-        _, indices = inputs
-        return keras.ops.not_equal(indices, self.mask_value)
-    
+        _, sequence = inputs
+        sequence = self._sequence_splitter(sequence)
+        return keras.ops.not_equal(sequence, '')
 
-@keras.saving.register_keras_serializable(package='molcraft')
-class AminoAcidType(descriptors.Descriptor):
-
-    def __init__(self, vocab=None, **kwargs):
-        vocab = [
-            "A", "C", "D", "E", "F", "G", "H", "I", "K", "L", 
-            "M", "N", "P", "Q", "R", "S", "T", "V", "W", "Y",
-        ]
-        super().__init__(vocab=vocab, **kwargs)
-
-    def call(self, mol: chem.Mol) -> list[str]:
-        residue = residues_reverse.get(mol.canonical_smiles)
-        if not residue:
-            raise KeyError(f'Could not find {mol.canonical_smiles} in `residues_reverse`.')
-        mol = chem.remove_hs(mol)
-        return _extract_residue_type(residues_reverse[mol.canonical_smiles])
-    
-def sequence_split(sequence: str):
-    patterns = [
-        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
-    ]
-    return [match.group(0) for match in re.finditer("|".join(patterns), sequence)]
 
 residues = {
     "A": "N[C@@H](C)C(=O)O",
@@ -252,13 +261,21 @@ residues = {
 }
 
 residues_reverse = {}
-def register_peptide_residues(residues: dict[str, str]):
-    for residue, smiles in residues.items():
-        residues[residue] = Chem.MolToSmiles(Chem.MolFromSmiles(smiles))
+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)
+register_peptide_residues(residues, canonicalize=False)
 
 def _extract_residue_type(residue_tag: str) -> str:
-    pattern = r"(?<!\[)[A-Z](?![\w-])"
-    return [match.group(0) for match in re.finditer(pattern, residue_tag)][0]
+    pattern = r"(?<!\[)[A-Z](?![^\[]*\])"
+    return [match.group(0) for match in re.finditer(pattern, residue_tag)][0]
+
+special_residues = {}
+for key, value in residues.items():
+    special_residues[key + '-[X]'] = value.rstrip('O')
+
+register_peptide_residues(special_residues, canonicalize=False)

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/molcraft/features.py

@@ -155,9 +155,11 @@ class Distance(EdgeFeature):
         encode_oov: bool = True, 
         **kwargs,
     ) -> None:
-        if max_distance is None:
-            max_distance = 20
-        vocab = list(range(max_distance + 1))
+        vocab = kwargs.pop('vocab', None)
+        if not vocab:
+            if max_distance is None:
+                max_distance = 20
+            vocab = list(range(max_distance + 1))
         super().__init__(
             vocab=vocab, 
             allow_oov=allow_oov, 

{molcraft-0.1.0a2 → molcraft-0.1.0a4}/molcraft/featurizers.py

@@ -186,9 +186,11 @@ class MolGraphFeaturizer(Featurizer):
         if default_molecule_features:
             molecule_features = [
                 descriptors.MolWeight(),
-                descriptors.MolTPSA(),
-                descriptors.MolLogP(),
+                descriptors.TPSA(),
+                descriptors.CrippenLogP(),
+                descriptors.CrippenMolarRefractivity(),
                 descriptors.NumHeavyAtoms(),
+                descriptors.NumHeteroAtoms(),
                 descriptors.NumHydrogenDonors(),
                 descriptors.NumHydrogenAcceptors(),
                 descriptors.NumRotatableBonds(),
@@ -219,7 +221,7 @@ class MolGraphFeaturizer(Featurizer):
         
         atom_feature = self.atom_features(mol)
         bond_feature = self.bond_features(mol)
-        context_feature = self.context_feature(mol)
+        molecule_feature = self.molecule_feature(mol)
         molecule_size = self.num_atoms(mol)
         
         if isinstance(context, dict):
@@ -241,8 +243,14 @@ class MolGraphFeaturizer(Featurizer):
         else:
             context = {'size': molecule_size}
 
-        if context_feature is not None:
-            context['feature'] = context_feature
+        if molecule_feature is not None:
+            if 'feature' in context:
+                warn(
+                    'Found both inputted and computed context feature. '
+                    'Overwriting inputted context feature with computed '
+                    'context feature (based on `molecule_features`).'
+                )
+            context['feature'] = molecule_feature
 
         node = {}
         node['feature'] = atom_feature
@@ -288,7 +296,7 @@ class MolGraphFeaturizer(Featurizer):
                 edge['feature'] = self._expand_bond_features(
                     mol, paths, bond_feature,
                 )
-            edge['length'] = np.eye(self.radius + 1)[edge['length']]
+            edge['length'] = np.eye(self.radius + 1, dtype=self.feature_dtype)[edge['length']]
 
         if self.super_atom:
             node, edge = self._add_super_atom(node, edge)
@@ -315,13 +323,13 @@ class MolGraphFeaturizer(Featurizer):
         )
         return bond_feature.astype(self.feature_dtype)
     
-    def context_feature(self, mol: chem.Mol) -> np.ndarray:
+    def molecule_feature(self, mol: chem.Mol) -> np.ndarray:
         if self._molecule_features is None:
             return None
-        context_feature: np.ndarray = np.concatenate(
+        molecule_feature: np.ndarray = np.concatenate(
             [f(mol) for f in self._molecule_features], axis=-1
         )
-        return context_feature.astype(self.feature_dtype)
+        return molecule_feature.astype(self.feature_dtype)
 
     def num_atoms(self, mol: chem.Mol) -> np.ndarray:
         return np.asarray(mol.num_atoms, dtype=self.index_dtype)
@@ -361,9 +369,7 @@ class MolGraphFeaturizer(Featurizer):
     ) -> tuple[dict[str, np.ndarray]]:
         num_super_nodes = 1 
         num_nodes = node['feature'].shape[0]
-        node = _add_super_nodes(
-            node, num_super_nodes, self.feature_dtype
-        )
+        node = _add_super_nodes(node, num_super_nodes)
         edge = _add_super_edges(
             edge, num_nodes, num_super_nodes, self.feature_dtype, self.index_dtype
         )
@@ -402,6 +408,7 @@ class MolGraphFeaturizer(Featurizer):
         return cls(**config)
     
 
+@keras.saving.register_keras_serializable(package='molcraft')
 class MolGraphFeaturizer3D(MolGraphFeaturizer):
 
     """Molecular 3d-graph featurizer.
@@ -543,7 +550,7 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
                 'of the `Featurizer` or input a 3D representation of the molecule. '
             )
 
-        context_feature = self.context_feature(mol)
+        molecule_feature = self.molecule_feature(mol)
         molecule_size = self.num_atoms(mol) + int(self.super_atom)
 
         if isinstance(context, dict):
@@ -565,8 +572,14 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
         else:
             context = {'size': molecule_size}
 
-        if context_feature is not None:
-            context['feature'] = context_feature
+        if molecule_feature is not None:
+            if 'feature' in context:
+                warn(
+                    'Found both inputted and computed context feature. '
+                    'Overwriting inputted context feature with computed '
+                    'context feature (based on `molecule_features`).'
+                )
+            context['feature'] = molecule_feature
             
         node = {}
         node['feature'] = self.atom_features(mol)
@@ -586,7 +599,7 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
                 radius=self.radius, 
                 sparse=False, 
                 self_loops=self.self_loops, 
-                dtype=np.bool
+                dtype=bool
             )
             edge_conformer['source'], edge_conformer['target'] = np.where(adjacency_matrix)
             edge_conformer['source'] = edge_conformer['source'].astype(self.index_dtype)
@@ -603,7 +616,7 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
                 )
                 node_conformer['coordinate'] = np.concatenate(
                     [node_conformer['coordinate'], conformer.centroid[None]], axis=0
-                )
+                ).astype(self.feature_dtype)
             tensor_list.append(
                 tensors.GraphTensor(context, node_conformer, edge_conformer)
             )
@@ -627,7 +640,7 @@ class MolGraphFeaturizer3D(MolGraphFeaturizer):
         config['conformer_generator'] = keras.saving.deserialize_keras_object(
             config['conformer_generator']
         )
-        return super().from_config(**config)
+        return super().from_config(config)
     
 
 def save_featurizer(
@@ -669,12 +682,15 @@ def load_featurizer(
 def _add_super_nodes(
     node: dict[str, np.ndarray], 
     num_super_nodes: int = 1,
-    feature_dtype: str = 'float32',
 ) -> dict[str, np.ndarray]:
     node = copy.deepcopy(node)
-    node['super'] = np.array([False] * len(node['feature']) + [True] * num_super_nodes)
+    node['super'] = np.array(
+        [False] * len(node['feature']) + [True] * num_super_nodes,
+        dtype=bool
+    )
     super_node_feature = np.zeros(
-        [num_super_nodes, node['feature'].shape[-1]], dtype=feature_dtype
+        [num_super_nodes, node['feature'].shape[-1]], 
+        dtype=node['feature'].dtype
     )
     node['feature'] = np.concatenate([node['feature'], super_node_feature])
     return node 
@@ -694,31 +710,38 @@ def _add_super_edges(
         np.tile(np.arange(num_nodes), [num_super_nodes])
     )
     edge['source'] = np.concatenate(
-        [
-            edge['source'], 
-            node_indices,
-            super_node_indices,
-        ]
-    )
-    edge['source'] = edge['source'].astype(index_dtype)
+        [edge['source'], node_indices, super_node_indices]
+    ).astype(index_dtype)
+   
     edge['target'] = np.concatenate(
-        [
-            edge['target'], 
-            super_node_indices, 
-            node_indices
-        ]
-    )
-    edge['target'] = edge['target'].astype(index_dtype)
+        [edge['target'], super_node_indices, node_indices]
+    ).astype(index_dtype)
+
     if 'feature' in edge:
-        edge['super'] = np.asarray([False] * edge['feature'].shape[0] + [True] * (num_super_nodes * num_nodes * 2))
-        edge['feature'] = np.concatenate([edge['feature'], np.zeros((num_super_nodes * num_nodes * 2, edge['feature'].shape[-1]))])
+        num_edges = int(edge['feature'].shape[0])
+        num_super_edges = int(num_super_nodes * num_nodes * 2)
+        edge['super'] = np.asarray(
+            ([False] * num_edges + [True] * num_super_edges),
+            dtype=bool
+        )
+        edge['feature'] = np.concatenate(
+            [
+                edge['feature'], 
+                np.zeros(
+                    shape=(num_super_edges, edge['feature'].shape[-1]),
+                    dtype=edge['feature'].dtype
+                )
+            ]
+        )
+
     if 'length' in edge:
         edge['length'] = np.pad(edge['length'], [(0, 0), (1, 0)])
-        zero_array = np.zeros((num_nodes * num_super_nodes * 2,), dtype='int32')
+        zero_array = np.zeros([num_nodes * num_super_nodes * 2], dtype='int32')
         edge_length_dim = edge['length'].shape[1]
         virtual_edge_length = np.eye(edge_length_dim)[zero_array]
         edge['length'] = np.concatenate([edge['length'], virtual_edge_length])
         edge['length'] = edge['length'].astype(feature_dtype)
+
     return edge