molcraft 0.1.0a15__tar.gz → 0.1.0a17__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: molcraft
-Version: 0.1.0a15
+Version: 0.1.0a17
 Summary: Graph Neural Networks for Molecular Machine Learning
 Author-email: Alexander Kensert <alexander.kensert@gmail.com>
 License: MIT License
@@ -35,6 +35,7 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: tensorflow>=2.16
+Requires-Dist: tensorflow-text>=2.16
 Requires-Dist: rdkit>=2023.9.5
 Requires-Dist: pandas>=1.0.3
 Requires-Dist: ipython>=8.12.0
@@ -42,9 +43,9 @@ 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">
+<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. 
+**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML.
 
 > [!NOTE]  
 > In progress.
@@ -82,11 +83,12 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.BondType(),
         features.IsRotatable(),
     ],
-    super_atom=True,
+    super_node=True,
     self_loops=True,
+    include_hydrogens=False,
 )
 
-graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])
+graph = featurizer([('N[C@@H](C)C(=O)O', 2.5), ('N[C@@H](CS)C(=O)O', 1.5)])
 print(graph)
 
 model = models.GraphModel.from_layers(
@@ -94,13 +96,13 @@ model = models.GraphModel.from_layers(
         layers.Input(graph.spec),
         layers.NodeEmbedding(dim=128),
         layers.EdgeEmbedding(dim=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.Readout(mode='mean'),
-        keras.layers.Dense(units=1024, activation='relu'),
-        keras.layers.Dense(units=1024, activation='relu'),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.Readout(),
+        keras.layers.Dense(units=1024, activation='elu'),
+        keras.layers.Dense(units=1024, activation='elu'),
         keras.layers.Dense(1)
     ]
 )

{molcraft-0.1.0a15 → molcraft-0.1.0a17}/README.md

@@ -1,6 +1,6 @@
-<img src="https://github.com/akensert/molcraft/blob/main/docs/_static/molcraft-logo.png" alt="molcraft-logo">
+<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. 
+**Deep Learning on Molecules**: A Minimalistic GNN package for Molecular ML.
 
 > [!NOTE]  
 > In progress.
@@ -38,11 +38,12 @@ featurizer = featurizers.MolGraphFeaturizer(
         features.BondType(),
         features.IsRotatable(),
     ],
-    super_atom=True,
+    super_node=True,
     self_loops=True,
+    include_hydrogens=False,
 )
 
-graph = featurizer([('N[C@@H](C)C(=O)O', 2.0), ('N[C@@H](CS)C(=O)O', 1.0)])
+graph = featurizer([('N[C@@H](C)C(=O)O', 2.5), ('N[C@@H](CS)C(=O)O', 1.5)])
 print(graph)
 
 model = models.GraphModel.from_layers(
@@ -50,13 +51,13 @@ model = models.GraphModel.from_layers(
         layers.Input(graph.spec),
         layers.NodeEmbedding(dim=128),
         layers.EdgeEmbedding(dim=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.GraphTransformer(units=128),
-        layers.Readout(mode='mean'),
-        keras.layers.Dense(units=1024, activation='relu'),
-        keras.layers.Dense(units=1024, activation='relu'),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.GraphConv(units=128),
+        layers.Readout(),
+        keras.layers.Dense(units=1024, activation='elu'),
+        keras.layers.Dense(units=1024, activation='elu'),
         keras.layers.Dense(1)
     ]
 )

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

@@ -1,4 +1,4 @@
-__version__ = '0.1.0a15'
+__version__ = '0.1.0a17'
 
 import os
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
@@ -6,7 +6,6 @@ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 from molcraft import chem
 from molcraft import features
 from molcraft import descriptors
-from molcraft import conformers 
 from molcraft import featurizers
 from molcraft import layers 
 from molcraft import models 

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

@@ -0,0 +1,194 @@
+import re
+import keras
+import numpy as np
+import tensorflow as tf
+import tensorflow_text as tf_text
+
+from molcraft import featurizers
+from molcraft import tensors
+from molcraft import layers
+from molcraft import models 
+from molcraft import chem
+
+
+# TODO: Add regex pattern for residue (C-term mod + N-term mod)?
+# TODO: Add regex pattern for residue (C-term mod + N-term mod + mod)?
+residue_pattern: str = "|".join([
+    r'(\[[A-Za-z0-9]+\]-[A-Z]\[[A-Za-z0-9]+\])', # residue (N-term mod + mod)
+    r'([A-Z]\[[A-Za-z0-9]+\]-\[[A-Za-z0-9]+\])', # residue (C-term mod + mod)
+    r'([A-Z]-\[[A-Za-z0-9]+\])', # residue (C-term mod)
+    r'(\[[A-Za-z0-9]+\]-[A-Z])', # residue (N-term mod)
+    r'([A-Z]\[[A-Za-z0-9]+\])', # residue (mod)
+    r'([A-Z])', # residue (no mod)
+])
+
+default_residues: dict[str, str] = {
+    "A": "N[C@@H](C)C(=O)O",
+    "C": "N[C@@H](CS)C(=O)O",
+    "D": "N[C@@H](CC(=O)O)C(=O)O",
+    "E": "N[C@@H](CCC(=O)O)C(=O)O",
+    "F": "N[C@@H](Cc1ccccc1)C(=O)O",
+    "G": "NCC(=O)O",
+    "H": "N[C@@H](CC1=CN=C-N1)C(=O)O",
+    "I": "N[C@@H](C(CC)C)C(=O)O",
+    "K": "N[C@@H](CCCCN)C(=O)O",
+    "L": "N[C@@H](CC(C)C)C(=O)O",
+    "M": "N[C@@H](CCSC)C(=O)O",
+    "N": "N[C@@H](CC(=O)N)C(=O)O",
+    "P": "N1[C@@H](CCC1)C(=O)O",
+    "Q": "N[C@@H](CCC(=O)N)C(=O)O",
+    "R": "N[C@@H](CCCNC(=N)N)C(=O)O",
+    "S": "N[C@@H](CO)C(=O)O",
+    "T": "N[C@@H](C(O)C)C(=O)O",
+    "V": "N[C@@H](C(C)C)C(=O)O",
+    "W": "N[C@@H](CC(=CN2)C1=C2C=CC=C1)C(=O)O",
+    "Y": "N[C@@H](Cc1ccc(O)cc1)C(=O)O",
+}
+
+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')
+    
+    
+class Peptide(chem.Mol):
+
+    @classmethod
+    def from_sequence(cls, sequence: str, **kwargs) -> 'Peptide':
+        sequence = [
+            match.group(0) for match in re.finditer(residue_pattern, sequence)
+        ]
+        peptide_smiles = []
+        for i, residue in enumerate(sequence):
+            if i < len(sequence) - 1:
+                residue_smiles = registered_residues[residue + '*']
+            else:
+                residue_smiles = registered_residues[residue]
+            peptide_smiles.append(residue_smiles)
+        peptide_smiles = ''.join(peptide_smiles)
+        return super().from_encoding(peptide_smiles, **kwargs)
+
+
+@keras.saving.register_keras_serializable(package='proteomics')
+class ResidueEmbedding(keras.layers.Layer):
+
+    def __init__(
+        self, 
+        featurizer: featurizers.MolGraphFeaturizer,
+        embedder: models.GraphModel, 
+        residues: dict[str, str] | None = None,
+        **kwargs
+    ) -> None:
+        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.ragged_split = SequenceSplitter(pad=False)
+        self.split = SequenceSplitter(pad=True)
+        self.use_cached_embeddings = tf.Variable(False)
+        self.supports_masking = True
+        
+    @property
+    def residues(self) -> dict[str, str]:
+        return self._residue_dict
+
+    @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]
+        self.mapping = tf.lookup.StaticHashTable(
+            tf.lookup.KeyValueTensorInitializer(
+                keys=residue_keys, 
+                values=range(num_residues)
+            ),
+            default_value=oov_value,
+        )
+        self.graph = tf.stack([
+            self.featurizer(residues[residue]) for residue in residue_keys
+        ], axis=0)
+        self.cached_embeddings = tf.Variable(
+            initial_value=tf.zeros((num_residues, self.embedding_dim))
+        )
+        _ = self.cache_and_get_embeddings()
+
+    def build(self, input_shape) -> None:
+        self.residues = self._residue_dict
+        super().build(input_shape)
+
+    def call(self, sequences: tf.Tensor, training: bool = None) -> tf.Tensor:
+        if training is False:
+            self.use_cached_embeddings.assign(True)
+        else:
+            self.use_cached_embeddings.assign(False)
+        embeddings = tf.cond(
+            pred=self.use_cached_embeddings,
+            true_fn=lambda: self.cached_embeddings,
+            false_fn=lambda: self.cache_and_get_embeddings(),
+        )
+        sequences = self.ragged_split(sequences)
+        sequences = keras.ops.concatenate([
+            tf.strings.join([sequences[:, :-1], '*']), sequences[:, -1:]
+        ], axis=1)
+        indices = self.mapping.lookup(sequences)
+        return tf.gather(embeddings, indices).to_tensor()
+    
+    def cache_and_get_embeddings(self) -> tf.Tensor:
+        embeddings = self.embedder(self.graph)
+        self.cached_embeddings.assign(embeddings)
+        return embeddings
+    
+    def compute_mask(
+        self, 
+        inputs: tensors.GraphTensor, 
+        mask: bool | None = None
+    ) -> tf.Tensor | None:
+        sequences = self.split(inputs)
+        return keras.ops.not_equal(sequences, '')
+                
+    def get_config(self) -> dict:
+        config = super().get_config()
+        config.update({
+            'featurizer': keras.saving.serialize_keras_object(self.featurizer),
+            'embedder': keras.saving.serialize_keras_object(self.embedder),
+            'residues': self._residue_dict,
+        })
+        return config
+    
+    @classmethod
+    def from_config(cls, config: dict) -> 'ResidueEmbedding':
+        config['featurizer'] = keras.saving.deserialize_keras_object(config['featurizer'])
+        config['embedder'] = keras.saving.deserialize_keras_object(config['embedder'])
+        return super().from_config(config)
+    
+
+@keras.saving.register_keras_serializable(package='proteomics')
+class SequenceSplitter(keras.layers.Layer): 
+
+    def __init__(self, pad: bool, **kwargs):
+        super().__init__(**kwargs)
+        self.pad = pad 
+
+    def call(self, inputs: tf.Tensor) -> tf.Tensor | tf.RaggedTensor:
+        inputs = tf_text.regex_split(inputs, residue_pattern, residue_pattern)
+        if self.pad:
+            inputs = inputs.to_tensor()
+        return inputs
+    
+
+registered_residues: dict[str, str] = {}
+register_residues(default_residues)

{molcraft-0.1.0a15 → molcraft-0.1.0a17}/molcraft/chem.py

@@ -19,8 +19,6 @@ class Mol(Chem.Mol):
     @classmethod
     def from_encoding(cls, encoding: str, explicit_hs: bool = False, **kwargs) -> 'Mol':
         rdkit_mol = get_mol(encoding, **kwargs)
-        if not rdkit_mol:
-            return None
         if explicit_hs:
             rdkit_mol = Chem.AddHs(rdkit_mol) 
         rdkit_mol.__class__ = cls 
@@ -102,21 +100,13 @@ class Mol(Chem.Mol):
 
     def get_conformer(self, index: int = 0) -> 'Conformer':
         if self.num_conformers == 0:
-            warnings.warn(
-                'Molecule has no conformer. To embed conformer(s), invoke the `embed` method, '
-                'and optionally followed by `minimize()` to perform force field minimization.',
-                stacklevel=2
-            )
+            warnings.warn('Molecule has no conformer.')
             return None
         return Conformer.cast(self.GetConformer(index))
     
     def get_conformers(self) -> list['Conformer']:
         if self.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.',
-                stacklevel=2
-            )
+            warnings.warn('Molecule has no conformer.')
             return []
         return [Conformer.cast(x) for x in self.GetConformers()]
      
@@ -222,11 +212,10 @@ def get_mol(
     else:
         mol = Chem.MolFromSmiles(encoding, sanitize=False)
     if mol is not None:
-        return sanitize_mol(mol, strict, assign_stereo_chemistry)
-    raise ValueError(
-        f"{encoding} is invalid; "
-        f"make sure {encoding} is a valid SMILES or InChI string."
-    )
+        mol = sanitize_mol(mol, strict, assign_stereo_chemistry)
+    if mol is not None:
+        return mol
+    raise ValueError(f'Could not obtain `chem.Mol` from {encoding}.')
 
 def get_adjacency_matrix(
     mol: Chem.Mol,
@@ -402,8 +391,9 @@ def embed_conformers(
     mol: Mol, 
     num_conformers: int, 
     method: str = 'ETKDGv3',
+    random_seed: int | None = None, 
     **kwargs
-) -> None:
+) -> Mol:
     available_embedding_methods = {
         'ETDG': rdDistGeom.ETDG(),
         'ETKDG': rdDistGeom.ETKDG(),
@@ -423,6 +413,9 @@ def embed_conformers(
     for key, value in kwargs.items():
         setattr(embedding_method, key, value)
 
+    if random_seed is not None:
+        embedding_method.randomSeed = random_seed
+
     success = rdDistGeom.EmbedMultipleConfs(
         mol, numConfs=num_conformers, params=embedding_method
     )
@@ -440,6 +433,8 @@ def embed_conformers(
             fallback_embedding_method.useRandomCoords = True
             fallback_embedding_method.maxAttempts = max_attempts
             fallback_embedding_method.clearConfs = False
+            if random_seed is not None:
+                fallback_embedding_method.randomSeed = random_seed
             success = rdDistGeom.EmbedMultipleConfs(
                 mol, numConfs=(num_conformers - num_successes), params=fallback_embedding_method
             )
@@ -459,7 +454,7 @@ def optimize_conformers(
     num_threads: bool = 1,
     ignore_interfragment_interactions: bool = True,
     vdw_threshold: float = 10.0,
-):
+) -> Mol:
     available_force_field_methods = [
         'MMFF', 'MMFF94', 'MMFF94s', 'UFF'
     ]
@@ -502,7 +497,7 @@ def prune_conformers(
     keep: int = 1,
     threshold: float = 0.0,
     energy_force_field: str = 'UFF',
-):
+) -> Mol:
     if mol.num_conformers == 0:
         warnings.warn(
             'Molecule has no conformers. To embed conformers, invoke the `embed` method, '
@@ -539,7 +534,7 @@ def _uff_optimize_conformers(
     vdw_threshold: float = 10.0,
     ignore_interfragment_interactions: bool = True,
     **kwargs,
-) -> Mol:
+) -> tuple[list[float], list[bool]]:
     """Universal Force Field Minimization.
     """
     results = rdForceFieldHelpers.UFFOptimizeMoleculeConfs(
@@ -560,7 +555,7 @@ def _mmff_optimize_conformers(
     variant: str = 'MMFF94',
     ignore_interfragment_interactions: bool = True,
     **kwargs,
-) -> Mol:
+) -> tuple[list[float], list[bool]]:
     """Merck Molecular Force Field Minimization.
     """
     if not rdForceFieldHelpers.MMFFHasAllMoleculeParams(mol):

{molcraft-0.1.0a15 → molcraft-0.1.0a17}/molcraft/datasets.py

@@ -11,7 +11,7 @@ def split(
     test_size: float | None = None,
     groups: str | np.ndarray = None,
     shuffle: bool = False, 
-    random_state: int | None = None,
+    random_seed: int | None = None,
 ) -> tuple[np.ndarray | pd.DataFrame, ...]:
     """Splits the dataset into subsets.
 
@@ -28,7 +28,7 @@ def split(
             The groups to perform the splitting on.
         shuffle:
             Whether the dataset should be shuffled prior to splitting.
-        random_state:
+        random_seed:
             The random state/seed. Only applicable if shuffling.
     """
     if not isinstance(data, (pd.DataFrame, np.ndarray)):
@@ -69,7 +69,7 @@ def split(
     train_size += remainder
 
     if shuffle:
-        np.random.seed(random_state)
+        np.random.seed(random_seed)
         np.random.shuffle(indices)
 
     train_mask = np.isin(groups, indices[:train_size])
@@ -84,7 +84,7 @@ def cv_split(
     num_splits: int = 10,
     groups: str | np.ndarray = None,
     shuffle: bool = False, 
-    random_state: int | None = None,
+    random_seed: int | None = None,
 ) -> typing.Iterator[
         tuple[np.ndarray | pd.DataFrame, np.ndarray | pd.DataFrame]
     ]:
@@ -99,7 +99,7 @@ def cv_split(
             The groups to perform the splitting on.
         shuffle:
             Whether the dataset should be shuffled prior to splitting.
-        random_state:
+        random_seed:
             The random state/seed. Only applicable if shuffling.
     """
     if not isinstance(data, (pd.DataFrame, np.ndarray)):
@@ -119,7 +119,7 @@ def cv_split(
             f'the data size or the number of groups ({size}).'
         )
     if shuffle:
-        np.random.seed(random_state)
+        np.random.seed(random_seed)
         np.random.shuffle(indices)
 
     indices_splits = np.array_split(indices, num_splits)

{molcraft-0.1.0a15 → molcraft-0.1.0a17}/molcraft/descriptors.py

@@ -91,3 +91,17 @@ class NumRings(Descriptor):
     def call(self, mol: chem.Mol) -> np.ndarray:
         return rdMolDescriptors.CalcNumRings(mol) 
 
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class AtomCount(Descriptor):
+
+    def __init__(self, atom_type: str, **kwargs):
+        super().__init__(**kwargs)
+        self.atom_type = atom_type
+
+    def call(self, mol: chem.Mol) -> np.ndarray:
+        count = 0
+        for atom in mol.atoms:
+            if atom.GetSymbol() == self.atom_type:
+                count += 1
+        return count

{molcraft-0.1.0a15 → molcraft-0.1.0a17}/molcraft/features.py

@@ -41,11 +41,7 @@ class Feature(abc.ABC):
     
     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`.'
-            )
+            raise TypeError(f'Input to {self.name} must be a `chem.Mol` instance.')
         features = self.call(mol)
         if len(features) != mol.num_atoms and len(features) != mol.num_bonds:
             raise ValueError(
@@ -119,59 +115,6 @@ class Feature(abc.ABC):
         return np.asarray([value], dtype=self.dtype)
     
 
-@keras.saving.register_keras_serializable(package='molcraft')
-class EdgeFeature(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`.'
-            )
-        features = self.call(mol)
-        if len(features) != int(mol.num_atoms**2):
-            raise ValueError(
-                f'The number of features computed by {self.name} does not '
-                'match the number of node pairs in the `chem.Mol` object. '
-                f'Make sure the list of items returned by {self.name}(input) '
-                'correspond to node/atom pairs: '
-                '[(0, 0), (0, 1), ..., (0, N), (1, 0), ... (N, N)], '
-                'where N denotes the number of nodes/atoms.'
-            )
-        func = (
-            self._featurize_categorical if self.vocab else 
-            self._featurize_floating
-        )
-        return np.asarray([func(x) for x in features], dtype=self.dtype)
-        
-
-@keras.saving.register_keras_serializable(package='molcraft')
-class Distance(EdgeFeature):
-    
-    def __init__(
-        self, 
-        max_distance: int = None, 
-        allow_oov: int = True,
-        encode_oov: bool = True, 
-        **kwargs,
-    ) -> None:
-        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, 
-            encode_oov=encode_oov, 
-            **kwargs
-        )
-
-    def call(self, mol: chem.Mol) -> list[int]:
-        return [int(x) for x in chem.get_distances(mol).reshape(-1)]
-    
-
 @keras.saving.register_keras_serializable(package='molcraft')
 class AtomType(Feature):
     def call(self, mol: chem.Mol) -> list[int, float, str]: 
@@ -340,6 +283,55 @@ class IsRotatable(Feature):
         return chem.rotatable_bonds(mol)
 
 
+@keras.saving.register_keras_serializable(package='molcraft')
+class PairFeature(Feature):
+
+    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.')
+        features = self.call(mol)
+        if len(features) != int(mol.num_atoms**2):
+            raise ValueError(
+                f'The number of features computed by {self.name} does not '
+                'match the number of node/atom pairs in the `chem.Mol` object. '
+                f'Make sure the list of items returned by {self.name}(input) '
+                'correspond to node/atom pairs: '
+                '[(0, 0), (0, 1), ..., (0, N), (1, 0), ... (N, N)], '
+                'where N denotes the number of nodes/atoms.'
+            )
+        func = (
+            self._featurize_categorical if self.vocab else 
+            self._featurize_floating
+        )
+        return np.asarray([func(x) for x in features], dtype=self.dtype)
+        
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class PairDistance(PairFeature):
+    
+    def __init__(
+        self, 
+        max_distance: int = None, 
+        allow_oov: int = True,
+        encode_oov: bool = True, 
+        **kwargs,
+    ) -> None:
+        vocab = kwargs.pop('vocab', None)
+        if not vocab:
+            if max_distance is None:
+                max_distance = 10
+            vocab = list(range(max_distance + 1))
+        super().__init__(
+            vocab=vocab, 
+            allow_oov=allow_oov, 
+            encode_oov=encode_oov, 
+            **kwargs
+        )
+
+    def call(self, mol: chem.Mol) -> list[int]:
+        return [int(x) for x in chem.get_distances(mol).reshape(-1)]
+    
+    
 default_vocabulary = {
     'AtomType': [
         '*', 'H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na',