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

import numpy as np
import pandas as pd
from math import sin
import matplotlib.pyplot as plt
# %pylab inline

# - Maximize the log probability of y given the parameters of the kernel
# - Reference: https://math.stackexchange.com/questions/1030534/gradients-of-marginal-likelihood-of-gaussian-process-with-squared-exponential-co/1072701#1072701
# - $lnP(y|x,θ)=−1/2*ln|K|−1/2*ytK−1y−N2ln2π$

# - Define the kernel

def rbf( a, b, length=2.0, train=False,
                            train_noise = 5e-3, theta=1.5):
        """Calculated the squarred exponential kernel.
        Adds a noise term for the covariance of the training data
        Adjusting the length changes the difference where points will have a positive covariance
        Returns a covaraince Matrix.
        Vertical scale controls the vertical scale of the function"""


        if train == False:
            # ensure a and b are numpy arrays
            a = np.array(a)
            b = np.array(b)
            sqdist = np.sum(a**2,1).reshape(-1,1) + np.sum(b**2,1) - 2*np.dot(a, b.T)
            return theta*np.exp(-.5 * (1/length) * sqdist)

        else:
            # ensure a and b are numpy arrays
            a = np.array(a)
            b = np.array(b)
            noisy_observations = train_noise*np.eye(len(a))
            sqdist = np.sum(a**2,1).reshape(-1,1) + np.sum(b**2,1) - 2*np.dot(a, b.T)
            return theta*np.exp(-.5 * (1/length) * sqdist) + noisy_observations


# - Define the RBF gradients

def rbf_grad_theta(a,b, theta = 1.5, length = 2.0):
    """Computes the derivative of the RBF kernel wrt the theta parameter.
    Only use with the training points (points that have a corresponding y-value)"""
    sqdist = np.sum(a**2,1).reshape(-1,1) + np.sum(b**2,1) - 2*np.dot(a, b.T)

    return 2*theta*np.exp( - sqdist / (2*length**2 ))


def rbf_grad_length(a,b, theta = 1.5, length = 2.0 ):
    """COmputes the derivative of the RBF kernel wrt the length parameter.
    Only use with the training points (points that have a corresponding y-value
    """
    sqdist = np.sum(a**2,1).reshape(-1,1) + np.sum(b**2,1) - 2*np.dot(a, b.T)
    
    return theta**2 * np.exp(-sqdist / (2*length**2)  )* ( sqdist / length**3)


# - Define test points

# +
test_x = np.array(np.linspace(-100,100,1_000)).reshape(-1,1)
train_x = np.array([[-5,10,40]]).reshape(-1,1)

train_y_func = lambda x: 1000*sin(.1*x)/abs(x+10*x+5)
train_y_func_two_d = lambda a,b : (a*b)/(a-1/(b+1)) - b
train_y_numbers = np.array([train_y_func(_) for _ in train_x]).reshape(-1,1)
y_var = 100
max_covariance = 5

# -

plt.figure(figsize=(10,8))
plt.title('Demo 1D Function')
plt.plot(np.linspace(-100,100,1_000),np.array([train_y_func(_) for _ in test_x]).reshape(-1,1) );

# ### Test the gradients
# - returns a matrix of first derivatives

rbf_grad_theta(train_x,test_x)

rbf_grad_length(train_x,test_x)


# # Define the gradient of the marginal log-likelihood and the marginal log-likelihood

def marginal_log_likehood_grad(covariance_matrix, y_values, cov_gradient_matrix):
    """
    covaraince_matrix: the covaraince given the kernel 
    y_values: returning y-vals from the underlying function
    cov_gradient_matrix: the jacobian matrix for the given parameter
    
    return the gradient of the marginal log-likelihood"""
    
    return 1/2 * y_values.T @ np.linalg.inv(covariance_matrix) @  cov_gradient_matrix @np.linalg.inv(covariance_matrix)\
            @ y_values -\
            .5* np.trace(np.linalg.inv(covariance_matrix) @ cov_gradient_matrix)


def marginal_log_likehood(covariance_matrix, y_values):
    """Return the marginal likelihood of the y-values over the latent variables"""
    from math import pi
    return -.5*np.log(abs(covariance_matrix))- (.5*y_values.T @ np.linalg.inv(covariance_matrix) @ y_values)  - \
        len(y_values)/2*np.log(2*pi)


cov_train_train = rbf(train_x,train_x)

length_grad = rbf_grad_length(train_x,train_x)
theta_grad = rbf_grad_theta(train_x,train_x)

marginal_log_likehood_grad(cov_train_train,train_y_numbers, length_grad )

marginal_log_likehood_grad(cov_train_train,train_y_numbers, theta_grad  )

# ## Hill climbing optimization

# +
original_theta = 1.5
original_length = 5
learning_rate = .1
thetas = []
lengths = []
marginal_likelihood_nums = []

for i in range(200):
    if i == 0:
        cov_train_train = rbf(train_x,train_x , theta=original_theta, length = original_length)
        length_grad = rbf_grad_length(train_x,train_x, theta=original_theta, length = original_length)
        theta_grad = rbf_grad_theta(train_x,train_x, theta=original_theta, length = original_length)
        
        new_theta = original_theta+learning_rate* marginal_log_likehood_grad(cov_train_train,train_y_numbers, 
                                                                             theta_grad )

        new_length = original_length+ learning_rate * marginal_log_likehood_grad(cov_train_train,train_y_numbers, 
                                                                                 length_grad)
        marginal_likelihood_nums.append(np.sum(marginal_log_likehood(cov_train_train,train_y_numbers)))
        
    else:
        # hill climb
        cov_train_train = rbf(train_x,train_x , theta=new_theta, length = new_length)
        length_grad = rbf_grad_length(train_x,train_x, theta=new_theta, length = new_length)
        theta_grad = rbf_grad_theta(train_x,train_x, theta=new_theta, length = new_length)
        new_theta = new_theta+learning_rate* marginal_log_likehood_grad(cov_train_train,train_y_numbers, 
                                                                             theta_grad )
        new_length = new_length+ learning_rate * marginal_log_likehood_grad(cov_train_train,train_y_numbers, 
                                                                                 length_grad)
        marginal_likelihood_nums.append(np.sum((marginal_log_likehood(cov_train_train,train_y_numbers))))
        
    print(f"New theta = {new_theta}")
    print(f"New length = {new_length}")
    print()
# -

plt.plot(marginal_likelihood_nums, label='Marginal Likelihood Over Time');
plt.legend()
plt.ylabel('Maringal Likelihood')
plt.xlabel('Iteration No.');

# +
demo_marginal_likelihoods = []
combos = []
for l in np.linspace(2,40,40):
    for t in np.linspace(2,40,40):
        cov_train_train = rbf(train_x,train_x , theta=t, length = l)
        try:
            mar_likelihood = np.sum((marginal_log_likehood(cov_train_train,train_y_numbers)))
            demo_marginal_likelihoods.append(mar_likelihood)
            combos.append((l,t))
        except:
            pass
                                
                                
        
    
# -

plt.plot(demo_marginal_likelihoods )

np.argmax(demo_marginal_likelihoods)

combos[np.argmax(demo_marginal_likelihoods)]



StarcoderdataJupyterScriptsDedupFiltered