.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/neighbors/06_pair_outputs_lj.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_neighbors_06_pair_outputs_lj.py: Targeted Lennard-Jones Pair Outputs =================================== This example demonstrates four neighbor-list features used together: 1. ``target_indices`` for compact source rows 2. Per-neighbor vectors and distances 3. An inline Warp ``pair_fn`` that computes Lennard-Jones energies and forces 4. ``CompiledPairFn`` for CUDA fixed-shape ``torch.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. .. GENERATED FROM PYTHON SOURCE LINES 33-52 .. code-block:: Python 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, ) .. GENERATED FROM PYTHON SOURCE LINES 53-55 System setup ============ .. GENERATED FROM PYTHON SOURCE LINES 55-129 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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 .. GENERATED FROM PYTHON SOURCE LINES 130-132 Dispatch estimate ================= .. GENERATED FROM PYTHON SOURCE LINES 132-151 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none 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 .. GENERATED FROM PYTHON SOURCE LINES 152-154 Lennard-Jones pair function =========================== .. GENERATED FROM PYTHON SOURCE LINES 154-289 .. code-block:: Python @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, ) .. GENERATED FROM PYTHON SOURCE LINES 290-292 Inspect active pair outputs =========================== .. GENERATED FROM PYTHON SOURCE LINES 292-310 .. code-block:: Python 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") .. rst-class:: sphx-glr-script-out .. code-block:: none 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 .. GENERATED FROM PYTHON SOURCE LINES 311-313 CompiledPairFn fullgraph path ============================= .. GENERATED FROM PYTHON SOURCE LINES 313-463 .. code-block:: Python 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!") .. rst-class:: sphx-glr-script-out .. code-block:: none 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! .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.822 seconds) .. _sphx_glr_download_examples_neighbors_06_pair_outputs_lj.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 06_pair_outputs_lj.ipynb <06_pair_outputs_lj.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 06_pair_outputs_lj.py <06_pair_outputs_lj.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 06_pair_outputs_lj.zip <06_pair_outputs_lj.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_