gp2Scale#
gp2Scale is a special setting in fvgp that combines non-stationary, compactly-supported kernels, HPC distributed computing, and sparse linear algebra to allow scale-up of exact GPs to millions of data points. gp2Scale holds the world record in this category! Here we run a moderately-sized GP, just because we assume you might run this locally.
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.
##first install the newest version of fvgp
#!pip install fvgp~=4.8.7
#!pip install imate
Setup#
import numpy as np
import matplotlib.pyplot as plt
from fvgp import GP
from dask.distributed import Client
import sys
%load_ext autoreload
%autoreload 2
#further control plotting
from loguru import logger
logger.disable("fvgp")
client = Client() ##this is the client you can make locally like this or
#your HPC team can provide a script to get it. We included an example to get gp2Scale going
#on Perlmutter
#It's good practice to make sure to wait for all the workers to be ready
client.wait_for_workers(4)
Preparing the data and some other inputs#
def f1(x):
return ((np.sin(5. * x) + np.cos(10. * x) + (2.* (x-0.4)**2) * np.cos(100. * x)))
input_dim = 1
N = 2000
x_data = np.random.rand(N,input_dim)
y_data = f1(x_data).reshape(len(x_data))
hps_n = 2
hps_bounds = np.array([[0.1,1.], ##signal var of Wendland kernel
[0.001,0.04]]) ##length scale for Wendland kernel
Standard MCMC training#
from fvgp.kernels import wendland_anisotropic_gp2Scale_cpu
def kernel(x1,x2,hps):
return wendland_anisotropic_gp2Scale_cpu(x1,x2,hps)
init_hps = np.array([0.2, 0.02])
my_gp2S = GP(x_data,y_data, kernel_function=kernel,
init_hyperparameters = init_hps, #compute_device = 'gpu', #you can use gpus here
gp2Scale = True, gp2Scale_batch_size= 1000, gp2Scale_distribution="blockwise" ,
dask_client = client, linalg_mode = "Chol",
)
#optional data update
x_update = np.random.rand(5,input_dim)
y_update = f1(x_update).reshape(len(x_update))
my_gp2S.update_gp_data(x_update, y_update, append = True, rank_n_update = True)
print("Bayesian Optimization")
my_gp2S.train(hyperparameter_bounds = hps_bounds, max_iter = 20, info=True, method = "bo")
print("Trained likelihood BO: ", my_gp2S.log_likelihood())
print("MCMC")
my_gp2S.train(hyperparameter_bounds = hps_bounds, max_iter = 20, info=True, method = "mcmc")
print("Trained likelihood MCMC: ", my_gp2S.log_likelihood())
/home/marcus/Coding/fvGP/fvgp/gp.py:536: UserWarning: No noise function or measurement noise provided. Noise variances will be set to (0.01 * mean(|y_data|))^2.
self.likelihood = GPlikelihood(self.data,
Bayesian Optimization
fvGP bo: evaluating a space-filling design of 6 points in 2 hyperparameters (budget 20)
fvGP bo: design complete, best f(x)= -6016.822833873692
fvGP bo evaluation 7 of at most 20: f(x)= -6717.094181691152, best= -6717.094181691152, EI was 387.2484767340749
fvGP bo evaluation 8 of at most 20: f(x)= -6763.67124717356, best= -6763.67124717356, EI was 84.1596001727191
fvGP bo evaluation 9 of at most 20: f(x)= -6601.007082053854, best= -6763.67124717356, EI was 43.02015324233612
fvGP bo evaluation 10 of at most 20: f(x)= -6711.642931488907, best= -6763.67124717356, EI was 5.792428528176185
fvGP bo evaluation 11 of at most 20: f(x)= -6757.190165253495, best= -6763.67124717356, EI was 154.13420839781406
fvGP bo evaluation 12 of at most 20: f(x)= -6774.327434440128, best= -6774.327434440128, EI was 39.907822179018225
fvGP bo evaluation 13 of at most 20: f(x)= -6694.413723378116, best= -6774.327434440128, EI was 6.253562880815704
fvGP bo evaluation 14 of at most 20: f(x)= -6775.950510497239, best= -6775.950510497239, EI was 3.1057630314289018
fvGP bo evaluation 15 of at most 20: f(x)= -6523.688758252579, best= -6775.950510497239, EI was 0.03072534111569434
fvGP bo evaluation 16 of at most 20: f(x)= -6776.380315275647, best= -6776.380315275647, EI was 0.771727171854196
fvGP bo evaluation 17 of at most 20: f(x)= -6493.963915644109, best= -6776.380315275647, EI was 2.7736089738593374e-07
fvGP bo evaluation 18 of at most 20: f(x)= -6759.821460196373, best= -6776.380315275647, EI was 1.0619979690262273e-10
fvGP bo evaluation 19 of at most 20: f(x)= -6660.103737659295, best= -6776.380315275647, EI was 2.575620447929114e-10
fvGP bo evaluation 20 of at most 20: f(x)= -6766.527384299113, best= -6776.380315275647, EI was 0.014621153610502514
fvGP bo finished after 20 evaluations (budget): f(x)= -6776.380315275647 at [0.18506157 0.04 ]
Trained likelihood BO: 6776.380315275647
MCMC
Starting likelihood. f(x)= 5311.48816038792
Finished 10 out of 20 iterations. f(x)= 5697.8003718924565
Trained likelihood MCMC: 5783.537711105288
A custom MCMC#
from fvgp import ProposalDistribution
init_s = (np.diag(hps_bounds[:,1]-hps_bounds[:,0])/100.)**2
def obj_func(hps,args):
return my_gp2S.log_likelihood(hyperparameters=hps[0:2])
from fvgp import gpMCMC
def proposal_distribution(x0, hps, obj):
cov = obj.prop_args["prop_Sigma"]
proposal_hps = np.zeros((len(x0)))
proposal_hps = np.random.multivariate_normal(
mean = x0, cov = cov, size = 1).reshape(len(x0))
return proposal_hps
def in_bounds(v,bounds):
if any(v<bounds[:,0]) or any(v>bounds[:,1]): return False
return True
def prior_function(theta,bounds,args):
if in_bounds(theta, bounds):
return 0.
else:
return -np.inf
pd = ProposalDistribution([0,1] ,proposal_dist=proposal_distribution,
init_prop_Sigma = init_s, adapt_callable="normal")
my_mcmc = gpMCMC(obj_func, bounds=hps_bounds, prior_function=prior_function, proposal_distributions=[pd],)
logger.disable("fvgp")
hps = np.random.uniform(
low = hps_bounds[:,0],
high = hps_bounds[:,1],
size = len(hps_bounds))
mcmc_result = my_mcmc.run_mcmc(x0=hps, n_updates=110, break_condition="default", info = True)
my_gp2S.set_hyperparameters(mcmc_result["x"][-1])
Starting likelihood. f(x)= 2993.796299916403
Finished 10 out of 110 iterations. f(x)= 3297.493614936596
Finished 20 out of 110 iterations. f(x)= 3565.042450551092
Finished 30 out of 110 iterations. f(x)= 3754.3358447164146
Finished 40 out of 110 iterations. f(x)= 4228.883327513869
Finished 50 out of 110 iterations. f(x)= 4582.15306558611
Finished 60 out of 110 iterations. f(x)= 4948.885408432042
Finished 70 out of 110 iterations. f(x)= 5125.138404968175
Finished 80 out of 110 iterations. f(x)= 5278.500405120966
Finished 90 out of 110 iterations. f(x)= 5522.750494256003
Finished 100 out of 110 iterations. f(x)= 5762.610795630837
Posterior evaluation#
x_pred = np.linspace(0,1,1000) ##for big GPs, this is usually not a good idea, but in 1d, we can still do it
##It's better to do predictions only for a handful of points at a time.
mean1 = my_gp2S.posterior_mean(x_pred.reshape(1000,1))["m(x)"]
var1 = my_gp2S.posterior_covariance(x_pred.reshape(1000,1), variance_only=False)["v(x)"]
plt.figure(figsize = (16,10))
plt.plot(x_pred,mean1, label = "posterior mean", linewidth = 4)
plt.plot(x_pred,f1(x_pred), label = "latent function", linewidth = 4)
plt.fill_between(x_pred, mean1 - 3. * np.sqrt(var1), mean1 + 3. * np.sqrt(var1), alpha = 0.5, color = "grey", label = "var")
plt.scatter(x_data,y_data, color = 'black')
<matplotlib.collections.PathCollection at 0x7fcb4021f150>
