e3j.core.TensorProduct

class e3j.core.TensorProduct(source: tuple[O3Space, O3Space] | tuple[str, str], target: O3Space | str | None, coef: BCOO | None = None, sort: bool = True, unroll: int | tuple[int, int, int] = 1, config: Config | None = None, layout: str | Layout = 'LEADING_CHANNELS', mode: str | TPMode = 'OUTER', normalization: str | TPNormalization = 'NONE')

Bilinear tensor products.

This class implements (equivariant) bilinear maps of the form:

z[:,i] = Σⱼₖ c[i,j,k] * x[:,j] * y[:,k],

defined by a 3D coefficient array c in sparse BCOO format.

Evaluation is implemented by:

  1. a pull-back of x and y by coefficient indices j and k,

  2. a product of x[:,j], y[:,k] and coefficient values c[i,j,k],

  3. a sparse reduction step to aggregate output coordinates z[:,i].

Methods

__call__(x, y[, coef])

Evaluate bilinear map on pair of inputs.

__init__(source, target[, coef, sort, ...])

Stack Clebsch Gordan coefficients or use explicitly given ones.

aggregate(values[, layout])

Aggregate a values-like vector on output coordinates.

clebsch_gordan(target, source_1, source_2[, ...])

Stack Clebsch-Gordan coefficients.

dense_eval(x, y, coef)

Evaluate bilinear map on pair of inputs.

fused_eval(x, y, coef[, unroll])

Evaluate bilinear map on pair of inputs.

infer_target(source[, target, sort])

Infer target representation from Clebsch-Gordan rules.

sort()

Sort irreducible output blocks.

sparse_eval(x, y, coef)

Evaluate bilinear map on pair of inputs.

Attributes

aggregation_method

Cache self._aggregation_method with default.

coef

A writable cache descriptor.

indices

is_dense

is_fused

nnz

nnz_ratio

shape

target_matrix

Return aggregation matrix of shape (nnz, dim_out).

values

__call__(x: Array, y: Array, coef: Array | None = None) → Array

Evaluate bilinear map on pair of inputs.

__init__(source: tuple[O3Space, O3Space] | tuple[str, str], target: O3Space | str | None, coef: BCOO | None = None, sort: bool = True, unroll: int | tuple[int, int, int] = 1, config: Config | None = None, layout: str | Layout = 'LEADING_CHANNELS', mode: str | TPMode = 'OUTER', normalization: str | TPNormalization = 'NONE')

Stack Clebsch Gordan coefficients or use explicitly given ones.

Parameters:

  • source (representations of both inputs.)

  • target (output representation, inferred by default.) – If coefficients are None, then out is interpreted as a momentum and parity filter which should only contain irreducible blocks of multiplicity 1.

  • coef (sparse array, constructed by stacking Clebsch-Gordan) – coefficients by default.

  • sort (whether to sort the output coordinates by grouping momenta and) – parities, True by default.

  • layout (specifies the channel axis, TRAILING_CHANNELS is faster.)

  • mode (specifies how to mix the channels, can take three values:) –

    • OUTER : tensor product of channels in np.outer fashion. Also used to broadcast one operand on the channel axis.

    • INNER : scalar product of channels, summed over after the Clebsch-Gordan tensor product.

    • MAP : channel-wise tensor products. Only useful with trailing channels layout, since leading axes are mapped over by default.

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.

static clebsch_gordan(target: O3Space, source_1: O3Space, source_2: O3Space, normalization: TPNormalization = TPNormalization.NONE) → BCOO

Stack Clebsch-Gordan coefficients.

Returns a 3D-array whose dimensions appear in the same order as the arguments.

dense_eval(x: Array, y: Array, coef: Array) → Array

Evaluate bilinear map on pair of inputs.

fused_eval(x: Array, y: Array, coef: Array, unroll: int | tuple[int, int] = 1) → Array

Evaluate bilinear map on pair of inputs.

classmethod infer_target(source: tuple[O3Space, O3Space], target: O3Space | None = None, sort: bool = False) → O3Space

Infer target representation from Clebsch-Gordan rules.

The optional target argument should only contain irreducible representations with multiplicity 1, and ValueError will be raised otherwise.

sort() → TensorProduct

Sort irreducible output blocks.

In contrast with e3nn, the reordering of output coordinates is performed once and for all on the full coefficient tensor, avoiding the associated overhead at evaluation time.

sparse_eval(x: Array, y: Array, coef: Array) → Array

Evaluate bilinear map on pair of inputs.

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