{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "xbsmj1hcj1g",
   "metadata": {},
   "source": [
    "# Building a Neural Network: $C_p(T)$ for Nitrogen\n",
    "\n",
    "**CHEG 667-013 — LLMs for Engineers**\n",
    "\n",
    "In this notebook we fit the heat capacity of N₂ gas using three approaches:\n",
    "1. A polynomial fit (the classical approach)\n",
    "2. A neural network built from scratch in numpy\n",
    "3. The same network in PyTorch\n",
    "\n",
    "This makes the ML concepts behind LLMs — weights, loss, gradient descent, overfitting — concrete and tangible."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "szrl41l3xbq",
   "metadata": {},
   "source": [
    "## 1. Load and plot the data\n",
    "\n",
    "The data is from the [NIST Chemistry WebBook](https://webbook.nist.gov/): isobaric heat capacity of N₂ at 1 bar, 300–2000 K."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "t4lqkcoeyil",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "data = np.loadtxt(\"data/n2_cp.csv\", delimiter=\",\", skiprows=1)\n",
    "T_raw = data[:, 0]   # Temperature (K)\n",
    "Cp_raw = data[:, 1]  # Cp (kJ/kg/K)\n",
    "\n",
    "plt.figure(figsize=(8, 5))\n",
    "plt.plot(T_raw, Cp_raw, 'ko', markersize=6)\n",
    "plt.xlabel('Temperature (K)')\n",
    "plt.ylabel('$C_p$ (kJ/kg/K)')\n",
    "plt.title('$C_p(T)$ for N$_2$ at 1 bar — NIST WebBook')\n",
    "plt.show()\n",
    "\n",
    "print(f\"{len(T_raw)} data points, T range: {T_raw.min():.0f} – {T_raw.max():.0f} K\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1jyrgsvp7op",
   "metadata": {},
   "source": [
    "## 2. Polynomial fit (baseline)\n",
    "\n",
    "Textbooks fit $C_p(T)$ with a polynomial: $C_p = a + bT + cT^2 + dT^3$. This is a **4-parameter** model. Let's fit it with `numpy.polyfit` and see how well it does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4smvu4z2oro",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit a cubic polynomial\n",
    "coeffs = np.polyfit(T_raw, Cp_raw, 3)\n",
    "poly = np.poly1d(coeffs)\n",
    "\n",
    "T_fine = np.linspace(T_raw.min(), T_raw.max(), 200)\n",
    "Cp_poly = poly(T_fine)\n",
    "\n",
    "# Compute residuals\n",
    "Cp_poly_at_data = poly(T_raw)\n",
    "mse_poly = np.mean((Cp_poly_at_data - Cp_raw) ** 2)\n",
    "\n",
    "plt.figure(figsize=(8, 5))\n",
    "plt.plot(T_raw, Cp_raw, 'ko', markersize=6, label='NIST data')\n",
    "plt.plot(T_fine, Cp_poly, 'b-', linewidth=2, label=f'Cubic polynomial (4 params)')\n",
    "plt.xlabel('Temperature (K)')\n",
    "plt.ylabel('$C_p$ (kJ/kg/K)')\n",
    "plt.title('Polynomial fit')\n",
    "plt.legend()\n",
    "plt.show()\n",
    "\n",
    "print(f\"Polynomial coefficients: {coeffs}\")\n",
    "print(f\"MSE: {mse_poly:.8f}\")\n",
    "print(f\"Parameters: 4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97y7mrcekji",
   "metadata": {},
   "source": [
    "## 3. Neural network from scratch (numpy)\n",
    "\n",
    "Now let's build a one-hidden-layer neural network. The architecture:\n",
    "\n",
    "```\n",
    "Input (1: T) -> Hidden (10 neurons, tanh) -> Output (1: Cp)\n",
    "```\n",
    "\n",
    "We need to:\n",
    "1. **Normalize** the data to [0, 1] so the network trains efficiently\n",
    "2. **Forward pass**: compute predictions from input through each layer\n",
    "3. **Loss**: mean squared error between predictions and data\n",
    "4. **Backpropagation**: compute gradients of the loss w.r.t. each weight using the chain rule\n",
    "5. **Gradient descent**: update weights in the direction that reduces the loss\n",
    "\n",
    "This is exactly what nanoGPT's `train.py` does — just at a much larger scale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "365o7bqbwkr",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalize inputs and outputs to [0, 1]\n",
    "T_min, T_max = T_raw.min(), T_raw.max()\n",
    "Cp_min, Cp_max = Cp_raw.min(), Cp_raw.max()\n",
    "\n",
    "T = (T_raw - T_min) / (T_max - T_min)\n",
    "Cp = (Cp_raw - Cp_min) / (Cp_max - Cp_min)\n",
    "\n",
    "X = T.reshape(-1, 1)    # (N, 1) input matrix\n",
    "Y = Cp.reshape(-1, 1)   # (N, 1) target matrix\n",
    "N = X.shape[0]\n",
    "\n",
    "# Network setup\n",
    "H = 10  # hidden neurons\n",
    "\n",
    "np.random.seed(42)\n",
    "W1 = np.random.randn(1, H) * 0.5   # input -> hidden weights\n",
    "b1 = np.zeros((1, H))               # hidden biases\n",
    "W2 = np.random.randn(H, 1) * 0.5   # hidden -> output weights\n",
    "b2 = np.zeros((1, 1))               # output bias\n",
    "\n",
    "print(f\"Parameters: W1({W1.shape}) + b1({b1.shape}) + W2({W2.shape}) + b2({b2.shape})\")\n",
    "print(f\"Total: {W1.size + b1.size + W2.size + b2.size} parameters for {N} data points\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5w1ezs9t2w6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop\n",
    "learning_rate = 0.01\n",
    "epochs = 5000\n",
    "log_interval = 500\n",
    "losses_np = []\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    # Forward pass\n",
    "    Z1 = X @ W1 + b1           # hidden pre-activation  (N, H)\n",
    "    A1 = np.tanh(Z1)           # hidden activation       (N, H)\n",
    "    Y_pred = A1 @ W2 + b2      # output                  (N, 1)\n",
    "\n",
    "    # Loss (mean squared error)\n",
    "    error = Y_pred - Y\n",
    "    loss = np.mean(error ** 2)\n",
    "    losses_np.append(loss)\n",
    "\n",
    "    # Backpropagation (chain rule, working backward)\n",
    "    dL_dYpred = 2 * error / N\n",
    "    dL_dW2 = A1.T @ dL_dYpred\n",
    "    dL_db2 = np.sum(dL_dYpred, axis=0, keepdims=True)\n",
    "    dL_dA1 = dL_dYpred @ W2.T\n",
    "    dL_dZ1 = dL_dA1 * (1 - A1 ** 2)   # tanh derivative\n",
    "    dL_dW1 = X.T @ dL_dZ1\n",
    "    dL_db1 = np.sum(dL_dZ1, axis=0, keepdims=True)\n",
    "\n",
    "    # Gradient descent update\n",
    "    W2 -= learning_rate * dL_dW2\n",
    "    b2 -= learning_rate * dL_db2\n",
    "    W1 -= learning_rate * dL_dW1\n",
    "    b1 -= learning_rate * dL_db1\n",
    "\n",
    "    if epoch % log_interval == 0 or epoch == epochs - 1:\n",
    "        print(f\"Epoch {epoch:5d}  Loss: {loss:.6f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "onel9r0kjk",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict on a fine grid and convert back to physical units\n",
    "T_fine_norm = np.linspace(0, 1, 200).reshape(-1, 1)\n",
    "A1_fine = np.tanh(T_fine_norm @ W1 + b1)\n",
    "Cp_nn_norm = A1_fine @ W2 + b2\n",
    "Cp_nn = Cp_nn_norm * (Cp_max - Cp_min) + Cp_min\n",
    "T_fine_K = T_fine_norm * (T_max - T_min) + T_min\n",
    "\n",
    "# MSE in original units for comparison with polynomial\n",
    "Cp_nn_at_data = np.tanh(X @ W1 + b1) @ W2 + b2\n",
    "Cp_nn_at_data = Cp_nn_at_data * (Cp_max - Cp_min) + Cp_min\n",
    "mse_nn = np.mean((Cp_nn_at_data.flatten() - Cp_raw) ** 2)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))\n",
    "\n",
    "ax1.plot(T_raw, Cp_raw, 'ko', markersize=6, label='NIST data')\n",
    "ax1.plot(T_fine, Cp_poly, 'b-', linewidth=2, label=f'Polynomial (4 params, MSE={mse_poly:.2e})')\n",
    "ax1.plot(T_fine_K.flatten(), Cp_nn.flatten(), 'r-', linewidth=2, label=f'NN numpy (31 params, MSE={mse_nn:.2e})')\n",
    "ax1.set_xlabel('Temperature (K)')\n",
    "ax1.set_ylabel('$C_p$ (kJ/kg/K)')\n",
    "ax1.set_title('$C_p(T)$ for N$_2$ at 1 bar')\n",
    "ax1.legend()\n",
    "\n",
    "ax2.semilogy(losses_np)\n",
    "ax2.set_xlabel('Epoch')\n",
    "ax2.set_ylabel('MSE (normalized)')\n",
    "ax2.set_title('Training loss — numpy NN')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea9z35qm9u8",
   "metadata": {},
   "source": [
    "## 4. Neural network in PyTorch\n",
    "\n",
    "The same network, but PyTorch handles backpropagation automatically. Compare the training loop above to the one below — `loss.backward()` replaces all of our manual gradient calculations.\n",
    "\n",
    "This is the same API used in nanoGPT's `model.py` — `nn.Linear`, activation functions, `optimizer.step()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3qxnrtyxqgz",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Prepare data as PyTorch tensors\n",
    "X_t = torch.tensor((T_raw - T_min) / (T_max - T_min), dtype=torch.float32).reshape(-1, 1)\n",
    "Y_t = torch.tensor((Cp_raw - Cp_min) / (Cp_max - Cp_min), dtype=torch.float32).reshape(-1, 1)\n",
    "\n",
    "# Define the network\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(1, H),    # input -> hidden (W1, b1)\n",
    "    nn.Tanh(),           # activation\n",
    "    nn.Linear(H, 1),    # hidden -> output (W2, b2)\n",
    ")\n",
    "\n",
    "print(model)\n",
    "print(f\"Total parameters: {sum(p.numel() for p in model.parameters())}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ydl3ycnypps",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "loss_fn = nn.MSELoss()\n",
    "losses_torch = []\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    Y_pred_t = model(X_t)\n",
    "    loss = loss_fn(Y_pred_t, Y_t)\n",
    "    losses_torch.append(loss.item())\n",
    "\n",
    "    optimizer.zero_grad()   # reset gradients\n",
    "    loss.backward()         # automatic differentiation\n",
    "    optimizer.step()        # update weights\n",
    "\n",
    "    if epoch % log_interval == 0 or epoch == epochs - 1:\n",
    "        print(f\"Epoch {epoch:5d}  Loss: {loss.item():.6f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bg0kvnk4ho",
   "metadata": {},
   "source": [
    "## 5. Compare all three approaches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "h2dfstoh8gd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# PyTorch predictions\n",
    "T_fine_t = torch.linspace(0, 1, 200).reshape(-1, 1)\n",
    "with torch.no_grad():\n",
    "    Cp_torch_norm = model(T_fine_t)\n",
    "Cp_torch = Cp_torch_norm.numpy() * (Cp_max - Cp_min) + Cp_min\n",
    "\n",
    "# MSE for PyTorch model\n",
    "with torch.no_grad():\n",
    "    Cp_torch_at_data = model(X_t).numpy() * (Cp_max - Cp_min) + Cp_min\n",
    "mse_torch = np.mean((Cp_torch_at_data.flatten() - Cp_raw) ** 2)\n",
    "\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))\n",
    "\n",
    "# Left: all three fits\n",
    "ax1.plot(T_raw, Cp_raw, 'ko', markersize=6, label='NIST data')\n",
    "ax1.plot(T_fine, Cp_poly, 'b-', linewidth=2, label=f'Polynomial (4 params)')\n",
    "ax1.plot(T_fine_K.flatten(), Cp_nn.flatten(), 'r--', linewidth=2, label=f'NN numpy (31 params)')\n",
    "ax1.plot(T_fine_K.flatten(), Cp_torch.flatten(), 'g-', linewidth=2, alpha=0.8, label=f'NN PyTorch (31 params)')\n",
    "ax1.set_xlabel('Temperature (K)')\n",
    "ax1.set_ylabel('$C_p$ (kJ/kg/K)')\n",
    "ax1.set_title('$C_p(T)$ for N$_2$ at 1 bar')\n",
    "ax1.legend()\n",
    "\n",
    "# Right: training loss comparison\n",
    "ax2.semilogy(losses_np, label='numpy (gradient descent)')\n",
    "ax2.semilogy(losses_torch, label='PyTorch (Adam)')\n",
    "ax2.set_xlabel('Epoch')\n",
    "ax2.set_ylabel('MSE (normalized)')\n",
    "ax2.set_title('Training loss comparison')\n",
    "ax2.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "print(f\"MSE — Polynomial: {mse_poly:.2e}  |  NN numpy: {mse_nn:.2e}  |  NN PyTorch: {mse_torch:.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "xyw3sr20brn",
   "metadata": {},
   "source": [
    "## 6. Extrapolation\n",
    "\n",
    "How do the models behave *outside* the training range? This is a key test — and where the differences become stark."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fi3iq2sjh6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extrapolate beyond the training range\n",
    "T_extrap = np.linspace(100, 2500, 300)\n",
    "T_extrap_norm = ((T_extrap - T_min) / (T_max - T_min)).reshape(-1, 1)\n",
    "\n",
    "# Polynomial extrapolation\n",
    "Cp_poly_extrap = poly(T_extrap)\n",
    "\n",
    "# Numpy NN extrapolation\n",
    "A1_extrap = np.tanh(T_extrap_norm @ W1 + b1)\n",
    "Cp_nn_extrap = (A1_extrap @ W2 + b2) * (Cp_max - Cp_min) + Cp_min\n",
    "\n",
    "# PyTorch NN extrapolation\n",
    "with torch.no_grad():\n",
    "    Cp_torch_extrap = model(torch.tensor(T_extrap_norm, dtype=torch.float32)).numpy()\n",
    "Cp_torch_extrap = Cp_torch_extrap * (Cp_max - Cp_min) + Cp_min\n",
    "\n",
    "plt.figure(figsize=(10, 6))\n",
    "plt.plot(T_raw, Cp_raw, 'ko', markersize=6, label='NIST data')\n",
    "plt.plot(T_extrap, Cp_poly_extrap, 'b-', linewidth=2, label='Polynomial')\n",
    "plt.plot(T_extrap, Cp_nn_extrap.flatten(), 'r--', linewidth=2, label='NN numpy')\n",
    "plt.plot(T_extrap, Cp_torch_extrap.flatten(), 'g-', linewidth=2, alpha=0.8, label='NN PyTorch')\n",
    "plt.axvline(T_raw.min(), color='gray', linestyle=':', alpha=0.5, label='Training range')\n",
    "plt.axvline(T_raw.max(), color='gray', linestyle=':', alpha=0.5)\n",
    "plt.xlabel('Temperature (K)')\n",
    "plt.ylabel('$C_p$ (kJ/kg/K)')\n",
    "plt.title('Extrapolation beyond training data')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "yb2s18keiw",
   "metadata": {},
   "source": [
    "## 7. Exercises\n",
    "\n",
    "Try these in new cells below:\n",
    "\n",
    "1. **Change the number of hidden neurons** (`H`). Try 2, 5, 20, 50. How does the fit change? At what point does adding neurons stop helping?\n",
    "\n",
    "2. **Activation functions**: In the PyTorch model, replace `nn.Tanh()` with `nn.ReLU()` or `nn.Sigmoid()`. How does the fit change?\n",
    "\n",
    "3. **Optimizer comparison**: Replace `Adam` with `torch.optim.SGD(model.parameters(), lr=0.01)`. How does training speed compare?\n",
    "\n",
    "4. **Remove normalization**: Use `T_raw` and `Cp_raw` directly (no scaling to [0,1]). What happens? Can you fix it by adjusting the learning rate?\n",
    "\n",
    "5. **Overfitting**: Set `H = 100` and train for 20,000 epochs. Does it fit the training data well? Look at the extrapolation — is it reasonable?\n",
    "\n",
    "6. **Higher-order polynomial**: Try `np.polyfit(T_raw, Cp_raw, 10)`. How does it compare to the cubic? How does it extrapolate?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}