e3j.utils.sparse.SparseMixin

class e3j.utils.sparse.SparseMixin

Properties and methods depending on a coef attribute.

The coef attribute can in fact be either of type:

• jax.experimental.sparse.BCOO, or

• np.ndarray : in this case some properties and methods will fail, such as .aggregate(), .nnz…

To switch between possible equivalent methods – typically, a matrix product or scatter-add operation in a final reduction stage, child classes such as TensorProduct or Bigotimes override the aggregation_method descriptor during __init__. The default attempts to read $E3J_AGGREGATION_METHOD from environment.

Methods

__init__()
aggregate(values[, layout])	Aggregate a values-like vector on output coordinates.

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

__init__()

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.

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