PDE-SGD-NoFit.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use this method to just get the value of the function at specific points without trying to fit the parameter in the entire domain.\n",
    "\n",
    "We follow the same process as in the other case, but we just ask to fit at one point at a time (still use SGD for the boundary, which has to be respected by the solution at any point.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir('Optimizers/Stochastic-Gradient-Descent/python/')\n",
    "import SGD as SGD \n",
    "os.chdir('../../../')\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('nbAgg')\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example\n",
    "\n",
    "Let's try solve the diffusion equation\n",
    "$$\n",
    "\\dfrac{\\partial u(x,t)}{\\partial t} - \\alpha^2 \\dfrac{\\partial^2 u(x,t)}{\\partial x^2}= 0\n",
    "$$\n",
    "for $x \\in [0,1]$ and $t>0$. In order to find a solution, we need\n",
    "1. initial condition that fix  $u(x,0)=\\sin( \\pi \\ x)$\n",
    "2. boundary conditions that fix $u(0,t)=u(1,t)=0$\n",
    "\n",
    "\n",
    "Notice that the solution is  \n",
    "$$\n",
    "u(x,t) =  \\sin( \\pi \\ x) \\ e^{-(\\pi \\ \\alpha )^2 t}\n",
    "$$\n",
    "\n",
    "For simplicity, we'll set $a=\\dfrac{1}{2}$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir('Optimizers/Stochastic-Gradient-Descent/python/')\n",
    "import SGD as SGD \n",
    "os.chdir('../../../')\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('nbAgg')\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As a sanity check, we choose\n",
    "$$\n",
    "u=w_0 \\ \\sin(w_1 \\ x) e^{ w_2 \\ t}\n",
    "$$\n",
    "and as to find $u$ at a few points for $t$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class modelFunc(SGD.modelBase):\n",
    "    def __init__(self,dimensions,w0,h=1e-3):\n",
    "        SGD.modelBase.__init__(self,dimensions,w0)\n",
    "\n",
    "        self.ds_ldx_i=0#derivative of the l^th wrt the i^th input\n",
    "        self.d2s_ldx_ij=0#second derivative of the l^th wrt the i^th and j^th inputs\n",
    "        self.d2s_ldx_ii=0#second derivative of the l^th wrt the i^th input\n",
    "        \n",
    "        self.h=h\n",
    "        \n",
    "\n",
    "    def __call__(self):\n",
    "        self.signal[0]=self.w[0]*np.sin(self.w[1]*self.input[0]) * np.exp(self.w[2]*self.input[1])\n",
    "#         self.signal[0]=self.w[0]+self.w[1]*self.input[0] \n",
    "#         self.signal[0]+=self.w[2]*self.input[1] \n",
    "#         self.signal[0]+=self.w[3]*self.input[0]*self.input[1]\n",
    "    \n",
    "    \n",
    "    \n",
    "    def derivative_ij(self,l,i,j):\n",
    "        heff1=self.h+np.abs(self.input[i])*self.h\n",
    "        heff2=self.h+np.abs(self.input[j])*self.h\n",
    "        \n",
    "        self.input[i]+=heff1\n",
    "        self.input[j]+=heff2\n",
    "        self()\n",
    "        f_ff=self.signal[l]\n",
    "        \n",
    "        self.input[i]-=2*heff1\n",
    "        self.input[j]-=2*heff2\n",
    "        self()\n",
    "        f_bb=self.signal[l]\n",
    "        \n",
    "        self.input[i]+=2*heff1\n",
    "        self()\n",
    "        f_fb=self.signal[l]\n",
    "        \n",
    "        self.input[i]-=2*heff1\n",
    "        self.input[j]+=2*heff2\n",
    "        self()\n",
    "        f_bf=self.signal[l]\n",
    "        \n",
    "        self.input[i]+=heff1\n",
    "        self.input[j]-=heff2\n",
    "        \n",
    "        \n",
    "        self.d2s_ldx_ij = (f_ff+f_bb-f_fb-f_bf)/(4*heff1*heff2)\n",
    "        \n",
    "        \n",
    "    def derivative_ii(self,l,i):\n",
    "        heff=self.h+np.abs(self.input[i])*self.h\n",
    "\n",
    "        self.input[i]-=heff\n",
    "        self()\n",
    "        f_b=self.signal[l]\n",
    "        \n",
    "        self.input[i]+=2*heff\n",
    "        self()\n",
    "        f_f=self.signal[l]\n",
    "        \n",
    "        self.input[i]-=heff\n",
    "        self()\n",
    "        f_0=self.signal[l]\n",
    "\n",
    "        self.d2s_ldx_ii = (f_f+f_b-2*f_0)/(heff**2)\n",
    "    \n",
    "    \n",
    "    def derivative_i(self,l,i):\n",
    "        heff=self.h+np.abs(self.input[i])*self.h\n",
    "        \n",
    "        self.input[i]-=heff\n",
    "        self()\n",
    "        f0=self.signal[l]\n",
    "\n",
    "        self.input[i]+=2*heff\n",
    "        self()\n",
    "        f1=self.signal[l]\n",
    "        \n",
    "        self.input[i]-=heff\n",
    "\n",
    "        self.ds_ldx_i = (f1-f0)/(2*heff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The boundary conditions\n",
    "\n",
    "Boundary conditions is given as in the class below. \n",
    "\n",
    "It would be convinient to define a function that returns a random point inside the boundary, which will be used to generate points that will be used to train the model.\n",
    "\n",
    "\n",
    "I also think that the contribution of the boundary conditions to the loss should be included here (as in the ODE case). \n",
    "\n",
    "In contrast to the ```Boundary``` class of the ODE example, we choose the points to be generated \"on the fly\" and not binge passed in the initializations, because in general this may be dificult to do (teh boundary condition can be very complicated). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Boundary:\n",
    "    def __init__(self,model,points):\n",
    "        self.model=model\n",
    "        self.points=points\n",
    "        self.NPoints=len(points)\n",
    "        self.currentPoint=0\n",
    "        \n",
    "        self.NConditions=3\n",
    "        self.insidePoint=[0 for _ in range(self.model.dimensions[0])]# this will hold the inputs inside the boundary\n",
    "        self.boundaryPoints=[[0 for _1 in range(self.model.dimensions[0])] for _2 in range(self.NConditions)]# this will hold the inputs on the boundary for all conditions       \n",
    "        self.boundary_diff=[[0 for _1 in range(self.model.dimensions[1])] for _2 in range(self.NConditions)]# this will hold the difference of the boundary conditions for all signals\n",
    "\n",
    "    #get a random point in the region of interest. In this case, its not a random point, it is the point self.points[self.currentPoint]\n",
    "    def randomPoint(self):        \n",
    "        self.insidePoint=self.points[self.currentPoint]\n",
    "            \n",
    "    def randomBoundaryPoints(self):\n",
    "        '''get a  point on the boundary for each boundary condition'''\n",
    "        \n",
    "        \n",
    "        x1=[0,np.random.rand()*2]\n",
    "        x2=[1,np.random.rand()*2]\n",
    "        x3=[np.random.rand(),0]\n",
    "        \n",
    "        self.boundaryPoints=[x1,x2,x3]\n",
    "\n",
    " \n",
    "    def randomBoundaryConditions(self):\n",
    "        tmp_p=self.model.input[:]\n",
    "        \n",
    "        self.randomBoundaryPoints()\n",
    "        \n",
    "        \n",
    "        for _c in range(self.NConditions):\n",
    "            #get difference for the all conditions\n",
    "            self.model.setInput(self.boundaryPoints[_c])\n",
    "            self.model()\n",
    "            self.boundary_diff[_c][0]=self.model.signal[0]\n",
    "\n",
    "        self.boundary_diff[2][0]+=-np.sin(np.pi*self.boundaryPoints[2][0])\n",
    "        \n",
    "        self.model.setInput(tmp_p)\n",
    "        \n",
    "    def randomBoundaryLoss(self):\n",
    "        '''average loss of the boundary conditions'''\n",
    "        self.randomBoundaryConditions()\n",
    "        avgBoundaryLoss=0\n",
    "        for _c in range(self.NConditions):\n",
    "            for l in range(self.model.dimensions[1]):\n",
    "                avgBoundaryLoss+=self.boundary_diff[_c][l]**2\n",
    "        \n",
    "        return avgBoundaryLoss/(1.*(self.NConditions*self.model.dimensions[1]))\n",
    "        \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The PDE\n",
    "\n",
    "Just define a class that holds everything that is relevant. Basically, you want have a place that can give you the loss easily.\n",
    "\n",
    "Overload the ```__call__``` function to return $\\rm lhs-rhs$. Define the loss and its derivarative over \n",
    "$\\{\\bf{w}\\}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DifferentialEquation:\n",
    "    def __init__(self,Model,Boundary,h=1e-2):\n",
    "        self.RHS=[lambda x:0]\n",
    "\n",
    "        self.model=Model\n",
    "        self.boundary=Boundary\n",
    "        \n",
    "        self.DE=[0 for _ in range(self.model.dimensions[1])]#this should be LHS-RHS, which we want to be 0 (see definition of __call__)\n",
    "        self.h=h\n",
    "        self.dQdw=0\n",
    "        \n",
    "    def __call__(self):\n",
    "        self.model.derivative_i(0,1)\n",
    "        self.model.derivative_ii(0,0)\n",
    "        \n",
    "        \n",
    "        dudt = self.model.ds_ldx_i\n",
    "        d2udx2 = self.model.d2s_ldx_ii\n",
    "        alpha=0.5\n",
    "        \n",
    "        self.DE[0]= dudt - alpha**2 * d2udx2 - self.RHS[0](self.model.input)\n",
    "    \n",
    "    \n",
    "    def randomDataPoint(self):\n",
    "        self.boundary.randomPoint()\n",
    "        self.model.setInput(self.boundary.insidePoint)\n",
    "\n",
    "        \n",
    "    def averageLoss(self):\n",
    "        '''get the average total loss'''\n",
    "        \n",
    "        avrgLoss=0\n",
    "        #loss from the PDE at the random point \n",
    "        self()\n",
    "        for l in range(self.model.dimensions[1]):\n",
    "            avrgLoss+=(self.DE[l]**2)/(1.*self.model.dimensions[1])\n",
    "        \n",
    "        #loss for the boundary conditions\n",
    "        avrgLoss+=self.boundary.randomBoundaryLoss()\n",
    "        \n",
    "        return avrgLoss/(2.)\n",
    "        \n",
    "\n",
    "    def grad(self,i):\n",
    "        '''Get the gradient of the averge loss wrt w[i] at current point + initial conditon'''\n",
    "        '''we should find a way to do it analytically (at least the dQds part)'''\n",
    "        \n",
    "        heff=self.h*np.abs(self.model.w[i])+self.h\n",
    "\n",
    "        self.model.w[i]-=heff\n",
    "        Q0=self.averageLoss()\n",
    "\n",
    "        self.model.w[i]+=2*heff\n",
    "        Q1=self.averageLoss()\n",
    "\n",
    "        self.dQdw=(Q1-Q0)/(2.*heff)\n",
    "\n",
    "        self.model.w[i]-=heff\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "model=modelFunc([2,1], np.random.rand(3)*2-1 )\n",
    "\n",
    "#fit only for one point\n",
    "S=Boundary(model,[[0.5,0.5]])\n",
    "\n",
    "PDE=DifferentialEquation(model,S)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sgd=SGD.VanillaSGD(PDE,alpha=1e-2)\n",
    "# sgd=SGD.RMSpropSGD(PDE,gamma=1-1e-2,epsilon=1e-5,alpha=1e-1)\n",
    "# sgd=SGD.AdaDeltaSGD(PDE,gamma=0.99,epsilon=1e-5,alpha=1)\n",
    "# sgd=SGD.AdamSGD(PDE,beta_m=0.9,beta_v=0.999,epsilon=1e-8,alpha=1e-1)\n",
    "sgd=SGD.AdaMaxSGD(PDE,beta_m=0.9,beta_v=0.999,epsilon=1e-8,alpha=1e-1)\n",
    "# sgd=SGD.NAdamSGD(PDE,beta_m=0.9,beta_v=0.999,epsilon=1e-8,alpha=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(38.91981101036072, 50001)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "time0=time.time()\n",
    "\n",
    "sgd.run(abs_tol=1e-3, rel_tol=1e-3, step_break=5000,max_step=50000)\n",
    "\n",
    "time.time()-time0,len(sgd.steps)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The loss for a lot of random points is small!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.3031791466089078e-06"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "maxLoss=0\n",
    "for i in range(1000):\n",
    "    PDE.randomDataPoint()\n",
    "    loss=PDE.averageLoss()\n",
    "    if loss>maxLoss:\n",
    "        maxLoss=loss\n",
    "maxLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "we expect: 0.291212933214021 we got:  0.2925959253762391\n"
     ]
    }
   ],
   "source": [
    "x,t=S.points[0]\n",
    "model.setInput(S.points[0])\n",
    "model()\n",
    "\n",
    "\n",
    "print('we expect:', np.sin(np.pi*x)*np.exp(-(np.pi/2.)**2 * t) ,'we got: ',model.signal[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}