`kups.relaxation.transforms.linesearch` ¶

Per-system line searches that reproduce ASE's implementations.

Two transforms that rescale an incoming descent direction d by a per-system step length t. Like every other transform in this package, the search is taken per system: φ(t), φ'(t) and the accept / backtrack decision are arrays of shape (n_systems,) keyed on the system ids, so a system that is already satisfied never throttles the step of one that is not, and a batched run is bit-identical to running each system on its own.

The two searches reproduce ASE's two line searches step-for-step (a single system matches ASE to floating-point tolerance):

:class:ScaleByBacktrackingLinesearch is :class:ase.utils.linesearcharmijo LineSearchArmijo: quadratic-interpolation backtracking enforcing the Armijo sufficient-decrease condition, with the t ← max(t_quad, t/10) floor.
:class:ScaleByMoreThuenteLinesearch is :class:ase.utils.linesearch LineSearch: the MINPACK More–Thuente search (dcsrch/dcstep) targeting the strong Wolfe conditions, which ASE's LBFGS(use_line_search=True) uses and which pairs naturally with :class:ScaleByAseLbfgs.

ASE bakes a per-atom maxstep clamp into the search (via determine_step); here that clamp is left to the separate :class:MaxStepSize transform, so determine_step is the identity and the searches reproduce the pure algorithms. ASE also fixes the initial trial step at 1.0; t_init = 1.0 (the default) reproduces ASE exactly.

API convention¶

Following the optax composability pattern, the updates passed to :meth:update is the descent direction d produced by the preceding transforms (e.g. -H⁻¹∇L once the L-BFGS preconditioner is sign-flipped), and the raw gradient ∇L arrives as the grad keyword. The search emits t · d per system, so it belongs at the tail of a chain:

.. code-block:: python

from kups.relaxation.optimizer import chain
from kups.relaxation.transforms import ScaleByAseLbfgs, ScaleByMoreThuenteLinesearch
import optax

optimizer = chain(
    ScaleByAseLbfgs(memory_size=10),   # H⁻¹∇L
    optax.scale(-1.0),                 # descent direction d = -H⁻¹∇L
    ScaleByMoreThuenteLinesearch(),           # t · d
)

Each step :class:kups.relaxation.propagator.RelaxationPropagator supplies the current per-system energies (energies), the raw gradient (grad) and a value_and_grad_fn that returns the per-system energies and gradient at a trial point. A system whose direction is not a descent direction (∇L · d ≥ 0) is left unmoved (t = 0).

`ValueAndGradFn = Callable[[PyTree], tuple[Table[SupportsSorting, Array], PyTree]]` ¶

Maps trial params to (per-system energies, gradient pytree).

`LineSearchState` ¶

Line-search state.

Attributes:

Name	Type	Description
`index_prefix`	`PyTree`	Tree prefix whose leaves are `Index` objects, captured at init and used to take every reduction per system.
`prev_phi0`	`Array`	Previous step's per-system energies, `NaN` until a step records them. :class:`ScaleByBacktrackingLinesearch` uses them for its Nocedal & Wright eq. 3.60 initial-step estimate; :class:`ScaleByMoreThuenteLinesearch` leaves them untouched.

Source code in src/kups/relaxation/transforms/linesearch.py

@dataclass
class LineSearchState:
    """Line-search state.

    Attributes:
        index_prefix: Tree prefix whose leaves are ``Index`` objects, captured at
            init and used to take every reduction per system.
        prev_phi0: Previous step's per-system energies, ``NaN`` until a step
            records them. :class:`ScaleByBacktrackingLinesearch` uses them for its
            Nocedal & Wright eq. 3.60 initial-step estimate;
            :class:`ScaleByMoreThuenteLinesearch` leaves them untouched.
    """

    index_prefix: PyTree
    prev_phi0: Array

`ScaleByBacktrackingLinesearch` ¶

Bases: Optimizer[Params, LineSearchState]

Per-system Armijo backtracking line search — ASE LineSearchArmijo.

Rescales the incoming descent direction by a per-system step t: it shrinks t by quadratic interpolation (floored at t/10) until the step meets the Armijo sufficient-decrease condition, deciding per system. The first trial uses a Nocedal & Wright eq. 3.60 estimate from the previous step's energy (clamped to [a_min, a_max], else t_init). See the module docstring for the chain convention.

Attributes:

Name	Type	Description
`c1`	`float`	Armijo sufficient-decrease constant (`0 < c1 < 0.5`).
`a_min`	`float`	Smallest allowed step; the search freezes a system that reaches it. Also the lower clamp on the eq. 3.60 initial-step estimate.
`a_max`	`float`	Upper clamp on the eq. 3.60 initial-step estimate; above it the first trial falls back to `t_init`.
`max_steps`	`int`	Maximum backtracking iterations.
`t_init`	`float`	Initial trial step when no estimate applies, rounded to `1.0` when within `0.5`. `1.0` suits Newton-scaled directions.

Source code in src/kups/relaxation/transforms/linesearch.py

@dataclass
class ScaleByBacktrackingLinesearch[Params](Optimizer[Params, LineSearchState]):
    """Per-system Armijo backtracking line search — ASE ``LineSearchArmijo``.

    Rescales the incoming descent direction by a per-system step ``t``: it shrinks
    ``t`` by quadratic interpolation (floored at ``t/10``) until the step meets the
    Armijo sufficient-decrease condition, deciding per system. The first trial uses
    a Nocedal & Wright eq. 3.60 estimate from the previous step's energy (clamped
    to ``[a_min, a_max]``, else ``t_init``). See the module docstring for the chain
    convention.

    Attributes:
        c1: Armijo sufficient-decrease constant (``0 < c1 < 0.5``).
        a_min: Smallest allowed step; the search freezes a system that reaches it.
            Also the lower clamp on the eq. 3.60 initial-step estimate.
        a_max: Upper clamp on the eq. 3.60 initial-step estimate; above it the
            first trial falls back to ``t_init``.
        max_steps: Maximum backtracking iterations.
        t_init: Initial trial step when no estimate applies, rounded to ``1.0``
            when within ``0.5``. ``1.0`` suits Newton-scaled directions.
    """

    c1: float = field(static=True, default=0.1)
    a_min: float = field(static=True, default=1e-10)
    a_max: float = field(static=True, default=2.0)
    max_steps: int = field(static=True, default=50)
    t_init: float = field(static=True, default=1.0)

    @override
    def init(
        self, parameters: Params, index_prefix: PyTree | None = None
    ) -> LineSearchState:
        return _init_state(parameters, index_prefix)

    @override
    def update(
        self,
        updates: Params,
        state: LineSearchState,
        params: Params | None = None,
        **kwargs: Any,
    ) -> tuple[Params, LineSearchState]:
        idx, keys, value_and_grad_fn, phi0, dphi0 = _setup(
            updates, state, params, kwargs
        )
        t = _armijo(
            params,
            updates,
            idx,
            keys,
            phi0,
            dphi0,
            value_and_grad_fn,
            state.prev_phi0,
            c1=self.c1,
            a_min=self.a_min,
            a_max=self.a_max,
            max_steps=self.max_steps,
            t_init=self.t_init,
        )
        new_state = LineSearchState(index_prefix=state.index_prefix, prev_phi0=phi0)
        return tree_scale_per_row(updates, Table(keys, t), idx), new_state

`ScaleByMoreThuenteLinesearch` ¶

Bases: Optimizer[Params, LineSearchState]

Per-system More–Thuente line search — ASE LineSearch.

Brackets and refines a step meeting the strong Wolfe conditions via the MINPACK dcsrch/dcstep algorithm, deciding per system. This is the search ASE's LBFGS(use_line_search=True) uses; it pairs naturally with :class:ScaleByAseLbfgs, whose secant updates rely on the curvature condition. See the module docstring for the chain convention.

Attributes:

Name	Type	Description
`c1`	`float`	Armijo sufficient-decrease constant (`0 < c1 < c2 < 1`).
`c2`	`float`	Curvature constant (`c1 < c2 < 1`); smaller demands a flatter slope.
`t_init`	`float`	Initial trial step.
`stpmin`	`float`	Smallest step the search will return.
`stpmax`	`float`	Largest step the search will return.
`xtol`	`float`	Relative bracket-width tolerance ending the search.
`xtrapl`	`float`	Lower extrapolation factor while bracketing.
`xtrapu`	`float`	Upper extrapolation factor while bracketing.
`max_steps`	`int`	Maximum combined bracketing + zoom iterations.

Source code in src/kups/relaxation/transforms/linesearch.py

@dataclass
class ScaleByMoreThuenteLinesearch[Params](Optimizer[Params, LineSearchState]):
    """Per-system More–Thuente line search — ASE ``LineSearch``.

    Brackets and refines a step meeting the strong Wolfe conditions via the
    MINPACK ``dcsrch``/``dcstep`` algorithm, deciding per system. This is the
    search ASE's ``LBFGS(use_line_search=True)`` uses; it pairs naturally with
    :class:`ScaleByAseLbfgs`, whose secant updates rely on the curvature
    condition. See the module docstring for the chain convention.

    Attributes:
        c1: Armijo sufficient-decrease constant (``0 < c1 < c2 < 1``).
        c2: Curvature constant (``c1 < c2 < 1``); smaller demands a flatter slope.
        t_init: Initial trial step.
        stpmin: Smallest step the search will return.
        stpmax: Largest step the search will return.
        xtol: Relative bracket-width tolerance ending the search.
        xtrapl: Lower extrapolation factor while bracketing.
        xtrapu: Upper extrapolation factor while bracketing.
        max_steps: Maximum combined bracketing + zoom iterations.
    """

    c1: float = field(static=True, default=0.23)
    c2: float = field(static=True, default=0.46)
    t_init: float = field(static=True, default=1.0)
    stpmin: float = field(static=True, default=1e-8)
    stpmax: float = field(static=True, default=50.0)
    xtol: float = field(static=True, default=1e-14)
    xtrapl: float = field(static=True, default=1.1)
    xtrapu: float = field(static=True, default=4.0)
    max_steps: int = field(static=True, default=50)

    @override
    def init(
        self, parameters: Params, index_prefix: PyTree | None = None
    ) -> LineSearchState:
        return _init_state(parameters, index_prefix)

    @override
    def update(
        self,
        updates: Params,
        state: LineSearchState,
        params: Params | None = None,
        **kwargs: Any,
    ) -> tuple[Params, LineSearchState]:
        idx, keys, value_and_grad_fn, phi0, dphi0 = _setup(
            updates, state, params, kwargs
        )
        t = _more_thuente(
            params,
            updates,
            idx,
            keys,
            phi0,
            dphi0,
            value_and_grad_fn,
            c1=self.c1,
            c2=self.c2,
            t_init=self.t_init,
            stpmin=self.stpmin,
            stpmax=self.stpmax,
            xtol=self.xtol,
            xtrapl=self.xtrapl,
            xtrapu=self.xtrapu,
            max_steps=self.max_steps,
        )
        return tree_scale_per_row(updates, Table(keys, t), idx), state

kups.relaxation.transforms.linesearch ¶

API convention¶

ValueAndGradFn = Callable[[PyTree], tuple[Table[SupportsSorting, Array], PyTree]] ¶

LineSearchState ¶

ScaleByBacktrackingLinesearch ¶

ScaleByMoreThuenteLinesearch ¶

`kups.relaxation.transforms.linesearch` ¶

`ValueAndGradFn = Callable[[PyTree], tuple[Table[SupportsSorting, Array], PyTree]]` ¶

`LineSearchState` ¶

`ScaleByBacktrackingLinesearch` ¶

`ScaleByMoreThuenteLinesearch` ¶