neuralop.layers.spherical_convolution.SphericalConv

class neuralop.layers.spherical_convolution.SphericalConv(in_channels, out_channels, n_modes, max_n_modes=None, bias=True, n_layers=1, separable=False, output_scaling_factor: int | float | List[float | int] | None = None, rank=0.5, factorization='cp', implementation='reconstructed', fixed_rank_modes=False, joint_factorization=False, decomposition_kwargs={}, init_std='auto', sht_norm='ortho', sht_grids='equiangular', device=None, dtype=torch.float32)[source]

Spherical Convolution, base class for the SFNO [1]

Parameters:

sht_normstr, {‘ortho’}

sht_gridsstr or str list, default is “equiangular”, {“equiangular”, “legendre-gauss”}

• If str, the same grid is used for all layers

• If list, should have n_layers + 1 values, corresponding to the input and output grid of each layer e.g. for 1 layer, [“input_grid”, “output_grid”]

See SpectralConv for full list of other parameters

References

[1]

Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere, Boris Bonev, Thorsten Kurth, Christian Hundt, Jaideep Pathak, Maximilian Baust, Karthik Kashinath, Anima Anandkumar, ICML 2023.

Attributes:

n_modes

Methods

forward(x[, indices, output_shape])	Generic forward pass for the Factorized Spectral Conv
get_conv(indices)	Returns a sub-convolutional layer from the joint parametrize main-convolution
transform(x[, layer_index, output_shape])	Transforms an input x for a skip connection, by default just an identity map

transform(x, layer_index=0, output_shape=None)[source]

Transforms an input x for a skip connection, by default just an identity map

If your function transforms the input then you should also implement this transform method so the skip connection can also work.

Typical usecases are:

• Your upsample or downsample the input in the Spectral conv: the skip connection has to be similarly scaled. This allows you to deal with it however you want (e.g. avoid aliasing)

• You perform a change of basis in your Spectral Conv, again, this needs to be applied to the skip connection too.

forward(x, indices=0, output_shape=None)[source]

Generic forward pass for the Factorized Spectral Conv

Parameters:

xtorch.Tensor

input activation of size (batch_size, channels, d1, …, dN)

indicesint, default is 0

if joint_factorization, index of the layers for n_layers > 1

Returns:

tensorized_spectral_conv(x)

get_conv(indices)[source]

Returns a sub-convolutional layer from the joint parametrize main-convolution

The parametrization of sub-convolutional layers is shared with the main one.