.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/electrostatics/11_pme_energy_derivative_training.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_electrostatics_11_pme_energy_derivative_training.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_electrostatics_11_pme_energy_derivative_training.py:


Energy-Derivative Training Contract (Forces, Stress, Charge Gradients)
======================================================================

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``.

The legacy ``compute_forces`` / ``compute_virial`` / ``compute_charge_gradients``
/ ``hybrid_forces`` flags on ``particle_mesh_ewald`` and ``ewald_summation`` are
deprecated and emit a ``DeprecationWarning``. They remain available for
compatibility in v0.4.0, so this example also checks that the autograd results
match the legacy direct outputs, which is the migration check.

Monopole Ewald/PME/slab entry points accept ``energy_reduction="atom"`` (default)
or ``"system"``. Atom mode returns per-atom energies ``(N,)``; system mode
returns per-system totals ``(B,)`` for batched per-system losses. Direct-output
fields (forces, charge gradients, virials) keep their existing shapes. On Torch
CUDA, eager atom mode may synchronize when proving a materialized uniform
cotangent; ``energy_reduction="system"`` is structurally sync-free.

In this example you will learn:

- Forces from energy autograd: ``F = -grad(E.sum(), positions)``
- Force-loss training (double-backward) with ``create_graph=True``
- Geometry-dependent charges ``q(R)``: the full force includes the
  ``dE/dq . dq/dR`` charge-model chain-rule term
- Strain-first virial/stress and stress-loss training
- Charge gradients from energy autograd: ``dE/dq = grad(E.sum(), charges)``
- ``energy_reduction="system"`` for batched per-system energy layout ``(B,)``

.. important::
    This script is intended as an API demonstration. Do not use this script
    for performance benchmarking; refer to the `benchmarks` folder instead.

.. GENERATED FROM PYTHON SOURCE LINES 54-58

Setup and Imports
-----------------
We use ``particle_mesh_ewald`` throughout; ``ewald_summation`` follows the
exact same energy-derivative contract.

.. GENERATED FROM PYTHON SOURCE LINES 58-69

.. code-block:: Python


    from __future__ import annotations

    import warnings

    import numpy as np
    import torch

    from nvalchemiops.torch.interactions.electrostatics import particle_mesh_ewald
    from nvalchemiops.torch.neighbors import neighbor_list as neighbor_list_fn





.. 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 70-72

Configure Device
----------------

.. GENERATED FROM PYTHON SOURCE LINES 72-83

.. code-block:: Python


    if torch.cuda.is_available():
        device = torch.device("cuda:0")
        print("Using CUDA device")
        print(f"  {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device("cpu")
        print("Using CPU")

    dtype = torch.float64





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Using CUDA device
      NVIDIA H100 PCIe




.. GENERATED FROM PYTHON SOURCE LINES 84-88

Create a NaCl Crystal System
----------------------------
A small NaCl rock-salt supercell (2x2x2 = 16 atoms) is enough to demonstrate
every derivative path. Charges use the same floating dtype as the geometry.

.. GENERATED FROM PYTHON SOURCE LINES 88-138

.. code-block:: Python



    def create_nacl_system(n_cells: int = 2, lattice_constant: float = 5.64):
        """Create a NaCl crystal supercell."""
        base_positions = np.array([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]])
        base_charges = np.array([1.0, -1.0])

        positions = []
        charges = []
        for i in range(n_cells):
            for j in range(n_cells):
                for k in range(n_cells):
                    offset = np.array([i, j, k])
                    for pos, charge in zip(base_positions, base_charges):
                        positions.append((pos + offset) * lattice_constant)
                        charges.append(charge)

        positions = torch.tensor(np.array(positions), dtype=dtype, device=device)
        charges = torch.tensor(np.array(charges), dtype=dtype, device=device)
        cell = torch.eye(3, dtype=dtype, device=device) * lattice_constant * n_cells
        cell = cell.unsqueeze(0)
        pbc = torch.tensor([[True, True, True]], dtype=torch.bool, device=device)
        return positions, charges, cell, pbc


    positions, charges, cell, pbc = create_nacl_system(n_cells=2)

    # Rattle the perfect lattice so the per-atom forces are nonzero -- otherwise the
    # centrosymmetric crystal gives zero forces and the autograd-vs-direct force
    # check below would compare 0 to 0.
    torch.manual_seed(0)
    positions = positions + 0.05 * torch.randn_like(positions)

    print(f"\nSystem: {len(positions)} atoms NaCl crystal (rattled)")

    # Build the real-space neighbor list once; reuse it for every call below.
    neighbor_list, neighbor_ptr, neighbor_shifts = neighbor_list_fn(
        positions, 8.0, cell=cell, pbc=pbc, return_neighbor_list=True
    )

    # Fixed PME parameters keep the example deterministic and fast.
    pme_kwargs = dict(
        alpha=0.35,
        mesh_dimensions=(32, 32, 32),
        spline_order=4,
        neighbor_list=neighbor_list,
        neighbor_ptr=neighbor_ptr,
        neighbor_shifts=neighbor_shifts,
    )





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    System: 16 atoms NaCl crystal (rattled)




.. GENERATED FROM PYTHON SOURCE LINES 139-143

Forces From Energy Autograd
---------------------------
With no deprecated flag, the call returns the per-atom energy only. The full
force is the negative gradient of the total energy w.r.t. positions.

.. GENERATED FROM PYTHON SOURCE LINES 143-152

.. code-block:: Python


    positions_f = positions.detach().requires_grad_(True)
    energy = particle_mesh_ewald(positions_f, charges, cell, **pme_kwargs)
    print(f"\nenergy shape: {tuple(energy.shape)}  (per-atom, energy_reduction='atom')")

    forces = -torch.autograd.grad(energy.sum(), positions_f)[0]
    print(f"forces shape: {tuple(forces.shape)}")
    print(f"max force magnitude: {forces.norm(dim=1).max().item():.6f}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    energy shape: (16,)  (per-atom, energy_reduction='atom')
    forces shape: (16, 3)
    max force magnitude: 0.005254




.. GENERATED FROM PYTHON SOURCE LINES 153-155

Migration check: the autograd force equals the legacy ``compute_forces=True``
direct output (the deprecated flag still works, with a ``DeprecationWarning``).

.. GENERATED FROM PYTHON SOURCE LINES 155-171

.. code-block:: Python


    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            category=DeprecationWarning,
            message=".*compute_forces.*",
        )
        _, forces_direct = particle_mesh_ewald(
            positions, charges, cell, compute_forces=True, **pme_kwargs
        )

    force_diff = (forces - forces_direct).abs().max().item()
    print(f"max |autograd force - direct force|: {force_diff:.2e}")
    if force_diff >= 1e-6:
        raise RuntimeError(f"autograd force does not match direct force: {force_diff:.2e}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    max |autograd force - direct force|: 0.00e+00




.. GENERATED FROM PYTHON SOURCE LINES 172-177

Force-Loss Training (Double-Backward)
-------------------------------------
To train on a force loss, build the force with ``create_graph=True`` so the
subsequent ``loss.backward()`` can differentiate through the force. We use a
trivial linear charge model ``q(R)`` to stand in for a learned charge head.

.. GENERATED FROM PYTHON SOURCE LINES 177-202

.. code-block:: Python


    torch.manual_seed(0)
    weight = torch.randn(3, dtype=dtype, device=device, requires_grad=True)


    def charge_model(pos: torch.Tensor) -> torch.Tensor:
        """Toy geometry-dependent charge model with enforced neutrality."""
        raw = pos @ weight
        return raw - raw.mean()


    positions_q = positions.detach().requires_grad_(True)
    charges_qr = charge_model(positions_q)  # kept in the graph -> full q(R) force

    energy = particle_mesh_ewald(positions_q, charges_qr, cell, **pme_kwargs)
    forces_qr = -torch.autograd.grad(energy.sum(), positions_q, create_graph=True)[0]

    target_forces = torch.zeros_like(forces_qr)
    force_loss = (forces_qr - target_forces).pow(2).sum()
    force_loss.backward()  # differentiates through the force construction

    print(f"\nq(R) force-loss backward OK; weight.grad shape: {tuple(weight.grad.shape)}")
    if not torch.isfinite(weight.grad).all():
        raise RuntimeError("q(R) force-loss produced non-finite parameter gradients")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    q(R) force-loss backward OK; weight.grad shape: (3,)




.. GENERATED FROM PYTHON SOURCE LINES 203-208

Because ``charges = charge_model(positions)`` stays connected to ``positions``,
the autograd force includes both the fixed-charge term and the
``dE/dq . dq/dR`` chain-rule term. The legacy ``compute_forces=True`` output is
only the fixed-charge partial and does not include the charge-model term --
this is the central reason direct force output on the full API is deprecated.

.. GENERATED FROM PYTHON SOURCE LINES 210-216

Strain-First Virial and Stress
------------------------------
``strain`` is not a PME argument. Build a differentiable strain tensor, deform
positions and cell by ``I + strain``, and let autograd map gradients back to
strain. The virial is ``W = -dE/d(strain)`` and tensile-positive stress is
``dE/d(strain) / V``; see ``conventions.md`` for the project-wide sign.

.. GENERATED FROM PYTHON SOURCE LINES 216-237

.. code-block:: Python


    num_systems = cell.shape[0]
    positions_s = positions.detach().requires_grad_(True)
    strain = torch.zeros(num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True)
    eye = torch.eye(3, device=device, dtype=dtype).unsqueeze(0)
    deform = eye + strain

    # Single system: every atom maps to system 0.
    batch_idx = torch.zeros(positions_s.shape[0], dtype=torch.int32, device=device)
    positions_def = torch.einsum("ni,nij->nj", positions_s, deform[batch_idx])
    cell_def = torch.einsum("bij,bjk->bik", cell, deform)

    energy = particle_mesh_ewald(positions_def, charges, cell_def, **pme_kwargs)
    grad_strain = torch.autograd.grad(energy.sum(), strain)[0]
    virial = -grad_strain
    volume = torch.abs(torch.linalg.det(cell_def))
    stress = grad_strain / volume[:, None, None]  # tensile-positive Cauchy

    print(f"\nvirial shape: {tuple(virial.shape)}")
    print(f"stress shape: {tuple(stress.shape)}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    virial shape: (1, 3, 3)
    stress shape: (1, 3, 3)




.. GENERATED FROM PYTHON SOURCE LINES 238-240

Migration check: the strain-first virial equals the legacy
``compute_virial=True`` direct virial (both are ``-dE/d(strain)``).

.. GENERATED FROM PYTHON SOURCE LINES 240-258

.. code-block:: Python


    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            category=DeprecationWarning,
            message=".*compute_virial.*",
        )
        _, virial_direct = particle_mesh_ewald(
            positions, charges, cell, compute_virial=True, **pme_kwargs
        )

    virial_diff = (virial - virial_direct).abs().max().item()
    print(f"max |strain-first virial - direct virial|: {virial_diff:.2e}")
    if virial_diff >= 1e-5:
        raise RuntimeError(
            f"strain-first virial does not match direct virial: {virial_diff:.2e}"
        )





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    max |strain-first virial - direct virial|: 1.11e-16




.. GENERATED FROM PYTHON SOURCE LINES 259-263

Stress-Loss Training (Double-Backward)
--------------------------------------
Stress training uses the same strain-first recipe with ``create_graph=True``,
so the stress loss back-propagates to model parameters.

.. GENERATED FROM PYTHON SOURCE LINES 263-288

.. code-block:: Python


    weight_s = torch.randn(3, dtype=dtype, device=device, requires_grad=True)
    positions_s = positions.detach().requires_grad_(True)
    strain = torch.zeros(num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True)
    deform = torch.eye(3, device=device, dtype=dtype).unsqueeze(0) + strain
    positions_def = torch.einsum("ni,nij->nj", positions_s, deform[batch_idx])
    cell_def = torch.einsum("bij,bjk->bik", cell, deform)

    charges_qr = positions_def @ weight_s
    charges_qr = charges_qr - charges_qr.mean()

    energy = particle_mesh_ewald(positions_def, charges_qr, cell_def, **pme_kwargs)
    grad_strain = torch.autograd.grad(energy.sum(), strain, create_graph=True)[0]
    virial = -grad_strain
    volume = torch.abs(torch.linalg.det(cell_def))
    stress = grad_strain / volume[:, None, None]

    stress_loss = stress.pow(2).sum()
    stress_loss.backward()

    stress_grad_ok = bool(torch.isfinite(weight_s.grad).all())
    print(f"\nstress-loss backward OK; weight_s.grad finite: {stress_grad_ok}")
    if not stress_grad_ok:
        raise RuntimeError("stress-loss produced non-finite parameter gradients")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    stress-loss backward OK; weight_s.grad finite: True




.. GENERATED FROM PYTHON SOURCE LINES 289-298

Combined Energy + Force + Stress Loss (Performance Tip)
-------------------------------------------------------
When a single loss mixes energy, forces, AND stress, take the forces and the
virial from **one** ``torch.autograd.grad`` call over both ``positions`` and
``strain`` -- NOT two separate calls. Each ``create_graph=True`` ``grad`` call
builds its own first-derivative graph node, and ``loss.backward()`` then runs the
(O(K*N)) reciprocal second-derivative once per node. Fusing the two into a single
call avoids duplicate reciprocal double-backward work; the gradients are
bit-identical to the two-call form.

.. GENERATED FROM PYTHON SOURCE LINES 298-330

.. code-block:: Python

    weight_m = torch.randn(3, dtype=dtype, device=device, requires_grad=True)
    positions_m = positions.detach().requires_grad_(True)
    strain_m = torch.zeros(
        num_systems, 3, 3, device=device, dtype=dtype, requires_grad=True
    )
    deform_m = torch.eye(3, device=device, dtype=dtype).unsqueeze(0) + strain_m
    positions_md = torch.einsum("ni,nij->nj", positions_m, deform_m[batch_idx])
    cell_md = torch.einsum("bij,bjk->bik", cell, deform_m)
    charges_m = positions_md @ weight_m
    charges_m = charges_m - charges_m.mean()

    energy = particle_mesh_ewald(positions_md, charges_m, cell_md, **pme_kwargs)
    # ONE combined grad call -> one double-backward (do this instead of separate
    # ``grad(E, positions)`` and ``grad(E, strain)`` calls):
    grad_pos, grad_strain = torch.autograd.grad(
        energy.sum(), (positions_m, strain_m), create_graph=True
    )
    # ``positions_m`` are the undeformed reference coordinates, so ``forces_m`` are
    # reference-frame forces. Differentiate with respect to ``positions_md`` instead
    # when training against deformed-coordinate force targets.
    forces_m = -grad_pos
    virial_m = -grad_strain
    volume_m = torch.abs(torch.linalg.det(cell_md))
    stress_m = grad_strain / volume_m[:, None, None]

    mixed_loss = energy.sum() + forces_m.pow(2).sum() + stress_m.pow(2).sum()
    mixed_loss.backward()
    mixed_grad_ok = bool(torch.isfinite(weight_m.grad).all())
    print(f"\nmixed E+F+stress backward OK; weight_m.grad finite: {mixed_grad_ok}")
    if not mixed_grad_ok:
        raise RuntimeError("mixed-loss produced non-finite parameter gradients")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    mixed E+F+stress backward OK; weight_m.grad finite: True




.. GENERATED FROM PYTHON SOURCE LINES 331-334

Charge Gradients From Energy Autograd
-------------------------------------
``dE/dq`` is an ordinary gradient of the energy w.r.t. charges.

.. GENERATED FROM PYTHON SOURCE LINES 334-360

.. code-block:: Python


    charges_g = charges.detach().requires_grad_(True)
    energy = particle_mesh_ewald(positions, charges_g, cell, **pme_kwargs)
    charge_grad = torch.autograd.grad(energy.sum(), charges_g)[0]
    print(f"\ncharge gradient shape: {tuple(charge_grad.shape)}")

    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            category=DeprecationWarning,
            message=".*compute_forces.*",
        )
        _, _, charge_grad_direct = particle_mesh_ewald(
            positions,
            charges,
            cell,
            compute_forces=True,
            compute_charge_gradients=True,
            **pme_kwargs,
        )

    cg_diff = (charge_grad - charge_grad_direct).abs().max().item()
    print(f"max |autograd dE/dq - direct dE/dq|: {cg_diff:.2e}")
    if cg_diff >= 1e-6:
        raise RuntimeError(f"autograd dE/dq does not match direct dE/dq: {cg_diff:.2e}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    charge gradient shape: (16,)
    max |autograd dE/dq - direct dE/dq|: 1.80e-09




.. GENERATED FROM PYTHON SOURCE LINES 361-381

Summary
-------
This example demonstrated the energy-derivative training contract:

1. **Forces** -- ``F = -torch.autograd.grad(E.sum(), positions)[0]``; with
   ``create_graph=True`` for force-loss training.
2. **q(R) forces** -- keep ``charges = charge_model(positions)`` in the graph so
   the full ``dE/dR`` includes the ``dE/dq . dq/dR`` charge-model term.
3. **Virial / stress** -- strain-first: deform positions and cell by
   ``I + strain``, then ``virial = -grad(E.sum(), strain)`` and
   ``stress = grad(E.sum(), strain) / volume[:, None, None]``.
4. **Charge gradients** -- ``dE/dq = torch.autograd.grad(E.sum(), charges)[0]``.
5. **Combined E + F + stress loss** -- take forces and virial from a single
   ``torch.autograd.grad(E.sum(), (positions, strain), create_graph=True)`` call,
   not two separate calls, so the reciprocal double-backward runs once.
6. **Batched per-system layout** -- pass ``energy_reduction="system"`` to get
   ``(B,)`` energies for per-system losses without manual ``scatter_add``.

Each autograd result matched the corresponding (deprecated) direct kernel
output, confirming the migration is numerically exact.

.. GENERATED FROM PYTHON SOURCE LINES 381-383

.. code-block:: Python


    print("\nEnergy-derivative training contract example complete!")




.. rst-class:: sphx-glr-script-out

 .. code-block:: none


    Energy-derivative training contract example complete!





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 10.444 seconds)


.. _sphx_glr_download_examples_electrostatics_11_pme_energy_derivative_training.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

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

      :download:`Download Jupyter notebook: 11_pme_energy_derivative_training.ipynb <11_pme_energy_derivative_training.ipynb>`

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

      :download:`Download Python source code: 11_pme_energy_derivative_training.py <11_pme_energy_derivative_training.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 11_pme_energy_derivative_training.zip <11_pme_energy_derivative_training.zip>`


.. only:: html

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

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