e3tools package

Subpackages

Module contents

e3tools.pack_irreps(unpacked_tuples: Iterator[Tuple[int, Irrep, Tensor]]) Tensor[source]

Pack fields into a single tensor.

e3tools.radius(x: Float[Tensor, 'N D'], y: Float[Tensor, 'M D'], r: float, batch_x: Int64[Tensor, 'N'] | None = None, batch_y: Int64[Tensor, 'M'] | None = None, ignore_same_index: bool = False, chunk_size: int | None = None) Int64[Tensor, '2 E'][source]

For each element in y find all points in x within distance r

e3tools.radius_graph(x: Float[Tensor, 'N D'], r: float, batch: Int64[Tensor, 'N'] | None = None, chunk_size: int | None = None, loop: bool = False) Int64[Tensor, '2 E'][source]

Returns the edge index given by applying a radial cutoff of r to x.

e3tools.scatter(src: Tensor, index: Tensor, dim: int, dim_size: int | None = None, reduce: str = 'sum') Tensor[source]

Collects elements at the indices index of a source tensor src.

e3tools.unpack_irreps(x: Tensor, irreps: Irreps) Iterator[Tuple[int, Irrep, Tensor]][source]

Given a packed irreps vector of dimension […, irreps.dim] yield tuples (mul, ir, field) where field has dimension […, mul, 2*l+1] for each irrep in irreps