`nvalchemiops.jax.lbfgs`: Geometry Optimization#

JAX bindings for the batched L-BFGS geometry optimizer.

L-BFGS reaches a given force tolerance in far fewer force evaluations than the FIRE optimizers – the cost that dominates relaxation with a machine-learned potential.

JAX arrays are immutable, so unlike the PyTorch binding these entry points return a new state. Both LBFGSState and LBFGSCellState are registered pytrees, so a state crosses jax.jit as one argument and one donate_argnums entry donates every field:

state = lbfgs_prepare_state(num_atoms, num_systems)

@functools.partial(jax.jit, donate_argnums=(0, 1))
def relax_step(positions, state, forces):
    return lbfgs_step_coord(positions, forces, state, batch_idx)

for _ in range(max_steps):
    forces = model(positions)
    if float(jnp.linalg.norm(forces, axis=1).max()) < force_tol:
        break
    positions, state = relax_step(positions, state, forces)

Test before stepping, as above: the returned positions are a new, unevaluated point, and the forces you passed describe the point before it (kept in x_base/force_base). The optimizer owns no tolerance and has no terminal status, matching FIRE2. Reading a norm back costs a host synchronization, so test every few steps if that matters.

These operations are not differentiable: jax.grad through a step fails rather than returning a silently wrong answer.

State#

lbfgs_prepare_state() allocates, initializes and validates the whole state; in a functional setting resetting and allocating are the same thing, so calling it again is how you restart. The arrays stay yours – the state is a plain dataclass with a validate() method; see nvalchemiops.dynamics.optimizers.lbfgs for shapes and required initial contents.

Every array follows the coordinate dtype, so an fp32 state is fp32 end to end and needs no JAX_ENABLE_X64.

Donation and pointer stability#

Every mutable array is an input-output alias, so XLA may reuse each input buffer for the matching output. Both consequences are about performance:

Donate the state. Without donation JAX copies, doubling peak memory and moving the pointers. On the variable-cell path donate state but not cell_state: its five chart fields are read-only and come back unchanged, so XLA cannot reuse their buffers.
Keep topology out of the donated set. batch_idx never changes, so close over it.

Under JaxCallableGraphMode.WARP the step replays as a CUDA graph. The capture is keyed on input addresses, so a fresh forces array each step gives a small working set rather than one graph per call; measurements settle at four or five. If the count grows without bound, pass graph_mode="warp_staged", which keys on the call instead at the cost of one copy per staged array.

Scalars are baked into the compiled call, so changing maxstep between steps triggers a recompilation.

Coordinate Relaxation#

nvalchemiops.jax.lbfgs.lbfgs_prepare_state(num_dofs, num_systems, *, dtype=jnp.float64, history_size=6)[source]#

Allocate, initialize and validate a complete optimizer state.

Call this once, before the first step. The three fields that do not start at zero are set for you – alpha_step to one, iteration to minus one. Calling it again is how you reset.

LBFGSState is a registered pytree, so the whole thing can be passed through jax.jit, donated, and carried by lax.while_loop. You are not obliged to use this function: build one from arrays you own and call LBFGSState.validate() yourself.

Parameters:

num_dofs (int) – Degrees of freedom the optimizer moves. On the variable-cell path this is num_atoms + 2 * num_systems, not the atom count.
num_systems (int) – Independent systems in the batch.
dtype (optional) – Coordinate precision, float32 or float64. Every per-system scalar follows it, so this selects an end-to-end fp32 or fp64 state – and an fp32 state needs no JAX_ENABLE_X64.
history_size (int, optional) – Stored curvature pairs m; 3 to 7 is the usual range.

Return type:

LBFGSState

nvalchemiops.jax.lbfgs.lbfgs_step_coord(positions, forces, state, batch_idx, *, maxstep=0.2, curvature_eps=None, graph_mode='warp')[source]#

Advance one batched L-BFGS step, consuming one force evaluation.

Returns new arrays; nothing is mutated in place. Donate positions and state so XLA can reuse their memory.

Parameters:

positions (jax.Array, shape (num_atoms, 3)) – Current geometry.
forces (jax.Array, shape (num_atoms, 3)) – Forces at positions. Forces, not gradients.
state (LBFGSState) – From lbfgs_prepare_state(), or built from your own arrays.
batch_idx (jax.Array, shape (num_atoms,), dtype int32) – Sorted system index per atom. Static topology, so close over it.
maxstep (float, optional) – Largest distance an atom may move in one step. Zero disables the trust region.
curvature_eps (float, optional) – Relative threshold below which a curvature pair is discarded. Defaults to 1e-6 for float32 coordinates and 1e-10 for float64: ys is accumulated at the coordinate precision, so one value cannot serve both.
graph_mode ({"warp", "warp_staged", "none"}, optional) – How the step is captured. "warp" replays a CUDA graph and is the default; "warp_staged" keys the capture on the call rather than on buffer addresses, which helps if the capture count grows without bound; "none" disables capture.

Returns:

positions (jax.Array) – The advanced geometry.
state (LBFGSState) – The advanced state. The input state is unchanged.

Raises:

ValueError – If this call’s inputs are incompatible with the state.

Return type:

tuple[Array, LBFGSState]

Variable-Cell Relaxation#

Coordinates and cell are mapped into one packed coordinate vector, so the two-loop recursion couples them automatically. Build ext_atom_ptr and ext_batch_idx with extend_atom_ptr() and atom_ptr_to_batch_idx(), which handle ragged batches as well as uniform ones.

nvalchemiops.jax.lbfgs.lbfgs_prepare_cell_state(atom_ptr, cell, *, cell_force_scale=1.0, dtype=jnp.float64, alignment_atol=None)[source]#

Build a complete, ready-to-step variable-cell chart from atom topology.

One call: give it the ordinary atom_ptr and aligned cells, and it derives the packed topology, captures the reference chart and computes kappa. Alignment is checked before the reference cell is captured.

Ragged batches are unaffected – atom_ptr already carries each system’s atom count, so nothing here assumes an even split.

Parameters:

atom_ptr (jax.Array, shape (num_systems + 1,), dtype int32) – CSR-style atom pointer for the batch.
cell (jax.Array, shape (num_systems, 3, 3)) – Reference lattice per system, aligned.
cell_force_scale (float, optional) – Multiplier on the atom count in kappa.
dtype (optional) – Coordinate precision; must match the coordinates and the state.
alignment_atol (float, optional) – Absolute tolerance for entries strictly above the cell diagonal. Defaults to 1e-4 for float32 cells and 1e-8 for float64 cells.

Return type:

LBFGSCellState

nvalchemiops.jax.lbfgs: Geometry Optimization#

State#

Donation and pointer stability#

Coordinate Relaxation#

Variable-Cell Relaxation#

`nvalchemiops.jax.lbfgs`: Geometry Optimization#