Reduction of Tensors in Irreps

class e3nn.o3.ReducedTensorProducts(formula, filter_ir_out=None, filter_ir_mid=None, eps=1e-09, **irreps)[source]

Bases: CodeGenMixin, Module

reduce a tensor with symmetries into irreducible representations

Parameters
  • formula (str) – String made of letters - and = that represent the indices symmetries of the tensor. For instance ij=ji means that the tensor has two indices and if they are exchanged, its value is the same. ij=-ji means that the tensor change its sign if the two indices are exchanged.

  • filter_ir_out (list of e3nn.o3.Irrep, optional) – Optional, list of allowed irrep in the output

  • filter_ir_mid (list of e3nn.o3.Irrep, optional) – Optional, list of allowed irrep in the intermediary operations

  • **kwargs (dict of e3nn.o3.Irreps) – each letter present in the formula has to be present in the irreps dictionary, unless it can be inferred by the formula. For instance if the formula is ij=ji you can provide the representation of i only: ReducedTensorProducts('ij=ji', i='1o').

irreps_in[source]

input representations

Type

list of e3nn.o3.Irreps

irreps_out[source]

output representation

Type

e3nn.o3.Irreps

change_of_basis[source]

tensor of shape (irreps_out.dim, irreps_in[0].dim, ..., irreps_in[-1].dim)

Type

torch.Tensor

Examples

>>> tp = ReducedTensorProducts('ij=-ji', i='1o')
>>> x = torch.tensor([1.0, 0.0, 0.0])
>>> y = torch.tensor([0.0, 1.0, 0.0])
>>> tp(x, y) + tp(y, x)
tensor([0., 0., 0.])
>>> tp = ReducedTensorProducts('ijkl=jikl=ikjl=ijlk', i="1e")
>>> tp.irreps_out
1x0e+1x2e+1x4e
>>> tp = ReducedTensorProducts('ij=ji', i='1o')
>>> x, y = torch.randn(2, 3)
>>> a = torch.einsum('zij,i,j->z', tp.change_of_basis, x, y)
>>> b = tp(x, y)
>>> assert torch.allclose(a, b, atol=1e-3, rtol=1e-3)

Methods:

forward(*xs)

Defines the computation performed at every call.

forward(*xs)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.