.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/neighbors/01_simple_neighbor_list.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_01_simple_neighbor_list.py: 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. .. GENERATED FROM PYTHON SOURCE LINES 35-58 .. code-block:: Python 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, ) .. rst-class:: sphx-glr-script-out .. code-block:: none 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. .. GENERATED FROM PYTHON SOURCE LINES 59-62 Set up the computation device ============================= Choose device (CPU or GPU if available) .. GENERATED FROM PYTHON SOURCE LINES 62-113 .. 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}") _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] .. rst-class:: sphx-glr-script-out .. code-block:: none Using device: cuda Using dtype: torch.float32 .. GENERATED FROM PYTHON SOURCE LINES 114-117 Create random systems of different sizes ======================================== We'll compare performance on small vs large systems .. GENERATED FROM PYTHON SOURCE LINES 117-153 .. code-block:: Python 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} Å") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== CREATING TEST SYSTEMS ====================================================================== Small system: 100 atoms in 10.0³ box Large system: 4096 atoms in 40.0³ box Cutoff distance: 5.0 Å .. GENERATED FROM PYTHON SOURCE LINES 154-157 Method 1: Cell List Algorithm (O(N) - best for large systems) ============================================================= The cell list algorithm uses spatial decomposition for efficient neighbor finding .. GENERATED FROM PYTHON SOURCE LINES 157-204 .. code-block:: Python 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()}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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] .. GENERATED FROM PYTHON SOURCE LINES 205-208 Method 2: Naive Algorithm (O(N²) - best for small systems) ========================================================== The naive algorithm computes all pairwise distances .. GENERATED FROM PYTHON SOURCE LINES 208-234 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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]) .. GENERATED FROM PYTHON SOURCE LINES 235-238 Comparing Cell List vs Naive Algorithm ====================================== Compare performance and verify results match .. GENERATED FROM PYTHON SOURCE LINES 238-279 .. code-block:: Python 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." ) .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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. .. GENERATED FROM PYTHON SOURCE LINES 280-285 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=...``. .. GENERATED FROM PYTHON SOURCE LINES 285-364 .. code-block:: Python 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") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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 .. GENERATED FROM PYTHON SOURCE LINES 365-370 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. .. GENERATED FROM PYTHON SOURCE LINES 370-399 .. code-block:: Python 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. .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== METHOD 4: PREPARED STATE ====================================================================== Resolved method: naive Resolved strategy: naive_scalar Initialized: True Prepared pair count: 5152 Named result matches tuple: True .. GENERATED FROM PYTHON SOURCE LINES 400-403 Method 5: Low-level Cell List API with Caching ============================================== For advanced users: separate build and query phases for efficient caching .. GENERATED FROM PYTHON SOURCE LINES 403-472 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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 .. GENERATED FROM PYTHON SOURCE LINES 473-476 Understanding Half-Fill Mode ============================ Control whether to store both (i,j) and (j,i) pairs .. GENERATED FROM PYTHON SOURCE LINES 476-499 .. code-block:: Python 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") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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 .. GENERATED FROM PYTHON SOURCE LINES 500-502 Validate neighbor distances =========================== .. GENERATED FROM PYTHON SOURCE LINES 502-542 .. code-block:: Python 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") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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 .. GENERATED FROM PYTHON SOURCE LINES 543-546 Differentiable distances and vectors ==================================== The returned distance and vector tensors can participate in autograd. .. GENERATED FROM PYTHON SOURCE LINES 546-588 .. code-block:: Python 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.") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== 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 .. GENERATED FROM PYTHON SOURCE LINES 589-590 .. code-block:: Python print("\nExample completed successfully!") .. rst-class:: sphx-glr-script-out .. code-block:: none Example completed successfully! .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 3.511 seconds) .. _sphx_glr_download_examples_neighbors_01_simple_neighbor_list.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 01_simple_neighbor_list.ipynb <01_simple_neighbor_list.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 01_simple_neighbor_list.py <01_simple_neighbor_list.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 01_simple_neighbor_list.zip <01_simple_neighbor_list.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_