neuralop.layers.discrete_continuous_convolution.EquidistantDiscreteContinuousConv2d

class neuralop.layers.discrete_continuous_convolution.EquidistantDiscreteContinuousConv2d(in_channels: int, out_channels: int, in_shape: Tuple[int], out_shape: Tuple[int], kernel_shape: int | List[int], basis_type: str = 'piecewise_linear', domain_length: Tuple[float] | None = None, periodic: bool | None = False, groups: int | None = 1, bias: bool | None = True, radius_cutoff: float | None = None, **kwargs)[source]

Discrete-continuous convolutions (DISCO) on equidistant 2d grids as implemented for [1]. This implementation maps to 2d convolution kernels which makes it more efficient than the unstructured implementation above. Due to the mapping to an equidistant grid, the domain lengths need to be specified in order to compute the effective resolution and the corresponding cutoff radius. Forward call expects an input of shape (batch_size, in_channels, in_shape[0], in_shape[1]).

Parameters:

in_channels: int

input channels to DISCO convolution

out_channels: int

output channels of DISCO convolution

in_shape: Tuple[int]

shape of the (regular) input grid.

out_shape: torch.Tensor or str

shape of the (regular) output grid. Note that the side lengths of out_shape must be less than or equal to the side lengths of in_shape, and must be integer divisions of the corresponding in_shape side lengths.

kernel_shapeUnion[int, List[int]]

Dimensions of the convolution kernel, either one int or a two-int tuple. * If one int k, the kernel will be a square of shape (k,k), meaning the convolution will be ‘isotropic’: both directions are equally scaled in feature space.

• If two ints (k1,k2), the kernel will be a rectangle of shape (k1,k2), meaning the convolution

will be ‘anisotropic’: one direction will be compressed or stretched in feature space.

basis_type: str literal ``{‘piecewise_linear’, ‘morlet’, ‘zernike’}``

choice of basis functions to use for convolution filter tensor.

domain_length: torch.Tensor, optional

extent/length of the physical domain. Assumes square domain [-1, 1]^2 by default

periodic: bool, optional

whether the domain is periodic, by default False

groups: int, optional

number of groups in the convolution, by default 1

bias: bool, optional

whether to use a bias, by default True

radius_cutoff: float, optional

cutoff radius for the kernel. For a point x on the input grid, every point y on the output grid with ||x - y|| <= radius_cutoff will be affected by the value at x. By default, set to 2 / sqrt(# of output points)

Methods

forward(x)	Forward call.
get_local_filter_matrix()	Returns the precomputed local convolution filter matrix Psi.

References

[1]

Liu-Schiaffini M., Berner J., Bonev B., Kurth T., Azizzadenesheli K., Anandkumar A.; Neural Operators with Localized Integral and Differential Kernels; arxiv:2402.16845

get_local_filter_matrix()[source]

Returns the precomputed local convolution filter matrix Psi. Psi parameterizes the kernel function as triangular basis functions evaluated on pairs of points on the convolution’s input and output grids, such that Psi[l, i, j] is the l-th basis function evaluated on point i in the output grid and point j in the input grid.

forward(x: Tensor) → Tensor[source]

Forward call. Expects an input of shape batch_size x in_channels x in_shape[0] x in_shape[1].