xnn.cnn.models.spherical.S2Convolution#

class xnn.cnn.models.spherical.S2Convolution(n_in, n_out, b_in, b_out, points)[source]#

Bases: Module

S^2 correlation layer (eq 4, computed by the Fourier theorem of Appendix D).

The filter of every (input, output) channel pair is a set of values on the points of the sphere; its spectrum is formed with s2_kernel_basis(), multiplied with the input spectrum as the outer product f^l psi^l+ and transformed back onto the SO(3) grid of bandwidth b_out. The filter values are drawn uniformly in [-1, 1] and scaled by 1 / sqrt(n_points n_in b_out^4 / b_in^2), the initialization of the reference implementation.

Parameters:

  • n_in (int) – Input and output channels.

  • n_out (int) – Input and output channels.

  • b_in (int) – Bandwidth of the input (S^2 grid of 2 b_in samples per axis) and of the output (SO(3) grid of 2 b_out samples per axis; b_out <= b_in).

  • b_out (int) – Bandwidth of the input (S^2 grid of 2 b_in samples per axis) and of the output (SO(3) grid of 2 b_out samples per axis; b_out <= b_in).

  • points (Tensor) – Filter support, (n_points, 2) pairs (beta, alpha).

filter_spectrum()[source]#

psi^l_m of every channel pair, (n_in, n_out, b_out, 2 b_out - 1).

Return type:

Tensor

forward(x)[source]#

(B, n_in, 2 b_in, 2 b_in) -> (B, n_out, 2 b_out, 2 b_out, 2 b_out).

Parameters:

x (Tensor)

Return type:

Tensor

extra_repr()[source]#

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

Return type:

str