:mod:`nvalchemiops.torch`: Dynamics Optimizers
===================================================

.. currentmodule:: nvalchemiops.torch

The dynamics module provides PyTorch bindings for GPU-accelerated geometry
optimization algorithms.

.. tip::
    For the underlying framework-agnostic Warp kernels and full MD integrators,
    see :doc:`../warp/dynamics`.

.. automodule:: nvalchemiops.torch
    :no-members:
    :no-inherited-members:

FIRE2 Optimizer
---------------

PyTorch adapter for the FIRE2 (Fast Inertial Relaxation Engine v2) geometry optimizer.
These functions accept PyTorch tensors, allocate scratch buffers via PyTorch's CUDA
caching allocator, and call the pure-Warp FIRE2 kernels.

Coordinate-Only Optimization
^^^^^^^^^^^^^^^^^^^^^^^^^^^^

.. autofunction:: nvalchemiops.torch.fire2.fire2_step_coord

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

For optimizing both atomic coordinates and simulation cell parameters simultaneously.

.. autofunction:: nvalchemiops.torch.fire2.fire2_step_coord_cell

Extended Array Interface
^^^^^^^^^^^^^^^^^^^^^^^^

For advanced use cases where you manage packed extended arrays directly.

.. autofunction:: nvalchemiops.torch.fire2.fire2_step_extended

L-BFGS Optimizer
----------------

Quasi-Newton relaxation with a ``maxstep`` trust region. Each step consumes one
force evaluation, updates the history, restarts a non-descending direction and
takes one bounded step. As with FIRE2, convergence is the caller's: there is no
tolerance and no terminal status. No energy is read either -- the step length
comes from the trust region.

.. note::
   Call :func:`~nvalchemiops.torch.lbfgs.lbfgs_prepare_state` once to allocate,
   initialize and validate the whole state; calling it again is how you reset.
   The tensors remain yours -- :class:`~nvalchemiops.dynamics.optimizers.lbfgs.LBFGSState`
   is a plain dataclass, so you can build one from tensors you already own.

.. autofunction:: nvalchemiops.torch.lbfgs.lbfgs_prepare_state
.. autofunction:: nvalchemiops.torch.lbfgs.lbfgs_step_coord

Variable-cell relaxation maps coordinates and cell into one packed coordinate
vector, so the two-loop recursion couples them automatically. Build
``ext_atom_ptr`` and ``ext_batch_idx`` with
:func:`~nvalchemiops.dynamics.utils.cell_filter.extend_atom_ptr` and
:func:`~nvalchemiops.batch_utils.atom_ptr_to_batch_idx`, which handle ragged
batches as well as uniform ones.

.. autofunction:: nvalchemiops.torch.lbfgs.lbfgs_prepare_cell_state
.. autofunction:: nvalchemiops.torch.lbfgs.lbfgs_step_coord_cell