xnn.gnn.models.mace.U_matrix_real#

xnn.gnn.models.mace.U_matrix_real(irreps_in, irreps_out, correlation, normalization='component', filter_ir_mid=None, dtype=None)[source]#

Symmetric coupling basis of correlation copies of irreps_in -> irreps_out.

Assembles the generalized Clebsch-Gordan U basis for the MACE symmetric contraction: it couples correlation copies of irreps_in (via _wigner_nj()) and stacks the paths reaching each target irrep along a trailing n_paths axis. The result is bit-identical to mace-torch.

Parameters:

irreps_in (e3nn.o3.Irreps) – The input irreps; correlation copies are coupled together.
irreps_out (e3nn.o3.Irreps) – The target output irreps to keep coupling paths for.
correlation (int) – Correlation order (number of copies of irreps_in to couple), i.e. the body order minus one. When correlation == 4 the intermediate irreps are filtered (upstream restriction) for tractability.
normalization ({"component", "norm"}, optional) – Clebsch-Gordan normalization convention passed through to _wigner_nj(). Defaults to "component".
filter_ir_mid (list of e3nn.o3.Irrep or None, optional) – Intermediate-irrep filter passed to _wigner_nj(); overridden internally when correlation == 4.
dtype (torch.dtype or None, optional) – Dtype for the coupling tensors. None uses the tensor default.

Returns:

Alternating [ir, U, ir, U, ...] pairs, one per contiguous run of coupling paths reaching an irrep of irreps_out; each U has shape ([out.dim,] in.dim, ..., in.dim, n_paths) (leading output axis squeezed away for scalar outputs, so callers typically take the final element). When no coupling path exists, a single [label, zeros] pair with one all-zero path is returned instead.

Return type:

list