{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Guided exercise: photon energy regression from calorimeter images\n", "\n", "This notebook is a guided exercise inspired by the **Google School of Computing 2026 technical evaluation task** on calorimeter energy reconstruction.\n", "\n", "The goal is to study a simple regression problem:\n", "\n", "$$\n", "f_\\theta : \\mathbb{R}^{225} \\to \\mathbb{R},\n", "$$\n", "\n", "mapping a flattened **15 × 15 calorimeter image** to the **true photon energy**.\n", "\n", "We will:\n", "- inspect the dataset and the target distribution;\n", "- visualize calorimeter showers as 2D images;\n", "- define regression metrics;\n", "- build and compare three baselines:\n", " - a **sum-of-cells baseline**;\n", " - a **linear model**;\n", " - a small **MLP**;\n", "- compare their performance using both standard regression metrics and physics-motivated relative-error metrics.\n", "\n", "> The main question of this notebook is:\n", ">\n", "> **How much of this problem is already captured by a simple linear response, and how much requires non-linearity?**" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 0. Conceptual connection with the previous notebook\n", "\n", "In the classification notebook, we studied how models learn to **separate classes**.\n", "\n", "Here the task is different:\n", "- we are no longer predicting a discrete label;\n", "- we are learning a **continuous target**.\n", "\n", "So this is a regression problem.\n", "\n", "A useful conceptual bridge is the following:\n", "- in classification, the model output was a **score** used to rank events by class compatibility;\n", "- in regression, the model output is a **continuous prediction**, here the reconstructed photon energy.\n", "\n", "In both cases, the core ML question is the same:\n", "\n", "> **What structure in the input allows the model to make better predictions?**" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Problem setup\n", "\n", "We are given a file called `toy_photon_dataset.npz` containing:\n", "\n", "- `X`: array of shape `(N, 225)` \n", "- `y`: array of shape `(N,)`\n", "\n", "where:\n", "- each event is a flattened **15 × 15** calorimeter grid;\n", "- each input feature corresponds to the energy deposited in one calorimeter cell;\n", "- the target is the **true photon energy**.\n", "\n", "The dataset is simplified:\n", "- isolated single-photon events;\n", "- no pile-up or multi-particle overlaps;\n", "- stochastic fluctuations and per-cell noise are included.\n", "\n", "This makes it a clean setting to study the basics of **supervised regression**." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import random\n", "from pathlib import Path\n", "\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "\n", "import torch\n", "import torch.nn as nn\n", "from torch.utils.data import TensorDataset, DataLoader\n", "\n", "from sklearn.model_selection import train_test_split" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def set_seed(seed=42):\n", " random.seed(seed)\n", " np.random.seed(seed)\n", " torch.manual_seed(seed)\n", "\n", "set_seed(42)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Load the dataset\n", "\n", "Adjust the path below if needed.\n", "\n", "The notebook expects the file `toy_photon_dataset.npz` to be available either:\n", "- in the same directory as this notebook, or\n", "- one level above." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "candidate_paths = [\n", " Path(\"toy_photon_dataset.npz\"),\n", " Path(\"../toy_photon_dataset.npz\"),\n", "]\n", "\n", "dataset_path = None\n", "for path in candidate_paths:\n", " if path.exists():\n", " dataset_path = path\n", " break\n", "\n", "if dataset_path is None:\n", " raise FileNotFoundError(\n", " \"Could not find toy_photon_dataset.npz. Place it next to the notebook or in the parent directory.\"\n", " )\n", "\n", "data = np.load(dataset_path)\n", "X = data[\"X\"].astype(np.float32)\n", "y = data[\"y\"].astype(np.float32)\n", "\n", "print(f\"Loaded dataset from: {dataset_path}\")\n", "print(f\"X shape: {X.shape}\")\n", "print(f\"y shape: {y.shape}\")\n", "print(f\"Energy range: [{y.min():.2f}, {y.max():.2f}] GeV\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. First look at the data\n", "\n", "Before training any model, let us inspect:\n", "- the target distribution;\n", "- the total deposited energy;\n", "- the event-by-event ratio between total deposited energy and true energy;\n", "- a few example showers.\n", "\n", "This already helps us guess whether a simple baseline might work well." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "grid_size = 15\n", "X_grid = X.reshape(-1, grid_size, grid_size)\n", "\n", "E_sum = X.sum(axis=1)\n", "sampling_fraction = E_sum / y\n", "\n", "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5))\n", "\n", "axes[0].hist(y, bins=80, alpha=0.85)\n", "axes[0].set_xlabel(r\"$E_{\\mathrm{true}}$ [GeV]\")\n", "axes[0].set_ylabel(\"Events\")\n", "axes[0].set_title(\"True energy distribution\")\n", "\n", "axes[1].hist(E_sum, bins=80, alpha=0.85)\n", "axes[1].set_xlabel(r\"$E_{\\mathrm{sum}}$ [GeV]\")\n", "axes[1].set_ylabel(\"Events\")\n", "axes[1].set_title(\"Total deposited energy\")\n", "\n", "axes[2].hist(sampling_fraction, bins=80, alpha=0.85)\n", "axes[2].set_xlabel(r\"$E_{\\mathrm{sum}} / E_{\\mathrm{true}}$\")\n", "axes[2].set_ylabel(\"Events\")\n", "axes[2].set_title(\"Sampling fraction\")\n", "\n", "plt.tight_layout()\n", "plt.show()\n", "\n", "print(f\"Mean sampling fraction = {sampling_fraction.mean():.4f}\")\n", "print(f\"Std sampling fraction = {sampling_fraction.std():.4f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Discussion\n", "\n", "A useful first question is:\n", "\n", "- if the calorimeter response were perfectly linear and noiseless,\n", " would the **sum of all cell energies** already be enough to reconstruct the photon energy?\n", "\n", "This leads naturally to our first baseline:\n", "- use only the total deposited energy;\n", "- then ask whether a linear model or an MLP can improve upon it.\n", "\n", "That is a very instructive progression:\n", "1. **Physics-motivated baseline**: total deposited energy;\n", "2. **ML baseline**: linear regression;\n", "3. **More flexible model**: MLP." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "energy_ranges = [(5, 30), (30, 80), (80, 140), (140, 200)]\n", "\n", "fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n", "for ax, (lo, hi) in zip(axes, energy_ranges):\n", " mask = (y >= lo) & (y < hi)\n", " mean_image = X_grid[mask].mean(axis=0)\n", " im = ax.imshow(mean_image, interpolation=\"nearest\")\n", " ax.set_title(f\"{lo} ≤ E < {hi} GeV\")\n", " ax.set_xlabel(\"column\")\n", " ax.set_ylabel(\"row\")\n", " plt.colorbar(im, ax=ax, shrink=0.8)\n", "\n", "plt.suptitle(\"Average calorimeter shower by true energy range\", y=1.03)\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n", "for ax, (lo, hi) in zip(axes, energy_ranges):\n", " idx = np.where((y >= lo) & (y < hi))[0][0]\n", " im = ax.imshow(X_grid[idx], interpolation=\"nearest\")\n", " ax.set_title(f\"Single event\\nE = {y[idx]:.1f} GeV\")\n", " ax.set_xlabel(\"column\")\n", " ax.set_ylabel(\"row\")\n", " plt.colorbar(im, ax=ax, shrink=0.8)\n", "\n", "plt.suptitle(\"Single-event examples\", y=1.03)\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 1\n", "Before training any model, answer:\n", "\n", "1. Do you expect the total deposited energy $E_{\\mathrm{sum}}$ to be a strong baseline for this problem? Why?\n", "2. What kinds of effects could make $E_{\\mathrm{sum}}$ insufficient on an event-by-event basis?\n", "3. Looking at the average and single-event images, do you expect shower shape to contain useful extra information?" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4. Train/validation split\n", "\n", "We now split the dataset into:\n", "- **training set**: used to optimize model parameters;\n", "- **validation set**: used to evaluate performance on unseen events.\n", "\n", "We keep the setup simple and reproducible." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "X_train, X_val, y_train, y_val = train_test_split(\n", " X, y, test_size=0.20, random_state=42\n", ")\n", "\n", "print(f\"Training events: {len(y_train):,}\")\n", "print(f\"Validation events: {len(y_val):,}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. Regression metrics\n", "\n", "We will use both standard ML metrics and physics-motivated metrics.\n", "\n", "### Standard metrics\n", "- **MSE**: mean squared error\n", "- **RMSE**: root mean squared error\n", "- **MAE**: mean absolute error\n", "\n", "### Relative energy error\n", "For calorimeter reconstruction, it is often more informative to study\n", "\n", "$$\n", "\\frac{\\Delta E}{E} = \\frac{E_{\\mathrm{pred}} - E_{\\mathrm{true}}}{E_{\\mathrm{true}}}.\n", "$$\n", "\n", "From this quantity, we can report:\n", "- the **bias**: mean of $\\Delta E / E$;\n", "- the **resolution**: for example the standard deviation of $\\Delta E / E$;\n", "- the **p68**: 68th percentile of $|\\Delta E / E|$.\n", "\n", "In this notebook, relative-error metrics will be reported in **percent**." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def regression_metrics(y_true, y_pred):\n", " y_true = np.asarray(y_true)\n", " y_pred = np.asarray(y_pred)\n", "\n", " mse = np.mean((y_pred - y_true) ** 2)\n", " rmse = np.sqrt(mse)\n", " mae = np.mean(np.abs(y_pred - y_true))\n", "\n", " rel = (y_pred - y_true) / y_true\n", "\n", " metrics = {\n", " \"MSE\": mse,\n", " \"RMSE\": rmse,\n", " \"MAE\": mae,\n", " \"bias_percent\": 100.0 * np.mean(rel),\n", " \"resolution_std_percent\": 100.0 * np.std(rel),\n", " \"p68_abs_percent\": 100.0 * np.percentile(np.abs(rel), 68),\n", " }\n", " return metrics, rel" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 6. Baseline 0: sum-of-cells reconstruction\n", "\n", "Let us start with the most physics-motivated baseline:\n", "\n", "$$\n", "E_{\\mathrm{pred}} = \\alpha \\, E_{\\mathrm{sum}} + \\beta,\n", "$$\n", "\n", "where\n", "\n", "$$\n", "E_{\\mathrm{sum}} = \\sum_{\\mathrm{cells}} E_i.\n", "$$\n", "\n", "This is useful because it makes the physics question explicit:\n", "\n", "> Is the problem mostly about measuring the total deposited energy, or does the shower pattern carry additional corrective information?\n", "\n", "We will fit the calibration parameters $\\alpha$ and $\\beta$ using the training set." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "E_sum_train = X_train.sum(axis=1)\n", "E_sum_val = X_val.sum(axis=1)\n", "\n", "alpha, beta = np.polyfit(E_sum_train, y_train, deg=1)\n", "y_pred_sum = alpha * E_sum_val + beta\n", "\n", "print(f\"Calibration: E_pred = {alpha:.6f} * E_sum + {beta:.6f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 7. Baseline model A: linear regression in PyTorch\n", "\n", "Our first ML model is simply a linear map:\n", "\n", "$$\n", "\\hat{E} = W \\cdot X + b.\n", "$$\n", "\n", "This is slightly more flexible than the sum baseline:\n", "- instead of one global coefficient, each cell gets its own learned weight;\n", "- the output is a weighted sum over all cells plus a bias term.\n", "\n", "So this model can learn:\n", "- different effective weights for different cells;\n", "- simple spatial corrections;\n", "- but still **no non-linear correlations**." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class LinearRegressor(nn.Module):\n", " def __init__(self, input_dim=225):\n", " super().__init__()\n", " self.linear = nn.Linear(input_dim, 1)\n", "\n", " def forward(self, x):\n", " return self.linear(x).squeeze(-1)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 8. Baseline model B: a small MLP\n", "\n", "The second ML model is a simple fully-connected neural network.\n", "\n", "This introduces **non-linearity**, which may help because:\n", "- shower-shape information may matter;\n", "- the relation between cell pattern and true energy may not be perfectly linear;\n", "- local fluctuations might be better corrected by a flexible model." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class MLPRegressor(nn.Module):\n", " def __init__(self, input_dim=225, hidden_dims=(128, 64)):\n", " super().__init__()\n", " self.net = nn.Sequential(\n", " nn.Linear(input_dim, hidden_dims[0]),\n", " nn.ReLU(),\n", " nn.Linear(hidden_dims[0], hidden_dims[1]),\n", " nn.ReLU(),\n", " nn.Linear(hidden_dims[1], 1),\n", " )\n", "\n", " def forward(self, x):\n", " return self.net(x).squeeze(-1)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 9. Data loaders and training utilities\n", "\n", "We will use the same training loop for both ML models so that the comparison is fair." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "X_train_t = torch.from_numpy(X_train)\n", "X_val_t = torch.from_numpy(X_val)\n", "y_train_t = torch.from_numpy(y_train)\n", "y_val_t = torch.from_numpy(y_val)\n", "\n", "train_dataset = TensorDataset(X_train_t, y_train_t)\n", "val_dataset = TensorDataset(X_val_t, y_val_t)\n", "\n", "train_loader = DataLoader(train_dataset, batch_size=512, shuffle=True)\n", "val_loader = DataLoader(val_dataset, batch_size=2048, shuffle=False)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def train_model(\n", " model,\n", " train_loader,\n", " val_loader,\n", " epochs=30,\n", " lr=1e-3,\n", " device=None,\n", "):\n", " if device is None:\n", " device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", " model = model.to(device)\n", " criterion = nn.MSELoss()\n", " optimizer = torch.optim.Adam(model.parameters(), lr=lr)\n", "\n", " history = {\n", " \"train_loss\": [],\n", " \"val_loss\": [],\n", " }\n", "\n", " for epoch in range(epochs):\n", " model.train()\n", " train_loss_sum = 0.0\n", " train_count = 0\n", "\n", " for xb, yb in train_loader:\n", " xb = xb.to(device)\n", " yb = yb.to(device)\n", "\n", " optimizer.zero_grad()\n", " pred = model(xb)\n", " loss = criterion(pred, yb)\n", " loss.backward()\n", " optimizer.step()\n", "\n", " batch_size = xb.size(0)\n", " train_loss_sum += loss.item() * batch_size\n", " train_count += batch_size\n", "\n", " train_loss = train_loss_sum / train_count\n", "\n", " model.eval()\n", " val_loss_sum = 0.0\n", " val_count = 0\n", " with torch.no_grad():\n", " for xb, yb in val_loader:\n", " xb = xb.to(device)\n", " yb = yb.to(device)\n", " pred = model(xb)\n", " loss = criterion(pred, yb)\n", "\n", " batch_size = xb.size(0)\n", " val_loss_sum += loss.item() * batch_size\n", " val_count += batch_size\n", "\n", " val_loss = val_loss_sum / val_count\n", "\n", " history[\"train_loss\"].append(train_loss)\n", " history[\"val_loss\"].append(val_loss)\n", "\n", " if epoch in [0, 4, 9, 19, epochs - 1]:\n", " print(\n", " f\"Epoch {epoch + 1:2d}/{epochs} | \"\n", " f\"train loss = {train_loss:.4f} | \"\n", " f\"val loss = {val_loss:.4f}\"\n", " )\n", "\n", " return model, history" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def predict_model(model, X_array, device=None, batch_size=4096):\n", " if device is None:\n", " device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", " model.eval()\n", " model = model.to(device)\n", "\n", " X_tensor = torch.from_numpy(X_array.astype(np.float32))\n", " loader = DataLoader(TensorDataset(X_tensor), batch_size=batch_size, shuffle=False)\n", "\n", " preds = []\n", " with torch.no_grad():\n", " for (xb,) in loader:\n", " xb = xb.to(device)\n", " pred = model(xb).cpu().numpy()\n", " preds.append(pred)\n", "\n", " return np.concatenate(preds, axis=0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 10. Train the linear model" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "set_seed(42)\n", "linear_model = LinearRegressor(input_dim=225)\n", "\n", "linear_model, linear_history = train_model(\n", " linear_model,\n", " train_loader,\n", " val_loader,\n", " epochs=30,\n", " lr=1e-3,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 11. Train the MLP" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "set_seed(42)\n", "mlp_model = MLPRegressor(input_dim=225, hidden_dims=(128, 64))\n", "\n", "mlp_model, mlp_history = train_model(\n", " mlp_model,\n", " train_loader,\n", " val_loader,\n", " epochs=30,\n", " lr=1e-3,\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 12. Compare training curves\n", "\n", "A first question is whether the MLP learns faster or reaches a lower validation loss.\n", "\n", "This gives a first indication of:\n", "- underfitting;\n", "- overfitting;\n", "- optimization stability." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "epochs_linear = np.arange(1, len(linear_history[\"train_loss\"]) + 1)\n", "epochs_mlp = np.arange(1, len(mlp_history[\"train_loss\"]) + 1)\n", "\n", "plt.figure(figsize=(8, 5))\n", "plt.plot(epochs_linear, linear_history[\"train_loss\"], label=\"Linear - train\")\n", "plt.plot(epochs_linear, linear_history[\"val_loss\"], label=\"Linear - val\")\n", "plt.plot(epochs_mlp, mlp_history[\"train_loss\"], label=\"MLP - train\")\n", "plt.plot(epochs_mlp, mlp_history[\"val_loss\"], label=\"MLP - val\")\n", "plt.xlabel(\"Epoch\")\n", "plt.ylabel(\"MSE loss\")\n", "plt.title(\"Training and validation loss\")\n", "plt.legend()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2\n", "By looking at the loss curves:\n", "\n", "1. Which ML model reaches the lower validation loss?\n", "2. Do you see signs of underfitting or overfitting?\n", "3. Does the MLP seem to learn something beyond the linear model?" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 13. Predictions and numerical comparison" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "y_pred_linear = predict_model(linear_model, X_val)\n", "y_pred_mlp = predict_model(mlp_model, X_val)\n", "\n", "sum_metrics, rel_sum = regression_metrics(y_val, y_pred_sum)\n", "linear_metrics, rel_linear = regression_metrics(y_val, y_pred_linear)\n", "mlp_metrics, rel_mlp = regression_metrics(y_val, y_pred_mlp)\n", "\n", "all_metrics = {\n", " \"Sum baseline\": sum_metrics,\n", " \"Linear model\": linear_metrics,\n", " \"MLP\": mlp_metrics,\n", "}\n", "\n", "for model_name, metrics in all_metrics.items():\n", " print(model_name)\n", " for k, v in metrics.items():\n", " print(f\" {k:>24s} : {v:.4f}\")\n", " print()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "metric_names = [\"RMSE\", \"MAE\", \"bias_percent\", \"resolution_std_percent\", \"p68_abs_percent\"]\n", "sum_values = [sum_metrics[m] for m in metric_names]\n", "linear_values = [linear_metrics[m] for m in metric_names]\n", "mlp_values = [mlp_metrics[m] for m in metric_names]\n", "\n", "x = np.arange(len(metric_names))\n", "width = 0.26\n", "\n", "plt.figure(figsize=(11, 5))\n", "plt.bar(x - width, sum_values, width=width, label=\"Sum baseline\")\n", "plt.bar(x, linear_values, width=width, label=\"Linear\")\n", "plt.bar(x + width, mlp_values, width=width, label=\"MLP\")\n", "plt.xticks(x, metric_names, rotation=20)\n", "plt.ylabel(\"Metric value\")\n", "plt.title(\"Validation metrics comparison\")\n", "plt.legend()\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Discussion\n", "\n", "This comparison is especially informative because the three models probe different levels of complexity:\n", "\n", "- **Sum baseline**: only the total deposited energy matters;\n", "- **Linear model**: individual cells can have different weights;\n", "- **MLP**: the model can learn non-linear corrections based on the full shower pattern.\n", "\n", "So if the MLP improves only slightly over the linear model, that tells us something important:\n", "- much of the problem may already be close to linear.\n", "\n", "If instead the MLP improves clearly, that suggests:\n", "- shower-shape information contains useful non-linear corrections." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 14. Predicted vs true energy" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n", "\n", "for ax, pred, title in zip(\n", " axes,\n", " [y_pred_sum, y_pred_linear, y_pred_mlp],\n", " [\"Sum baseline\", \"Linear model\", \"MLP\"],\n", "):\n", " h = ax.hist2d(y_val, pred, bins=100, cmin=1)\n", " ax.plot([y_val.min(), y_val.max()], [y_val.min(), y_val.max()], \"--\", linewidth=1.5)\n", " ax.set_xlabel(r\"$E_{\\mathrm{true}}$ [GeV]\")\n", " ax.set_ylabel(r\"$E_{\\mathrm{pred}}$ [GeV]\")\n", " ax.set_title(title)\n", " plt.colorbar(h[3], ax=ax, label=\"Events\")\n", "\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 3\n", "By looking at the predicted-vs-true plots:\n", "\n", "1. Which model follows the diagonal more closely?\n", "2. Do you notice any energy range where the predictions degrade?\n", "3. Does the sum baseline already perform surprisingly well, or not?\n", "4. What visual differences do you see between the linear model and the MLP?" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 15. Relative error distributions\n", "\n", "Standard regression metrics summarize performance, but the distribution of the relative energy error is often more directly interpretable in calorimeter reconstruction." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axes = plt.subplots(1, 3, figsize=(18, 4.8), sharey=True)\n", "\n", "for ax, rel, title in zip(\n", " axes,\n", " [rel_sum, rel_linear, rel_mlp],\n", " [\"Sum baseline\", \"Linear model\", \"MLP\"],\n", "):\n", " bias = 100 * np.mean(rel)\n", " ax.hist(100 * rel, bins=100, alpha=0.85)\n", " ax.axvline(bias, linestyle=\"--\", linewidth=2, label=f\"bias = {bias:.2f}%\")\n", " ax.set_xlabel(r\"$100 \\times \\Delta E / E$ [%]\")\n", " ax.set_title(title)\n", " ax.legend()\n", "\n", "axes[0].set_ylabel(\"Events\")\n", "\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "plt.figure(figsize=(8, 5))\n", "plt.hist(100 * rel_sum, bins=100, alpha=0.45, label=\"Sum baseline\", density=True)\n", "plt.hist(100 * rel_linear, bins=100, alpha=0.45, label=\"Linear\", density=True)\n", "plt.hist(100 * rel_mlp, bins=100, alpha=0.45, label=\"MLP\", density=True)\n", "plt.xlabel(r\"$100 \\times \\Delta E / E$ [%]\")\n", "plt.ylabel(\"Density\")\n", "plt.title(\"Relative energy error comparison\")\n", "plt.legend()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 4\n", "By looking at the relative-error distributions:\n", "\n", "1. Which model has the narrowest core?\n", "2. Which model appears least biased?\n", "3. Why is $\\Delta E / E$ often more informative than looking only at MSE or MAE?\n", "4. Is the MLP learning a fundamentally new behavior, or mostly refining the simpler baselines?" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 16. Resolution as a function of true energy\n", "\n", "A useful diagnostic is to check whether the models behave similarly across the full energy range.\n", "\n", "We will compute, in bins of true energy:\n", "- the mean of $\\Delta E / E$;\n", "- the standard deviation of $\\Delta E / E$." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def binned_relative_stats(y_true, rel, bins):\n", " centers = 0.5 * (bins[:-1] + bins[1:])\n", " means = np.full(len(centers), np.nan)\n", " stds = np.full(len(centers), np.nan)\n", "\n", " for i in range(len(centers)):\n", " mask = (y_true >= bins[i]) & (y_true < bins[i + 1])\n", " if mask.sum() > 10:\n", " means[i] = np.mean(rel[mask])\n", " stds[i] = np.std(rel[mask])\n", "\n", " return centers, means, stds\n", "\n", "energy_bins = np.linspace(y_val.min(), y_val.max(), 16)\n", "\n", "centers, mean_rel_sum, std_rel_sum = binned_relative_stats(y_val, rel_sum, energy_bins)\n", "_, mean_rel_linear, std_rel_linear = binned_relative_stats(y_val, rel_linear, energy_bins)\n", "_, mean_rel_mlp, std_rel_mlp = binned_relative_stats(y_val, rel_mlp, energy_bins)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig, axes = plt.subplots(1, 2, figsize=(13, 4.8), sharex=True)\n", "\n", "axes[0].plot(centers, 100 * mean_rel_sum, marker=\"o\", label=\"Sum baseline\")\n", "axes[0].plot(centers, 100 * mean_rel_linear, marker=\"o\", label=\"Linear\")\n", "axes[0].plot(centers, 100 * mean_rel_mlp, marker=\"o\", label=\"MLP\")\n", "axes[0].axhline(0.0, linestyle=\"--\", linewidth=1)\n", "axes[0].set_xlabel(r\"$E_{\\mathrm{true}}$ [GeV]\")\n", "axes[0].set_ylabel(r\"Mean($\\Delta E / E$) [%]\")\n", "axes[0].set_title(\"Bias vs true energy\")\n", "axes[0].legend()\n", "\n", "axes[1].plot(centers, 100 * std_rel_sum, marker=\"o\", label=\"Sum baseline\")\n", "axes[1].plot(centers, 100 * std_rel_linear, marker=\"o\", label=\"Linear\")\n", "axes[1].plot(centers, 100 * std_rel_mlp, marker=\"o\", label=\"MLP\")\n", "axes[1].set_xlabel(r\"$E_{\\mathrm{true}}$ [GeV]\")\n", "axes[1].set_ylabel(r\"Std($\\Delta E / E$) [%]\")\n", "axes[1].set_title(\"Resolution vs true energy\")\n", "axes[1].legend()\n", "\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 5\n", "From the binned relative-error plots:\n", "\n", "1. Is the bias stable across the full energy range?\n", "2. Which model gives the best resolution as a function of energy?\n", "3. Does one model improve mostly at low energy, high energy, or everywhere?\n", "4. If the model mostly learns a correction to the total deposited energy, what does that suggest about the detector response?" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 17. Model size\n", "\n", "Even though this notebook is not focused on timing, it is still useful to compare the size of the ML models.\n", "\n", "This helps build intuition about model complexity." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def count_parameters(model):\n", " return sum(p.numel() for p in model.parameters() if p.requires_grad)\n", "\n", "print(f\"Linear model parameters: {count_parameters(linear_model):,}\")\n", "print(f\"MLP parameters: {count_parameters(mlp_model):,}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 18. Final discussion\n", "\n", "We now have a clean baseline hierarchy:\n", "\n", "1. **Sum baseline** \n", " A simple calibrated estimator based only on total deposited energy.\n", "\n", "2. **Linear model** \n", " A model that can weight individual cells differently, but remains linear.\n", "\n", "3. **MLP** \n", " A non-linear model that can learn corrections based on the full shower pattern.\n", "\n", "This is a useful ML lesson:\n", "\n", "- if the target depends mostly on a simple global quantity, a very simple model may already work well;\n", "- if spatial structure and fluctuations contain extra information, a more flexible model may improve the reconstruction.\n", "\n", "This is also a useful detector-physics lesson:\n", "\n", "- if the sum baseline is already very strong, the calorimeter response is close to a simple calibrated energy measurement;\n", "- if the MLP improves clearly, then the shower image carries additional corrective information beyond the total deposited energy." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 6\n", "Answer:\n", "\n", "1. Which model would you choose as the best baseline, and why?\n", "2. If the MLP performs better, what kind of extra information do you think it is learning?\n", "3. If the linear model is already very competitive, what does that tell you about the structure of the problem?\n", "4. How does this regression task differ conceptually from the classification problem in the previous notebook?\n", "5. What would be your next model or analysis improvement?" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Challenges\n", "\n", "### Challenge A\n", "Use only the scalar feature $E_{\\mathrm{sum}}$ as input to a one-dimensional linear regression model and compare it explicitly with the calibrated sum baseline.\n", "\n", "### Challenge B\n", "Normalize the inputs and retrain the linear model and the MLP. Does this change the optimization or final performance?\n", "\n", "### Challenge C\n", "Try a smaller and a larger MLP. How does the performance change with model capacity?\n", "\n", "### Challenge D\n", "Study the distribution of the residuals\n", "$$\n", "E_{\\mathrm{pred}} - E_{\\mathrm{true}}\n", "$$\n", "instead of $\\Delta E / E$. What information is easier to see in one variable than in the other?\n", "\n", "### Challenge E\n", "Estimate uncertainties on the resolution metrics using bootstrap resampling.\n", "\n", "### Challenge F\n", "Extra datasets are provided with additional pile-up activity.\n", "\n", "1. Compare the sum baseline, linear model, and MLP on the clean and pile-up datasets.\n", "2. Which model degrades the most when pile-up is introduced?\n", "3. Does the relative advantage of the MLP increase in the presence of pile-up?\n", "4. What does this suggest about the role of shower-shape information in noisy environments?\n", "\n", "### Challenge G\n", "Try a CNN using the 15 × 15 image directly. Does locality help?" ] }, { "cell_type": "markdown", "id": "73b0af60", "metadata": {}, "source": [] }, { "cell_type": "markdown", "id": "a072f211", "metadata": {}, "source": [] } ], "metadata": { "kernelspec": { "display_name": "rare", "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.12" } }, "nbformat": 4, "nbformat_minor": 5 }