trust_region

A scaled trust-region quasi-Newton optimiser.

At each iteration a quadratic model m(p) = g.p + p.B.p / 2 is minimised inside a ball of radius delta; the step is taken, and the ratio of the actual reduction to the predicted one decides whether to keep it and what to do with the radius. The implementation follows klasp's trust_region.hpp, generalised from rigid bodies to any Coordinates.

Four details carry most of the robustness, each of them a failure mode that had to be fixed rather than a refinement:

• Scaled coordinates. Degrees of freedom are measured in Angstroms of atomic displacement (Coordinates.scale), so one radius and one gradient tolerance mean the same thing for an atom, a cell strain and anything added later. Unscaled, the radius is set by whichever freedom carries the largest units.

• Powell damping. Plain BFGS must skip the curvature update whenever y.s <= 0, which is where the model is worst -- and in a trust region those are also the steps that get rejected, so a model that causes a rejection never learns from it. Damping blends y towards Bs by just enough to keep the update positive definite, so every step contributes.

• The radius shrinks only on rejection. Shrinking whenever rho < 1/4 regardless of acceptance is the textbook rule; measured against this objective it over-reacts to a noisy energy and collapses the radius.

• A floor on the radius, and restarts. Repeated rejection drives the radius geometrically to zero; once a step is too small to change the energy, every later trial is rejected and the structure is stuck. On reaching the floor the model is reset from the accepted point, and after a few restarts it pins at the floor and keeps inching.

A fifth is about the calculator rather than the algorithm: below the smallest energy difference a calculator can resolve, the reduction ratio carries no information and is not computed. See _judge.

Relaxation dataclass

The outcome of a relaxation.

Attributes:

Name	Type	Description
converged	bool	whether every convergence criterion was met
steps	int	iterations taken
energy	float	final energy in eV
measures	dict	final convergence measures, e.g. fmax in eV/A and smax in GPa
evaluations	int	calculator evaluations used
history	list	one Step per iteration
structure	object	the relaxed structure
result	object	the calculator result at the final geometry
model	object	the CurvatureModel the run ended with, so a later stage can start from what this one learned
stages	list	for a staged relaxation, the Relaxation of each stage

Source code in chmpy/opt/trust_region.py

@dataclass
class Relaxation:
    """The outcome of a relaxation.

    Attributes:
        converged: whether every convergence criterion was met
        steps: iterations taken
        energy: final energy in eV
        measures: final convergence measures, e.g. fmax in eV/A and smax in GPa
        evaluations: calculator evaluations used
        history: one `Step` per iteration
        structure: the relaxed structure
        result: the calculator result at the final geometry
        model: the `CurvatureModel` the run ended with, so a later stage can
            start from what this one learned
        stages: for a staged relaxation, the `Relaxation` of each stage
    """

    converged: bool
    steps: int
    energy: float
    measures: dict
    evaluations: int
    history: list = field(default_factory=list)
    structure: object = None
    result: object = None
    stalled: bool = False
    model: object = None
    stages: list = field(default_factory=list)

    @property
    def accepted_steps(self) -> int:
        "How many trial steps were kept"
        return sum(1 for step in self.history if step.accepted)

    def __repr__(self) -> str:
        state = "converged" if self.converged else "NOT converged"
        measures = ", ".join(f"{k}={v:.4g}" for k, v in self.measures.items())
        return (
            f"<Relaxation {state} in {self.steps} steps "
            f"({self.evaluations} evaluations), E={self.energy:.6f} eV, {measures}>"
        )

    def __str__(self) -> str:
        lines = [repr(self), f"  {self.accepted_steps} of {self.steps} steps accepted"]
        if len(self.stages) > 1:
            lines.extend(f"    {stage!r}" for stage in self.stages)
        return "\n".join(lines)

accepted_steps property

How many trial steps were kept

Step dataclass

One iteration, for the trajectory.

Source code in chmpy/opt/trust_region.py

@dataclass
class Step:
    """One iteration, for the trajectory."""

    step: int
    energy: float
    measures: dict
    radius: float
    rho: float
    accepted: bool
    restarted: bool = False
    reanchored: bool = False

TrustRegion

Relax a structure by a scaled dogleg trust region with damped BFGS.

Parameters:

Name	Type	Description	Default
coordinates		the degrees of freedom to vary	required
calculator		what to evaluate energies with	required
options		tuning, or None for the defaults	None
pressure	float	external hydrostatic pressure in GPa, applied when the coordinates include a cell	0.0
hessian		(n_dof, n_dof) starting model of the curvature, in scaled coordinates. The default is the identity, which is a reasonable model precisely because the coordinates are scaled to Angstroms. Passing a better one -- an analytic or previously converged Hessian -- is the single biggest saving available on a hard system.	None
model		the curvature model, or None to build one. Anything with matvec, solve, update, reset and rescale will do; see chmpy.opt.curvature.	None

Source code in chmpy/opt/trust_region.py

class TrustRegion:
    """Relax a structure by a scaled dogleg trust region with damped BFGS.

    Args:
        coordinates: the degrees of freedom to vary
        calculator: what to evaluate energies with
        options: tuning, or None for the defaults
        pressure: external hydrostatic pressure in GPa, applied when the
            coordinates include a cell
        hessian: (n_dof, n_dof) starting model of the curvature, in scaled
            coordinates. The default is the identity, which is a reasonable
            model precisely because the coordinates are scaled to Angstroms.
            Passing a better one -- an analytic or previously converged
            Hessian -- is the single biggest saving available on a hard system.
        model: the curvature model, or None to build one. Anything with
            `matvec`, `solve`, `update`, `reset` and `rescale` will do; see
            `chmpy.opt.curvature`.
    """

    def __init__(
        self,
        coordinates,
        calculator,
        options=None,
        pressure: float = 0.0,
        hessian=None,
        model=None,
    ):
        self.coordinates = coordinates
        self.calculator = calculator
        self.options = options or TrustRegionOptions()
        self.pressure = float(pressure)
        self.hessian = None if hessian is None else np.array(hessian, dtype=float)
        self.model = model

    def run(
        self,
        fmax: float = 0.01,
        smax: float = 0.05,
        steps: int = 200,
        progress=None,
    ) -> Relaxation:
        """Relax until converged or out of steps.

        Args:
            fmax: force convergence in eV/A, on the largest component
            smax: stress convergence in GPa, on the largest component. Ignored
                when the coordinates hold no cell degrees of freedom.
            steps: maximum iterations
            progress: True to print a line per step, or a callable given a
                `chmpy.opt.progress.Progress` for each, carrying the `Step`

        Returns:
            Relaxation
        """
        report = reporter(progress)
        options = self.options
        coordinates = self.coordinates
        tolerances = {"fmax": fmax, "smax": smax}
        started = self.calculator.stats.calls

        x = coordinates.get()
        scale = coordinates.scale()
        result = self._evaluate(x)
        energy = self._energy(result)
        gradient = self._gradient(result, scale)

        n = coordinates.n_dof
        model = self.model
        if model is None:
            model = for_size(n, initial=self.hessian)
        if model.n != n:
            raise ValueError(
                f"the curvature model has {model.n} degrees of freedom, but "
                f"the coordinates have {n}"
            )
        LOG.debug("relaxing %d degrees of freedom with %r", n, model)
        radius = options.delta0
        restarts = 0
        stalled = False
        inconsistent = False
        floor_rejections = 0
        history = []
        converged = self._converged(result, tolerances)

        for iteration in range(1, steps + 1):
            if converged:
                break

            step, predicted = _dogleg(model, gradient, radius)

            reanchored = False
            fraction = coordinates.step_fraction(x, step / scale)
            if fraction < REANCHOR_FRACTION:
                # The step has run into a bound on the parameterisation. Move
                # the reference here -- the geometry does not change, so the
                # result and its gradient still stand -- and ask again. Waiting
                # for an accepted step to do this would wait for ever: at a
                # bound the step is clipped to nothing, so nothing is ever
                # accepted and the radius shrinks to the floor instead.
                anchored = coordinates.reanchor(x, force=True)
                if anchored is not None:
                    reanchored = True
                    x = anchored
                    scale, previous = coordinates.scale(), scale
                    model.rescale(previous, scale)
                    gradient = self._gradient(result, scale)
                    step, predicted = _dogleg(model, gradient, radius)
                    fraction = coordinates.step_fraction(x, step / scale)
            if fraction < 1.0:
                step = step * fraction
                predicted = model.quadratic(gradient, step)

            if predicted <= 0.0 and np.linalg.norm(step) <= 1e-14:
                # nothing left to try: the model proposes no move at all
                stalled = True
                break

            trial_x = x + step / scale
            trial_result = self._evaluate(trial_x)
            trial_energy = self._energy(trial_result)
            trial_gradient = self._gradient(trial_result, scale)

            noise = (
                self.calculator.energy_noise(energy)
                if options.energy_noise is None
                else options.energy_noise
            )
            rho, accept, unusable = _judge(
                trial_energy, energy, predicted, options.eta, options.rho_max, noise
            )

            # Both branches learn from the step: y = g(x + p) - g(x) is valid
            # curvature whether or not the point is kept, and throwing away the
            # rejected ones is what makes a bad model self-perpetuating.
            if np.isfinite(trial_energy):
                model.update(step, trial_gradient - gradient)

            restarted = False
            if accept:
                x, energy, gradient, result = (
                    trial_x,
                    trial_energy,
                    trial_gradient,
                    trial_result,
                )
                if rho > 0.75 and np.linalg.norm(step) > 0.8 * radius:
                    radius = min(options.grow * radius, options.delta_max)
                anchored = coordinates.reanchor(x)
                if anchored is not None:
                    # the geometry has not changed, only its parameterisation,
                    # so the result still stands and the gradient is re-derived
                    reanchored = True
                    x = anchored
                    scale, previous = coordinates.scale(), scale
                    model.rescale(previous, scale)
                    gradient = self._gradient(result, scale)
                converged = self._converged(result, tolerances)
            else:
                coordinates.set(x)
                radius, restart = _shrink(radius, unusable, restarts, options)
                if restart:
                    restarts += 1
                    restarted = True
                    model.reset()

            entry = Step(
                iteration,
                energy,
                coordinates.measures(result),
                radius,
                rho,
                accept,
                restarted,
                reanchored,
            )
            history.append(entry)
            report(
                "relax",
                "step",
                _format(entry),
                index=iteration - 1,
                total=steps,
                step=entry,
            )

            # rejected steps at the minimum radius, after every restart, mean
            # the energy rises along the gradient: nothing left to try
            if accept:
                floor_rejections = 0
            elif radius <= options.delta_min and restarts >= options.max_restarts:
                floor_rejections += 1
                if floor_rejections >= FLOOR_REJECTIONS:
                    stalled = inconsistent = True
                    break

        coordinates.set(x)
        _warn_if_forbidden(coordinates.measures(result), fmax)
        if inconsistent:
            LOG.warning(
                "the optimiser stopped: %d steps in a row at the minimum trust "
                "radius raised the energy, so the gradient does not match the "
                "energy. Check the calculator with `calculator.check_gradients`.",
                FLOOR_REJECTIONS,
            )
        elif stalled:
            LOG.warning(
                "the optimiser stalled: the model proposes no step at all, which "
                "usually means the parameterisation has run into a bound it "
                "cannot be re-anchored out of"
            )
        return Relaxation(
            converged=converged,
            steps=len(history),
            energy=energy,
            measures=coordinates.measures(result),
            evaluations=self.calculator.stats.calls - started,
            history=history,
            structure=coordinates.system,
            result=result,
            stalled=stalled,
            model=model,
        )

    # -- internals -----------------------------------------------------------

    def _evaluate(self, x):
        self.coordinates.set(x)
        return self.calculator(self.coordinates.system, self.coordinates.wanted)

    def _energy(self, result) -> float:
        """The objective: the enthalpy when a pressure is applied."""
        if self.pressure:
            from chmpy.calc.result import EV_PER_ANGSTROM3_TO_GPA

            return (
                result.energy
                + (self.pressure / EV_PER_ANGSTROM3_TO_GPA)
                * self.coordinates.system.volume
            )
        return result.energy

    def _gradient(self, result, scale) -> np.ndarray:
        gradient = self.coordinates.gradient(result)
        if self.pressure:
            from .coordinates import pressure_term

            gradient = gradient + pressure_term(self.coordinates, self.pressure)
        return gradient / scale

    def _converged(self, result, tolerances) -> bool:
        measures = self.coordinates.measures(result)
        return all(
            value <= tolerances[name]
            for name, value in measures.items()
            if name in tolerances
        )

run(fmax=0.01, smax=0.05, steps=200, progress=None)

Relax until converged or out of steps.

Parameters:

Name	Type	Description	Default
fmax	float	force convergence in eV/A, on the largest component	0.01
smax	float	stress convergence in GPa, on the largest component. Ignored when the coordinates hold no cell degrees of freedom.	0.05
steps	int	maximum iterations	200
progress		True to print a line per step, or a callable given a chmpy.opt.progress.Progress for each, carrying the Step	None

Returns:

Type	Description
Relaxation	Relaxation

Source code in chmpy/opt/trust_region.py

def run(
    self,
    fmax: float = 0.01,
    smax: float = 0.05,
    steps: int = 200,
    progress=None,
) -> Relaxation:
    """Relax until converged or out of steps.

    Args:
        fmax: force convergence in eV/A, on the largest component
        smax: stress convergence in GPa, on the largest component. Ignored
            when the coordinates hold no cell degrees of freedom.
        steps: maximum iterations
        progress: True to print a line per step, or a callable given a
            `chmpy.opt.progress.Progress` for each, carrying the `Step`

    Returns:
        Relaxation
    """
    report = reporter(progress)
    options = self.options
    coordinates = self.coordinates
    tolerances = {"fmax": fmax, "smax": smax}
    started = self.calculator.stats.calls

    x = coordinates.get()
    scale = coordinates.scale()
    result = self._evaluate(x)
    energy = self._energy(result)
    gradient = self._gradient(result, scale)

    n = coordinates.n_dof
    model = self.model
    if model is None:
        model = for_size(n, initial=self.hessian)
    if model.n != n:
        raise ValueError(
            f"the curvature model has {model.n} degrees of freedom, but "
            f"the coordinates have {n}"
        )
    LOG.debug("relaxing %d degrees of freedom with %r", n, model)
    radius = options.delta0
    restarts = 0
    stalled = False
    inconsistent = False
    floor_rejections = 0
    history = []
    converged = self._converged(result, tolerances)

    for iteration in range(1, steps + 1):
        if converged:
            break

        step, predicted = _dogleg(model, gradient, radius)

        reanchored = False
        fraction = coordinates.step_fraction(x, step / scale)
        if fraction < REANCHOR_FRACTION:
            # The step has run into a bound on the parameterisation. Move
            # the reference here -- the geometry does not change, so the
            # result and its gradient still stand -- and ask again. Waiting
            # for an accepted step to do this would wait for ever: at a
            # bound the step is clipped to nothing, so nothing is ever
            # accepted and the radius shrinks to the floor instead.
            anchored = coordinates.reanchor(x, force=True)
            if anchored is not None:
                reanchored = True
                x = anchored
                scale, previous = coordinates.scale(), scale
                model.rescale(previous, scale)
                gradient = self._gradient(result, scale)
                step, predicted = _dogleg(model, gradient, radius)
                fraction = coordinates.step_fraction(x, step / scale)
        if fraction < 1.0:
            step = step * fraction
            predicted = model.quadratic(gradient, step)

        if predicted <= 0.0 and np.linalg.norm(step) <= 1e-14:
            # nothing left to try: the model proposes no move at all
            stalled = True
            break

        trial_x = x + step / scale
        trial_result = self._evaluate(trial_x)
        trial_energy = self._energy(trial_result)
        trial_gradient = self._gradient(trial_result, scale)

        noise = (
            self.calculator.energy_noise(energy)
            if options.energy_noise is None
            else options.energy_noise
        )
        rho, accept, unusable = _judge(
            trial_energy, energy, predicted, options.eta, options.rho_max, noise
        )

        # Both branches learn from the step: y = g(x + p) - g(x) is valid
        # curvature whether or not the point is kept, and throwing away the
        # rejected ones is what makes a bad model self-perpetuating.
        if np.isfinite(trial_energy):
            model.update(step, trial_gradient - gradient)

        restarted = False
        if accept:
            x, energy, gradient, result = (
                trial_x,
                trial_energy,
                trial_gradient,
                trial_result,
            )
            if rho > 0.75 and np.linalg.norm(step) > 0.8 * radius:
                radius = min(options.grow * radius, options.delta_max)
            anchored = coordinates.reanchor(x)
            if anchored is not None:
                # the geometry has not changed, only its parameterisation,
                # so the result still stands and the gradient is re-derived
                reanchored = True
                x = anchored
                scale, previous = coordinates.scale(), scale
                model.rescale(previous, scale)
                gradient = self._gradient(result, scale)
            converged = self._converged(result, tolerances)
        else:
            coordinates.set(x)
            radius, restart = _shrink(radius, unusable, restarts, options)
            if restart:
                restarts += 1
                restarted = True
                model.reset()

        entry = Step(
            iteration,
            energy,
            coordinates.measures(result),
            radius,
            rho,
            accept,
            restarted,
            reanchored,
        )
        history.append(entry)
        report(
            "relax",
            "step",
            _format(entry),
            index=iteration - 1,
            total=steps,
            step=entry,
        )

        # rejected steps at the minimum radius, after every restart, mean
        # the energy rises along the gradient: nothing left to try
        if accept:
            floor_rejections = 0
        elif radius <= options.delta_min and restarts >= options.max_restarts:
            floor_rejections += 1
            if floor_rejections >= FLOOR_REJECTIONS:
                stalled = inconsistent = True
                break

    coordinates.set(x)
    _warn_if_forbidden(coordinates.measures(result), fmax)
    if inconsistent:
        LOG.warning(
            "the optimiser stopped: %d steps in a row at the minimum trust "
            "radius raised the energy, so the gradient does not match the "
            "energy. Check the calculator with `calculator.check_gradients`.",
            FLOOR_REJECTIONS,
        )
    elif stalled:
        LOG.warning(
            "the optimiser stalled: the model proposes no step at all, which "
            "usually means the parameterisation has run into a bound it "
            "cannot be re-anchored out of"
        )
    return Relaxation(
        converged=converged,
        steps=len(history),
        energy=energy,
        measures=coordinates.measures(result),
        evaluations=self.calculator.stats.calls - started,
        history=history,
        structure=coordinates.system,
        result=result,
        stalled=stalled,
        model=model,
    )

TrustRegionOptions dataclass

Tuning for TrustRegion. The defaults are klasp's measured ones.

Attributes:

Name	Type	Description
eta	float	smallest reduction ratio that still accepts a step
rho_max	float	above this the model is so wrong the step is treated as failed
delta0	float	initial and post-restart radius, in Angstroms of displacement
delta_max	float	largest radius
delta_min	float	floor below which the radius has collapsed
energy_noise	float | None	smallest energy difference worth believing, in eV. None asks the calculator.
max_restarts	int	how many times to reset the model before pinning at the floor
grow	float	factor to grow the radius by on a good step at the boundary
shrink	float	factor to shrink by on a rejected step
shrink_hard	float	factor for a step that was not merely bad but unusable

Source code in chmpy/opt/trust_region.py

@dataclass
class TrustRegionOptions:
    """Tuning for `TrustRegion`. The defaults are klasp's measured ones.

    Attributes:
        eta: smallest reduction ratio that still accepts a step
        rho_max: above this the model is so wrong the step is treated as failed
        delta0: initial and post-restart radius, in Angstroms of displacement
        delta_max: largest radius
        delta_min: floor below which the radius has collapsed
        energy_noise: smallest energy difference worth believing, in eV. None
            asks the calculator.
        max_restarts: how many times to reset the model before pinning at the
            floor
        grow: factor to grow the radius by on a good step at the boundary
        shrink: factor to shrink by on a rejected step
        shrink_hard: factor for a step that was not merely bad but unusable
    """

    eta: float = 0.10
    rho_max: float = 10.0
    #: smallest believable energy difference, eV. None asks the calculator
    #: (`Calculator.energy_noise`), which is almost always what you want.
    energy_noise: float | None = None
    delta0: float = 0.20
    delta_max: float = 1.0
    delta_min: float = 1e-4
    max_restarts: int = 3
    grow: float = 2.0
    shrink: float = 0.5
    shrink_hard: float = 0.25