Ctrl+K

Note

Go to the end to download the full example code.

Simple Neighbor List Example#

This example demonstrates how to use the basic neighbor list functions in nvalchemiops on random systems. We’ll cover:

  1. cell_list: O(N) algorithm using spatial cell lists (best for large systems)

  2. naive_neighbor_list: O(N²) algorithm (best for small systems)

  3. neighbor_list: Unified wrapper with automatic method selection (RECOMMENDED)

  4. Comparison between algorithms

  5. Neighbor matrix vs neighbor list (COO) formats

  6. build_cell_list + query_cell_list: Lower-level API with caching

  7. return_distances / return_vectors: differentiable per-pair geometry (autograd)

The neighbor list construction efficiently finds all atom pairs within a cutoff distance, which is essential for molecular simulations and materials science calculations.

import time

import torch

from nvalchemiops.torch.neighbors import estimate_neighbor_list_costs, neighbor_list
from nvalchemiops.torch.neighbors.cell_list import (
    build_cell_list,
    cell_list,
    estimate_cell_list_sizes,
    query_cell_list,
)
from nvalchemiops.torch.neighbors.naive import naive_neighbor_list
from nvalchemiops.torch.neighbors.neighbor_utils import (
    allocate_cell_list,
    estimate_max_neighbors,
    get_neighbor_list_from_neighbor_matrix,
)

Set up the computation device#

Choose device (CPU or GPU if available)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32

print(f"Using device: {device}")
print(f"Using dtype: {dtype}")

_USE_CUDA_EVENTS = device.type == "cuda"


def _timed_loop(fn, n_iter=10, warmup=10):
    """Run *fn* for *warmup* + *n_iter* iterations and return avg ms."""
    for _ in range(warmup):
        fn()
    if _USE_CUDA_EVENTS:
        start_ev = torch.cuda.Event(enable_timing=True)
        end_ev = torch.cuda.Event(enable_timing=True)
        torch.cuda.synchronize()
        start_ev.record()
        for _ in range(n_iter):
            fn()
        end_ev.record()
        torch.cuda.synchronize()
        return start_ev.elapsed_time(end_ev) / n_iter
    else:
        t0 = time.perf_counter()
        for _ in range(n_iter):
            fn()
        return (time.perf_counter() - t0) / n_iter * 1000.0


def _dispatch_report(label, positions, cell, pbc, cutoff, **kwargs):
    """Print sorted Torch neighbor-list method costs and return the cheapest."""
    batch_ptr = torch.tensor(
        [0, positions.shape[0]],
        dtype=torch.int32,
        device=positions.device,
    )
    report = estimate_neighbor_list_costs(
        batch_ptr,
        cell,
        pbc,
        cutoff,
        positions_dtype=positions.dtype,
        **kwargs,
    )
    print(f"\n{label}")
    for method_name, cost in report:
        print(f"  {method_name:24s} estimated cost (arbitrary units): {cost:.3g}")
    return report[0][0]

Using device: cuda
Using dtype: torch.float32

Create random systems of different sizes#

We’ll compare performance on small vs large systems

print("\n" + "=" * 70)
print("CREATING TEST SYSTEMS")
print("=" * 70)

# Small system: good for naive algorithm
small_num_atoms = 100
small_box_size = 10.0

# Larger system: demonstrates cell-list scaling without expensive O(N²) timing
large_num_atoms = 4096
large_box_size = 40.0

cutoff = 5.0

torch.manual_seed(42)

# Small system
small_positions = (
    torch.rand(small_num_atoms, 3, dtype=dtype, device=device) * small_box_size
)
small_cell = (torch.eye(3, dtype=dtype, device=device) * small_box_size).unsqueeze(0)

# Large system
large_positions = (
    torch.rand(large_num_atoms, 3, dtype=dtype, device=device) * large_box_size
)
large_cell = (torch.eye(3, dtype=dtype, device=device) * large_box_size).unsqueeze(0)

# Periodic boundary conditions (all directions periodic)
pbc = torch.tensor([True, True, True], device=device).unsqueeze(0)

print(f"\nSmall system: {small_num_atoms} atoms in {small_box_size}³ box")
print(f"Large system: {large_num_atoms} atoms in {large_box_size}³ box")
print(f"Cutoff distance: {cutoff} Å")

======================================================================
CREATING TEST SYSTEMS
======================================================================

Small system: 100 atoms in 10.0³ box
Large system: 4096 atoms in 40.0³ box
Cutoff distance: 5.0 Å

Method 1: Cell List Algorithm (O(N) - best for large systems)#

The cell list algorithm uses spatial decomposition for efficient neighbor finding

print("\n" + "=" * 70)
print("METHOD 1: CELL LIST ALGORITHM (O(N))")
print("=" * 70)

# On larger systems, spatial decomposition avoids all-pairs distance checks.
print(f"\n--- Large System ({large_num_atoms:,} atoms) ---")

# Return neighbor matrix format (default)
neighbor_matrix, num_neighbors, shifts = cell_list(
    large_positions,
    cutoff,
    large_cell,
    pbc,
)

print(f"Returned neighbor matrix: shape {neighbor_matrix.shape}")
print(f"  neighbor_matrix: (num_atoms, max_neighbors) = {neighbor_matrix.shape}")
print(f"  num_neighbors: (num_atoms,) = {num_neighbors.shape}")
print(f"  shifts: (num_atoms, max_neighbors, 3) = {shifts.shape}")
print(f"Total neighbor pairs: {num_neighbors.sum()}")
print(f"Average neighbors per atom: {num_neighbors.float().mean():.2f}")

# Convert neighbor matrix to neighbor list (COO) format for GNNs
neighbor_list_coo, neighbor_list_ptr, shifts_coo = (
    get_neighbor_list_from_neighbor_matrix(
        neighbor_matrix, num_neighbors, shifts, fill_value=large_num_atoms
    )
)  # fill_value is the number of atoms in the system

source_atoms = neighbor_list_coo[0]
target_atoms = neighbor_list_coo[1]
print(f"\nReturned neighbor list (COO format): shape {neighbor_list_coo.shape}")
print(f"  neighbor_list: (2, num_pairs) = {neighbor_list_coo.shape}")
print(f"  source_atoms: {source_atoms.shape}")
print(f"  target_atoms: {target_atoms.shape}")
print(f"  shifts: (num_pairs, 3) = {shifts_coo.shape}")
print(f"Total pairs (COO): {neighbor_list_coo.shape[1]}")

# Show some example neighbors
if neighbor_list_coo.shape[1] > 0:
    print("\nFirst 5 neighbor pairs:")
    for idx in range(min(5, neighbor_list_coo.shape[1])):
        src, tgt = source_atoms[idx].item(), target_atoms[idx].item()
        shift = shifts_coo[idx]
        print(f"  Atom {src} -> Atom {tgt}, shift: {shift.tolist()}")

======================================================================
METHOD 1: CELL LIST ALGORITHM (O(N))
======================================================================

--- Large System (4,096 atoms) ---
Returned neighbor matrix: shape torch.Size([4096, 112])
  neighbor_matrix: (num_atoms, max_neighbors) = torch.Size([4096, 112])
  num_neighbors: (num_atoms,) = torch.Size([4096])
  shifts: (num_atoms, max_neighbors, 3) = torch.Size([4096, 112, 3])
Total neighbor pairs: 136604
Average neighbors per atom: 33.35

Returned neighbor list (COO format): shape torch.Size([2, 136604])
  neighbor_list: (2, num_pairs) = torch.Size([2, 136604])
  source_atoms: torch.Size([136604])
  target_atoms: torch.Size([136604])
  shifts: (num_pairs, 3) = torch.Size([136604, 3])
Total pairs (COO): 136604

First 5 neighbor pairs:
  Atom 0 -> Atom 1620, shift: [0, 0, 0]
  Atom 0 -> Atom 595, shift: [0, 0, 0]
  Atom 0 -> Atom 3771, shift: [0, 0, 0]
  Atom 0 -> Atom 1745, shift: [0, 0, 0]
  Atom 0 -> Atom 2094, shift: [0, 0, 0]

Method 2: Naive Algorithm (O(N²) - best for small systems)#

The naive algorithm computes all pairwise distances

print("\n" + "=" * 70)
print("METHOD 2: NAIVE ALGORITHM (O(N²))")
print("=" * 70)

# On small system (where naive algorithm is competitive)
print("\n--- Small System (100 atoms) ---")

# Return neighbor matrix format
neighbor_matrix_naive, num_neighbors_naive, shifts_naive = naive_neighbor_list(
    small_positions, cutoff, cell=small_cell, pbc=pbc
)

print(f"Returned neighbor matrix: shape {neighbor_matrix_naive.shape}")
print(f"Total neighbor pairs: {num_neighbors_naive.sum()}")
print(f"Average neighbors per atom: {num_neighbors_naive.float().mean():.2f}")

# Or return neighbor list (COO) format
neighbor_list_naive, neighbor_ptr_naive, shifts_naive_coo = naive_neighbor_list(
    small_positions, cutoff, cell=small_cell, pbc=pbc, return_neighbor_list=True
)

print(f"\nReturned neighbor list (COO format): shape {neighbor_list_naive.shape}")
print(f"Total pairs (COO): {neighbor_list_naive.shape[1]}")
print(f"Neighbor ptr shape: {neighbor_ptr_naive.shape}")

======================================================================
METHOD 2: NAIVE ALGORITHM (O(N²))
======================================================================

--- Small System (100 atoms) ---
Returned neighbor matrix: shape torch.Size([100, 112])
Total neighbor pairs: 5152
Average neighbors per atom: 51.52

Returned neighbor list (COO format): shape torch.Size([2, 5152])
Total pairs (COO): 5152
Neighbor ptr shape: torch.Size([101])

Comparing Cell List vs Naive Algorithm#

Compare performance and verify results match

print("\n" + "=" * 70)
print("ALGORITHM COMPARISON")
print("=" * 70)

print("\n--- Small System (100 atoms) ---")

# Cell list on small system
_, num_neighbors_cell_small, _ = cell_list(small_positions, cutoff, small_cell, pbc)
cell_time_small = _timed_loop(
    lambda: cell_list(small_positions, cutoff, small_cell, pbc)
)

# Naive on small system
_, num_neighbors_naive_small, _ = naive_neighbor_list(
    small_positions, cutoff, cell=small_cell, pbc=pbc
)
naive_time_small = _timed_loop(
    lambda: naive_neighbor_list(small_positions, cutoff, cell=small_cell, pbc=pbc)
)

print(f"Cell list:  {cell_time_small:.3f} ms, {num_neighbors_cell_small.sum()} pairs")
print(f"Naive:      {naive_time_small:.3f} ms, {num_neighbors_naive_small.sum()} pairs")
print(
    f"Results match: {torch.equal(num_neighbors_cell_small, num_neighbors_naive_small)}"
)

print(f"\n--- Large System ({large_num_atoms:,} atoms) ---")

# Cell list on large system
_, num_neighbors_cell_large, _ = cell_list(large_positions, cutoff, large_cell, pbc)
cell_time_large = _timed_loop(
    lambda: cell_list(large_positions, cutoff, large_cell, pbc)
)

print(f"Cell list:  {cell_time_large:.3f} ms, {num_neighbors_cell_large.sum()} pairs")
print(
    "Naive large-system timing skipped: the O(N²) path is only demonstrated "
    "on the small system above."
)

======================================================================
ALGORITHM COMPARISON
======================================================================

--- Small System (100 atoms) ---
Cell list:  0.797 ms, 5152 pairs
Naive:      0.539 ms, 5152 pairs
Results match: True

--- Large System (4,096 atoms) ---
Cell list:  0.728 ms, 136604 pairs
Naive large-system timing skipped: the O(N²) path is only demonstrated on the small system above.

Method 4: Low-level Cell List API with Caching#

For advanced users: separate build and query phases for efficient caching

print("\n" + "=" * 70)
print("METHOD 4: LOW-LEVEL CELL LIST API (BUILD + QUERY)")
print("=" * 70)

# This approach is useful when you want to:
# 1. Build cell list once, query multiple times with different cutoffs
# 2. Cache cell list for MD simulations with rebuild detection
# 3. Have fine control over memory allocation

positions = large_positions
cell_tensor = large_cell

print("\n--- Step 1: Estimate Memory Requirements ---")
max_total_cells, neighbor_search_radius = estimate_cell_list_sizes(
    cell_tensor, pbc, cutoff
)
print(f"Estimated max cells: {max_total_cells}")
print(f"Neighbor search radius: {neighbor_search_radius}")

print("\n--- Step 2: Allocate Memory ---")
cell_list_cache = allocate_cell_list(
    total_atoms=positions.shape[0],
    max_total_cells=max_total_cells,
    neighbor_search_radius=neighbor_search_radius,
    device=device,
)

(
    cells_per_dimension,
    neighbor_search_radius,
    atom_periodic_shifts,
    atom_to_cell_mapping,
    atoms_per_cell_count,
    cell_atom_start_indices,
    cell_atom_list,
) = cell_list_cache

print(f"Allocated cell list cache with {len(cell_list_cache)} components")

print("\n--- Step 3: Build Cell List ---")
build_cell_list(positions, cutoff, cell_tensor, pbc, *cell_list_cache)
print(f"Built cell list with {cells_per_dimension.tolist()} cells per dimension")

print("\n--- Step 4: Query Neighbors ---")
# Allocate output tensors
max_neighbors = estimate_max_neighbors(cutoff, safety_factor=20.0)
neighbor_matrix_query = torch.full(
    (positions.shape[0], max_neighbors), -1, dtype=torch.int32, device=device
)
neighbor_shifts_query = torch.zeros(
    (positions.shape[0], max_neighbors, 3), dtype=torch.int32, device=device
)
num_neighbors_query = torch.zeros(positions.shape[0], dtype=torch.int32, device=device)

query_cell_list(
    positions,
    cutoff,
    cell_tensor,
    pbc,
    *cell_list_cache,
    neighbor_matrix_query,
    neighbor_shifts_query,
    num_neighbors_query,
)

print(f"Query found {num_neighbors_query.sum()} neighbor pairs")
print(f"Average neighbors per atom: {num_neighbors_query.float().mean():.2f}")

======================================================================
METHOD 4: LOW-LEVEL CELL LIST API (BUILD + QUERY)
======================================================================

--- Step 1: Estimate Memory Requirements ---
Estimated max cells: 512
Neighbor search radius: tensor([1, 1, 1], device='cuda:0', dtype=torch.int32)

--- Step 2: Allocate Memory ---
Allocated cell list cache with 7 components

--- Step 3: Build Cell List ---
Built cell list with [8, 8, 8] cells per dimension

--- Step 4: Query Neighbors ---
/lustre/fs1/portfolios/coreai/projects/coreai_modulus_alchemi/users/kinlongkelvi/code/nvalchemi-toolkit-ops/examples/neighbors/01_simple_neighbor_list.py:409: DeprecationWarning: The 'safety_factor' argument to estimate_max_neighbors is deprecated; it scales the estimate identically to 'atomic_density'. Set 'atomic_density' instead.
  max_neighbors = estimate_max_neighbors(cutoff, safety_factor=20.0)
Query found 136604 neighbor pairs
Average neighbors per atom: 33.35

Understanding Half-Fill Mode#

Control whether to store both (i,j) and (j,i) pairs

print("\n" + "=" * 70)
print("HALF-FILL MODE")
print("=" * 70)

positions = small_positions
cell_tensor = small_cell

# Full neighbor list: both (i,j) and (j,i) stored
print("\n--- Full Neighbor List (half_fill=False) ---")
neighbor_matrix_full, num_neighbors_full, _ = cell_list(
    positions, cutoff, cell_tensor, pbc, half_fill=False
)
print(f"Total pairs: {num_neighbors_full.sum()}")

# Half-fill: only (i,j) where i < j (or with non-zero periodic shift)
print("\n--- Half-Fill Neighbor List (half_fill=True) ---")
neighbor_matrix_half, num_neighbors_half, _ = cell_list(
    positions, cutoff, cell_tensor, pbc, half_fill=True
)
print(f"Total pairs: {num_neighbors_half.sum()}")
print(f"Reduction: {num_neighbors_half.sum() / num_neighbors_full.sum():.1%} of full")

======================================================================
HALF-FILL MODE
======================================================================

--- Full Neighbor List (half_fill=False) ---
Total pairs: 5152

--- Half-Fill Neighbor List (half_fill=True) ---
Total pairs: 2576
Reduction: 50.0% of full

Validate neighbor distances#

print("\n" + "=" * 70)
print("DISTANCE VALIDATION")
print("=" * 70)

# Use neighbor list format for easy iteration
neighbor_list, neighbor_ptr, shifts = cell_list(
    small_positions, cutoff, small_cell, pbc, return_neighbor_list=True
)

if neighbor_list.shape[1] > 0:
    # Calculate distances for first 10 pairs
    n_check = min(10, neighbor_list.shape[1])
    source_atoms = neighbor_list[0, :n_check]
    target_atoms = neighbor_list[1, :n_check]
    pair_shifts = shifts[:n_check]

    # Compute Cartesian shifts
    cartesian_shifts = pair_shifts.float() @ small_cell.squeeze(0)

    # Calculate distances
    pos_i = small_positions[source_atoms]
    pos_j = small_positions[target_atoms]
    distances = torch.norm(pos_j - pos_i + cartesian_shifts, dim=1)

    print(f"\nFirst {n_check} neighbor distances:")
    for idx in range(n_check):
        src = source_atoms[idx].item()
        tgt = target_atoms[idx].item()
        dist = distances[idx].item()
        print(f"  Pair {idx}: atom {src} -> {tgt}, distance = {dist:.4f} Å")

    max_distance = distances.max().item()
    print(f"\nMax distance in pairs: {max_distance:.4f} Å")
    print(f"Cutoff distance: {cutoff:.4f} Å")
    if max_distance <= cutoff + 1e-5:
        print("✓ All neighbor distances are within the cutoff")
    else:
        print("✗ Some distances are outside the cutoff")

======================================================================
DISTANCE VALIDATION
======================================================================

First 10 neighbor distances:
  Pair 0: atom 0 -> 59, distance = 2.9541 Å
  Pair 1: atom 0 -> 54, distance = 3.7794 Å
  Pair 2: atom 0 -> 30, distance = 2.4518 Å
  Pair 3: atom 0 -> 55, distance = 3.4643 Å
  Pair 4: atom 0 -> 32, distance = 4.0754 Å
  Pair 5: atom 0 -> 22, distance = 3.9844 Å
  Pair 6: atom 0 -> 33, distance = 4.9199 Å
  Pair 7: atom 0 -> 97, distance = 3.2057 Å
  Pair 8: atom 0 -> 38, distance = 4.3813 Å
  Pair 9: atom 0 -> 90, distance = 4.8307 Å

Max distance in pairs: 4.9199 Å
Cutoff distance: 5.0000 Å
✓ All neighbor distances are within the cutoff

Differentiable distances and vectors#

The returned distance and vector tensors can participate in autograd.

print("\n" + "=" * 70)
print("DIFFERENTIABLE DISTANCES AND VECTORS")
print("=" * 70)

positions_for_gradients = small_positions.clone().detach().requires_grad_(True)

(
    neighbor_matrix,
    neighbor_counts,
    neighbor_shifts,
    neighbor_distances,
    neighbor_vectors,
) = cell_list(
    positions_for_gradients,
    cutoff,
    small_cell,
    pbc,
    return_distances=True,
    return_vectors=True,
)

neighbor_slots = torch.arange(neighbor_distances.shape[1], device=device)
valid_slots = neighbor_slots[None, :] < neighbor_counts[:, None]

if valid_slots.any():
    geometry_loss = (
        neighbor_distances[valid_slots].pow(2).mean()
        + 1.0e-4 * neighbor_vectors[valid_slots].pow(2).sum()
    )
    (position_gradients,) = torch.autograd.grad(
        geometry_loss,
        positions_for_gradients,
    )

    print(f"Neighbor matrix shape: {neighbor_matrix.shape}")
    print(f"Neighbor shifts shape: {neighbor_shifts.shape}")
    print(f"Geometry loss: {geometry_loss.item():.6f}")
    print(f"Position gradient norm: {position_gradients.norm().item():.6f}")
else:
    print("No neighbor pairs found for the differentiable geometry example.")

======================================================================
DIFFERENTIABLE DISTANCES AND VECTORS
======================================================================
Neighbor matrix shape: torch.Size([100, 112])
Neighbor shifts shape: torch.Size([100, 112, 3])
Geometry loss: 23.116781
Position gradient norm: 0.243981

print("\nExample completed successfully!")

Example completed successfully!

Total running time of the script: (0 minutes 0.204 seconds)

Gallery generated by Sphinx-Gallery

On this page
Show Source