{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6f71ca5c",
   "metadata": {},
   "source": [
    "# Tutorial 1: Physics Informed Neural Networks on PINA"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ef4949c9",
   "metadata": {},
   "source": [
    "In this tutorial, we will demonstrate a typical use case of PINA on a toy problem. Specifically, the tutorial aims to introduce the following topics:\n",
    "\n",
    "* Defining a PINA Problem,\n",
    "* Building a `pinn` object,\n",
    "* Sampling points in a domain\n",
    "\n",
    "These are the three main steps needed **before** training a Physics Informed Neural Network (PINN). We will show each step in detail, and at the end, we will solve the problem."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1bd1904d",
   "metadata": {},
   "source": [
    "## PINA Problem"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cf9c96e3",
   "metadata": {},
   "source": [
    "### Initialize the `Problem` class"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8a819659",
   "metadata": {},
   "source": [
    "Problem definition in the PINA framework is done by building a python `class`, which inherits from one or more problem classes (`SpatialProblem`, `TimeDependentProblem`, `ParametricProblem`) depending on the nature of the problem. Below is an example:\n",
    "#### Simple Ordinary Differential Equation\n",
    "Consider the following:\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\begin{cases}\n",
    "\\frac{d}{dx}u(x) &=  u(x) \\quad x\\in(0,1)\\\\\n",
    "u(x=0) &= 1 \\\\\n",
    "\\end{cases}\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "with the analytical solution $u(x) = e^x$. In this case, our ODE depends only on the spatial variable $x\\in(0,1)$ , meaning that our `Problem` class is going to be inherited from the `SpatialProblem` class:\n",
    "\n",
    "```python\n",
    "from pina.problem import SpatialProblem\n",
    "from pina import CartesianProblem\n",
    "\n",
    "class SimpleODE(SpatialProblem):\n",
    "    \n",
    "    output_variables = ['u']\n",
    "    spatial_domain = CartesianProblem({'x': [0, 1]})\n",
    "\n",
    "    # other stuff ...\n",
    "```\n",
    "\n",
    "Notice that we define `output_variables` as a list of symbols, indicating the output variables of our equation (in this case only $u$). The `spatial_domain` variable indicates where the sample points are going to be sampled in the domain, in this case $x\\in[0,1]$."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "4e0a22bc",
   "metadata": {},
   "source": [
    "What about if our equation is also time dependent? In this case, our `class` will inherit from both `SpatialProblem` and `TimeDependentProblem`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2373a925",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pina.problem import SpatialProblem, TimeDependentProblem\n",
    "from pina import CartesianDomain\n",
    "\n",
    "class TimeSpaceODE(SpatialProblem, TimeDependentProblem):\n",
    "    \n",
    "    output_variables = ['u']\n",
    "    spatial_domain = CartesianDomain({'x': [0, 1]})\n",
    "    temporal_domain = CartesianDomain({'t': [0, 1]})\n",
    "\n",
    "    # other stuff ..."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ad8566b8",
   "metadata": {},
   "source": [
    "where we have included the `temporal_domain` variable, indicating the time domain wanted for the solution.\n",
    "\n",
    "In summary, using PINA, we can initialize a problem with a class which inherits from three base classes: `SpatialProblem`, `TimeDependentProblem`, `ParametricProblem`, depending on the type of problem we are considering. For reference:\n",
    "* `SpatialProblem` $\\rightarrow$ a differential equation with spatial variable(s)\n",
    "* `TimeDependentProblem` $\\rightarrow$ a time-dependent differential equation\n",
    "* `ParametricProblem` $\\rightarrow$ a parametrized differential equation"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "592a4c43",
   "metadata": {},
   "source": [
    "### Write the `Problem` class\n",
    "\n",
    "Once the `Problem` class is initialized, we need to represent the differential equation in PINA. In order to do this, we need to load the PINA operators from `pina.operators` module. Again, we'll consider Equation (1) and represent it in PINA:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f2608e2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pina.problem import SpatialProblem\n",
    "from pina.operators import grad\n",
    "from pina import Condition, CartesianDomain\n",
    "from pina.equation.equation import Equation\n",
    "\n",
    "import torch\n",
    "\n",
    "\n",
    "class SimpleODE(SpatialProblem):\n",
    "\n",
    "    output_variables = ['u']\n",
    "    spatial_domain = CartesianDomain({'x': [0, 1]})\n",
    "\n",
    "    # defining the ode equation\n",
    "    def ode_equation(input_, output_):\n",
    "\n",
    "        # computing the derivative\n",
    "        u_x = grad(output_, input_, components=['u'], d=['x'])\n",
    "\n",
    "        # extracting the u input variable\n",
    "        u = output_.extract(['u'])\n",
    "\n",
    "        # calculate the residual and return it\n",
    "        return u_x - u\n",
    "\n",
    "    # defining the initial condition\n",
    "    def initial_condition(input_, output_):\n",
    "        \n",
    "        # setting the initial value\n",
    "        value = 1.0\n",
    "\n",
    "        # extracting the u input variable\n",
    "        u = output_.extract(['u'])\n",
    "\n",
    "        # calculate the residual and return it\n",
    "        return u - value\n",
    "\n",
    "    # conditions to hold\n",
    "    conditions = {\n",
    "        'x0': Condition(location=CartesianDomain({'x': 0.}), equation=Equation(initial_condition)),\n",
    "        'D': Condition(location=CartesianDomain({'x': [0, 1]}), equation=Equation(ode_equation)),\n",
    "    }\n",
    "\n",
    "    # sampled points (see below)\n",
    "    input_pts = None\n",
    "\n",
    "    # defining the true solution\n",
    "    def truth_solution(self, pts):\n",
    "        return torch.exp(pts.extract(['x']))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7cf64d01",
   "metadata": {},
   "source": [
    "After we define the `Problem` class, we need to write different class methods, where each method is a function returning a residual. These functions are the ones minimized during PINN optimization, given the initial conditions. For example, in the domain $[0,1]$, the ODE equation (`ode_equation`) must be satisfied. We represent this by returning the difference between subtracting the variable `u` from its gradient (the residual), which we hope to minimize to 0. This is done for all conditions (`ode_equation`, `initial_condition`). \n",
    "\n",
    "Once we have defined the function, we need to tell the neural network where these methods are to be applied. To do so, we use the `Condition` class. In the `Condition` class, we pass the location points and the function we want minimized on those points (other possibilities are allowed, see the documentation for reference) as parameters.\n",
    "\n",
    "Finally, it's possible to define a `truth_solution` function, which can be useful if we want to plot the results and see how the real solution compares to the expected (true) solution. Notice that the `truth_solution` function is a method of the `PINN` class, but is not mandatory for problem definition.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "22e502dd",
   "metadata": {},
   "source": [
    "## Build the `PINN` object"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "075f43f5",
   "metadata": {},
   "source": [
    "The basic requirements for building a `PINN` model are a `Problem` and a model. We have just covered the `Problem` definition. For the model parameter, one can use either the default models provided in PINA or a custom model. We will not go into the details of model definition (see Tutorial2 and Tutorial3 for more details on model definition)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3bb4dc9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pina.model import FeedForward\n",
    "from pina import PINN\n",
    "\n",
    "# initialize the problem\n",
    "problem = SimpleODE()\n",
    "\n",
    "# build the model\n",
    "model = FeedForward(\n",
    "    layers=[10, 10],\n",
    "    func=torch.nn.Tanh,\n",
    "    output_dimensions=len(problem.output_variables),\n",
    "    input_dimensions=len(problem.input_variables)\n",
    ")\n",
    "\n",
    "# create the PINN object\n",
    "pinn = PINN(problem, model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8d2cb313",
   "metadata": {},
   "source": [
    "Creating the `PINN` object is fairly simple. Different optional parameters include: optimizer, batch size, ... (see [documentation](https://mathlab.github.io/PINA/) for reference)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "78b30f95",
   "metadata": {},
   "source": [
    "## Sample points in the domain "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "53c783e8",
   "metadata": {},
   "source": [
    "Once the `PINN` object is created, we need to generate the points for starting the optimization. To do so, we use the `sample` method of the `CartesianDomain` class. Below are three examples of sampling methods on the $[0,1]$ domain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "09ce5c3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sampling 20 points in [0, 1] through discretization\n",
    "pinn.problem.discretise_domain(n=20, mode='grid', variables=['x'])\n",
    "\n",
    "# sampling 20 points in (0, 1) through latin hypercube samping\n",
    "pinn.problem.discretise_domain(n=20, mode='latin', variables=['x'])\n",
    "\n",
    "# sampling 20 points in (0, 1) randomly\n",
    "pinn.problem.discretise_domain(n=20, mode='random', variables=['x'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "27a287db",
   "metadata": {},
   "source": [
    "### Very simple training and plotting\n",
    "\n",
    "Once we have defined the PINA model, created a network, and sampled points in the domain, we have everything necessary for training a PINN. To do so, we make use of the `Trainer` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f8b4f496",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/u/n/ndemo/.local/lib/python3.9/site-packages/torch/cuda/__init__.py:546: UserWarning: Can't initialize NVML\n",
      "  warnings.warn(\"Can't initialize NVML\")\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/u/n/ndemo/.local/lib/python3.9/site-packages/lightning/pytorch/loops/utilities.py:72: PossibleUserWarning: `max_epochs` was not set. Setting it to 1000 epochs. To train without an epoch limit, set `max_epochs=-1`.\n",
      "  rank_zero_warn(\n",
      "2023-10-17 10:02:21.318700: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2023-10-17 10:02:21.345355: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-10-17 10:02:23.572602: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "/opt/sissa/apps/intelpython/2022.0.2/intelpython/latest/lib/python3.9/site-packages/scipy/__init__.py:138: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.26.0)\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion} is required for this version of \"\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name        | Type    | Params\n",
      "----------------------------------------\n",
      "0 | _loss       | MSELoss | 0     \n",
      "1 | _neural_net | Network | 141   \n",
      "----------------------------------------\n",
      "141       Trainable params\n",
      "0         Non-trainable params\n",
      "141       Total params\n",
      "0.001     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e99075a1776436eb94b80f7dfbbc794",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=1000` reached.\n"
     ]
    }
   ],
   "source": [
    "from pina import Trainer\n",
    "\n",
    "# initialize trainer\n",
    "trainer = Trainer(pinn)\n",
    "\n",
    "# train the model\n",
    "trainer.train()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
  },
  "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}