e3j.core.Bigotimes

class e3j.core.Bigotimes(source: list[O3Space | str], target: O3Space | str, coef: Array | None = None, l_max: int | None = None, sort: bool = True, config: Config | None = None)

Multilinear tensor products.

Methods

__call__(*xs)	Evaluate N-linear product on N inputs.
__init__(source, target[, coef, l_max, ...])	Compute and stack generalized CG coefficients, or use given ones.
aggregate(values[, layout])	Aggregate a values-like vector on output coordinates.
clebsch_gordan(target, *sources[, l_max])	Compute generalized CG coefficients by successive compositions.
infer_target(source[, target, l_max, sort])	Infer output representation without computing CG coefficients.
pullback(prod[, position])	Flatten a composition of tensor products.
sort()	Sort irreducible output blocks.

Attributes

aggregation_method	Cache self._aggregation_method with default.
coef	A writable cache descriptor.
indices
nnz
nnz_ratio
shape
target_matrix	Return aggregation matrix of shape (nnz, dim_out).
values

__call__(*xs: Array) → Array

Evaluate N-linear product on N inputs.

By default, this method uses a dense-dense matrix multiplication for the aggregation. The aggregation method can be changed for a sparse on by setting the environment variable E3J_AGGREGATION_METHOD to “scatter”.

__init__(source: list[O3Space | str], target: O3Space | str, coef: Array | None = None, l_max: int | None = None, sort: bool = True, config: Config | None = None)

Compute and stack generalized CG coefficients, or use given ones.

aggregate(values, layout: Layout | str = 'E3NN')

Aggregate a values-like vector on output coordinates.

Parameters:

• values (np.ndarray) – a vector of length self.nnz, e.g. the product of CG coefficients with input coordinates.

• layout – array layout for scatter-based aggregation.

Returns:

a vector of shape target.dim.

Return type:

np.ndarray

property aggregation_method: Aggregation

Cache self._aggregation_method with default.

classmethod clebsch_gordan(target: O3Space | str, *sources: Iterable[str], l_max: int | None = None) → Bigotimes

Compute generalized CG coefficients by successive compositions.

Returns an (N+1)-dimensional array whose dimensions appear in the same order as the arguments.

N.B. This factory is called by __init__ when coef = None.

classmethod infer_target(source: tuple[str, ...], target: str | None = None, l_max: int | None = None, sort: bool = True) → O3Space

Infer output representation without computing CG coefficients.

This method is useful e.g. for the e3j.linen.Bigotimes flax wrapper.

pullback(prod: Bigotimes, position: int = 0) → Bigotimes

Flatten a composition of tensor products.

By associativity of the binary tensor product, any succession of multilinear tensor products can be flattened down to a single operation.:

x0 ... xk
  ` | /
   `|/
    y0 ... yn          <=>         x0 ... xk y1 ... yn
     `  |  /                         `    |   |    /
      ` | /                            `___`_/___/
       `|/                                  |
        z                                   z

Composing tensor products requires to repeat and multiply coefficients of the two l.h.s. trees, as any coefficient downstream is pulled-back by all the upstream coefficients matching its y0 coordinate.

Parameters:

• prod (Bigotimes) – an upstream tensor product.

• position (int) – the position where prod’s output is fed in as argument.

Returns:

composed – the m-ary tensor product obtained by applying prod at position, where:

m = self.arity + prod.arity - 1

Return type:

Bigotimes

sort() → Bigotimes

Sort irreducible output blocks.

property target_matrix: BCOO | Array

Return aggregation matrix of shape (nnz, dim_out).

By default, the aggregation matrix is a dense array. To get a sparse array, set the environment variable E3J_AGGREGATION_METHOD to “sparse”.

Example

Number of coefficients above each output coordinate:

>>> num_coefs = jnp.ones(op.nnz) @ op.target_matrix