2-intro_FNO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "59194c45-83c9-4a77-a1b0-185eca26afd5",
   "metadata": {},
   "source": [
    "# Check the dependencies "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "af7a5c4c-b3a5-4f32-aee9-55290566ff56",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tl.__version__='0.8.0'\n",
      "tltorch.__version__='0.3.0'\n",
      "no.__version__='0.1.0'\n"
     ]
    }
   ],
   "source": [
    "import tensorly as tl\n",
    "import tltorch\n",
    "import neuralop as no\n",
    "\n",
    "print(f'{tl.__version__=}')\n",
    "print(f'{tltorch.__version__=}')\n",
    "print(f'{no.__version__=}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a36bb3e2-c158-497c-babe-5eead700cbf1",
   "metadata": {
    "tags": []
   },
   "source": [
    "# FFT and Spectral Convolution\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4efa0d7f-e39c-496e-891d-6b34c62fbd9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from neuralop.models.fno_block import FactorizedSpectralConv\n",
    "from neuralop.models import TFNO2d\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2c8c3eb7-82e1-4df7-b6e6-3f34331637c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fourier_conv = FactorizedSpectralConv(in_channels=3, out_channels=10, n_modes=(4, 4),\n",
    "                                      factorization=None, implementation='reconstructed')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4eaf645a-b7f5-4dcc-b8de-fb388ccc9b26",
   "metadata": {},
   "outputs": [],
   "source": [
    "in_data = torch.randn((2, 3, 16, 16))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "016b33e0-88a6-4215-99e0-19da4f8fd5f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "out = fourier_conv(in_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "36d0f546-9fa9-4936-a6b6-19d7bde03639",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 10, 16, 16])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4936746b-5abb-4a8b-9e74-238502c65930",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FactorizedSpectralConv(\n",
       "  (weight): ModuleList(\n",
       "    (0): ComplexDenseTensor(shape=torch.Size([3, 10, 2, 2]), rank=None)\n",
       "    (1): ComplexDenseTensor(shape=torch.Size([3, 10, 2, 2]), rank=None)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fourier_conv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a616d68d-677a-4e6f-abd5-9e631ebf7fb6",
   "metadata": {},
   "source": [
    "The way the spectral convolution works is that it multiplies (complex) coefficients with (complex) weights, learned end-to-end."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c8d9860-d43d-47f3-a6aa-c7ed4522684e",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Tensorized Spectral Convolutions\n",
    "\n",
    "It is possible to express the weights of one or more layers as in factorized form, as a low-rank decomposition of the full weights.\n",
    "\n",
    "`neuralop` comes with support for tensorization out of the box, you can simply specify, e.g., to use a Tucker factorization, `factorization='tucker'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b3f919de-97c2-4f0b-bb40-8e47cd2c1e0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fourier_conv = FactorizedSpectralConv(in_channels=3, out_channels=10, n_modes=(4, 4),\n",
    "                                      factorization='tucker', implementation='reconstructed')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "91a7aa04-9cc3-4f8c-b34f-54fbc625b718",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FactorizedSpectralConv(\n",
       "  (weight): ModuleList(\n",
       "    (0): ComplexTuckerTensor(shape=(3, 10, 2, 2), rank=(1, 5, 1, 1))\n",
       "    (1): ComplexTuckerTensor(shape=(3, 10, 2, 2), rank=(1, 5, 1, 1))\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fourier_conv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8df876d-72e1-40cd-9a86-330a57dc0e8d",
   "metadata": {},
   "source": [
    "## Efficient forward pass\n",
    "\n",
    "When factorizing the weights, have two main options during the forward pass:\n",
    "1. reconstruct the full weights and use that for the forward pass \n",
    "2. contract the input directly with the factorized weights to predict the output\n",
    "\n",
    "When the factorized weights are small, the second option can lead to large speedups or memory reduction, particularly when coupled with checkpointing. \n",
    "\n",
    "In `neuralop`, you can use those simply by specifying `implementation='reconstructed'` or `implementation='factorized'`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a0667a6b-1efe-47e0-8908-29c5fb0cf45a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fourier_conv = FactorizedSpectralConv(in_channels=3, out_channels=10, n_modes=(4, 4),\n",
    "                                      factorization='tucker', implementation='factorized')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec3ab24a-09fe-4864-b2ed-e96b54792e9f",
   "metadata": {},
   "source": [
    "# Full Tensorized Fourier Neural Operator \n",
    "\n",
    "The full architecture is composed of \n",
    "\n",
    "i) a lifting layer taking the number of input channels and lifting that to the desired number of hidden channels\n",
    "ii) a number of spectral convolutions, as shown above\n",
    "iii) a projection layer projecting back from the number of hidden channels to the desired number of output channels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d51aec17-2cf4-40c4-9452-84a4b5259db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "tfno = TFNO2d(n_modes_height=16, n_modes_width=16, hidden_channels=16, \n",
    "              factorization=None, skip='linear')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c87127e5-d24c-4096-be3a-8872a853a132",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TFNO2d(\n",
       "  (convs): FactorizedSpectralConv2d(\n",
       "    (weight): ModuleList(\n",
       "      (0): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "      (1): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "      (2): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "      (3): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "      (4): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "      (5): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "      (6): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "      (7): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    )\n",
       "  )\n",
       "  (fno_skips): ModuleList(\n",
       "    (0): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (1): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (2): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "    (3): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "  )\n",
       "  (lifting): Lifting(\n",
       "    (fc): Conv2d(3, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       "  (projection): Projection(\n",
       "    (fc1): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "    (fc2): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfno"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e70efec-bf3c-48ac-b53a-59800055f1b9",
   "metadata": {},
   "source": [
    "## Lifting layer\n",
    "\n",
    "Increasing the number of channels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1deead74-bd3d-4aa9-8d2c-cfd9ab0763d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Lifting(\n",
       "  (fc): Conv2d(3, 16, kernel_size=(1, 1), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfno.lifting"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08844bac-9335-4ac4-afc8-f1d67c3e31bb",
   "metadata": {},
   "source": [
    "## Spectral convolutions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f2bc28dc-1226-4ed3-b757-3c42357d276a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "FactorizedSpectralConv2d(\n",
       "  (weight): ModuleList(\n",
       "    (0): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    (1): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    (2): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    (3): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    (4): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    (5): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    (6): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "    (7): ComplexDenseTensor(shape=torch.Size([16, 16, 8, 8]), rank=None)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfno.convs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c7d9882-13db-447d-affd-07ef17256e1c",
   "metadata": {},
   "source": [
    "## Skip connections: recovering non-periodicity\n",
    "\n",
    "Recall the FNO architecture has skip connections: the FFT transformation will loose non-periodic information that has to be reinjected through skip connections. These skip connections also help with learning.\n",
    "\n",
    "![FNO_layer](./images/fourier_layer.png)\n",
    "\n",
    "Here, linear layer (represented by weight W in the image). We can also use Identity skip (`skip='identity'`) or soft-gated connections (`skip='soft-gating'`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f063e3bf-34e5-4d7f-83f9-b3522aa6430b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ModuleList(\n",
       "  (0): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "  (1): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "  (2): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "  (3): Conv2d(16, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfno.fno_skips"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "070e930e-38b6-4d3c-b62a-3ca700294c99",
   "metadata": {},
   "source": [
    "## Projection: going back to the target number of channels \n",
    "\n",
    "Finally, the projection layer takes the hidden dimension to projection_channels and to the actual number of output channels (here, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "88344f47-a7e8-458e-9fbb-775804fbbaad",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Projection(\n",
       "  (fc1): Conv2d(16, 256, kernel_size=(1, 1), stride=(1, 1))\n",
       "  (fc2): Conv2d(256, 1, kernel_size=(1, 1), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tfno.projection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aae1ab6-852c-4720-9b3b-5791c2b42872",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}