# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Bayesian Optimisation Xmas
# This notebook implements a xmas-themed guessing game to try and find the maximum of a black-box function.
# It is part of this [repository](https://github.com/JohnReid/BayesOptJournalClub).

# +
from dataclasses import astuple
import ipywidgets as widgets
from IPython.display import display
import numpy as np
from scipy.interpolate import interp1d
import pandas as pd
import altair as alt
import tensorflow as tf

# For Gaussian processes
import gpflow
from gpflow.utilities import print_summary, set_trainable

# For Bayesian optimisation
import trieste
from trieste.bayesian_optimizer import OptimizationResult
from trieste.utils.objectives import mk_observer
from trieste.acquisition.rule import OBJECTIVE, EfficientGlobalOptimization
from trieste.acquisition import NegativeLowerConfidenceBound, ExpectedImprovement

# Supporting code in the repository
import blackbox as bb
from importlib import reload
reload(bb);
# -

# ## Santa's snowy scenario
# It is Xmas Eve and Santa has to deliver his presents to all the children around the world.
# Unfortunately it has been a really cold winter and some houses are completely blanketed by snow.
#
# Santa does have a shovel and could dig down to the chimney of each house to deliver his presents if only he knew where the chimneys were.
# Santa does know that the chimney is the highest point in the skyline of each house.
#
# Fortunately Rudolph can help.
# He can measure the depth of the snow at any point with his antlers but unfortunately this takes some time and the measurements are noisy.
# Furthermore Rudolph demands a boiled sweet (his favourite treat) every time he measures the snow depth and Santa only has so many sweets.
#
# Santa has one night to deliver his presents, how can he find the chimneys in time?
#
# ### Skyline
# Let's see what a typical house's skyline looks like.
# `x` represents the location and `depth` measures the depth below the snowline.

x, depth = bb.random_submerged_skyline(100)
house = pd.DataFrame(dict(x=x, depth=depth))
house_chart = (
    alt.Chart(house)
    .mark_line()
    .encode(x='x', y='depth'))
house_chart.properties(width=800, height=300)

# ### Grid-search strategy
# One strategy to find the chimney is to lay down a grid and ask Rudolph to measure the height at each grid point.
#
# Let's fix the standard deviation of the Rudolph's measurement error and choose a number of grid points:

error_sd = widgets.FloatSlider(
    value=.04,
    min=0.02,
    max=.1,
    step=0.01,
    description='error sd:',
    readout_format='.2f',
)
npoints = widgets.IntSlider(
    value=20,
    min=3,
    max=100,
    step=1,
    description='grid size:',
    readout_format='d'
)
display(error_sd)
display(npoints)

# Rudolph makes his measurements.

measurements = pd.DataFrame(dict(x=np.linspace(house['x'].min(), house['x'].max(), npoints.value)))
f = interp1d(house['x'], house['depth'])
measurements['f'] = f(measurements['x'])
measurements['y'] = measurements['f'] + bb.rng.normal(0, error_sd.value, size=npoints.value)
measurements_chart = (
    alt.Chart(measurements)
    .mark_point(color='red')
    .encode(x='x', y='y'))
(measurements_chart + house_chart).properties(width=800, height=300)

# We can see that there are many wasted measurements in locations where we are reasonably sure the maximum is not located.
# In addition it is still difficult to be sure where the chimney is.
#
# How can we do better?
#
# ## Bayesian optimisation
# This is where Bayesian optimisation (BO) comes in.
# Bayesian optimisation is going to do better than grid (or random) search by choosing where to measure the depth of snow in some sort of *optimal* way.
#
# Bayesian optimisation requires:
# - a *prior on the underlying function* to be optimised (some understanding of houses' skylines and their depths)
# - a *prior on the measurement error* (some understanding of how good Rudolph is at measuring depth)
#
# Using some underlying theory, Bayesian optimisation then iterates around the following loop:
# - *choose the next measurement location* so as to minimise the total number of measurements required (the acquisition function)
# - ask the black-box function for a *noisy measurement* at that location
# - *update its model* of the underlying function and measurement error (computes a posterior given prior and data)
#
# Some stopping criterion is used to avoid iterating indefinitely.
# Bayesian optimisation returns the posterior for the underlying function to use as you see fit.
# In particular you may wish to obtain a posterior over the location of the maximum of the function.
#
# ## Optimisation guessing game
# However, before we try Bayesian optimisation, we will try to optimise a black-box function ourselves.
# Your task will be to find the maximum of the underlying function using as few function evaluations as possible.
#
# Unfortunately we do not have a suitable Bayesian model for house skylines so we will use a Gaussian process prior on functions instead.
# Some typical functions are shown below to give you an idea what the underlying function might look like.

blackbox1 = bb.GPBlackBox(ndim=1)
xx = blackbox1.xgrid(npoints=100)[:, 0]
prior_samples = pd.concat([
    pd.DataFrame(dict(x=xx, f=bb.prior_draw(blackbox1.kernel, xx), sample=s))
    for s in range(3)])
prior_chart = (
    alt.Chart(prior_samples)
    .mark_line()
    .encode(x='x', y='f', color='sample:N'))
prior_chart.properties(width=800, height=300)

# ### A function of one input
# Now it is your turn.
# You have to optimise a function of one input with domain `[-1, 1]`.
# To start you off you already have three evaluations.
# Remember these are noisy versions of the underlying objective function.

blackbox1 = bb.GPBlackBox(ndim=1)
y = blackbox1([-.6])[0][0]
y = blackbox1([.6])[0][0]
blackbox1.plot_xy().properties(width=800, height=300)

# Now you should choose which input point (`x0`) to evaluate the function at next:

x0 = widgets.FloatSlider(
    value=0.2,
    min=bb.DOMAIN_MIN,
    max=bb.DOMAIN_MAX,
    step=0.01,
    description='x0:',
    readout_format='.2f',
)
display(x0)

# Repeatedly adjust the slider above and execute the following cell to make new evaluations and plot all the evaluations `y` so far.

y = blackbox1([x0.value])[0][0]
print(f'Evaluated black box at {x0.value}; result={y}')
blackbox1.plot_xy().properties(width=800, height=300)

# Now go back and choose another point at which to evaluate the black box function.
#
# If you have finished evaluating the function at different points and you are confident where the maximum is, you can make a guess before executing the cells below. The cell below shows the function as a line and the noisy data we received as evaluations of it:

f1 = blackbox1.sample_f(100)
chart1f = (
    alt.Chart(f1)
    .mark_line()
    .encode(x='x0', y='f'))
chart1y = blackbox1.plot_xy()
chart1 = alt.layer(chart1y, chart1f)
chart1.properties(width=800, height=300)

# Note that we didn't actually sample the underlying function `f` until the cell above despite having already seen the noisy observations `y`.
# If you execute the cell above again, you will have a different sample for `f` that is also consistent with our noisy observations `y`.
# We can do this because Gaussian processes have a nice consistency property under marginalisation.
#
# ### Bayesian optimisation for 1-dimensional input
# Now we will try Bayesian optimisation using the [Trieste](https://secondmind-labs.github.io/trieste/) package which uses [GPflow](https://gpflow.readthedocs.io/en/master/) for Gaussian process regression.
# This section is based on the Trieste tutorial.

# +
model1 = gpflow.models.GPR((f1['x0'].to_numpy().reshape((-1, 1)), f1['f'].to_numpy().reshape((-1, 1))),
                           kernel=gpflow.kernels.Matern52(lengthscales=.4),
                           noise_variance=1.1e-6)

def blackbox_bo1(x):
    mean, var = model1.predict_y(tf.reshape(x, (-1, 1)))
    return - bb.rng.normal(loc=mean, scale=np.sqrt(blackbox1.noise_variance))


# -

# #### Initial sample over search space

# +
observer = mk_observer(blackbox_bo1, OBJECTIVE)
lower_bound = tf.cast([bb.DOMAIN_MIN], gpflow.default_float())
upper_bound = tf.cast([bb.DOMAIN_MAX], gpflow.default_float())
search_space = trieste.space.Box(lower_bound, upper_bound)

num_initial_points = 5
initial_query_points = search_space.sample(num_initial_points)
initial_data = observer(initial_query_points)
# -

# #### Model the objective function

kernel = gpflow.kernels.Matern52(lengthscales=0.4 * tf.ones(1,))
objective_model = gpflow.models.GPR(astuple(initial_data[OBJECTIVE]),
                                    kernel=kernel,
                                    noise_variance=blackbox1.noise_variance)
set_trainable(objective_model.kernel.lengthscales, False)
set_trainable(objective_model.likelihood, False)
print_summary(objective_model)

model = {OBJECTIVE: trieste.models.create_model_interface(
    {
        "model": objective_model,
        "optimizer": gpflow.optimizers.Scipy(),
        "optimizer_args": {"options": dict(maxiter=100)},
    }
)}

# #### Acquisition function
# Could try either of the two below.
# `ExpectedImprovement` is the default in Trieste.
# `NegativeLowerConfidenceBound` is probably more robust to model mis-specification.

acq_fn = NegativeLowerConfidenceBound()
# acq_fn = ExpectedImprovement()
rule = trieste.acquisition.rule.EfficientGlobalOptimization(acq_fn.using(OBJECTIVE))

# #### Optimisation loop

bo = trieste.bayesian_optimizer.BayesianOptimizer(observer, search_space)
result: OptimizationResult = bo.optimize(15, initial_data, model, acquisition_rule=rule)
if result.error is not None:
    raise result.error
dataset = result.datasets[OBJECTIVE]

# #### Explore results
# First use the optimiser's model to predict the underlying function, `f`, with uncertainty.

mean, var = model[OBJECTIVE].model.predict_f(f1['x0'].to_numpy().reshape((-1, 1)))
bo_f1 = pd.DataFrame(dict(x0=f1['x0'],
                          f=- mean.numpy()[:, 0],
                          sd=np.sqrt(var.numpy()[:, 0])))

# Next retrieve the evaluations the optimiser received from the black box.

observations = pd.DataFrame(dict(x0=dataset.query_points.numpy()[:, 0],
                                 y=- dataset.observations.numpy()[:, 0]))

# Now plot the optimiser's estimate of the underlying function in green (with error bars +/- two standard deviations),
# the black box evaluations in red and the true underlying function in blue.

chart_bo_f1 = (
    alt.Chart(bo_f1)
    .mark_line(color='green')
    .encode(x='x0:Q', y='f:Q'))
chart_bo_ferror = (
    alt.Chart(bo_f1.assign(fmin=bo_f1['f'] - 2 * bo_f1['sd'], fmax=bo_f1['f'] + 2 * bo_f1['sd']))
    .mark_errorbar(color='green')
    .encode(x="x0:Q", y="fmin:Q", y2="fmax:Q"))
chart_obs = (
    alt.Chart(observations)
    .mark_circle(color='red', size=60)
    .encode(x='x0:Q', y='y:Q'))
(chart_obs + chart_bo_ferror + chart_bo_f1 + chart1f).properties(width=800, height=300)

# ### 2-dimensional input
# If you wish, you can try the same optimisation problem but for a function of two variables.

blackbox2 = bb.GPBlackBox(ndim=2)
x0 = widgets.FloatSlider(
    value=0.5,
    min=bb.DOMAIN_MIN,
    max=bb.DOMAIN_MAX,
    step=0.01,
    description='x0:',
    readout_format='.2f',
)
x1 = widgets.FloatSlider(
    value=0.5,
    min=bb.DOMAIN_MIN,
    max=bb.DOMAIN_MAX,
    step=0.01,
    description='x1:',
    readout_format='.2f',
)
w = widgets.Box([x0, x1])

# Choose which input point (x0, x1) to evaluate the function at:

display(w)

# Evaluate the function and plot all the evaluations so far:

y = blackbox2([x0.value, x1.value])[0][0]
print(f'Evaluated black box at ({x0.value}, {x1.value}); result={y}')
blackbox2.plot_xy()

# If you have finished evaluating the function at different points and you are confident where the maximum is, you can make a guess before executing the cells below.
#
# Now show the underlying function f (without noise) as a heatmap and the noisy data we received as evaluations of it:

f2 = blackbox2.sample_f(51)
chart2f = (
    alt.Chart(f2)
    .mark_square(size=60)
    .encode(x=alt.X('x0:Q', scale=alt.Scale(domain=bb.DOMAIN)),
            y=alt.Y('x1:Q', scale=alt.Scale(domain=bb.DOMAIN)),
            color=alt.Color('f:Q', scale=alt.Scale(scheme=bb.COLOURSCHEME, domainMid=0))))
chart2y = blackbox2.plot_xy()
chart2 = alt.layer(chart2f, chart2y)
chart2.properties(width=800, height=300)



StarcoderdataJupyterScriptsDedupFiltered