Uncertainty Quantification#

[1]:

# Install chemprop from GitHub if running in Google Colab
import os

if os.getenv("COLAB_RELEASE_TAG"):
    try:
        import chemprop
    except ImportError:
        !git clone https://github.com/chemprop/chemprop.git
        %cd chemprop
        !pip install .
        %cd examples

Import packages#

[2]:

import pandas as pd
import numpy as np
import torch
import pandas as pd
from pathlib import Path

from lightning import pytorch as pl
from lightning.pytorch.callbacks import ModelCheckpoint

from chemprop import data, models, nn, uncertainty
from chemprop.models import save_model, load_model
from chemprop.cli.conf import NOW
from chemprop.cli.predict import find_models

%load_ext autoreload
%autoreload 2

Training#

Loda data#

[3]:

chemprop_dir = Path.cwd().parent
input_path = (
    chemprop_dir / "tests" / "data" / "regression" / "mol" / "mol.csv"
)  # path to your data .csv file
df_input = pd.read_csv(input_path)
smis = df_input.loc[:, "smiles"].values
ys = df_input.loc[:, ["lipo"]].values
all_data = [data.MoleculeDatapoint.from_smi(smi, y) for smi, y in zip(smis, ys)]

[4]:

mols = [d.mol for d in all_data]  # RDkit Mol objects are use for structure based splits
train_indices, val_indices, test_indices = data.make_split_indices(mols, "random", (0.8, 0.1, 0.1))
train_data, val_data, test_data = data.split_data_by_indices(
    all_data, train_indices, val_indices, test_indices
)

The return type of make_split_indices has changed in v2.1 - see help(make_split_indices)

[5]:

train_dset = data.MoleculeDataset(train_data[0])
scaler = train_dset.normalize_targets()

val_dset = data.MoleculeDataset(val_data[0])
val_dset.normalize_targets(scaler)

test_dset = data.MoleculeDataset(test_data[0])

[6]:

train_loader = data.build_dataloader(train_dset)
val_loader = data.build_dataloader(val_dset, shuffle=False)
test_loader = data.build_dataloader(test_dset, shuffle=False)

Constructs MPNN#

A Message passing constructs molecular graphs using message passing to learn node-level hidden representations.
An Aggregation is responsible for constructing a graph-level representation from the set of node-level representations after message passing.
A FFN takes the aggregated representations and make target predictions. To obtain uncertainty predictions, the FFN must be modified accordingly.

For regression:
- ffn = nn.RegressionFFN()
- ffn = nn.MveFFN()
- ffn = nn.EvidentialFFN()
For classification:
- ffn = nn.BinaryClassificationFFN()
- ffn = nn.BinaryDirichletFFN()
- ffn = nn.MulticlassClassificationFFN()
- ffn = nn.MulticlassDirichletFFN()
For spectral:
- ffn = nn.SpectralFFN() # will be available in future version

[7]:

mp = nn.BondMessagePassing()
agg = nn.MeanAggregation()
output_transform = nn.UnscaleTransform.from_standard_scaler(scaler)
# Change to other predictor if needed.
ffn = nn.MveFFN(output_transform=output_transform)
mpnn = models.MPNN(mp, agg, ffn, batch_norm=False)
mpnn

[7]:

MPNN(
  (message_passing): BondMessagePassing(
    (W_i): Linear(in_features=86, out_features=300, bias=False)
    (W_h): Linear(in_features=300, out_features=300, bias=False)
    (W_o): Linear(in_features=372, out_features=300, bias=True)
    (dropout): Dropout(p=0.0, inplace=False)
    (tau): ReLU()
    (V_d_transform): Identity()
    (graph_transform): Identity()
  )
  (agg): MeanAggregation()
  (bn): Identity()
  (predictor): MveFFN(
    (ffn): MLP(
      (0): Sequential(
        (0): Linear(in_features=300, out_features=300, bias=True)
      )
      (1): Sequential(
        (0): ReLU()
        (1): Dropout(p=0.0, inplace=False)
        (2): Linear(in_features=300, out_features=2, bias=True)
      )
    )
    (criterion): MVELoss(task_weights=[[1.0]])
    (output_transform): UnscaleTransform()
  )
  (X_d_transform): Identity()
  (metrics): ModuleList(
    (0): MSE(task_weights=[[1.0]])
    (1): MVELoss(task_weights=[[1.0]])
  )
)

Set up trainer#

[8]:

model_output_dir = Path(f"chemprop_training/{NOW}")
monitor_mode = "min" if mpnn.metrics[0].higher_is_better else "max"
checkpointing = ModelCheckpoint(
    model_output_dir / "checkpoints",
    "best-{epoch}-{val_loss:.2f}",
    "val_loss",
    mode=monitor_mode,
    save_last=True,
)

[9]:

trainer = pl.Trainer(
    logger=False,
    enable_checkpointing=True,
    enable_progress_bar=False,
    accelerator="cpu",
    callbacks=[checkpointing],
    devices=1,
    max_epochs=20,
)

GPU available: False, used: False
TPU available: False, using: 0 TPU cores

Start training#

[10]:

trainer.fit(mpnn, train_loader, val_loader)

Loading `train_dataloader` to estimate number of stepping batches.
/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.

┏━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━┳━━━━━━━┓
┃   ┃ Name            ┃ Type               ┃ Params ┃ Mode  ┃ FLOPs ┃
┡━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━╇━━━━━━━┩
│ 0 │ message_passing │ BondMessagePassing │  227 K │ train │     0 │
│ 1 │ agg             │ MeanAggregation    │      0 │ train │     0 │
│ 2 │ bn              │ Identity           │      0 │ train │     0 │
│ 3 │ predictor       │ MveFFN             │ 90.9 K │ train │     0 │
│ 4 │ X_d_transform   │ Identity           │      0 │ train │     0 │
│ 5 │ metrics         │ ModuleList         │      0 │ train │     0 │
└───┴─────────────────┴────────────────────┴────────┴───────┴───────┘

Trainable params: 318 K
Non-trainable params: 0
Total params: 318 K
Total estimated model params size (MB): 1
Modules in train mode: 24
Modules in eval mode: 0
Total FLOPs: 0

/home/knathan/anaconda3/envs/chemprop/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.
`Trainer.fit` stopped: `max_epochs=20` reached.

Save the best model#

[11]:

best_model_path = checkpointing.best_model_path
model = mpnn.__class__.load_from_checkpoint(best_model_path)
p_model = model_output_dir / "best.pt"
save_model(p_model, model)

Predicting#

Change model input here#

[12]:

chemprop_dir = Path.cwd().parent
test_path = chemprop_dir / "tests" / "data" / "regression" / "mol" / "mol.csv"
df_test = pd.read_csv(test_path)
test_dset = data.MoleculeDataset(test_data[0])
test_loader = data.build_dataloader(test_dset, shuffle=False)
df_test

[12]:

	smiles	lipo
0	Cn1c(CN2CCN(CC2)c3ccc(Cl)cc3)nc4ccccc14	3.54
1	COc1cc(OC)c(cc1NC(=O)CSCC(=O)O)S(=O)(=O)N2C(C)...	-1.18
2	COC(=O)[C@@H](N1CCc2sccc2C1)c3ccccc3Cl	3.69
3	OC[C@H](O)CN1C(=O)C(Cc2ccccc12)NC(=O)c3cc4cc(C...	3.37
4	Cc1cccc(C[C@H](NC(=O)c2cc(nn2C)C(C)(C)C)C(=O)N...	3.10
...	...	...
95	CC(C)N(CCCNC(=O)Nc1ccc(cc1)C(C)(C)C)C[C@H]2O[C...	2.20
96	CCN(CC)CCCCNc1ncc2CN(C(=O)N(Cc3cccc(NC(=O)C=C)...	2.04
97	CCSc1c(Cc2ccccc2C(F)(F)F)sc3N(CC(C)C)C(=O)N(C)...	4.49
98	COc1ccc(Cc2c(N)n[nH]c2N)cc1	0.20
99	CCN(CCN(C)C)S(=O)(=O)c1ccc(cc1)c2cnc(N)c(n2)C(...	2.00

100 rows × 2 columns

[13]:

# use the validation set from the training as the calibration set as an example
cal_dset = data.MoleculeDataset(val_data[0])
cal_loader = data.build_dataloader(cal_dset, shuffle=False)

Constructs uncertainty estimator#

An uncertianty estimator can make model predictions and associated uncertainty predictions.

Available options can be found in uncertainty.UncertaintyEstimatorRegistry.

[14]:

print(uncertainty.UncertaintyEstimatorRegistry)

ClassRegistry {
    'none': <class 'chemprop.uncertainty.estimator.NoUncertaintyEstimator'>,
    'mve': <class 'chemprop.uncertainty.estimator.MVEEstimator'>,
    'ensemble': <class 'chemprop.uncertainty.estimator.EnsembleEstimator'>,
    'classification': <class 'chemprop.uncertainty.estimator.ClassEstimator'>,
    'evidential-total': <class 'chemprop.uncertainty.estimator.EvidentialTotalEstimator'>,
    'evidential-epistemic': <class 'chemprop.uncertainty.estimator.EvidentialEpistemicEstimator'>,
    'evidential-aleatoric': <class 'chemprop.uncertainty.estimator.EvidentialAleatoricEstimator'>,
    'dropout': <class 'chemprop.uncertainty.estimator.DropoutEstimator'>,
    'classification-dirichlet': <class 'chemprop.uncertainty.estimator.ClassificationDirichletEstimator'>,
    'multiclass-dirichlet': <class 'chemprop.uncertainty.estimator.MulticlassDirichletEstimator'>,
    'quantile-regression': <class 'chemprop.uncertainty.estimator.QuantileRegressionEstimator'>
}

[15]:

unc_estimator = uncertainty.MVEEstimator()

Constructs uncertainty calibrator#

An uncertianty calibrator can calibrate the predicted uncertainties.

Available options can be found in uncertainty.UncertaintyCalibratorRegistry.

For regression:

ZScalingCalibrator
ZelikmanCalibrator
MVEWeightingCalibrator
RegressionConformalCalibrator

For binary classification:

PlattCalibrator
IsotonicCalibrator
MultilabelConformalCalibrator

For multiclass classification:

MulticlassConformalCalibrator
AdaptiveMulticlassConformalCalibrator
IsotonicMulticlassCalibrator

[16]:

print(uncertainty.UncertaintyCalibratorRegistry)

ClassRegistry {
    'zscaling': <class 'chemprop.uncertainty.calibrator.ZScalingCalibrator'>,
    'zelikman-interval': <class 'chemprop.uncertainty.calibrator.ZelikmanCalibrator'>,
    'mve-weighting': <class 'chemprop.uncertainty.calibrator.MVEWeightingCalibrator'>,
    'conformal-regression': <class 'chemprop.uncertainty.calibrator.RegressionConformalCalibrator'>,
    'platt': <class 'chemprop.uncertainty.calibrator.PlattCalibrator'>,
    'isotonic': <class 'chemprop.uncertainty.calibrator.IsotonicCalibrator'>,
    'conformal-multilabel': <class 'chemprop.uncertainty.calibrator.MultilabelConformalCalibrator'>,
    'conformal-multiclass': <class 'chemprop.uncertainty.calibrator.MulticlassConformalCalibrator'>,
    'conformal-adaptive': <class 'chemprop.uncertainty.calibrator.AdaptiveMulticlassConformalCalibrator'>,
    'isotonic-multiclass': <class 'chemprop.uncertainty.calibrator.IsotonicMulticlassCalibrator'>
}

[17]:

unc_calibrator = uncertainty.ZScalingCalibrator()

Constructs uncertainty evaluator#

An uncertianty evaluator can evaluates the quality of uncertainty estimates.

Available options can be found in uncertainty.UncertaintyEvaluatorRegistry.