Kernel product helper

pykeops.torch - The kernel_product helper:

Summary

Kernel([name, formula_sum, routine_sum, …])

Defines a new Kernel identifier for kernel_product().

kernel_product(params, x, y, *bs[, mode, …])

Math-friendly wrapper around the torch.Genred constructor.

Syntax

class pykeops.torch.Kernel(name=None, formula_sum=None, routine_sum=None, formula_log=None, routine_log=None)[source]

Defines a new Kernel identifier for kernel_product().

Keyword Arguments

name (string) –

Computation identifier. The kernel name should be built from a small set of atomic formulas, acting on arbitrary pairs of variables and combined using:

  • integer constants,

  • the addition +,

  • the product *,

  • the integer exponentiation **k.

Parameters and variables. Every kernel name is associated to a list of atomic formulas (that will require parameters) and a list of pairs of variables, ordered as they are in the name string. Both parameters and variables will be required as inputs by kernel_product(). A few examples:

  • "gaussian(x,y)" : one formula and one pair of variables.

  • "gaussian(x,y) * linear(u,v)**2" : two formulas and two pairs of variables.

  • "cauchy(x,y) + gaussian(x,y) * (1 + cauchy(u,v)**2): three formulas (cauchy, gaussian and cauchy once again) with two pairs of variables ((x,y) first, (u,v) second)

Note that by convention, pairs of variables should be denoted by single-letter, non-overlapping duets: "gaussian(x',yy)" or "gaussian(x,y) + cauchy(y,z)" are not supported.

Atomic formulas. As of today, the pre-defined kernel names are:

  • linear(x,y), the \(L^2\) scalar product:

    \[k(x,y)=\langle x,y\rangle.\]
  • gaussian(x,y), the standard RBF kernel:

    \[k(x,y)=\exp(-\langle x-y, G\, (x-y)\rangle).\]
  • laplacian(x,y), the pointy exponential kernel:

    \[k(x,y)=\exp(-\sqrt{\langle x-y, G\, (x-y)\rangle}).\]
  • cauchy(x,y), a heavy-tail kernel:

    \[k(x,y)=1/(1+\langle x-y, G\, (x-y)\rangle).\]
  • inverse_multiquadric(x,y), a very heavy-tail kernel:

    \[k(x,y)=1/\sqrt{1+\langle x-y, G\, (x-y)\rangle}.\]
  • distance(x,y), arbitrary Euclidean norms:

    \[k(x,y)=\sqrt{\langle x-y, G\, (x-y)\rangle}.\]

Defining your own formulas is also possible, and documented in the second part of this example.

Parameters. With the exception of the linear kernel (which accepts None as its parameter), all these kernels act on arbitrary vectors of dimension D and are parametrized by a variable G that can represent :

Parameter \(G\)

Dimension of the tensor G

scalar

dim-1 vector

diagonal matrix

dim-D vector

symmetric D-by-D matrix

dim-D*D vector

j-varying scalar

N-by-1 array

j-varying diagonal matrix

N-by-D array

j-varying symmetric D-by-D matrix

N-by-D*D array

If required by the user, a kernel-id can thus be used to represent non-uniform, non-radial kernels as documented in the anisotropic_kernels example.

Example

>>> M, N = 1000, 2000 # number of "i" and "j" indices
>>> # Generate the data as pytorch tensors.
>>> #
>>> # First, the "i" variables:
>>> x = torch.randn(M,3) # Positions,    in R^3
>>> u = torch.randn(M,3) # Orientations, in R^3 (for example)
>>> #
>>> # Then, the "j" ones:
>>> y = torch.randn(N,3) # Positions,    in R^3
>>> v = torch.randn(N,3) # Orientations, in R^3
>>> #
>>> # The signal b_j, supported by the (y_j,v_j)'s
>>> b = torch.randn(N,4)
>>> #
>>> # Pre-defined kernel: using custom expressions is also possible!
>>> # Notice that the parameter sigma is a dim-1 vector, *not* a scalar:
>>> sigma  = torch.tensor([.5])
>>> params = {
...     # The "id" is defined using a set of special function names
...     "id"      : Kernel("gaussian(x,y) * (linear(u,v)**2) "),
...     # gaussian(x,y) requires a standard deviation; linear(u,v) requires no parameter
...     "gamma"   : ( 1./sigma**2 , None ) ,
... }
>>> #
>>> # Don't forget to normalize the orientations:
>>> u = torch.nn.functional.normalize(u, p=2, dim=1)
>>> v = torch.nn.functional.normalize(v, p=2, dim=1)
>>> #
>>> # We're good to go! Notice how we grouped together the "i" and "j" features:
>>> a = kernel_product(params, (x,u), (y,v), b)
>>> print(a)
pykeops.torch.kernel_product(params, x, y, *bs, mode=None, backend=None, dtype='float32', cuda_type=None)[source]

Math-friendly wrapper around the torch.Genred constructor.

This routine allows you to compute kernel dot products (aka. as discrete convolutions) with arbitrary formulas, using a Sum or a LogSumExp reduction operation. It is syntactic sugar, meant to ease the implementation of mixture models on point clouds.

Its use is explained in the documentation and showcased in the anisotropic kernels and GMM-fitting tutorials.

Having created a kernel id, and with a few torch tensors at hand, you can feed the kernel_product() numerical routine with the appropriate input. More precisely, if Kernel("my_kernel_name...") defines a kernel with F formulas and V variable pairs, kernel_product() will accept the following arguments:

Parameters
  • params (dict) –

    Describes the kernel being used. It should have the following attributes :

    • "id" (Kernel): Describes the formulas for \(k(x_i,y_j)\), \(c(x_i,y_j)\) and the number of features F (locations, location+directions, etc.) being used.

    • "gamma" (tuple of Tensors): Parameterizes the kernel. Typically, something along the lines of (.5/σ**2, .5/τ**2, .5/κ**2) for Gaussian or Laplacian kernels…

  • x (Tensor, or F-tuple of Tensors) – The F features associated to points \(x_i\), for \(i\in [0,M)\). All feature Tensors should be 2d, with the same number of lines \(M\) and an arbitrary number of columns.

  • y (Tensor, or F-tuple of Tensors) – The F features associated to points \(y_j\), for \(j\in [0,N)\). All feature Tensors should be 2d, with the same number of lines \(N\) and an arbitrary number of columns, compatible with x.

  • b ((N,E) Tensor) – The weights, or signal vectors \(b_j\) associated to the points \(y_j\).

  • a_log ((M,1) Tensor, optional) – If mode is one of the "log_*" reductions, specifies the scalar variable \(\text{Alog}_i\).

  • b_log ((N,1) Tensor, optional) – If mode is one of the "log_*" reductions, specifies the scalar variable \(\text{Blog}_j\).

Keyword Arguments
  • mode (string) –

    Specifies the reduction operation. The supported values are:

    • "sum" (default):

      \[v_i = \sum_j k(x_i, y_j) \cdot b_j\]
    • "lse":

      \[v_i = \log \sum_j \exp(c(x_i, y_j) + b_j )\]

      with \(c(x_i,y_j) = \log(k(x_i,y_j) )\)

    • "log_scaled":

      \[v_i = \sum_j \exp(c(x_i,y_j) + \text{Alog}_i + \text{Blog}_j ) \cdot b_j\]
    • "log_scaled_lse":

      \[v_i = \log \sum_j \exp(c(x_i,y_j) + \text{Alog}_i + \text{Blog}_j + b_j )\]

  • backend (string, default="auto") –

    Specifies the implementation to run. The supported values are:

    • "auto": to use libkp’s CPU or CUDA routines.

    • "pytorch": to fall back on a reference tensorized implementation.

  • dtype (string, default = "float32") –

    Specifies the numerical dtype of the input and output arrays. The supported values are:

    • dtype = "float32",

    • dtype = "float64".

Returns

The output scalar or vector \(v_i\) sampled on the \(x_i\)’s.

Return type

(M,E) Tensor

Example

>>> # Generate the data as pytorch tensors
>>> x = torch.randn(1000,3, requires_grad=True)
>>> y = torch.randn(2000,3, requires_grad=True)
>>> b = torch.randn(2000,2, requires_grad=True)
>>> #
>>> # Pre-defined kernel: using custom expressions is also possible!
>>> # Notice that the parameter sigma is a dim-1 vector, *not* a scalar:
>>> sigma  = torch.tensor([.5], requires_grad=True)
>>> params = {
...    "id"      : Kernel("gaussian(x,y)"),
...    "gamma"   : .5/sigma**2,
... }
>>> #
>>> # Depending on the inputs' types, 'a' is a CPU or a GPU variable.
>>> # It can be differentiated wrt. x, y, b and sigma.
>>> a = kernel_product(params, x, y, b)
>>> print(a)
tensor([[-0.0898, -0.3760],
        [-0.8888, -1.3352],
        [ 1.0236, -1.3245],
        ...,
        [ 2.5233, -2.6578],
        [ 1.3097,  4.3967],
        [ 0.4095, -0.3039]], grad_fn=<GenredAutogradBackward>)