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 state-free wrapper with automatic method selection

4. Comparison between algorithms

5. Neighbor matrix vs neighbor list (COO) formats

6. prepare_neighbor_list: Reusable state for repeated Torch execution

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

8. 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,
    prepare_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,
)

Warp UserWarning: Kernel cache artifacts from a previous Warp version were found in '/tmp/warp-cache'. These will be ignored. You can safely delete them.

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 3771, shift: [0, 0, 0]
  Atom 0 -> Atom 2255, shift: [0, 0, 0]
  Atom 0 -> Atom 1620, shift: [0, 0, 0]
  Atom 0 -> Atom 2094, shift: [0, 0, 0]
  Atom 0 -> Atom 1158, 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:  1.781 ms, 5152 pairs
Naive:      1.171 ms, 5152 pairs
Results match: True

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

Method 3: Unified State-Free neighbor_list Wrapper

The neighbor_list() wrapper provides a unified API and can choose a method from geometry-derived cost estimates. The same estimate can be printed once and reused by passing the selected method as method=....

print("\n" + "=" * 70)
print("METHOD 3: UNIFIED STATE-FREE neighbor_list() WRAPPER")
print("=" * 70)

print("\n--- Cost-Model Method Dispatch ---")

small_method = _dispatch_report(
    "Small system dispatch estimate:",
    small_positions,
    small_cell,
    pbc,
    cutoff,
)
nm_auto_small, num_auto_small, shifts_auto_small = neighbor_list(
    small_positions,
    cutoff,
    cell=small_cell,
    pbc=pbc,
    method=small_method,
)
print(f"  Total pairs: {num_auto_small.sum()}")
print(f"  Selected method: {small_method}")

large_method = _dispatch_report(
    "Large system dispatch estimate:",
    large_positions,
    large_cell,
    pbc,
    cutoff,
)
nm_auto_large, num_auto_large, shifts_auto_large = neighbor_list(
    large_positions,
    cutoff,
    cell=large_cell,
    pbc=pbc,
    method=large_method,
)
print(f"  Total pairs: {num_auto_large.sum()}")
print(f"  Selected method: {large_method}")

print("\n--- Explicit Method Selection ---")
# You can also explicitly specify the method
nm_explicit, num_explicit, shifts_explicit = neighbor_list(
    small_positions, cutoff, cell=small_cell, pbc=pbc, method="cell_list"
)
print(f"Explicitly using cell_list on small system: {num_explicit.sum()} pairs")

print("\n--- Passing Arguments to Underlying Methods ---")
# Pass kwargs to the underlying algorithm (e.g., max_neighbors)
nm_kwargs, num_kwargs, shifts_kwargs = neighbor_list(
    small_positions,
    cutoff,
    cell=small_cell,
    pbc=pbc,
    method="naive",
    max_neighbors=50,  # Passed to naive_neighbor_list
    half_fill=True,  # Also passed through
)
print(f"Using kwargs (max_neighbors=50, half_fill=True): {num_kwargs.sum()} pairs")
print(f"Neighbor matrix shape: {nm_kwargs.shape}")

print("\n--- Return Format Options ---")
# Get neighbor list (COO) format directly
nlist_coo, neighbor_ptr_coo, shifts_coo = neighbor_list(
    small_positions, cutoff, cell=small_cell, pbc=pbc, return_neighbor_list=True
)
print(f"Neighbor list (COO) format: {nlist_coo.shape}")
print(f"  Source atoms: {nlist_coo[0].shape}")
print(f"  Target atoms: {nlist_coo[1].shape}")
print(f"  Neighbor ptr: {neighbor_ptr_coo.shape}")

print("\n--- Benefits of the Wrapper ---")
print("✓ Unified API for all neighbor list methods")
print("✓ Cost-model dispatch based on geometry and requested outputs")
print("✓ Easy to switch between methods for testing")
print("✓ Familiar result formats; the tuple follows the selected route")
print("✓ Pass-through of method-specific kwargs")

======================================================================
METHOD 3: UNIFIED STATE-FREE neighbor_list() WRAPPER
======================================================================

--- Cost-Model Method Dispatch ---

Small system dispatch estimate:
  naive_tile               estimated cost (arbitrary units): 3.48e+05
  naive_scalar             estimated cost (arbitrary units): 3.48e+05
  cell_list_atom_centric   estimated cost (arbitrary units): 1.2e+06
  cell_list_pair_centric   estimated cost (arbitrary units): 1.36e+06
  Total pairs: 5152
  Selected method: naive_tile

Large system dispatch estimate:
  naive_tile               estimated cost (arbitrary units): 9.89e+05
  naive_scalar             estimated cost (arbitrary units): 9.94e+05
  cell_list_atom_centric   estimated cost (arbitrary units): 1.49e+06
  cell_list_pair_centric   estimated cost (arbitrary units): 1.81e+06
  Total pairs: 136604
  Selected method: naive_tile

--- Explicit Method Selection ---
Explicitly using cell_list on small system: 5152 pairs

--- Passing Arguments to Underlying Methods ---
Using kwargs (max_neighbors=50, half_fill=True): 2576 pairs
Neighbor matrix shape: torch.Size([100, 50])

--- Return Format Options ---
Neighbor list (COO) format: torch.Size([2, 5152])
  Source atoms: torch.Size([5152])
  Target atoms: torch.Size([5152])
  Neighbor ptr: torch.Size([101])

--- Benefits of the Wrapper ---
✓ Unified API for all neighbor list methods
✓ Cost-model dispatch based on geometry and requested outputs
✓ Easy to switch between methods for testing
✓ Familiar result formats; the tuple follows the selected route
✓ Pass-through of method-specific kwargs

Method 4: Prepared State for Repeated Torch Execution

Prepared state is the recommended high-level Torch workflow when fixed metadata and capacities are reused. The ordinary result tuple is unchanged, while the state owns reusable route buffers and exposes the latest result by name.

print("\n" + "=" * 70)
print("METHOD 4: PREPARED STATE")
print("=" * 70)

prepared_state = prepare_neighbor_list(
    small_positions,
    cutoff,
    cell=small_cell,
    pbc=pbc,
    method="naive",
    strategy="scalar",
    max_neighbors=128,
)
prepared_result = neighbor_list(small_positions, state=prepared_state)
prepared_matrix, prepared_counts, prepared_shifts = prepared_result

print(f"Resolved method: {prepared_state.method}")
print(f"Resolved strategy: {prepared_state.strategy}")
print(f"Initialized: {bool(prepared_state.initialized.all())}")
print(f"Prepared pair count: {int(prepared_counts.sum())}")
print(
    f"Named result matches tuple: {prepared_state.neighbor_matrix is prepared_matrix}"
)

# With ``method=None`` and ``strategy="auto"``, Torch resolves its ordinary
# selector once during preparation. Inspect the properties above to see the
# result. Explicit choices pin the implementation and raise when unsupported.

======================================================================
METHOD 4: PREPARED STATE
======================================================================
Resolved method: naive
Resolved strategy: naive_scalar
Initialized: True
Prepared pair count: 5152
Named result matches tuple: True

Method 5: Low-level Cell List API with Caching

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

print("\n" + "=" * 70)
print("METHOD 5: 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 5: 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 ---
/tmp/tmp268gp85d/6702138e404d40ea79dec62f68c2fa67e858580e/examples/neighbors/01_simple_neighbor_list.py:449: 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 -> 54, distance = 3.7794 Å
  Pair 1: atom 0 -> 55, distance = 3.4643 Å
  Pair 2: atom 0 -> 90, distance = 4.8307 Å
  Pair 3: atom 0 -> 33, distance = 4.9199 Å
  Pair 4: atom 0 -> 30, distance = 2.4518 Å
  Pair 5: atom 0 -> 32, distance = 4.0754 Å
  Pair 6: atom 0 -> 58, distance = 1.8651 Å
  Pair 7: atom 0 -> 39, distance = 3.0348 Å
  Pair 8: atom 0 -> 80, distance = 4.4274 Å
  Pair 9: atom 0 -> 59, distance = 2.9541 Å

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!