{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6a739a84",
   "metadata": {},
   "source": [
    "# Tutorial: Two dimensional Wave problem with hard constraint\n",
    "\n",
    "In this tutorial we present how to solve the wave equation using hard constraint PINNs. For doing so we will build a costum `torch` model and pass it to the `PINN` solver.\n",
    "\n",
    "First of all, some useful imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d93daba0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "from pina.problem import SpatialProblem, TimeDependentProblem\n",
    "from pina.operators import laplacian, grad\n",
    "from pina.geometry import CartesianDomain\n",
    "from pina.solvers import PINN\n",
    "from pina.trainer import Trainer\n",
    "from pina.equation import Equation\n",
    "from pina.equation.equation_factory import FixedValue\n",
    "from pina import Condition, Plotter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2316f24e",
   "metadata": {},
   "source": [
    "## The problem definition "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc2bbf62",
   "metadata": {},
   "source": [
    "The problem is written in the following form:\n",
    "\n",
    "\\begin{equation}\n",
    "\\begin{cases}\n",
    "\\Delta u(x,y,t) = \\frac{\\partial^2}{\\partial t^2} u(x,y,t) \\quad \\text{in } D, \\\\\\\\\n",
    "u(x, y, t=0) = \\sin(\\pi x)\\sin(\\pi y), \\\\\\\\\n",
    "u(x, y, t) = 0 \\quad \\text{on } \\Gamma_1 \\cup \\Gamma_2 \\cup \\Gamma_3 \\cup \\Gamma_4,\n",
    "\\end{cases}\n",
    "\\end{equation}\n",
    "\n",
    "where $D$ is a square domain $[0,1]^2$, and $\\Gamma_i$, with $i=1,...,4$, are the boundaries of the square, and the velocity in the standard wave equation is fixed to one."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbc50741",
   "metadata": {},
   "source": [
    "Now, the wave problem is written in PINA code as a class, inheriting from `SpatialProblem` and `TimeDependentProblem` since we deal with spatial, and time dependent variables. The equations are written as `conditions` that should be satisfied in the corresponding domains. `truth_solution` is the exact solution which will be compared with the predicted one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b60176c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Wave(TimeDependentProblem, SpatialProblem):\n",
    "    output_variables = ['u']\n",
    "    spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})\n",
    "    temporal_domain = CartesianDomain({'t': [0, 1]})\n",
    "\n",
    "    def wave_equation(input_, output_):\n",
    "        u_t = grad(output_, input_, components=['u'], d=['t'])\n",
    "        u_tt = grad(u_t, input_, components=['dudt'], d=['t'])\n",
    "        nabla_u = laplacian(output_, input_, components=['u'], d=['x', 'y'])\n",
    "        return nabla_u - u_tt\n",
    "\n",
    "    def initial_condition(input_, output_):\n",
    "        u_expected = (torch.sin(torch.pi*input_.extract(['x'])) *\n",
    "                      torch.sin(torch.pi*input_.extract(['y'])))\n",
    "        return output_.extract(['u']) - u_expected\n",
    "\n",
    "    conditions = {\n",
    "        'gamma1': Condition(location=CartesianDomain({'x': [0, 1], 'y':  1, 't': [0, 1]}), equation=FixedValue(0.)),\n",
    "        'gamma2': Condition(location=CartesianDomain({'x': [0, 1], 'y': 0, 't': [0, 1]}), equation=FixedValue(0.)),\n",
    "        'gamma3': Condition(location=CartesianDomain({'x':  1, 'y': [0, 1], 't': [0, 1]}), equation=FixedValue(0.)),\n",
    "        'gamma4': Condition(location=CartesianDomain({'x': 0, 'y': [0, 1], 't': [0, 1]}), equation=FixedValue(0.)),\n",
    "        't0': Condition(location=CartesianDomain({'x': [0, 1], 'y': [0, 1], 't': 0}), equation=Equation(initial_condition)),\n",
    "        'D': Condition(location=CartesianDomain({'x': [0, 1], 'y': [0, 1], 't': [0, 1]}), equation=Equation(wave_equation)),\n",
    "    }\n",
    "\n",
    "    def wave_sol(self, pts):\n",
    "        return (torch.sin(torch.pi*pts.extract(['x'])) *\n",
    "                torch.sin(torch.pi*pts.extract(['y'])) *\n",
    "                torch.cos(torch.sqrt(torch.tensor(2.))*torch.pi*pts.extract(['t'])))\n",
    "\n",
    "    truth_solution = wave_sol\n",
    "\n",
    "problem = Wave()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03557e0c-1f82-4dad-b611-5d33fddfd0ef",
   "metadata": {},
   "source": [
    "## Hard Constraint Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "356fe363",
   "metadata": {},
   "source": [
    "After the problem, a **torch** model is needed to solve the PINN. Usually, many models are already implemented in **PINA**, but the user has the possibility to build his/her own model in `torch`. The hard constraint we impose is on the boundary of the spatial domain. Specifically, our solution is written as:\n",
    "\n",
    "$$ u_{\\rm{pinn}} = xy(1-x)(1-y)\\cdot NN(x, y, t), $$\n",
    "\n",
    "where $NN$ is the neural net output. This neural network takes as input the coordinates (in this case $x$, $y$ and $t$) and provides the unknown field $u$. By construction, it is zero on the boundaries. The residuals of the equations are evaluated at several sampling points (which the user can manipulate using the method `discretise_domain`) and the loss minimized by the neural network is the sum of the residuals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9fbbb74f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HardMLP(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, output_dim):\n",
    "        super().__init__()\n",
    "\n",
    "        self.layers = torch.nn.Sequential(torch.nn.Linear(input_dim, 40),\n",
    "                                          torch.nn.ReLU(),\n",
    "                                          torch.nn.Linear(40, 40),\n",
    "                                          torch.nn.ReLU(),\n",
    "                                          torch.nn.Linear(40, output_dim))\n",
    "        \n",
    "    # here in the foward we implement the hard constraints\n",
    "    def forward(self, x):\n",
    "        hard = x.extract(['x'])*(1-x.extract(['x']))*x.extract(['y'])*(1-x.extract(['y']))\n",
    "        return hard*self.layers(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f79fc901-4720-4fac-8b72-84ac5f7d2ec3",
   "metadata": {},
   "source": [
    "## Train and Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b465bebd",
   "metadata": {},
   "source": [
    "In this tutorial, the neural network is trained for 1000 epochs with a learning rate of 0.001 (default in `PINN`). Training takes approximately 3 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0be8e7f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "Missing logger folder: /Users/dariocoscia/Desktop/PINA/tutorials/tutorial3/lightning_logs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 999: : 1it [00:00, 84.47it/s, v_num=0, gamma1_loss=0.000, gamma2_loss=0.000, gamma3_loss=0.000, gamma4_loss=0.000, t0_loss=0.0419, D_loss=0.0307, mean_loss=0.0121]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=1000` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 999: : 1it [00:00, 68.69it/s, v_num=0, gamma1_loss=0.000, gamma2_loss=0.000, gamma3_loss=0.000, gamma4_loss=0.000, t0_loss=0.0419, D_loss=0.0307, mean_loss=0.0121]\n"
     ]
    }
   ],
   "source": [
    "# generate the data\n",
    "problem.discretise_domain(1000, 'random', locations=['D', 't0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])\n",
    "\n",
    "# crete the solver\n",
    "pinn = PINN(problem, HardMLP(len(problem.input_variables), len(problem.output_variables)))\n",
    "\n",
    "# create trainer and train\n",
    "trainer = Trainer(pinn, max_epochs=1000, accelerator='cpu', enable_model_summary=False) # we train on CPU and avoid model summary at beginning of training (optional)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2a5c405",
   "metadata": {},
   "source": [
    "Notice that the loss on the boundaries of the spatial domain is exactly zero, as expected! After the training is completed one can now plot some results using the `Plotter` class of **PINA**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c086c05f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 77316 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=0.5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 76880 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=1\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 73916 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plotter = Plotter()\n",
    "\n",
    "# plotting at fixed time t = 0.0\n",
    "print('Plotting at t=0')\n",
    "plotter.plot(pinn, fixed_variables={'t': 0.0})\n",
    "\n",
    "# plotting at fixed time t = 0.5\n",
    "print('Plotting at t=0.5')\n",
    "plotter.plot(pinn, fixed_variables={'t': 0.5})\n",
    "\n",
    "# plotting at fixed time t = 1.\n",
    "print('Plotting at t=1')\n",
    "plotter.plot(pinn, fixed_variables={'t': 1.0})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35e51649",
   "metadata": {},
   "source": [
    "The results are not so great, and we can clearly see that as time progress the solution get worse.... Can we do better?\n",
    "\n",
    "A valid option is to impose the initial condition as hard constraint as well. Specifically, our solution is written as:\n",
    "\n",
    "$$ u_{\\rm{pinn}} = xy(1-x)(1-y)\\cdot NN(x, y, t)\\cdot t + \\cos(\\sqrt{2}\\pi t)sin(\\pi x)\\sin(\\pi y), $$\n",
    "\n",
    "Let us build the network first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "33e43412",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HardMLPtime(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, input_dim, output_dim):\n",
    "        super().__init__()\n",
    "\n",
    "        self.layers = torch.nn.Sequential(torch.nn.Linear(input_dim, 40),\n",
    "                                          torch.nn.ReLU(),\n",
    "                                          torch.nn.Linear(40, 40),\n",
    "                                          torch.nn.ReLU(),\n",
    "                                          torch.nn.Linear(40, output_dim))\n",
    "        \n",
    "    # here in the foward we implement the hard constraints\n",
    "    def forward(self, x):\n",
    "        hard_space = x.extract(['x'])*(1-x.extract(['x']))*x.extract(['y'])*(1-x.extract(['y']))\n",
    "        hard_t = torch.sin(torch.pi*x.extract(['x'])) * torch.sin(torch.pi*x.extract(['y'])) * torch.cos(torch.sqrt(torch.tensor(2.))*torch.pi*x.extract(['t']))\n",
    "        return hard_space * self.layers(x) * x.extract(['t']) + hard_t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d3dc67b",
   "metadata": {},
   "source": [
    "Now let's train with the same configuration as thre previous test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f4bc6be2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0: : 0it [00:00, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 999: : 1it [00:00, 52.10it/s, v_num=1, gamma1_loss=1.97e-15, gamma2_loss=0.000, gamma3_loss=2.14e-15, gamma4_loss=0.000, t0_loss=0.000, D_loss=1.25e-7, mean_loss=2.09e-8]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=1000` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 999: : 1it [00:00, 45.78it/s, v_num=1, gamma1_loss=1.97e-15, gamma2_loss=0.000, gamma3_loss=2.14e-15, gamma4_loss=0.000, t0_loss=0.000, D_loss=1.25e-7, mean_loss=2.09e-8]\n"
     ]
    }
   ],
   "source": [
    "# generate the data\n",
    "problem.discretise_domain(1000, 'random', locations=['D', 't0', 'gamma1', 'gamma2', 'gamma3', 'gamma4'])\n",
    "\n",
    "# crete the solver\n",
    "pinn = PINN(problem, HardMLPtime(len(problem.input_variables), len(problem.output_variables)))\n",
    "\n",
    "# create trainer and train\n",
    "trainer = Trainer(pinn, max_epochs=1000, accelerator='cpu', enable_model_summary=False) # we train on CPU and avoid model summary at beginning of training (optional)\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0f80cb8",
   "metadata": {},
   "source": [
    "We can clearly see that the loss is way lower now. Let's plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "019767e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=0\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 52140 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=0.5\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 70856 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plotting at t=1\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 66204 characters]",
      "text/plain": [
       "<Figure size 1600x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plotter = Plotter()\n",
    "\n",
    "# plotting at fixed time t = 0.0\n",
    "print('Plotting at t=0')\n",
    "plotter.plot(pinn, fixed_variables={'t': 0.0})\n",
    "\n",
    "# plotting at fixed time t = 0.5\n",
    "print('Plotting at t=0.5')\n",
    "plotter.plot(pinn, fixed_variables={'t': 0.5})\n",
    "\n",
    "# plotting at fixed time t = 1.\n",
    "print('Plotting at t=1')\n",
    "plotter.plot(pinn, fixed_variables={'t': 1.0})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7338109",
   "metadata": {},
   "source": [
    "We can see now that the results are way better! This is due to the fact that previously the network was  not learning correctly the initial conditon, leading to a poor solution when the time evolved. By imposing the initial condition the network is able to correctly solve the problem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61195b1f",
   "metadata": {},
   "source": [
    "## What's next?\n",
    "\n",
    "Nice you have completed the two dimensional Wave tutorial of **PINA**! There are multiple directions you can go now:\n",
    "\n",
    "1. Train the network for longer or with different layer sizes and assert the finaly accuracy\n",
    "\n",
    "2. Propose new types of hard constraints in time, e.g. $$ u_{\\rm{pinn}} = xy(1-x)(1-y)\\cdot NN(x, y, t)(1-\\exp(-t)) + \\cos(\\sqrt{2}\\pi t)sin(\\pi x)\\sin(\\pi y), $$\n",
    "\n",
    "3. Exploit extrafeature training for model 1 and 2\n",
    "\n",
    "4. Many more..."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "56be7540488f3dc66429ddf54a0fa9de50124d45fcfccfaf04c4c3886d735a3a"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
