Non-Stationary Kernel¶

In [ ]:

import tensorflow as tf
import numpy as np
import gpflow
gpflow.config.set_default_float(np.float32)

from sgptools.utils.metrics import *
from sgptools.utils.misc import *
from sgptools.models.continuous_sgp import *
from sgptools.kernels.neural_kernel import init_neural_kernel
from sgptools.utils.gpflow import get_model_params

from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt

np.random.seed(10)
tf.random.set_seed(10)

import tensorflow as tf import numpy as np import gpflow gpflow.config.set_default_float(np.float32) from sgptools.utils.metrics import * from sgptools.utils.misc import * from sgptools.models.continuous_sgp import * from sgptools.kernels.neural_kernel import init_neural_kernel from sgptools.utils.gpflow import get_model_params from sklearn.preprocessing import StandardScaler import matplotlib.pyplot as plt np.random.seed(10) tf.random.set_seed(10)

Generate synthetic non-stationary data¶

In [2]:

def non_stationary_function(X, Y):
    Z1 = np.sin(X/10) + np.cos(Y/10)
    Z2 = np.cos(X/6) + np.cos(Y/6)
    Z = np.concatenate([Z1[:, ::2], Z2[:, ::2]], 
                       axis=1)
    return Z

# Generate data
train_data_dims = (40, 20)
x = np.linspace(0, 200, train_data_dims[0])
y = np.linspace(0, 100, train_data_dims[1])
X, Y = np.meshgrid(x, y)
Z = non_stationary_function(X, Y)

X_train = np.stack([X.ravel(), Y.ravel()], 
                   axis=1).astype(np.float32)
Y_train = Z.ravel()[:, None].astype(np.float32)
print('Train Data:', X_train.shape, Y_train.shape)

test_data_dims = (100, 50)
x = np.linspace(0, 200, test_data_dims[0])
y = np.linspace(0, 100, test_data_dims[1])
X, Y = np.meshgrid(x, y)
Z = non_stationary_function(X, Y)

X_test = np.stack([X.ravel(), Y.ravel()], 
                  axis=1).astype(np.float32)
Y_test = Z.ravel()[:, None].astype(np.float32)
print('Test Data:', X_test.shape, Y_test.shape)

# Normalize the data
Xscalar = StandardScaler()
X_train = Xscalar.fit_transform(X_train)
X_test = Xscalar.transform(X_test)

yScalar = StandardScaler()
Y_train = yScalar.fit_transform(Y_train)
Y_test = yScalar.transform(Y_test)

# Plot the data
plt.figure()
plt.imshow(Y_train.reshape(train_data_dims[1], 
                           train_data_dims[0]), 
           cmap='jet')
plt.gca().set_aspect('equal')
plt.title('Train Data')
plt.show()

plt.figure()
plt.imshow(Y_test.reshape(test_data_dims[1], 
                          test_data_dims[0]), 
           cmap='jet')
plt.gca().set_aspect('equal')
plt.title('Test Data')
plt.show()

def non_stationary_function(X, Y): Z1 = np.sin(X/10) + np.cos(Y/10) Z2 = np.cos(X/6) + np.cos(Y/6) Z = np.concatenate([Z1[:, ::2], Z2[:, ::2]], axis=1) return Z # Generate data train_data_dims = (40, 20) x = np.linspace(0, 200, train_data_dims[0]) y = np.linspace(0, 100, train_data_dims[1]) X, Y = np.meshgrid(x, y) Z = non_stationary_function(X, Y) X_train = np.stack([X.ravel(), Y.ravel()], axis=1).astype(np.float32) Y_train = Z.ravel()[:, None].astype(np.float32) print('Train Data:', X_train.shape, Y_train.shape) test_data_dims = (100, 50) x = np.linspace(0, 200, test_data_dims[0]) y = np.linspace(0, 100, test_data_dims[1]) X, Y = np.meshgrid(x, y) Z = non_stationary_function(X, Y) X_test = np.stack([X.ravel(), Y.ravel()], axis=1).astype(np.float32) Y_test = Z.ravel()[:, None].astype(np.float32) print('Test Data:', X_test.shape, Y_test.shape) # Normalize the data Xscalar = StandardScaler() X_train = Xscalar.fit_transform(X_train) X_test = Xscalar.transform(X_test) yScalar = StandardScaler() Y_train = yScalar.fit_transform(Y_train) Y_test = yScalar.transform(Y_test) # Plot the data plt.figure() plt.imshow(Y_train.reshape(train_data_dims[1], train_data_dims[0]), cmap='jet') plt.gca().set_aspect('equal') plt.title('Train Data') plt.show() plt.figure() plt.imshow(Y_test.reshape(test_data_dims[1], test_data_dims[0]), cmap='jet') plt.gca().set_aspect('equal') plt.title('Test Data') plt.show()

Train Data: (800, 2) (800, 1)
Test Data: (5000, 2) (5000, 1)

Learn kernel parameters from train data¶

In [3]:

# Model / inference hyperparameters
num_inducing = 100
lr = 1e-2
max_steps = 5000
Q = 3

# Initialize non-stationary kernel model
print('Initializing model...')
Z = get_inducing_pts(X_train, num_inducing)
model = init_neural_kernel(X_train, Y_train, 
                           n_inits=5,
                           inducing_variable=Z,
                           Q=Q, 
                           hidden_sizes=(4, 4))

# Create optimizer
learning_rate = tf.keras.optimizers.schedules.ExponentialDecay(
                      lr, 
                      decay_steps=50, 
                      decay_rate=0.98)
loss = optimize_model(model, 
                      max_steps=max_steps,
                      lr=learning_rate,
                      optimizer='tf')

# Plot loss
plt.plot(loss)
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.show()

kernels = []
noice_vars = []
labels = ['Non Stationary Kernel', 'Stationary Kernel']

kernels.append(model.kernel)
noice_vars.append(model.likelihood.variance.numpy())

# Get stationary kernel
_, noise_variance, kernel = get_model_params(X_train, Y_train, 
                                             lengthscales=[1.0, 0.5])
kernels.append(kernel)
noice_vars.append(noise_variance)

# Model / inference hyperparameters num_inducing = 100 lr = 1e-2 max_steps = 5000 Q = 3 # Initialize non-stationary kernel model print('Initializing model...') Z = get_inducing_pts(X_train, num_inducing) model = init_neural_kernel(X_train, Y_train, n_inits=5, inducing_variable=Z, Q=Q, hidden_sizes=(4, 4)) # Create optimizer learning_rate = tf.keras.optimizers.schedules.ExponentialDecay( lr, decay_steps=50, decay_rate=0.98) loss = optimize_model(model, max_steps=max_steps, lr=learning_rate, optimizer='tf') # Plot loss plt.plot(loss) plt.xlabel('Iteration') plt.ylabel('Loss') plt.show() kernels = [] noice_vars = [] labels = ['Non Stationary Kernel', 'Stationary Kernel'] kernels.append(model.kernel) noice_vars.append(model.likelihood.variance.numpy()) # Get stationary kernel _, noise_variance, kernel = get_model_params(X_train, Y_train, lengthscales=[1.0, 0.5]) kernels.append(kernel) noice_vars.append(noise_variance)

Initializing model...
Initializing neural spectral kernel...

2024-05-05 12:08:35.281864: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1960] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...

Best init: -1165.898315

╒═════════════════════════╤═══════════╤══════════════════╤═════════╤═════════════╤═════════╤═════════╤═══════════════════╕
│ name                    │ class     │ transform        │ prior   │ trainable   │ shape   │ dtype   │ value             │
╞═════════════════════════╪═══════════╪══════════════════╪═════════╪═════════════╪═════════╪═════════╪═══════════════════╡
│ GPR.kernel.variance     │ Parameter │ Softplus         │         │ True        │ ()      │ float32 │ 0.59679           │
├─────────────────────────┼───────────┼──────────────────┼─────────┼─────────────┼─────────┼─────────┼───────────────────┤
│ GPR.kernel.lengthscales │ Parameter │ Softplus         │         │ True        │ (2,)    │ float32 │ [1.16673 0.37082] │
├─────────────────────────┼───────────┼──────────────────┼─────────┼─────────────┼─────────┼─────────┼───────────────────┤
│ GPR.likelihood.variance │ Parameter │ Softplus + Shift │         │ True        │ ()      │ float32 │ 0.52522           │
╘═════════════════════════╧═══════════╧══════════════════╧═════════╧═════════════╧═════════╧═════════╧═══════════════════╛

Get Continuous-SGP solution with a baseline stationary kernel and the learned neural kernel function¶

In [4]:

for i, (kernel, noice_var) in enumerate(zip(kernels, noice_vars)):
        print(labels[i])
        for num_placements in [4, 9]:
                # Get initial inducing points
                Xu_init = get_inducing_pts(X_train, num_placements)*0.5
                Xu_init.astype(np.float32)

                # Setup SGP and optimize it
                sgpr, _ = continuous_sgp(num_placements, 
                                         X_train,
                                         noise_variance=noice_var,
                                         kernel=kernel,
                                         Xu_init=Xu_init)
                sgp_sol_sp = sgpr.inducing_variable.Z.numpy()

                # Get the GP predictions
                Xu_X, Xu_y = cont2disc(sgp_sol_sp, candidates=X_test,
                                       candidate_labels=Y_test)
                gpr = gpflow.models.GPR((Xu_X, Xu_y),
                                        noise_variance=noice_var,
                                        kernel=kernel)
                y_pred, y_var = gpr.predict_f(X_test)
                y_pred = y_pred.numpy()

                rmse = np.sqrt(np.mean((y_pred - Y_test)**2))

                # Plot the results
                plt.figure()
                plt.imshow(y_pred.reshape(test_data_dims[1], test_data_dims[0]),
                        cmap='jet', origin='lower',
                        extent=[np.min(X_test[:, 0]), np.max(X_test[:, 0]),
                                np.min(X_test[:, 1]), np.max(X_test[:, 1])])
                plt.gca().set_aspect('equal')
                plt.scatter(sgp_sol_sp[:, 0], sgp_sol_sp[:, 1], c='k', marker='p', s=100)
                plt.title('Number of Sensors: {}    RMSE: {:.3f}'.format(num_placements, rmse))
                plt.xlabel('X')
                plt.ylabel('Y')
                plt.show()

for i, (kernel, noice_var) in enumerate(zip(kernels, noice_vars)): print(labels[i]) for num_placements in [4, 9]: # Get initial inducing points Xu_init = get_inducing_pts(X_train, num_placements)*0.5 Xu_init.astype(np.float32) # Setup SGP and optimize it sgpr, _ = continuous_sgp(num_placements, X_train, noise_variance=noice_var, kernel=kernel, Xu_init=Xu_init) sgp_sol_sp = sgpr.inducing_variable.Z.numpy() # Get the GP predictions Xu_X, Xu_y = cont2disc(sgp_sol_sp, candidates=X_test, candidate_labels=Y_test) gpr = gpflow.models.GPR((Xu_X, Xu_y), noise_variance=noice_var, kernel=kernel) y_pred, y_var = gpr.predict_f(X_test) y_pred = y_pred.numpy() rmse = np.sqrt(np.mean((y_pred - Y_test)**2)) # Plot the results plt.figure() plt.imshow(y_pred.reshape(test_data_dims[1], test_data_dims[0]), cmap='jet', origin='lower', extent=[np.min(X_test[:, 0]), np.max(X_test[:, 0]), np.min(X_test[:, 1]), np.max(X_test[:, 1])]) plt.gca().set_aspect('equal') plt.scatter(sgp_sol_sp[:, 0], sgp_sol_sp[:, 1], c='k', marker='p', s=100) plt.title('Number of Sensors: {} RMSE: {:.3f}'.format(num_placements, rmse)) plt.xlabel('X') plt.ylabel('Y') plt.show()

Non Stationary Kernel

Stationary Kernel