:orphan:

:orphan:

Examples
========

End-to-end examples demonstrating ALCHEMI Toolkit functionality.



.. raw:: html

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

.. thumbnail-parent-div-open

.. thumbnail-parent-div-close

.. raw:: html

    </div>

Basic Examples
==============

These examples introduce the core nvalchemi-toolkit workflow to users
familiar with tools like ASE.  Each script focuses on one feature set
and runs end-to-end on a single GPU in under 60 seconds.

**01 — Data Structures**: AtomicData and Batch API.
**02 — Geometry Optimization**: FIRE optimizer with NeighborListHook and ConvergenceHook.
**03 — ASE Integration**: Loading ASE structures, FreezeAtomsHook on a surface system.
**04 — NVE MD**: Microcanonical dynamics, WrapPeriodicHook, EnergyDriftMonitorHook.
**05 — NVT MD**: Langevin thermostat, thermalization, LoggingHook to CSV.



.. raw:: html

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

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example walks through the full API of AtomicData and Batch: construction, properties, indexing, mutation, device movement, and serialization.">

.. only:: html

  .. image:: /examples/basic/images/thumb/sphx_glr_01_data_structures_thumb.png
    :alt:

  :doc:`/examples/basic/01_data_structures`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">AtomicData and Batch: Graph-structured molecular data</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates geometry optimization using the FIRE algorithm on a batch of argon clusters described by the Lennard-Jones potential.">

.. only:: html

  .. image:: /examples/basic/images/thumb/sphx_glr_02_geometry_optimization_thumb.png
    :alt:

  :doc:`/examples/basic/02_geometry_optimization`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">FIRE Geometry Optimization with Lennard-Jones Argon</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example shows how to load real atomic structures from ASE into the nvalchemi-toolkit workflow.  ASE provides a rich library of molecular and crystal builders, file I/O, and visualization tools; nvalchemi-toolkit consumes the resulting ase.Atoms objects via AtomicData.from_atoms.">

.. only:: html

  .. image:: /examples/basic/images/thumb/sphx_glr_03_ase_integration_thumb.png
    :alt:

  :doc:`/examples/basic/03_ase_integration`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">ASE Integration: Real Molecular Structures in nvalchemi-toolkit</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="The NVE ensemble (constant Number of atoms, Volume, and Energy) is the most fundamental MD ensemble.  The total energy E = KE + PE is conserved by the equations of motion, so any drift in E is a measure of integration error.">

.. only:: html

  .. image:: /examples/basic/images/thumb/sphx_glr_04_nve_energy_conservation_thumb.png
    :alt:

  :doc:`/examples/basic/04_nve_energy_conservation`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Microcanonical (NVE) Molecular Dynamics</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="The NVT ensemble (constant Number of atoms, Volume, and Temperature) samples the canonical distribution.  A thermostat couples the system to a heat bath, injecting or removing kinetic energy to maintain the target temperature.">

.. only:: html

  .. image:: /examples/basic/images/thumb/sphx_glr_05_nvt_langevin_thumb.png
    :alt:

  :doc:`/examples/basic/05_nvt_langevin`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Canonical (NVT) Molecular Dynamics with Langevin Thermostat</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>

Intermediate Examples
=====================

These examples assume familiarity with the basic tier and introduce
the storage layer, performance monitoring, and more complex pipeline
patterns.

**01 — Multi-Stage Pipeline**: FusedStage composition, LoggingHook CSV output,
step-budget migration, fused hooks for global status monitoring.

**02 — Trajectory I/O**: Writing trajectories to Zarr, reading back with
DataLoader, round-trip validation.

**03 — NPT MD**: Pressure-controlled dynamics with the MTK barostat,
LJ stress computation, cell fluctuation monitoring.

**04 — Inflight Batching**: SizeAwareSampler, Mode 2 FusedStage run (batch=None),
system_id tracking, ConvergedSnapshotHook collecting results.

**05 — Safety and Monitoring**: NaNDetectorHook, MaxForceClampHook,
EnergyDriftMonitorHook, ProfilerHook — defensive MD patterns.



.. raw:: html

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

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example demonstrates how to compose multiple dynamics stages into a single GPU-resident pipeline using FusedStage. A typical MD workflow chains geometry relaxation (FIRE) into equilibration (NVT) and then production sampling (NVE).  Running all three stages on a shared batch avoids repeated CPU–GPU data transfers between stages.">

.. only:: html

  .. image:: /examples/intermediate/images/thumb/sphx_glr_01_multistage_pipeline_thumb.png
    :alt:

  :doc:`/examples/intermediate/01_multistage_pipeline`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Multi-Stage Dynamics Pipelines with FusedStage</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Zarr is a chunked, compressed array format designed for large scientific datasets.  nvalchemi uses it as the on-disk representation for trajectories: each dynamics snapshot is serialised into a CSR-style layout (pointer arrays + concatenated fields) so that random-access reads and incremental appends are both efficient.">

.. only:: html

  .. image:: /examples/intermediate/images/thumb/sphx_glr_02_trajectory_zarr_io_thumb.png
    :alt:

  :doc:`/examples/intermediate/02_trajectory_zarr_io`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Writing and Replaying Trajectories with Zarr</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="The NPT ensemble keeps temperature T and pressure P constant while letting the simulation cell fluctuate in response to the instantaneous virial stress.  This is the correct ensemble for studying phase transitions, equations of state, and crystal relaxations under applied pressure.">

.. only:: html

  .. image:: /examples/intermediate/images/thumb/sphx_glr_03_npt_variable_cell_thumb.png
    :alt:

  :doc:`/examples/intermediate/03_npt_variable_cell`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Pressure-Controlled (NPT) Molecular Dynamics</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="The problem: a production relaxation campaign may require running thousands of structures through FIRE → NVT stages.  Loading all of them into GPU memory at once is impossible; processing one at a time wastes GPU throughput.">

.. only:: html

  .. image:: /examples/intermediate/images/thumb/sphx_glr_04_inflight_batching_thumb.png
    :alt:

  :doc:`/examples/intermediate/04_inflight_batching`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Processing Large Datasets with Inflight Batching</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Numerical instabilities are a fact of life in machine-learning molecular dynamics.  ML potentials are trained on a finite region of configuration space; geometries outside that region can produce enormous forces, NaN gradients, or energy drift that silently corrupts a long trajectory.">

.. only:: html

  .. image:: /examples/intermediate/images/thumb/sphx_glr_05_safety_and_monitoring_thumb.png
    :alt:

  :doc:`/examples/intermediate/05_safety_and_monitoring`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Defensive MD: Safety Hooks and Performance Monitoring</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>

Advanced Examples
=================

These examples are for users who want to extend the nvalchemi-toolkit
framework.  They require understanding of the intermediate tier.

**01 — Biased Potential**: BiasedPotentialHook for harmonic COM restraints
and umbrella sampling patterns.

**02 — Custom Hook**: Implementing the Hook protocol with a full
radial distribution function accumulator.

**03 — Custom Convergence**: ConvergenceHook with multiple criteria and
custom_op for arbitrary convergence logic.

**04 — MACE NVT**: Using a real MACE MLIP for NVT dynamics; automatic
neighbor list wiring via ModelCard; LJ fallback for CI.

**05 — Custom Integrator**: Subclassing BaseDynamics to implement a
velocity-rescaling thermostat; the pre_update/post_update contract;
_init_state for stateful integrators.



.. raw:: html

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

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Standard MD explores configuration space according to the Boltzmann distribution, which concentrates sampling near free-energy minima.  Many physically important processes — protein folding, nucleation, diffusion across a barrier — are rare events on MD timescales.  Biased sampling adds an external potential that encourages the system to explore regions it would not visit spontaneously.">

.. only:: html

  .. image:: /examples/advanced/images/thumb/sphx_glr_01_biased_potential_thumb.png
    :alt:

  :doc:`/examples/advanced/01_biased_potential`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Biased Sampling with BiasedPotentialHook</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="The Hook protocol is nvalchemi-toolkit&#x27;s primary extension point.  Any object that has stage, frequency, and __call__(batch, dynamics) satisfies the protocol — no base class is required.  This duck-typing approach makes hooks easy to write and easy to test in isolation.">

.. only:: html

  .. image:: /examples/advanced/images/thumb/sphx_glr_02_custom_hook_thumb.png
    :alt:

  :doc:`/examples/advanced/02_custom_hook`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Writing a Custom Hook: Radial Distribution Function</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="ConvergenceHook detects when geometry optimisation or MD has reached a desired stopping condition and removes converged systems from the active batch so subsequent steps spend compute only on unconverged systems.">

.. only:: html

  .. image:: /examples/advanced/images/thumb/sphx_glr_03_custom_convergence_thumb.png
    :alt:

  :doc:`/examples/advanced/03_custom_convergence`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Multi-Criteria Convergence with Custom Operators</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="MACE (Message-passing Atomic Cluster Expansion) is an equivariant message-passing neural network potential with competitive accuracy on bulk and molecular systems within its training distribution.  The MACE-MP family of foundation models is pre-trained on the Materials Project and provides reasonable zero-shot predictions across a broad range of compositions, though accuracy should be validated for the specific system of interest before production use — particularly for chemistries or structures not represented in the Materials Project training set.">

.. only:: html

  .. image:: /examples/advanced/images/thumb/sphx_glr_04_mace_nvt_thumb.png
    :alt:

  :doc:`/examples/advanced/04_mace_nvt`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">NVT MD with MACE (with LJ Fallback)</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="BaseDynamics is the single extension point for writing new molecular dynamics integrators.  Subclasses implement two methods — pre_update(batch) and post_update(batch) — that bracket the model forward pass (force computation).">

.. only:: html

  .. image:: /examples/advanced/images/thumb/sphx_glr_05_custom_integrator_thumb.png
    :alt:

  :doc:`/examples/advanced/05_custom_integrator`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Building a Custom Integrator by Subclassing BaseDynamics</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Lennard-Jones and other short-range potentials truncate interactions at a cutoff radius, which works well for neutral systems with rapidly decaying forces.  Ionic systems — salts, molten metals, charged proteins — require long-range Coulomb interactions that decay as 1/r and cannot be safely truncated without severe artefacts.">

.. only:: html

  .. image:: /examples/advanced/images/thumb/sphx_glr_06_ewald_electrostatics_thumb.png
    :alt:

  :doc:`/examples/advanced/06_ewald_electrostatics`

.. raw:: html

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


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Real force fields for ionic systems combine multiple physical contributions:">

.. only:: html

  .. image:: /examples/advanced/images/thumb/sphx_glr_07_composable_model_composition_thumb.png
    :alt:

  :doc:`/examples/advanced/07_composable_model_composition`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Composable Model Composition (LJ + Ewald)</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>

Distributed Pipeline Examples
==============================

These examples demonstrate multi-GPU distributed simulation pipelines
using :class:`~nvalchemi.dynamics.DistributedPipeline`.  They require
multiple GPUs and must be launched with ``torchrun``.

.. warning::

   These examples are **not executed** during the Sphinx documentation
   build.  To run them, use ``torchrun`` as shown in each example.

Architecture Overview
---------------------

A :class:`~nvalchemi.dynamics.DistributedPipeline` maps GPU ranks to
dynamics stages.  Systems flow between stages via fixed-size NCCL
communication buffers:

.. graphviz::

   digraph distributed_pipeline {
       rankdir=LR;
       node [shape=box, style="rounded,filled", fillcolor="#e8f4fd",
             fontname="Helvetica", fontsize=11];
       edge [fontname="Helvetica", fontsize=10];

       rank0 [label="Rank 0: FIRE\n(upstream)"];
       rank1 [label="Rank 1: Langevin\n(downstream + sink)"];

       rank0 -> rank1 [label="NCCL"];
   }

Key concepts:

- **Upstream ranks** (``prior_rank=None``): hold a
  :class:`~nvalchemi.dynamics.SizeAwareSampler` and push graduated
  (converged) systems to the next rank.
- **Downstream ranks** (``next_rank=None``): receive systems from the
  prior rank and write results to a sink.
- **BufferConfig**: must be set to a fixed size on all ranks; NCCL
  requires identical message sizes every communication step.
- ``torchrun --nproc_per_node=N`` launches one process per GPU; each
  process runs only the stage assigned to its rank.

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

**01 — Parallel FIRE → Langevin** (4 GPUs required):

.. code-block:: bash

   torchrun --nproc_per_node=4 examples/distributed/01_distributed_pipeline.py

   # CPU/debug mode (set backend="gloo" in the script first):
   torchrun --nproc_per_node=4 --master_port=29500 examples/distributed/01_distributed_pipeline.py

**02 — Monitoring with LoggingHook, ProfilerHook, and ZarrData** (4 GPUs required):

.. code-block:: bash

   torchrun --nproc_per_node=4 examples/distributed/02_distributed_monitoring.py

After running example 02, per-rank CSV logs and Zarr trajectory stores are
written to the working directory.  Rank 0 also prints a collated summary.

Example Descriptions
--------------------

**01 — Distributed Pipeline**
   Two independent FIRE → NVTLangevin sub-pipelines running on 4 GPUs.
   Demonstrates DistributedPipeline wiring, BufferConfig, and HostMemory sinks.

**02 — Distributed Monitoring**
   Same topology as example 01, augmented with per-rank LoggingHook and
   ProfilerHook for observability, and ZarrData sinks for persistent
   trajectory storage.  Shows post-run log collation on rank 0.



.. raw:: html

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

.. thumbnail-parent-div-open

.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="This example orchestrates two independent FIRE → NVTLangevin pipelines running in parallel across 4 GPUs using DistributedPipeline.">

.. only:: html

  .. image:: /examples/distributed/images/thumb/sphx_glr_01_distributed_pipeline_thumb.png
    :alt:

  :doc:`/examples/distributed/01_distributed_pipeline`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Distributed Multi-GPU Pipeline: Parallel FIRE → Langevin</div>
    </div>


.. raw:: html

    <div class="sphx-glr-thumbcontainer" tooltip="Building on 01_distributed_pipeline, this example adds comprehensive observability to the FIRE → NVTLangevin topology:">

.. only:: html

  .. image:: /examples/distributed/images/thumb/sphx_glr_02_distributed_monitoring_thumb.png
    :alt:

  :doc:`/examples/distributed/02_distributed_monitoring`

.. raw:: html

      <div class="sphx-glr-thumbnail-title">Monitoring a Distributed Pipeline: Per-Rank Logging and Profiling</div>
    </div>


.. thumbnail-parent-div-close

.. raw:: html

    </div>


.. toctree::
   :hidden:
   :includehidden:


   /examples/basic/index.rst
   /examples/intermediate/index.rst
   /examples/advanced/index.rst
   /examples/distributed/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>`_