Skip to content

Geometry and rotations

Atomic geometry determines which atoms interact and how their relative directions enter an equivariant model. Neighbor search selects pairs within a cutoff. Spherical harmonics represent directions in spherical components, and Wigner matrices describe how those components transform under rotations.

These functions accept Torch tensors on CPU or CUDA. Spherical harmonics and Wigner matrices support differentiation through second derivatives. Local-frame transforms use the rotation matrices to move network features between global and local frames; the corresponding NumPy operations are documented under Physics.

The returned indices and cell shifts identify the selected pairs; distances and displacement vectors describe their geometry.

neighbor_list 

neighbor_list(
    quantities: str,
    positions: Tensor,
    cell: Tensor,
    pbc: Tensor,
    cutoff: float,
    batch_ptr: Tensor | None = None,
    *,
    algorithm: Literal["auto", "brute_force", "cell_list"] = "auto",
    cpu_threads: int | None = None,
    sorted: bool = False,
    half_list: bool = False,
    include_self: bool = False,
) -> tuple[Tensor, ...]

Build a Torch neighbor list within a strict distance cutoff.

Parameters:

  • quantities (str) –

    Returned quantities in order: "i" and "j" select source and target indices, "P" paired indices, "S" integer cell shifts, "d" distances, and "D" displacement vectors.

  • positions (Tensor) –

    Cartesian positions with shape (n_atoms, 3) for one structure or (n_total_atoms, 3) for a batch. Tensors may reside on CPU or CUDA and must have dtype float32 or float64.

  • cell (Tensor) –

    Cell vectors stored as rows, with shape (3, 3) for one structure or (n_structures, 3, 3) for a batch. Its dtype and device must match positions.

  • pbc (Tensor) –

    Periodic-axis flags with shape (3,) for one structure or (n_structures, 3) for a batch, dtype bool, and the same device as positions.

  • cutoff (float) –

    Strict, finite, positive distance cutoff in the same length unit as positions and cell.

  • batch_ptr (Tensor | None, default: None ) –

    Optional int64 structure boundaries with shape (n_structures + 1,). None denotes one structure.

  • algorithm (Literal['auto', 'brute_force', 'cell_list'], default: 'auto' ) –

    Native search algorithm.

  • cpu_threads (int | None, default: None ) –

    CPU thread count. None uses one CPU thread and leaves the CUDA option unspecified.

  • sorted (bool, default: False ) –

    Whether to sort pairs by source index.

  • half_list (bool, default: False ) –

    Whether to retain one member of each reverse-pair class.

  • include_self (bool, default: False ) –

    Whether to include zero-shift self pairs.

Returns:

  • Tensor –

    One tensor per character in quantities. For n_edges pairs, i and

  • ... –

    j have shape (n_edges,), P has shape (n_edges, 2), S has

  • tuple[Tensor, ...] –

    shape (n_edges, 3), d has shape (n_edges,), and D has shape

  • tuple[Tensor, ...] –

    (n_edges, 3). Results remain on the input device. Distances and

  • tuple[Tensor, ...] –

    displacements preserve the floating dtype and remain differentiable

  • tuple[Tensor, ...] –

    with respect to positions and cell while the discrete neighbor

  • tuple[Tensor, ...] –

    identity is fixed.

Spherical functions and rotations

spherical_harmonics 

spherical_harmonics(l_max: int, directions: Tensor) -> Tensor

Evaluate normalized real spherical harmonics through l_max.

The harmonics follow the Wikipedia real spherical harmonic convention and are ordered in consecutive \(l\) blocks, with m = -l, ..., l within each block.

Parameters:

  • l_max (int) –

    Highest \(l\) to calculate. Results include every \(l\) from zero through l_max.

  • directions (Tensor) –

    Nonzero Cartesian directions with shape (..., 3) on CPU or CUDA and dtype float32 or float64.

Returns:

  • Tensor –

    Harmonic values with shape (..., (l_max + 1) ** 2), preserving the

  • Tensor –

    dtype and device of directions and supporting reverse-mode

  • Tensor –

    differentiation through second derivatives.

wigner_D 

wigner_D(l_max: int, rotation: Tensor) -> Tensor

Return flattened real Wigner \(D\) matrices through l_max.

The matrices use the Wikipedia real spherical-harmonic basis, with each block ordered by m = -l, ..., l and the l = 1 block ordered as \((y,z,x)\). The input acts actively on Cartesian column vectors and the result satisfies \(Y_l(R^\mathsf{T}x)=Y_l(x)D^{(l)}(R)\) for row harmonic values.

Parameters:

  • l_max (int) –

    Highest non-negative l. This is a static model configuration under torch.compile.

  • rotation (Tensor) –

    Proper rotation matrices in SO(3), with shape (..., 3, 3), dtype float32 or float64, and device CPU or CUDA. Matrix entries are treated as independent variables when differentiating.

Returns:

  • Tensor –

    Consecutive row-major flattened blocks

  • Tensor –

    \(D^{(0)},\ldots,D^{(l_{\max})}\),

  • Tensor –

    with shape

  • Tensor –

    (..., sum((2 * l + 1) ** 2 for l in range(l_max + 1))). Values

  • Tensor –

    preserve the input dtype and device and support reverse-mode

  • Tensor –

    differentiation through second derivatives.