1D Heat PINNs.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.autograd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "class Net1D(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super(Net1D, self).__init__()\n",
    "        self.hidden = nn.Linear(input_size, hidden_size)\n",
    "        self.output = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.tanh(self.hidden(x))\n",
    "        x = self.output(x)\n",
    "        return x\n",
    "\n",
    "def compute_loss(net, x, t):\n",
    "    xt = torch.cat([x, t], dim=1)\n",
    "    xt.requires_grad = True\n",
    "    \n",
    "    u = net(xt)\n",
    "    grads = torch.autograd.grad(u, xt, create_graph=True, grad_outputs=torch.ones_like(u))\n",
    "    u_t = grads[0][:, 1:2]\n",
    "    u_x = grads[0][:, 0:1]\n",
    "    u_xx = torch.autograd.grad(u_x, xt, create_graph=True, grad_outputs=torch.ones_like(u_x))[0][:, 0:1]\n",
    "    \n",
    "    # PDE Loss\n",
    "    pde_loss = (u_t - u_xx).pow(2).mean()\n",
    "    \n",
    "    # Boundary condition losses\n",
    "    u_left = net(torch.cat([torch.zeros_like(t), t], dim=1))\n",
    "    u_right = net(torch.cat([torch.ones_like(t), t], dim=1))\n",
    "    bc_loss = u_left.pow(2).mean() + u_right.pow(2).mean()\n",
    "    \n",
    "    # Initial condition loss\n",
    "    u_init = net(torch.cat([x, torch.zeros_like(x)], dim=1))\n",
    "    ic_loss = (u_init - torch.sin(np.pi * x)).pow(2).mean()\n",
    "    \n",
    "    # Total Loss\n",
    "    loss = pde_loss + bc_loss + ic_loss\n",
    "    \n",
    "    return loss\n",
    "\n",
    "# Training\n",
    "net = Net1D(2, 20, 1)\n",
    "optimizer = torch.optim.Adam(net.parameters(), lr=0.01)\n",
    "\n",
    "training_steps = [0, 5000, 10000,15000]\n",
    "solutions = []\n",
    "\n",
    "for steps in training_steps:\n",
    "    for epoch in range(steps):\n",
    "        optimizer.zero_grad()\n",
    "        x_rand = torch.rand(10, 1)\n",
    "        t_rand = torch.rand(10, 1)\n",
    "        loss = compute_loss(net, x_rand, t_rand)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # Store solution\n",
    "    x = torch.linspace(0, 1, 100).reshape(-1, 1)\n",
    "    t = torch.tensor([[0.1]])\n",
    "    xt = torch.cat([x, t.repeat(100, 1)], dim=1)\n",
    "    u_pred = net(xt).detach().numpy()\n",
    "    solutions.append(u_pred)\n",
    "\n",
    "# Analytical solution\n",
    "x = torch.linspace(0, 1, 100).reshape(-1, 1)\n",
    "u_true = np.exp(-np.pi**2 * 0.1) * np.sin(np.pi * x.numpy())\n",
    "\n",
    "# Plotting\n",
    "for i, steps in enumerate(training_steps):\n",
    "    plt.plot(x.numpy(), solutions[i], label=f\"PINN {steps} steps\")\n",
    "\n",
    "plt.plot(x.numpy(), u_true, 'k--', label=\"Analytical\")\n",
    "plt.legend()\n",
    "plt.title(\"1D Heat Equation: PINN vs Analytical\")\n",
    "plt.xlabel('x')\n",
    "plt.ylabel('u')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MATH420",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.15 (main, Nov 24 2022, 14:39:17) [MSC v.1916 64 bit (AMD64)]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "19b73810dcafedcbea3fadce48c320cc17c41811b51a3f83b2c30c62e8a41340"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}