Skip to content

Domain Padding

Domain Padding¤

Fast Fourier Transforms (FFT) assume periodic boundary conditions. When learning on non-periodic domains, DomainPadding pads the domain before the spectral convolutions and unpads it afterward, severely mitigating boundary artifacts.

neojax.nn.domain_padding.DomainPadding ¤

Applies domain padding to a signal.

Domain Padding helps to mitigate boundary artifacts. In the context of FNO, it is often used to handle non-periodic boundary conditions by padding the spatial dimensions before the spectral convolution and cropping the result back to the original resolution. Padding is applied symmetrically to each dimension, i.e., a 20% (0.2) padding pads an input by 20% on each side along an axis.

Parameters:

Name Type Description Default
padding float | Sequence[float]

Percentage of padding to apply to each spatial dimension on both sides. Can be a single float in [0, 1] (applied to all dims) or a sequence of floats.

 required
mode str

The type of padding to apply (e.g. "constant", "edge"). Defaults to "constant" which pads with zeros. See https://docs.jax.dev/en/latest/_autosummary/jax.numpy.pad.html#jax.numpy.pad for all padding modes.

 'constant'
resolution_scaling_factor float | int | Sequence[float | int] | None

Scaling factor(s) for layers called between pad and unpad. Used to strip the padding after the input has been scaled. Default is 1, i.e., no scaling. If resolution_scaling_factor is a sequence, each entry is the scaling factor along that spatial dimension.

 1
Internal Attributes

These fields store the internal layers state (and weights).

  • padding (float | tuple[float]): The stored padding ratios for each dimension.
  • mode (str): The padding mode.
  • resolution_scaling_factor (float | int | tuple[float | int, ...]):
is_scaled class-attribute ¤

bool(x) -> bool

Returns True when the argument x is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

mode class-attribute ¤

str(object='') -> str str(bytes_or_buffer[, encoding[, errors]]) -> str

Create a new string object from the given object. If encoding or errors is specified, then the object must expose a data buffer that will be decoded using the given encoding and error handler. Otherwise, returns the result of object.str() (if defined) or repr(object). encoding defaults to sys.getdefaultencoding(). errors defaults to 'strict'.

pad(x: Inexact[Array, 'c ...']) -> Inexact[Array, 'c ...'] ¤

Pads the input array based on the configured ratios.

Parameters:

Name Type Description Default
x Inexact[Array, 'c ...']

The input array of shape (channels, d1, ..., dN).

 required

Returns:

Type Description
Inexact[Array, 'c ...']

The padded array.
unpad(x: Inexact[Array, 'c ...'], original_shape: tuple[int, ...]) -> Inexact[Array, 'c ...'] ¤

Crops the padded array back to its original resolution.

Parameters:

Name Type Description Default
x Inexact[Array, 'c ...']

The padded array.

 required
original_shape tuple[int, ...]

The original array shape.

 required

Returns:

Type Description
Inexact[Array, 'c ...']

The cropped array of original resolution.
__call__(x: Inexact[Array, 'c ...']) -> Inexact[Array, 'c ...'] ¤

Pads the input array based on the configured ratios.

Parameters:

Name Type Description Default
x Inexact[Array, 'c ...']

The input array of shape (channels, d1, ..., dN).

 required

Returns:

Type Description
Inexact[Array, 'c ...']

The padded array.