{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e80567a6",
   "metadata": {},
   "source": [
    "# Tutorial: Two dimensional Darcy flow using the Fourier Neural Operator"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8762bbe5",
   "metadata": {},
   "source": [
    "In this tutorial we are going to solve the Darcy flow problem in two dimensions, presented in [*Fourier Neural Operator for\n",
    "Parametric Partial Differential Equation*](https://openreview.net/pdf?id=c8P9NQVtmnO). First of all we import the modules needed for the tutorial. Importing `scipy` is needed for input output operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5f2744dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install scipy  # install scipy\n",
    "from scipy import io\n",
    "import torch\n",
    "from pina.model import FNO, FeedForward  # let's import some models\n",
    "from pina import Condition, LabelTensor\n",
    "from pina.solvers import SupervisedSolver\n",
    "from pina.trainer import Trainer\n",
    "from pina.problem import AbstractProblem\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cf5b181",
   "metadata": {},
   "source": [
    "## Data Generation\n",
    "\n",
    "We will focus on solving the a specfic PDE, the **Darcy Flow** equation. The Darcy PDE is a second order, elliptic PDE with the following form:\n",
    "\n",
    "$$\n",
    "-\\nabla\\cdot(k(x, y)\\nabla u(x, y)) = f(x) \\quad (x, y) \\in D.\n",
    "$$\n",
    "\n",
    "Specifically, $u$ is the flow pressure, $k$ is the permeability field and $f$ is the forcing function. The Darcy flow can parameterize a variety of systems including flow through porous media, elastic materials and heat conduction. Here you will define the domain as a 2D unit square Dirichlet boundary conditions. The dataset is taken from the authors original reference.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2ffb8a4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download the dataset\n",
    "data = io.loadmat(\"Data_Darcy.mat\")\n",
    "\n",
    "# extract data (we use only 100 data for train)\n",
    "k_train = LabelTensor(torch.tensor(data['k_train'], dtype=torch.float).unsqueeze(-1), ['u0'])\n",
    "u_train = LabelTensor(torch.tensor(data['u_train'], dtype=torch.float).unsqueeze(-1), ['u'])\n",
    "k_test =  LabelTensor(torch.tensor(data['k_test'], dtype=torch.float).unsqueeze(-1), ['u0'])\n",
    "u_test= LabelTensor(torch.tensor(data['u_test'], dtype=torch.float).unsqueeze(-1), ['u'])\n",
    "x = torch.tensor(data['x'], dtype=torch.float)[0]\n",
    "y = torch.tensor(data['y'], dtype=torch.float)[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a9defd4",
   "metadata": {},
   "source": [
    "Let's visualize some data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c8501b6f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20632 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.subplot(1, 2, 1)\n",
    "plt.title('permeability')\n",
    "plt.imshow(k_train.squeeze(-1)[0])\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.title('field solution')\n",
    "plt.imshow(u_train.squeeze(-1)[0])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89a77ff1",
   "metadata": {},
   "source": [
    "We now create the neural operator class. It is a very simple class, inheriting from `AbstractProblem`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8b27d283",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralOperatorSolver(AbstractProblem):\n",
    "    input_variables = k_train.labels\n",
    "    output_variables = u_train.labels\n",
    "    conditions = {'data' : Condition(input_points=k_train, \n",
    "                                     output_points=u_train)}\n",
    "\n",
    "# make problem\n",
    "problem = NeuralOperatorSolver()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1096cc20",
   "metadata": {},
   "source": [
    "## Solving the problem with a FeedForward Neural Network\n",
    "\n",
    "We will first solve the problem using a Feedforward neural network. We will use the `SupervisedSolver` for solving the problem, since we are training using supervised learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e34f18b0",
   "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 9: : 100it [00:00, 357.28it/s, v_num=1, mean_loss=0.108]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=10` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: : 100it [00:00, 354.81it/s, v_num=1, mean_loss=0.108]\n"
     ]
    }
   ],
   "source": [
    "# make model\n",
    "model = FeedForward(input_dimensions=1, output_dimensions=1)\n",
    "\n",
    "\n",
    "# make solver\n",
    "solver = SupervisedSolver(problem=problem, model=model)\n",
    "\n",
    "# make the trainer and train\n",
    "trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) # we train on CPU and avoid model summary at beginning of training (optional)\n",
    "trainer.train()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b2c35be",
   "metadata": {},
   "source": [
    "The final loss is pretty high... We can calculate the error by importing `LpLoss`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0e2a6aa4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final error training 56.04%\n",
      "Final error testing 56.01%\n"
     ]
    }
   ],
   "source": [
    "from pina.loss import LpLoss\n",
    "\n",
    "# make the metric\n",
    "metric_err = LpLoss(relative=True)\n",
    "\n",
    "\n",
    "err = float(metric_err(u_train.squeeze(-1), solver.neural_net(k_train).squeeze(-1)).mean())*100\n",
    "print(f'Final error training {err:.2f}%')\n",
    "\n",
    "err = float(metric_err(u_test.squeeze(-1), solver.neural_net(k_test).squeeze(-1)).mean())*100\n",
    "print(f'Final error testing {err:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b5e5aa6",
   "metadata": {},
   "source": [
    "## Solving the problem with a Fuorier Neural Operator (FNO)\n",
    "\n",
    "We will now move to solve the problem using a FNO. Since we are learning operator this approach is better suited, as we shall see."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "9af523a5",
   "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 9: : 100it [00:02, 47.76it/s, v_num=4, mean_loss=0.00106] "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=10` reached.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9: : 100it [00:02, 47.65it/s, v_num=4, mean_loss=0.00106]\n"
     ]
    }
   ],
   "source": [
    "# make model\n",
    "lifting_net = torch.nn.Linear(1, 24)\n",
    "projecting_net = torch.nn.Linear(24, 1)\n",
    "model = FNO(lifting_net=lifting_net,\n",
    "            projecting_net=projecting_net,\n",
    "            n_modes=8,\n",
    "            dimensions=2,\n",
    "            inner_size=24,\n",
    "            padding=8)\n",
    "\n",
    "\n",
    "# make solver\n",
    "solver = SupervisedSolver(problem=problem, model=model)\n",
    "\n",
    "# make the trainer and train\n",
    "trainer = Trainer(solver=solver, max_epochs=10, accelerator='cpu', enable_model_summary=False, batch_size=10) # we train on CPU and avoid model summary at beginning of training (optional)\n",
    "trainer.train()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84964cb9",
   "metadata": {},
   "source": [
    "We can clearly see that the final loss is lower. Let's see in testing.. Notice that the number of parameters is way higher than a `FeedForward` network. We suggest to use GPU or TPU for a speed up in training, when many data samples are used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "58e2db89",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Final error training 4.83%\n",
      "Final error testing 5.16%\n"
     ]
    }
   ],
   "source": [
    "err = float(metric_err(u_train.squeeze(-1), solver.neural_net(k_train).squeeze(-1)).mean())*100\n",
    "print(f'Final error training {err:.2f}%')\n",
    "\n",
    "err = float(metric_err(u_test.squeeze(-1), solver.neural_net(k_test).squeeze(-1)).mean())*100\n",
    "print(f'Final error testing {err:.2f}%')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26e3a6e4",
   "metadata": {},
   "source": [
    "As we can see the loss is way lower!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba1dfa4b",
   "metadata": {},
   "source": [
    "## What's next?\n",
    "\n",
    "We have made a very simple example on how to use the `FNO` for learning neural operator. Currently in **PINA** we implement 1D/2D/3D cases. We suggest to extend the tutorial using more complex problems and train for longer, to see the full potential of neural operators."
   ]
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
