docs/source/examples/GPOptimizer_gp2ScaleTest.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "56dcab7c",
   "metadata": {},
   "source": [
    "# gp2Scale via the gpOptimizer \n",
    "gp2Scale is a special setting in gpCAM that combines non-stationary, compactly-supported kernels, HPC distributed computing, and sparse random linear algebra to allow scale-up of exact GPs to millions of data points. Here we run a moderately-sized GP, just because we assume you might run this locally.\n",
    "\n",
    "I hope it is clear how cool it is what is happening here. If you have a dask client that points to a remote cluster with 500 GPUs, you will distribute the covariance matrix computation across those. The full matrix is sparse and will be fast to work with in downstream operations. The algorithm only makes use of naturally-occuring sparsity, so the result is exact in contrast to Vecchia or inducing-point methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8a50af7",
   "metadata": {},
   "outputs": [],
   "source": [
    "##first install the newest version of fvgp\n",
    "#!pip install gpcam==8.0.3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c08e81e2",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a20595d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from gpcam import GPOptimizer\n",
    "from dask.distributed import Client\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "client = Client() ##this is the client you can make locally like this or \n",
    "#your HPC team can provide a script to get it. We included an example to get gp2Scale going\n",
    "#on NERSC's Perlmutter\n",
    "\n",
    "\n",
    "#It's good practice to make sure to wait for all the workers to be ready\n",
    "client.wait_for_workers(4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6165ec13",
   "metadata": {},
   "source": [
    "## Preparing the Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee0b9649",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f1(x):\n",
    "    return ((np.sin(5. * x) + np.cos(10. * x) + (2.* (x-0.4)**2) * np.cos(100. * x)))\n",
    "\n",
    "input_dim = 1\n",
    "N = 10000\n",
    "x_data = np.random.rand(N,input_dim)\n",
    "y_data = f1(x_data)\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3b77be1-02ad-4c49-a53d-1c777c69dbcd",
   "metadata": {},
   "source": [
    "## Setting up the GPOptimizer with gp2Scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0805781d",
   "metadata": {},
   "outputs": [],
   "source": [
    "hps_n = 2\n",
    "\n",
    "hps_bounds = np.array([[0.1,10.],      ##signal var of Wendland kernel\n",
    "                       [0.001,0.02]])  ##length scale for Wendland kernel\n",
    "\n",
    "init_hps = np.random.uniform(size = len(hps_bounds), low = hps_bounds[:,0], high = hps_bounds[:,1])\n",
    "\n",
    "my_gp2S = GPOptimizer(x_data,y_data,init_hyperparameters=init_hps, \n",
    "            gp2Scale = True, gp2Scale_batch_size= 1000, gp2Scale_dask_client = client, info = True\n",
    "            )\n",
    "\n",
    "my_gp2S.train(hyperparameter_bounds=hps_bounds, max_iter = 2)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec650a00-3a1c-4cab-beb5-03968888f713",
   "metadata": {},
   "source": [
    "## Vizualizing the Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de466547",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_pred = np.linspace(0,1,100) ##for big GPs, this is usually not a good idea, but in 1d, we can still do it\n",
    "                              ##It's better to do predictions only for a handful of points.\n",
    "\n",
    "mean1 = my_gp2S.posterior_mean(x_pred.reshape(-1,1))[\"f(x)\"]\n",
    "var1 =  my_gp2S.posterior_covariance(x_pred.reshape(-1,1), variance_only=False)[\"v(x)\"]\n",
    "\n",
    "print(my_gp2S.hyperparameters)\n",
    "\n",
    "plt.figure(figsize = (16,10))\n",
    "plt.plot(x_pred,mean1, label = \"posterior mean\", linewidth = 4)\n",
    "plt.plot(x_pred,f1(x_pred), label = \"latent function\", linewidth = 4)\n",
    "plt.fill_between(x_pred, mean1 - 3. * np.sqrt(var1), mean1 + 3. * np.sqrt(var1), alpha = 0.5, color = \"grey\", label = \"var\")\n",
    "plt.scatter(x_data,y_data, color = 'black')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a327330a",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_gp2S.ask(np.array([[0,1]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4d79582",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}