###############################################################################
# .. note::
# This operator uses a conjugate gradient solver and assumes
# that **formula** defines a **symmetric**, positive and definite
# **linear** reduction with respect to the alias ``"b"``
# specified trough the third argument.
#
# Apply our solver on arbitrary point clouds:
#
print(
"Solving a Gaussian linear system, with {} points in dimension {}.".format(
N, D))
start = time.time()
K_xx = keops.exp(-keops.sum((Vi(x) - Vj(x))**2, dim=2) / (2 * sigma**2))
cfun = keops.solve(K_xx, Vi(b), alpha=alpha, call=False)
c = cfun()
end = time.time()
print("Timing (KeOps implementation):", round(end - start, 5), "s")
###############################################################################
# Compare with a straightforward PyTorch implementation:
#
start = time.time()
K_xx = alpha * torch.eye(N) + torch.exp(-torch.sum(
(x[:, None, :] - x[None, :, :])**2, dim=2) / (2 * sigma**2))
c_py = torch.solve(b, K_xx)[0]
end = time.time()
print("Timing (PyTorch implementation):", round(end - start, 5), "s")
print("Relative error = ", (torch.norm(c - c_py) / torch.norm(c_py)).item())
########################################################################
# The ``[]`` operator can be used to do element selection or slicing
# (Elem or Extract operation in KeOps).
res = (x_i[:2] * y_j[2:] - x_i[2:] * y_j[:2]).sqnorm2().sum(axis=1)
########################################################################
# Kernel inversion : let's do a gaussian kernel inversion. Note that
# we have to use both :func:`Vi <pykeops.torch.Vi>` and :func:`Vj <pykeops.torch.Vj>` helpers on the same tensor ``x`` here.
#
e_i = Vi(torch.rand(M, D).type(tensor))
x_j = Vj(x)
D2xx = LazyTensor.sum((x_i - x_j)**2)
sigma = 0.25
Kxx = (-D2xx / sigma**2).exp()
res = LazyTensor.solve(Kxx, e_i, alpha=0.1)
#########################################################################
# Use of loops or vector operations for sums of kernels
# -----------------------------------------------------
#############################################################################
# Let us now perform again a kernel convolution, but replacing the gaussian
# kernel by a sum of 4 gaussian kernels with different sigma widths.
# This can be done as follows with a for loop:
sigmas = tensor([0.5, 1.0, 2.0, 4.0])
b_j = Vj(torch.rand(N, D).type(tensor))
Kxy = 0
for sigma in sigmas:
Kxy += LazyTensor.exp(-D2xy / sigma**2)
gamma = (Kxy * b_j).sum_reduction(axis=1)