e3j.core.ScalarMixing

class e3j.core.ScalarMixing(source: ~e3j.spaces.o3.O3Space, layout: ~e3j.utils.options.Layout = <factory>)

Rescale equivariant features by as many scalars.

Computes output features \(x'\) lying in the same space as input features \(x\),

\[x'_{nlm} = s_{nl} \cdot x_{nlm},\]

where \(s\) denotes an array of num_irreps scalars, \(x\) an array of num_irreps irreducible features and \(n,l,m\) denote channel, degree and spin respectively.

Note

This is the same operation as scalars * irreps_array in e3nn. As a bilinear coupling, it is a particular case of an equivariant tensor product with diagonal coefficients.

Parameters:
  • source – The input space for the equivariant features \(x\).

  • layout – The channel axis Layout, equal to config().layout by default.

Methods

__call__(scalars, features)

Mix scalars with equivariant features.

Attributes

mix_indices

Return index map from equivariant coordinates to scalars.

num_irreps

source

layout

__call__(scalars: Array, features: Array) Array

Mix scalars with equivariant features.

property mix_indices: Array

Return index map from equivariant coordinates to scalars.