neighbor_algorithm

This is summary documentation for hippynn.layers.pairs.csr_pairs.neighbor_algorithm module. Click here for full documentation.

Links to Full Documentation:

This file was written with assistance from an LLM.

Neighbor list construction over triclinic periodic cells using PyTorch + CSR.

This module builds neighbor pairs for batched molecular/atomistic systems with triclinic cells. It uses a light CSR “algebra” (see csrtable) to stage row-wise joins and filtering:

Pipeline (high level)

  1. build_initial_data() — wrap raw inputs into CSRs for systems/atoms.

  2. normalize_atoms() — wrap positions to fractional [0,1)^3 and record integer image offsets.

  3. build_image_offsets() — compute per-system image-shift stencils large enough for cutoff.

  4. build_image_atoms() — cartesian-expand primary atoms × image shifts.

  5. voxelize_images() — voxel grid per system with edge ≈ cutoff and assign image-atoms to voxels.

  6. voxel_adjacency() — build voxel–voxel edges via a fixed stencil (includes self-edges).

  7. expand_pairs() — expand voxel edges into candidate atom pairs.

  8. prune_pairs() — filter by cutoff; drop same-atom self-edges for primary images.

  9. calc_neighbors() — orchestration entry point that returns final pairs.

All tensor operations are batched and device/dtype agnostic. Index tensors are torch.long. Distances/displacements remain differentiable w.r.t. input positions.

Functions

build_image_atoms(atomCSR, imageCSR, systemCSR)

Expand primary atoms across image offsets (row-wise outer join).

build_image_offsets(systemCSR[, ...])

Compute per-system periodic image offsets sufficient for cutoff.

build_initial_data(positions, nonblank, ...)

Package raw inputs (positions/mask/cells/cutoff) into CSR containers.

calc_neighbors(positions, nonblank, cells, ...)

Full neighbor-list pipeline (batched triclinic PBC).

calculate_aabb(positions, system_id, n_systems)

Compute per-system axis-aligned bounding boxes (AABB).

expand_pairs(voxel_atomCSR, first_vox, ...)

Expand voxel edges to candidate atom pairs.

normalize_atoms(atomCSR, systemCSR)

Wrap atoms into the primary triclinic cell and record image offsets.

prune_pairs(pairs)

Apply geometric/pruning predicates to candidate pairs.

voxel_adjacency(voxelCSR, systemCSR)

Build voxel–voxel edges using a fixed stencil (includes self-edges).

voxelize_images(image_atomCSR, systemCSR)

Partition image-atoms into near-cubic voxels of edge ≈ cutoff.