.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/neighborlist/03_rebuild_neighborlist_detection.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_neighborlist_03_rebuild_neighborlist_detection.py: Neighbor List Rebuild Detection Example ======================================= This example demonstrates how to use rebuild detection functions in nvalchemiops to efficiently determine when neighbor lists need to be reconstructed during molecular dynamics simulations. We'll cover: - cell_list_needs_rebuild: Detect when atoms move between spatial cells - neighbor_list_needs_rebuild: Detect when atoms exceed skin distance - Skin distance approach for efficient neighbor list caching - Integration with build_cell_list + query_cell_list for MD workflows Rebuild detection is crucial for MD performance - neighbor lists are expensive to compute but only need updating when atoms have moved significantly. Smart rebuild detection can improve simulation performance by 2-10x. .. GENERATED FROM PYTHON SOURCE LINES 33-47 .. code-block:: Python import numpy as np import torch from nvalchemiops.neighborlist import ( allocate_cell_list, build_cell_list, cell_list_needs_rebuild, estimate_cell_list_sizes, estimate_max_neighbors, neighbor_list_needs_rebuild, query_cell_list, ) .. GENERATED FROM PYTHON SOURCE LINES 48-50 Set up the computation device and simulation parameters ======================================================= .. GENERATED FROM PYTHON SOURCE LINES 50-69 .. code-block:: Python device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.float32 print(f"Using device: {device}") print(f"Using dtype: {dtype}") # Simulation parameters num_atoms = 128 box_size = 10.0 cutoff = 2.5 skin_distance = 0.5 # Buffer distance to avoid frequent rebuilds total_cutoff = cutoff + skin_distance print("\nSimulation Parameters:") print(f" System: {num_atoms} atoms in {box_size}³ box") print(f" Neighbor cutoff: {cutoff} Å") print(f" Skin distance: {skin_distance} Å") print(f" Total cutoff (neighbor + skin): {total_cutoff} Å") .. rst-class:: sphx-glr-script-out .. code-block:: none Using device: cuda Using dtype: torch.float32 Simulation Parameters: System: 128 atoms in 10.0³ box Neighbor cutoff: 2.5 Å Skin distance: 0.5 Å Total cutoff (neighbor + skin): 3.0 Å .. GENERATED FROM PYTHON SOURCE LINES 70-72 Create initial system configuration =================================== .. GENERATED FROM PYTHON SOURCE LINES 72-96 .. code-block:: Python print("\n" + "=" * 70) print("INITIAL SYSTEM SETUP") print("=" * 70) # Create simple cubic lattice n_side = int(np.ceil(num_atoms ** (1 / 3))) lattice_spacing = box_size / n_side # Generate lattice positions d = (torch.arange(n_side, dtype=dtype, device=device) + 0.5) * lattice_spacing di, dj, dk = torch.meshgrid(d, d, d, indexing="ij") positions_lattice = torch.stack([di.flatten(), dj.flatten(), dk.flatten()], dim=1) initial_positions = positions_lattice[:num_atoms].clone() # System setup cell = (torch.eye(3, dtype=dtype, device=device) * box_size).unsqueeze(0) pbc = torch.tensor([True, True, True], device=device) print(f"Created lattice with spacing {lattice_spacing:.3f} Å") print( f"Initial position range: {initial_positions.min().item():.3f} to {initial_positions.max().item():.3f}" ) .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== INITIAL SYSTEM SETUP ====================================================================== Created lattice with spacing 1.667 Å Initial position range: 0.833 to 9.167 .. GENERATED FROM PYTHON SOURCE LINES 97-99 Build initial neighbor list with skin distance =============================================== .. GENERATED FROM PYTHON SOURCE LINES 99-167 .. code-block:: Python print("\n" + "=" * 70) print("BUILDING INITIAL NEIGHBOR LIST") print("=" * 70) # Estimate memory requirements max_total_cells, neighbor_search_radius = estimate_cell_list_sizes( cell, pbc, total_cutoff ) print("Memory estimates:") print(f" Max cells: {max_total_cells}") print(f" Neighbor search radius: {neighbor_search_radius}") # Allocate cell list cache cell_list_cache = allocate_cell_list( total_atoms=num_atoms, 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 # Build cell list with total_cutoff (including skin) build_cell_list(initial_positions, total_cutoff, cell, pbc, *cell_list_cache) print("\nBuilt cell list:") print(f" Cells per dimension: {cells_per_dimension.tolist()}") print(f" Neighbor search radius: {neighbor_search_radius.tolist()}") # Query to get initial neighbors (using actual cutoff, not total) max_neighbors = estimate_max_neighbors(total_cutoff) neighbor_matrix = torch.full( (num_atoms, max_neighbors), -1, dtype=torch.int32, device=device ) neighbor_shifts = torch.zeros( (num_atoms, max_neighbors, 3), dtype=torch.int32, device=device ) num_neighbors_arr = torch.zeros(num_atoms, dtype=torch.int32, device=device) query_cell_list( initial_positions, cutoff, cell, pbc, *cell_list_cache, neighbor_matrix, neighbor_shifts, num_neighbors_arr, ) print(f"\nInitial neighbor list (cutoff={cutoff}):") print(f" Total pairs: {num_neighbors_arr.sum()}") print(f" Avg neighbors per atom: {num_neighbors_arr.float().mean():.2f}") # Save reference for rebuild detection reference_positions = initial_positions.clone() reference_atom_to_cell_mapping = atom_to_cell_mapping.clone() .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== BUILDING INITIAL NEIGHBOR LIST ====================================================================== Memory estimates: Max cells: 27 Neighbor search radius: tensor([1, 1, 1], device='cuda:0', dtype=torch.int32) Built cell list: Cells per dimension: [3, 3, 3] Neighbor search radius: [1, 1, 1] Initial neighbor list (cutoff=2.5): Total pairs: 1906 Avg neighbors per atom: 14.89 .. GENERATED FROM PYTHON SOURCE LINES 168-170 Simulate atomic motion and test rebuild detection ================================================= .. GENERATED FROM PYTHON SOURCE LINES 170-255 .. code-block:: Python print("\n" + "=" * 70) print("SIMULATING ATOMIC MOTION") print("=" * 70) # Simulate a sequence of small displacements n_steps = 20 displacement_per_step = 0.15 # Small displacement per step rebuild_count = 0 print(f"\nSimulating {n_steps} MD steps:") print(f" Displacement per step: {displacement_per_step} Å") print(f" Skin distance: {skin_distance} Å") print() old_positions = reference_positions.clone() for step in range(n_steps): # Apply random small displacement displacement = ( torch.rand(num_atoms, 3, device=device, dtype=dtype) - 0.5 ) * displacement_per_step current_positions = old_positions + displacement # Apply periodic boundary conditions current_positions = current_positions % box_size # Check if cell list needs rebuild (atoms moved between cells) cell_rebuild_needed = cell_list_needs_rebuild( current_positions=current_positions, atom_to_cell_mapping=reference_atom_to_cell_mapping, cells_per_dimension=cells_per_dimension, cell=cell, pbc=pbc, ) # Check if neighbor list needs rebuild (exceeded skin distance) neighbor_rebuild_needed = neighbor_list_needs_rebuild( reference_positions, current_positions, skin_distance, ) # Calculate max atomic displacement for reference displacements = current_positions - reference_positions # Account for PBC displacements = displacements - torch.round(displacements / box_size) * box_size max_displacement = torch.norm(displacements, dim=1).max().item() status = "" if cell_rebuild_needed.item() or neighbor_rebuild_needed.item(): # Rebuild! rebuild_count += 1 status = "REBUILD" # Rebuild cell list build_cell_list(current_positions, total_cutoff, cell, pbc, *cell_list_cache) # Update reference reference_positions = current_positions.clone() reference_atom_to_cell_mapping = atom_to_cell_mapping.clone() print( f"Step {step:2d}: max_disp={max_displacement:.4f} Å " f"cell_rebuild={cell_rebuild_needed.item()}, " f"neighbor_rebuild={neighbor_rebuild_needed.item()} {status}" ) # Query neighbors (always use actual cutoff, not total_cutoff) query_cell_list( current_positions, cutoff, cell, pbc, *cell_list_cache, neighbor_matrix, neighbor_shifts, num_neighbors_arr, ) old_positions = current_positions.clone() print("\nRebuild Statistics:") print(f" Total rebuilds: {rebuild_count} / {n_steps} steps") print(f" Rebuild rate: {rebuild_count / n_steps * 100:.1f}%") print(f" Performance gain: ~{n_steps / max(1, rebuild_count):.1f}x") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== SIMULATING ATOMIC MOTION ====================================================================== Simulating 20 MD steps: Displacement per step: 0.15 Å Skin distance: 0.5 Å Step 0: max_disp=0.1233 Å cell_rebuild=False, neighbor_rebuild=False Step 1: max_disp=0.2101 Å cell_rebuild=False, neighbor_rebuild=False Step 2: max_disp=0.2361 Å cell_rebuild=False, neighbor_rebuild=False Step 3: max_disp=0.3088 Å cell_rebuild=False, neighbor_rebuild=False Step 4: max_disp=0.3121 Å cell_rebuild=False, neighbor_rebuild=False Step 5: max_disp=0.3638 Å cell_rebuild=False, neighbor_rebuild=False Step 6: max_disp=0.4079 Å cell_rebuild=False, neighbor_rebuild=False Step 7: max_disp=0.4622 Å cell_rebuild=False, neighbor_rebuild=False Step 8: max_disp=0.5180 Å cell_rebuild=False, neighbor_rebuild=True REBUILD Step 9: max_disp=0.1204 Å cell_rebuild=False, neighbor_rebuild=False Step 10: max_disp=0.1787 Å cell_rebuild=False, neighbor_rebuild=False Step 11: max_disp=0.2426 Å cell_rebuild=False, neighbor_rebuild=False Step 12: max_disp=0.3269 Å cell_rebuild=False, neighbor_rebuild=False Step 13: max_disp=0.3055 Å cell_rebuild=False, neighbor_rebuild=False Step 14: max_disp=0.3369 Å cell_rebuild=False, neighbor_rebuild=False Step 15: max_disp=0.3679 Å cell_rebuild=False, neighbor_rebuild=False Step 16: max_disp=0.4242 Å cell_rebuild=False, neighbor_rebuild=False Step 17: max_disp=0.4376 Å cell_rebuild=False, neighbor_rebuild=False Step 18: max_disp=0.4848 Å cell_rebuild=False, neighbor_rebuild=False Step 19: max_disp=0.5234 Å cell_rebuild=False, neighbor_rebuild=True REBUILD Rebuild Statistics: Total rebuilds: 2 / 20 steps Rebuild rate: 10.0% Performance gain: ~10.0x .. GENERATED FROM PYTHON SOURCE LINES 256-258 Demonstrate large atomic motion causing rebuild =============================================== .. GENERATED FROM PYTHON SOURCE LINES 258-301 .. code-block:: Python print("\n" + "=" * 70) print("LARGE DISPLACEMENT TEST") print("=" * 70) # Reset to initial configuration current_positions = initial_positions.clone() reference_positions = initial_positions.clone() # Build fresh cell list build_cell_list(current_positions, total_cutoff, cell, pbc, *cell_list_cache) reference_atom_to_cell_mapping = atom_to_cell_mapping.clone() print("\nTesting with increasing displacements:") for displacement_magnitude in [0.1, 0.3, 0.5, 0.7, 1.0]: # Apply displacement to a few atoms displaced_positions = reference_positions.clone() displaced_positions[:10] += displacement_magnitude # Check rebuild need cell_rebuild = cell_list_needs_rebuild( current_positions=displaced_positions, atom_to_cell_mapping=reference_atom_to_cell_mapping, cells_per_dimension=cells_per_dimension, cell=cell, pbc=pbc, ) neighbor_rebuild = neighbor_list_needs_rebuild( reference_positions, displaced_positions, skin_distance, ) rebuild_status = "YES" if (cell_rebuild.item() or neighbor_rebuild.item()) else "NO" print( f" Displacement {displacement_magnitude:.1f} Å: " f"cell={cell_rebuild.item()}, neighbor={neighbor_rebuild.item()} " f"-> Rebuild: {rebuild_status}" ) print("\nExample completed successfully!") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== LARGE DISPLACEMENT TEST ====================================================================== Testing with increasing displacements: Displacement 0.1 Å: cell=False, neighbor=False -> Rebuild: NO Displacement 0.3 Å: cell=False, neighbor=True -> Rebuild: YES Displacement 0.5 Å: cell=False, neighbor=True -> Rebuild: YES Displacement 0.7 Å: cell=False, neighbor=True -> Rebuild: YES Displacement 1.0 Å: cell=True, neighbor=True -> Rebuild: YES Example completed successfully! .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.034 seconds) .. _sphx_glr_download_examples_neighborlist_03_rebuild_neighborlist_detection.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 03_rebuild_neighborlist_detection.ipynb <03_rebuild_neighborlist_detection.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 03_rebuild_neighborlist_detection.py <03_rebuild_neighborlist_detection.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 03_rebuild_neighborlist_detection.zip <03_rebuild_neighborlist_detection.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_