Complex networks solve NLS.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 10480,
     "status": "ok",
     "timestamp": 1620063731706,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "KLKYhTgqt5ia",
    "outputId": "f7b362ce-dd34-4950-9dcb-a68e70f5d473"
   },
   "outputs": [],
   "source": [
    "# For enabling running in google colab\n",
    "\n",
    "# !pip3 install pyDOE\n",
    "# !pip3 install complexPyTorch\n",
    "# !pip3 install cplxmodule\n",
    "\n",
    "# import sys; sys.path.insert(1, '/content/drive/MyDrive/Colab Notebooks/')\n",
    "# import os\n",
    "# from glob import glob\n",
    "# from google.colab import drive\n",
    "# drive.mount('/content/drive')\n",
    "# PATH = '/content/drive/MyDrive/Colab Notebooks/'\n",
    "# DATA_PATH = PATH + 'data/NLS.mat'\n",
    "# from helper.utils import *\n",
    "# from helper.models import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "executionInfo": {
     "elapsed": 394,
     "status": "ok",
     "timestamp": 1620067508294,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "J1GaVK_KvnSV"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.io as io\n",
    "from pyDOE import lhs\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from complexPyTorch.complexLayers import ComplexLinear\n",
    "\n",
    "import cplxmodule\n",
    "# complex valued tensor class\n",
    "from cplxmodule import cplx\n",
    "from cplxmodule.nn import RealToCplx, CplxToReal, CplxSequential, CplxToCplx\n",
    "from cplxmodule.nn import CplxLinear, CplxModReLU, CplxAdaptiveModReLU\n",
    "\n",
    "from utils import *\n",
    "from models import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "executionInfo": {
     "elapsed": 633,
     "status": "ok",
     "timestamp": 1620067508625,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "O1AhOKnfvA9x"
   },
   "outputs": [],
   "source": [
    "# torch device\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(\"You're running on\", device)\n",
    "\n",
    "# Doman bounds\n",
    "lb = np.array([-5.0, 0.0])\n",
    "ub = np.array([5.0, np.pi/2])\n",
    "\n",
    "N0 = 50\n",
    "N_b = 50\n",
    "N_f = 20000\n",
    "\n",
    "DATA_PATH = './experimental_data/NLS.mat'\n",
    "data = io.loadmat(DATA_PATH)\n",
    "\n",
    "t = data['tt'].flatten()[:,None]\n",
    "x = data['x'].flatten()[:,None]\n",
    "Exact = data['uu']\n",
    "Exact_u = np.real(Exact)\n",
    "Exact_v = np.imag(Exact)\n",
    "\n",
    "X, T = np.meshgrid(x,t)\n",
    "\n",
    "X_star = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "Exact_h = np.sqrt(Exact_u**2 + Exact_v**2)\n",
    "h_star = Exact_h.T.flatten()[:,None]\n",
    "\n",
    "idx_x = np.random.choice(x.shape[0], N0, replace=False)\n",
    "x0 = x[idx_x, :]\n",
    "Exact_u0 = np.real(Exact)[idx_x, 0:1]\n",
    "Exact_v0 = np.imag(Exact)[idx_x, 0:1]\n",
    "h0 = np.hstack((Exact_u0, Exact_v0))\n",
    "\n",
    "idx_t = np.random.choice(t.shape[0], N_b, replace=False)\n",
    "tb = t[idx_t,:]\n",
    "\n",
    "X_f = lb + (ub-lb)*lhs(2, N_f) # -> PDE Loss\n",
    "X0 = np.concatenate((x0, 0*x0), 1) # (x0, 0)\n",
    "X_lb = np.concatenate((0*tb + lb[0], tb), 1) # (lb[0], tb)\n",
    "X_ub = np.concatenate((0*tb + ub[0], tb), 1) # (ub[0], tb)\n",
    "\n",
    "# PDE Loss\n",
    "X_f = to_tensor(X_f, True).to(device)\n",
    "\n",
    "# NN(X0) approx. h0\n",
    "X0 = to_tensor(X0, False).to(device)\n",
    "h0 = to_tensor(h0, False).to(device)\n",
    "\n",
    "# NN(X_lb) approx NN(X_ub)\n",
    "# NN(X_lb)_x approx NN(X_ub)_x, _x = diff wrt the first dimension\n",
    "X_lb = to_tensor(X_lb, True).to(device)\n",
    "X_ub = to_tensor(X_ub, True).to(device)\n",
    "\n",
    "lb = to_tensor(lb, False).to(device)\n",
    "ub = to_tensor(ub, False).to(device)\n",
    "\n",
    "X_star = to_tensor(X_star, False).to(device)\n",
    "# h_star = to_complex_tensor(h_star, False).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 545,
     "status": "ok",
     "timestamp": 1620067508626,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "ed_u5FKLz8lR",
    "outputId": "91d4791e-3978-4d6e-d349-475226484afe"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.9/site-packages/torch/nn/modules/container.py:587: UserWarning: Setting attributes on ParameterDict is not supported.\n",
      "  warnings.warn(\"Setting attributes on ParameterDict is not supported.\")\n"
     ]
    }
   ],
   "source": [
    "class Act(cnn.CplxToCplx):\n",
    "    def __init__(self,):\n",
    "        super(Act, self).__init__()\n",
    "        pass\n",
    "    def forward(self, inp):\n",
    "        return torch.tanh(cplx2tensor(inp))\n",
    "    \n",
    "class ImaginaryDimensionAdder(nn.Module):\n",
    "    def __init__(self,):\n",
    "        super(ImaginaryDimensionAdder, self).__init__(); pass\n",
    "    def forward(self, real_tensor):\n",
    "        added = cat(real_tensor[:, 0:1], torch.zeros(real_tensor.shape[0], 1))\n",
    "        for i in range(1, real_tensor.shape[1]):\n",
    "            added = cat(added, real_tensor[:, i:i+1], torch.zeros(real_tensor.shape[0], 1))\n",
    "        return added\n",
    "        \n",
    "# act = Act\n",
    "# act = CplxAdaptiveModReLU\n",
    "inp_dimension = 2\n",
    "act = CplxToCplx[torch.tanh]\n",
    "complex_model = CplxSequential(\n",
    "    CplxLinear(inp_dimension, 100, bias=True),\n",
    "    act(),\n",
    "    CplxLinear(100, 100, bias=True),\n",
    "    act(),\n",
    "    CplxLinear(100, 100, bias=True),\n",
    "    act(),\n",
    "    CplxLinear(100, 100, bias=True),\n",
    "    act(),\n",
    "    CplxLinear(100, 1, bias=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "executionInfo": {
     "elapsed": 471,
     "status": "ok",
     "timestamp": 1620067508628,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "9_baKQ8_XgeM"
   },
   "outputs": [],
   "source": [
    "m = nn.Sequential(\n",
    "        ImaginaryDimensionAdder(),\n",
    "        RealToCplx(),\n",
    "        complex_model,\n",
    "#         CplxToReal(),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "executionInfo": {
     "elapsed": 657,
     "status": "ok",
     "timestamp": 1620067508893,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "DRKNm8YHvHvF"
   },
   "outputs": [],
   "source": [
    "class ComplexPhysicsInformedNN(nn.Module):\n",
    "    def __init__(self, model, lb, ub, scale=False):\n",
    "        super(ComplexPhysicsInformedNN, self).__init__()\n",
    "        self.model = model\n",
    "        self.lb = lb\n",
    "        self.ub = ub\n",
    "        self.scale = scale\n",
    "    \n",
    "    def forward(self, X):\n",
    "        if self.scale: \n",
    "            return self.model(self.neural_net_scale(X))\n",
    "        return self.model(X)\n",
    "\n",
    "    def predict(self, X_test):\n",
    "        return CplxToReal()(self.forward(self.preprocess(*dimension_slicing(X_test))))\n",
    "    \n",
    "    def neural_net_scale(self, inp):\n",
    "        return (2.0*(inp-self.lb)/(self.ub-self.lb))-1.0\n",
    "\n",
    "    def preprocess(self, spatial, time):\n",
    "        return cat(spatial, time)\n",
    "    \n",
    "    def loss(self, X_f, X0, h0, X_lb, X_ub):\n",
    "        loss = self.net_f(*dimension_slicing(X_f))\n",
    "        h0_pred = self.predict(X0); u0 = h0_pred[:, 0:1]; v0 = h0_pred[:, 1:2]\n",
    "        loss += F.mse_loss(u0, h0[:, 0:1])+F.mse_loss(v0, h0[:, 1:2])\n",
    "        u_lb, v_lb, u_lb_x, v_lb_x = self.net_h(*dimension_slicing(X_lb))\n",
    "        u_ub, v_ub, u_ub_x, v_ub_x = self.net_h(*dimension_slicing(X_ub))\n",
    "        loss += F.mse_loss(u_lb, u_ub)\n",
    "        loss += F.mse_loss(v_lb, v_ub)\n",
    "        loss += F.mse_loss(u_lb_x, u_ub_x)\n",
    "        loss += F.mse_loss(v_lb_x, v_ub_x)\n",
    "        return loss\n",
    "    \n",
    "    def net_h(self, x, t):\n",
    "        X = cat(x, t)\n",
    "        h = self.forward(X)\n",
    "        u = h.real\n",
    "        v = h.imag\n",
    "        return u, v, self.diff(u, x), self.diff(v, x)\n",
    "    \n",
    "    def net_f(self, x, t):\n",
    "        u, v, u_x, v_x = self.net_h(x, t)\n",
    "        u_t, v_t = self.diff(u, t), self.diff(v, t)\n",
    "        u_xx, v_xx = self.diff(u_x, x), self.diff(v_x, x)\n",
    "        f_u = u_t + 0.5*v_xx + (u**2 + v**2)*v\n",
    "        f_v = v_t - 0.5*u_xx - (u**2 + v**2)*u\n",
    "        return (f_u**2).mean()+(f_v**2).mean()\n",
    "\n",
    "    def diff(self, func, inp):\n",
    "        return grad(func, inp, create_graph=True, retain_graph=True, grad_outputs=torch.ones(func.shape, dtype=func.dtype).to(device))[0]\n",
    "\n",
    "    def complex_mse(self, v1, v2):\n",
    "        assert v1.shape == v2.shape\n",
    "        assert v1.shape[1] == 1\n",
    "        return F.mse_loss(v1.real, v2.real)+F.mse_loss(v2.imag, v2.imag)\n",
    "\n",
    "    def add_imag_dim(self, v1):\n",
    "        z = torch.zeros(v1.shape).requires_grad_(False).to(device)\n",
    "        return torch.complex(v1, z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "executionInfo": {
     "elapsed": 568,
     "status": "ok",
     "timestamp": 1620067508894,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "5W_ojN1MaNrE"
   },
   "outputs": [],
   "source": [
    "layers = [2, 100, 100, 100, 100, 1]\n",
    "model = TorchComplexMLP(layers)\n",
    "cpinn = ComplexPhysicsInformedNN(model=m, lb=lb, ub=ub, scale=False).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "executionInfo": {
     "elapsed": 489,
     "status": "ok",
     "timestamp": 1620067508894,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "xrkxRWKTl_Ci"
   },
   "outputs": [],
   "source": [
    "# def add_imag_dim(v1):\n",
    "#   z = torch.zeros(v1.shape).requires_grad_(False).to(device)\n",
    "#   return torch.complex(v1, z)\n",
    "\n",
    "# def tdiff(func, inp):\n",
    "#   return grad(func, inp, create_graph=True, retain_graph=True, grad_outputs=torch.ones(func.shape, dtype=func.dtype).to(device))[0]\n",
    "\n",
    "# x, t = dimension_slicing(X_f)\n",
    "# a, b = add_imag_dim(x), add_imag_dim(t)\n",
    "# func = model(cat(a, b))\n",
    "# tdiff(func.real, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 202033,
     "status": "ok",
     "timestamp": 1620067712730,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "Zr6-TbbBaUYq",
    "outputId": "47bb9b61-60a4-4403-ded0-83965f16e935"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1st Stage optimization\n",
      "4.6230244636535645\n",
      "Test score 0.650792658358732\n",
      "1.516767144203186\n",
      "Test score 0.637255628148042\n",
      "0.4470001757144928\n",
      "Test score 0.38790888955214015\n",
      "0.4434736967086792\n",
      "Test score 0.3515484679506801\n",
      "0.3543202877044678\n",
      "Test score 0.3445935947824777\n",
      "0.30723142623901367\n",
      "Test score 0.32058153904209896\n",
      "0.2450341433286667\n",
      "Test score 0.2691525724699435\n",
      "0.19877596199512482\n",
      "Test score 0.23350022649095842\n",
      "0.18728980422019958\n",
      "Test score 0.2102306620028725\n",
      "0.14753946661949158\n",
      "Test score 0.1845154373208915\n",
      "0.1359558403491974\n",
      "Test score 0.17411084423822376\n",
      "0.12862499058246613\n",
      "Test score 0.17026527523281149\n",
      "0.12386464327573776\n",
      "Test score 0.16799392200253402\n",
      "0.13405485451221466\n",
      "Test score 0.16969658058343734\n",
      "0.1481952965259552\n",
      "Test score 0.16786675706378817\n",
      "0.11653044819831848\n",
      "Test score 0.16724235252911387\n",
      "0.11698071658611298\n",
      "Test score 0.16533738237580153\n",
      "0.11031896620988846\n",
      "Test score 0.16287644742867283\n",
      "0.10831142216920853\n",
      "Test score 0.1622804010094372\n",
      "0.10557455569505692\n",
      "Test score 0.1606056187803276\n",
      "0.1276548206806183\n",
      "Test score 0.16314274478808855\n",
      "0.103033147752285\n",
      "Test score 0.16264822361446335\n",
      "0.10158327966928482\n",
      "Test score 0.15905516659336333\n",
      "0.09989886730909348\n",
      "Test score 0.15739608152333673\n",
      "0.09618782997131348\n",
      "Test score 0.15541197325570294\n",
      "0.09193132817745209\n",
      "Test score 0.15414232228113686\n",
      "0.11196734756231308\n",
      "Test score 0.15695345644698422\n",
      "0.10365200787782669\n",
      "Test score 0.1620163706627989\n",
      "0.0968431830406189\n",
      "Test score 0.15262867127768642\n",
      "0.09354035556316376\n",
      "Test score 0.15222195645703968\n",
      "0.08716655522584915\n",
      "Test score 0.1493405413776157\n",
      "0.08446399122476578\n",
      "Test score 0.1476583663071633\n",
      "0.08183867484331131\n",
      "Test score 0.14556720588114705\n",
      "0.08032017946243286\n",
      "Test score 0.14461033768140052\n",
      "0.07900132983922958\n",
      "Test score 0.14315838141482184\n",
      "0.07775754481554031\n",
      "Test score 0.14204570918057255\n",
      "0.076674684882164\n",
      "Test score 0.1409275490059915\n",
      "0.07565320283174515\n",
      "Test score 0.13995075594271392\n",
      "0.07466904073953629\n",
      "Test score 0.13895443277527736\n",
      "0.07373373955488205\n",
      "Test score 0.13808793554951077\n",
      "0.3527108132839203\n",
      "Test score 0.18384162073125299\n",
      "0.18248170614242554\n",
      "Test score 0.13899716850342972\n",
      "0.07591311633586884\n",
      "Test score 0.14581447010646145\n",
      "0.08222564309835434\n",
      "Test score 0.1401058854466491\n",
      "0.07758887857198715\n",
      "Test score 0.13788637092963912\n",
      "0.07336293905973434\n",
      "Test score 0.13853343154235118\n",
      "0.07204492390155792\n",
      "Test score 0.13612300650798398\n",
      "0.07115643471479416\n",
      "Test score 0.13551444022452355\n",
      "0.0702342689037323\n",
      "Test score 0.13475635306172418\n",
      "0.06941799074411392\n",
      "Test score 0.13401769965084737\n",
      "0.06868080794811249\n",
      "Test score 0.133230093367155\n",
      "0.06798055022954941\n",
      "Test score 0.13254051940850275\n",
      "0.06731913238763809\n",
      "Test score 0.13192660201291748\n",
      "0.06669590622186661\n",
      "Test score 0.13132126900936622\n",
      "0.06611087173223495\n",
      "Test score 0.13076527299868915\n",
      "0.06556373089551926\n",
      "Test score 0.13022625777936914\n",
      "0.06505285203456879\n",
      "Test score 0.1297211719119867\n",
      "0.06457625329494476\n",
      "Test score 0.1292409164234977\n",
      "0.06413166970014572\n",
      "Test score 0.12879282216178506\n",
      "0.06371670961380005\n",
      "Test score 0.12837508482903043\n",
      "0.06332904100418091\n",
      "Test score 0.1279875490362227\n",
      "0.06299221515655518\n",
      "Test score 0.127653018913616\n",
      "0.17177069187164307\n",
      "Test score 0.14594856265247305\n",
      "0.1326584368944168\n",
      "Test score 0.13129976822634676\n",
      "0.06794916838407516\n",
      "Test score 0.13371335841431398\n",
      "0.0741088017821312\n",
      "Test score 0.1347683734025356\n",
      "0.06958914548158646\n",
      "Test score 0.13188618249814602\n",
      "0.06451507657766342\n",
      "Test score 0.12917376617404225\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-9-4fa103212c47>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      7\u001b[0m     \u001b[0mcpinn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrain\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      8\u001b[0m     \u001b[0moptimizer1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 9\u001b[0;31m     \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcpinn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_f\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_lb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_ub\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     10\u001b[0m     \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0moptimizer1\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-6-3353ffcb1a0f>\u001b[0m in \u001b[0;36mloss\u001b[0;34m(self, X_f, X0, h0, X_lb, X_ub)\u001b[0m\n\u001b[1;32m     23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     24\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_f\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_lb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_ub\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 25\u001b[0;31m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnet_f\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mdimension_slicing\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_f\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     26\u001b[0m         \u001b[0mh0_pred\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpredict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mu0\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mh0_pred\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m;\u001b[0m \u001b[0mv0\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mh0_pred\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     27\u001b[0m         \u001b[0mloss\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmse_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mu0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh0\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m+\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmse_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh0\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-6-3353ffcb1a0f>\u001b[0m in \u001b[0;36mnet_f\u001b[0;34m(self, x, t)\u001b[0m\n\u001b[1;32m     44\u001b[0m         \u001b[0mu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mu_x\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv_x\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnet_h\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     45\u001b[0m         \u001b[0mu_t\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv_t\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiff\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiff\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 46\u001b[0;31m         \u001b[0mu_xx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv_xx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiff\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mu_x\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdiff\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mv_x\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     47\u001b[0m         \u001b[0mf_u\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mu_t\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;36m0.5\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mv_xx\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mu\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;36m2\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mv\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     48\u001b[0m         \u001b[0mf_v\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mv_t\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;36m0.5\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mu_xx\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mu\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;36m2\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mu\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-6-3353ffcb1a0f>\u001b[0m in \u001b[0;36mdiff\u001b[0;34m(self, func, inp)\u001b[0m\n\u001b[1;32m     50\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     51\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mdiff\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 52\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mgrad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_outputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mones\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     53\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     54\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mcomplex_mse\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.9/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mgrad\u001b[0;34m(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)\u001b[0m\n\u001b[1;32m    224\u001b[0m         \u001b[0mretain_graph\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    225\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 226\u001b[0;31m     return Variable._execution_engine.run_backward(\n\u001b[0m\u001b[1;32m    227\u001b[0m         \u001b[0moutputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_outputs_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    228\u001b[0m         inputs, allow_unused, accumulate_grad=False)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "best_state_dict = None; best_loss = 1000.0\n",
    "learning_rate1, learning_rate2 = 1e-3, 1e-1\n",
    "optimizer1 = torch.optim.Adam(cpinn.parameters(), lr=learning_rate1)\n",
    "epochs1 = 1000\n",
    "print(\"1st Stage optimization\")\n",
    "for i in range(epochs1):\n",
    "    cpinn.train()\n",
    "    optimizer1.zero_grad()\n",
    "    loss = cpinn.loss(X_f, X0, h0, X_lb, X_ub)\n",
    "    loss.backward(retain_graph=False)\n",
    "    optimizer1.step()\n",
    "\n",
    "    track = loss.item()\n",
    "    if track < best_loss:\n",
    "      best_loss = track\n",
    "      best_state_dict = cpinn.state_dict()\n",
    "\n",
    "    if i%10==0: \n",
    "      cpinn.eval()\n",
    "      print(track)\n",
    "      h_star_pred = cpinn.predict(X_star)\n",
    "      h_star_pred = torch.sqrt(h_star_pred[:, 0:1]**2 + h_star_pred[:, 1:2]**2)\n",
    "      h_star_pred = (h_star_pred.detach().cpu().numpy())\n",
    "      print('Test score', mean_squared_error(h_star_pred, h_star))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "executionInfo": {
     "elapsed": 1808591,
     "status": "error",
     "timestamp": 1620070327200,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "7MZxRsYRaLog",
    "outputId": "3b2e18d1-36e6-41f4-c906-34b5f6fc7111"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2st Stage optimization\n",
      "0.0005419803783297539\n",
      "Test score 0.0018069956771203238\n",
      "0.00015163954230956733\n",
      "Test score 0.00030664691319947805\n",
      "7.657348032807931e-05\n",
      "Test score 0.00011338522830886114\n",
      "4.5973127271281555e-05\n",
      "Test score 5.549353836509127e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n",
      "3.10495997837279e-05\n",
      "Test score 1.8905030561149737e-05\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "ignored",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-27-7e72317e6522>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     15\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_grad\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     16\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 17\u001b[0;31m     \u001b[0moptimizer2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     18\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     19\u001b[0m     \u001b[0mtrack\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mitem\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/optim/optimizer.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     87\u001b[0m                 \u001b[0mprofile_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m\"Optimizer.step#{}.step\"\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mformat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__name__\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     88\u001b[0m                 \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprofiler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrecord_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mprofile_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 89\u001b[0;31m                     \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     90\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mwrapper\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     91\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/optim/lbfgs.py\u001b[0m in \u001b[0;36mstep\u001b[0;34m(self, closure)\u001b[0m\n\u001b[1;32m    424\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    425\u001b[0m                     loss, flat_grad, t, ls_func_evals = _strong_wolfe(\n\u001b[0;32m--> 426\u001b[0;31m                         obj_func, x_init, t, d, loss, flat_grad, gtd)\n\u001b[0m\u001b[1;32m    427\u001b[0m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_add_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    428\u001b[0m                 \u001b[0mopt_cond\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflat_grad\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mabs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmax\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m<=\u001b[0m \u001b[0mtolerance_grad\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/optim/lbfgs.py\u001b[0m in \u001b[0;36m_strong_wolfe\u001b[0;34m(obj_func, x, t, d, f, g, gtd, c1, c2, tolerance_change, max_ls)\u001b[0m\n\u001b[1;32m    145\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    146\u001b[0m         \u001b[0;31m# Evaluate new point\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m         \u001b[0mf_new\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mg_new\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mobj_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    148\u001b[0m         \u001b[0mls_func_evals\u001b[0m \u001b[0;34m+=\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m         \u001b[0mgtd_new\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mg_new\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/optim/lbfgs.py\u001b[0m in \u001b[0;36mobj_func\u001b[0;34m(x, t, d)\u001b[0m\n\u001b[1;32m    421\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    422\u001b[0m                     \u001b[0;32mdef\u001b[0m \u001b[0mobj_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 423\u001b[0;31m                         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_directional_evaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    424\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    425\u001b[0m                     loss, flat_grad, t, ls_func_evals = _strong_wolfe(\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/optim/lbfgs.py\u001b[0m in \u001b[0;36m_directional_evaluate\u001b[0;34m(self, closure, x, t, d)\u001b[0m\n\u001b[1;32m    275\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_directional_evaluate\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mclosure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    276\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_add_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0md\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 277\u001b[0;31m         \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    278\u001b[0m         \u001b[0mflat_grad\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_gather_flat_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    279\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_set_param\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py\u001b[0m in \u001b[0;36mdecorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     25\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     26\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__class__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 27\u001b[0;31m                 \u001b[0;32mreturn\u001b[0m \u001b[0mfunc\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     28\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mcast\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mF\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdecorate_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-27-7e72317e6522>\u001b[0m in \u001b[0;36mclosure\u001b[0;34m()\u001b[0m\n\u001b[1;32m     13\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_grad_enabled\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0moptimizer2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzero_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     14\u001b[0m       \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcpinn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mX_f\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mh0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_lb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mX_ub\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 15\u001b[0;31m       \u001b[0;32mif\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrequires_grad\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mretain_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     16\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mloss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     17\u001b[0m     \u001b[0moptimizer2\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mclosure\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    243\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    244\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 245\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    246\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    247\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.7/dist-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    145\u001b[0m     Variable._execution_engine.run_backward(\n\u001b[1;32m    146\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 147\u001b[0;31m         allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag\n\u001b[0m\u001b[1;32m    148\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    149\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# if best_state_dict is not None: cpinn.load_state_dict(best_state_dict)\n",
    "\n",
    "optimizer2 = torch.optim.LBFGS(cpinn.parameters(), lr=learning_rate2, \n",
    "                               max_iter=200, max_eval=200, \n",
    "                               history_size=120, line_search_fn='strong_wolfe')\n",
    "\n",
    "epochs2 = 50000\n",
    "cpinn.train()\n",
    "print(\"2st Stage optimization\")\n",
    "\n",
    "for i in range(epochs1):\n",
    "    def closure():\n",
    "      if torch.is_grad_enabled(): optimizer2.zero_grad()\n",
    "      loss = cpinn.loss(X_f, X0, h0, X_lb, X_ub)\n",
    "      if loss.requires_grad: loss.backward(retain_graph=False)\n",
    "      return loss\n",
    "    optimizer2.step(closure)\n",
    "\n",
    "    track = closure().item()\n",
    "    if track < best_loss:\n",
    "      best_loss = track\n",
    "      best_state_dict = cpinn.state_dict()\n",
    "\n",
    "    if i%10==0: \n",
    "      cpinn.eval()\n",
    "      print(track)\n",
    "      h_star_pred = cpinn.predict(X_star)\n",
    "      h_star_pred = torch.sqrt(h_star_pred[:, 0:1]**2 + h_star_pred[:, 1:2]**2)\n",
    "      h_star_pred = (h_star_pred.detach().cpu().numpy())\n",
    "      print('Test score', mean_squared_error(h_star_pred, h_star))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 553,
     "status": "ok",
     "timestamp": 1620070371554,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "km1jgf1ngR7I",
    "outputId": "ca866b2c-eb1b-4263-a047-90eaf710c07b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test score 1.8905041546390686e-05\n"
     ]
    }
   ],
   "source": [
    "h_star_pred = cpinn.predict(X_star)\n",
    "h_star_pred = torch.sqrt(h_star_pred[:, 0:1]**2 + h_star_pred[:, 1:2]**2)\n",
    "h_star_pred = (h_star_pred.detach().cpu().numpy())\n",
    "\n",
    "print('Test score', mean_squared_error(h_star_pred, h_star))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 433,
     "status": "ok",
     "timestamp": 1620070397546,
     "user": {
      "displayName": "Pongpisit Thanasutives",
      "photoUrl": "",
      "userId": "18350641563056089823"
     },
     "user_tz": -540
    },
    "id": "LeL2a2O7uA5v",
    "outputId": "cbdfc2c7-95c9-4f45-9f3b-160469f7c431"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0048614245690865995"
      ]
     },
     "execution_count": 30,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.linalg.norm(h_star-h_star_pred,2)/np.linalg.norm(h_star,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iOwW9urkub75"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "authorship_tag": "ABX9TyOLQN9Yffqns1SB6tPLxz7R",
   "collapsed_sections": [],
   "name": "Complex networks solve NLS.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}