mindscience.e3nn.so2_conv.SO2Convolution

class mindscience.e3nn.so2_conv.SO2Convolution(irreps_in, irreps_out)[source]

SO(2)-equivariant convolution layer for complex-valued features on the circle group.

This layer maps between two Irreps spaces that describe how the inputs/outputs transform under planar rotations. It keeps the m-quantum number (the index that labels the SO(2) irreducible representations) diagonal, so that each m-block is processed independently. For \(m = 0\) (scalar part) a real dense layer is used; for \(m > 0\) a complex-valued \(1\times 1\) convolution (implemented as a real \(2\times 2\) weight matrix acting on the real/imaginary parts) is applied.

Parameters

irreps_in (Union[str, Irreps]) – Input irreps, e.g. "3x0e + 2x1o" (3 scalars + 2 vectors).
irreps_out (Union[str, Irreps]) – Output irreps, e.g. "5x0e + 1x1o".

Inputs:

x (list[Tensor]): list of real tensors, each of shape (batch, mul, 2l+1) representing the complex-valued SO(2) features for each irrep of order l. The last dimension indexes the magnetic quantum number m = -l ... +l.
x_edge (Tensor): real tensor of shape (edges, features) containing edge (radial) attributes that are broadcast and combined with the SO(2) features during convolution.

Outputs:

tuple (Tensor): A tuple of real tensors, one for each irrep in irreps_out. Each tensor has shape (batch, mul, 2l+1) and contains the complex-valued SO(2) features for the corresponding irrep of order l. The last dimension indexes the magnetic quantum number m = -l ... +l.

Examples

>>> import mindspore as ms
>>> from mindscience.e3nn.so2_conv import SO2Convolution
>>> conv = SO2Convolution("2x0e + 1x1o", "2x0e + 1x1o")
>>> x = [ms.ops.randn(4, 2, 1),
...      ms.ops.randn(4, 1, 3)]
>>> x_edge = ms.ops.randn(4, 10)
>>> out = conv(x, x_edge)
>>> len(out), out[0].shape, out[1].shape
(2, (4, 2, 1), (4, 1, 3))