Lab 4 solutions: a neural network for solubility

Lab 4: a neural network for solubility

Goal. Compare a small PyTorch MLP with ridge regression on ESOL, using training-only preprocessing, validation-based early stopping and model selection, then a final held-out test report.

Work through the numbered tasks. The student version replaces solution cells with # TODO prompts and retains checks. Use the stated variable names so each check can give useful feedback. All fits run on CPU; no GPU or paid service is required. Budget 90–120 minutes for the exercises, discussion and report.

The experiment deliberately uses a random split to connect with Lab 2. It does not establish performance on new chemical scaffolds. The winning validation score is part of model selection, not an unbiased final estimate.

# Setup: locally use the course environment; Colab installs missing packages only.
import importlib.util
import subprocess
import sys
try:
    import google.colab
    IN_COLAB = True
except ImportError:
    IN_COLAB = False
packages = {"numpy": "numpy", "pandas": "pandas", "matplotlib": "matplotlib",
            "sklearn": "scikit-learn", "torch": "torch==2.9.1"}
missing = [package for module, package in packages.items()
           if importlib.util.find_spec(module) is None]
if missing and IN_COLAB:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])
elif missing:
    raise RuntimeError("Install the course requirements first: " + ", ".join(missing))
print("Python:", sys.version.split()[0], "Colab:", IN_COLAB)

Python: 3.12.16 Colab: False

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sklearn
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from copy import deepcopy

torch.set_num_threads(1)
torch.manual_seed(42)
print("numpy:", np.__version__, "sklearn:", sklearn.__version__, "torch:", torch.__version__)
print("Device: CPU; seed: 42")

numpy: 2.5.4 sklearn: 1.9.1 torch: 2.9.1+cpu
Device: CPU; seed: 42

Data and provenance

ESOL contains 1,128 molecules and measured logS in log10(mol/L). We reuse five descriptors from Lab 2. The file is downloaded rather than redistributed; see the course data/README.md for the original Delaney publication and MoleculeNet provenance.

The loader first checks local parent directories. A raw-GitHub fallback supports checkouts that explicitly supply the data; the canonical DeepChem download handles the public course repository, where ESOL is not bundled. The SHA-256 must match before parsing any source.

from pathlib import Path
from urllib.request import urlopen
import hashlib
import io

expected_sha = "8c06a76f0c6487d29ab0f903e6a7a7139f189ab3c1178f159c8be8964602f189"
local_paths = [p / "data" / "esol.csv" for p in (Path.cwd(), *Path.cwd().parents)]
urls = ["https://raw.githubusercontent.com/Awguhst/MolModels/main/data/esol.csv",
        "https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/delaney-processed.csv"]
sources = [p for p in local_paths if p.is_file()] + urls
print("ESOL sources prepared.")

ESOL sources prepared.

for source in sources:
    try:
        if isinstance(source, Path):
            payload = source.read_bytes()
        else:
            with urlopen(source, timeout=30) as response:
                payload = response.read()
        if hashlib.sha256(payload).hexdigest() != expected_sha:
            continue
        esol = pd.read_csv(io.BytesIO(payload))
        break
    except OSError:
        continue
else:
    raise RuntimeError("No checksum-verified ESOL source was available; retry the data download.")
print("Verified ESOL:", esol.shape)

Verified ESOL: (1128, 10)

Task 1 · Features and target

Create features, X and y using Molecular Weight, Number of H-Bond Donors, Number of Rings, Number of Rotatable Bonds and Polar Surface Area, in that order. Use the measured solubility column as the target. Exclude the published ESOL prediction: it is another model’s output, not an independent descriptor. Check missing values, preserve row alignment and display a small preview.

# Task 1: define features, X and measured y in a fixed order.
features = ["Molecular Weight", "Number of H-Bond Donors", "Number of Rings",
            "Number of Rotatable Bonds", "Polar Surface Area"]
X = esol[features].copy()
y = esol["measured log solubility in mols per litre"].copy()
print("missing inputs / targets:", X.isna().sum().sum(), y.isna().sum())
X.head()

missing inputs / targets: 0 0

index	Molecular Weight	Number of H-Bond Donors	Number of Rings	Number of Rotatable Bonds	Polar Surface Area
0	457.432	7	3	7	202.32
1	201.225	1	2	2	42.24
2	152.237	0	0	4	17.07
3	278.354	0	5	0	0.00
4	84.143	0	1	0	0.00

assert X.shape == (1128, 5), "Use the five requested descriptors."
assert X.index.equals(y.index), "Keep feature and target rows aligned."
assert np.isfinite(X.to_numpy()).all() and np.isfinite(y).all(), "Inspect non-finite measurements."
print("Check passed: complete, aligned feature and target arrays.")

Check passed: complete, aligned feature and target arrays.

Task 2 · Split before preprocessing

Use an 80/20 development/test split with random_state=42. Split development again, reserving 25% for validation with the same seed. Define X_train, X_val, X_test, y_train, y_val, y_test. No fitting or model selection may use test labels before Task 8. This random split is for teaching; similar molecular structures may cross its boundary.

# Task 2: create disjoint training, validation and test subsets.
from sklearn.model_selection import train_test_split

X_dev, X_test, y_dev, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
X_train, X_val, y_train, y_val = train_test_split(
    X_dev, y_dev, test_size=0.25, random_state=42)
print("training / validation / test:", len(X_train), len(X_val), len(X_test))

training / validation / test: 676 226 226

index_sets = [set(frame.index) for frame in (X_train, X_val, X_test)]
assert not index_sets[0] & index_sets[1], "Training and validation overlap."
assert not index_sets[0] & index_sets[2], "Training and test overlap."
assert not index_sets[1] & index_sets[2], "Validation and test overlap."
assert set.union(*index_sets) == set(X.index), "Every row must appear once."
print("Check passed: disjoint and complete split.")

Check passed: disjoint and complete split.

Task 3 · Mean baseline and tuned ridge

Fit dummy on the training subset. Define ridge_search as a pipeline of StandardScaler and Ridge, searched over alpha values 0.01, 0.1, 1, 10 and 100. Use five shuffled training folds (KFold, seed 42), negative RMSE scoring and n_jobs=1. Scaling must be inside the searched pipeline so each fold learns its own preprocessing.

The refitted best ridge estimator uses the training subset, not validation. Compute ridge_val_rmse on validation data. Keep the test labels unopened.

# Task 3: fit a dummy baseline and tune ridge on training folds only.
from sklearn.dummy import DummyRegressor
from sklearn.linear_model import Ridge
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import GridSearchCV, KFold
from sklearn.metrics import root_mean_squared_error, mean_absolute_error, r2_score

dummy = DummyRegressor().fit(X_train, y_train)
cv = KFold(n_splits=5, shuffle=True, random_state=42)
ridge_search = GridSearchCV(make_pipeline(StandardScaler(), Ridge()),
    {"ridge__alpha": [0.01, 0.1, 1, 10, 100]}, cv=cv,
    scoring="neg_root_mean_squared_error", n_jobs=1)
ridge_search.fit(X_train, y_train)
ridge_val_rmse = root_mean_squared_error(y_val, ridge_search.predict(X_val))
print("ridge alpha:", ridge_search.best_params_, "validation RMSE:", ridge_val_rmse)

ridge alpha: {'ridge__alpha': 10} validation RMSE: 1.132674310664084

assert ridge_search.best_estimator_.named_steps["standardscaler"].n_samples_seen_ == len(X_train)
assert np.isfinite(ridge_val_rmse) and ridge_val_rmse > 0
print("Check passed: ridge refit uses training rows and validation RMSE is finite.")

Check passed: ridge refit uses training rows and validation RMSE is finite.

Task 4 · Tensors and preprocessing

Fit scaler on X_train. Make Xtr, Xva, ytr, yva tensors with torch.float32. Keep targets as a one-dimensional vector in the original logS units. Define loader with batch size 32, shuffling and a generator seeded with 42. The test features will be transformed only during final evaluation.

# Task 4: training-fitted scaling and matching tensors.
scaler = StandardScaler().fit(X_train)
Xtr = torch.tensor(scaler.transform(X_train), dtype=torch.float32)
Xva = torch.tensor(scaler.transform(X_val), dtype=torch.float32)
ytr = torch.tensor(y_train.to_numpy(), dtype=torch.float32)
yva = torch.tensor(y_val.to_numpy(), dtype=torch.float32)
loader = DataLoader(TensorDataset(Xtr, ytr), batch_size=32, shuffle=True,
                    generator=torch.Generator().manual_seed(42))
print("shapes:", tuple(Xtr.shape), tuple(ytr.shape), "batches:", len(loader))

shapes: (676, 5) (676,) batches: 22

assert Xtr.shape == (len(ytr), 5) and ytr.ndim == 1
assert Xtr.dtype == ytr.dtype == torch.float32
assert scaler.n_samples_seen_ == len(X_train), "Fit scaling on training rows only."
assert np.allclose(scaler.mean_, X_train.mean()), "Unexpected scaling statistics."
print("Check passed: shapes, types and scaler fitting boundary.")

Check passed: shapes, types and scaler fitting boundary.

Task 5 · Network and loop contracts

Write make_model(width, dropout) with a linear 5-to-width layer, ReLU, dropout and a width-to-1 unrestricted output. Write train_epoch(model, loader, optimiser) using MSE: train mode, clear gradients, forward pass, squeeze only axis 1, backward, step. Weight batch means by sample counts.

Write mse_on(model, inputs, targets) using eval mode and torch.no_grad(). Do not update anything during validation. The target is continuous logS, so do not add a sigmoid.

# Task 5a: construct a small regression network.
def make_model(width, dropout):
    return nn.Sequential(nn.Linear(5, width), nn.ReLU(), nn.Dropout(dropout), nn.Linear(width, 1))

print(make_model(width=16, dropout=0.0))

Sequential(
  (0): Linear(in_features=5, out_features=16, bias=True)
  (1): ReLU()
  (2): Dropout(p=0.0, inplace=False)
  (3): Linear(in_features=16, out_features=1, bias=True)
)

# Task 5b: one training epoch with a sample-weighted loss.
def train_epoch(model, loader, optimiser):
    model.train()
    total, count = 0.0, 0
    for inputs, targets in loader:
        optimiser.zero_grad(set_to_none=True)
        prediction = model(inputs).squeeze(1)
        assert prediction.shape == targets.shape
        loss = nn.functional.mse_loss(prediction, targets)
        loss.backward()
        optimiser.step()
        total += loss.item() * len(targets)
        count += len(targets)
    return total / count

print("Training function defined.")

Training function defined.

# Task 5c: evaluation without gradient tracking or parameter changes.
def mse_on(model, inputs, targets):
    model.eval()
    with torch.no_grad():
        prediction = model(inputs).squeeze(1)
        assert prediction.shape == targets.shape
        return nn.functional.mse_loss(prediction, targets).item()

print("Evaluation function defined.")

Evaluation function defined.

probe = make_model(16, 0.1)
assert probe(Xtr[:1]).shape == (1, 1), "Preserve a batch of one."
before = deepcopy(probe.state_dict())
value = mse_on(probe, Xva, yva)
assert all(torch.equal(before[k], probe.state_dict()[k]) for k in before)
assert not probe.training and np.isfinite(value)
print("Check passed: evaluation is finite and does not change model weights.")

Check passed: evaluation is finite and does not change model weights.

Task 6 · Early stopping

Use the supplied run_candidate setup and finish its companion fit_with_stopping function. Train for at most 140 epochs. After each epoch, record comparable evaluation-mode training and validation MSE. Save an independent copy of the weights whenever validation MSE improves by at least 0.0001. Stop after 15 non-improving epochs, restore the selected checkpoint and return (model, history, best_epoch).

A checkpoint dictionary that still refers to changing tensors is not a saved model. Use deepcopy. The helper resets the model seed and batch-order generator for each candidate and creates a fresh optimiser.

# Experiment setup: all candidates receive the same split, seed and maximum budget.
def run_candidate(width, dropout, decay):
    torch.manual_seed(42)
    model = make_model(width, dropout)
    optimiser = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=decay)
    fresh_loader = DataLoader(TensorDataset(Xtr, ytr), batch_size=32, shuffle=True,
                             generator=torch.Generator().manual_seed(42))
    return fit_with_stopping(model, fresh_loader, optimiser)

print("Candidate setup defined.")

Candidate setup defined.

# Task 6: preserve and restore the checkpoint chosen on validation MSE.
def fit_with_stopping(model, loader, optimiser):
    best_loss, best_epoch, stale = float("inf"), 0, 0
    history = []
    for epoch in range(1, 141):
        train_epoch(model, loader, optimiser)
        train_mse, val_mse = mse_on(model, Xtr, ytr), mse_on(model, Xva, yva)
        history.append((train_mse, val_mse))
        if val_mse < best_loss - 1e-4:
            best_loss, best_epoch, stale = val_mse, epoch, 0
            best_state = deepcopy(model.state_dict())
        else:
            stale += 1
        if stale >= 15:
            break
    model.load_state_dict(best_state)
    model.eval()
    return model, history, best_epoch

print("Early-stopping function defined.")

Early-stopping function defined.

Task 7 · Bounded selection on validation data

Fit exactly the two configurations below. For each, store its restored model, history, selected epoch and validation RMSE in runs. Select the MLP with the lowest validation RMSE as best_mlp_name and best_mlp.

Then create validation_results for ridge and this MLP and set selected_name to its lowest-RMSE row. This choice is frozen before Task 8. The dummy baseline provides context but is not part of this prespecified final family choice. If neither learned model beats it, the report must say so rather than claim useful learning.

# Task 7: select architecture and model family using validation only.
configs = {"MLP 16": (16, 0.0, 0.0), "MLP 32 regularised": (32, 0.1, 0.01)}
runs = {}
for name, settings in configs.items():
    fitted, curve, epoch = run_candidate(*settings)
    runs[name] = {"model": fitted, "history": curve, "epoch": epoch,
                  "validation RMSE": np.sqrt(mse_on(fitted, Xva, yva))}
best_mlp_name = min(runs, key=lambda name: runs[name]["validation RMSE"])
best_mlp = runs[best_mlp_name]["model"]
validation_results = pd.Series({"ridge": ridge_val_rmse,
    "MLP": runs[best_mlp_name]["validation RMSE"]}, name="validation RMSE")
selected_name = validation_results.idxmin()
print(pd.DataFrame({name: {k: v for k, v in run.items() if k in ["epoch", "validation RMSE"]}
                    for name, run in runs.items()}).T)
print(validation_results, "\nSelected before test:", selected_name)

                    epoch  validation RMSE
MLP 16               81.0         0.827936
MLP 32 regularised  128.0         0.840231
ridge    1.132674
MLP      0.827936
Name: validation RMSE, dtype: float64 
Selected before test: MLP

assert set(runs) == set(configs) and selected_name in {"ridge", "MLP"}
for run in runs.values():
    assert 1 <= run["epoch"] <= len(run["history"]) <= 140
    restored = mse_on(run["model"], Xva, yva)
    assert np.isclose(restored, run["history"][run["epoch"] - 1][1], rtol=1e-5)
print("Check passed: each selected checkpoint was restored.")

Check passed: each selected checkpoint was restored.

The selected validation score reflects architecture and checkpoint selection, so it is optimistic for the procedure. The upcoming test estimate addresses that remaining question. The validation ranking can reverse on test observations; that does not authorise a new model choice.

curve = np.sqrt(np.asarray(runs[best_mlp_name]["history"]))
fig, ax = plt.subplots()
ax.plot(np.arange(1, len(curve) + 1), curve[:, 0], label="training (eval mode)")
ax.plot(np.arange(1, len(curve) + 1), curve[:, 1], label="validation")
ax.axvline(runs[best_mlp_name]["epoch"], linestyle="--", label="restored checkpoint")
ax.set(xlabel="epoch", ylabel="RMSE (log10 mol/L)", title=best_mlp_name)
ax.legend()
plt.show()

Task 8 · Final evaluation, with the choice frozen

Define frozen_selection = selected_name before predictions. Transform test features with scaler, use the restored MLP in eval/no-grad mode, and collect predictions for dummy, ridge and the selected MLP. Create test_results with MAE, RMSE and R2. Set selected_prediction using frozen_selection, not by sorting the test scores.

This prespecified comparison is the final test evaluation. The models are retained from training-only fitting; we do not refit using validation rows. Error analysis after this point may suggest future experiments, but revising the model would require fresh independent evaluation.

# Task 8: evaluate once and preserve the validation-selected model.
frozen_selection = selected_name
Xte = torch.tensor(scaler.transform(X_test), dtype=torch.float32)
best_mlp.eval()
with torch.no_grad():
    mlp_test_prediction = best_mlp(Xte).squeeze(1).numpy()
predictions = {"dummy": dummy.predict(X_test), "ridge": ridge_search.predict(X_test),
               "MLP": mlp_test_prediction}
test_results = pd.DataFrame({name: {"MAE": mean_absolute_error(y_test, pred),
    "RMSE": root_mean_squared_error(y_test, pred), "R2": r2_score(y_test, pred)}
    for name, pred in predictions.items()}).T
selected_prediction = predictions[frozen_selection]
print("Model chosen before seeing test scores:", frozen_selection)
test_results.round(3)

Model chosen before seeing test scores: MLP

index	MAE	RMSE	R2
dummy	1.727	2.175	-0.001
ridge	0.910	1.221	0.685
MLP	0.725	0.986	0.794

assert frozen_selection == selected_name == validation_results.idxmin()
assert np.array_equal(selected_prediction, predictions[frozen_selection])
assert selected_prediction.shape == y_test.shape and np.isfinite(selected_prediction).all()
print("Check passed: the final prediction follows the validation choice.")

Check passed: the final prediction follows the validation choice.

fig, ax = plt.subplots()
ax.scatter(y_test, selected_prediction, alpha=0.6)
limits = [min(y_test.min(), selected_prediction.min()) - 0.5,
          max(y_test.max(), selected_prediction.max()) + 0.5]
ax.plot(limits, limits, linestyle="--", color="0.4", label="perfect prediction")
ax.set(xlabel="measured logS (log10 mol/L)", ylabel="predicted logS (log10 mol/L)",
       title=f"Validation-selected model: {frozen_selection}")
ax.legend()
plt.show()

Task 9 · Write the result, including its limits

Complete a paragraph using your actual results:

We used five fixed descriptors and a reproducible random train/validation/test split of … / … / … molecules. Ridge alpha was selected by five-fold training CV; two MLP candidates used validation early stopping with a 140-epoch maximum and patience 15. The validation-selected model was … . Its final test RMSE was … log10(mol/L), compared with … for the mean baseline and … for the other prespecified learned model. The result supports … under this random-split protocol, but it does not establish … . Our next independent evaluation would … .

Include the chosen MLP width, dropout, decay and epoch even if ridge was selected. Do not claim a neural-network advantage merely because its training loss is lower.

# Task 9: collect the settings needed to reproduce this experiment.
experiment = {"seed": 42, "features": features, "selected_family": frozen_selection,
    "ridge_parameters": ridge_search.best_params_, "mlp_candidate": best_mlp_name,
    "mlp_settings_width_dropout_decay": configs[best_mlp_name],
    "mlp_selected_epoch": runs[best_mlp_name]["epoch"], "learning_rate": 0.003,
    "batch_size": 32, "patience": 15, "maximum_epochs": 140,
    "torch": torch.__version__, "sklearn": sklearn.__version__, "device": "cpu"}
experiment

{'seed': 42,
 'features': ['Molecular Weight',
  'Number of H-Bond Donors',
  'Number of Rings',
  'Number of Rotatable Bonds',
  'Polar Surface Area'],
 'selected_family': 'MLP',
 'ridge_parameters': {'ridge__alpha': 10},
 'mlp_candidate': 'MLP 16',
 'mlp_settings_width_dropout_decay': (16, 0.0, 0.0),
 'mlp_selected_epoch': 81,
 'learning_rate': 0.003,
 'batch_size': 32,
 'patience': 15,
 'maximum_epochs': 140,
 'torch': '2.9.1+cpu',
 'sklearn': '1.9.1',
 'device': 'cpu'}

Reflection

1. Was the network’s improvement, if any, large enough to justify its extra tuning and training complexity? Distinguish what the observed test difference says from what a single split cannot establish.
2. Why can the best validation score be too optimistic? What would nested CV estimate if model selection and early stopping were repeated inside every outer training fold?
3. How would a scaffold-based split alter the scientific question? What leakage might a random molecular split permit even when every scaler is fitted correctly?
4. What would varying only the random initialisation seed measure? How does that differ from changing the data split? Why is neither standard deviation automatically a confidence interval for a new population?
5. Which objects besides the state dictionary are required to reproduce a prediction from a raw descriptor row? Explain why saving only the weights is incomplete.

Extension, on development data only: prespecify three seeds and repeat each candidate; compare mean validation RMSE and variability. Keep the test labels sealed until you have fixed how the final model is chosen. If you have already studied this test set, treat this as a new development exercise and obtain independent data for a fresh final claim.

PyTorch API references: saving checkpoints, AdamW, and reproducibility.