<jupyter_start><jupyter_text>module name here> API details.<jupyter_code>#hide
from nbdev.showdoc import *
#export
import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.distributions as dist
import torch.nn as nn
import torch.nn.functional as F
import gpytorch
#export
from pp_lvm.simulations import TrueParameters, simulate_data, transform_x
#export
device = 'cuda' if torch.cuda.is_available() else 'cpu'<jupyter_output><empty_output><jupyter_text>Helper functions<jupyter_code>#export
def add_jitter_covar(M, eps=1e-6):
    identity_matrix = torch.eye(M.shape[-1]).to(M.device)
    M = M + eps * identity_matrix
    return M
#export
def get_covariance_matrix_from_RBFkernel(M, lengthscale=20, eps=1e-6, device=device):
    x1 = M.unsqueeze(-2).div(lengthscale)
    x2 = M.unsqueeze(-1).div(lengthscale)
    K = torch.pow(x1-x2, 2).div(-2).exp()
    # Adding jitter for numerical stability
    K = add_jitter_covar(K, eps=eps)
#     K += 1e-6*torch.eye(M.shape[-1]).to(device)
    return K

def get_covariance_matrix_from_RBFkernel_new(M1, M2, lengthscale=20, eps=1e-6, device=device):
    assert (M1.shape[:-1] == M2.shape[:-1])
    x1 = M1.unsqueeze(-1).div(lengthscale)
    x2 = M2.unsqueeze(-2).div(lengthscale)
    K = torch.pow(x1-x2, 2).div(-2).exp()
    # Adding jitter for numerical stability
    if (eps > 0) and (M1.shape[-1] == M2.shape[-1]):
        K = add_jitter_covar(K, eps=eps)
    return K
#export
def plot_predictions(train_x, train_y, test_x, observed_pred, num_plots=1):
    """
    train_x: shape=(BS,S)
    train_y: shape=(BS,S)
    test_x: shape=(BS,S')
    observed_pred: shape=(BS,S')
    """
    train_x = Model_X.reshape_T(train_x, train_x.shape).cpu()
    train_y = Model_X.reshape_X(train_y, train_x.shape).cpu()
    test_x = Model_X.reshape_T(test_x, test_x.shape).cpu()

    with torch.no_grad():
        for s in range(num_plots):
            # Initialize plot
            f, ax = plt.subplots(1, 1, figsize=(8, 6))
            # Get upper and lower confidence bounds
            lower, upper = observed_pred.confidence_region()
            lower = lower.cpu()
            upper = upper.cpu()
            # Plot training data as black stars
            ax.plot(train_x.numpy()[s].flatten(), train_y.detach().numpy()[s], 'k*')
            # Plot predictive means as blue line
            ax.plot(test_x.numpy()[s].flatten(), observed_pred.mean.cpu().numpy()[s], 'b')

            ax.plot(test_x[s].flatten().numpy(), transform_x(test_x[s].flatten()).numpy(), 'r')
            # Shade between the lower and upper confidence bounds
            ax.fill_between(test_x[s].flatten().numpy(), lower[s].numpy(), upper[s].numpy(), alpha=0.5)
            ax.set_xlim([train_x[s].min()-50, train_x[s].max()+50])
            ax.set_ylim([-1.2, 1.2])
            ax.legend(['Observed Data', 'Mean', 'True', 'Confidence'])
    plt.show()<jupyter_output><empty_output><jupyter_text>Model  Exact ModelThis is where the mean and kernel functions are setup<jupyter_code>#export
class ExactGPModelLayer(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, BS):
        super().__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([BS]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([BS]), lengthscale_prior=gpytorch.priors.GammaPrior(250., 1.)),
            # gpytorch.kernels.RBFKernel(batch_shape=torch.Size([BS])),
            # gpytorch.kernels.RBFKernel(batch_shape=torch.Size([BS]), lengthscale_prior=gpytorch.priors.SmoothedBoxPrior(20, 100)),
            batch_shape=torch.Size([BS])
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        # import pdb; pdb.set_trace()
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
# likelihood = gpytorch.likelihoods.GaussianLikelihood(batch_shape=torch.Size([8]))
# model = ExactGPModelLayer(None, None, likelihood, 8)
# # model = ExactGPModelLayer(T.view(*T.shape, 1), X, likelihood, 8)

# def make_training_inputs(train_inputs):
#     return tuple(tri.unsqueeze(-1) if tri.ndimension() == 1 else tri for tri in train_inputs)

# def f():
# #     model.set_train_data(make_training_inputs((T.view(*T.shape,1),)), X, strict=0)
#     model.set_train_data((T.unsqueeze(-1)[:3],), X[:3], strict=0)
#     # model()
#     import pdb; pdb.set_trace()
#     return model(T.view(*T.shape, 1))
# f()<jupyter_output><empty_output><jupyter_text>Model Wrapper as an Estimator<jupyter_code>#export
class Model_X:
    def __init__(self, T, X, device=device):
        self.Tshape = T.shape
        self.Xshape = X.shape
        BS = T.shape[0]

        train_x = self.reshape_T(T, self.Tshape)
        train_y = self.reshape_X(X, self.Tshape)
        # initialize likelihood and model
        self.likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_prior=gpytorch.priors.NormalPrior(.01,1),
                                                                  batch_shape=torch.Size([BS])).to(device)
        self.model = ExactGPModelLayer(train_x, train_y, self.likelihood, BS).to(device)

    @staticmethod
    def reshape_T(T, Tshape):
        train_x = T.view(*Tshape, 1)
        return train_x

    @staticmethod
    def reshape_X(X, Tshape):
        train_y = X.view(Tshape)
        return train_y

    def fit(self, T, X, training_iter=1000, lr=0.1, verbose=True):
        train_x, train_y = self.reshape_T(T, self.Tshape), self.reshape_X(X, self.Tshape)

        self.model.train()
        self.likelihood.train()

        optimizer = torch.optim.Adam([
            {'params': self.model.parameters()},  # Includes GaussianLikelihood parameters
        ], lr=lr)

        mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.likelihood, self.model)

        for i in range(training_iter):
            # Zero gradients from previous iteration
            optimizer.zero_grad()
            # Output from model
            output = self.model(train_x)
            # Calc loss and backprop gradients
            loss = -mll(output, train_y).sum()
            loss.backward(retain_graph=True)
            if verbose:
                if i % (training_iter//10) ==0 :
                    print(f'Iter {i + 1}/{training_iter} - Loss: {loss.item()}, \
                    mean lengthscale: {self.model.covar_module.base_kernel.lengthscale.mean().item()}')
            optimizer.step()

        self.model.eval()
        self.likelihood.eval()

    def get_posterior(self, test_x):
        """
        test_x: shape: (BS, S, 1), BS: should be the same batch size as the training data or 1 (will broadcast)

        Output:

        """
        # Get into evaluation (predictive posterior) mode
        self.model.eval()
        self.likelihood.eval()

        with gpytorch.settings.fast_pred_var():
            observed_pred = self.likelihood(self.model(test_x))

        return observed_pred

    def predict(self, test_x):
        """
        test_x: shape: (BS, S, 1), BS: should be the same batch size as the training data or 1 (will broadcast)
        """
        # Get into evaluation (predictive posterior) mode
        self.model.eval()
        self.likelihood.eval()

        # Test points are regularly spaced along [0,1]
        # Make predictions by feeding model through likelihood
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            # test_x = torch.linspace(0, T.max(), 100).view(1,100,1).expand(10,100,1)
            observed_pred = self.likelihood(self.model(test_x))

        return observed_pred
#export
def inference_X(T, X, _model_x=None, training_iter=2000, lr=0.1, device=device, verbose=True):
    """
    T: shape=(BS,S)
    X: shape=(BS,S)
    """
    if _model_x is None:
        _model_x = Model_X(T, X, device=device)
        _model_x.fit(T, X, training_iter=training_iter, lr=lr, verbose=verbose)
    return _model_x
def get_covariance_matrix_from_RBFkernel_new(M1, M2,lengthscale=20, eps=1e-6, device=device):
#     print(M1.shape, M2.shape)
    assert (M1.shape[:-1] == M2.shape[:-1])
    x1 = M1.unsqueeze(-1).div(lengthscale)
    x2 = M2.unsqueeze(-2).div(lengthscale)
    K = torch.pow(x1-x2, 2).div(-2).exp()
    # Adding jitter for numerical stability
    if (M1.shape[-1] == M2.shape[-1]):
        K = add_jitter_covar(K, eps=eps)
    return K
def _GP_Exact(train_x, test_x, lengthscale=20, noise=1e-2):
    K = get_covariance_matrix_from_RBFkernel_new(train_x, train_x, lengthscale=lengthscale, eps=noise)
    L = torch.cholesky(K)
#     K_inv = torch.inverse(K)
    K_star = get_covariance_matrix_from_RBFkernel_new(train_x, test_x, eps=0)
#     F = torch.matmul(K_star.transpose(-1,-2), K_inv)
#     F = torch.solve(K_star, K).solution.transpose(-1,-2)
    F = torch.cholesky_solve(K_star, L).transpose(-1,-2)
    return K,K_star,F<jupyter_output><empty_output><jupyter_text>TESTS<jupyter_code>def test_res(l=20):
    device = 'cpu'
    M = torch.rand(2,10).to(device)
    
    C1 = get_covariance_matrix_from_RBFkernel(M, lengthscale=l, device=device)
    
    kernel = gpytorch.kernels.RBFKernel().initialize(lengthscale=l).to(device)
    C2 = kernel(M.unsqueeze(-1), M.unsqueeze(-1)).evaluate().detach()
    
    assert (C1.shape==C2.shape)
    print(torch.max(torch.abs(C1-C2)))
    assert (torch.max(torch.abs(C1-C2)) < 1e-4)
    print("Tests passed")
    
test_res()<jupyter_output>tensor(9.5367e-07)
Tests passed<jupyter_text>Test Run<jupyter_code>def _simulate_data_wrapper(S=8, N=400):
    TrueParams = TrueParameters(S=S, N=N, device=device)
    _,_,T,X,_,_ = simulate_data(TrueParams)
    print(f"Original shapes: T={T.shape}, X={X.shape}")
    return T,X
def train_model(T,X):
    gp_training_iters = 2000

    train_x = T[:,::10]
    train_y = X[:,::10]
    test_x = T
    print(f"Training data shapes: T={train_x.shape}, X={train_y.shape}\n")

    posterior_model_x = inference_X(train_x, train_y, training_iter=gp_training_iters, lr=1, device=device, verbose=True)
    return posterior_model_x
# class ExactGP:
#     def __init__(self, train_x, train_y):
#         self.lengthscale
    
#     def __call__(self):
# T,X = _simulate_data_wrapper()
def train_model_analytical(T,X, lengthscale=60, noise=.2):
    train_x = T[:,::20]
    train_y = X[:,::20]
    _, K_star,F = _GP_Exact(train_x, T, lengthscale=lengthscale, noise=noise)
    mu = torch.matmul(F,train_y.unsqueeze(-1)).squeeze(-1)
    covar_matrix = get_covariance_matrix_from_RBFkernel_new(T, T, eps=1e-6) - torch.matmul(F, K_star)
#     covar_matrix = add_jitter_covar(torch.relu(covar_matrix))
    return dist.MultivariateNormal(mu, covar_matrix)
# dist_post_gpytorch = train_model(T,X)
# dist_post_analytical = train_model_analytical(T, X, 30, .2)
# c=plt.imshow(dist_post_analytical.covariance_matrix[0].cpu())
# plt.colorbar(c)
# plot_predictions_analytical(T[:,::20], X[:,::20], T, dist_post_analytical)
# def confidence_region(dist_x):
#     mean_x, std_x = dist_x.mean, dist_x.stddev
# #     std_x = torch.stack([torch.diag(_cov) for _cov in covar_x], dim=0).sqrt()
#     return (mean_x - 2*std_x, mean_x + 2*std_x)

# def plot_predictions_analytical(train_x, train_y, test_x, dist_posterior):
#     s = 0
#     f, ax = plt.subplots(1, 1, figsize=(8, 6))
#     # Get upper and lower confidence bounds
#     lower, upper = confidence_region(dist_posterior)
#     lower = lower.cpu()
#     upper = upper.cpu()

#     # Plot training data as black stars
#     ax.plot(train_x.cpu().numpy()[s].flatten(), train_y.cpu().detach().numpy()[s], 'k*')
#     # Plot predictive means as blue line
#     ax.plot(T.cpu().numpy()[s].flatten(), dist_posterior.mean.cpu().numpy()[s], 'b')

#     ax.plot(T[s].cpu().flatten().numpy(), transform_x(T[s].cpu().flatten()).numpy(), 'r')
#     #             # Shade between the lower and upper confidence bounds
#     ax.fill_between(T[s].cpu().flatten().numpy(), lower[s].numpy(), upper[s].numpy(), alpha=0.5)
#     #             ax.set_xlim([train_x[s].min()-50, train_x[s].max()+50])
#     ax.set_ylim([-1.2, 1.2])
#     ax.legend(['Observed Data', 'Mean', 'True', 'Confidence'])
# posterior_model_x = train_model(T,X)<jupyter_output><empty_output><jupyter_text>Predict<jupyter_code># observed_predictions = posterior_model_x.predict(T.unsqueeze(-1))<jupyter_output><empty_output><jupyter_text>Visualize the fit<jupyter_code># plot_predictions(train_x, train_y, test_x, observed_predictions, num_plots=2)
from nbdev.export import *
notebook2script()<jupyter_output>Converted 00_core.ipynb.
Converted 01_simulations.ipynb.
Converted 02_models.ipynb.
Converted 03_GP.ipynb.
Converted index.ipynb.
{'fraction_non_alphanumeric': 0.14547744627253403, 'fraction_numerical': 0.019208061087931983, 'mean_word_length': 3.41264327891629, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '591880', 'license': 'Apache-2.0', 'n_tokens_mistral': 4794, 'n_tokens_neox': 4647, 'n_words': 823, 'path': '03_GP.ipynb', 'repo_name': 'shaabhishek/pp_lvm', 'size': '12703'}