TC.py

import torch
import gpytorch
import math
from matplotlib import pyplot as plt
from tqdm import tqdm
import numpy as np
import tqdm

#torch.set_default_dtype(torch.float64)

class GPModelWithDerivatives(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GPModelWithDerivatives, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMeanGrad()
        self.base_kernel = gpytorch.kernels.RBFKernelGrad()
        self.base_kernel.lengthscale = 0.6
        self.covar_module = gpytorch.kernels.ScaleKernel(self.base_kernel)

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultitaskMultivariateNormal(mean_x, covar_x)


def i_vs_dim(num_dims,alpha=0.07):
    noise_constraint = gpytorch.constraints.Positive()
    likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_dims + 1, noise_constraint = gpytorch.constraints.Positive())  # Value + Derivative
    likelihood.noise = torch.tensor([1e-4])
    likelihood.task_noises = torch.tensor([1e-4] * (num_dims + 1))

    likelihood#.cuda()

    stop_iters = []
    for j in tqdm(range(10)):  # Repeat the experiment 10 times
        train_x = None
        train_y = None
      
        model = GPModelWithDerivatives(train_x, train_y, likelihood)
        model#.cuda()
        model.eval()
      
        #model.cuda(cuda1)
      
        x = torch.tensor([[0.0] * num_dims])#.cuda()
        for i in tqdm(range(1000)):
            pred = model(x)
            func_and_grad = pred.rsample()  # get f(x) and df/dx
            y = func_and_grad[-1, 0]
            dy = func_and_grad[-1, 1:]  # For higher dimensions, the gradient is d dimensional
          

            if train_x is None:
                train_x = x  # First training point
                train_y = func_and_grad  # First train_y (both function and derivative)
            else:
                train_x = torch.cat((train_x, x), dim=0)
                train_y = torch.cat((train_y, func_and_grad[-1, :].unsqueeze(-2)), dim=0)

            # Update x
            #train_x = train_x - alpha * dy
            x = x - alpha * dy

            if torch.norm(dy) < 1e-3:
                #print(f'Convergence reached after {i} iterations')
                stop_iters.append(i)
                break

            # Update model with new data
            model = GPModelWithDerivatives(train_x, train_y, likelihood)
            model#.cuda()
            model.eval()
            
    # Computes cond(K + \sigma^2 I)
    covar_check = model.covar_module(train_x).add_jitter(likelihood.noise).evaluate()
    cond_num = torch.linalg.cond(covar_check)
    print(f'Covarince: \n{covar_check}')
    print(f'Condition number: {cond_num}')  
  
    experiment_mean = torch.tensor(stop_iters).float().mean()
    print(f'Mean convergence itertations: {experiment_mean}')

    return experiment_mean, num_dims


#condition_number = []
#covariance_check = []
iterations=[]
dims=[]

with torch.no_grad(), gpytorch.settings.max_cholesky_size(10000), gpytorch.settings.max_preconditioner_size(800), gpytorch.settings.eval_cg_tolerance(0.001):  
    for i in range(1,3):
        experiment_mean, num_dims = i_vs_dim(i)
        iterations.append(experiment_mean)
        dims.append(num_dims)


fig, ax = plt.subplots(1, 1, figsize=(8,5))
ax.plot(dims, iterations, '--rv', markersize=12, lw=3)
ax.set_xticks([j+1 for j in range(len(dims))])
#ax.set_xticks([1, 2, 3, 4])
#ax.set_yticks(np.linspace(0, 400, 6))
#ax.set_yticks(np.linspace(30, iterations[len(dims)], 5))
ax.set_xlabel('dimension', fontsize = 15)
ax.set_ylabel('iteration', fontsize = 15)
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.tight_layout()

for i in range(len(dims)):
  ax.annotate('%.1f' % iterations[i] ,xy=(dims[i], iterations[i]), 
            xytext = (dims[i], iterations[i]))

plt.savefig(f'TC in high dimensions.png', bbox_inches=0, transparent=True, dpi=300)