Design Class of Tensor Markov GP

import torch
from dmgp.utils.sparse_design.design_class import SparseGridDesign, n_sum_k
from dmgp.utils.sparse_design.design_class import HyperbolicCrossDesign

# test Hyperbolic Cross Design
design_class = HyperbolicCrossDesign
deg = 3
input_bd = [0,1]

print( design_class(dyadic_sort=True, return_neighbors=True)(deg=deg, input_lb=0, input_ub=1).points )
print( design_class(dyadic_sort=True, return_neighbors=True)(deg=deg, input_lb=0, input_ub=1).lefts)
print( design_class(dyadic_sort=True, return_neighbors=True)(deg=deg, input_lb=0, input_ub=1).rights )
print( design_class(dyadic_sort=True, return_neighbors=True)(deg=deg, input_lb=0, input_ub=1).indices_sort )

# initial setting
d = 2
eta = 4
input_lb = 0
input_ub = 1
design_class = HyperbolicCrossDesign

# generate sparse grid design
sg = SparseGridDesign(d, eta, input_lb=input_lb, input_ub=input_ub, design_class=design_class).gen_sg(dyadic_sort=True, return_neighbors=True)
x_tot = sg.pts_tot
id_prt = sg.id_prt # use self.tot_pts[ id_x_prt[t_sum, prt], : ] to extract grid points in each smolyak iter
pts_prt = sg.pts_prt # use self.pts_prt[t_sum, prt] to extract a d-dimensional list, each entry is one-dim points forming the sgdesign
pts_prt_set = sg.pts_prt_set

pts_set = sg.pts_set
n_pts = sg.n_pts

design_str_prt = sg.design_str_prt


indices_prt = sg.indices_prt
indices_prt_set = sg.indices_prt_set

indices_tot = sg.indices_tot
indices_set = sg.indices_set
print(n_pts)

pts_set

design_str_prt[eta,2][0].points

pts_prt[eta,2]

indices_prt[eta,2]

print(indices_set)

indices_select = indices_prt_set[3,1]
print(indices_select)

list_full = [tuple(l) for l in indices_set.tolist()] # [n_u, d] size list of tuple
list_select = [tuple(l) for l in indices_select.tolist()] # [n_arrow, d] size list of tuple

index_dict = dict((value, idx) for idx,value in enumerate(list_full))
index_select = torch.tensor( [index_dict[x] for x in list_select] )
index_select

i = torch.tensor([[0, 1, 2, 0],
                   [1, 1, 0, 2]])
v = torch.tensor([1, 2, 3, 4], dtype=float)
sp_test = torch.sparse_coo_tensor(i, v, [5, 5])
a = torch.randn((5,5), dtype=float)
print(a)
print(sp_test)
print(sp_test @ a)

sp_id = sp_test._indices().T
print(sp_id)
print(index_select)

index_select[sp_id]

sp_test_new = torch.sparse_coo_tensor(index_select[sp_test._indices()], sp_test._values(), [17, 17])
sp_test_new