Change2DisentangledBasis

class Change2DisentangledBasis(in_rep, learnable=False)

Bases: Module

Map features to the isotypic/irrep-spectral basis.

For \(\mathbf{x}\in\mathcal{X}\) with representation \(\rho_{\mathcal{X}}\), this module applies \(\mathbf{Q}^{-1}\) from isotypic_decomp_rep():

\[\hat{\mathbf{x}} = \mathbf{Q}^{-1}\mathbf{x}, \qquad \rho_{\mathcal{X}} = \mathbf{Q}\left( \bigoplus_{k\in[1,n_{\text{iso}}]} \bigoplus_{i\in[1,n_k]} \hat{\rho}_k \right)\mathbf{Q}^T.\]

Hence, coordinates in out_rep are grouped by isotypic subspace (same irrep type contiguous). The map is linear and \(\mathbb{G}\)-equivariant:

\[\hat{\rho}_{\mathcal{X}}(g)\,\hat{\mathbf{x}} = \mathbf{Q}^{-1}\rho_{\mathcal{X}}(g)\mathbf{x}, \quad \hat{\rho}_{\mathcal{X}}(g) = \mathbf{Q}^{-1}\rho_{\mathcal{X}}(g)\mathbf{Q}.\]

Parameters:

• in_rep (Representation) – Representation \(\rho_{\text{in}}\) describing the input feature space.

• learnable (bool, optional) – If True, the change-of-basis matrix is a trainable parameter. Defaults to False.

forward(x)

Apply the basis change to isotypic coordinates.

Parameters:

x (Tensor) – Input whose last dimension equals in_rep.size; arbitrary leading dimensions allowed.

Returns:

Tensor with the same leading shape and last dimension out_rep.size (same as in_rep), expressed in the isotypic basis. If the input is already in that basis, the tensor is returned unchanged.

Return type:

Tensor