Note
Go to the end to download the full example code.
Targeted Lennard-Jones Pair Outputs#
This example demonstrates four neighbor-list features used together:
target_indicesfor compact source rowsPer-neighbor vectors and distances
An inline Warp
pair_fnthat computes Lennard-Jones energies and forcesCompiledPairFnfor CUDA fixed-shapetorch.compile(fullgraph=True)calls
The compiled section is skipped when CUDA is unavailable.
The system is a 4x4x4 FCC Argon box using the same parameters as the dynamics examples: epsilon = 0.0104 eV, sigma = 3.40 Å, lattice constant = 5.26 Å, and cutoff = 2.5 * sigma.
import torch
import warp as wp
from nvalchemiops.torch.neighbors import (
compile_pair_fn,
estimate_neighbor_list_costs,
neighbor_list,
)
from nvalchemiops.torch.neighbors.cell_list import (
allocate_query_sort_scratch,
cell_list,
estimate_cell_list_sizes,
)
from nvalchemiops.torch.neighbors.neighbor_utils import (
allocate_cell_list,
estimate_max_neighbors,
)
System setup#
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32
EPSILON_AR = 0.0104
SIGMA_AR = 3.40
LATTICE_A_AR = 5.26
CUTOFF = 2.5 * SIGMA_AR
NUM_UNIT_CELLS = 4
print("=" * 70)
print("TARGETED LENNARD-JONES PAIR OUTPUTS")
print("=" * 70)
print(f"Using device: {device}")
print(f"Using dtype: {dtype}")
def create_fcc_argon_box(
num_unit_cells: int,
lattice_constant: float,
*,
dtype: torch.dtype,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Create an FCC Argon box with periodic boundary conditions."""
basis = torch.tensor(
[
[0.0, 0.0, 0.0],
[0.5, 0.5, 0.0],
[0.5, 0.0, 0.5],
[0.0, 0.5, 0.5],
],
dtype=dtype,
device=device,
)
unit_axis = torch.arange(num_unit_cells, dtype=dtype, device=device)
grid = torch.stack(
torch.meshgrid(unit_axis, unit_axis, unit_axis, indexing="ij"),
dim=-1,
).reshape(-1, 3)
positions = (grid[:, None, :] + basis[None, :, :]).reshape(-1, 3)
positions = positions * lattice_constant
box_length = float(num_unit_cells) * lattice_constant
cell = (torch.eye(3, dtype=dtype, device=device) * box_length).unsqueeze(0)
pbc = torch.tensor([[True, True, True]], dtype=torch.bool, device=device)
return positions, cell, pbc
positions, cell, pbc = create_fcc_argon_box(
NUM_UNIT_CELLS,
LATTICE_A_AR,
dtype=dtype,
device=device,
)
num_atoms = positions.shape[0]
target_indices = torch.arange(0, num_atoms, 4, dtype=torch.int32, device=device)
batch_ptr = torch.tensor([0, num_atoms], dtype=torch.int32, device=device)
argon_density = 4.0 / (LATTICE_A_AR**3)
max_neighbors = estimate_max_neighbors(
CUTOFF,
atomic_density=argon_density,
safety_factor=1.5,
)
print(f"\nFCC Argon box: {num_atoms} atoms")
print(f"Box length: {cell[0, 0, 0].item():.2f} Å")
print(f"LJ parameters: epsilon={EPSILON_AR:.4f} eV, sigma={SIGMA_AR:.2f} Å")
print(f"Cutoff: {CUTOFF:.2f} Å")
print(f"Targeted source rows: {target_indices.numel()} of {num_atoms}")
print(f"Estimated max_neighbors: {max_neighbors}")
======================================================================
TARGETED LENNARD-JONES PAIR OUTPUTS
======================================================================
Using device: cuda
Using dtype: torch.float32
/lustre/fs1/portfolios/coreai/projects/coreai_modulus_alchemi/users/kinlongkelvi/code/nvalchemi-toolkit-ops/examples/neighbors/06_pair_outputs_lj.py:115: 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(
FCC Argon box: 256 atoms
Box length: 21.04 Å
LJ parameters: epsilon=0.0104 eV, sigma=3.40 Å
Cutoff: 8.50 Å
Targeted source rows: 64 of 256
Estimated max_neighbors: 112
Dispatch estimate#
cost_report = estimate_neighbor_list_costs(
batch_ptr,
cell,
pbc,
CUTOFF,
target_indices=target_indices,
return_vectors=True,
return_distances=True,
use_pair_fn=True,
positions_dtype=positions.dtype,
)
print("\nTorch dispatch estimate for targeted LJ outputs:")
for method_name, cost in cost_report:
print(f" {method_name:24s} estimated cost (arbitrary units): {cost:.3g}")
selected_method = cost_report[0][0]
print(f"Selected method: {selected_method}")
Torch dispatch estimate for targeted LJ outputs:
cell_list_pair_centric estimated cost (arbitrary units): 3.4e+04
cell_list_atom_centric estimated cost (arbitrary units): 6.61e+04
naive_scalar estimated cost (arbitrary units): 9.14e+04
Selected method: cell_list_pair_centric
Lennard-Jones pair function#
@wp.func
def lj_pair_fn(
r_ij: wp.vec3f,
distance: wp.float32,
pair_params: wp.array2d(dtype=wp.float32),
i: int,
j: int,
):
"""Compute Lennard-Jones pair energy and force from per-atom parameters."""
epsilon = wp.sqrt(pair_params[i, 0] * pair_params[j, 0])
sigma = 0.5 * (pair_params[i, 1] + pair_params[j, 1])
sr = sigma / distance
sr2 = sr * sr
sr6 = sr2 * sr2 * sr2
sr12 = sr6 * sr6
energy = 4.0 * epsilon * (sr12 - sr6)
force = (24.0 * epsilon * (sr6 - 2.0 * sr12) / (distance * distance)) * r_ij
return energy, force
def allocate_lj_pair_output_buffers(
num_rows: int,
max_neighbors: int,
fill_value: int,
*,
dtype: torch.dtype,
device: torch.device,
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
"""Allocate fixed-shape matrix outputs for the LJ pair-output example."""
neighbor_matrix = torch.full(
(num_rows, max_neighbors),
fill_value,
dtype=torch.int32,
device=device,
)
neighbor_shifts = torch.zeros(
(num_rows, max_neighbors, 3),
dtype=torch.int32,
device=device,
)
neighbor_counts = torch.zeros(num_rows, dtype=torch.int32, device=device)
neighbor_vectors = torch.zeros(
(num_rows, max_neighbors, 3),
dtype=dtype,
device=device,
)
neighbor_distances = torch.zeros(
(num_rows, max_neighbors),
dtype=dtype,
device=device,
)
pair_energies = torch.zeros(
(num_rows, max_neighbors),
dtype=dtype,
device=device,
)
pair_forces = torch.zeros(
(num_rows, max_neighbors, 3),
dtype=dtype,
device=device,
)
return (
neighbor_matrix,
neighbor_shifts,
neighbor_counts,
neighbor_vectors,
neighbor_distances,
pair_energies,
pair_forces,
)
pair_params = torch.empty((num_atoms, 2), dtype=dtype, device=device)
pair_params[:, 0] = EPSILON_AR
pair_params[:, 1] = SIGMA_AR
num_targets = target_indices.numel()
fill_value = num_atoms
(
neighbor_matrix,
neighbor_shifts,
neighbor_counts,
neighbor_vectors,
neighbor_distances,
pair_energies,
pair_forces,
) = allocate_lj_pair_output_buffers(
num_targets,
max_neighbors,
fill_value,
dtype=dtype,
device=device,
)
(
neighbor_matrix,
neighbor_counts,
neighbor_shifts,
neighbor_distances,
neighbor_vectors,
pair_energies,
pair_forces,
) = neighbor_list(
positions,
CUTOFF,
cell=cell,
pbc=pbc,
method=selected_method,
target_indices=target_indices,
fill_value=fill_value,
max_neighbors=max_neighbors,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_shifts,
num_neighbors=neighbor_counts,
return_vectors=True,
return_distances=True,
pair_fn=lj_pair_fn,
pair_params=pair_params,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
Inspect active pair outputs#
neighbor_slots = torch.arange(neighbor_matrix.shape[1], device=device)
valid_slots = neighbor_slots[None, :] < neighbor_counts[:, None]
source_atoms = target_indices[:, None].expand_as(neighbor_matrix)[valid_slots].long()
target_atoms = neighbor_matrix[valid_slots].long()
active_distances = neighbor_distances[valid_slots]
active_energies = pair_energies[valid_slots]
active_forces = pair_forces[valid_slots]
active_pair_count = int(active_energies.numel())
active_energy_sum = active_energies.sum()
print("\nTargeted pair-output statistics:")
print(f" Compact output rows: {num_targets}")
print(f" Active targeted pairs: {active_pair_count}")
print(f" Average targeted neighbors: {neighbor_counts.float().mean().item():.2f}")
print(f" Targeted directed LJ energy sum: {active_energy_sum.item():.6f} eV")
Targeted pair-output statistics:
Compact output rows: 64
Active targeted pairs: 4992
Average targeted neighbors: 78.00
Targeted directed LJ energy sum: -10.781156 eV
CompiledPairFn fullgraph path#
if device.type != "cuda":
print("\nCompiledPairFn fullgraph demo skipped: CUDA is required.")
else:
compiled_lj_pair_fn = compile_pair_fn(lj_pair_fn, name="lj_pair")
(
compiled_neighbor_matrix,
compiled_neighbor_shifts,
compiled_neighbor_counts,
compiled_neighbor_vectors,
compiled_neighbor_distances,
compiled_pair_energies,
compiled_pair_forces,
) = allocate_lj_pair_output_buffers(
num_targets,
max_neighbors,
fill_value,
dtype=dtype,
device=device,
)
max_cells, radius = estimate_cell_list_sizes(
cell,
pbc.reshape(3),
CUTOFF,
min_cells_per_dimension=4,
)
(
cells_per_dimension,
neighbor_search_radius,
atom_periodic_shifts,
atom_to_cell_mapping,
atoms_per_cell_count,
cell_atom_start_indices,
cell_atom_list,
) = allocate_cell_list(num_atoms, max_cells, radius, device)
sorted_positions, sorted_shifts = allocate_query_sort_scratch(
num_atoms,
dtype=dtype,
device=device,
)
@torch.compile(fullgraph=True)
def compiled_targeted_lj_outputs(
positions: torch.Tensor,
neighbor_matrix: torch.Tensor,
neighbor_shifts: torch.Tensor,
neighbor_counts: torch.Tensor,
neighbor_vectors: torch.Tensor,
neighbor_distances: torch.Tensor,
pair_energies: torch.Tensor,
pair_forces: torch.Tensor,
) -> tuple[
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
torch.Tensor,
]:
"""Run the fixed-shape CompiledPairFn cell-list path under fullgraph."""
return cell_list(
positions,
CUTOFF,
cell,
pbc,
max_neighbors=max_neighbors,
fill_value=fill_value,
neighbor_matrix=neighbor_matrix,
neighbor_matrix_shifts=neighbor_shifts,
num_neighbors=neighbor_counts,
cells_per_dimension=cells_per_dimension,
neighbor_search_radius=neighbor_search_radius,
atom_periodic_shifts=atom_periodic_shifts,
atom_to_cell_mapping=atom_to_cell_mapping,
atoms_per_cell_count=atoms_per_cell_count,
cell_atom_start_indices=cell_atom_start_indices,
cell_atom_list=cell_atom_list,
strategy="atom_centric",
atom_centric_path="direct",
sorted_positions=sorted_positions,
sorted_shifts=sorted_shifts,
target_indices=target_indices,
return_vectors=True,
return_distances=True,
neighbor_vectors=neighbor_vectors,
neighbor_distances=neighbor_distances,
pair_fn=compiled_lj_pair_fn,
pair_params=pair_params,
pair_energies=pair_energies,
pair_forces=pair_forces,
)
(
compiled_neighbor_matrix,
compiled_neighbor_counts,
compiled_neighbor_shifts,
compiled_neighbor_distances,
compiled_neighbor_vectors,
compiled_pair_energies,
compiled_pair_forces,
) = compiled_targeted_lj_outputs(
positions,
compiled_neighbor_matrix,
compiled_neighbor_shifts,
compiled_neighbor_counts,
compiled_neighbor_vectors,
compiled_neighbor_distances,
compiled_pair_energies,
compiled_pair_forces,
)
compiled_neighbor_slots = torch.arange(
compiled_neighbor_matrix.shape[1], device=device
)
compiled_valid_slots = (
compiled_neighbor_slots[None, :] < compiled_neighbor_counts[:, None]
)
compiled_active_energies = compiled_pair_energies[compiled_valid_slots]
compiled_active_pair_count = int(compiled_active_energies.numel())
if compiled_active_pair_count != active_pair_count:
raise RuntimeError(
"CompiledPairFn path produced a different active-pair count "
f"({compiled_active_pair_count}) than the eager path "
f"({active_pair_count})."
)
print("\nCompiledPairFn fullgraph statistics:")
print(f" Active targeted pairs: {compiled_active_pair_count}")
print(
" Compiled directed LJ energy sum: "
f"{compiled_active_energies.sum().item():.6f} eV"
)
print("\nFirst targeted LJ pairs:")
sample_count = min(5, source_atoms.numel())
for pair_index in range(sample_count):
source = int(source_atoms[pair_index].item())
target = int(target_atoms[pair_index].item())
distance = float(active_distances[pair_index].item())
energy = float(active_energies[pair_index].item())
force_norm = float(active_forces[pair_index].norm().item())
print(
f" Atom {source} -> {target}: "
f"distance={distance:.4f} Å, energy={energy:.6f} eV, "
f"|force|={force_norm:.6f} eV/Å"
)
print("\nExample completed successfully!")
CompiledPairFn fullgraph statistics:
Active targeted pairs: 4992
Compiled directed LJ energy sum: -10.781155 eV
First targeted LJ pairs:
Atom 0 -> 243: distance=6.4422 Å, energy=-0.000880 eV, |force|=0.000801 eV/Å
Atom 0 -> 253: distance=6.4422 Å, energy=-0.000880 eV, |force|=0.000801 eV/Å
Atom 0 -> 202: distance=8.3168 Å, energy=-0.000193 eV, |force|=0.000139 eV/Å
Atom 0 -> 206: distance=3.7194 Å, energy=-0.010110 eV, |force|=0.006540 eV/Å
Atom 0 -> 193: distance=3.7194 Å, energy=-0.010110 eV, |force|=0.006540 eV/Å
Example completed successfully!
Total running time of the script: (0 minutes 0.822 seconds)