imbirik / liestationarykernels Goto Github PK

x = B.linspace(0, 2*3.14159, 101)
xx = np.meshgrid(x,x)
X = B.stack(*xx, axis=-1).reshape((-1, 2))

import sys
import torch
from lie_stationary_kernels.spaces.torus import Torus
from lie_stationary_kernels.spectral_kernel import RandomSpectralKernel, RandomFourierFeatureKernel, RandomPhaseKernel
from lie_stationary_kernels.spectral_kernel import EigenbasisSumKernel, EigenbasisKernel
from lie_stationary_kernels.prior_approximation import RandomFourierApproximation
from lie_stationary_kernels.prior_approximation import RandomPhaseApproximation
from lie_stationary_kernels.spectral_measure import MaternSpectralMeasure, SqExpSpectralMeasure
import matplotlib.pyplot as plt
import os
import lab as B
import numpy as np

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="1"
torch.autograd.set_detect_anomaly(True)
torch.cuda.set_device("cuda:1")
dtype = torch.float64
device = 'cuda' if torch.cuda.is_available() else 'cpu'

order = 100
space = Torus(n=1, order=order)
print(space.dim)

num = 300
theta = B.linspace(0, 2*3.14159, num).reshape(-1,1)

Y = np.sin(theta)
X = np.cos(theta)
Z = np.zeros_like(Y)

lengthscale, nu = 1, 1.0/2
measure = MaternSpectralMeasure(space.dim, lengthscale, nu)
circle_kernel = EigenbasisSumKernel(measure, space)
circle_sampler = RandomPhaseApproximation(circle_kernel)
index = 20

point = theta[index].reshape(1,-1)
print(theta.shape, point.shape)
circle_matrix = circle_kernel(torch.tensor(theta, device=device, dtype=dtype), torch.tensor(point, device=device, dtype=dtype)).detach().cpu().numpy()

off_set = np.min(circle_matrix)
print(off_set)
k_point = (circle_matrix - off_set).reshape(-1, 1)

fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(X, Y, Z)
ax.scatter(X[0], Y[0], Z[0],  c='r', marker = 'D', s=50)
ax.scatter(X, Y, k_point,  c='k', marker = 'o', s=50)
ax.axis('off')
plt.show()