.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/neighbors/05_jax_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_05_jax_neighbor_list.py: JAX Neighbor List Example ========================= This example demonstrates how to use the JAX neighbor list API in nvalchemiops for computing neighbor lists in periodic systems. In this example you will learn: - How to use the unified ``neighbor_list()`` API with JAX arrays - Matrix format vs COO (list) format outputs - Comparing ``naive_neighbor_list`` and ``cell_list`` algorithms - Using ``half_fill`` mode for symmetric neighbor lists - Building compact partial lists with ``target_indices`` - Validating neighbor distances are within cutoff - Preparing and threading reusable state through ``jax.jit`` - ``jax.jit`` compilation of matrix and fixed-capacity COO outputs - Estimating dispatch cost with ``estimate_neighbor_list_costs`` / ``suggest_neighbor_list_method`` - Evaluating an inline Warp ``pair_fn`` (per-pair energy and force) .. important:: This example is for educational purposes. Do not use it for performance benchmarking, as the code includes print statements and small system sizes that are not representative of production workloads. .. GENERATED FROM PYTHON SOURCE LINES 43-74 .. code-block:: Python import sys try: import jax import jax.numpy as jnp except ImportError: print( "This example requires JAX. Install with: pip install 'nvalchemi-toolkit-ops[jax]'" ) sys.exit(0) try: import warp as wp from nvalchemiops.jax.neighbors import ( check_neighbor_list_state, compute_naive_num_shifts, estimate_neighbor_list_costs, neighbor_list, prepare_neighbor_list, suggest_neighbor_list_method, ) except Exception as exc: print( f"JAX/Warp backend unavailable ({exc}). This example requires a CUDA-backed runtime." ) sys.exit(0) from nvalchemiops.jax.neighbors.cell_list import cell_list from nvalchemiops.jax.neighbors.naive import naive_neighbor_list .. 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 75-79 Setup ===== JAX handles device placement automatically. We'll create a random periodic system to demonstrate the neighbor list API. .. GENERATED FROM PYTHON SOURCE LINES 79-107 .. code-block:: Python print("=" * 70) print("JAX NEIGHBOR LIST EXAMPLE") print("=" * 70) # System parameters num_atoms = 200 box_size = 15.0 cutoff = 5.0 # Create random atomic positions using JAX random key = jax.random.PRNGKey(42) positions = jax.random.uniform(key, (num_atoms, 3), dtype=jnp.float32) * box_size # Create a cubic periodic cell: (1, 3, 3) shape cell = jnp.eye(3, dtype=jnp.float32)[None, ...] * box_size # Enable periodic boundary conditions in all directions: (1, 3) shape pbc = jnp.array([[True, True, True]]) print("\nSystem configuration:") print(f" Number of atoms: {num_atoms}") print(f" Box size: {box_size} Å") print(f" Cutoff distance: {cutoff} Å") print(f" Positions shape: {positions.shape}") print(f" Cell shape: {cell.shape}") print(f" PBC shape: {pbc.shape}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== JAX NEIGHBOR LIST EXAMPLE ====================================================================== System configuration: Number of atoms: 200 Box size: 15.0 Å Cutoff distance: 5.0 Å Positions shape: (200, 3) Cell shape: (1, 3, 3) PBC shape: (1, 3) .. GENERATED FROM PYTHON SOURCE LINES 108-113 Unified API - Matrix Format (default) ===================================== The ``neighbor_list()`` function provides a consistent entry point for JAX neighbor-list construction while preserving the same matrix-format outputs used by the direct algorithms. .. GENERATED FROM PYTHON SOURCE LINES 113-140 .. code-block:: Python print("\n" + "=" * 70) print("UNIFIED API - MATRIX FORMAT") print("=" * 70) # Call the unified API (returns matrix format by default) neighbor_matrix, num_neighbors, shifts = neighbor_list( positions, cutoff, cell=cell, pbc=pbc ) print("\nReturned neighbor matrix format:") print(f" neighbor_matrix shape: {neighbor_matrix.shape}") print(f" num_neighbors shape: {num_neighbors.shape}") print(f" shifts shape: {shifts.shape}") print("\nStatistics:") print(f" Total neighbor pairs: {int(num_neighbors.sum())}") print(f" Average neighbors per atom: {float(num_neighbors.mean()):.2f}") print(f" Max neighbors for any atom: {int(num_neighbors.max())}") print(f" Min neighbors for any atom: {int(num_neighbors.min())}") # Show first few neighbors of atom 0 print("\nFirst 5 neighbors of atom 0:") for i in range(min(5, int(num_neighbors[0]))): neighbor_idx = int(neighbor_matrix[0, i]) shift = shifts[0, i].tolist() print(f" Neighbor {i}: atom {neighbor_idx}, shift {shift}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== UNIFIED API - MATRIX FORMAT ====================================================================== Returned neighbor matrix format: neighbor_matrix shape: (200, 112) num_neighbors shape: (200,) shifts shape: (200, 112, 3) Statistics: Total neighbor pairs: 6166 Average neighbors per atom: 30.83 Max neighbors for any atom: 43 Min neighbors for any atom: 13 First 5 neighbors of atom 0: Neighbor 0: atom 5, shift [0, 0, 0] Neighbor 1: atom 85, shift [0, 0, 0] Neighbor 2: atom 1, shift [0, 0, 0] Neighbor 3: atom 79, shift [0, 0, 0] Neighbor 4: atom 14, shift [0, 0, 0] .. GENERATED FROM PYTHON SOURCE LINES 141-145 Unified API - COO Format ======================== The COO (coordinate) format is often preferred for graph neural networks. Set ``return_neighbor_list=True`` to get this format. .. GENERATED FROM PYTHON SOURCE LINES 145-176 .. code-block:: Python print("\n" + "=" * 70) print("UNIFIED API - COO FORMAT") print("=" * 70) # Get neighbor list in COO format neighbor_list_coo, neighbor_ptr, shifts_coo = neighbor_list( positions, cutoff, cell=cell, pbc=pbc, return_neighbor_list=True ) print("\nReturned COO format:") print(f" neighbor_list shape: {neighbor_list_coo.shape} (2 x num_pairs)") print(f" neighbor_ptr shape: {neighbor_ptr.shape} (CSR pointers)") print(f" shifts shape: {shifts_coo.shape}") source_atoms = neighbor_list_coo[0] target_atoms = neighbor_list_coo[1] print("\nStatistics:") print(f" Total pairs: {neighbor_list_coo.shape[1]}") print(f" Source atoms range: [{int(source_atoms.min())}, {int(source_atoms.max())}]") print(f" Target atoms range: [{int(target_atoms.min())}, {int(target_atoms.max())}]") # Show first few pairs print("\nFirst 5 neighbor pairs:") for i in range(min(5, neighbor_list_coo.shape[1])): src = int(source_atoms[i]) tgt = int(target_atoms[i]) shift = shifts_coo[i].tolist() print(f" Pair {i}: atom {src} -> atom {tgt}, shift {shift}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== UNIFIED API - COO FORMAT ====================================================================== Returned COO format: neighbor_list shape: (2, 6166) (2 x num_pairs) neighbor_ptr shape: (201,) (CSR pointers) shifts shape: (6166, 3) Statistics: Total pairs: 6166 Source atoms range: [0, 199] Target atoms range: [0, 199] First 5 neighbor pairs: Pair 0: atom 0 -> atom 182, shift [0, 0, 0] Pair 1: atom 0 -> atom 197, shift [0, 0, 0] Pair 2: atom 0 -> atom 5, shift [0, 0, 0] Pair 3: atom 0 -> atom 85, shift [0, 0, 0] Pair 4: atom 0 -> atom 184, shift [0, 0, 0] .. GENERATED FROM PYTHON SOURCE LINES 177-185 Prepared Execution ================== Preparation performs route selection, validation, capacity sizing, and reusable allocation eagerly. The immutable JAX state is then threaded through the compiled function while the scientific result keeps the selected route's tuple. This is the convenient unified interface for compiled loops. A method-specific compiled function with explicit fixed buffers and launch metadata is the lower-overhead option when latency matters, as shown later. .. GENERATED FROM PYTHON SOURCE LINES 185-264 .. code-block:: Python print("\n" + "=" * 70) print("PREPARED JAX EXECUTION") print("=" * 70) prepared_state = prepare_neighbor_list( positions, cutoff, cell=cell, pbc=pbc, method="naive", strategy="scalar", max_neighbors=128, ) @jax.jit def compiled_prepared(current_positions, current_cell, current_state): """Execute a prepared route and return its successor state.""" return neighbor_list( current_positions, cell=current_cell, state=current_state, ) (prepared_matrix, prepared_counts, prepared_shifts), prepared_state = compiled_prepared( positions, cell, prepared_state ) check_neighbor_list_state(prepared_state) print(f"Prepared method: {prepared_state.method}") print(f"Compilation eligible: {prepared_state.supports_compilation}") print(f"Prepared matrix shape: {prepared_matrix.shape}") print(f"Prepared pair count: {int(prepared_counts.sum())}") print(f"Prepared shift shape: {prepared_shifts.shape}") # Automatic preparation may prefer an equivalent compilation-eligible route. # Read these properties after preparation to make that decision visible. Pass # explicit method and strategy values, as shown above, when implementation identity # matters; an unsupported explicit combination raises instead of changing it. print(f"Resolved prepared strategy: {prepared_state.strategy}") # Selective states require one Boolean rebuild flag per system. Initialize a # system before asking a later call to preserve it. The false flag below keeps # the previously initialized result while still returning a successor state. selective_state = prepare_neighbor_list( positions, cutoff, cell=cell, pbc=pbc, method="naive", max_neighbors=128, selective=True, ) @jax.jit def compiled_selective(current_positions, current_state, rebuild_flags): """Execute a selective prepared route and return its successor state.""" return neighbor_list( current_positions, state=current_state, rebuild_flags=rebuild_flags, ) selective_results, selective_state = compiled_selective( positions, selective_state, jnp.array([True]) ) check_neighbor_list_state(selective_state) preserved_results, selective_state = compiled_selective( positions, selective_state, jnp.array([False]) ) check_neighbor_list_state(selective_state) print( "Selective preservation kept counts: " f"{bool(jnp.array_equal(selective_results[1], preserved_results[1]))}" ) .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== PREPARED JAX EXECUTION ====================================================================== Prepared method: naive Compilation eligible: True Prepared matrix shape: (200, 128) Prepared pair count: 6166 Prepared shift shape: (200, 128, 3) Resolved prepared strategy: scalar Selective preservation kept counts: True .. GENERATED FROM PYTHON SOURCE LINES 265-272 Algorithm Comparison ==================== The nvalchemiops library provides direct access to two main algorithms: - ``naive_neighbor_list``: O(N²) all-pairs distance checks - ``cell_list``: spatial decomposition for larger systems Both should produce identical results. .. GENERATED FROM PYTHON SOURCE LINES 272-297 .. code-block:: Python print("\n" + "=" * 70) print("ALGORITHM COMPARISON") print("=" * 70) # Direct call to naive algorithm nm_naive, num_naive, shifts_naive = naive_neighbor_list( positions, cutoff, cell=cell, pbc=pbc ) # Direct call to cell list algorithm nm_cell, num_cell, shifts_cell = cell_list(positions, cutoff, cell=cell, pbc=pbc) print("\nNaive algorithm (O(N²)):") print(f" Total pairs: {int(num_naive.sum())}") print(f" Average neighbors: {float(num_naive.mean()):.2f}") print("\nCell list algorithm (O(N)):") print(f" Total pairs: {int(num_cell.sum())}") print(f" Average neighbors: {float(num_cell.mean()):.2f}") # Verify they find the same number of pairs per atom pairs_match = jnp.allclose(num_naive, num_cell) print(f"\nResults match: {pairs_match}") # .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== ALGORITHM COMPARISON ====================================================================== Naive algorithm (O(N²)): Total pairs: 6166 Average neighbors: 30.83 Cell list algorithm (O(N)): Total pairs: 6166 Average neighbors: 30.83 Results match: True .. GENERATED FROM PYTHON SOURCE LINES 298-301 Distance Validation =================== Let's verify that all neighbor pairs are actually within the cutoff distance. .. GENERATED FROM PYTHON SOURCE LINES 301-351 .. code-block:: Python print("\n" + "=" * 70) print("DISTANCE VALIDATION") print("=" * 70) # Get neighbor list in COO format for easy distance computation nlist, nptr, nshifts = naive_neighbor_list( positions, cutoff, cell=cell, pbc=pbc, return_neighbor_list=True ) if nlist.shape[1] > 0: # Extract source and target positions src_idx = nlist[0] tgt_idx = nlist[1] pos_src = positions[src_idx] pos_tgt = positions[tgt_idx] # Compute Cartesian shift from lattice shift # shifts are in lattice coordinates, multiply by cell vectors cell_squeezed = cell.squeeze(0) # (3, 3) cartesian_shifts = jnp.einsum( "ij,jk->ik", nshifts.astype(jnp.float32), cell_squeezed ) # Compute distances: r_j - r_i + shift diff = pos_tgt - pos_src + cartesian_shifts distances = jnp.linalg.norm(diff, axis=1) print(f"\nComputed distances for {len(distances)} neighbor pairs:") print(f" Min distance: {float(distances.min()):.4f} Å") print(f" Max distance: {float(distances.max()):.4f} Å") print(f" Mean distance: {float(distances.mean()):.4f} Å") print(f" Cutoff: {cutoff} Å") # Check if all distances are within cutoff (with small tolerance) within_cutoff = jnp.all(distances <= cutoff + 1e-5) print(f"\n All distances within cutoff: {within_cutoff}") # Show distribution of first 10 distances print("\nFirst 10 neighbor distances:") for i in range(min(10, len(distances))): src = int(src_idx[i]) tgt = int(tgt_idx[i]) dist = float(distances[i]) print(f" Atom {src} -> {tgt}: {dist:.4f} Å") else: print("\nNo neighbor pairs found (empty system or cutoff too small)") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== DISTANCE VALIDATION ====================================================================== Computed distances for 6166 neighbor pairs: Min distance: 0.4676 Å Max distance: 4.9992 Å Mean distance: 3.7425 Å Cutoff: 5.0 Å All distances within cutoff: True First 10 neighbor distances: Atom 0 -> 75: 4.7814 Å Atom 0 -> 79: 2.8016 Å Atom 0 -> 85: 2.4960 Å Atom 0 -> 88: 4.4487 Å Atom 0 -> 93: 4.6019 Å Atom 0 -> 145: 2.7335 Å Atom 0 -> 148: 3.8481 Å Atom 0 -> 151: 3.7375 Å Atom 0 -> 155: 4.6850 Å Atom 0 -> 156: 4.9675 Å .. GENERATED FROM PYTHON SOURCE LINES 352-355 JIT compilation =============== Demonstrate the fixed-capacity method-specific JIT boundary. .. GENERATED FROM PYTHON SOURCE LINES 355-441 .. code-block:: Python print("\n" + "=" * 70) print("JIT compilation example") print("=" * 70) max_neighbors = 128 # This specialization closes over cell, pbc, cutoff, and matching shift metadata. # Recompute the metadata and create another specialization if any of them changes. shift_range, num_shifts, max_shifts = compute_naive_num_shifts(cell, cutoff, pbc) neighbor_matrix = jnp.full((num_atoms, max_neighbors), num_atoms, dtype=jnp.int32) num_neighbors = jnp.zeros((num_atoms,), dtype=jnp.int32) neighbor_shift_matrix = jnp.zeros((num_atoms, max_neighbors, 3), dtype=jnp.int32) @jax.jit def compiled_naive_neighbors( positions, neighbor_matrix, num_neighbors, neighbor_shift_matrix, ): """Build a fixed-capacity periodic neighbor matrix inside ``jax.jit``.""" return naive_neighbor_list( positions, cutoff, cell=cell, pbc=pbc, max_neighbors=max_neighbors, neighbor_matrix=neighbor_matrix, num_neighbors=num_neighbors, neighbor_matrix_shifts=neighbor_shift_matrix, shift_range_per_dimension=shift_range, num_shifts_per_system=num_shifts, max_shifts_per_system=max_shifts, ) neighbor_matrix, num_neighbors, neighbor_shift_matrix = compiled_naive_neighbors( positions, neighbor_matrix, num_neighbors, neighbor_shift_matrix, ) print(f"Returned neighbor matrix shape: {neighbor_matrix.shape}") if int(jnp.max(num_neighbors)) > max_neighbors: raise RuntimeError("neighbor matrix capacity overflow; grow it outside jax.jit") coo_capacity = num_atoms * max_neighbors @jax.jit def compiled_naive_coo(positions): """Build padded COO arrays with device-side recovery metadata.""" return naive_neighbor_list( positions, cutoff, cell=cell, pbc=pbc, max_neighbors=max_neighbors, return_neighbor_list=True, coo_capacity=coo_capacity, shift_range_per_dimension=shift_range, num_shifts_per_system=num_shifts, max_shifts_per_system=max_shifts, ) fixed_coo, fixed_ptr, fixed_shifts, required_counts, metadata_valid = ( compiled_naive_coo(positions) ) if not bool(metadata_valid): raise RuntimeError("refresh launch metadata outside jax.jit") stored_counts = fixed_ptr[1:] - fixed_ptr[:-1] if bool(jnp.any(stored_counts != required_counts)): required_width = int(jnp.max(required_counts, initial=0)) required_capacity = int(jnp.sum(required_counts)) raise RuntimeError( f"COO output is incomplete; retry with max_neighbors >= {required_width} " f"and coo_capacity >= {required_capacity}" ) num_pairs = int(fixed_ptr[-1]) print(f"Returned fixed COO shape: {fixed_coo.shape}") print(f"Valid COO prefix: {num_pairs} pairs") print(f"Fixed shift shape: {fixed_shifts.shape}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== JIT compilation example ====================================================================== Returned neighbor matrix shape: (200, 128) Returned fixed COO shape: (2, 25600) Valid COO prefix: 6166 pairs Fixed shift shape: (25600, 3) .. GENERATED FROM PYTHON SOURCE LINES 442-450 Cost-model dispatch =================== ``estimate_neighbor_list_costs`` and ``suggest_neighbor_list_method`` expose the geometry cost model that ``neighbor_list(method=None)`` uses internally. Call them once on per-system geometry (``batch_ptr``, ``cell``, ``pbc``, ``cutoff``) and pass the returned name as an explicit ``method=`` so repeated builds skip the auto-dispatch host read. They synchronize on the host (a small selector kernel runs on the device and its result is read back), so call them outside ``jax.jit``. .. GENERATED FROM PYTHON SOURCE LINES 450-472 .. code-block:: Python print("\n" + "=" * 70) print("COST-MODEL DISPATCH") print("=" * 70) batch_ptr = jnp.array([0, num_atoms], dtype=jnp.int32) cost_report = estimate_neighbor_list_costs(batch_ptr, cell, pbc, cutoff) print("\nFeasible methods (cheapest first):") for method_name, cost in cost_report: print(f" {method_name:24s} estimated cost (arbitrary units): {cost:.3g}") suggested_method = suggest_neighbor_list_method(batch_ptr, cell, pbc, cutoff) print(f"\nSuggested method: {suggested_method}") # Reuse the suggestion as an explicit ``method=`` on the eager entry point. # For JIT, select the matching direct public method and close its strategy over # a fixed-capacity wrapper like ``compiled_naive_neighbors`` above. nm_suggested, num_suggested, _ = neighbor_list( positions, cutoff, cell=cell, pbc=pbc, method=suggested_method ) print(f"Total pairs via suggested method: {int(num_suggested.sum())}") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== COST-MODEL DISPATCH ====================================================================== Feasible methods (cheapest first): naive_tile estimated cost (arbitrary units): 3.5e+05 naive_scalar estimated cost (arbitrary units): 3.5e+05 cell_list_atom_centric estimated cost (arbitrary units): 1.2e+06 cell_list_pair_centric estimated cost (arbitrary units): 1.73e+06 Suggested method: naive_tile Total pairs via suggested method: 6166 .. GENERATED FROM PYTHON SOURCE LINES 473-478 Partial neighbor lists with ``target_indices`` ============================================== ``target_indices`` builds neighbors only for selected central atoms. Matrix rows are compact: row ``r`` maps to atom ``target_indices[r]``. In COO mode, ``neighbor_list[0]`` also stores compact row ids, not original atom ids. .. GENERATED FROM PYTHON SOURCE LINES 478-513 .. code-block:: Python print("\n" + "=" * 70) print("PARTIAL NEIGHBOR LISTS (target_indices)") print("=" * 70) target_indices = jnp.array([0, 4, 9], dtype=jnp.int32) partial_nm, partial_counts, _ = neighbor_list( positions, cutoff, cell=cell, pbc=pbc, method="cell_list_atom_centric", max_neighbors=128, target_indices=target_indices, ) print(f"\nSelected atoms: {target_indices.tolist()}") print(f"Partial matrix shape: {partial_nm.shape}") for row, atom in enumerate(target_indices.tolist()): print(f" compact row {row} -> atom {atom}: {int(partial_counts[row])} neighbors") partial_coo, partial_ptr, _ = neighbor_list( positions, cutoff, cell=cell, pbc=pbc, method="cell_list_atom_centric", max_neighbors=128, target_indices=target_indices, return_neighbor_list=True, ) num_partial_pairs = int(partial_ptr[-1]) compact_rows = jnp.unique(partial_coo[0, :num_partial_pairs]) print(f"COO compact source rows present: {compact_rows.tolist()}") print(f"COO pointer length: {partial_ptr.shape[0]} (num_targets + 1)") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== PARTIAL NEIGHBOR LISTS (target_indices) ====================================================================== Selected atoms: [0, 4, 9] Partial matrix shape: (3, 128) compact row 0 -> atom 0: 33 neighbors compact row 1 -> atom 4: 32 neighbors compact row 2 -> atom 9: 37 neighbors COO compact source rows present: [0, 1, 2] COO pointer length: 4 (num_targets + 1) .. GENERATED FROM PYTHON SOURCE LINES 514-523 Inline pair potentials with ``pair_fn`` ======================================= A Warp ``pair_fn`` evaluates a pairwise potential as neighbors are enumerated, returning per-pair energy and force in the same pass (no second loop over the list). On JAX, ``pair_fn`` is available through the direct algorithm bindings and compatible unified-dispatch paths. The ``pair_energies`` / ``pair_forces`` buffers are auto-allocated, appended to the return tuple, and forward-only (use ``return_distances`` / ``return_vectors`` for differentiable geometry). .. GENERATED FROM PYTHON SOURCE LINES 523-575 .. code-block:: Python print("\n" + "=" * 70) print("INLINE PAIR POTENTIALS (pair_fn)") print("=" * 70) @wp.func def lj_pair_fn( r_ij: wp.vec3f, distance: wp.float32, pair_params: wp.array2d(dtype=wp.float32), i: int, j: int, ): """Lennard-Jones per-pair energy and force from per-atom (epsilon, sigma).""" 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 # Per-atom (epsilon, sigma) table, shape (num_atoms, 2), float32. pair_params = jnp.stack( [ jnp.full((num_atoms,), 0.0104, dtype=jnp.float32), # epsilon jnp.full((num_atoms,), 3.40, dtype=jnp.float32), # sigma ], axis=1, ) # pair_fn returns auto-allocated energy/force buffers appended after the matrix # outputs: (neighbor_matrix, num_neighbors, shifts, pair_energies, pair_forces). nm_pair, num_pair, _shifts_pair, pair_energies, pair_forces = naive_neighbor_list( positions, cutoff, cell=cell, pbc=pbc, max_neighbors=128, pair_fn=lj_pair_fn, pair_params=pair_params, ) print("\nPair-output buffers (matrix-aligned with the neighbor matrix):") print(f" pair_energies shape: {pair_energies.shape}") print(f" pair_forces shape: {pair_forces.shape}") print(" Auto-allocated, returned, and forward-only.") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== INLINE PAIR POTENTIALS (pair_fn) ====================================================================== Pair-output buffers (matrix-aligned with the neighbor matrix): pair_energies shape: (200, 128) pair_forces shape: (200, 128, 3) Auto-allocated, returned, and forward-only. .. GENERATED FROM PYTHON SOURCE LINES 576-591 Summary ======= This example demonstrated the JAX neighbor list API in nvalchemiops: - **Unified API**: ``neighbor_list()`` provides a single entry point - **Prepared execution**: Conveniently thread reusable state through ``jax.jit`` - **Matrix format**: Dense (N, max_neighbors) format for neighbor indices - **COO format**: Sparse (2, num_pairs) format for graph neural networks - **Algorithm choice**: Direct naive and cell-list calls for comparison - **Half-fill mode**: Store only unique pairs to save memory - **Distance validation**: Verify all pairs are within cutoff - **Cost-model dispatch**: ``estimate``/``suggest`` helpers pick a method - **Partial lists**: Compact ``target_indices`` rows for selected atoms - **Inline pair_fn**: Per-pair energy/force during enumeration - **Compiled COO**: Padded arrays plus raw counts and metadata validity .. GENERATED FROM PYTHON SOURCE LINES 591-609 .. code-block:: Python print("\n" + "=" * 70) print("SUMMARY") print("=" * 70) print("\nKey takeaways:") print(" - Use neighbor_list() for eager dispatch and capacity management") print(" - Use prepared state for a convenient unified compiled interface") print(" - Use method-specific compiled functions when lower call overhead matters") print(" - Inspect state.method and state.strategy after automatic preparation") print(" - Use return_neighbor_list=True for COO format (GNNs)") print(" - Add coo_capacity for fixed-shape COO inside jax.jit") print(" - Use half_fill=True to store only unique pairs") print(" - naive_neighbor_list performs O(N²) all-pairs checks") print(" - cell_list uses spatial decomposition") print(" - suggest_neighbor_list_method picks a method from geometry") print(" - target_indices builds compact partial neighbor lists") print(" - pair_fn evaluates a pairwise potential during enumeration") print("\nExample completed successfully!") .. rst-class:: sphx-glr-script-out .. code-block:: none ====================================================================== SUMMARY ====================================================================== Key takeaways: - Use neighbor_list() for eager dispatch and capacity management - Use prepared state for a convenient unified compiled interface - Use method-specific compiled functions when lower call overhead matters - Inspect state.method and state.strategy after automatic preparation - Use return_neighbor_list=True for COO format (GNNs) - Add coo_capacity for fixed-shape COO inside jax.jit - Use half_fill=True to store only unique pairs - naive_neighbor_list performs O(N²) all-pairs checks - cell_list uses spatial decomposition - suggest_neighbor_list_method picks a method from geometry - target_indices builds compact partial neighbor lists - pair_fn evaluates a pairwise potential during enumeration Example completed successfully! .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 37.315 seconds) .. _sphx_glr_download_examples_neighbors_05_jax_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: 05_jax_neighbor_list.ipynb <05_jax_neighbor_list.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 05_jax_neighbor_list.py <05_jax_neighbor_list.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 05_jax_neighbor_list.zip <05_jax_neighbor_list.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_