xnn.cnn.models.spherical#

Spherical CNN (Cohen, Geiger, Koehler, Welling, ICLR 2018).

Rotation-equivariant convolutions for signals on the sphere S^2 and on the rotation group SO(3), built from the manuscript (arXiv:1801.10130) and the conventions of its reference implementation (jonas-koehler/s2cnn), without using its code:

  • signals are sampled on the Driscoll-Healy grids of the paper: S^2 at beta_j = pi (2j + 1) / 4b, alpha_k = 2 pi k / 2b (j, k < 2b) and SO(3) at the ZYZ Euler angles (beta_j, alpha_k, gamma_l) for a bandwidth b (Appendix A; the axes of a signal tensor are beta, alpha and beta, alpha, gamma);

  • the generalized Fourier transform (eq 8, 21) expands a signal in the Wigner D-functions D^l_{mn}(alpha, beta, gamma) = e^{i m alpha} d^l_{mn}(beta) e^{i n gamma} (Y^l_m = D^l_{m0} on the sphere, footnote 3), with the Wigner d-matrices d^l(beta) = exp(-i beta J_y) evaluated by the eigendecomposition of the spin matrix J_y (exact in float64 at any degree) and the quadrature weights of the grid, and its inverse is eq 9; the transforms are a translational FFT over the alpha (and gamma) axis followed by a contraction over beta with the sampled d (Sec. 4.1);

  • the correlations [psi * f](R) = <L_R psi, f> (eq 4 and 6) are computed in the spectrum by the Fourier theorems of Appendix D, f^l psi^l+ (an outer product of S^2 spectra, a matrix product of SO(3) spectra), with the filter psi parameterized by its values on a small set of grid points (a local patch around the north pole or identity, or a ring around the equator) and transformed with the same basis functions;

  • so3_integrate() is the invariant integral over SO(3) (eq 11), so3_rotate() / s2_rotate() rotate a signal exactly in the spectrum (L_R f of eq 1 and 5), which is what the equivariance checks of Sec. 5.1 use.

SphericalCNN is the interatomic potential of Sec. 5.4: every atom carries the spherical signals of SphericalGrid (one Coulomb-like potential channel per species on a sphere around the atom), a stack of ResNet blocks of S^2 / SO(3) correlations maps them to features on SO(3), the integral over SO(3) gives the invariant per-atom features, and the readout the energy; the paper’s DeepSet readout is available as an option. Needs no e3nn.

Functions

default_schedule(n_features, n_blocks, bandwidth)

Channels and bandwidths of n_blocks ResNet blocks: channels growing linearly to n_features, bandwidths falling from bandwidth to 2 (the shape of Table 3 of the paper).

euler_to_matrix(alpha, beta, gamma)

The rotation matrix Z(alpha) Y(beta) Z(gamma) (eq 10), float64.

matrix_to_euler(R)

ZYZ Euler angles (alpha, beta, gamma) of a rotation matrix.

quadrature_weights(b)

Quadrature weights w_j of the SO(3) grid for the normalized Haar measure (eq 11): int f dR = sum_{j,k,l} w_j f(beta_j, alpha_k, gamma_l).

s2_equatorial_grid([max_beta, n_alpha, n_beta])

Points (beta, alpha) on n_beta rings around the equator (the non-local filter support of the SHREC17 network).

s2_grid_points(b[, dtype])

Unit vectors x(alpha, beta) = Z(alpha) Y(beta) n of the S^2 grid (eq 12), shape (2b, 2b, 3) indexed [beta, alpha].

s2_kernel_basis(points, b)

2b conj(Y^l_m) at filter points (beta, alpha): the matrix that turns filter values into the filter spectrum of an S^2 correlation.

s2_near_identity_grid([max_beta, n_alpha, ...])

Points (beta, alpha) on n_beta rings around the north pole, the local filter support of the paper's networks.

s2_rotate(x, alpha, beta, gamma)

[L_R f](x) = f(R^{-1} x) (eq 1) of spherical signals (..., 2b, 2b), computed exactly in the spectrum.

so3_alphas(b[, dtype])

The 2b azimuths alpha_k = 2 pi k / 2b (also the gamma grid).

so3_betas(b[, dtype])

The 2b polar angles beta_j = pi (2j + 1) / 4b of the grids.

so3_equatorial_grid([max_beta, max_gamma, ...])

Rotations (beta, alpha, gamma) on rings around the equator.

so3_integrate(x)

Integral of SO(3) signals (..., 2b, 2b, 2b) over the normalized Haar measure (eq 11): the invariant pooling of the paper's networks.

so3_kernel_basis(points, b)

conj(D^l_{mn}) at filter rotations (beta, alpha, gamma): the matrix that turns filter values into the filter spectrum of an SO(3) correlation.

so3_near_identity_grid([max_beta, ...])

Rotations (beta, alpha, gamma) on rings around the identity: every rotation of a ring is at the same distance from the identity.

so3_rotate(x, alpha, beta, gamma)

[L_R f](Q) = f(R^{-1} Q) (eq 5) of SO(3) signals (..., 2b, 2b, 2b), for R = Z(alpha) Y(beta) Z(gamma), computed exactly in the spectrum.

wigner_D(l, alpha, beta, gamma)

D^l_{mn}(alpha, beta, gamma) = e^{i m alpha} d^l_{mn}(beta) e^{i n gamma} (complex128, (..., 2l + 1, 2l + 1)), the basis functions of the SO(3) Fourier transform in the convention of the reference code.

wigner_d(l, beta)

Wigner d-matrix d^l_{mn}(beta) = <l m| exp(-i beta J_y) |l n>.

wigner_d_table(b_grid, b_spec)

d^l_{mn}(beta_j) on the grid of bandwidth b_grid for every l < b_spec, as a zero-padded table [j, l, m + b_spec - 1, n + b_spec - 1] of shape (2 b_grid, b_spec, 2 b_spec - 1, 2 b_spec - 1) (float64).

Classes

S2Convolution(n_in, n_out, b_in, b_out, points)

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

S2Transform(b_grid, b_spec)

Spherical harmonic transform between the S^2 grid of bandwidth b_grid and spectra of bandwidth b_spec (b_spec <= b_grid).

SO3Convolution(n_in, n_out, b_in, b_out, points)

SO(3) correlation layer (eq 6, computed as the matrix product f^l psi^l+ of Appendix D).

SO3Transform(b_grid, b_spec)

Fourier transform between the SO(3) grid of bandwidth b_grid and spectra of bandwidth b_spec (b_spec <= b_grid).

SphericalCNN([species, cutoff, radius, ...])

Spherical CNN interatomic potential (Cohen et al., ICLR 2018, Sec.

SphericalResBlock(n_in, n_out, b_in, b_out, ...)

ResNet block of the QM7 network (Sec.
Show Source