Cholesky Inverse in Tensor Markov GP

import torch
from dmgp.utils.sparse_design.design_class import HyperbolicCrossDesign, SparseGridDesign
from dmgp.kernels.laplace_kernel import LaplaceProductKernel
from dmgp.utils.operators.chol_inv import mk_chol_inv, tmk_chol_inv

design_fun = HyperbolicCrossDesign(dyadic_sort=True, return_neighbors=True)
deg = 3
input_bd = [0,1]

dyadic_design = design_fun(deg=deg, input_bd=input_bd)
print(dyadic_design.points)
print(dyadic_design.lefts)
print(dyadic_design.rights)
print(dyadic_design.indices_sort)

markov_kernel = LaplaceProductKernel(lengthscale=1.)
Rinv = mk_chol_inv(dyadic_design=dyadic_design, markov_kernel=markov_kernel, upper=True)

Rinv_dense = Rinv.to_dense()
print(Rinv_dense[:, :5])

ker_input = dyadic_design.points
K_true = markov_kernel(ker_input, ker_input)

R_true = torch.linalg.cholesky(K_true, upper=True)
Rinv_true = torch.linalg.inv(R_true)
print(Rinv_true[:,:5])

print(torch.allclose(Rinv_dense, Rinv_true))
print((Rinv_dense-Rinv_true).norm())

Kinv_true = torch.linalg.inv(K_true)
print(Kinv_true[:5,:5])
Kinv_sp = Rinv_dense @ Rinv_dense.T
print(Kinv_sp[:5,:5])

Cholesky inverse for sparse grids

# initial setting
d = 2 # dimension
eta = 3 # level
input_bd = [[0,1]]*d # None
design_class = HyperbolicCrossDesign
dyadic_sort = True
indices_sort = True

# generate sparse grid design
sg = SparseGridDesign(d, eta, input_bd=input_bd, design_class=design_class).gen_sg(dyadic_sort=True, return_neighbors=True)
sg.pts_set

tensor_markov_kernel = LaplaceProductKernel(lengthscale=1.)
Rinv = tmk_chol_inv(sparse_grid_design=sg,
                    tensor_markov_kernel=tensor_markov_kernel,
                    upper = True)
Rinv_dense = Rinv.to_dense()
print(Rinv_dense)

ker_input = sg.pts_set
K_true = tensor_markov_kernel(ker_input, ker_input)
# print(K_true)

R_true = torch.linalg.cholesky(K_true, upper=True)
Rinv_true = torch.linalg.inv(R_true)
print(Rinv_true)

Kinv_true = torch.linalg.inv(K_true)
print(Kinv_true[:5,:5])

Kinv_sp = Rinv_dense @ Rinv_dense.T
print(Kinv_sp[:5,:5])