dnn_reinit_2d.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "dnn_reinit_2d.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyOvmFtf8p8xofRXJRl+8lO8"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "t7TlHimL0sSs"
      },
      "source": [
        "import math\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kRDpnral09Jv"
      },
      "source": [
        "class Implicit(nn.Module):\n",
        "  def __init__(self, fun, dimension, radius=1):\n",
        "    super(Implicit, self).__init__()\n",
        "\n",
        "    d_in = dimension\n",
        "    dims = [512, 512, 512, 512, 512, 512, 512, 512]\n",
        "    beta = 100\n",
        "    skip_in = [4]\n",
        "    radius_init = radius\n",
        "\n",
        "    dims = [d_in] + dims + [1]\n",
        "\n",
        "    self.num_layers = len(dims)\n",
        "    self.skip_in = skip_in\n",
        "\n",
        "    for layer in range(0, self.num_layers - 1):\n",
        "\n",
        "      if layer + 1 in skip_in:\n",
        "        out_dim = dims[layer + 1] - d_in\n",
        "      else:\n",
        "        out_dim = dims[layer + 1]\n",
        "\n",
        "      lin = nn.Linear(dims[layer], out_dim)\n",
        "\n",
        "      if layer == self.num_layers - 2:\n",
        "        torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[layer]), std=0.00001)\n",
        "        torch.nn.init.constant_(lin.bias, -radius_init)\n",
        "      else:\n",
        "        torch.nn.init.constant_(lin.bias, 0.0)\n",
        "        torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))\n",
        "\n",
        "      setattr(self, \"lin\" + str(layer), lin)\n",
        "\n",
        "    self.activation = nn.Softplus(beta=beta)\n",
        "    #self.activation = nn.ReLU()\n",
        "    self.fun = fun\n",
        "\n",
        "\n",
        "  def forward(self, input):\n",
        "    x = input\n",
        "    for layer in range(0, self.num_layers - 1):\n",
        "      lin = getattr(self, \"lin\" + str(layer))\n",
        "      if layer in self.skip_in:\n",
        "        x = torch.cat([x, input], -1) / np.sqrt(2)\n",
        "      x = lin(x)\n",
        "      if layer < self.num_layers - 2:\n",
        "        x = self.activation(x)\n",
        "    fun_sign = torch.tanh(0.1*self.fun(input))\n",
        "    fun_sign = fun_sign.reshape(x.shape)\n",
        "    x = x * fun_sign\n",
        "    return x"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nGG90eQD1FWp"
      },
      "source": [
        "def grad(y, x, grad_outputs=None):\n",
        "  if grad_outputs is None:\n",
        "    grad_outputs = torch.ones_like(y)\n",
        "  grad = torch.autograd.grad(y, [x], grad_outputs=grad_outputs, create_graph=True)[0]\n",
        "  return grad\n",
        "\n",
        "def div(y, x):\n",
        "  div = 0.\n",
        "  for i in range(y.shape[-1]):\n",
        "    div += torch.autograd.grad(y[..., i], x, torch.ones_like(y[..., i]), create_graph=True)[0][..., i:i+1]\n",
        "  return div\n",
        "\n",
        "def laplacian(y, x):\n",
        "  grad = grad(y, x)\n",
        "  return div(grad, x)\n",
        "\n",
        "def pLaplacian(y, x, p=2):\n",
        "  g = grad(y, x)\n",
        "  #g_n = g.norm(2, dim=1)\n",
        "  g_n = torch.linalg.norm(g, 2, dim=1)\n",
        "  g_n = g_n**(p-2)\n",
        "  g_n = torch.reshape(g_n, (g.shape[0],1))\n",
        "  g = g_n * g\n",
        "  return div(g, x)"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tN9UjU7hULXV"
      },
      "source": [
        "# Uniform distribution in 2d\n",
        "def uniform_data(num_samples, x_bounds, y_bounds, device='cpu'):\n",
        "  x_min, x_max = x_bounds\n",
        "  y_min, y_max = y_bounds\n",
        "  u_x = torch.FloatTensor(num_samples, 1).uniform_(x_min, x_max)\n",
        "  u_y = torch.FloatTensor(num_samples, 1).uniform_(y_min, y_max)\n",
        "  u = torch.cat((u_x, u_y), dim=-1)  \n",
        "  u = u.to(device)\n",
        "  return u"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jsRZxbme1SlO"
      },
      "source": [
        "def trainPPoisson(num_iters, fun, xbounds, ybounds, p=2, device='cpu'):\n",
        "  model = Implicit(fun=fun, dimension=2).to(device)\n",
        "  optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
        "\n",
        "  n_samples = 256\n",
        "  loss = 0\n",
        "\n",
        "  # Train network\n",
        "  for i in range(0, num_iters):\n",
        "    # Uniform samples\n",
        "    u = uniform_data(n_samples, xbounds, ybounds, device)\n",
        "    u.requires_grad = True\n",
        "    \n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    f_d = model(u)\n",
        "    \n",
        "    # p-Laplacian\n",
        "    lap = pLaplacian(f_d, u, p)\n",
        "    lap_loss = torch.mean((lap+1)**2)\n",
        "    loss = lap_loss\n",
        "\n",
        "    loss.backward()\n",
        "\n",
        "    #clip_grad = 1.0\n",
        "    #torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip_grad)\n",
        "\n",
        "    optimizer.step()\n",
        "  \n",
        "  return model"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ig3YTg7SBCIb"
      },
      "source": [
        "def trainEikonal(num_iters, fun, xbounds, ybounds, device='cpu'):\n",
        "  model = Implicit(fun=fun, dimension=2).to(device)\n",
        "  optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)\n",
        "\n",
        "  n_samples = 1024\n",
        "  loss = 0\n",
        "\n",
        "  # Train network\n",
        "  for i in range(0, num_iters):\n",
        "    # Uniform samples\n",
        "    u = uniform_data(n_samples, xbounds, ybounds, device)\n",
        "    u.requires_grad = True\n",
        "    \n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    f_d = model(u)\n",
        "    \n",
        "    g_d = grad(f_d, u)\n",
        "    g_norm = (g_d.norm(2, dim=1) - 1)**2\n",
        "    eik_loss = torch.mean(g_norm)\n",
        "    loss = eik_loss\n",
        "\n",
        "    loss.backward()\n",
        "\n",
        "    #clip_grad = 1.0\n",
        "    #torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=clip_grad)\n",
        "\n",
        "    optimizer.step()\n",
        "  \n",
        "  return model"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SmM57zeO18WE"
      },
      "source": [
        "# The implicit surface (a unit disk)\n",
        "def f(p):\n",
        "  return 1.0 - p[:,0]**2 - p[:,1]**2"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dlaujKrDUV5H"
      },
      "source": [
        "# For visualization\n",
        "def showZeroLevel(x, y, f):\n",
        "  plt.figure(figsize=(8,4))\n",
        "  h = plt.contour(x, y, f, levels=[0.0], colors='b')\n",
        "  h.ax.axis('equal')\n",
        "  plt.title('Zero level-set')\n",
        "  plt.show()\n",
        "\n",
        "def showContourPlot(x, y, f): \n",
        "  plt.figure(figsize=(8,4))\n",
        "  h = plt.contourf(x, y, f)\n",
        "  h.ax.axis('equal')\n",
        "  plt.title('Filled Contour Plot')\n",
        "  plt.show()"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-z4TZE9vUZtd"
      },
      "source": [
        "# Create a grid from torch tensors\n",
        "def torchLinearGrid(xbounds, ybounds, grid_res, device='cpu'):\n",
        "  xmin = xbounds[0]\n",
        "  xmax = xbounds[1]\n",
        "  ymin = ybounds[0]\n",
        "  ymax = ybounds[1]\n",
        "\n",
        "  dx = xmax - xmin\n",
        "  dy = ymax - ymin\n",
        "\n",
        "  resx = grid_res[0]\n",
        "  resy = grid_res[1]\n",
        "\n",
        "  x = torch.arange(xmin, xmax, step=dx/float(resx))\n",
        "  y = torch.arange(ymin, ymax, step=dy/float(resy))\n",
        "\n",
        "  xx, yy = torch.meshgrid(x, y)\n",
        "\n",
        "  xx = xx.to(device)\n",
        "  yy = yy.to(device)\n",
        "\n",
        "  dimg = resx * resy\n",
        "  xy = torch.stack((xx, yy), dim=-1).reshape(dimg,2)\n",
        "\n",
        "  return xy"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TcWiD_S6z_mp"
      },
      "source": [
        "def torchLinearSampling(model, xy):\n",
        "  d = model(xy)\n",
        "  return d"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TtnZ7kch0QfA"
      },
      "source": [
        "def saveVTK(filename, xy, res, field):\n",
        "  resx, resy = res\n",
        "  resz = 1\n",
        "  \n",
        "  # set the z coord to 0 \n",
        "  xyz = torch.zeros((xy.shape[0], 3))\n",
        "  xyz[:,0:2] = xy\n",
        "\n",
        "  field_title = 'VALUE'\n",
        "\n",
        "  with open(filename, 'w') as f:\n",
        "    f.write('# vtk DataFile Version 3.0\\n')\n",
        "    f.write('vtk output\\n')\n",
        "    f.write('ASCII\\n')\n",
        "    f.write('DATASET STRUCTURED_GRID\\n')\n",
        "    f.write('DIMENSIONS ' + str(resx) + ' ' + str(resy) + ' ' + str(resz) +'\\n')\n",
        "    f.write('POINTS ' + str(resx*resy*resz) + ' double\\n')\n",
        "\n",
        "    np.savetxt(f, xyz.detach().cpu().numpy())\n",
        "    \n",
        "    f.write('\\n\\n')\n",
        "\n",
        "    f.write('POINT_DATA ' + str(resx*resy*resz) + '\\n')\n",
        "    f.write('SCALARS ' + field_title + ' double' + '\\n')\n",
        "    f.write('LOOKUP_TABLE default\\n')\n",
        "        \n",
        "    np.savetxt(f, field.detach().cpu().numpy())\n",
        "    f.write('\\n')"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vGYWh0L3AR41"
      },
      "source": [
        "# Domain boundary\n",
        "xbounds = (-2.0, 2.0)\n",
        "ybounds = (-2.0, 2.0)"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Gbbg-F9B1ZY-",
        "outputId": "c9215c42-6110-4f35-e507-8bd56117ca1c"
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(device)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "cuda\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i63MLRP603gf"
      },
      "source": [
        "xy = torchLinearGrid(xbounds, ybounds, (64, 64), device=device)"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a75RkCNa1FKv"
      },
      "source": [
        "yt = f(xy)"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IFNm-UAs13gN"
      },
      "source": [
        "saveVTK('2d_disk_rfunction.vtk', xy, (64,64), yt)"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tLgM0sAu2Csd"
      },
      "source": [
        "# p-Poisson problem, p = 8\n",
        "max_iteration = 15000\n",
        "p = 8\n",
        "model_p8 = trainPPoisson(max_iteration, fun=f, xbounds=xbounds, ybounds=ybounds, p=p, device=device)"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r5imNjcg2Q9k"
      },
      "source": [
        "yt_pLap_p8 = model_p8(xy)"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0X1cFUih2YpN"
      },
      "source": [
        "saveVTK('2d_disk_pLap_p8.vtk', xy, (64,64), yt_pLap_p8)"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IWVVJPVhTa2m"
      },
      "source": [
        "# Eikonal problem\n",
        "max_iteration = 15000\n",
        "model_eik = trainEikonal(max_iteration, fun=f, xbounds=xbounds, ybounds=ybounds, device=device)"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GYAI4i1-2ona"
      },
      "source": [
        "yt_eik = model_eik(xy)"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kwab4I1r2q87"
      },
      "source": [
        "saveVTK('2d_disk_eikonal.vtk', xy, (64,64), yt_eik)"
      ],
      "execution_count": 22,
      "outputs": []
    }
  ]
}