{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook to test and learn about the Ensemble GP (Experimental)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#HAVE fvGP installed? Otherwise run this cell which installs the latest fvGP\n",
    "#!pip install fvgp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "#import autograd.numpy as np  # Thinly-wrapped numpy\n",
    "from fvgp.ensemble_gp import EnsembleGP\n",
    "from fvgp.ensemble_gp import hyperparameters\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# An ensemble GP means that we allow the prior to be a sum of Gaussian priors. That means we have to define a set of prior-mean and kernel functions\n",
    "In our case we assume 2 GPs, and therefore 2 prior mean and kernel functions\n",
    "Feel free to play with those mean functions. Prior mean functions can be any real function, but remember, kernel funcitons have to be positive semi-definite. In terms of test functions, we define here a 1d step function and eggholders function in 2d."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def step(x):\n",
    "    i = np.where(x[:,0] < 0.0)\n",
    "    r = np.zeros((len(x[:,0])))\n",
    "    r[i] = 1.0\n",
    "    return r + np.sin(5.0 * x[:,0])/50.0\n",
    "def eggholder(x):\n",
    "    x1 = x[:,0]\n",
    "    x2 = x[:,1]\n",
    "    term1 = -(x2 + 47.) * np.sin(np.sqrt(abs(x2 + x1/2.+47.)))\n",
    "    term2 = -x1 * np.sin(np.sqrt(abs(x1-(x2+47.))))\n",
    "    return (term1 + term2)/1000.0\n",
    "\n",
    "def mean_func1(obj,x,hyperparameters):\n",
    "    return np.ones(len(x)) * hyperparameters[-1]\n",
    "\n",
    "def mean_func2(obj,x,hyperparameters):\n",
    "    return np.ones(len(x)) * hyperparameters[-1]\n",
    "\n",
    "def kernel1(x1,x2,hps,obj):\n",
    "    d = 0\n",
    "    for i in range(len(x1[0])):\n",
    "        d += abs(np.subtract.outer(x1[:,i],x2[:,i]))\n",
    "        #a = np.array([x1[:,i],]*len(x2)).T\n",
    "        #b = np.array([x2[:,i],]*len(x1))\n",
    "        #d += abs(a-b) #+ 1e-6\n",
    "    d = np.sqrt(d)\n",
    "    #return hps[0] * obj.exponential_kernel(d,hps[1])\n",
    "    return hps[0] * obj.matern_kernel_diff1(d,hps[1])\n",
    "\n",
    "def kernel2(x1,x2,hps,obj):\n",
    "    d = 0\n",
    "    for i in range(len(x1[0])):\n",
    "        d += abs(np.subtract.outer(x1[:,i],x2[:,i]))\n",
    "        #a = np.array([x1[:,i],]*len(x2)).T\n",
    "        #b = np.array([x2[:,i],]*len(x1))\n",
    "        #d += abs(a-b) #+ 1e-6\n",
    "    d = np.sqrt(d)\n",
    "    #return hps[0] * obj.exponential_kernel(d,hps[1])\n",
    "    return hps[0] * obj.matern_kernel_diff1(d,hps[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### For ensemble GPs, the hyperparameters are more complicated than for single GPs and therefore we define an extra object for the hyperparameters. That object keeps track of the structure of the hyperparameters vectorized and unvectorized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights = np.ones(2) / 2.0   ####those are two weights for the two GPs\n",
    "weights_bounds = np.array([[0.000001,1.0],[0.000001,1.0]]) ##those are the bounds we will use to optimize the weights\n",
    "hps = [np.array([1e6,10000.0,100.]),np.array([1e6,10000.0,0.0])] ##here we have the hyperparameters: (in order)\n",
    "                                                           ##GP1 signal variance, length scale, prior mean\n",
    "                                                           ##same for GP2\n",
    "hps_bounds = [np.array([[0.001,10.],[0.01,10000.],[-2000.0,2000.0]]),\n",
    "              np.array([[0.001,10.],[0.01,10000.],[-2000.0,2000.0]])] ##those are bounds for training, same order\n",
    "hps_obj = hyperparameters(weights,weights_bounds,hps,hps_bounds)  ##creating the hps obj instance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ...creating some data..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(400, 2)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<mpl_toolkits.mplot3d.art3d.Path3DCollection at 0x7f70e4f16fd0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 192052 characters]\n",
      "text/plain": [
       "<Figure size 1080x1080 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "##Random Data wit Eggholder or Step Function\n",
    "N = 400\n",
    "points = np.empty((N,2))   ##needst two columns for eggholder\n",
    "#points = np.random.uniform(low = -3, high = 3, size = points.shape) # for step function\n",
    "points = np.random.uniform(low = -600, high = 600, size = points.shape) # for step function\n",
    "#values = step(points)\n",
    "values = eggholder(points)\n",
    "\n",
    "print(points.shape)\n",
    "\n",
    "\n",
    "#plt.scatter(points,values)\n",
    "\n",
    "\n",
    "fig = plt.figure(figsize = (15,15))\n",
    "ax = fig.add_subplot(projection='3d')\n",
    "ax.scatter(points[:,0], points[:,1], values, c=values, cmap=plt.cm.coolwarm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ...initializing the ensemble GP and training it ... oh and checking some derivatives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5.e-01 5.e-01 1.e+06 1.e+04 1.e+02 1.e+06 1.e+04 0.e+00]\n",
      "(400, 2)\n",
      "(400,)\n",
      "[array([1.e+06, 1.e+04, 1.e+02]), array([1000000.,   10000.,       0.])]\n",
      "likelihood before training:  306.11522748139805\n"
     ]
    }
   ],
   "source": [
    "print(hps_obj.vectorized_hps)\n",
    "print(points.shape)\n",
    "print(values.shape)\n",
    "print(hps_obj.hps)\n",
    "\n",
    "#the first parameter has to be \"1\" for the step function and \"2\" for eggholder\n",
    "my_gp = EnsembleGP(len(points[0]),points,values,2,hps_obj,variances = np.ones(values.shape) * 0.01,\n",
    "                gp_kernel_functions = [kernel1,kernel2],\n",
    "                gp_mean_functions = [mean_func1,mean_func2],\n",
    "                compute_device = \"cpu\")\n",
    "print(\"likelihood before training: \", my_gp.ensemble_log_likelihood(hps_obj.vectorized_hps))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ensemble fvGP submitted to global optimization\n",
      "bounds are [[ 1.e-06  1.e+00]\n",
      " [ 1.e-06  1.e+00]\n",
      " [ 1.e-03  1.e+01]\n",
      " [ 1.e-02  1.e+04]\n",
      " [-2.e+03  2.e+03]\n",
      " [ 1.e-03  1.e+01]\n",
      " [ 1.e-02  1.e+04]\n",
      " [-2.e+03  2.e+03]]\n",
      "maximum number of iterations:  2\n",
      "termination tolerance:  0.0001\n",
      "differential_evolution step 1: f(x)= 12075.8\n",
      "differential_evolution step 2: f(x)= 3003.33\n",
      "fvGP found weights  [0.4    0.5999]\n",
      "and hyperparameters:  [array([  10.        , 4259.01761973, -116.09355957]), array([ 2.84861982e+00,  7.68427338e+03, -9.36517634e+02])]\n",
      " with likelihood:  2409.498559712622  via global optimization\n",
      "new weights after training:  [0.4    0.5999]\n",
      "new hps     after training:  [array([  10.        , 4259.01761973, -116.09355957]), array([ 2.84861982e+00,  7.68427338e+03, -9.36517634e+02])]\n",
      "GPs updated\n"
     ]
    }
   ],
   "source": [
    "my_gp.train(hps_obj,\n",
    "                pop_size = 20,\n",
    "                tolerance = 0.0001,\n",
    "                max_iter = 2, dask_client = None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "current hps:  [ 4.00e-01  6.00e-01  1.00e+01  4.26e+03 -1.16e+02  2.85e+00  7.68e+03 -9.37e+02]\n",
      "GRADIENT TEST\n",
      "[ 4.00e-01  6.00e-01  1.00e+01  4.26e+03 -1.16e+02  2.85e+00  7.68e+03 -9.37e+02]\n",
      "compare:  -2.499999936844688 -2.5000000000159157\n",
      "compare:  0.0 -0.0\n",
      "compare:  -69.32736050657695 -69.3264532791129\n",
      "compare:  0.00949853529164102 0.009452495423411249\n",
      "compare:  -11.61043928732397 -11.610446434067\n",
      "compare:  0.0 -0.0\n",
      "compare:  0.0 0.0\n",
      "compare:  0.0 -0.0\n",
      "grad:  [-2.50e+00 -0.00e+00 -6.93e+01  9.45e-03 -1.16e+01 -0.00e+00  0.00e+00 -0.00e+00]\n",
      "===========================\n",
      "gradient correctly computed\n",
      "===========================\n"
     ]
    }
   ],
   "source": [
    "#from autograd import jacobian\n",
    "#from autograd import elementwise_grad as el\n",
    "np.set_printoptions(precision=2,linewidth=100)\n",
    "print(\"current hps: \",hps_obj.vectorized_hps)\n",
    "grad = my_gp.ensemble_log_likelihood_grad(hps_obj.vectorized_hps)\n",
    "print(\"GRADIENT TEST\")\n",
    "e = 1e-6\n",
    "print(hps_obj.vectorized_hps)\n",
    "for i in range(len(hps_obj.vectorized_hps)):\n",
    "    v1 = np.array(hps_obj.vectorized_hps)\n",
    "    v2 = np.array(hps_obj.vectorized_hps)\n",
    "    v1[i] = v1[i] - e\n",
    "    v2[i] = v2[i] + e\n",
    "    #print(my_gp.ensemble_log_likelihood(v1))\n",
    "    #print(my_gp.ensemble_log_likelihood(v2))\n",
    "    g = (my_gp.ensemble_log_likelihood(v2) - my_gp.ensemble_log_likelihood(v1))/(2.*e)\n",
    "    print(\"compare: \", g, grad[i])\n",
    "    #print(\"v1: \",v1[i],\" v2: \",v2[i])\n",
    "    assert abs(grad[i] - g) < 0.01, \"grad wrong\"\n",
    "print(\"grad: \",grad)\n",
    "print(\"===========================\")\n",
    "print(\"gradient correctly computed\")\n",
    "print(\"===========================\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------\n",
      "----------\n",
      "hess w: \n",
      "[[ 6.25  0.  ]\n",
      " [ 0.   -0.  ]]\n",
      "-------------\n",
      "3\n",
      "[[1.16e+02 1.86e+02 1.16e+00]\n",
      " [1.86e+02 2.40e+02 8.81e-07]\n",
      " [1.16e+00 8.81e-07 1.00e-01]]\n",
      "[[0. 0. 0.]\n",
      " [0. 0. 0.]\n",
      " [0. 0. 0.]]\n",
      "[[ 3.79e+04 -2.31e+00  1.15e+02]\n",
      " [-2.31e+00 -6.22e+01  2.61e-06]\n",
      " [ 1.15e+02  2.61e-06  3.51e-01]]\n",
      "-------------\n",
      "HESSIAN TEST\n",
      "[[6.25e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00]\n",
      " [0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00]\n",
      " [0.00e+00 0.00e+00 1.36e+01 3.69e-02 1.16e+00 0.00e+00 0.00e+00 0.00e+00]\n",
      " [0.00e+00 0.00e+00 3.69e-02 6.90e-02 9.09e-05 0.00e+00 0.00e+00 0.00e+00]\n",
      " [0.00e+00 0.00e+00 1.16e+00 9.09e-05 1.01e-01 0.00e+00 0.00e+00 0.00e+00]\n",
      " [0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00]\n",
      " [0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00]\n",
      " [0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00 0.00e+00]]\n"
     ]
    }
   ],
   "source": [
    "print(\"----------\")\n",
    "hess = my_gp.ensemble_log_likelihood_hess(hps_obj.vectorized_hps)\n",
    "print(\"----------\")\n",
    "\n",
    "\n",
    "print(\"hess w: \")\n",
    "print(hess[0])\n",
    "\n",
    "print(\"-------------\")\n",
    "print(len(hess[1]))\n",
    "print((hess[1][0]))\n",
    "print((hess[1][1]))\n",
    "print((hess[1][2]))\n",
    "print(\"-------------\")\n",
    "\n",
    "\n",
    "print(\"HESSIAN TEST\")\n",
    "e = 1e-4\n",
    "h = np.zeros((len(hps_obj.vectorized_hps),len(hps_obj.vectorized_hps)))\n",
    "for i in range(len(hps_obj.vectorized_hps)):\n",
    "    for j in range(len(hps_obj.vectorized_hps)):\n",
    "        v1 = np.array(hps_obj.vectorized_hps)\n",
    "        v2 = np.array(hps_obj.vectorized_hps)\n",
    "        v3 = np.array(hps_obj.vectorized_hps)\n",
    "        v4 = np.array(hps_obj.vectorized_hps)\n",
    "        v1[i] = v1[i] - e\n",
    "        v1[j] = v1[j] - e\n",
    "        \n",
    "        v2[i] = v2[i] + e\n",
    "        v2[j] = v2[j] + e\n",
    "        \n",
    "        \n",
    "        v3[i] = v3[i] - e\n",
    "        v3[j] = v3[j] + e\n",
    "        \n",
    "        \n",
    "        v4[i] = v4[i] + e\n",
    "        v4[j] = v4[j] - e\n",
    "\n",
    "        h[i,j] = (my_gp.ensemble_log_likelihood(v2) - my_gp.ensemble_log_likelihood(v4) - \n",
    "                  my_gp.ensemble_log_likelihood(v3) + my_gp.ensemble_log_likelihood(v1))/(4.*e*e)\n",
    "\n",
    "print(h)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ...training asynchronously... "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import time\n",
    "my_gp.train_async(hps_obj,\n",
    "                pop_size = 20,   ####not active, HGDL will use all available \n",
    "                                 ####workers as communicated by the dask_client\n",
    "                tolerance = 0.0001,\n",
    "                max_iter = 100000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "my_gp.update_hyperparameters()\n",
    "print(\"current likelihood: \", my_gp.ensemble_log_likelihood(hps_obj.vectorized_hps))\n",
    "print(my_gp.opt.optima.list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CAREFUL: The next command will kill your training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_gp.kill_training()\n",
    "print(\"current likelihood: \", my_gp.ensemble_log_likelihood(hps_obj.vectorized_hps))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### I hope you saw how the Gaussian priors are places at 0 and 1, just like we would expect from the data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Depending on the randomly chosen data you might see that the result is not Gaussian\n",
    "if you can't see it, run the notebook again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "#plot for 1d\n",
    "import matplotlib.pyplot as plt\n",
    "plot_n  = 1000\n",
    "x_input = np.empty((plot_n,1))\n",
    "x_input[:,0] = np.linspace(-2,2,plot_n)\n",
    "y = func(x_input)\n",
    "post = my_gp.posterior(x_input, res = 1000)\n",
    "pdf = post[\"pdf\"]\n",
    "\n",
    "fig = plt.figure(figsize = (10,10))\n",
    "plt.imshow((np.flip(np.array(pdf).T,axis = 0)),cmap=plt.cm.viridis)\n",
    "plt.show()\n",
    "mean1 = post[\"f(x)\"][0]\n",
    "mean2 = post[\"f(x)\"][1]\n",
    "cov1 = post[\"v(x)\"][0]\n",
    "cov2 = post[\"v(x)\"][1]\n",
    "plt.plot(np.linspace(-3,3,1000),mean1, label = \"GP1\")\n",
    "plt.plot(np.linspace(-3,3,1000),mean2, label = \"GP2\")\n",
    "plt.legend()\n",
    "plt.fill_between(np.linspace(-3,3,1000),mean1 - 3.0 * np.sqrt(cov1), mean1 + 3.0 * np.sqrt(cov1), alpha=0.2)\n",
    "plt.fill_between(np.linspace(-3,3,1000),mean2 - 3.0 * np.sqrt(cov2), mean2 + 3.0 * np.sqrt(cov2), alpha=0.2)\n",
    "plt.scatter(points,values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#plot for 2d\n",
    "import plotly.graph_objects as go\n",
    "import numpy as np\n",
    "def plot(x,y,z,data = None):\n",
    "    fig = go.Figure()\n",
    "    fig.add_trace(go.Surface(x = x, y = y,z=z))\n",
    "    if data is not None: \n",
    "        fig.add_trace(go.Scatter3d(x=data[:,0], y=data[:,1], z=data[:,2],\n",
    "                                   mode='markers'))\n",
    "\n",
    "    fig.update_layout(title='Posterior Mean', autosize=True,\n",
    "                  width=800, height=800,\n",
    "                  margin=dict(l=65, r=50, b=65, t=90))\n",
    "\n",
    "\n",
    "    fig.show()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "plot_n  = 50\n",
    "x_pred = np.zeros((plot_n**2,2))\n",
    "x = np.linspace(-600,600,plot_n)\n",
    "y = np.linspace(-600,600,plot_n)\n",
    "x,y = np.meshgrid(x,y)\n",
    "counter = 0\n",
    "for i in  range(plot_n):\n",
    "    for j in range(plot_n):\n",
    "        x_pred[counter] = np.array([x[i,j],y[i,j]])\n",
    "        counter += 1\n",
    "  \n",
    "post = my_gp.posterior(x_pred, res = 1000)\n",
    "pdf = post[\"pdf\"]\n",
    "\n",
    "\n",
    "mean1 = post[\"f(x)\"][0].reshape(plot_n,plot_n)\n",
    "mean2 = post[\"f(x)\"][1].reshape(plot_n,plot_n)\n",
    "cov1 = post[\"v(x)\"][0].reshape(plot_n,plot_n)\n",
    "cov2 = post[\"v(x)\"][1].reshape(plot_n,plot_n)\n",
    "plot(x,y,mean1)\n",
    "plot(x,y,mean2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "fvgp",
   "language": "python",
   "name": "fvgp"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
