{
 "cells": [
  {
   "cell_type": "raw",
   "id": "5b41bf46",
   "metadata": {},
   "source": [
    "---\n",
    "title: \"Lab 4 solutions: a neural network for solubility\"\n",
    "chapter: 4\n",
    "lab-solutions: true\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33765b2d",
   "metadata": {},
   "source": [
    "# Lab 4: a neural network for solubility\n",
    "\n",
    "**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.\n",
    "\n",
    "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.\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42740306",
   "metadata": {
    "tags": [
     "setup"
    ]
   },
   "outputs": [],
   "source": [
    "# Setup: locally use the course environment; Colab installs missing packages only.\n",
    "import importlib.util\n",
    "import subprocess\n",
    "import sys\n",
    "try:\n",
    "    import google.colab\n",
    "    IN_COLAB = True\n",
    "except ImportError:\n",
    "    IN_COLAB = False\n",
    "packages = {\"numpy\": \"numpy\", \"pandas\": \"pandas\", \"matplotlib\": \"matplotlib\",\n",
    "            \"sklearn\": \"scikit-learn\", \"torch\": \"torch==2.9.1\"}\n",
    "missing = [package for module, package in packages.items()\n",
    "           if importlib.util.find_spec(module) is None]\n",
    "if missing and IN_COLAB:\n",
    "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", *missing])\n",
    "elif missing:\n",
    "    raise RuntimeError(\"Install the course requirements first: \" + \", \".join(missing))\n",
    "print(\"Python:\", sys.version.split()[0], \"Colab:\", IN_COLAB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e280cf66",
   "metadata": {
    "tags": [
     "setup"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import sklearn\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from copy import deepcopy\n",
    "\n",
    "torch.set_num_threads(1)\n",
    "torch.manual_seed(42)\n",
    "print(\"numpy:\", np.__version__, \"sklearn:\", sklearn.__version__, \"torch:\", torch.__version__)\n",
    "print(\"Device: CPU; seed: 42\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e732a459",
   "metadata": {},
   "source": [
    "## Data and provenance\n",
    "\n",
    "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.\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "451f01f1",
   "metadata": {
    "tags": [
     "setup"
    ]
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from urllib.request import urlopen\n",
    "import hashlib\n",
    "import io\n",
    "\n",
    "expected_sha = \"8c06a76f0c6487d29ab0f903e6a7a7139f189ab3c1178f159c8be8964602f189\"\n",
    "local_paths = [p / \"data\" / \"esol.csv\" for p in (Path.cwd(), *Path.cwd().parents)]\n",
    "urls = [\"https://raw.githubusercontent.com/Awguhst/MolModels/main/data/esol.csv\",\n",
    "        \"https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/delaney-processed.csv\"]\n",
    "sources = [p for p in local_paths if p.is_file()] + urls\n",
    "print(\"ESOL sources prepared.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c16b5b",
   "metadata": {
    "tags": [
     "setup"
    ]
   },
   "outputs": [],
   "source": [
    "for source in sources:\n",
    "    try:\n",
    "        if isinstance(source, Path):\n",
    "            payload = source.read_bytes()\n",
    "        else:\n",
    "            with urlopen(source, timeout=30) as response:\n",
    "                payload = response.read()\n",
    "        if hashlib.sha256(payload).hexdigest() != expected_sha:\n",
    "            continue\n",
    "        esol = pd.read_csv(io.BytesIO(payload))\n",
    "        break\n",
    "    except OSError:\n",
    "        continue\n",
    "else:\n",
    "    raise RuntimeError(\"No checksum-verified ESOL source was available; retry the data download.\")\n",
    "print(\"Verified ESOL:\", esol.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ec09027",
   "metadata": {},
   "source": [
    "## Task 1 · Features and target\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef50ab59",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 1: define features, X and measured y in a fixed order.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c92b486b",
   "metadata": {
    "tags": [
     "check"
    ]
   },
   "outputs": [],
   "source": [
    "assert X.shape == (1128, 5), \"Use the five requested descriptors.\"\n",
    "assert X.index.equals(y.index), \"Keep feature and target rows aligned.\"\n",
    "assert np.isfinite(X.to_numpy()).all() and np.isfinite(y).all(), \"Inspect non-finite measurements.\"\n",
    "print(\"Check passed: complete, aligned feature and target arrays.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfae3fcb",
   "metadata": {},
   "source": [
    "## Task 2 · Split before preprocessing\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1cba335",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 2: create disjoint training, validation and test subsets.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18154cf0",
   "metadata": {
    "tags": [
     "check"
    ]
   },
   "outputs": [],
   "source": [
    "index_sets = [set(frame.index) for frame in (X_train, X_val, X_test)]\n",
    "assert not index_sets[0] & index_sets[1], \"Training and validation overlap.\"\n",
    "assert not index_sets[0] & index_sets[2], \"Training and test overlap.\"\n",
    "assert not index_sets[1] & index_sets[2], \"Validation and test overlap.\"\n",
    "assert set.union(*index_sets) == set(X.index), \"Every row must appear once.\"\n",
    "print(\"Check passed: disjoint and complete split.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1341be90",
   "metadata": {},
   "source": [
    "## Task 3 · Mean baseline and tuned ridge\n",
    "\n",
    "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.\n",
    "\n",
    "The refitted best ridge estimator uses the training subset, not validation. Compute `ridge_val_rmse` on validation data. Keep the test labels unopened."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8bae32",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 3: fit a dummy baseline and tune ridge on training folds only.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bee39477",
   "metadata": {
    "tags": [
     "check"
    ]
   },
   "outputs": [],
   "source": [
    "assert ridge_search.best_estimator_.named_steps[\"standardscaler\"].n_samples_seen_ == len(X_train)\n",
    "assert np.isfinite(ridge_val_rmse) and ridge_val_rmse > 0\n",
    "print(\"Check passed: ridge refit uses training rows and validation RMSE is finite.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76fc4667",
   "metadata": {},
   "source": [
    "## Task 4 · Tensors and preprocessing\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c9223b6",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 4: training-fitted scaling and matching tensors.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a20d4dcf",
   "metadata": {
    "tags": [
     "check"
    ]
   },
   "outputs": [],
   "source": [
    "assert Xtr.shape == (len(ytr), 5) and ytr.ndim == 1\n",
    "assert Xtr.dtype == ytr.dtype == torch.float32\n",
    "assert scaler.n_samples_seen_ == len(X_train), \"Fit scaling on training rows only.\"\n",
    "assert np.allclose(scaler.mean_, X_train.mean()), \"Unexpected scaling statistics.\"\n",
    "print(\"Check passed: shapes, types and scaler fitting boundary.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "126ea657",
   "metadata": {},
   "source": [
    "## Task 5 · Network and loop contracts\n",
    "\n",
    "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.\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74c39911",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 5a: construct a small regression network.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "360b57cc",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 5b: one training epoch with a sample-weighted loss.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f850f43",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 5c: evaluation without gradient tracking or parameter changes.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0a3d379",
   "metadata": {
    "tags": [
     "check"
    ]
   },
   "outputs": [],
   "source": [
    "probe = make_model(16, 0.1)\n",
    "assert probe(Xtr[:1]).shape == (1, 1), \"Preserve a batch of one.\"\n",
    "before = deepcopy(probe.state_dict())\n",
    "value = mse_on(probe, Xva, yva)\n",
    "assert all(torch.equal(before[k], probe.state_dict()[k]) for k in before)\n",
    "assert not probe.training and np.isfinite(value)\n",
    "print(\"Check passed: evaluation is finite and does not change model weights.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91fc7159",
   "metadata": {},
   "source": [
    "## Task 6 · Early stopping\n",
    "\n",
    "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)`.\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0cfeab6",
   "metadata": {
    "tags": [
     "setup"
    ]
   },
   "outputs": [],
   "source": [
    "# Experiment setup: all candidates receive the same split, seed and maximum budget.\n",
    "def run_candidate(width, dropout, decay):\n",
    "    torch.manual_seed(42)\n",
    "    model = make_model(width, dropout)\n",
    "    optimiser = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=decay)\n",
    "    fresh_loader = DataLoader(TensorDataset(Xtr, ytr), batch_size=32, shuffle=True,\n",
    "                             generator=torch.Generator().manual_seed(42))\n",
    "    return fit_with_stopping(model, fresh_loader, optimiser)\n",
    "\n",
    "print(\"Candidate setup defined.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9adcc56",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 6: preserve and restore the checkpoint chosen on validation MSE.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2de5330",
   "metadata": {},
   "source": [
    "## Task 7 · Bounded selection on validation data\n",
    "\n",
    "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`.\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f18409bf",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 7: select architecture and model family using validation only.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82541807",
   "metadata": {
    "tags": [
     "check"
    ]
   },
   "outputs": [],
   "source": [
    "assert set(runs) == set(configs) and selected_name in {\"ridge\", \"MLP\"}\n",
    "for run in runs.values():\n",
    "    assert 1 <= run[\"epoch\"] <= len(run[\"history\"]) <= 140\n",
    "    restored = mse_on(run[\"model\"], Xva, yva)\n",
    "    assert np.isclose(restored, run[\"history\"][run[\"epoch\"] - 1][1], rtol=1e-5)\n",
    "print(\"Check passed: each selected checkpoint was restored.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a580cb02",
   "metadata": {},
   "source": [
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45177a55",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#| fig-alt: \"Training and validation RMSE of the selected MLP, with a line at its restored checkpoint.\"\n",
    "curve = np.sqrt(np.asarray(runs[best_mlp_name][\"history\"]))\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(np.arange(1, len(curve) + 1), curve[:, 0], label=\"training (eval mode)\")\n",
    "ax.plot(np.arange(1, len(curve) + 1), curve[:, 1], label=\"validation\")\n",
    "ax.axvline(runs[best_mlp_name][\"epoch\"], linestyle=\"--\", label=\"restored checkpoint\")\n",
    "ax.set(xlabel=\"epoch\", ylabel=\"RMSE (log10 mol/L)\", title=best_mlp_name)\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb95dfbe",
   "metadata": {},
   "source": [
    "## Task 8 · Final evaluation, with the choice frozen\n",
    "\n",
    "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.\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6d9b5d5",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 8: evaluate once and preserve the validation-selected model.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1191e2e",
   "metadata": {
    "tags": [
     "check"
    ]
   },
   "outputs": [],
   "source": [
    "assert frozen_selection == selected_name == validation_results.idxmin()\n",
    "assert np.array_equal(selected_prediction, predictions[frozen_selection])\n",
    "assert selected_prediction.shape == y_test.shape and np.isfinite(selected_prediction).all()\n",
    "print(\"Check passed: the final prediction follows the validation choice.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab5d06d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#| fig-alt: \"Parity plot of test measured and predicted logS for the model selected before test evaluation.\"\n",
    "fig, ax = plt.subplots()\n",
    "ax.scatter(y_test, selected_prediction, alpha=0.6)\n",
    "limits = [min(y_test.min(), selected_prediction.min()) - 0.5,\n",
    "          max(y_test.max(), selected_prediction.max()) + 0.5]\n",
    "ax.plot(limits, limits, linestyle=\"--\", color=\"0.4\", label=\"perfect prediction\")\n",
    "ax.set(xlabel=\"measured logS (log10 mol/L)\", ylabel=\"predicted logS (log10 mol/L)\",\n",
    "       title=f\"Validation-selected model: {frozen_selection}\")\n",
    "ax.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "782e2788",
   "metadata": {},
   "source": [
    "## Task 9 · Write the result, including its limits\n",
    "\n",
    "Complete a paragraph using your actual results:\n",
    "\n",
    "> 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 ... .\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ba2e39d",
   "metadata": {
    "tags": [
     "solution"
    ]
   },
   "outputs": [],
   "source": [
    "# Task 9: collect the settings needed to reproduce this experiment.\n",
    "# TODO: your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfd1be0b",
   "metadata": {},
   "source": [
    "## Reflection\n",
    "\n",
    "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.\n",
    "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?\n",
    "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?\n",
    "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?\n",
    "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.\n",
    "\n",
    "**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.\n",
    "\n",
    "PyTorch API references: [saving checkpoints](https://docs.pytorch.org/tutorials/beginner/saving_loading_models.html), [AdamW](https://docs.pytorch.org/docs/2.9/generated/torch.optim.AdamW.html), and [reproducibility](https://docs.pytorch.org/docs/2.9/notes/randomness.html)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}