advection_diffusion_simplified.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PDE 3 - Advection-Diffusion\n",
    "\n",
    "\n",
    "#### Problem Setup\n",
    "\n",
    "\\begin{equation} \n",
    "0 = au_x + bu_y + cu_{xx} + du_{yy} - u_t \\; \\; \\; \\; \\; \\; \\; \\; \\; (*)\n",
    "\\end{equation}\n",
    "\n",
    "For the generation of our training data we use the same method as in PDE-Net in Tensorflow (for details please refer to that chapter in the paper) with:\n",
    "\n",
    "$a = b = 2$ <br> $c = d = 0.5$ <br>\n",
    "\n",
    "$u: \\mathbb{R}^3 \\rightarrow \\mathbb{R}$ is some solution to $(*)$. <br>\n",
    "$f: \\mathbb{R}^3 \\rightarrow \\mathbb{R}, \\;f(x,y,t) = 0$ <br>\n",
    "$X_i := (x_i, y_i, t_i) \\in [0,1] \\times [0,1] \\times [0, 0.135] \\subseteq \\mathbb{R}^3$ for $i \\in \\{1, \\dotsc, n\\}$\n",
    "\n",
    "and our known function values will be $\\{u(X_i), f(X_i)\\}_{i \\in \\{1, \\dotsc, n\\}}$.\n",
    "\n",
    "We assume that $u$ can be represented as a Gaussian process with SE kernel.\n",
    "\n",
    "$u \\sim \\mathcal{GP}(0, k_{uu}(X_i, X_j; \\theta))$, where $\\theta = \\{\\sigma, l_x, l_y, l_t\\}$.\n",
    "\n",
    "Set the linear operator to:\n",
    "\n",
    "$\\mathcal{L}_X^{\\phi} := a\\partial_x + b\\partial_y + c\\partial_{x,x} + d\\partial_{y,y} - \\partial_t$\n",
    "\n",
    "so that\n",
    "\n",
    "$\\mathcal{L}_X^{\\phi} u = f = 0$\n",
    "\n",
    "Problem at hand: Estimate $\\phi:=\\{a\\}$ with $\\{b,c,d\\} = \\{2, 0.5, 0.5\\}$ fixed.\n",
    "\n",
    "\n",
    "#### Step 1: Simulate data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import numpy as np\n",
    "import sympy as sp\n",
    "import numpy.fft as fft\n",
    "from scipy.linalg import solve_triangular\n",
    "import scipy.optimize as opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Global variables: x, y, t, n, y_u, y_f, s\n",
    "\n",
    "# Number of data samples\n",
    "n = 80\n",
    "\n",
    "# Noise of our data:\n",
    "s = 1e-7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "options = {'mesh_size': [250, 250], # How large is the (regular) 2D-grid of function values for \n",
    "                                    # a fixed t. Keep mesh_size[0] = mesh_size[1]\n",
    "           'layers': 9,             # Layers of the NN. Also counting the initial layer!\n",
    "           'dt': 0.0015,            # Time discretization. We step dt*(layers - 1) forward in time.\n",
    "           'batch_size': 1,         # We take a batch of sub-grids in space\n",
    "           'noise_level': 0.0,      # Can add some noise to the data (not taken 1 to 1, gets \n",
    "                                    # multiplied by stddev)\n",
    "           'downsample_by': 1,      # Size of sub-grids (in space) * downsample_by = mesh_size\n",
    "           }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initgen(mesh_size, freq=3, boundary='Periodic'):\n",
    "    \"\"\"\n",
    "    Returns function values for t=0 on a regular grid of size 'mesh_size' in [0, 2*pi]x[0, 2*pi] \n",
    "    as a matrix\n",
    "    \"\"\"\n",
    "    # Default: (mesh_size, freq, boundary) = ([250, 250], 4, 'Periodic')\n",
    "    if np.iterable(freq):\n",
    "        return freq\n",
    "    # 250x250 normally distributed variables IFFTed and FFTed:\n",
    "    x = _initgen_periodic(mesh_size, freq=freq)\n",
    "    x = x * 100\n",
    "    if boundary.upper() == 'DIRICHLET':\n",
    "        dim = x.ndim\n",
    "        for i in range(dim):\n",
    "            y = np.arange(mesh_size[i]) / mesh_size[i]\n",
    "            y = y * (1 - y)\n",
    "            s = np.ones(dim, dtype=np.int32)\n",
    "            s[i] = mesh_size[i]\n",
    "            y = np.reshape(y, s)\n",
    "            x = x * y\n",
    "        x = x[[slice(1, None), ] * dim]\n",
    "        x = x * 16\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _initgen_periodic(mesh_size, freq=3):\n",
    "    # np.random.seed(50)\n",
    "    # Default: (mesh_size, freq) = ([250, 250], 4)\n",
    "    dim = len(mesh_size)\n",
    "    # Default: 250x250-matrix of normally distributed variables\n",
    "    x = np.random.randn(*mesh_size)\n",
    "    coe = fft.ifftn(x)\n",
    "    # set frequency of generated initial value\n",
    "    # Array of random ints in [freq, 2*freq - 1]\n",
    "    freqs = np.random.randint(freq, 2 * freq, size=[dim, ])\n",
    "    # freqs = [10,10]\n",
    "    for i in range(dim):\n",
    "        perm = np.arange(dim, dtype=np.int32)\n",
    "        perm[i] = 0\n",
    "        perm[0] = i\n",
    "        # Permutes for i = 1 and does nothing for i = 0.\n",
    "        coe = coe.transpose(*perm)\n",
    "        coe[freqs[i] + 1:-freqs[i]] = 0\n",
    "        coe = coe.transpose(*perm)\n",
    "    x = fft.fftn(coe)\n",
    "    assert np.linalg.norm(x.imag) < 1e-8\n",
    "    x = x.real\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_input_2(input, pad_by):\n",
    "    \"\"\"\n",
    "    We increase the size of input for all j by pad_by on each side of the matrix\n",
    "    by inserting values from the opposite side\n",
    "    \"\"\"\n",
    "    mesh_size = input.shape[0]\n",
    "\n",
    "    B = np.eye(mesh_size, dtype=np.float32)\n",
    "    for i in range(pad_by):\n",
    "        a = np.zeros(mesh_size, dtype=np.float32)\n",
    "        a[mesh_size - i - 1] = 1\n",
    "        B = np.concatenate(([a], B), axis=0)\n",
    "    for i in range(pad_by):\n",
    "        a = np.zeros(mesh_size, dtype=np.float32)\n",
    "        a[i] = 1\n",
    "        B = np.concatenate((B, [a]), axis=0)\n",
    "\n",
    "    return B @ input @ B.T\n",
    "\n",
    "def downsample(sample, scale):\n",
    "    \"\"\"\n",
    "    Returns a regular somewhat random sub-grid of sample, whose size is reduced by a \n",
    "    factor of 'scale'.\n",
    "    \"\"\"\n",
    "    # np.random.seed(50)\n",
    "    idx1 = slice(np.random.randint(scale), None, scale)\n",
    "    idx2 = slice(np.random.randint(scale), None, scale)\n",
    "    # idx1 = slice(1, None, scale)\n",
    "    # idx2 = slice(0, None, scale)\n",
    "\n",
    "    for kwarg in sample:\n",
    "        sample[kwarg] = sample[kwarg][idx1, idx2]\n",
    "    return sample\n",
    "\n",
    "def addNoise(sample, noise, layers):\n",
    "    # Adding noise to u0\n",
    "    mean = sample['u0'].mean()\n",
    "    stdvar = np.sqrt(((sample['u0'] - mean) ** 2).mean())\n",
    "    size = sample['u0'].shape\n",
    "    startnoise = np.random.standard_normal(size)\n",
    "    sample['u0'] = sample['u0'] + noise * stdvar * startnoise\n",
    "\n",
    "    # Adding noise to ut, t > 0\n",
    "    for l in range(1, layers):\n",
    "        arg = 'u' + str(l)\n",
    "        size = sample[arg].shape\n",
    "        endnoise = np.random.standard_normal(size)\n",
    "        sample[arg] = sample[arg] + noise * stdvar * endnoise\n",
    "\n",
    "    return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate(options):\n",
    "    \"\"\"\n",
    "    Generating data / function-values on a regular grid of space-time, adding noise and taking a \n",
    "    batch of down-sampled regular sub-grids of this grid. This batch will contain the samples to \n",
    "    train our network with.\n",
    "\n",
    "    :param options: The dictionary of user-specified options (cf. main.py). Contains e.g. the \n",
    "                    grid-dimensions and noise\n",
    "    :return: A batch (as a list) of samples (as dictionaries), that in turn consist of (noisy) \n",
    "             function values on down-sampled sub-grids for all dt-layers.\n",
    "    \"\"\"\n",
    "    # u_t = a*u_x + b*u_y + c*u_{xx} + d*u_{yy}\n",
    "\n",
    "    # Variable declarations\n",
    "    nx = options['mesh_size'][0]\n",
    "    ny = options['mesh_size'][1]\n",
    "    nt = options['layers']\n",
    "    batch_size = options['batch_size']\n",
    "    dt = options['dt']\n",
    "    noise_level = options['noise_level']\n",
    "    downsample_by = options['downsample_by']\n",
    "\n",
    "    # Really good with a = b = 2\n",
    "    a = 2\n",
    "    b = 2\n",
    "    c = 0.5\n",
    "    d = 0.5\n",
    "\n",
    "    dx = 2 * np.pi / (nx - 1)\n",
    "    dy = 2 * np.pi / (ny - 1)\n",
    "\n",
    "    ## Needed for plotting:\n",
    "    # x = np.linspace(0, 2*np.pi, num = nx)\n",
    "    # y = np.linspace(0, 2*np.pi, num = ny)\n",
    "    # X, Y = np.meshgrid(x, y)\n",
    "\n",
    "    batch = []\n",
    "\n",
    "    for i in range(batch_size):\n",
    "        ############ Change the following lines to implement your own data ############\n",
    "\n",
    "        # Assign initial function:\n",
    "        u = initgen(options['mesh_size'], freq=4, boundary='Periodic')\n",
    "\n",
    "        ## Plotting the initial function:\n",
    "        # fig = plt.figure(figsize=(11,7), dpi=100)\n",
    "        # ax = fig.gca(projection='3d')\n",
    "        # surf = ax.plot_surface(X, Y, u[:], cmap=cm.viridis)\n",
    "        #\n",
    "        # plt.show()\n",
    "\n",
    "        sample = {}\n",
    "        sample['u0'] = u\n",
    "\n",
    "        for n in range(nt - 1):\n",
    "            un = pad_input_2(u, 2)[1:, 1:]\n",
    "\n",
    "            u = (un[1:-1, 1:-1] \n",
    "                 + c * dt / dx ** 2 * (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, 0:-2])\n",
    "                 + d * dt / dy ** 2 * (un[2:, 1: -1] - 2 * un[1:-1, 1:-1] + un[0:-2, 1:-1])\n",
    "                 + a * dt / dx * (un[1:-1, 2:] - un[1:-1, 1:-1])\n",
    "                 + b * dt / dy * (un[2:, 1:-1] - un[1:-1, 1:-1]))[:-1, :-1]\n",
    "\n",
    "            sample['u' + str(n + 1)] = u\n",
    "\n",
    "        downsample(sample, downsample_by)\n",
    "        addNoise(sample, noise_level, nt)\n",
    "\n",
    "        batch.append(sample)\n",
    "\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note, that we only take $n$ data samples with their corresponding function values from the grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simulate_data():\n",
    "    batch = generate(options)\n",
    "    \n",
    "    x_val_arr = []\n",
    "    y_val_arr = []\n",
    "    t_val_arr = []\n",
    "    val_arr = []\n",
    "    \n",
    "    for i in range(n): \n",
    "        t_rand = np.random.randint(options['layers'])\n",
    "        # Data is in [0,1]^2:\n",
    "        x_rand = np.random.randint(options['mesh_size'][0]//(2*np.pi))\n",
    "        y_rand = np.random.randint(options['mesh_size'][1]//(2*np.pi))\n",
    "\n",
    "        val = batch[0]['u' + str(t_rand)][x_rand, y_rand]\n",
    "\n",
    "        x_val = 2 * np.pi / (options['mesh_size'][0] - 1) * x_rand\n",
    "        y_val = 2 * np.pi / (options['mesh_size'][1] - 1) * y_rand\n",
    "        t_val = t_rand * options['dt']\n",
    "        \n",
    "        x_val_arr.append(x_val)\n",
    "        y_val_arr.append(y_val)\n",
    "        t_val_arr.append(t_val)\n",
    "        val_arr.append(val)\n",
    "        \n",
    "    return (np.array(x_val_arr), np.array(y_val_arr), np.array(t_val_arr), np.array(val_arr), \n",
    "            np.zeros(len(val_arr)))\n",
    "(x,y,t,y_u,y_f) = simulate_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2: Evaluate kernels\n",
    "\n",
    "$k_{uu}(X_i, X_j; \\theta) = \\sigma \\cdot exp(-\\frac{1}{2l_x}(x_i-x_j)^2 - \\frac{1}{2l_y}(y_i-y_j)^2 - \\frac{1}{2l_t}(t_i-t_j)^2)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_i, x_j, y_i, y_j, t_i, t_j, sigma, l_x, l_y, l_t, a,b,c,d = sp.symbols('x_i x_j y_i y_j t_i t_j \\\n",
    "                                                                         sigma l_x l_y l_t a b c d')\n",
    "kuu_sym = sigma*sp.exp(-1/(2*l_x)*((x_i - x_j)**2) - 1/(2*l_y)*((y_i - y_j)**2) \n",
    "                       - 1/(2*l_t)*((t_i - t_j)**2))\n",
    "kuu_fn = sp.lambdify((x_i, x_j, y_i, y_j, t_i, t_j, sigma, l_x, l_y, l_t), kuu_sym, \"numpy\")\n",
    "def kuu(x, y, t, sigma, l_x, l_y, l_t):\n",
    "    k = np.zeros((x.size, x.size))\n",
    "    for i in range(x.size):\n",
    "        for j in range(x.size):\n",
    "            k[i,j] = kuu_fn(x[i], x[j], y[i], y[j], t[i], t[j], sigma, l_x, l_y, l_t)\n",
    "    return k"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$k_{ff}(X_i,X_j;\\theta,\\phi)$ <br>\n",
    "$= \\mathcal{L}_{X_i}^{\\phi} \\mathcal{L}_{X_j}^{\\phi} k_{uu}(X_i, X_j; \\theta)$ <br>\n",
    "$= a^2\\partial_{x_i, x_j}k_{uu} + ab \\partial_{x_i, y_j}k_{uu} + ac \\partial_{x_i, x_j, x_j}k_{uu} + ad \\partial_{x_i, y_j, y_j}k_{uu} - a \\partial_{x_i, t_j}k_{uu}$ <br> \n",
    "$+ ba\\partial_{y_i, x_j}k_{uu} + b^2\\partial_{y_i, y_j}k_{uu} + bc\\partial_{y_i, x_j, x_j}k_{uu} + bd\\partial_{y_i, y_j, y_j}k_{uu} - b\\partial_{y_i, t_j}k_{uu}$ <br>\n",
    "$+ ca\\partial_{x_i, x_i, x_j}k_{uu}+ cb\\partial_{x_i, x_i, y_j}k_{uu}+ c^2\\partial_{x_i, x_i, x_j, x_j}k_{uu}+ cd\\partial_{x_i, x_i, y_j, y_j}k_{uu}- c\\partial_{x_i, x_i, t_j}k_{uu}$ <br>\n",
    "$+ da\\partial_{y_i, y_i,x_j}k_{uu}+ db\\partial_{y_i, y_i, y_j}k_{uu}+ dc\\partial_{y_i, y_i, x_j, x_j}k_{uu}+ d^2\\partial_{y_i, y_i, y_j, y_j}k_{uu}- d\\partial_{y_i, y_i, t_j}k_{uu}$ <br>\n",
    "$- a\\partial_{t_i, x_j}k_{uu}- b\\partial_{t_i, y_j}k_{uu}- c\\partial_{t_i, x_j, x_j}k_{uu}- d\\partial_{t_i, y_j, y_j}k_{uu}+ \\partial_{t_i, t_j}k_{uu}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "kff_sym = a**2*sp.diff(kuu_sym, x_i, x_j) \\\n",
    "        + a*b*sp.diff(kuu_sym, x_i, y_j) \\\n",
    "        + a*c*sp.diff(kuu_sym, x_i, x_j, x_j) \\\n",
    "        + a*d*sp.diff(kuu_sym, x_i, y_j, y_j) \\\n",
    "        - a*sp.diff(kuu_sym, x_i, t_j) \\\n",
    "        + b*a*sp.diff(kuu_sym, y_i, x_j) \\\n",
    "        + b**2*sp.diff(kuu_sym, y_i, y_j) \\\n",
    "        + b*c*sp.diff(kuu_sym, y_i, x_j, x_j) \\\n",
    "        + b*d*sp.diff(kuu_sym, y_i, y_j, y_j) \\\n",
    "        - b*sp.diff(kuu_sym, y_i, t_j) \\\n",
    "        + c*a*sp.diff(kuu_sym, x_i, x_i, x_j) \\\n",
    "        + c*b*sp.diff(kuu_sym, x_i, x_i, y_j) \\\n",
    "        + c**2*sp.diff(kuu_sym, x_i, x_i, x_j, x_j) \\\n",
    "        + c*d*sp.diff(kuu_sym, x_i, x_i, y_j, y_j) \\\n",
    "        - c*sp.diff(kuu_sym, x_i, x_i, t_j) \\\n",
    "        + d*a*sp.diff(kuu_sym, y_i, y_i, x_j) \\\n",
    "        + d*b*sp.diff(kuu_sym, y_i, y_i, y_j) \\\n",
    "        + d*c*sp.diff(kuu_sym, y_i, y_i, x_j, x_j) \\\n",
    "        + d**2*sp.diff(kuu_sym, y_i, y_i, y_j, y_j) \\\n",
    "        - d*sp.diff(kuu_sym, y_i, y_i, t_j) \\\n",
    "        - a*sp.diff(kuu_sym, t_i, x_j) \\\n",
    "        - b*sp.diff(kuu_sym, t_i, y_j) \\\n",
    "        - c*sp.diff(kuu_sym, t_i, x_j, x_j) \\\n",
    "        - d*sp.diff(kuu_sym, t_i, y_j, y_j) \\\n",
    "        + sp.diff(kuu_sym, t_i, t_j)\n",
    "kff_fn = sp.lambdify((x_i, x_j, y_i, y_j, t_i, t_j, sigma, l_x, l_y, l_t, a,b,c,d), kff_sym, \"numpy\")\n",
    "def kff(x, y, t, sigma, l_x, l_y, l_t, a,b,c,d):\n",
    "    k = np.zeros((x.size, x.size))\n",
    "    for i in range(x.size):\n",
    "        for j in range(x.size):\n",
    "            k[i,j] = kff_fn(x[i], x[j], y[i], y[j], t[i], t[j], sigma, l_x, l_y, l_t, a,b,c,d)\n",
    "    return k"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$k_{fu}(X_i,X_j;\\theta,\\phi) \\\\\n",
    "= \\mathcal{L}_{X_i}^{\\phi} k_{uu}(X_i, X_j; \\theta) \\\\\n",
    "= a\\partial_{x_i}k_{uu} + b \\partial_{y_i}k_{uu} + c \\partial_{x_i, x_i}k_{uu} + d \\partial_{y_i, y_i}k_{uu} -  \\partial_{t_i}k_{uu}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "kfu_sym = a*sp.diff(kuu_sym, x_i) \\\n",
    "        + b*sp.diff(kuu_sym, y_i) \\\n",
    "        + c*sp.diff(kuu_sym, x_i, x_i) \\\n",
    "        + d*sp.diff(kuu_sym, y_i, y_i) \\\n",
    "        - sp.diff(kuu_sym, t_i)\n",
    "kfu_fn = sp.lambdify((x_i, x_j, y_i, y_j, t_i, t_j, sigma, l_x, l_y, l_t, a,b,c,d), kfu_sym, \"numpy\")\n",
    "def kfu(x, y, t, sigma, l_x, l_y, l_t, a,b,c,d):\n",
    "    k = np.zeros((x.size, x.size))\n",
    "    for i in range(x.size):\n",
    "        for j in range(x.size):\n",
    "            k[i,j] = kfu_fn(x[i], x[j], y[i], y[j], t[i], t[j], sigma, l_x, l_y, l_t, a,b,c,d)\n",
    "    return k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kuf(x, y, t, sigma, l_x, l_y, l_t, a,b,c,d):\n",
    "    return kfu(x, y, t, sigma, l_x, l_y, l_t, a,b,c,d).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3: Compute NLML (with Cholesky decomposition)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implementing the covariance matrix K and its inverse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def K(sigma, l_x, l_y, l_t, a,b,c,d, s):\n",
    "    K_mat = np.block([\n",
    "        [kuu(x, y, t, sigma, l_x, l_y, l_t)+s*np.eye(n),kuf(x, y, t, sigma, l_x, l_y, l_t, a,b,c,d)],\n",
    "        [kfu(x, y, t, sigma, l_x, l_y, l_t,a,b,c,d),kff(x,y,t,sigma,l_x,l_y,l_t,a,b,c,d)+s*np.eye(n)]\n",
    "    ])\n",
    "    return K_mat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def K_inv_and_det(sigma, l_x, l_y, l_t, a,b,c,d, s):\n",
    "    \n",
    "    K_inv = np.zeros((2*n, 2*n))\n",
    "    log_sum = 0\n",
    "    \n",
    "    # Use Cholesky, if possible. Otherwise use SVD.\n",
    "        \n",
    "    try:\n",
    "        L = np.linalg.cholesky(K(sigma, l_x, l_y, l_t, a,b,c,d, s))\n",
    "        L_inv = solve_triangular(L, np.identity(2*n), lower=True) # Slight performance boost \n",
    "                                                                  # over np.linalg.inv\n",
    "        K_inv = (L_inv.T).dot(L_inv)\n",
    "\n",
    "        for i in range(2*n):\n",
    "            log_sum = log_sum + np.log(np.abs(L[i,i]))\n",
    "    except np.linalg.LinAlgError:\n",
    "        # Inverse of K via SVD\n",
    "        u, s_mat, vt = np.linalg.svd(K(sigma, l_x, l_y, l_t, a,b,c,d, s))\n",
    "        K_inv = (vt.T).dot(np.linalg.inv(np.diag(s_mat))).dot(u.T)  \n",
    "\n",
    "        # Calculating the log of the determinant of K\n",
    "        # Singular values are always positive.\n",
    "        for i in range(s_mat.size):\n",
    "            log_sum = log_sum + np.log(s_mat[i])\n",
    "     \n",
    "    return K_inv, log_sum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implementing normalized negative log-likelihood function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def nlml(params):\n",
    "\n",
    "    b_par = 2\n",
    "    c_par = 0.5\n",
    "    d_par = 0.5\n",
    "    \n",
    "    # Exponentiation to enable unconstrained optimization\n",
    "    sigma_exp = np.exp(params[0]) \n",
    "    l_x_exp = np.exp(params[1])\n",
    "    l_y_exp = np.exp(params[2]) \n",
    "    l_t_exp = np.exp(params[3])\n",
    "    # a = params[4]\n",
    "    y_con = np.concatenate((y_u, y_f))\n",
    "    \n",
    "    A,b = K_inv_and_det(sigma_exp, l_x_exp, l_y_exp, l_t_exp, params[4],b_par,c_par,d_par, s)\n",
    "        \n",
    "    val = b + y_con @ A @ y_con\n",
    "    return val"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 4: Optimize hyperparameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**1. Nelder-Mead**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def callbackf(params):\n",
    "    print(params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def minimize_restarts(x,y,y_u,y_f,n=5): \n",
    "    all_results = []\n",
    "    for it in range(0,n):\n",
    "        all_results.append(opt.minimize(nlml, np.random.rand(5), callback = callbackf, \n",
    "                                        method=\"Nelder-Mead\", \n",
    "                                        options={'maxfev':5000, 'fatol':0.001, 'xatol':0.001}))\n",
    "    return min(all_results, key = lambda x: x.fun)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t0 = time.time()\n",
    "m = minimize_restarts(x, y, y_u, y_f, 3)\n",
    "t_Nelder = time.time() - t0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1865.7036004066467"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "t_Nelder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The inferred parameter is: 2.1634004301602183\n"
     ]
    }
   ],
   "source": [
    "print('The inferred parameter is:', m.x[4])"
   ]
  },
  {
   "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}