nifty.re.multi_grid.kernel module

class ICRKernel(grid, covariance, *, window_size=None, _cim=None)

Bases: Kernel

Full ICR implementation taking an arbitrary grid and a covariance function.

__init__(grid, covariance, *, window_size=None, _cim=None)

compute_matrices(index, level)

Kernel matrices for a given index and level.

Computes and returns the matrix elements required to refine index onto it’s children. These matrix elements are to be combined with array elements gathered from the output of self.get_output_input_indices.

property covariance_outer

get_kernel_distmat(index, level)

Distance matrix between fine and coarse points.

Distances between all points involved in one local refinement operation. In addition to the distance matrix, the number of features in the output grid is returned, defining the number of columns in the distance matrix that belong to the fine points.

get_output_input_indices(index, level)

Index mappings defining matrix elements to refine index.

Each index at level is refined onto children gout at level + 1. The indices of neighbors at level and inputs at level + 1 needed to gather the array elements of the input are returned.

replace(*, covariance=None, window_size=None, _cim=None)

property window_size

class Kernel(grid, *, _cim: _CompressedIndexMap | None = None)

Bases: object

• Apply linear operators with inputs accross an arbitrary grid

• A Kernel could be as simple as coarse graining children to a parent; or a full ICR or MSC implementation

• Fully jax transformable to allow seamlessly being used in larger models

__init__(grid, *, _cim: _CompressedIndexMap | None = None)

compress(*args, **kwargs)

Optimize the kernel matrices for the given grid.

This links kernel matrices by similarity and subsequently looks them up via tables. Compressing the kernel in this way has the potential to drastically reduce the memory requirements of applying it and make the inference of kernels faster.

compress_indices(*, rtol=1e-05, atol=1e-10, buffer_size=10000, use_distances=True, distance_norm=functools.partial(<PjitFunction of <function norm>>, axis=0))

Link kernel matrices to preserve memory and increase speed.

You probably don’t want to access this function directly. Instead use compress or replace and compress_matrices.

compress_matrices()

Recompute the kernel matrices while keeping the tables for speedy lookups fixed.

This is useful for updating the covariance function of a kernel without recreating the tables in an expensive optimization.

property compressed: bool

compute_matrices(index, level)

Compute kernel matrices from scratch.

get_matrices(index, level)

Compute or lookup kernel matrices.

get_output_input_indices(index, level: int) → Tuple[Tuple[NDArray[int64], int], Tuple[Tuple[NDArray[int64], int], ...]]

property grid: Grid

lookup_matrices(index, level)

Efficient retrieval of kernel matrices for a compressed kernel.

replace(*, _cim=None, **kwargs)

apply_kernel(x, *, kernel, indices=None)

Applies the kernel to values on an entire multigrid

refinement_matrices(cov, n_fsz: int)

tree_flatten(self)

tree_unflatten(aux, children)