From 56506e0c92272ccd2ba09e940000e6c9a44a4d4c Mon Sep 17 00:00:00 2001
From: Maxim Ziatdinov <ziatdinovmax@gmail.com>
Date: Fri, 8 Nov 2024 09:26:19 -0800
Subject: [PATCH] Update example

---
 examples/pbnn_example1d.ipynb | 62 ++++++++++++++++++++++++++++++-----
 1 file changed, 54 insertions(+), 8 deletions(-)

diff --git a/examples/pbnn_example1d.ipynb b/examples/pbnn_example1d.ipynb
index f1752da..89082d5 100644
--- a/examples/pbnn_example1d.ipynb
+++ b/examples/pbnn_example1d.ipynb
@@ -1,5 +1,21 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Partially Bayesian Neural Networks - PBNNs\n",
+    "\n",
+    "*Prepared by Maxim Ziatdinov (October 2024)*"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This notebook compares active learning with a standard Gaussian process (GP) and with a partially Bayesian neural network (PBNN) on a dataset with discontinuities and non-stationarities."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -33,6 +49,13 @@
     "## Data"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Create a measurement function and view the ground truth:"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 2,
@@ -69,6 +92,13 @@
     "## GP-based active learning"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Run standard Gaussian process-based active learning:"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,
@@ -1054,6 +1084,13 @@
     "    y_measured = np.append(y_measured, y_next)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Compute MSE and NLPD scores:"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 4,
@@ -1071,6 +1108,13 @@
     "## Partial BNN-based active learning"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The ```PartialBNN``` combines deterministic and Bayesian layers. For untrained networks, it first performs MAP estimation using regularized SGD with stochastic weight averaging to find good deterministic weights. These MAP-optimized weights then serve as the means for the Gaussian priors in the Bayesian layers, around which MCMC sampling using No U-Turn Sampler explores the posterior distribution. Alternatively, pre-trained weights can be provided directly at the model initialization stage."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 6,
@@ -2368,6 +2412,13 @@
     "    y_measured = np.append(y_measured, y_next)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Compute MSE and NLPD scores for PBNN:"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 7,
@@ -2382,7 +2433,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Compare results"
+    "## Compare results\n",
+    "\n",
+    "(the lower the better)"
    ]
   },
   {
@@ -2418,13 +2471,6 @@
     "plt.show()"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "code",
    "execution_count": null,