Neural Fingerprints

We create atom, bond, and edge tensors from molecule SMILES using chemml.chem.tensorize_molecules in order to build neural fingerprints using chemml.models.NeuralGraphHidden and chemml.models.NeuralGraphOutput modules. These neural fingerprints are then used as features to train a simple feed forward neural network to predict densities of small organic compounds using tensorflow.

Here we import a sample dataset from ChemML library which has the SMILES codes for 500 small organic molecules with their densities in \(kg/m^3\).

import numpy as np
from chemml.datasets import load_organic_density
molecules, target, dragon_subset = load_organic_density()
target = np.asarray(target['density_Kg/m3'])

Building chemml.chem.Molecule objects from molecule SMILES.

from chemml.chem import Molecule
mol_objs_list = []
for smi in molecules['smiles']:
    mol = Molecule(smi, 'smiles')
    mol.hydrogens('add')
    mol.to_xyz('MMFF', maxIters=10000, mmffVariant='MMFF94s')
    mol_objs_list.append(mol)

Molecule tensors can be used to create neural graph fingerprints using chemml.models

from chemml.chem import tensorise_molecules
xatoms, xbonds, xedges = tensorise_molecules(molecules=mol_objs_list, max_degree=5,
                                        max_atoms=None, n_jobs=-1, batch_size=100, verbose=True)

Tensorising molecules in batches of 100 ...
500/500 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 9s 17ms/step
Merging batch tensors ...    [DONE]

Splitting and preprocessing the data

from sklearn.model_selection import ShuffleSplit
from sklearn.preprocessing import StandardScaler
y_scale = StandardScaler()
rs = ShuffleSplit(n_splits=1, test_size=.20, random_state=42)

for train, test in rs.split(mol_objs_list):
    xatoms_train = xatoms[train]
    xatoms_test = xatoms[test]
    xbonds_train = xbonds[train]
    xbonds_test = xbonds[test]
    xedges_train = xedges[train]
    xedges_test = xedges[test]
    target_train = target[train]
    target_test = target[test]
    target_train = y_scale.fit_transform(target_train.reshape(-1,1))

print('Training data:\n')
print('Atoms: ',xatoms_train.shape)
print('Bonds: ',xbonds_train.shape)
print('Edges: ',xedges_train.shape)
print('Target: ',target_train.shape)

print('\nTesting data:\n')
print('Atoms: ',xatoms_test.shape)
print('Bonds: ',xbonds_test.shape)
print('Edges: ',xedges_test.shape)
print('Target: ',target_test.shape)

Training data:

Atoms:  (400, 57, 62)
Bonds:  (400, 57, 5, 6)
Edges:  (400, 57, 5)
Target:  (400, 1)

Testing data:

Atoms:  (100, 57, 62)
Bonds:  (100, 57, 5, 6)
Edges:  (100, 57, 5)
Target:  (100,)

Building the Neural Fingerprints

The atom, bond, and edge tensors are used here to build 200 neural fingerprints of width 8 (i.e., the size atomic neighborhood which will be considered in the convolution process).

from chemml.models import NeuralGraphHidden, NeuralGraphOutput
from tensorflow.keras.layers import Input, add
import tensorflow as tf
tf.random.set_seed(42)

conv_width = 8
fp_length = 200

num_molecules = xatoms_train.shape[0]
max_atoms = xatoms_train.shape[1]
max_degree = xbonds_train.shape[2]
num_atom_features = xatoms_train.shape[-1]
num_bond_features = xbonds_train.shape[-1]

# Creating input layers for atoms ,bonds and edge information
atoms0 = Input(name='atom_inputs', shape=(max_atoms, num_atom_features),batch_size=None)
bonds = Input(name='bond_inputs', shape=(max_atoms, max_degree, num_bond_features),batch_size=None)
edges = Input(name='edge_inputs', shape=(max_atoms, max_degree), dtype='int32',batch_size=None)

# Defining the convolved atom feature layers
atoms1 = NeuralGraphHidden(conv_width, activation='relu', use_bias=False)([atoms0, bonds, edges])
atoms2 = NeuralGraphHidden(conv_width, activation='relu', use_bias=False)([atoms1, bonds, edges])

# Defining the outputs of each (convolved) atom feature layer to fingerprint
fp_out0 = NeuralGraphOutput(fp_length, activation='softmax')([atoms0,bonds,edges])
fp_out1 = NeuralGraphOutput(fp_length, activation='softmax')([atoms1,bonds,edges])
fp_out2 = NeuralGraphOutput(fp_length, activation='softmax')([atoms2,bonds,edges])

# Sum outputs to obtain fingerprint
final_fp = add([fp_out0, fp_out1, fp_out2])
print('Neural Fingerprint Shape: ',final_fp.shape)

Neural Fingerprint Shape:  (None, 200)

Building and training the neural network

Here, we build and train a simple feed forward neural network using tensorflow.keras and provide our neural fingerprints as features.

from tensorflow.keras import Model
from tensorflow.keras.layers import Dense

# Build and compile model for regression.
dense_layer0 = Dense(128,activation='relu',name='dense_layer0',
                     kernel_regularizer=tf.keras.regularizers.l2(0.01))(final_fp)
dense_layer1 = Dense(64,activation='relu',name='dense_layer1',
                     kernel_regularizer=tf.keras.regularizers.l2(0.01))(dense_layer0)
dense_layer2 = Dense(32,activation='relu',name='dense_layer2',
                     kernel_regularizer=tf.keras.regularizers.l2(0.01))(dense_layer1)

main_prediction = Dense(1, activation='linear', name='main_prediction')(dense_layer1)
model = Model(inputs=[atoms0, bonds, edges], outputs=[main_prediction])
model.compile(optimizer='adam', loss='mae')

# Show summary
model.summary()

model.fit([xatoms_train, xbonds_train, xedges_train], target_train, epochs=50,
          steps_per_epoch=None, batch_size=None,verbose=False,validation_split=0.1)

WARNING:tensorflow:TensorFlow GPU support is not available on native Windows for TensorFlow >= 2.11. Even if CUDA/cuDNN are installed, GPU will not be used. Please use WSL2 or the TensorFlow-DirectML plugin.

Model: "functional"

┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)        ┃ Output Shape      ┃    Param # ┃ Connected to      ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩
│ atom_inputs         │ (None, 57, 62)    │          0 │ -                 │
│ (InputLayer)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ bond_inputs         │ (None, 57, 5, 6)  │          0 │ -                 │
│ (InputLayer)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ edge_inputs         │ (None, 57, 5)     │          0 │ -                 │
│ (InputLayer)        │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ neural_graph_hidden │ (None, 57, 8)     │      2,720 │ atom_inputs[0][0… │
│ (NeuralGraphHidden) │                   │            │ bond_inputs[0][0… │
│                     │                   │            │ edge_inputs[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ neural_graph_hidde… │ (None, 57, 8)     │        560 │ neural_graph_hid… │
│ (NeuralGraphHidden) │                   │            │ bond_inputs[0][0… │
│                     │                   │            │ edge_inputs[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ neural_graph_output │ (None, 200)       │     13,800 │ atom_inputs[0][0… │
│ (NeuralGraphOutput) │                   │            │ bond_inputs[0][0… │
│                     │                   │            │ edge_inputs[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ neural_graph_outpu… │ (None, 200)       │      3,000 │ neural_graph_hid… │
│ (NeuralGraphOutput) │                   │            │ bond_inputs[0][0… │
│                     │                   │            │ edge_inputs[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ neural_graph_outpu… │ (None, 200)       │      3,000 │ neural_graph_hid… │
│ (NeuralGraphOutput) │                   │            │ bond_inputs[0][0… │
│                     │                   │            │ edge_inputs[0][0] │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ add (Add)           │ (None, 200)       │          0 │ neural_graph_out… │
│                     │                   │            │ neural_graph_out… │
│                     │                   │            │ neural_graph_out… │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dense_layer0        │ (None, 128)       │     25,728 │ add[0][0]         │
│ (Dense)             │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ dense_layer1        │ (None, 64)        │      8,256 │ dense_layer0[0][… │
│ (Dense)             │                   │            │                   │
├─────────────────────┼───────────────────┼────────────┼───────────────────┤
│ main_prediction     │ (None, 1)         │         65 │ dense_layer1[0][… │
│ (Dense)             │                   │            │                   │
└─────────────────────┴───────────────────┴────────────┴───────────────────┘

 Total params: 57,129 (223.16 KB)

 Trainable params: 57,129 (223.16 KB)

 Non-trainable params: 0 (0.00 B)

<keras.src.callbacks.history.History at 0x1d239e1f920>

Predicting the density of the molecules in our test data and evaluating our model based on it.

from chemml.utils import regression_metrics

y_pred = model.predict([xatoms_test,xbonds_test,xedges_test])
y_pred = y_scale.inverse_transform(y_pred)
metrics_df = regression_metrics(target_test, list(y_pred.reshape(-1,)))
mae = metrics_df['MAE'].values[0]
r_2 = metrics_df['r_squared'].values[0]

print("Mean Absolute Error = {} kg/m^3".format(mae.round(3)))
print("R squared = {}".format(r_2.round(3)))

4/4 ━━━━━━━━━━━━━━━━━━━━ 9s 2s/step
Mean Absolute Error = 17.415 kg/m^3
R squared = 0.928

metrics_df

	ME	MAE	MSE	RMSE	MSLE	RMSLE	MAPE	MaxAPE	RMSPE	MPE	MaxAE	deltaMaxE	r_squared	std
0	-10.108292	17.415105	551.830271	23.491068	0.000387	0.019676	1.420658	9.103122	2.012218	-0.868725	91.540991	127.90665	0.928118	87.617825

Unified Neural Graph Fingerprint API with Engine Selection

ChemML now provides a unified wrapper NeuralGraphFingerprint that supports multiple computation backends (TensorFlow and PyTorch). You can easily switch between engines without changing the core logic.

from chemml.models import NeuralGraphFingerprint
import pandas as pd
from time import time

# Example 1: Using PyTorch backend
pt_start_time = time()
print("=" * 60)
print("Example 1: PyTorch Backend")
print("=" * 60)
ngf_pytorch = NeuralGraphFingerprint(
    engine='pytorch',
    conv_width=8,
    fp_length=200,
    n_conv_layers=3,
    hidden_sizes = [128, 64, 32],
    epochs=50,
    batch_size=32,
    learning_rate=0.001,
    verbose=False
)

ngf_pytorch.fit(xatoms_train, xbonds_train, xedges_train, target_train)
y_pred_pt = y_scale.inverse_transform(ngf_pytorch.predict(xatoms_test, xbonds_test, xedges_test).reshape(-1, 1))
print(f"PyTorch predictions shape: {y_pred_pt.shape}")
print(f"Actual values (first 5): {target_test[:5]}")
print(f"PyTorch predictions (first 5): {y_pred_pt[:5].flatten()}")

pt_end_time = time()
print(f"PyTorch training and prediction time: {pt_end_time - pt_start_time:.2f} seconds")

# Example 2: Using TensorFlow backend
tf_start_time = time()
print("\n" + "=" * 60)
print("Example 2: TensorFlow Backend")
print("=" * 60)
ngf_tensorflow = NeuralGraphFingerprint(
    engine='tensorflow',
    conv_width=8,
    fp_length=200,
    n_conv_layers=3,
    hidden_sizes = [128, 64, 32],
    epochs=50,
    batch_size=32,
    learning_rate=0.001,
    verbose=False
)

# Same data, different backend
ngf_tensorflow.fit(xatoms_train, xbonds_train, xedges_train, target_train)
y_pred_tf = y_scale.inverse_transform(ngf_tensorflow.predict(xatoms_test, xbonds_test, xedges_test).reshape(-1, 1))
print(f"TensorFlow predictions shape: {y_pred_tf.shape}")
print(f"Actual values (first 5): {target_test[:5]}")
print(f"TensorFlow predictions (first 5): {y_pred_tf[:5].flatten()}")

tf_end_time = time()
print(f"TensorFlow training and prediction time: {tf_end_time - tf_start_time:.2f} seconds")

metrics_pt = regression_metrics(target_test, list(y_pred_pt.reshape(-1,)))
metrics_tf = regression_metrics(target_test, list(y_pred_tf.reshape(-1,)))

metrics_pt['engine'] = 'PyTorch'
metrics_tf['engine'] = 'TensorFlow'

pd.concat([metrics_pt, metrics_tf], ignore_index=True)

============================================================
Example 1: PyTorch Backend
============================================================
PyTorch predictions shape: (100, 1)
Actual values (first 5): [1328.35 1143.8  1351.33 1466.57 1207.67]
PyTorch predictions (first 5): [1339.304  1137.8044 1331.6082 1462.5497 1215.6892]
PyTorch training and prediction time: 13.98 seconds

============================================================
Example 2: TensorFlow Backend
============================================================
TensorFlow predictions shape: (100, 1)
Actual values (first 5): [1328.35 1143.8  1351.33 1466.57 1207.67]
TensorFlow predictions (first 5): [1358.3076 1153.0071 1357.1093 1474.0562 1220.1882]
TensorFlow training and prediction time: 74.82 seconds

	ME	MAE	MSE	RMSE	MSLE	RMSLE	MAPE	MaxAPE	RMSPE	MPE	MaxAE	deltaMaxE	r_squared	std	engine
0	3.434481	14.846350	344.743988	18.567283	0.000230	0.015162	1.194169	4.399324	1.516716	0.239773	46.250522	92.256938	0.955093	87.617825	PyTorch
1	-7.184584	14.876365	355.668758	18.859182	0.000236	0.015351	1.197706	5.650654	1.552435	-0.601583	56.822974	97.304233	0.953670	87.617825	TensorFlow

Simplified Workflow: Direct SMILES-to-Model

The NeuralGraphFingerprint wrapper also supports end-to-end workflows directly from SMILES strings or Molecule objects, eliminating the need for manual tensorization steps.

This section demonstrates the complete pipeline from loading the dataset to model prediction using the NeuralGraphFingerprint API.

# Force aggressive module reload
import sys
import importlib
import numpy as np
from sklearn.preprocessing import StandardScaler

# Remove all chemml modules from cache
modules_to_remove = [mod for mod in sys.modules if mod.startswith('chemml')]
for mod in modules_to_remove:
    del sys.modules[mod]

# Step 1: Load dataset using load_organic_density
from chemml.datasets import load_organic_density
smiles_df, density_df, features_df = load_organic_density()

# Step 2: Prepare data
smiles_all = smiles_df['smiles'].tolist()
smiles_idx = np.arange(len(smiles_all))
targets_all = density_df['density_Kg/m3'].values.astype('float32')

# Step 3: Train/Test split
from sklearn.model_selection import train_test_split
smiles_train, smiles_test, y_train, y_test = train_test_split(
    smiles_idx, targets_all, test_size=0.2, random_state=42
)

yscale = StandardScaler()
y_train = yscale.fit_transform(y_train.reshape(-1, 1))

smiles_train = [smiles_all[i] for i in smiles_train]
smiles_test = [smiles_all[i] for i in smiles_test]
print(f"Training set: {len(smiles_train)} molecules")
print(f"Test set: {len(smiles_test)} molecules")

# Step 4: Create and train model directly from SMILES (no manual tensorization!)
from chemml.models import NeuralGraphFingerprint
from time import time

start_time = time()

ngf_complete = NeuralGraphFingerprint(
    engine='pytorch',
    conv_width=8,
    fp_length=200,
    n_conv_layers=3,
    hidden_sizes=[128,64,32],
    epochs=50,
    batch_size=32,
    learning_rate=0.001,
    verbose=False,
    random_seed=42
)

# Fit the model directly on SMILES strings
ngf_complete.fit(smiles_train, y=y_train, tensorize_kwargs={'n_jobs': 1, 'verbose': False})

training_time = time() - start_time
print(f"Training completed in {training_time:.2f} seconds")

# Step 5: Make predictions on test set
print("\nGenerating predictions from raw SMILES...")
start_time = time()
y_pred = yscale.inverse_transform(ngf_complete.predict(smiles_test, tensorize_kwargs={'n_jobs': 1, 'verbose': False}).reshape(-1, 1))
prediction_time = time() - start_time
print(f"Predictions generated for {len(y_pred)} molecules in {prediction_time:.2f} seconds")

# Step 6: Evaluate results
from chemml.utils import regression_metrics
metrics = regression_metrics(y_test, list(y_pred.reshape(-1,)))
display(metrics)

Training set: 400 molecules
Test set: 100 molecules
Training completed in 18.20 seconds

Generating predictions from raw SMILES...
Predictions generated for 100 molecules in 8.62 seconds

	ME	MAE	MSE	RMSE	MSLE	RMSLE	MAPE	MaxAPE	RMSPE	MPE	MaxAE	deltaMaxE	r_squared	std
0	0.077003	14.059737	354.052032	18.816271	0.000221	0.014878	1.113103	5.090196	1.49254	-0.031188	56.511597	107.698608	0.953881	87.617821