:orphan:

Examples Gallery
================

This gallery contains examples demonstrating the capabilities of ALCHEMI Toolkit-Ops.
Each example is a fully executable Python script that showcases different features
of the library, from basic neighbor list construction to advanced dispersion corrections.

All examples are designed to run on both CPU and GPU, automatically detecting available
hardware. The code can be downloaded as Python scripts or Jupyter notebooks.



.. raw:: html

    <div class="sphx-glr-thumbnails">

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates B-spline charge spreading and gathering operations used in Particle Mesh Ewald (PME) and related methods. B-splines provide smooth interpolation between particle positions and mesh grid points.">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_01_bspline_visualization_thumb.png
    :alt:

  :doc:`/examples/01_bspline_visualization`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">B-Spline Interpolation for Particle Mesh Methods</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Segment operations are the scatter/gather primitives behind graph pooling and message passing. Each one reduces or broadcasts a per-element array over variable-length segments defined by a sorted index — exactly the operation you need to pool per-atom quantities up to per-molecule ones (and to push per-molecule quantities back down to atoms).">

.. only:: html

  .. image:: /examples/images/thumb/sphx_glr_02_segment_ops_autograd_thumb.png
    :alt:

  :doc:`/examples/02_segment_ops_autograd`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Segment Ops: Pooling, Autograd, and torch.compile</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>


.. toctree::
   :hidden:

   /examples/01_bspline_visualization
   /examples/02_segment_ops_autograd

Dispersion Interactions
=======================

Examples demonstrating GPU-accelerated computation of DFT-D3
dispersion corrections with environment-dependent C6 coefficients.

These examples show how to:

* Compute dispersion energies and forces for molecules
* Process batches of crystal structures
* Integrate with PyTorch for differentiable workflows



.. raw:: html

    <div class="sphx-glr-thumbnails">

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compute the DFT-D3 dispersion energy and forces for a single molecular system using GPU-accelerated kernels.">

.. only:: html

  .. image:: /examples/dispersion/images/thumb/sphx_glr_01_dftd3_molecule_thumb.png
    :alt:

  :doc:`/examples/dispersion/01_dftd3_molecule`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">DFT-D3 Dispersion Correction for a Molecule</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates GPU-accelerated DFT-D3 calculations for batches of periodic crystal structures. Batch processing is essential for high-throughput computational materials screening and allows efficient parallelization across multiple structures simultaneously.">

.. only:: html

  .. image:: /examples/dispersion/images/thumb/sphx_glr_02_dftd3_batch_crystals_thumb.png
    :alt:

  :doc:`/examples/dispersion/02_dftd3_batch_crystals`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">DFT-D3 Dispersion Correction for Batched Crystals</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compute the DFT-D3 dispersion energy and forces for a single molecular system using the JAX API with GPU-accelerated Warp kernels.">

.. only:: html

  .. image:: /examples/dispersion/images/thumb/sphx_glr_03_jax_dftd3_molecule_thumb.png
    :alt:

  :doc:`/examples/dispersion/03_jax_dftd3_molecule`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">JAX DFT-D3 Dispersion Correction for a Molecule</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>

Dynamics Examples
=================

This directory contains examples demonstrating the molecular dynamics integrators
and geometry optimization algorithms in nvalchemiops.

All examples use realistic Lennard-Jones (LJ) argon systems to demonstrate
practical workflows with neighbor list management and periodic boundaries.

MD Integration Examples
-----------------------

01_langevin_integration.py
    Langevin (BAOAB) dynamics in the NVT ensemble.

    - Simulates liquid argon near the triple point (94.4 K)
    - BAOAB splitting for configurational sampling
    - Temperature monitoring and equilibration

02_velocity_verlet_integration.py
    Velocity Verlet dynamics in the NVE ensemble.

    - Symplectic, time-reversible integrator
    - Energy conservation diagnostics
    - Stability considerations for LJ systems

03_nph_integration.py
    NPH dynamics with MTK barostat.

    - Constant enthalpy and pressure ensemble
    - Martyna-Tobias-Klein extended-system barostat
    - Pressure coupling mechanics

04_npt_integration.py
    NPT dynamics with MTK barostat + Nosé-Hoover chain thermostat.

    - Isothermal-isobaric ensemble
    - Combined temperature and pressure control
    - Volume fluctuations under constant pressure

05_batched_langevin_dynamics.py
    Batched Langevin dynamics for multiple independent systems.

    - Multiple systems packed into single arrays
    - Per-system temperature targets
    - Improved GPU utilization for small systems

Geometry Optimization Examples
------------------------------

06_fire_optimization.py
    FIRE geometry optimization for a single LJ cluster.

    - Fast Inertial Relaxation Engine (FIRE) algorithm
    - Adaptive timestep and velocity mixing
    - Convergence monitoring (energy, max force)

07_fire_variable_cell.py
    Variable-cell FIRE optimization for joint atom + cell relaxation.

    - Cell alignment to upper-triangular form
    - Virial → stress → cell force conversion
    - Pack/unpack utilities for extended DOF arrays
    - External pressure equilibration

08_fire_batched.py
    Batched FIRE optimization with two indexing strategies:

    - **batch_idx mode**: Each atom tagged with system index
    - **atom_ptr mode (CSR)**: Atom ranges via pointers
    - Per-system FIRE parameters adapt independently
    - Uses ``nvalchemiops.batch_utils`` for reductions

09_fire2_optimization.py
    FIRE2 geometry optimization for a single LJ cluster.

    - Improved FIRE variant (Guenole et al., 2020)
    - Adaptive damping and coupled step/dt scaling
    - Requires ``batch_idx`` even for single-system mode
    - Hyperparameters are Python scalars (no per-system arrays)

10_fire2_batched.py
    Batched FIRE2 optimization with ``batch_idx`` batching.

    - Multiple independent LJ clusters optimized in parallel
    - Per-system FIRE2 parameters adapt independently
    - Only supports ``batch_idx`` mode (no ``atom_ptr`` variant)

11_fire2_variable_cell.py
    Variable-cell FIRE2 optimization for joint atom + cell relaxation.

    - Uses ``fire2_step_coord_cell`` for coupled coordinate/cell updates
    - Preserves fractional coordinates during cell-only motion
    - Simpler state than FIRE (no masses, fewer per-system arrays)

Key Concepts
------------

State Management
^^^^^^^^^^^^^^^^

All algorithms use ``@wp.struct`` containers for clean organization:

- ``MDState``: positions, velocities, forces, masses
- ``FIREState``: MDState + adaptive parameters (dt, alpha, n_positive)
- Integration parameters: ``VerletParams``, ``LangevinParams``, ``FIREParams``

Structs are dtype-agnostic: use ``wp.vec3f``/``wp.float32`` for single
precision or ``wp.vec3d``/``wp.float64`` for double precision.

Two-Pass Integrators
^^^^^^^^^^^^^^^^^^^^

MD integrators use a two-pass design for flexibility::

    # Pass 1: Update positions and half-step velocities
    velocity_verlet_position_update(state, params)

    # User computes forces at new positions
    compute_forces(state.positions, state.forces)

    # Pass 2: Complete velocity update
    velocity_verlet_velocity_finalize(state, state.forces, params)

This allows any force calculation method to be used between the passes.

Batching Modes
^^^^^^^^^^^^^^

For processing multiple systems efficiently:

- **batch_idx**: Array mapping each atom to its system index. Flexible for
  heterogeneous systems but requires atomic accumulation.

- **atom_ptr (CSR)**: Pointer array where system *s* owns atoms
  ``[atom_ptr[s], atom_ptr[s+1])``. More efficient for homogeneous batches.

Use ``nvalchemiops.batch_utils`` for conversions and per-system reductions.

Variable-Cell Optimization
^^^^^^^^^^^^^^^^^^^^^^^^^^

The ``cell_filter`` utilities enable joint optimization of positions and cell
with FIRE. FIRE2 variable-cell optimization should use the coupled
``fire2_step_coord_cell`` adapter, which applies the affine coordinate remap
implied by cell motion.

Generic cell-filter workflow for optimizers such as FIRE::

    # 1. Align cell to upper-triangular form
    positions, cell = align_cell(positions, cell)

    # 2. Pack atomic + cell DOFs into extended arrays
    extended_pos = pack_positions_with_cell(positions, cell, ...)
    extended_forces = pack_forces_with_cell(forces, cell_force, ...)

    # 3. Run optimizer on extended arrays
    fire_step(extended_pos, extended_vel, extended_forces, ...)

    # 4. Unpack results
    positions, cell = unpack_positions_with_cell(extended_pos, ...)

Running the Examples
--------------------

Make sure you have warp installed::

    pip install warp-lang

Run from the repository root::

    python examples/dynamics/01_langevin_integration.py
    python examples/dynamics/06_fire_optimization.py
    python examples/dynamics/08_fire_batched.py

Or run interactively in VS Code or Jupyter by executing cell-by-cell
(sections marked with ``# %%``).



.. raw:: html

    <div class="sphx-glr-thumbnails">

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates GPU-accelerated molecular dynamics using the Langevin (BAOAB) thermostat with a Lennard-Jones potential.">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_01_langevin_integration_thumb.png
    :alt:

  :doc:`/examples/dynamics/01_langevin_integration`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Langevin Dynamics with Lennard-Jones Potential</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates GPU-accelerated molecular dynamics using the Velocity Verlet integrator with a Lennard-Jones (LJ) potential.">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_02_velocity_verlet_integration_thumb.png
    :alt:

  :doc:`/examples/dynamics/02_velocity_verlet_integration`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Velocity Verlet Dynamics with Lennard-Jones Potential</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates isothermal-free (no thermostat) NPH dynamics using the Martyna-Tobias-Klein (MTK) barostat and the Lennard-Jones (LJ) potential.">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_03_nph_integration_thumb.png
    :alt:

  :doc:`/examples/dynamics/03_nph_integration`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">NPH Dynamics (MTK) with Lennard-Jones Potential</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates isothermal-isobaric (NPT) dynamics using the Martyna-Tobias-Klein (MTK) barostat coupled with a Nosé-Hoover chain (NHC) thermostat, driven by the Lennard-Jones (LJ) potential.">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_04_npt_integration_thumb.png
    :alt:

  :doc:`/examples/dynamics/04_npt_integration`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">NPT Dynamics (MTK + Nosé-Hoover Chain) with Lennard-Jones Potential</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates batched molecular dynamics: multiple independent systems are packed into a single set of arrays, and we integrate all systems in one go on the GPU.">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_05_batched_langevin_dynamics_thumb.png
    :alt:

  :doc:`/examples/dynamics/05_batched_langevin_dynamics`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Batched Langevin Dynamics (BAOAB) with Lennard-Jones Potential</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates geometry optimization with the FIRE optimizer using: - The package LJ implementation (neighbor-list accelerated) - The package FIRE kernels (nvalchemiops.dynamics.optimizers) - The shared example utilities in examples.dynamics._dynamics_utils">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_06_fire_optimization_thumb.png
    :alt:

  :doc:`/examples/dynamics/06_fire_optimization`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">FIRE Geometry Optimization (LJ Cluster)</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates joint optimization of atomic positions and cell parameters using the FIRE optimizer with the cell filter utilities on a realistic LJ argon crystal.">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_07_fire_variable_cell_thumb.png
    :alt:

  :doc:`/examples/dynamics/07_fire_variable_cell`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Variable-Cell FIRE Optimization with LJ Potential</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates batched geometry optimization using the FIRE optimizer with two different batching strategies:">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_08_fire_batched_thumb.png
    :alt:

  :doc:`/examples/dynamics/08_fire_batched`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Batched FIRE Optimization with LJ Clusters</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates geometry optimization with the FIRE2 optimizer (Guenole et al., 2020) using:">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_09_fire2_optimization_thumb.png
    :alt:

  :doc:`/examples/dynamics/09_fire2_optimization`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">FIRE2 Geometry Optimization (LJ Cluster)</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates batched geometry optimization using the FIRE2 optimizer (Guenole et al., 2020) with batch_idx batching.">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_10_fire2_batched_thumb.png
    :alt:

  :doc:`/examples/dynamics/10_fire2_batched`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Batched FIRE2 Optimization with LJ Clusters</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates joint optimization of atomic positions and cell parameters using the coupled FIRE2 PyTorch adapter (Guenole et al., 2020).">

.. only:: html

  .. image:: /examples/dynamics/images/thumb/sphx_glr_11_fire2_variable_cell_thumb.png
    :alt:

  :doc:`/examples/dynamics/11_fire2_variable_cell`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Variable-Cell FIRE2 Optimization with LJ Potential</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>

Electrostatics
==============

Examples demonstrating GPU-accelerated computation of long-range electrostatic
interactions in periodic systems using Coulomb, Ewald summation, and Particle
Mesh Ewald (PME).

These examples show how to:

* Compute direct Coulomb interactions (damped and undamped)
* Use Ewald summation for periodic systems with automatic parameter estimation
* Apply two-dimensional slab corrections for Ewald and PME interfacial systems
* Apply Particle Mesh Ewald (PME) for O(N log N) scaling
* Work with neighbor list and neighbor matrix formats
* Perform batch evaluation for multiple systems
* Leverage autograd for computing forces and gradients
* Compute multipole Ewald and PME totals with charges, dipoles, and
  quadrupoles (l_max = 0 / 1 / 2), including stress tensors via the
  cell gradient and force-loss-style training via the second-order
  backward Warp kernel
* Extract atom-centered multipole features by projecting the periodic
  potential onto receiver GTOs
* Amortize the position-independent reciprocal-space state with the
  multipole SCF cache and reuse it across many step evaluations
* Train on forces, stress, and charge gradients via the energy-derivative
  contract (the recommended replacement for the deprecated direct-output flags)

The full Torch Ewald/PME APIs support first- and second-order energy-derived
training workflows. The full JAX Ewald/PME APIs support first-order
energy-derived gradients for positions, charges, and row-vector displacement
virials. Higher-order JAX support is limited to tested position and charge
scalar losses; PME cell/stress/strain higher-order derivatives are unsupported.
Electrostatics does not expose public Hessian or Jacobian APIs.

Point-charge Ewald/PME examples use ``float64`` for accuracy-sensitive
reciprocal-space calculations and gradient checks. The APIs also support
``float32`` when throughput is the priority; keep all floating inputs and
precomputed metadata in a call on a consistent dtype.



.. raw:: html

    <div class="sphx-glr-thumbnails">

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compute Coulomb electrostatic interactions using GPU-accelerated kernels. The implementation supports both direct (undamped) Coulomb interactions and damped interactions used as the real-space component in Ewald/PME methods.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_01_coulomb_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/01_coulomb_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Coulomb Electrostatic Interactions</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compute long-range electrostatic interactions in periodic systems using the Ewald summation method. Ewald splits the slowly converging Coulomb sum into rapidly converging real-space and reciprocal-space components.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_02_ewald_summation_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/02_ewald_summation_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Ewald Summation for Long-Range Electrostatics</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compute long-range electrostatic interactions using the Particle Mesh Ewald (PME) method. PME achieves O(N log N) scaling by using FFT-based mesh interpolation for the reciprocal-space contribution.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_03_pme_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/03_pme_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Particle Mesh Ewald (PME) for Long-Range Electrostatics</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compute long-range electrostatic interactions using the Particle Mesh Ewald (PME) method with the JAX backend. PME achieves O(N log N) scaling through FFT-based mesh interpolation.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_04_jax_pme_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/04_jax_pme_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Particle Mesh Ewald (PME) with JAX</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to apply the Yeh-Berkowitz / Ballenegger two-dimensional slab correction to Ewald and Particle Mesh Ewald (PME) electrostatics. The correction is used for slab-like systems with two periodic directions and one non-periodic direction, such as interfaces with vacuum padding.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_05_slab_correction_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/05_slab_correction_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">2D Slab Correction for Ewald and PME</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to apply the Yeh-Berkowitz / Ballenegger two-dimensional slab correction to JAX Ewald and Particle Mesh Ewald (PME) electrostatics. The correction is used for slab-like systems with two periodic directions and one non-periodic direction, such as interfaces with vacuum padding.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_06_jax_slab_correction_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/06_jax_slab_correction_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">2D Slab Correction for Ewald and PME with JAX</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates the GTO-Ewald multipole electrostatics path for periodic systems with charges (l_max=0), dipoles (l_max=1), and quadrupoles (l_max=2). The Ewald path uses a direct k-space sum for the reciprocal piece.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_07_multipole_ewald_summation_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/07_multipole_ewald_summation_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Multipole Ewald Summation (charges + dipoles + quadrupoles)</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates the GTO-PME multipole electrostatics path for periodic systems with charges (l_max=0), dipoles (l_max=1), and quadrupoles (l_max=2). PME accelerates the reciprocal-space piece via B-spline interpolation + FFT, dropping the asymptotic cost from O(N · N_k) (Ewald direct k-space) to O(M log M + N · p³).">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_08_multipole_pme_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/08_multipole_pme_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Multipole Particle Mesh Ewald (charges + dipoles + quadrupoles)</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates multipole_electrostatic_features — a reciprocal-space projection that turns the periodic electrostatic potential of a set of source multipoles into LODE/ACE-style per-atom feature vectors. Unlike the energy paths (Ewald / PME), this entry point needs no neighbor list: every atom&#x27;s descriptor is computed directly from the k-space potential, projected onto a multi-σ basis of receiver Gaussian-type orbitals (GTOs).">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_09_multipole_features_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/09_multipole_features_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Atom-Centered Multipole Electrostatic Features</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates the two-phase prepare_cache + scf_step pattern for the reciprocal-space (&quot;k-space&quot;) formulation of periodic multipole electrostatics. Rather than the real-space + reciprocal split of multipole_ewald_summation, this path evaluates the periodic Coulomb interaction directly in reciprocal space — the GTO smearing of each multipole makes the k-sum converge — so its whole geometry setup is cell-only and can be built once and reused.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_10_multipole_scf_cache_example_thumb.png
    :alt:

  :doc:`/examples/electrostatics/10_multipole_scf_cache_example`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Multipole SCF Cache + Step (amortized fixed-cell workflow)</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates the recommended way to obtain forces, virial/stress, and charge gradients from PME (and Ewald) for machine-learning interatomic potential (MLIP) training: energy is the only differentiable output, and all derivatives are taken from it with torch.autograd.grad.">

.. only:: html

  .. image:: /examples/electrostatics/images/thumb/sphx_glr_11_pme_energy_derivative_training_thumb.png
    :alt:

  :doc:`/examples/electrostatics/11_pme_energy_derivative_training`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Energy-Derivative Training Contract (Forces, Stress, Charge Gradients)</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>

Neighbor Lists
==============

Examples demonstrating efficient O(N) neighbor list construction using GPU-accelerated
cell list algorithms.

These examples show how to:

* Build neighbor lists for single and batched systems
* Use dense or sparse COO output formats
* Inspect Torch cost-model dispatch estimates
* Request compact source rows with ``target_indices``
* Return per-neighbor vectors and distances, including differentiable geometry
* Compute inline Lennard-Jones pair energies and forces with ``pair_fn``
* Detect when neighbor lists need rebuilding
* Optimize performance with ``torch.compile``
* Integrate with molecular dynamics workflows



.. raw:: html

    <div class="sphx-glr-thumbnails">

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to use the basic neighbor list functions in nvalchemiops on random systems. We&#x27;ll cover:">

.. only:: html

  .. image:: /examples/neighbors/images/thumb/sphx_glr_01_simple_neighbor_list_thumb.png
    :alt:

  :doc:`/examples/neighbors/01_simple_neighbor_list`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Simple Neighbor List Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to use the batch neighbor list functions in nvalchemiops with multiple molecular and crystalline systems. We&#x27;ll cover:">

.. only:: html

  .. image:: /examples/neighbors/images/thumb/sphx_glr_02_batch_neighbor_list_thumb.png
    :alt:

  :doc:`/examples/neighbors/02_batch_neighbor_list`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Batch Neighbor List Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="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&#x27;ll cover:">

.. only:: html

  .. image:: /examples/neighbors/images/thumb/sphx_glr_03_rebuild_neighborlist_detection_thumb.png
    :alt:

  :doc:`/examples/neighbors/03_rebuild_neighborlist_detection`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Neighbor List Rebuild Detection Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates the significant performance benefits of using torch.compile with nvalchemiops neighbor list routines in a realistic molecular dynamics simulation. We&#x27;ll cover:">

.. only:: html

  .. image:: /examples/neighbors/images/thumb/sphx_glr_04_neighbors_list_torch_compile_performance_thumb.png
    :alt:

  :doc:`/examples/neighbors/04_neighbors_list_torch_compile_performance`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Torch.compile Performance Benefits Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to use the JAX neighbor list API in nvalchemiops for computing neighbor lists in periodic systems.">

.. only:: html

  .. image:: /examples/neighbors/images/thumb/sphx_glr_05_jax_neighbor_list_thumb.png
    :alt:

  :doc:`/examples/neighbors/05_jax_neighbor_list`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">JAX Neighbor List Example</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates four neighbor-list features used together:">

.. only:: html

  .. image:: /examples/neighbors/images/thumb/sphx_glr_06_pair_outputs_lj_thumb.png
    :alt:

  :doc:`/examples/neighbors/06_pair_outputs_lj`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Targeted Lennard-Jones Pair Outputs</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>


.. toctree::
   :hidden:
   :includehidden:


   /examples/dispersion/index.rst
   /examples/dynamics/index.rst
   /examples/electrostatics/index.rst
   /examples/neighbors/index.rst


.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-gallery

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download all examples in Python source code: examples_python.zip </examples/examples_python.zip>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download all examples in Jupyter notebooks: examples_jupyter.zip </examples/examples_jupyter.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_