e3j.utils.sparse.SparseMixin¶

class e3j.utils.sparse.SparseMixin¶

Properties and methods depending on a coef attribute.

The coef attribute is a SparseCOO array of numpy coefficients.

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

aggregate(values[, layout])

Aggregate a values-like vector on output coordinates.

Attributes

`aggregation_method`	Cache `self._aggregation_method` with default.
`coef`	Compute and cache coefficients.
`indices`
`nnz`	Number of non-zero coefficients.
`nnz_ratio`	Density of the coefficient tensor, i.e. `nnz` over its size.
`shape`
`target_matrix`	Return aggregation matrix of shape `(nnz, dim_out)`.
`values`

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

Aggregate a values-like vector on output coordinates.

Parameters:

values (jnp.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:

jnp.ndarray

property aggregation_method: Aggregation¶: Cache self._aggregation_method with default.

coef¶: Compute and cache coefficients.

property nnz: int¶: Number of non-zero coefficients.

property nnz_ratio: float¶: Density of the coefficient tensor, i.e. nnz over its size.

property target_matrix: BCOO | ndarray¶

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 = np.ones(op.nnz) @ op.target_matrix