Custom Covariance Function for Gaussian Process

Questions version agnostic
3

Thank you for your replies!
I implemented a custom kernel here but the sampling is super inefficient and compute time blows up with sample size. The idea is that the X consists of two columns, time steps and rewarded_array. The later indicates getting a reward, and triggers the reward-sensitive kernel to makes subsequent time-step outcome more similar to current outcome that gets rewarded.
My code is a brute force way of implementing this… I am guessing the bottleneck is inverting large matrix. Any ideas on how to optimize this for efficient sampling?

import pymc as pm
import numpy as np
import  pytensor.tensor as pt

# define the reward sensitive kernel
class RS_KERNEL(pm.gp.cov.Covariance):
    def __init__(self, sigma, l_rs, input_dim, active_dims=None):
        super(RS_KERNEL, self).__init__(input_dim, active_dims)
        self.sigma = sigma
        self.l_rs = l_rs

    def square_exp_kernel(self, r, sigma, l):
        """
        Compute the squared exponential kernel.
        """
        return sigma**2 * pt.exp(-r**2 / (2 * l**2))

    def diag(self, X):
        return pt.ones(X.shape[0])

    def full(self, X, Xs=None):
        if Xs is None:
            Xs = X

        # Compute squared exponential using the first column (T)
        r = pt.abs(X[:, 0][:, None] - Xs[:, 0])
        K_se = self.square_exp_kernel(r, self.sigma, self.l_rs)

        # Generate mask using the second column (rewarded_array)
        mask = pt.zeros_like(K_se)
        for i in range(X.shape[0]):
            if X[i, 1] == 1:
                mask = pt.set_subtensor(mask[i, :], 1)
                mask = pt.set_subtensor(mask[:, i], 1)

        return K_se * mask

# Example data with two columns
T = 1000
data = np.column_stack((np.arange(T), rewarded_trials))
# Now that we have the simulation process, we want to use the data to infer the parameters of the model.
import  pytensor.tensor as pt
class RS_KERNEL(pm.gp.cov.Covariance):
    def __init__(self, sigma, l_rs, input_dim, active_dims=None):
        super(RS_KERNEL, self).__init__(input_dim, active_dims)
        self.sigma = sigma
        self.l_rs = l_rs

    def square_exp_kernel(self, r, sigma, l):
        """
        Compute the squared exponential kernel.
        """
        return sigma**2 * pt.exp(-r**2 / (2 * l**2))

    def diag(self, X):
        return pt.ones(X.shape[0])

    def full(self, X, Xs=None):
        if Xs is None:
            Xs = X

        # Compute squared exponential using the first column (T)
        r = pt.abs(X[:, 0][:, None] - Xs[:, 0])
        K_se = self.square_exp_kernel(r, self.sigma, self.l_rs)

        # Generate mask using the second column (rewarded_array)
        mask = pt.zeros_like(K_se)
        for i in range(X.shape[0]):
            if X[i, 1] == 1:
                mask = pt.set_subtensor(mask[i, :], 1)
                mask = pt.set_subtensor(mask[:, i], 1)

        return K_se * mask
    


with pm.Model() as model:
    # Priors
    sig_g = pm.HalfCauchy('sig_g', beta=2.5)
    ALPHA = pm.Beta('ALPHA', alpha=1, beta=1)
    l_se = pm.Exponential('l_se', 1/20)
    l_rs = pm.Exponential('l_rs', 1/5)
    sigma_private = pm.HalfCauchy('sigma_private', beta=2.5)
    
    # Covariance functions
    cov_se = sig_g**2 * pm.gp.cov.ExpQuad(2, l_se, active_dims=[0])
    cov_rs = RS_KERNEL(sig_g, l_rs, 2, active_dims=[0,1])
    cov_private_noise = pm.gp.cov.WhiteNoise(sigma_private)
    
    # Combined covariance
    K = ALPHA * cov_se + (1 - ALPHA) * cov_rs + cov_private_noise

    # Gaussian Process
    gp = pm.gp.Marginal(cov_func=K)
    y_ = gp.marginal_likelihood('y', X=data, y=Y_result_final, noise=sigma_private)
    
    # Sampling
    trace = pm.sample(1000, chains=2, tune=1000, target_accept=0.9)