Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/
# Optimistix vs JAXopt for a small nonlinear least-squares problem.
#
# Each pairing of (nonlinear library, linear solver) is shown.
# We make exactly one step of the solver, to normalise across differences for how the
# trust region radius is handled, or differing termination conditions. This is performed
# on a laptop CPU, with Optimistix version 0.0.1 and JAXopt version 0.7.
#
# --------------------------------------------------------------------------------------
#
# Levenberg--Marquardt timing: taking exactly one step.
# ---------
# Optimistix (QR; default):
# Compile+run time: 8.616801428001054
# Run time: 0.002039344999502646
# ----
# Optimistix (normal CG):
# Compile+run time: 25.256638882001425
# Run time: 0.015216022002277896
# ----
# Optimistix (normal Cholesky):
# Compile+run time: 7.744616070001939
# Run time: 0.0015822620007384103
# ----
# JAXopt (normal CG; default):
# Compile+run time: 60.017054975000065
# Run time: 0.018481929000699893
# ----
# JAXopt (normal Cholesky):
# Compile+run time: 10.43396132700218
# Run time: 0.0017714580026222393

# --------------------------------------------------------------------------------------
#
# In conclusion, we see that Optimistix is consistently faster than JAXopt, at both
# compile time and runtime. It also has a more useful default.

import functools as ft
import timeit

import diffrax as dfx  # pyright: ignore
import equinox as eqx
import jax
import jax.numpy as jnp
import jax.scipy as jsp
import jaxopt  # pyright: ignore
import lineax as lx
from jaxtyping import Array, Float

import optimistix as optx


def vector_field(
    t, y: Float[Array, "2"], parameters: Float[Array, "4"]
) -> Float[Array, "2"]:
    """Lotka--Volterra equations."""
    prey, predator = y
    α, β, γ, δ = parameters
    d_prey = α * prey - β * prey * predator
    d_predator = -γ * predator + δ * prey * predator
    d_y = jnp.stack([d_prey, d_predator])
    return d_y


def solve(
    parameters: Float[Array, "4"], y0: Float[Array, "2"], saveat: dfx.SaveAt
) -> Float[Array, " ts"]:
    """Solve a single ODE."""
    term = dfx.ODETerm(vector_field)
    solver = dfx.Tsit5()
    t0 = saveat.subs.ts[0]
    t1 = saveat.subs.ts[-1]
    dt0 = 0.1
    sol = dfx.diffeqsolve(
        term,
        solver,
        t0,
        t1,
        dt0,
        y0,
        args=parameters,
        saveat=saveat,
        # support forward-mode autodiff, which is used by Levenberg--Marquardt
        adjoint=dfx.DirectAdjoint(),
    )
    return sol.ys


def get_data() -> tuple[Float[Array, "3 2"], Float[Array, "3 50"]]:
    """Simulate some training data."""
    # We consider three possible initial conditions.
    y0_a = jnp.array([9.0, 9.0])
    y0_b = jnp.array([10.0, 10.0])
    y0_c = jnp.array([11.0, 11.0])
    y0s = jnp.stack([y0_a, y0_b, y0_c])
    true_parameters = jnp.array([0.1, 0.02, 0.4, 0.02])
    saveat = dfx.SaveAt(ts=jnp.linspace(0, 30, 20))
    batch_solve = eqx.filter_jit(eqx.filter_vmap(solve, in_axes=(None, 0, None)))
    values = batch_solve(true_parameters, y0s, saveat)
    return y0s, values


def residuals(parameters, y0s__values):
    """The residuals, that we would like to minimise as a nonlinear least square."""
    y0s, values = y0s__values
    saveat = dfx.SaveAt(ts=jnp.linspace(0, 30, 20))
    batch_solve = eqx.filter_vmap(solve, in_axes=(None, 0, None))
    pred_values = batch_solve(parameters, y0s, saveat)
    return values - pred_values


# Lineax deliberately doesn't offer this as a solver.
# It's mostly just a worse version of QR. It runs at similar speed empirically (see the
# benchmarks below), but is substantially less accurate.
# Anyway, we have an implementation here to provide a fair comparison with JAXopt, which
# uses it as its default solver.
class NormalCholesky(lx.AbstractLinearSolver):
    def init(self, operator, options):
        del options
        matrix = operator.as_matrix()
        factor, lower = jsp.linalg.cho_factor(matrix.T @ matrix)
        assert lower is False
        packed_structures = lx.internal.pack_structures(operator)
        return matrix, factor, packed_structures

    def compute(self, state, vector, options):
        matrix, factor, packed_structures = state
        vector = lx.internal.ravel_vector(vector, packed_structures)
        solution = jsp.linalg.cho_solve((factor, False), matrix.T @ vector)
        return (
            lx.internal.unravel_solution(solution, packed_structures),
            lx.RESULTS.successful,
            {},
        )

    def transpose(self, state, options):
        assert False

    def allow_dependent_columns(self, operator):
        assert False

    def allow_dependent_rows(self, operator):
        assert False


# Default option for Optimistix. (QR linear solver.)
@jax.jit
def optx_qr(init_parameters, y0s, values):
    optx_solver = optx.LevenbergMarquardt(rtol=1e-8, atol=1e-8)
    sol = optx.least_squares(
        residuals,
        optx_solver,
        init_parameters,
        args=(y0s, values),
        max_steps=1,
        throw=False,
    )
    return sol.value


# Normal-CG linear solver. Chosen to match the JAXopt default below.
@jax.jit
def optx_normal_cg(init_parameters, y0s, values):
    # Explicitly set the number of iterations, to be sure we have the same number of
    # iterations between Optimistix and JAXopt.
    cg = lx.NormalCG(atol=0, rtol=0, max_steps=10)
    optx_solver = optx.LevenbergMarquardt(rtol=1e-8, atol=1e-8, linear_solver=cg)
    sol = optx.least_squares(
        residuals,
        optx_solver,
        init_parameters,
        args=(y0s, values),
        max_steps=1,
        throw=False,
    )
    return sol.value


# Normal-Cholesky linear solver. Chosen to match the fastest JAXopt option below (if you
# know what you're doing).
@jax.jit
def optx_normal_cholesky(init_parameters, y0s, values):
    optx_solver = optx.LevenbergMarquardt(
        rtol=1e-8, atol=1e-8, linear_solver=NormalCholesky()
    )
    sol = optx.least_squares(
        residuals,
        optx_solver,
        init_parameters,
        args=(y0s, values),
        max_steps=1,
        throw=False,
    )
    return sol.value


# Default options for JAXopt. (Nonlinear-CG linear solver.)
@jax.jit
def jaxopt_normal_cg(init_parameters, y0s, values):
    # JAXopt doesn't take extra parameters. Also, JAXopt requires a single vector as the
    # output -- multidimensional arrays and pytrees are not supported.
    def residual_fn(parameters):
        return residuals(parameters, (y0s, values)).reshape(-1)

    # Explicitly set the number of iterations, to be sure we have the same number of
    # iterations between Optimistix and JAXopt.
    cg = ft.partial(jaxopt.linear_solve.solve_cg, tol=0, maxiter=10)
    jaxopt_solver = jaxopt.LevenbergMarquardt(residual_fn, maxiter=1, solver=cg)
    return jaxopt_solver.run(init_parameters).params


# Normal-Cholesky linear solver. This is as fast as JAXopt seems to go. (If you know
# what you're doing.)
@jax.jit
def jaxopt_normal_cholesky(init_parameters, y0s, values):
    def residual_fn(parameters):
        return residuals(parameters, (y0s, values)).reshape(-1)

    jaxopt_solver = jaxopt.LevenbergMarquardt(
        residual_fn, solver="cholesky", materialize_jac=True, maxiter=1
    )
    return jaxopt_solver.run(init_parameters).params


init_parameters = jnp.zeros(4)
(y0s, values) = get_data()

run_optx_qr = lambda: optx_qr(init_parameters, y0s, values).block_until_ready()
run_optx_normal_cg = lambda: optx_normal_cg(
    init_parameters, y0s, values
).block_until_ready()
run_optx_normal_cholesky = lambda: optx_normal_cholesky(
    init_parameters, y0s, values
).block_until_ready()
run_jaxopt_normal_cg = lambda: jaxopt_normal_cg(
    init_parameters, y0s, values
).block_until_ready()
run_jaxopt_normal_cholesky = lambda: jaxopt_normal_cholesky(
    init_parameters, y0s, values
).block_until_ready()

optx_qr_compile_time = timeit.timeit(run_optx_qr, number=1)
optx_normal_cg_compile_time = timeit.timeit(run_optx_normal_cg, number=1)
optx_normal_cholesky_compile_time = timeit.timeit(run_optx_normal_cholesky, number=1)
jaxopt_normal_cg_compile_time = timeit.timeit(run_jaxopt_normal_cg, number=1)
jaxopt_normal_cholesky_compile_time = timeit.timeit(
    run_jaxopt_normal_cholesky, number=1
)

optx_qr_run_time = min(timeit.repeat(run_optx_qr, number=1, repeat=10))
optx_normal_cg_run_time = min(timeit.repeat(run_optx_normal_cg, number=1, repeat=10))
optx_normal_cholesky_run_time = min(
    timeit.repeat(run_optx_normal_cholesky, number=1, repeat=10)
)
jaxopt_normal_cg_run_time = min(
    timeit.repeat(run_jaxopt_normal_cg, number=1, repeat=10)
)
jaxopt_normal_cholesky_run_time = min(
    timeit.repeat(run_jaxopt_normal_cholesky, number=1, repeat=10)
)

# We do exactly one step so that we don't need to worry about their differing
# terimination conditions and other differences in the algorithm.
print("Levenberg--Marquardt timing: taking exactly one step.")
print("---------")
print("Optimistix (QR; default):")
print(f"Compile+run time: {optx_qr_compile_time}")
print(f"Run time: {optx_qr_run_time}")
print("----")
print("Optimistix (normal CG):")
print(f"Compile+run time: {optx_normal_cg_compile_time}")
print(f"Run time: {optx_normal_cg_run_time}")
print("----")
print("Optimistix (normal Cholesky):")
print(f"Compile+run time: {optx_normal_cholesky_compile_time}")
print(f"Run time: {optx_normal_cholesky_run_time}")
print("----")
print("JAXopt (normal CG; default):")
print(f"Compile+run time: {jaxopt_normal_cg_compile_time}")
print(f"Run time: {jaxopt_normal_cg_run_time}")
print("----")
print("JAXopt (normal Cholesky):")
print(f"Compile+run time: {jaxopt_normal_cholesky_compile_time}")
print(f"Run time: {jaxopt_normal_cholesky_run_time}")

benchmarks/levenberg-marquardt.py
patrick-kidger/optimistix
# In this comparison, we compare Optimistix against JAXopt for backpropagating through
# a vmap'd unrolled loop. JAXopt has serious performance issues in this case: its
# runtime grows with the maximum number of steps, even if these steps aren't actually
# used.
#
# Here's some typical printout (laptop CPU). Incidentally if we take it a step further,
# to a million steps, then JAXopt just crashes with an out-of-memory instead.
#
# Comparison of vmap + unrolled autodiff.
# ------------------------
# Optimistix:
# max_steps=10000 runtime=0.01200
# max_steps=100000 runtime=0.01127
# ---------
# JAXopt:
# max_steps=10000 runtime=0.29483
# max_steps=100000 runtime=3.15099

import timeit

import jax
import jax.numpy as jnp
import jax.random as jr
import jaxopt  # pyright: ignore

import optimistix as optx


A_sqrt = jr.normal(jr.PRNGKey(0), (100, 100))
A = A_sqrt.T @ A_sqrt  # now positive definite
x0 = jr.normal(jr.PRNGKey(1), (1, 100))
optx_times = {}
jaxopt_times = {}

for max_steps in (10_000, 100_000):

    @jax.jit
    @jax.vmap
    @jax.grad
    def run_optx(x0):
        def fn(x, _):
            return x.T @ A @ x

        solver = optx.GradientDescent(learning_rate=1e-3, rtol=1e-4, atol=1e-4)
        sol = optx.minimise(
            fn,
            solver,
            x0,
            adjoint=optx.RecursiveCheckpointAdjoint(),
            max_steps=max_steps,
        )
        return jnp.sum(sol.value)

    optx_times[max_steps] = min(
        timeit.repeat(lambda: run_optx(x0), number=1, repeat=10)
    )

    @jax.jit
    @jax.vmap
    @jax.grad
    def run_jaxopt(x0):
        def fn(x):
            return x.T @ A @ x

        solver = jaxopt.GradientDescent(
            fn,
            stepsize=1e-3,
            tol=1e-4,
            acceleration=False,
            maxiter=max_steps,
            implicit_diff=False,
        )
        return jnp.sum(solver.run(x0).params)

    jaxopt_times[max_steps] = min(
        timeit.repeat(lambda: run_jaxopt(x0), number=1, repeat=10)
    )

print("Comparison of vmap + unrolled autodiff.")
print("------------------------")
print("Optimistix:")
for k, v in optx_times.items():
    print(f"max_steps={k} runtime={v:.5f}")
print("---------")
print("JAXopt:")
for k, v in jaxopt_times.items():
    print(f"max_steps={k} runtime={v:.5f}")

benchmarks/vmap-unroll.py
patrick-kidger/optimistix
import matplotlib.pyplot as plt  # pyright: ignore
import numpy as np


def quadratic_bowl(x, y):
    return x**2 + 0.8 * y**2 - 0.3 * x * y


x_min, x_max = -10, 10
y_min, y_max = -10, 10
x_values, y_values = np.meshgrid(
    np.linspace(x_min, x_max, 100), np.linspace(y_min, y_max, 100)
)
z_values = quadratic_bowl(x_values, y_values)
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.plot_surface(x_values, y_values, z_values, cmap="coolwarm")
plt.xticks([-10, -5, 0, 5, 10])
plt.yticks([-10, -5, 0, 5, 10])
ax.set_zticks([0, 50, 100, 150, 200])
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([])
ax.set_xlabel("x", labelpad=-28)
ax.set_ylabel("y", labelpad=-28)
ax.set_zlabel("f(x, y)\n\n\n\n\n\n\n\n", labelpad=-30)
ax.zaxis.set_rotate_label(False)
ax.xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
ax.yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
ax.zaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
fig.set_facecolor("none")
plt.savefig("quadratic_bowl.png", transparent=True, bbox_inches="tight")

docs/_static/quadratic_bowl.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import importlib.metadata

from . import internal as internal
from ._adjoint import (
    AbstractAdjoint as AbstractAdjoint,
    ImplicitAdjoint as ImplicitAdjoint,
    RecursiveCheckpointAdjoint as RecursiveCheckpointAdjoint,
)
from ._fixed_point import (
    AbstractFixedPointSolver as AbstractFixedPointSolver,
    fixed_point as fixed_point,
)
from ._iterate import (
    AbstractIterativeSolver as AbstractIterativeSolver,
)
from ._least_squares import (
    AbstractLeastSquaresSolver as AbstractLeastSquaresSolver,
    least_squares as least_squares,
)
from ._minimise import (
    AbstractMinimiser as AbstractMinimiser,
    minimise as minimise,
)
from ._misc import (
    max_norm as max_norm,
    rms_norm as rms_norm,
    two_norm as two_norm,
)
from ._root_find import (
    AbstractRootFinder as AbstractRootFinder,
    root_find as root_find,
)
from ._search import (
    AbstractDescent as AbstractDescent,
    AbstractSearch as AbstractSearch,
    FunctionInfo as FunctionInfo,
)
from ._solution import RESULTS as RESULTS, Solution as Solution
from ._solver import (
    AbstractBFGS as AbstractBFGS,
    AbstractGaussNewton as AbstractGaussNewton,
    AbstractGradientDescent as AbstractGradientDescent,
    BacktrackingArmijo as BacktrackingArmijo,
    BestSoFarFixedPoint as BestSoFarFixedPoint,
    BestSoFarLeastSquares as BestSoFarLeastSquares,
    BestSoFarMinimiser as BestSoFarMinimiser,
    BestSoFarRootFinder as BestSoFarRootFinder,
    BFGS as BFGS,
    Bisection as Bisection,
    Chord as Chord,
    ClassicalTrustRegion as ClassicalTrustRegion,
    dai_yuan as dai_yuan,
    DampedNewtonDescent as DampedNewtonDescent,
    Dogleg as Dogleg,
    DoglegDescent as DoglegDescent,
    FixedPointIteration as FixedPointIteration,
    fletcher_reeves as fletcher_reeves,
    GaussNewton as GaussNewton,
    GradientDescent as GradientDescent,
    hestenes_stiefel as hestenes_stiefel,
    IndirectDampedNewtonDescent as IndirectDampedNewtonDescent,
    IndirectLevenbergMarquardt as IndirectLevenbergMarquardt,
    LearningRate as LearningRate,
    LevenbergMarquardt as LevenbergMarquardt,
    LinearTrustRegion as LinearTrustRegion,
    NelderMead as NelderMead,
    Newton as Newton,
    NewtonDescent as NewtonDescent,
    NonlinearCG as NonlinearCG,
    NonlinearCGDescent as NonlinearCGDescent,
    OptaxMinimiser as OptaxMinimiser,
    polak_ribiere as polak_ribiere,
    SteepestDescent as SteepestDescent,
)


__version__ = importlib.metadata.version("optimistix")

optimistix/__init__.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import types
from collections.abc import Callable
from typing import TypeVar

import equinox as eqx
import equinox.internal as eqxi
import jax
import jax.tree_util as jtu
import lineax as lx
from equinox.internal import ω
from jaxtyping import PyTree

from ._misc import tree_full_like


def _is_global_function(x):
    return isinstance(x, types.FunctionType) and x.__closure__ is None


_Inputs = TypeVar("_Inputs")
_Root = TypeVar("_Root")
_Residual = TypeVar("_Residual")


def implicit_jvp(
    fn_primal: Callable[[_Inputs], tuple[_Root, _Residual]],
    fn_rewrite: Callable[[_Root, _Residual, _Inputs], PyTree],
    inputs: _Inputs,
    tags: frozenset[object],
    linear_solver: lx.AbstractLinearSolver,
):
    """Rewrites gradients via the implicit function theorem.

    **Arguments:**

    - `fn_primal` is a function `inputs -> (root, residual)`.
    - `fn_rewrite` is a function `(root, residual, inputs) -> arbitrary`.
    - `inputs` is some input PyTree of the primal inputs to the computation.
    - `tags`: any Lineax tags (symmetric, diagonal, ...) for the matrix
        `d(fn_rewrite)/d(root)`.
    - `linear_solver`: an `lx.AbstractLinearSolver`, used to solve the linear problem
        on the backward pass.

    Note that due to limitations with JAX's custom autodiff, both `fn_primal` and
    `fn_rewrite` should be global functions (i.e. they should not capture any JAX array
    via closure, even if it does not participate in autodiff).

    **Returns:**

    This function returns `fn_primal(inputs)`. The first output is the output
    primal, whilst the second is auxiliary information.

    The primals have tangents `-(d(fn_rewrite)/d(root))^-1 d(fn_rewrite)/d(inputs)`,
    evaluated at `(root, residual, inputs)`.
    """
    assert _is_global_function(fn_primal)
    assert _is_global_function(fn_rewrite)
    root, residual = _implicit_impl(fn_primal, fn_rewrite, inputs, tags, linear_solver)
    return root, jtu.tree_map(eqxi.nondifferentiable_backward, residual)


@eqx.filter_custom_jvp
def _implicit_impl(fn_primal, fn_rewrite, inputs, tags, linear_solver):
    del fn_rewrite, tags, linear_solver
    return fn_primal(inputs)


def _assert_false(x):
    assert False


def _is_none(x):
    return x is None


def _for_jac(root, args):
    fn_rewrite, residual, _inputs = args
    iterate, inputs, while_loop = _inputs
    del iterate, while_loop
    return fn_rewrite(root, residual, inputs)


@_implicit_impl.def_jvp
def _implicit_impl_jvp(primals, tangents):
    fn_primal, fn_rewrite, inputs, tags, linear_solver = primals
    (
        t_fn_primal,
        t_fn_rewrite,
        t_inputs,
        t_tags,
        t_linear_solver,
    ) = tangents

    jtu.tree_map(_assert_false, (t_fn_primal, t_fn_rewrite, t_tags, t_linear_solver))
    del t_fn_primal, t_fn_rewrite, t_tags, t_linear_solver
    no_tangent = jtu.tree_map(_is_none, t_inputs, is_leaf=_is_none)
    nondiff, diff = eqx.partition(inputs, no_tangent, is_leaf=_is_none)

    root, residual = implicit_jvp(fn_primal, fn_rewrite, inputs, tags, linear_solver)

    def _for_jvp(_diff):
        _inputs = eqx.combine(_diff, nondiff)
        iterate, inputs, while_loop = _inputs
        del iterate, while_loop
        return fn_rewrite(root, residual, inputs)

    operator = lx.JacobianLinearOperator(
        _for_jac, root, (fn_rewrite, residual, inputs), tags=tags
    )
    _, jvp_diff = jax.jvp(_for_jvp, (diff,), (t_inputs,))

    t_root = (-lx.linear_solve(operator, jvp_diff, linear_solver).value ** ω).ω
    t_residual = tree_full_like(residual, 0)

    return (root, residual), (t_root, t_residual)

optimistix/_ad.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
import functools as ft
from collections.abc import Callable
from typing import Optional

import equinox as eqx
import equinox.internal as eqxi
import lineax as lx
from jaxtyping import Array, PyTree

from ._ad import implicit_jvp


class AbstractAdjoint(eqx.Module):
    """The abstract base class of all adjoints."""

    @abc.abstractmethod
    def apply(
        self,
        primal_fn: Callable,
        rewrite_fn: Callable,
        inputs: PyTree,
        tags: frozenset[object],
    ) -> PyTree[Array]:
        """Runs the main solver loop. Subclasses can override this to provide custom
        autodifferentiation behaviour; see for example the implementation of
        [`optimistix.ImplicitAdjoint`][].
        """


class RecursiveCheckpointAdjoint(AbstractAdjoint):
    """Backpropagate by differentiating through the iterates directly.

    Uses a binomial checkpointing scheme to keep memory usage low.

    !!! info

        Note that this cannot be forward-mode autodifferentiated. (E.g. using
        `jax.jvp`.)

    ??? cite "References"

        Selecting which steps at which to save checkpoints (and when this is done, which
        old checkpoint to evict) is important for minimising the amount of recomputation
        performed.

        The implementation here performs "online checkpointing", as the number of steps
        is not known in advance. This was developed in:

        ```bibtex
        @article{stumm2010new,
            author = {Stumm, Philipp and Walther, Andrea},
            title = {New Algorithms for Optimal Online Checkpointing},
            journal = {SIAM Journal on Scientific Computing},
            volume = {32},
            number = {2},
            pages = {836--854},
            year = {2010},
            doi = {10.1137/080742439},
        }

        @article{wang2009minimal,
            author = {Wang, Qiqi and Moin, Parviz and Iaccarino, Gianluca},
            title = {Minimal Repetition Dynamic Checkpointing Algorithm for Unsteady
                     Adjoint Calculation},
            journal = {SIAM Journal on Scientific Computing},
            volume = {31},
            number = {4},
            pages = {2549--2567},
            year = {2009},
            doi = {10.1137/080727890},
        }
        ```

        For reference, the classical "offline checkpointing" (also known as "treeverse",
        "recursive binary checkpointing", "revolve" etc.) was developed in:

        ```bibtex
        @article{griewank1992achieving,
            author = {Griewank, Andreas},
            title = {Achieving logarithmic growth of temporal and spatial complexity in
                     reverse automatic differentiation},
            journal = {Optimization Methods and Software},
            volume = {1},
            number = {1},
            pages = {35--54},
            year  = {1992},
            publisher = {Taylor & Francis},
            doi = {10.1080/10556789208805505},
        }

        @article{griewank2000revolve,
            author = {Griewank, Andreas and Walther, Andrea},
            title = {Algorithm 799: Revolve: An Implementation of Checkpointing for the
                     Reverse or Adjoint Mode of Computational Differentiation},
            year = {2000},
            publisher = {Association for Computing Machinery},
            volume = {26},
            number = {1},
            doi = {10.1145/347837.347846},
            journal = {ACM Trans. Math. Softw.},
            pages = {19--45},
        }
        ```
    """

    checkpoints: Optional[int] = None

    def apply(self, primal_fn, rewrite_fn, inputs, tags):
        del rewrite_fn, tags
        while_loop = ft.partial(
            eqxi.while_loop, kind="checkpointed", checkpoints=self.checkpoints
        )
        return primal_fn(inputs, while_loop)


def _primal_fn(inputs):
    primal_fn, inputs, while_loop = inputs
    return primal_fn(inputs, while_loop)


class ImplicitAdjoint(AbstractAdjoint):
    r"""Backpropagate via the [implicit function theorem](https://en.wikipedia.org/wiki/Implicit_function_theorem).

    For example, using the root-finding case by way of example: suppose we find the
    root `y(θ)` for which `f(y(θ), θ) = 0`.
    Then we can skip backpropagating through the solver by computing
    $\frac{\mathrm{d}y}{\mathrm{d}\theta} = - (\frac{\mathrm{d}f}{\mathrm{d}y})^{-1}\frac{\mathrm{d}f}{\mathrm{d}\theta}$
    via the implicit function theorem.

    For most problems this is the preferred technique for backpropogating through
    a nonlinear solve.
    """  # noqa: E501

    linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=None)

    def apply(self, primal_fn, rewrite_fn, inputs, tags):
        _inputs = (primal_fn, inputs, ft.partial(eqxi.while_loop, kind="lax"))
        return implicit_jvp(_primal_fn, rewrite_fn, _inputs, tags, self.linear_solver)


RecursiveCheckpointAdjoint.__init__.__doc__ = """**Arguments:**

- `checkpoints`: the number of checkpoints to save. The amount of memory used by the
    iterative solve will be roughly equal to the number of checkpoints
    multiplied by the size of `y0`. You can speed up backpropagation by allocating more
    checkpoints. (So it makes sense to set as many checkpoints as you have memory for.)
    This value can also be set to `None` (the default), in which case it will be set to
    `log(max_steps)`, for which a theoretical result is available guaranteeing that
    backpropagation will take `O(n log n)` time in the number of steps `n <= max_steps`.
"""

ImplicitAdjoint.__init__.__doc__ = """**Arguments:**

- `linear_solver`: the linear solver to solve the linear problem in the cotangent pass.
"""

optimistix/_adjoint.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, TypeVar, Union
from typing_extensions import TypeAlias

import equinox.internal as eqxi


Args: TypeAlias = Any
Aux = TypeVar("Aux")
Out = TypeVar("Out")
SolverState = TypeVar("SolverState")
SearchState = TypeVar("SearchState")
DescentState = TypeVar("DescentState")
Y = TypeVar("Y")

Fn: TypeAlias = Callable[[Y, Args], tuple[Out, Aux]]
NoAuxFn: TypeAlias = Callable[[Y, Args], Out]
MaybeAuxFn: TypeAlias = Union[Fn[Y, Out, Aux], NoAuxFn[Y, Out]]

sentinel: Any = eqxi.doc_repr(object(), "sentinel")

optimistix/_custom_types.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, cast, Generic, Optional, Union

import equinox as eqx
import jax
import jax.tree_util as jtu
from equinox.internal import ω
from jaxtyping import PyTree

from ._adjoint import AbstractAdjoint, ImplicitAdjoint
from ._custom_types import Args, Aux, Fn, MaybeAuxFn, SolverState, Y
from ._iterate import AbstractIterativeSolver, iterative_solve
from ._least_squares import AbstractLeastSquaresSolver
from ._minimise import AbstractMinimiser
from ._misc import inexact_asarray, NoneAux, OutAsArray
from ._root_find import AbstractRootFinder, root_find
from ._solution import Solution


class AbstractFixedPointSolver(AbstractIterativeSolver[Y, Y, Aux, SolverState]):
    """Abstract base class for all fixed point solvers."""


def _rewrite_fn(fixed_point, _, inputs):
    fixed_point_fn, _, _, args, *_ = inputs
    del inputs
    f_val, _ = fixed_point_fn(fixed_point, args)
    return (f_val**ω - fixed_point**ω).ω


class _ToRootFn(eqx.Module, Generic[Y, Aux]):
    fixed_point_fn: Fn[Y, Y, Aux]

    def __call__(self, y: Y, args: Args) -> tuple[Y, Aux]:
        out, aux = self.fixed_point_fn(y, args)
        if (
            eqx.tree_equal(jax.eval_shape(lambda: y), jax.eval_shape(lambda: out))
            is not True
        ):
            raise ValueError(
                "The input and output of `fixed_point_fn` must have the same structure"
            )
        return (out**ω - y**ω).ω, aux


@eqx.filter_jit
def fixed_point(
    fn: MaybeAuxFn[Y, Y, Aux],
    solver: Union[
        AbstractFixedPointSolver,
        AbstractRootFinder,
        AbstractLeastSquaresSolver,
        AbstractMinimiser,
    ],
    y0: Y,
    args: PyTree[Any] = None,
    options: Optional[dict[str, Any]] = None,
    *,
    has_aux: bool = False,
    max_steps: Optional[int] = 256,
    adjoint: AbstractAdjoint = ImplicitAdjoint(),
    throw: bool = True,
    tags: frozenset[object] = frozenset()
) -> Solution[Y, Aux]:
    """Find a fixed-point of a function.

    Given a nonlinear function `fn(y, args)` which returns a pytree of arrays of the
    same shape as `y`, this returns the value `z` such that `fn(z, args) = z`.

    **Arguments:**

    - `fn`: The function to find the fixed-point of. This should take two arguments
        `fn(y, args)`, and return a pytree of arrays of the same shape as the input `y`.
    - `solver`: The root-finder to use. This can be either an
        [`optimistix.AbstractFixedPointSolver`][] or
        [`optimistix.AbstractRootFinder`][], or
        [`optimistix.AbstractLeastSquaresSolver`][], or
        [`optimistix.AbstractMinimiser`][]. If `solver` is a root-finder then it will
        will attempt to find the root of `fn(y, args) - y`. If `solver` is a
        least-squares or minimisation algorithm, then it will attempt to minimise
        `sum((fn(y, args) - y)^2)`.
    - `y0`: An initial guess for what `y` may be. Used to start the iterative process of
        finding the fixed point; using good initial guesses is often important.
    - `args`: Passed as the `args` of `fn(y, args)`.
    - `options`: Individual solvers may accept additional runtime arguments.
        See each individual solver's documentation for more details.
    - `has_aux`: If `True`, then `fn` may return a pair, where the first element is its
        function value, and the second is just auxiliary data. Keyword only argument.
    - `max_steps`: The maximum number of steps the solver can take. Keyword only
        argument.
    - `adjoint`: The adjoint method used to compute gradients through the fixed-point
        solve. Keyword only argument.
    - `throw`: How to report any failures. (E.g. an iterative solver running out of
        steps, or encountering divergent iterates.) If `True` then a failure will raise
        an error. If `False` then the returned solution object will have a `result`
        field indicating whether any failures occured. (See [`optimistix.Solution`][].)
        Keyword only argument.
    - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the
        any structure of the Jacobian of `y -> fn(y, args) - y` with respect to y. (That
        is, the structure of the matrix `dfn/dy - I`.) Used with
        [`optimistix.ImplicitAdjoint`][] to implement the implicit function theorem as
        efficiently as possible. Keyword only argument.

    **Returns:**

    An [`optimistix.Solution`][] object.
    """

    if not has_aux:
        fn = NoneAux(fn)  # pyright: ignore
    fn = OutAsArray(fn)

    if isinstance(
        solver, (AbstractRootFinder, AbstractLeastSquaresSolver, AbstractMinimiser)
    ):
        del tags
        return root_find(
            _ToRootFn(fn),
            solver,
            y0,
            args,
            options,
            has_aux=True,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=throw,
        )
    else:
        y0 = jtu.tree_map(inexact_asarray, y0)
        fn = eqx.filter_closure_convert(fn, y0, args)  # pyright: ignore
        fn = cast(Fn[Y, Y, Aux], fn)
        f_struct, aux_struct = fn.out_struct
        if eqx.tree_equal(jax.eval_shape(lambda: y0), f_struct) is not True:
            raise ValueError(
                "The input and output of `fixed_point_fn` must have the same structure"
            )
        if options is None:
            options = {}
        return iterative_solve(
            fn,
            solver,
            y0,
            args,
            options,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=throw,
            tags=tags,
            f_struct=f_struct,
            aux_struct=aux_struct,
            rewrite_fn=_rewrite_fn,
        )

optimistix/_fixed_point.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
from collections.abc import Callable
from typing import Any, Generic, Optional, TYPE_CHECKING

import equinox as eqx
import equinox.internal as eqxi
import jax
import jax.core
import jax.numpy as jnp
import jax.tree_util as jtu
from equinox import AbstractVar
from jaxtyping import Array, Bool, PyTree, Scalar

from ._adjoint import AbstractAdjoint
from ._custom_types import Aux, Fn, Out, SolverState, Y
from ._solution import RESULTS, Solution


if TYPE_CHECKING:
    _Node = Any
else:
    _Node = eqxi.doc_repr(Any, "Node")


def _is_jaxpr(x):
    return isinstance(x, (jax.core.Jaxpr, jax.core.ClosedJaxpr))


def _is_array_or_jaxpr(x):
    return _is_jaxpr(x) or eqx.is_array(x)


class AbstractIterativeSolver(eqx.Module, Generic[Y, Out, Aux, SolverState]):
    """Abstract base class for all iterative solvers."""

    # Essentially every solver has an rtol+atol+norm. So for now we're just hardcoding
    # that every solver must have these variables, as they're needed when using a
    # minimiser or least-squares solver on a root-finding problem.
    rtol: AbstractVar[float]
    atol: AbstractVar[float]
    norm: AbstractVar[Callable[[PyTree], Scalar]]

    @abc.abstractmethod
    def init(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> SolverState:
        """Perform all initial computation needed to initialise the solver state.

        For example, the [`optimistix.Chord`][] method computes the Jacobian `df/dy`
        with respect to the initial guess `y`, and then uses it throughout the
        computation.

        **Arguments:**

        - `fn`: The function to iterate over. This is expected to take two argumetns
            `fn(y, args)` and return a pytree of arrays in the first element, and any
            auxiliary data in the second argument.
        - `y`: The value of `y` at the current (first) iteration.
        - `args`: Passed as the `args` of `fn(y, args)`.
        - `options`: Individual solvers may accept additional runtime arguments.
            See each individual solver's documentation for more details.
        - `f_struct`: A pytree of `jax.ShapeDtypeStruct`s of the same shape as the
            output of `fn`. This is used to initialise any information in the state
            which may rely on the pytree structure, array shapes, or dtype of the
            output of `fn`.
        - `aux_struct`: A pytree of `jax.ShapeDtypeStruct`s of the same shape as the
            auxiliary data returned by `fn`.
        - `tags`: exact meaning depends on whether this is a fixed point, root find,
            least squares, or minimisation problem; see their relevant entry points.

        **Returns:**

        A PyTree representing the initial state of the solver.
        """

    @abc.abstractmethod
    def step(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: SolverState,
        tags: frozenset[object],
    ) -> tuple[Y, SolverState, Aux]:
        """Perform one step of the iterative solve.

        **Arguments:**

        - `fn`: The function to iterate over. This is expected to take two argumetns
            `fn(y, args)` and return a pytree of arrays in the first element, and any
            auxiliary data in the second argument.
        - `y`: The value of `y` at the current (first) iteration.
        - `args`: Passed as the `args` of `fn(y, args)`.
        - `options`: Individual solvers may accept additional runtime arguments.
            See each individual solver's documentation for more details.
        - `state`: A pytree representing the state of a solver. The shape of this
            pytree is solver-dependent.
        - `tags`: exact meaning depends on whether this is a fixed point, root find,
            least squares, or minimisation problem; see their relevant entry points.

        **Returns:**

        A 3-tuple containing the new `y` value in the first element, the next solver
        state in the second element, and the aux output of `fn(y, args)` in the third
        element.
        """

    @abc.abstractmethod
    def terminate(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: SolverState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        """Determine whether or not to stop the iterative solve.

        **Arguments:**

        - `fn`: The function to iterate over. This is expected to take two argumetns
            `fn(y, args)` and return a pytree of arrays in the first element, and any
            auxiliary data in the second argument.
        - `y`: The value of `y` at the current iteration.
        - `args`: Passed as the `args` of `fn(y, args)`.
        - `options`: Individual solvers may accept additional runtime arguments.
            See each individual solver's documentation for more details.
        - `state`: A pytree representing the state of a solver. The shape of this
            pytree is solver-dependent.
        - `tags`: exact meaning depends on whether this is a fixed point, root find,
            least squares, or minimisation problem; see their relevant entry points.

        **Returns:**

        A 2-tuple containing a bool indicating whether or not to stop iterating in the
        first element, and an [`optimistix.RESULTS`][] object in the second element.
        """

    @abc.abstractmethod
    def postprocess(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: SolverState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        """Any final postprocessing to perform on the result of the solve.

        **Arguments:**

        - `fn`: The function to iterate over. This is expected to take two argumetns
            `fn(y, args)` and return a pytree of arrays in the first element, and any
            auxiliary data in the second argument.
        - `y`: The value of `y` at the last iteration.
        - `aux`: The auxiliary output at the last iteration.
        - `args`: Passed as the `args` of `fn(y, args)`.
        - `options`: Individual solvers may accept additional runtime arguments.
            See each individual solver's documentation for more details.
        - `state`: A pytree representing the final state of a solver. The shape of this
            pytree is solver-dependent.
        - `tags`: exact meaning depends on whether this is a fixed point, root find,
            least squares, or minimisation problem; see their relevant entry points.
        - `result`: as returned by the final call to `terminate`.

        **Returns:**

        A 3-tuple of:

        - `final_y`: the final `y` to return as the solution of the solve.
        - `final_aux`: the final `aux` to return as the auxiliary output of the solve.
        - `stats`: any additional information to place in the `sol.stats` dictionary.

        !!! info

            Most solvers will not need to use this, so that this method may be defined
            as:
            ```python
            def postprocess(self, fn, y, aux, args, options, state, tags, result):
                return y, aux, {}
            ```
        """


def _zero(x):
    if isinstance(x, jax.ShapeDtypeStruct):
        return jnp.zeros(x.shape, dtype=x.dtype)
    else:
        return x


def _iterate(inputs, while_loop):
    fn, solver, y0, args, options, max_steps, f_struct, aux_struct, tags = inputs
    del inputs
    static_leaf = lambda x: isinstance(x, eqxi.Static)
    f_struct = jtu.tree_map(lambda x: x.value, f_struct, is_leaf=static_leaf)
    aux_struct = jtu.tree_map(lambda x: x.value, aux_struct, is_leaf=static_leaf)
    init_aux = jtu.tree_map(_zero, aux_struct)
    init_state = solver.init(fn, y0, args, options, f_struct, aux_struct, tags)
    dynamic_init_state, static_state = eqx.partition(init_state, eqx.is_array)
    init_carry = (
        y0,
        jnp.array(0),
        dynamic_init_state,
        init_aux,
    )

    def cond_fun(carry):
        y, _, dynamic_state, _ = carry
        state = eqx.combine(static_state, dynamic_state)
        terminate, _ = solver.terminate(fn, y, args, options, state, tags)
        return jnp.invert(terminate)

    def body_fun(carry):
        y, num_steps, dynamic_state, _ = carry
        state = eqx.combine(static_state, dynamic_state)
        new_y, new_state, aux = solver.step(fn, y, args, options, state, tags)
        new_dynamic_state, new_static_state = eqx.partition(new_state, eqx.is_array)

        new_static_state_no_jaxpr = eqx.filter(
            new_static_state, _is_jaxpr, inverse=True
        )
        static_state_no_jaxpr = eqx.filter(state, _is_array_or_jaxpr, inverse=True)
        assert eqx.tree_equal(static_state_no_jaxpr, new_static_state_no_jaxpr) is True
        return new_y, num_steps + 1, new_dynamic_state, aux

    def buffers(carry):
        _, _, state, _ = carry
        return solver.buffers(state)

    final_carry = while_loop(cond_fun, body_fun, init_carry, max_steps=max_steps)
    final_y, num_steps, final_state, final_aux = final_carry
    _final_state = eqx.combine(static_state, final_state)
    terminate, result = solver.terminate(fn, final_y, args, options, _final_state, tags)
    result = RESULTS.where(
        (result == RESULTS.successful) & jnp.invert(terminate),
        RESULTS.nonlinear_max_steps_reached,
        result,
    )
    final_y, final_aux, stats = solver.postprocess(
        fn, final_y, final_aux, args, options, _final_state, tags, result
    )
    return final_y, (num_steps, result, final_state, final_aux, stats)


def iterative_solve(
    fn: Fn[Y, Out, Aux],
    # no type parameters, see https://github.com/microsoft/pyright/discussions/5599
    solver: AbstractIterativeSolver,
    y0: PyTree[Array],
    args: PyTree = None,
    options: Optional[dict[str, Any]] = None,
    *,
    max_steps: Optional[int],
    adjoint: AbstractAdjoint,
    throw: bool,
    tags: frozenset[object],
    f_struct: PyTree[jax.ShapeDtypeStruct],
    aux_struct: PyTree[jax.ShapeDtypeStruct],
    rewrite_fn: Callable,
) -> Solution[Y, Aux]:
    """Compute the iterates of an iterative numerical method.

    Given a nonlinear function `fn(y, args)` and an iterative method `solver`,
    this computes the iterates generated by `solver`. This generalises minimisation,
    least-squares, root-finding, and fixed-point iteration to any iterative
    numerical method applied to `fn(y, args)`.

    **Arguments:**

    - `fn`: The function to iterate over. This is expected to take two arguments
        `fn(y, args)` and return a pytree of arrays in the first element, and any
        auxiliary data in the second argument.
    - `solver`: The solver to use. This should be a subclass of
        [`optimistix.AbstractIterativeSolver`][].
    - `y0`: An initial guess for what `y` may be.
    - `args`: Passed as the `args` of `fn(y, args)`.
    - `options`: Individual solvers may accept additional runtime arguments.
        See each individual solver's documentation for more details.
    - `max_steps`: The maximum number of steps the solver can take. Keyword only
        argument.
    - `adjoint`: The adjoint method used to compute gradients through an iterative
        solve. Keyword only argument.
    - `throw`: How to report any failures. (E.g. an iterative solver running out of
        steps, or encountering divergent iterates.) If `True` then a failure will
        raise an error. If `False` then the returned solution object will have a
        `result` field indicating whether any failures occured. (See
        [`optimistix.Solution`][].) Keyword only argument.
    - `tags`: exact meaning depends on whether this is a fixed point, root find,
        least squares, or minimisation problem; see their relevant entry points.
    - `f_struct`: A pytree of `jax.ShapeDtypeStruct`s of the same shape as the
        output of `fn`. This is used to initialise any information in the state
        which may rely on the pytree structure, array shapes, or dtype of the
        output of `fn`. Keyword only argument.
    - `aux_struct`: A pytree of `jax.ShapeDtypeStruct`s of the same shape as the
        auxiliary data returned by `fn`. Keyword only argument.

    **Returns:**

    An [`optimistix.Solution`][] object.
    """

    f_struct = jtu.tree_map(eqxi.Static, f_struct)
    aux_struct = jtu.tree_map(eqxi.Static, aux_struct)
    inputs = fn, solver, y0, args, options, max_steps, f_struct, aux_struct, tags
    out, (num_steps, result, final_state, aux, stats) = adjoint.apply(
        _iterate, rewrite_fn, inputs, tags
    )
    stats = {"num_steps": num_steps, "max_steps": max_steps, **stats}
    sol = Solution(value=out, result=result, state=final_state, aux=aux, stats=stats)
    if throw:
        sol = result.error_if(sol, result != RESULTS.successful)
    return sol

optimistix/_iterate.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, cast, Generic, Optional, Union

import equinox as eqx
import jax
import jax.tree_util as jtu
from jaxtyping import PyTree, Scalar

from ._adjoint import AbstractAdjoint, ImplicitAdjoint
from ._custom_types import Args, Aux, Fn, MaybeAuxFn, Out, SolverState, Y
from ._iterate import AbstractIterativeSolver, iterative_solve
from ._minimise import AbstractMinimiser, minimise
from ._misc import inexact_asarray, NoneAux, OutAsArray, sum_squares
from ._solution import Solution


class AbstractLeastSquaresSolver(AbstractIterativeSolver[Y, Out, Aux, SolverState]):
    """Abstract base class for all least squares solvers."""


def _rewrite_fn(optimum, _, inputs):
    residual_fn, _, _, args, *_ = inputs
    del inputs

    def objective(_optimum):
        residual, _ = residual_fn(_optimum, args)
        return 0.5 * sum_squares(residual)

    return jax.grad(objective)(optimum)


class _ToMinimiseFn(eqx.Module, Generic[Y, Out, Aux]):
    residual_fn: Fn[Y, Out, Aux]

    def __call__(self, y: Y, args: Args) -> tuple[Scalar, Aux]:
        residual, aux = self.residual_fn(y, args)
        return 0.5 * sum_squares(residual), aux


@eqx.filter_jit
def least_squares(
    fn: MaybeAuxFn[Y, Out, Aux],
    # no type parameters, see https://github.com/microsoft/pyright/discussions/5599
    solver: Union[AbstractLeastSquaresSolver, AbstractMinimiser],
    y0: Y,
    args: PyTree[Any] = None,
    options: Optional[dict[str, Any]] = None,
    *,
    has_aux: bool = False,
    max_steps: Optional[int] = 256,
    adjoint: AbstractAdjoint = ImplicitAdjoint(),
    throw: bool = True,
    tags: frozenset[object] = frozenset(),
) -> Solution[Y, Aux]:
    r"""Solve a nonlinear least-squares problem.

    Given a nonlinear function `fn(y, args)` which returns a pytree of residuals,
    this returns the solution to $\min_y \sum_i \textrm{fn}(y, \textrm{args})_i^2$.

    **Arguments:**

    - `fn`: The residual function. This should take two arguments: `fn(y, args)` and
        return a pytree of arrays not necessarily of the same shape as the input `y`.
    - `solver`: The least-squares solver to use. This can be either an
        [`optimistix.AbstractLeastSquaresSolver`][] solver, or an
        [`optimistix.AbstractMinimiser`][]. If `solver` is an
        [`optimistix.AbstractMinimiser`][], then it will attempt to minimise the scalar
        loss $\min_y \sum_i \textrm{fn}(y, \textrm{args})_i^2$ directly.
    - `y0`: An initial guess for what `y` may be.
    - `args`: Passed as the `args` of `fn(y, args)`.
    - `options`: Individual solvers may accept additional runtime arguments.
        See each individual solver's documentation for more details.
    - `has_aux`: If `True`, then `fn` may return a pair, where the first element is its
        function value, and the second is just auxiliary data. Keyword only argument.
    - `max_steps`: The maximum number of steps the solver can take. Keyword only
        argument.
    - `adjoint`: The adjoint method used to compute gradients through the fixed-point
        solve. Keyword only argument.
    - `throw`: How to report any failures. (E.g. an iterative solver running out of
        steps, or encountering divergent iterates.) If `True` then a failure will raise
        an error. If `False` then the returned solution object will have a `result`
        field indicating whether any failures occured. (See [`optimistix.Solution`][].)
        Keyword only argument.
    - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the
        any structure of the Hessian of `y -> sum(fn(y, args)**2)` with respect to y.
        Used with [`optimistix.ImplicitAdjoint`][] to implement the implicit function
        theorem as efficiently as possible. Keyword only argument.

    **Returns:**

    An [`optimistix.Solution`][] object.
    """

    if not has_aux:
        fn = NoneAux(fn)  # pyright: ignore
    fn = OutAsArray(fn)

    if isinstance(solver, AbstractMinimiser):
        del tags
        return minimise(
            _ToMinimiseFn(fn),
            solver,
            y0,
            args,
            options,
            has_aux=True,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=throw,
        )
    else:
        y0 = jtu.tree_map(inexact_asarray, y0)
        fn = eqx.filter_closure_convert(fn, y0, args)  # pyright: ignore
        fn = cast(Fn[Y, Out, Aux], fn)
        f_struct, aux_struct = fn.out_struct
        if options is None:
            options = {}
        return iterative_solve(
            fn,
            solver,
            y0,
            args,
            options,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=throw,
            tags=tags,
            f_struct=f_struct,
            aux_struct=aux_struct,
            rewrite_fn=_rewrite_fn,
        )

optimistix/_least_squares.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, cast, Optional

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.tree_util as jtu
from jaxtyping import PyTree, Scalar

from ._adjoint import AbstractAdjoint, ImplicitAdjoint
from ._custom_types import Aux, Fn, MaybeAuxFn, SolverState, Y
from ._iterate import AbstractIterativeSolver, iterative_solve
from ._misc import inexact_asarray, NoneAux, OutAsArray
from ._solution import Solution


class AbstractMinimiser(AbstractIterativeSolver[Y, Scalar, Aux, SolverState]):
    """Abstract base class for all minimisers."""


def _rewrite_fn(minimum, _, inputs):
    minimise_fn, _, _, args, *_ = inputs
    del inputs

    def min_no_aux(x):
        f_val, _ = minimise_fn(x, args)
        return f_val

    return jax.grad(min_no_aux)(minimum)


@eqx.filter_jit
def minimise(
    fn: MaybeAuxFn[Y, Scalar, Aux],
    # no type parameters, see https://github.com/microsoft/pyright/discussions/5599
    solver: AbstractMinimiser,
    y0: Y,
    args: PyTree[Any] = None,
    options: Optional[dict[str, Any]] = None,
    *,
    has_aux: bool = False,
    max_steps: Optional[int] = 256,
    adjoint: AbstractAdjoint = ImplicitAdjoint(),
    throw: bool = True,
    tags: frozenset[object] = frozenset(),
) -> Solution[Y, Aux]:
    """Minimise a function.

    This minimises a nonlinear function `fn(y, args)` which returns a scalar value.

    **Arguments:**

    - `fn`: The objective function. This should take two arguments: `fn(y, args)` and
        return a scalar.
    - `solver`: The minimiser solver to use. This should be an
        [`optimistix.AbstractMinimiser`][].
    - `y0`: An initial guess for what `y` may be.
    - `args`: Passed as the `args` of `fn(y, args)`.
    - `options`: Individual solvers may accept additional runtime arguments.
        See each individual solver's documentation for more details.
    - `has_aux`: If `True`, then `fn` may return a pair, where the first element is its
        function value, and the second is just auxiliary data. Keyword only argument.
    - `max_steps`: The maximum number of steps the solver can take. Keyword only
        argument.
    - `adjoint`: The adjoint method used to compute gradients through the fixed-point
        solve. Keyword only argument.
    - `throw`: How to report any failures. (E.g. an iterative solver running out of
        steps, or encountering divergent iterates.) If `True` then a failure will raise
        an error. If `False` then the returned solution object will have a `result`
        field indicating whether any failures occured. (See [`optimistix.Solution`][].)
        Keyword only argument.
    - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the
        any structure of the Hessian of `fn` with respect to `y`. Used with
        [`optimistix.ImplicitAdjoint`][] to implement the implicit function theorem as
        efficiently as possible. Keyword only argument.

    **Returns:**

    An [`optimistix.Solution`][] object.
    """

    y0 = jtu.tree_map(inexact_asarray, y0)
    if not has_aux:
        fn = NoneAux(fn)  # pyright: ignore
    fn = OutAsArray(fn)
    fn = eqx.filter_closure_convert(fn, y0, args)  # pyright: ignore
    fn = cast(Fn[Y, Scalar, Aux], fn)
    f_struct, aux_struct = fn.out_struct
    if options is None:
        options = {}

    if not (
        isinstance(f_struct, jax.ShapeDtypeStruct)
        and f_struct.shape == ()
        and jnp.issubdtype(f_struct.dtype, jnp.floating)
    ):
        raise ValueError(
            "minimisation function must output a single floating-point scalar."
        )

    return iterative_solve(
        fn,
        solver,
        y0,
        args,
        options,
        max_steps=max_steps,
        adjoint=adjoint,
        throw=throw,
        tags=tags,
        aux_struct=aux_struct,
        f_struct=f_struct,
        rewrite_fn=_rewrite_fn,
    )

optimistix/_minimise.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools as ft
import math
from collections.abc import Callable
from typing import Any, Literal, overload, TYPE_CHECKING, TypeVar, Union

import equinox as eqx
import equinox.internal as eqxi
import jax
import jax.core
import jax.lax as lax
import jax.numpy as jnp
import jax.tree_util as jtu
from equinox.internal import ω
from jaxtyping import Array, ArrayLike, Bool, Inexact, PyTree, Scalar

from ._custom_types import Y


if TYPE_CHECKING:
    pass
else:
    pass


def default_floating_dtype():
    if jax.config.jax_enable_x64:  # pyright: ignore
        return jnp.float64
    else:
        return jnp.float32


@overload
def tree_full_like(
    struct: PyTree[Union[Array, jax.ShapeDtypeStruct]],
    fill_value: ArrayLike,
    allow_static: Literal[False] = False,
):
    ...


@overload
def tree_full_like(
    struct: PyTree, fill_value: ArrayLike, allow_static: Literal[True] = True
):

    ...


def tree_full_like(struct: PyTree, fill_value: ArrayLike, allow_static: bool = False):
    """Return a pytree with the same type and shape as the input with values
    `fill_value`.

    If `allow_static=True`, then any non-{array, struct}s are ignored and left alone.
    If `allow_static=False` then any non-{array, struct}s will result in an error.
    """
    fn = lambda x: jnp.full(x.shape, fill_value, x.dtype)
    if isinstance(fill_value, (int, float)):
        if fill_value == 0:
            fn = lambda x: jnp.zeros(x.shape, x.dtype)
        elif fill_value == 1:
            fn = lambda x: jnp.ones(x.shape, x.dtype)
    if allow_static:
        _fn = fn
        fn = (
            lambda x: _fn(x)
            if eqx.is_array(x) or isinstance(x, jax.ShapeDtypeStruct)
            else x
        )
    return jtu.tree_map(fn, struct)


def tree_dot(tree1: PyTree[ArrayLike], tree2: PyTree[ArrayLike]) -> Inexact[Array, ""]:
    """Compute the dot product of two pytrees of arrays with the same pytree
    structure."""
    leaves1, treedef1 = jtu.tree_flatten(tree1)
    leaves2, treedef2 = jtu.tree_flatten(tree2)
    if treedef1 != treedef2:
        raise ValueError("trees must have the same structure")
    assert len(leaves1) == len(leaves2)
    dots = []
    for leaf1, leaf2 in zip(leaves1, leaves2):
        dots.append(
            jnp.dot(
                jnp.reshape(leaf1, -1),
                jnp.conj(leaf2).reshape(-1),
                precision=jax.lax.Precision.HIGHEST,  # pyright: ignore
            )
        )
    if len(dots) == 0:
        return jnp.array(0, default_floating_dtype())
    else:
        return ft.reduce(jnp.add, dots)


def tree_where(
    pred: Bool[ArrayLike, ""], true: PyTree[ArrayLike], false: PyTree[ArrayLike]
) -> PyTree[Array]:
    """Return the `true` or `false` pytree depending on `pred`."""
    keep = lambda a, b: jnp.where(pred, a, b)
    return jtu.tree_map(keep, true, false)


def sum_squares(x: PyTree[ArrayLike]) -> Scalar:
    """Computes the square of the L2 norm of a PyTree of arrays.

    Considering the input `x` as a flat vector `(x_1, ..., x_n)`, then this computes
    `Σ_i x_i^2`
    """
    return tree_dot(x, x).real


@jax.custom_jvp
def two_norm(x: PyTree[ArrayLike]) -> Scalar:
    """Computes the L2 norm of a PyTree of arrays.

    Considering the input `x` as a flat vector `(x_1, ..., x_n)`, then this computes
    `sqrt(Σ_i x_i^2)`
    """
    leaves = jtu.tree_leaves(x)
    size = sum([jnp.size(xi) for xi in leaves])
    if size == 1:
        # Avoid needless squaring-and-then-rooting.
        for leaf in leaves:
            if jnp.size(leaf) == 1:
                return jnp.abs(jnp.reshape(leaf, ()))
        else:
            assert False
    else:
        return jnp.sqrt(sum_squares(x))


@two_norm.defjvp
def _two_norm_jvp(x, tx):
    (x,) = x
    (tx,) = tx
    out = two_norm(x)
    # Get zero gradient, rather than NaN gradient, in these cases.
    pred = (out == 0) | jnp.isinf(out)
    denominator = jnp.where(pred, 1, out)
    # We could also switch the dot and the division.
    # This approach is a bit more expensive (more divisions), but should be more
    # numerically stable (`x` and `denominator` should be of the same scale; `tx` is of
    # unknown scale).
    t_out = tree_dot((x**ω / denominator).ω, tx).real
    t_out = jnp.where(pred, 0, t_out)
    return out, t_out


def rms_norm(x: PyTree[ArrayLike]) -> Scalar:
    """Compute the RMS (root-mean-squared) norm of a PyTree of arrays.

    This is the same as the L2 norm, averaged by the size of the input `x`. Considering
    the input `x` as a flat vector `(x_1, ..., x_n)`, then this computes
    `sqrt((Σ_i x_i^2)/n)`
    """
    leaves = jtu.tree_leaves(x)
    size = sum([jnp.size(xi) for xi in leaves])
    return two_norm(x) / math.sqrt(size)


def max_norm(x: PyTree[ArrayLike]) -> Scalar:
    """Compute the L-infinity norm of a PyTree of arrays.

    This is the largest absolute elementwise value. Considering the input `x` as a flat
    vector `(x_1, ..., x_n)`, then this computes `max_i |x_i|`.
    """
    leaf_maxes = [jnp.max(jnp.abs(xi)) for xi in jtu.tree_leaves(x) if jnp.size(xi) > 0]
    if len(leaf_maxes) == 0:
        return jnp.array(0, default_floating_dtype())
    else:
        out = ft.reduce(jnp.maximum, leaf_maxes)
        return _zero_grad_at_zero(out)


@jax.custom_jvp
def _zero_grad_at_zero(x):
    return x


@_zero_grad_at_zero.defjvp
def _zero_grad_at_zero_jvp(primals, tangents):
    (out,) = primals
    (t_out,) = tangents
    return out, jnp.where(out == 0, 0, t_out)


def resolve_rcond(rcond, n, m, dtype):
    if rcond is None:
        return jnp.finfo(dtype).eps * max(n, m)
    else:
        return jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)


class NoneAux(eqx.Module):
    """Wrap a function `fn` so it returns a dummy aux value `None`

    NoneAux is used to give a consistent API between functions which have an aux
    and functions which do not, allowing us to avoid unnecessary aux handling.
    """

    fn: Callable

    def __call__(self, *args, **kwargs):
        return self.fn(*args, **kwargs), None


class OutAsArray(eqx.Module):
    """Wrap a minimisation/root-find/etc. function so that its mathematical outputs are
    all inexact arrays, and its auxiliary outputs are all arrays.
    """

    fn: Callable

    def __call__(self, *args, **kwargs):
        out, aux = self.fn(*args, **kwargs)
        out = jtu.tree_map(inexact_asarray, out)
        aux = jtu.tree_map(asarray, aux)
        return out, aux


def jacobian(fn, in_size, out_size, has_aux=False):
    """Compute the Jacobian of a function using forward or backward mode AD.

    `jacobian` chooses between forward and backwards autodiff depending on the input
    and output dimension of `fn`, as specified in `in_size` and `out_size`.
    """

    # Heuristic for which is better in each case
    # These could probably be tuned a lot more.
    if (in_size < 100) or (in_size <= 1.5 * out_size):
        return jax.jacfwd(fn, has_aux=has_aux)
    else:
        return jax.jacrev(fn, has_aux=has_aux)


def lin_to_grad(lin_fn, *primals):
    return jax.linear_transpose(lin_fn, *primals)(1.0)


def _asarray(dtype, x):
    return jnp.asarray(x, dtype=dtype)


# Work around JAX issue #15676
_asarray = jax.custom_jvp(_asarray, nondiff_argnums=(0,))


@_asarray.defjvp
def _asarray_jvp(dtype, x, tx):
    (x,) = x
    (tx,) = tx
    return _asarray(dtype, x), _asarray(dtype, tx)


def asarray(x):
    dtype = jnp.result_type(x)
    return _asarray(dtype, x)


def inexact_asarray(x):
    dtype = jnp.result_type(x)
    if not jnp.issubdtype(jnp.result_type(x), jnp.inexact):
        dtype = default_floating_dtype()
    return _asarray(dtype, x)


_F = TypeVar("_F")


def cauchy_termination(
    rtol: float,
    atol: float,
    norm: Callable[[PyTree], Scalar],
    y: Y,
    y_diff: Y,
    f: _F,
    f_diff: _F,
) -> Bool[Array, ""]:
    """Terminate if there is a small difference in both `y` space and `f` space, as
    determined by `rtol` and `atol`.

    Specifically, this checks that `y_diff < atol + rtol * y` and
    `f_diff < atol + rtol * f_prev`, terminating if both of these are true.
    """
    y_scale = (atol + rtol * ω(y).call(jnp.abs)).ω
    f_scale = (atol + rtol * ω(f).call(jnp.abs)).ω
    y_converged = norm((ω(y_diff).call(jnp.abs) / y_scale**ω).ω) < 1
    f_converged = norm((ω(f_diff).call(jnp.abs) / f_scale**ω).ω) < 1
    return y_converged & f_converged


class _JaxprEqual:
    def __init__(self, jaxpr: jax.core.Jaxpr):
        self.jaxpr = jaxpr

    def __hash__(self):
        return 0

    def __eq__(self, other):
        # Could also implement actual checks here.
        return type(self) is type(other)


def _wrap_jaxpr_impl(leaf):
    # But not jax.core.ClosedJaxpr, which contains constants that should be handled in
    # pytree style.
    if isinstance(leaf, jax.core.Jaxpr):
        return _JaxprEqual(leaf)
    else:
        return leaf


def _wrap_jaxpr(tree):
    return jtu.tree_map(_wrap_jaxpr_impl, tree)


def _unwrap_jaxpr_impl(leaf):
    if isinstance(leaf, _JaxprEqual):
        return leaf.jaxpr
    else:
        return leaf


def _unwrap_jaxpr(tree):
    return jtu.tree_map(_unwrap_jaxpr_impl, tree)


def filter_cond(pred, true_fun, false_fun, *operands):
    dynamic, static = eqx.partition(operands, eqx.is_array)

    def _true_fun(_dynamic):
        _operands = eqx.combine(_dynamic, static)
        _out = true_fun(*_operands)
        _dynamic_out, _static_out = eqx.partition(_out, eqx.is_array)
        _static_out = _wrap_jaxpr(_static_out)
        return _dynamic_out, eqxi.Static(_static_out)

    def _false_fun(_dynamic):
        _operands = eqx.combine(_dynamic, static)
        _out = false_fun(*_operands)
        _dynamic_out, _static_out = eqx.partition(_out, eqx.is_array)
        _static_out = _wrap_jaxpr(_static_out)
        return _dynamic_out, eqxi.Static(_static_out)

    dynamic_out, static_out = lax.cond(pred, _true_fun, _false_fun, dynamic)
    return eqx.combine(dynamic_out, _unwrap_jaxpr(static_out.value))


def verbose_print(*args: tuple[bool, str, Any]) -> None:
    string_pieces = []
    arg_pieces = []
    for display, name, value in args:
        if display:
            string_pieces.append(name + ": {}")
            arg_pieces.append(value)
    if len(string_pieces) > 0:
        string = ", ".join(string_pieces)
        jax.debug.print(string, *arg_pieces)

optimistix/_misc.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, cast, Optional, Union

import equinox as eqx
import jax.numpy as jnp
import jax.tree_util as jtu
from equinox import AbstractVar
from jaxtyping import PyTree

from ._adjoint import AbstractAdjoint, ImplicitAdjoint
from ._custom_types import Aux, Fn, MaybeAuxFn, Out, SolverState, Y
from ._iterate import AbstractIterativeSolver, iterative_solve
from ._least_squares import AbstractLeastSquaresSolver, least_squares
from ._minimise import AbstractMinimiser, minimise
from ._misc import inexact_asarray, NoneAux, OutAsArray, tree_full_like
from ._solution import Solution


class AbstractRootFinder(AbstractIterativeSolver[Y, Out, Aux, SolverState]):
    """Abstract base class for all root finders."""


def _rewrite_fn(root, _, inputs):
    root_fn, _, _, args, *_ = inputs
    del inputs
    f_val, _ = root_fn(root, args)
    return f_val


def _to_minimise_fn(root_fn, norm, y, args):
    root, aux = root_fn(y, args)
    return norm(root), (root, aux)


def _to_lstsq_fn(root_fn, y, args):
    root, aux = root_fn(y, args)
    return root, (root, aux)


# Adds an additional termination condition, that `fn(y, args)` be near zero.
class _ToRoot(AbstractIterativeSolver):
    solver: AbstractVar[AbstractIterativeSolver]

    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm

    def init(self, fn, y, args, options, f_struct, aux_struct, tags):
        orig_f_struct, _ = aux_struct
        init_state = self.solver.init(fn, y, args, options, f_struct, aux_struct, tags)
        f_inf = tree_full_like(orig_f_struct, jnp.inf)
        return (init_state, f_inf)

    def step(self, fn, y, args, options, state, tags):
        state, _ = state
        new_y, new_state, (f, aux) = self.solver.step(fn, y, args, options, state, tags)
        return new_y, (new_state, f), (f, aux)

    def terminate(self, fn, y, args, options, state, tags):
        state, f = state
        terminate, result = self.solver.terminate(fn, y, args, options, state, tags)
        # No rtol, because `rtol * 0 = 0`.
        near_zero = self.norm(f) < self.atol
        return terminate & near_zero, result

    def postprocess(self, fn, y, aux, args, options, state, tags, result):
        state, _ = state
        return self.solver.postprocess(fn, y, aux, args, options, state, tags, result)


class _MinimToRoot(AbstractMinimiser, _ToRoot):
    solver: AbstractMinimiser

    # Redeclare these three to work around the Equinox bug fixed here:
    # https://github.com/patrick-kidger/equinox/pull/544
    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm


class _LstsqToRoot(AbstractLeastSquaresSolver, _ToRoot):
    solver: AbstractLeastSquaresSolver

    # Redeclare these three to work around the Equinox bug fixed here:
    # https://github.com/patrick-kidger/equinox/pull/544
    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm


@eqx.filter_jit
def root_find(
    fn: MaybeAuxFn[Y, Out, Aux],
    # no type parameters, see https://github.com/microsoft/pyright/discussions/5599
    solver: Union[AbstractRootFinder, AbstractLeastSquaresSolver, AbstractMinimiser],
    y0: Y,
    args: PyTree = None,
    options: Optional[dict[str, Any]] = None,
    *,
    has_aux: bool = False,
    max_steps: Optional[int] = 256,
    adjoint: AbstractAdjoint = ImplicitAdjoint(),
    throw: bool = True,
    tags: frozenset[object] = frozenset(),
) -> Solution[Y, Aux]:
    """Solve a root-finding problem.

    Given a nonlinear function `fn(y, args)` which returns a pytree of arrays,
    this returns the value `z` such that `fn(z, args) = 0`.

    **Arguments:**

    - `fn`: The function to find the roots of. This should take two arguments:
        `fn(y, args)` and return a pytree of arrays not necessarily of the same shape
        as the input `y`.
    - `solver`: The root-finder to use. This should be an
        [`optimistix.AbstractRootFinder`][],
        [`optimistix.AbstractLeastSquaresSolver`][], or
        [`optimistix.AbstractMinimiser`][]. If it is a least-squares solver or a
        minimiser, then the value `sum(fn(y, args)^2)` is minimised.
    - `y0`: An initial guess for what `y` may be.
    - `args`: Passed as the `args` of `fn(y, args)`.
    - `options`: Individual solvers may accept additional runtime arguments.
        See each individual solver's documentation for more details.
    - `has_aux`: If `True`, then `fn` may return a pair, where the first element is its
        function value, and the second is just auxiliary data. Keyword only argument.
    - `max_steps`: The maximum number of steps the solver can take. Keyword only
        argument.
    - `adjoint`: The adjoint method used to compute gradients through the fixed-point
        solve. Keyword only argument.
    - `throw`: How to report any failures. (E.g. an iterative solver running out of
        steps, or encountering divergent iterates.) If `True` then a failure will raise
        an error. If `False` then the returned solution object will have a `result`
        field indicating whether any failures occured. (See [`optimistix.Solution`][].)
        Keyword only argument.
    - `tags`: Lineax [tags](https://docs.kidger.site/lineax/api/tags/) describing the
        any structure of the Jacobian of `fn` with respect to `y`. Used with some
        solvers (e.g. [`optimistix.Newton`][]), and with some adjoint methods (e.g.
        [`optimistix.ImplicitAdjoint`][]) to improve the efficiency of linear solves.
        Keyword only argument.

    **Returns:**

    An [`optimistix.Solution`][] object.
    """

    if not has_aux:
        fn = NoneAux(fn)  # pyright: ignore
    fn = OutAsArray(fn)

    if isinstance(solver, AbstractMinimiser):
        del tags
        sol = minimise(
            eqx.Partial(_to_minimise_fn, fn, solver.norm),
            _MinimToRoot(solver),  # pyright: ignore
            y0,
            args,
            options,
            has_aux=True,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=throw,
        )
        _, aux = sol.aux
        sol = eqx.tree_at(lambda s: s.aux, sol, aux)
        return sol
    elif isinstance(solver, AbstractLeastSquaresSolver):
        del tags
        return least_squares(
            eqx.Partial(_to_lstsq_fn, fn),
            _LstsqToRoot(solver),  # pyright: ignore
            y0,
            args,
            options,
            has_aux=True,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=throw,
        )
        _, aux = sol.aux
        sol = eqx.tree_at(lambda s: s.aux, sol, aux)
        return sol
    else:
        y0 = jtu.tree_map(inexact_asarray, y0)
        fn = eqx.filter_closure_convert(fn, y0, args)  # pyright: ignore
        fn = cast(Fn[Y, Out, Aux], fn)
        f_struct, aux_struct = fn.out_struct
        if options is None:
            options = {}
        return iterative_solve(
            fn,
            solver,
            y0,
            args,
            options,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=throw,
            tags=tags,
            f_struct=f_struct,
            aux_struct=aux_struct,
            rewrite_fn=_rewrite_fn,
        )

optimistix/_root_find.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""We use notions of "searches" and "descents" to generalise line searches, trust
regions, and learning rates.

For example, the classical gradient descent algorithm is given by setting the search
to be `optimistix.LearningRate`, and setting the descent to be
`optimistix.SteepestDescent`.

As another example, Levenberg--Marquardt is given by setting the search to be
`optimistix.ClassicalTrustRegion`, and setting the descent to be
`optimistix.DampedNewtonDescent`.

As a third example, BFGS is given by setting the search to be
`optimistix.BacktrackingArmijo`, and setting the descent to be
`optimistix.NewtonDescent`.

This gives us a flexible way to mix-and-match ideas across different optimisers.

Right now, these are used exclusively for minimisers and least-squares optimisers. (As
the latter are a special case of the former, with `to_minimise = 0.5 * residuals^2`)

For an in-depth discussion of how these pieces fit together, see the documentation at
`docs/api/searches/introduction.md`.
"""

import abc
from typing import ClassVar, Generic, Type, TypeVar

import equinox as eqx
import lineax as lx
from jaxtyping import Array, Bool, Scalar

from ._custom_types import (
    DescentState,
    Out,
    SearchState,
    Y,
)
from ._misc import sum_squares
from ._solution import RESULTS


class FunctionInfo(eqx.Module):
    """Different solvers (BFGS, Levenberg--Marquardt, ...) evaluate different quantities
    of the objective function. Some may compute gradient information, some may provide
    approximate Hessian information, etc.

    This enumeration-ish object captures the different variants.

    Available variants are
    `optimistix.FunctionInfo.{Eval, EvalGrad, EvalGradHessian, EvalGradHessianInv, Residual, ResidualJac}`.
    """  # noqa: E501

    Eval: ClassVar[Type["Eval"]]
    EvalGrad: ClassVar[Type["EvalGrad"]]
    EvalGradHessian: ClassVar[Type["EvalGradHessian"]]
    EvalGradHessianInv: ClassVar[Type["EvalGradHessianInv"]]
    Residual: ClassVar[Type["Residual"]]
    ResidualJac: ClassVar[Type["ResidualJac"]]

    @abc.abstractmethod
    def as_min(self) -> Scalar:
        """For a minimisation problem, returns f(y). For a least-squares problem,
        returns 0.5*residuals^2 -- i.e. its loss as a minimisation problem.
        """


# NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
class Eval(FunctionInfo):
    """Has a `.f` attribute describing `fn(y)`. Used when no gradient information is
    available.
    """

    f: Scalar

    def as_min(self):
        return self.f


# NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
class EvalGrad(FunctionInfo, Generic[Y]):
    """Has a `.f` attribute as with [`optimistix.FunctionInfo.Eval`][]. Also has a
    `.grad` attribute describing `d(fn)/dy`. Used with first-order solvers for
    minimisation problems. (E.g. gradient descent; nonlinear CG.)
    """

    f: Scalar
    grad: Y

    def as_min(self):
        return self.f


# NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
class EvalGradHessian(FunctionInfo, Generic[Y]):
    """Has `.f` and `.grad` attributes as with [`optimistix.FunctionInfo.EvalGrad`][].
    Also has a `.hessian` attribute describing (an approximation to) the Hessian of
    `fn` at `y`. Used with quasi-Newton minimisation algorithms, like BFGS.
    """

    f: Scalar
    grad: Y
    hessian: lx.AbstractLinearOperator

    def as_min(self):
        return self.f


# NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
class EvalGradHessianInv(FunctionInfo, Generic[Y]):
    """As [`optimistix.FunctionInfo.EvalGradHessian`][], but records the (approximate)
    inverse-Hessian instead. Has `.f` and `.grad` and `.hessian_inv` attributes.
    """

    f: Scalar
    grad: Y
    hessian_inv: lx.AbstractLinearOperator

    def as_min(self):
        return self.f


# NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
class Residual(FunctionInfo, Generic[Out]):
    """Has a `.residual` attribute describing `fn(y)`. Used with least squares problems,
    for which `fn` returns residuals.
    """

    residual: Out

    def as_min(self):
        return 0.5 * sum_squares(self.residual)


# NOT PUBLIC, despite lacking an underscore. This is so pyright gets the name right.
class ResidualJac(FunctionInfo, Generic[Y, Out]):
    """Records the Jacobian `d(fn)/dy` as a linear operator. Used for least squares
    problems, for which `fn` returns residuals. Has `.residual` and `.jac` and `.grad`
    attributes, where `residual = fn(y)`, `jac = d(fn)/dy` and
    `grad = jac^T residual`.

    Takes just `residual` and `jac` as `__init__`-time arguments, from which `grad` is
    computed.
    """

    residual: Out
    jac: lx.AbstractLinearOperator
    grad: Y

    def __init__(self, residual: Out, jac: lx.AbstractLinearOperator):
        self.residual = residual
        self.jac = jac
        # The gradient is used ubiquitously, so compute it once here, so that it can be
        # used without recomputation in both the descent and search.
        self.grad = jac.transpose().mv(residual)

    def as_min(self):
        return 0.5 * sum_squares(self.residual)


Eval.__qualname__ = "FunctionInfo.Eval"
EvalGrad.__qualname__ = "FunctionInfo.EvalGrad"
EvalGradHessian.__qualname__ = "FunctionInfo.EvalGradHessian"
EvalGradHessianInv.__qualname__ = "FunctionInfo.EvalGradHessianInv"
Residual.__qualname__ = "FunctionInfo.Residual"
ResidualJac.__qualname__ = "FunctionInfo.ResidualJac"
FunctionInfo.Eval = Eval
FunctionInfo.EvalGrad = EvalGrad
FunctionInfo.EvalGradHessian = EvalGradHessian
FunctionInfo.EvalGradHessianInv = EvalGradHessianInv
FunctionInfo.Residual = Residual
FunctionInfo.ResidualJac = ResidualJac


_FnInfo = TypeVar("_FnInfo", contravariant=True, bound=FunctionInfo)
_FnEvalInfo = TypeVar("_FnEvalInfo", contravariant=True, bound=FunctionInfo)


class AbstractDescent(eqx.Module, Generic[Y, _FnInfo, DescentState]):
    """The abstract base class for descents. A descent consumes a scalar (e.g. a step
    size), and returns the `diff` to take at point `y`, so that `y + diff` is the next
    iterate in a nonlinear optimisation problem.

    See [this documentation](./introduction.md) for more information.
    """

    @abc.abstractmethod
    def init(self, y: Y, f_info_struct: _FnInfo) -> DescentState:
        """Is called just once, at the very start of the entire optimisation problem.
        This is used to set up the evolving state for the descent.

        **Arguments:**

        - `y`: The initial guess for the optimisation problem.
        - `f_info`: An [`optimistix.FunctionInfo`][] describing information about `f`
            evaluated at `y`, and potentially the gradient of `f` at `y`, etc.

        **Returns:**

        The initial state, prior to doing any descents or searches.
        """

    @abc.abstractmethod
    def query(self, y: Y, f_info: _FnInfo, state: DescentState) -> DescentState:
        """Is called whenever the search decides to halt and accept its current iterate,
        and then query the descent for a new descent direction.

        This method can be used to precompute information that does not depend on the
        step size.

        For example, consider running a minimisation algorithm that performs multiple
        line searches. (E.g. nonlinear CG.) A single line search may reject many steps
        until it finds a location that it is happy with, and then accept that one. At
        that point, we must compute the direction for the next line search. That
        should be done in this method.

        **Arguments:**

        - `y`: the value of the current (just accepted) iterate.
        - `f_info`: An [`optimistix.FunctionInfo`][] describing information about `f`
            evaluated at `y`, the gradient of `f` at `y`, etc.
        - `state`: the evolving state of the repeated descents.

        **Returns:**

        The updated descent state.
        """

    @abc.abstractmethod
    def step(self, step_size: Scalar, state: DescentState) -> tuple[Y, RESULTS]:
        """Computes the descent of size `step_size`.

        **Arguments:**

        - `step_size`: a non-negative scalar describing the step size to take.
        - `state`: the evolving state of the repeated descents. This will typically
            convey most of the information about the problem, having been computed in
            [`optimistix.AbstractDescent.query`][].

        **Returns:**

        A 2-tuple of `(y_diff, result)`, describing the delta to apply in
        y-space, and any error if computing the update was not successful.
        """


class AbstractSearch(eqx.Module, Generic[Y, _FnInfo, _FnEvalInfo, SearchState]):
    """The abstract base class for all searches. (Which are our generalisation of
    line searches, trust regions, and learning rates.)

    See [this documentation](./introduction.md) for more information.
    """

    @abc.abstractmethod
    def init(self, y: Y, f_info_struct: _FnInfo) -> SearchState:
        """Is called just once, at the very start of the entire optimisation problem.
        This is used to set up the evolving state for the search.

        **Arguments:**

        - `y`: The initial guess for the optimisation problem.
        - `f_info`: An [`optimistix.FunctionInfo`][] describing information about the
            shape/dtype of `f` evaluated at `y`, and potentially the gradient of `f` at
            `y`, etc.

        **Returns:**

        The initial state, prior to doing any descents or searches.
        """

    @abc.abstractmethod
    def step(
        self,
        first_step: Bool[Array, ""],
        y: Y,
        y_eval: Y,
        f_info: _FnInfo,
        f_eval_info: _FnEvalInfo,
        state: SearchState,
    ) -> tuple[Scalar, Bool[Array, ""], RESULTS, SearchState]:
        """Performs a step within a search. For example, one step within a line search.

        **Arguments:**

        - `first_step`: True on the first step, then False afterwards. On the first
            step, then `y` and `f_info` will be meaningless dummy data, as no steps have
            been accepted yet. The search will normally need to special-case this case.
            On the first step, it is assumed that `accept=True` will be returned.
        - `y`: the value of the most recently "accepted" iterate, i.e. the last point at
            which the descent was `query`'d.
        - `y_eval`: the value of the most recent iterate within the search.
        - `y_diff`: equal to `y_eval - y`. Provided as a convenience, to save some
            computation, as this is commonly available from the calling solver.
        - `f_info`: An [`optimistix.FunctionInfo`][] describing information about `f`
            evaluated at `y`, the gradient of `f` at `y`, etc.
        - `f_eval_info`: An [`optimistix.FunctionInfo`][] describing information about
            `f` evaluated at `y`, the gradient of `f` at `y`, etc.
        - `state`: the evolving state of the repeated searches.

        **Returns:**

        A 4-tuple of `(step_size, accept, result, state)`, where:

        - `step_size`: the size of the next step to make. (Relative to `y`, not
            `y_eval`.) Will be passed to the descent.
        - `accept`: whether this step should be "accepted", in which case the descent
            will be queried for a new direction. Must be `True` if `first_step=True`.
        - `result`: an [`optimistix.RESULTS`][] object, for declaring any errors.
        - `state`: the updated state for the next search.
        """

optimistix/_search.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Any, Generic

import equinox as eqx
import lineax as lx
from jaxtyping import ArrayLike, PyTree

from ._custom_types import Aux, Y


# Extend `lineax.RESULTS` as we want to be able to use their error messages too.
class RESULTS(lx.RESULTS):  # pyright: ignore
    successful = ""
    nonlinear_max_steps_reached = (
        "The maximum number of steps was reached in the nonlinear solver. "
        "The problem may not be solveable (e.g., a root-find on a function that has no "
        "roots), or you may need to increase `max_steps`."
    )
    nonlinear_divergence = "Nonlinear solve diverged."


class Solution(eqx.Module, Generic[Y, Aux]):
    """The solution to a nonlinear solve.

    **Attributes:**

    - `value`: The solution to the solve.
    - `result`: An integer representing whether the solve was successful or not. This
        can be converted into a human-readable error message via
        `optimistix.RESULTS[result]`
    - `aux`: Any user-specified auxiliary data returned from the problem; defaults to
        `None` if there is no auxiliary data. Auxiliary outputs can be captured by
        setting a `has_aux=True` flag, e.g. `optx.root_find(fn, ..., has_aux=True)`.
    - `stats`: Statistics about the solve, e.g. the number of steps that were required.
    - `state`: The final internal state of the solver. The meaning of this is specific
        to each solver.
    """

    value: Y
    result: RESULTS
    aux: Aux
    stats: dict[str, PyTree[ArrayLike]]
    state: Any

optimistix/_solution.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .backtracking import BacktrackingArmijo as BacktrackingArmijo
from .best_so_far import (
    BestSoFarFixedPoint as BestSoFarFixedPoint,
    BestSoFarLeastSquares as BestSoFarLeastSquares,
    BestSoFarMinimiser as BestSoFarMinimiser,
    BestSoFarRootFinder as BestSoFarRootFinder,
)
from .bfgs import AbstractBFGS as AbstractBFGS, BFGS as BFGS
from .bisection import Bisection as Bisection
from .dogleg import Dogleg as Dogleg, DoglegDescent as DoglegDescent
from .fixed_point import FixedPointIteration as FixedPointIteration
from .gauss_newton import (
    AbstractGaussNewton as AbstractGaussNewton,
    GaussNewton as GaussNewton,
    NewtonDescent as NewtonDescent,
)
from .gradient_methods import (
    AbstractGradientDescent as AbstractGradientDescent,
    GradientDescent as GradientDescent,
    SteepestDescent as SteepestDescent,
)
from .learning_rate import LearningRate as LearningRate
from .levenberg_marquardt import (
    DampedNewtonDescent as DampedNewtonDescent,
    IndirectDampedNewtonDescent as IndirectDampedNewtonDescent,
    IndirectLevenbergMarquardt as IndirectLevenbergMarquardt,
    LevenbergMarquardt as LevenbergMarquardt,
)
from .nelder_mead import NelderMead as NelderMead
from .newton_chord import Chord as Chord, Newton as Newton
from .nonlinear_cg import (
    dai_yuan as dai_yuan,
    fletcher_reeves as fletcher_reeves,
    hestenes_stiefel as hestenes_stiefel,
    NonlinearCG as NonlinearCG,
    NonlinearCGDescent as NonlinearCGDescent,
    polak_ribiere as polak_ribiere,
)
from .optax import OptaxMinimiser as OptaxMinimiser
from .trust_region import (
    ClassicalTrustRegion as ClassicalTrustRegion,
    LinearTrustRegion as LinearTrustRegion,
)

optimistix/_solver/__init__.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import cast, Union
from typing_extensions import TypeAlias

import equinox as eqx
import jax.numpy as jnp
from equinox.internal import ω
from jaxtyping import Array, Bool, Scalar, ScalarLike

from .._custom_types import Y
from .._misc import (
    tree_dot,
)
from .._search import AbstractSearch, FunctionInfo
from .._solution import RESULTS


class _BacktrackingState(eqx.Module):
    step_size: Scalar


_FnInfo: TypeAlias = Union[
    FunctionInfo.EvalGrad,
    FunctionInfo.EvalGradHessian,
    FunctionInfo.EvalGradHessianInv,
    FunctionInfo.ResidualJac,
]
_FnEvalInfo: TypeAlias = FunctionInfo


class BacktrackingArmijo(AbstractSearch[Y, _FnInfo, _FnEvalInfo, _BacktrackingState]):
    """Perform a backtracking Armijo line search."""

    decrease_factor: ScalarLike = 0.5
    slope: ScalarLike = 0.1
    step_init: ScalarLike = 1.0

    def __post_init__(self):
        self.decrease_factor = eqx.error_if(
            self.decrease_factor,
            (self.decrease_factor <= 0)  # pyright: ignore
            | (self.decrease_factor >= 1),  # pyright: ignore
            "`BacktrackingArmoji(decrease_factor=...)` must be between 0 and 1.",
        )
        self.slope = eqx.error_if(
            self.slope,
            (self.slope <= 0) | (self.slope >= 1),  # pyright: ignore
            "`BacktrackingArmoji(slope=...)` must be between 0 and 1.",
        )
        self.step_init = eqx.error_if(
            self.step_init,
            self.step_init <= 0,  # pyright: ignore
            "`BacktrackingArmoji(step_init=...)` must be strictly greater than 0.",
        )

    def init(self, y: Y, f_info_struct: _FnInfo) -> _BacktrackingState:
        del f_info_struct
        return _BacktrackingState(step_size=jnp.array(self.step_init))

    def step(
        self,
        first_step: Bool[Array, ""],
        y: Y,
        y_eval: Y,
        f_info: _FnInfo,
        f_eval_info: _FnEvalInfo,
        state: _BacktrackingState,
    ) -> tuple[Scalar, Bool[Array, ""], RESULTS, _BacktrackingState]:

        if isinstance(
            f_info,
            (
                FunctionInfo.EvalGrad,
                FunctionInfo.EvalGradHessian,
                FunctionInfo.EvalGradHessianInv,
                FunctionInfo.ResidualJac,
            ),
        ):
            grad = f_info.grad
        else:
            raise ValueError(
                "Cannot use `BacktrackingArmijo` with this solver. This is because "
                "`BacktrackingArmijo` requires gradients of the target function, but "
                "this solver does not evaluate such gradients."
            )

        y_diff = (y_eval**ω - y**ω).ω
        predicted_reduction = tree_dot(grad, y_diff)
        # Terminate when the Armijo condition is satisfied. That is, `fn(y_eval)`
        # must do better than its linear approximation:
        # `fn(y_eval) < fn(y) + grad•y_diff`
        f_min = f_info.as_min()
        f_min_eval = f_eval_info.as_min()
        f_min_diff = f_min_eval - f_min  # This number is probably negative
        satisfies_armijo = f_min_diff <= self.slope * predicted_reduction
        has_reduction = predicted_reduction <= 0

        accept = first_step | (satisfies_armijo & has_reduction)
        step_size = jnp.where(
            accept, self.step_init, self.decrease_factor * state.step_size
        )
        step_size = cast(Scalar, step_size)
        return (
            step_size,
            accept,
            RESULTS.successful,
            _BacktrackingState(step_size=step_size),
        )


BacktrackingArmijo.__init__.__doc__ = """**Arguments:**

- `decrease_factor`: The rate at which to backtrack, i.e.
    `next_stepsize = decrease_factor * current_stepsize`. Must be between 0 and 1.
- `slope`: The slope of of the linear approximation to
    `f` that the backtracking algorithm must exceed to terminate. Larger
    means stricter termination criteria. Must be between 0 and 1.
- `step_init`: The first `step_size` the backtracking algorithm will
    try. Must be greater than 0.
"""

optimistix/_solver/backtracking.py
patrick-kidger/optimistix
import abc
from typing import Any, Generic

import equinox as eqx
import jax
import jax.numpy as jnp
from equinox import AbstractVar
from equinox.internal import ω
from jaxtyping import Array, Bool, PyTree, Scalar

from .._custom_types import Aux, Fn, Out, SolverState, Y
from .._fixed_point import AbstractFixedPointSolver
from .._iterate import AbstractIterativeSolver
from .._least_squares import AbstractLeastSquaresSolver
from .._minimise import AbstractMinimiser
from .._misc import sum_squares, tree_full_like, tree_where
from .._root_find import AbstractRootFinder
from .._solution import RESULTS


class _BestSoFarState(eqx.Module, Generic[Y, Aux, SolverState]):
    best_y: Y
    best_aux: Aux
    best_loss: Scalar
    state: SolverState


def _auxmented(fn, y, args):
    out, aux = fn(y, args)
    return out, (out, aux)


class _BestSoFarSolver(AbstractIterativeSolver, Generic[Y, Out, Aux]):
    solver: AbstractVar[AbstractIterativeSolver[Y, Out, tuple[Out, Aux], Any]]

    @abc.abstractmethod
    def _to_loss(self, y: Y, f: Out) -> Scalar:
        ...

    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm

    def init(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _BestSoFarState:
        aux = tree_full_like(aux_struct, 0)
        loss = jnp.array(jnp.inf)
        auxmented_fn = eqx.Partial(_auxmented, fn)
        state = self.solver.init(
            auxmented_fn, y, args, options, f_struct, (f_struct, aux_struct), tags
        )
        return _BestSoFarState(best_y=y, best_aux=aux, best_loss=loss, state=state)

    def step(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _BestSoFarState,
        tags: frozenset[object],
    ) -> tuple[Y, _BestSoFarState, Aux]:
        auxmented_fn = eqx.Partial(_auxmented, fn)
        new_y, new_state, (f, new_aux) = self.solver.step(
            auxmented_fn, y, args, options, state.state, tags
        )
        loss = self._to_loss(y, f)
        pred = loss < state.best_loss
        best_y = tree_where(pred, y, state.best_y)
        best_aux = tree_where(pred, new_aux, state.best_aux)
        best_loss = jnp.where(pred, loss, state.best_loss)
        new_state = _BestSoFarState(
            best_y=best_y, best_aux=best_aux, best_loss=best_loss, state=new_state
        )
        return new_y, new_state, new_aux

    def terminate(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _BestSoFarState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        auxmented_fn = eqx.Partial(_auxmented, fn)
        return self.solver.terminate(auxmented_fn, y, args, options, state.state, tags)

    def postprocess(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _BestSoFarState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        return state.best_y, state.best_aux, {}


class BestSoFarMinimiser(  # pyright: ignore
    _BestSoFarSolver[Y, Scalar, Aux], AbstractMinimiser[Y, Aux, _BestSoFarState]
):
    """Wraps another minimiser, to return the best-so-far value. That is, it makes a
    copy of the best `y` seen, and returns that.
    """

    solver: AbstractMinimiser[Y, tuple[Scalar, Aux], Any]

    # Explicitly declare to keep pyright happy.
    def __init__(self, solver: AbstractMinimiser[Y, tuple[Scalar, Aux], Any]):
        self.solver = solver

    def _to_loss(self, y: Y, f: Scalar) -> Scalar:
        return f

    # Redeclare these three to work around the Equinox bug fixed here:
    # https://github.com/patrick-kidger/equinox/pull/544
    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm


BestSoFarMinimiser.__init__.__doc__ = """**Arguments:**

- `solver`: the minimiser to wrap.  
"""


class BestSoFarLeastSquares(  # pyright: ignore
    _BestSoFarSolver[Y, Out, Aux],
    AbstractLeastSquaresSolver[Y, Out, Aux, _BestSoFarState],
):
    """Wraps another least-squares solver, to return the best-so-far value. That is, it
    makes a copy of the best `y` seen, and returns that.
    """

    solver: AbstractLeastSquaresSolver[Y, Out, tuple[Out, Aux], Any]

    # Explicitly declare to keep pyright happy.
    def __init__(
        self, solver: AbstractLeastSquaresSolver[Y, Out, tuple[Out, Aux], Any]
    ):
        self.solver = solver

    def _to_loss(self, y: Y, f: Out) -> Scalar:
        return sum_squares(f)

    # Redeclare these three to work around the Equinox bug fixed here:
    # https://github.com/patrick-kidger/equinox/pull/544
    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm


BestSoFarLeastSquares.__init__.__doc__ = """**Arguments:**

- `solver`: the least-squares solver to wrap.  
"""


class BestSoFarRootFinder(  # pyright: ignore
    _BestSoFarSolver[Y, Out, Aux], AbstractRootFinder[Y, Out, Aux, _BestSoFarState]
):
    """Wraps another root-finder, to return the best-so-far value. That is, it
    makes a copy of the best `y` seen, and returns that.
    """

    solver: AbstractRootFinder[Y, Out, tuple[Out, Aux], Any]

    # Explicitly declare to keep pyright happy.
    def __init__(self, solver: AbstractRootFinder[Y, Out, tuple[Out, Aux], Any]):
        self.solver = solver

    def _to_loss(self, y: Y, f: Out) -> Scalar:
        return sum_squares(f)

    # Redeclare these three to work around the Equinox bug fixed here:
    # https://github.com/patrick-kidger/equinox/pull/544
    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm


BestSoFarRootFinder.__init__.__doc__ = """**Arguments:**

- `solver`: the root-finder solver to wrap.  
"""


class BestSoFarFixedPoint(  # pyright: ignore
    _BestSoFarSolver[Y, Y, Aux], AbstractFixedPointSolver[Y, Aux, _BestSoFarState]
):
    """Wraps another fixed-point solver, to return the best-so-far value. That is, it
    makes a copy of the best `y` seen, and returns that.
    """

    solver: AbstractFixedPointSolver[Y, tuple[Y, Aux], Any]

    # Explicitly declare to keep pyright happy.
    def __init__(self, solver: AbstractFixedPointSolver[Y, tuple[Y, Aux], Any]):
        self.solver = solver

    def _to_loss(self, y: Y, f: Y) -> Scalar:
        return sum_squares((y**ω - f**ω).ω)

    # Redeclare these three to work around the Equinox bug fixed here:
    # https://github.com/patrick-kidger/equinox/pull/544
    @property  # pyright: ignore
    def rtol(self):
        return self.solver.rtol

    @property  # pyright: ignore
    def atol(self):
        return self.solver.atol

    @property  # pyright: ignore
    def norm(self):
        return self.solver.norm


BestSoFarFixedPoint.__init__.__doc__ = """**Arguments:**

- `solver`: the fixed-point solver to wrap.  
"""

optimistix/_solver/best_so_far.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, Generic, TypeVar

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.tree_util as jtu
import lineax as lx
from equinox import AbstractVar
from equinox.internal import ω
from jaxtyping import Array, Bool, PyTree, Scalar

from .._custom_types import Aux, DescentState, Fn, SearchState, Y
from .._minimise import AbstractMinimiser
from .._misc import (
    cauchy_termination,
    filter_cond,
    lin_to_grad,
    max_norm,
    tree_dot,
    tree_full_like,
)
from .._search import (
    AbstractDescent,
    AbstractSearch,
    FunctionInfo,
)
from .._solution import RESULTS
from .backtracking import BacktrackingArmijo
from .gauss_newton import NewtonDescent


def _identity_pytree(pytree: PyTree[Array]) -> lx.PyTreeLinearOperator:
    """Create an identity pytree `I` such that
    `pytree = lx.PyTreeLinearOperator(I).mv(pytree)`

    **Arguments**:

    - `pytree`: A pytree such that the output of `_identity_pytree` is the identity
        with respect to pytrees of the same shape as `pytree`.

    **Returns**:

    A `lx.PyTreeLinearOperator` with input and output shape the shape of `pytree`.
    """
    leaves, structure = jtu.tree_flatten(pytree)
    eye_structure = structure.compose(structure)
    eye_leaves = []
    for i1, l1 in enumerate(leaves):
        for i2, l2 in enumerate(leaves):
            if i1 == i2:
                eye_leaves.append(
                    jnp.eye(jnp.size(l1)).reshape(jnp.shape(l1) + jnp.shape(l2))
                )
            else:
                eye_leaves.append(jnp.zeros(jnp.shape(l1) + jnp.shape(l2)))

    # This has a Lineax positive_semidefinite tag. This is okay because the BFGS update
    # preserves positive-definiteness.
    return lx.PyTreeLinearOperator(
        jtu.tree_unflatten(eye_structure, eye_leaves),
        jax.eval_shape(lambda: pytree),
        lx.positive_semidefinite_tag,
    )


def _outer(tree1, tree2):
    def leaf_fn(x):
        return jtu.tree_map(lambda leaf: jnp.tensordot(x, leaf, axes=0), tree2)

    return jtu.tree_map(leaf_fn, tree1)


def _bfgs_update(f_eval, grad, prev_grad, hessian, hessian_inv, y_diff):
    grad_diff = (grad**ω - prev_grad**ω).ω
    inner = tree_dot(grad_diff, y_diff)

    def bfgs_update(hessian, hessian_inv):
        if hessian is None:
            assert hessian_inv is not None
            # Use Woodbury identity for rank-1 update of approximate Hessian.
            inv_mvp = hessian_inv.mv(grad_diff)
            mvp_inner = tree_dot(grad_diff, inv_mvp)
            diff_outer = _outer(y_diff, y_diff)
            mvp_outer = _outer(y_diff, inv_mvp)
            term1 = (((inner + mvp_inner) * (diff_outer**ω)) / (inner**2)).ω
            term2 = ((_outer(inv_mvp, y_diff) ** ω + mvp_outer**ω) / inner).ω
            new_hessian_inv = lx.PyTreeLinearOperator(
                (hessian_inv.pytree**ω + term1**ω - term2**ω).ω,
                output_structure=jax.eval_shape(lambda: prev_grad),
                tags=lx.positive_semidefinite_tag,
            )
            return None, new_hessian_inv
        else:
            assert hessian_inv is None
            # BFGS update to the operator directly
            mvp = hessian.mv(y_diff)
            term1 = (_outer(grad_diff, grad_diff) ** ω / inner).ω
            term2 = (_outer(mvp, mvp) ** ω / tree_dot(y_diff, mvp)).ω
            hessian = lx.PyTreeLinearOperator(
                (hessian.pytree**ω + term1**ω - term2**ω).ω,
                output_structure=jax.eval_shape(lambda: prev_grad),
                tags=lx.positive_semidefinite_tag,
            )
            return hessian, None

    def no_update(hessian, hessian_inv):
        return hessian, hessian_inv

    # In particular inner = 0 on the first step (as then state.grad=0), and so for
    # this we jump straight to the line search.
    # Likewise we get inner <= eps on convergence, and so again we make no update
    # to avoid a division by zero.
    inner_nonzero = inner > jnp.finfo(inner.dtype).eps
    hessian, hessian_inv = filter_cond(
        inner_nonzero, bfgs_update, no_update, hessian, hessian_inv
    )
    if hessian is None:
        return FunctionInfo.EvalGradHessianInv(f_eval, grad, hessian_inv)
    else:
        return FunctionInfo.EvalGradHessian(f_eval, grad, hessian)


_Hessian = TypeVar(
    "_Hessian", FunctionInfo.EvalGradHessian, FunctionInfo.EvalGradHessianInv
)


class _BFGSState(eqx.Module, Generic[Y, Aux, SearchState, DescentState, _Hessian]):
    # Updated every search step
    first_step: Bool[Array, ""]
    y_eval: Y
    search_state: SearchState
    # Updated after each descent step
    f_info: _Hessian
    aux: Aux
    descent_state: DescentState
    # Used for termination
    terminate: Bool[Array, ""]
    result: RESULTS


class AbstractBFGS(AbstractMinimiser[Y, Aux, _BFGSState], Generic[Y, Aux, _Hessian]):
    """Abstract BFGS (Broyden--Fletcher--Goldfarb--Shanno) minimisation algorithm.

    This is a quasi-Newton optimisation algorithm, whose defining feature is the way
    it progressively builds up a Hessian approximation using multiple steps of gradient
    information.

    This abstract version may be subclassed to choose alternative descent and searches.
    """

    rtol: AbstractVar[float]
    atol: AbstractVar[float]
    norm: AbstractVar[Callable[[PyTree], Scalar]]
    use_inverse: AbstractVar[bool]
    descent: AbstractVar[AbstractDescent[Y, _Hessian, Any]]
    search: AbstractVar[AbstractSearch[Y, _Hessian, FunctionInfo.Eval, Any]]

    def init(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: jax.ShapeDtypeStruct,
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _BFGSState:
        f = tree_full_like(f_struct, 0)
        grad = tree_full_like(y, 0)
        if self.use_inverse:
            hessian_inv = _identity_pytree(y)
            f_info = FunctionInfo.EvalGradHessianInv(f, grad, hessian_inv)
        else:
            hessian = _identity_pytree(y)
            f_info = FunctionInfo.EvalGradHessian(f, grad, hessian)
        f_info_struct = eqx.filter_eval_shape(lambda: f_info)
        return _BFGSState(
            first_step=jnp.array(True),
            y_eval=y,
            search_state=self.search.init(y, f_info_struct),
            f_info=f_info,
            aux=tree_full_like(aux_struct, 0),
            descent_state=self.descent.init(y, f_info_struct),
            terminate=jnp.array(False),
            result=RESULTS.successful,
        )

    def step(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _BFGSState,
        tags: frozenset[object],
    ) -> tuple[Y, _BFGSState, Aux]:
        f_eval, lin_fn, aux_eval = jax.linearize(
            lambda _y: fn(_y, args), state.y_eval, has_aux=True
        )
        step_size, accept, search_result, search_state = self.search.step(
            state.first_step,
            y,
            state.y_eval,
            state.f_info,
            FunctionInfo.Eval(f_eval),
            state.search_state,
        )

        def accepted(descent_state):
            (grad,) = lin_to_grad(lin_fn, y)
            y_diff = (state.y_eval**ω - y**ω).ω
            if self.use_inverse:
                hessian = None
                hessian_inv = state.f_info.hessian_inv
            else:
                hessian = state.f_info.hessian
                hessian_inv = None
            f_eval_info = _bfgs_update(
                f_eval, grad, state.f_info.grad, hessian, hessian_inv, y_diff
            )
            descent_state = self.descent.query(
                state.y_eval, f_eval_info, descent_state  # pyright: ignore
            )
            f_diff = (f_eval**ω - state.f_info.f**ω).ω
            terminate = cauchy_termination(
                self.rtol, self.atol, self.norm, state.y_eval, y_diff, f_eval, f_diff
            )
            return state.y_eval, f_eval_info, aux_eval, descent_state, terminate

        def rejected(descent_state):
            return y, state.f_info, state.aux, descent_state, jnp.array(False)

        y, f_info, aux, descent_state, terminate = filter_cond(
            accept, accepted, rejected, state.descent_state
        )

        y_descent, descent_result = self.descent.step(step_size, descent_state)
        y_eval = (y**ω + y_descent**ω).ω
        result = RESULTS.where(
            search_result == RESULTS.successful, descent_result, search_result
        )

        state = _BFGSState(
            first_step=jnp.array(False),
            y_eval=y_eval,
            search_state=search_state,
            f_info=f_info,
            aux=aux,
            descent_state=descent_state,
            terminate=terminate,
            result=result,
        )
        return y, state, aux

        def make_f_info_on_accept(f_eval, lin_fn):
            (grad,) = lin_to_grad(lin_fn, y)
            return f_info

    def terminate(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _BFGSState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        return state.terminate, state.result

    def postprocess(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _BFGSState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        return y, aux, {}


class BFGS(AbstractBFGS[Y, Aux, _Hessian]):
    """BFGS (Broyden--Fletcher--Goldfarb--Shanno) minimisation algorithm.

    This is a quasi-Newton optimisation algorithm, whose defining feature is the way
    it progressively builds up a Hessian approximation using multiple steps of gradient
    information.
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    use_inverse: bool
    descent: NewtonDescent
    search: BacktrackingArmijo

    def __init__(
        self,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree], Scalar] = max_norm,
        use_inverse: bool = True,
    ):
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.use_inverse = use_inverse
        self.descent = NewtonDescent(linear_solver=lx.Cholesky())
        # TODO(raderj): switch out `BacktrackingArmijo` with a better line search.
        self.search = BacktrackingArmijo()


BFGS.__init__.__doc__ = """**Arguments:**

- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
- `use_inverse`: The BFGS algorithm involves computing matrix-vector products of the
    form `B^{-1} g`, where `B` is an approximation to the Hessian of the function to be
    minimised. This means we can either (a) store the approximate Hessian `B`, and do a
    linear solve on every step, or (b) store the approximate Hessian inverse `B^{-1}`,
    and do a matrix-vector product on every step. Option (a) is generally cheaper for
    sparse Hessians (as the inverse may be dense). Option (b) is generally cheaper for
    dense Hessians (as matrix-vector products are cheaper than linear solves). The
    default is (b), denoted via `use_inverse=True`. Note that this is incompatible with
    line search methods like [`optimistix.ClassicalTrustRegion`][], which use the
    Hessian approximation `B` as part of their own computations.
"""

optimistix/_solver/bfgs.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, ClassVar, Literal, Union

import equinox as eqx
import jax
import jax.numpy as jnp
from jaxtyping import Array, Bool, Float, PyTree, Scalar

from .._custom_types import Aux, Fn
from .._root_find import AbstractRootFinder
from .._solution import RESULTS


class _BisectionState(eqx.Module):
    lower: Scalar
    upper: Scalar
    flip: Bool[Array, ""]
    error: Float[Array, ""]


class Bisection(AbstractRootFinder[Scalar, Scalar, Aux, _BisectionState]):
    """The bisection method of root finding. This may only be used with functions
    `R->R`, i.e. functions with scalar input and scalar output.

    This requires the following `options`:

    - `lower`: The lower bound on the interval which contains the root.
    - `upper`: The upper bound on the interval which contains the root.

    This algorithm works by considering the interval `[lower, upper]`, checking the
    sign of the evaluated function at the midpoint of the interval, and then keeping
    whichever half contains the root. This is then repeated. The iteration stops once
    the interval is sufficiently small.
    """

    rtol: float
    atol: float
    flip: Union[bool, Literal["detect"]] = "detect"
    # All norms are the same for scalars.
    norm: ClassVar[Callable[[PyTree], Scalar]] = jnp.abs

    def init(
        self,
        fn: Fn[Scalar, Scalar, Aux],
        y: Scalar,
        args: PyTree,
        options: dict[str, Any],
        f_struct: jax.ShapeDtypeStruct,
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _BisectionState:
        lower = jnp.asarray(options["lower"], f_struct.dtype)
        upper = jnp.asarray(options["upper"], f_struct.dtype)
        del options, aux_struct
        if jnp.shape(y) != () or jnp.shape(lower) != () or jnp.shape(upper) != ():
            raise ValueError(
                "Bisection can only be used to find the roots of a function taking a "
                "scalar input."
            )
        if not isinstance(f_struct, jax.ShapeDtypeStruct) or f_struct.shape != ():
            raise ValueError(
                "Bisection can only be used to find the roots of a function producing "
                "a scalar output."
            )
        if isinstance(self.flip, bool):
            # Make it possible to avoid the extra two function compilations.
            flip = jnp.array(self.flip)
        elif self.flip == "detect":
            lower_val, _ = fn(lower, args)
            upper_val, _ = fn(upper, args)
            lower_neg = lower_val < 0
            upper_neg = upper_val < 0
            flip = lower_val > upper_val
            flip = eqx.error_if(
                flip,
                lower_neg == upper_neg,
                msg="The root is not contained in [lower, upper]",
            )
        else:
            raise ValueError("`flip` may only be True, False, or 'detect'.")
        return _BisectionState(
            lower=lower,
            upper=upper,
            flip=flip,
            error=jnp.array(jnp.inf, f_struct.dtype),
        )

    def step(
        self,
        fn: Fn[Scalar, Scalar, Aux],
        y: Scalar,
        args: PyTree,
        options: dict[str, Any],
        state: _BisectionState,
        tags: frozenset[object],
    ) -> tuple[Scalar, _BisectionState, Aux]:
        del options
        error, aux = fn(y, args)
        negative = state.flip ^ (error < 0)
        new_lower = jnp.where(negative, y, state.lower)
        new_upper = jnp.where(negative, state.upper, y)
        new_y = new_lower + 0.5 * (new_upper - new_lower)
        new_state = _BisectionState(
            lower=new_lower, upper=new_upper, flip=state.flip, error=error
        )
        return new_y, new_state, aux

    def terminate(
        self,
        fn: Fn[Scalar, Scalar, Aux],
        y: Scalar,
        args: PyTree,
        options: dict[str, Any],
        state: _BisectionState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        del fn, args, options
        scale = self.atol + self.rtol * jnp.abs(y)
        y_small = jnp.abs(state.lower - state.upper) < scale
        f_small = jnp.abs(state.error) < self.atol
        return y_small & f_small, RESULTS.successful

    def postprocess(
        self,
        fn: Fn[Scalar, Scalar, Aux],
        y: Scalar,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _BisectionState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Scalar, Aux, dict[str, Any]]:
        return y, aux, {}


Bisection.__init__.__doc__ = """**Arguments:**

- `rtol`: Relative tolerance for terminating solve.
- `atol`: Absolute tolerance for terminating solve.
- `flip`: Can be set to any of:
    - `False`: specify that `fn(lower, args) < 0 < fn(upper, args)`.
    - `True`: specify that `fn(lower, args) > 0 > fn(upper, args)`.
    - `"detect"`: automatically check `fn(lower, args)` and `fn(upper, args)`. Note that
        this option may increase both runtime and compilation time.
"""

optimistix/_solver/bisection.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, Generic, Union

import equinox as eqx
import jax.lax as lax
import jax.numpy as jnp
import lineax as lx
from equinox.internal import ω
from jaxtyping import PyTree, Scalar

from .._custom_types import Aux, Out, Y
from .._misc import (
    max_norm,
    sum_squares,
    tree_dot,
    tree_full_like,
    tree_where,
    two_norm,
)
from .._root_find import AbstractRootFinder, root_find
from .._search import AbstractDescent, FunctionInfo
from .._solution import RESULTS
from .bisection import Bisection
from .gauss_newton import AbstractGaussNewton, newton_step
from .trust_region import ClassicalTrustRegion


class _DoglegDescentState(eqx.Module, Generic[Y]):
    newton: Y
    cauchy: Y
    newton_norm: Scalar
    cauchy_norm: Scalar
    result: RESULTS


class DoglegDescent(
    AbstractDescent[
        Y,
        Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
        _DoglegDescentState,
    ]
):
    """The Dogleg descent step, which switches between the Cauchy and the Newton
    descent directions.
    """

    linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=None)
    root_finder: AbstractRootFinder[Scalar, Scalar, None, Any] = Bisection(
        rtol=1e-3, atol=1e-3
    )
    trust_region_norm: Callable[[PyTree], Scalar] = two_norm

    def init(
        self,
        y: Y,
        f_info_struct: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
    ) -> _DoglegDescentState:
        # Dummy values; unused
        del f_info_struct
        return _DoglegDescentState(
            newton=y,
            cauchy=y,
            newton_norm=jnp.array(0.0),
            cauchy_norm=jnp.array(0.0),
            result=RESULTS.successful,
        )

    def query(
        self,
        y: Y,
        f_info: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
        state: _DoglegDescentState,
    ) -> _DoglegDescentState:
        del state
        # Compute `denom = grad^T Hess grad.`
        if isinstance(f_info, FunctionInfo.EvalGradHessian):
            denom = tree_dot(f_info.grad, f_info.hessian.mv(f_info.grad))
        elif isinstance(f_info, FunctionInfo.ResidualJac):
            # Use Gauss--Newton approximation `Hess ~ J^T J`
            denom = sum_squares(f_info.jac.mv(f_info.grad))
        else:
            raise ValueError(
                "`DoglegDescent` can only be used with least-squares solvers, or "
                "quasi-Newton minimisers which make approximations to the Hessian "
                "(like `optx.BFGS(use_inverse=False)`)"
            )
        denom_nonzero = denom > jnp.finfo(denom.dtype).eps
        safe_denom = jnp.where(denom_nonzero, denom, 1)
        # Compute `grad^T grad / (grad^T Hess grad)`
        scaling = jnp.where(denom_nonzero, sum_squares(f_info.grad) / safe_denom, 0.0)

        # Downhill towards the bottom of the quadratic basin.
        newton_sol, result = newton_step(f_info, self.linear_solver)
        newton = (-(newton_sol**ω)).ω
        newton_norm = self.trust_region_norm(newton_sol)

        # Downhill steepest descent.
        cauchy = (-scaling * f_info.grad**ω).ω
        cauchy_norm = self.trust_region_norm(cauchy)

        return _DoglegDescentState(
            newton=newton,
            cauchy=cauchy,
            newton_norm=newton_norm,
            cauchy_norm=cauchy_norm,
            result=result,
        )

    def step(self, step_size: Scalar, state: _DoglegDescentState) -> tuple[Y, RESULTS]:
        # For trust-region methods like `DoglegDescent`, the trust-region size directly
        # controls how large a step we take. This is actually somewhat annoying,
        # as a trust region algorithm has no understanding of the scale of a
        # problem unless initialised with a complex initialisation procedure.
        #
        # A simple, heuristic way around this is to scale the trust region `step_size`
        # so that a `step_size` of `1` corresponds to the full length of the Newton
        # step (anything greater than `1` will still accept the Newton step.)
        scaled_step_size = state.newton_norm * step_size

        def accept_scaled_cauchy(cauchy, newton):
            """Scale and return the Cauchy step."""
            norm_nonzero = state.cauchy_norm > jnp.finfo(state.cauchy_norm.dtype).eps
            safe_norm = jnp.where(norm_nonzero, state.cauchy_norm, 1)
            # Return zeros in degenerate case instead of inf because if `cauchy_norm` is
            # near zero, then so is the gradient and `delta` must be tiny to accept
            # this.
            normalised_cauchy = tree_where(
                norm_nonzero,
                ((cauchy**ω / safe_norm) * scaled_step_size).ω,
                tree_full_like(cauchy, 0),
            )
            return normalised_cauchy, RESULTS.successful

        def interpolate_cauchy_and_newton(cauchy, newton):
            """Find the point interpolating the Cauchy and Newton steps which
            intersects the trust region radius.
            """

            def interpolate(t):
                return (cauchy**ω + (t - 1) * (newton**ω - cauchy**ω)).ω

            # The vast majority of the time we expect users to use `two_norm`,
            # ie. the classic, elliptical trust region radius. In this case, we
            # compute the value of `t` to hit the trust region radius using by solving
            # a quadratic equation `a * t**2 + b * t + c = 0`
            # See section 4.1 of Nocedal & Wright "Numerical Optimization" for details.
            #
            # If they pass a norm other than `two_norm`, ie. they use a more exotic
            # trust region shape, we use a root find to approximately
            # find the value which hits the trust region radius.
            if self.trust_region_norm is two_norm:
                a = sum_squares((newton**ω - cauchy**ω).ω)
                inner_prod = tree_dot(cauchy, (newton**ω - cauchy**ω).ω)
                b = 2 * (inner_prod - a)
                c = state.cauchy_norm**2 - 2 * inner_prod + a - scaled_step_size**2
                quadratic_1 = jnp.clip(
                    0.5 * (-b + jnp.sqrt(b**2 - 4 * a * c)) / a, a_min=1, a_max=2
                )
                quadratic_2 = jnp.clip(
                    ((2 * c) / (-b - jnp.sqrt(b**2 - 4 * a * c))), a_min=1, a_max=2
                )
                # The quadratic formula is not numerically stable, and it is best to
                # use slightly different formulas when `b >=` and `b < 0`.
                # See https://github.com/fortran-lang/stdlib/issues/674 for a number of
                # references.
                t_interp = jnp.where(b >= 0, quadratic_1, quadratic_2)
                result = RESULTS.successful
            else:
                root_find_options = {"lower": jnp.array(1.0), "upper": jnp.array(2.0)}
                root_sol = root_find(
                    lambda t, _: (
                        self.trust_region_norm(interpolate(t)) - scaled_step_size
                    ),
                    self.root_finder,
                    y0=jnp.array(1.5),
                    options=root_find_options,
                    throw=False,
                )
                t_interp = root_sol.value
                result = root_sol.result
            return interpolate(t_interp), result

        def accept_newton(cauchy, newton):
            """Return the Newton step."""
            return newton, RESULTS.successful

        index = jnp.where(
            state.cauchy_norm > scaled_step_size, 0, jnp.where(step_size < 1, 1, 2)
        )
        y_diff, root_result = lax.switch(
            index,
            [accept_scaled_cauchy, interpolate_cauchy_and_newton, accept_newton],
            state.cauchy,
            state.newton,
        )
        result = RESULTS.where(
            state.result == RESULTS.successful, root_result, state.result
        )
        return y_diff, result


DoglegDescent.__init__.__doc__ = """**Arguments:**

- `linear_solver`: The linear solver used to compute the Newton step.
- `root_finder`: The root finder used to find the point where the trust-region
    intersects the dogleg path. This is ignored if
    `trust_region_norm=optimistix.two_norm`, for which there is an analytic formula
    instead.
- `trust_region_norm`: The norm used to determine the trust-region shape.
"""


class Dogleg(AbstractGaussNewton[Y, Out, Aux]):
    """Dogleg algorithm. Used for nonlinear least squares problems.

    Given a quadratic bowl that locally approximates the function to be minimised, then
    there are two different ways we might try to move downhill: in the steepest descent
    direction (as in gradient descent; this is also sometimes called the Cauchy
    direction), and in the direction of the minima of the quadratic bowl (as in Newton's
    method; correspondingly this is called the Newton direction).

    The distinguishing feature of this algorithm is the "dog leg" shape of its descent
    path, in which it begins by moving in the steepest descent direction, and then
    switches to moving in the Newton direction.
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    descent: DoglegDescent[Y]
    search: ClassicalTrustRegion[Y]
    verbose: frozenset[str]

    def __init__(
        self,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree], Scalar] = max_norm,
        linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=None),
        verbose: frozenset[str] = frozenset(),
    ):
        # We don't expose root_finder to the default API for Dogleg because
        # we assume the `trust_region_norm` norm is `two_norm`, which has
        # an analytic formula for the intersection with the dogleg path.
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.descent = DoglegDescent(linear_solver=linear_solver)
        self.search = ClassicalTrustRegion()
        self.verbose = verbose


Dogleg.__init__.__doc__ = """**Arguments:**

- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
- `linear_solver`: The linear solver used to compute the Newton part of the dogleg step.
- `verbose`: Whether to print out extra information about how the solve is proceeding.
    Should be a frozenset of strings, specifying what information to print out. Valid
    entries are `step`, `loss`, `accepted`, `step_size`, `y`. For example
    `verbose=frozenset({"loss", "step_size"})`.
"""

optimistix/_solver/dogleg.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any

import equinox as eqx
import jax
import jax.numpy as jnp
from equinox.internal import ω
from jaxtyping import Array, Bool, PyTree, Scalar

from .._custom_types import Aux, Fn, Y
from .._fixed_point import AbstractFixedPointSolver
from .._misc import max_norm
from .._solution import RESULTS


class _FixedPointState(eqx.Module):
    relative_error: Scalar


class FixedPointIteration(AbstractFixedPointSolver[Y, Aux, _FixedPointState]):
    """Repeatedly calls a function in search of a fixed point.

    This is one of the simplest ways to find a fixed point `y` of `f`: simply
    repeatedly call `y_{n+1}=f(y_n)` until `y_n` stops changing.

    Note that this is often not a very effective method, and root-finding algorithms are
    frequently preferred in practice.
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar] = max_norm

    def init(
        self,
        fn: Fn[Y, Y, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _FixedPointState:
        del fn, y, args, options, f_struct, aux_struct
        return _FixedPointState(jnp.array(jnp.inf))

    def step(
        self,
        fn: Fn[Y, Y, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _FixedPointState,
        tags: frozenset[object],
    ) -> tuple[Y, _FixedPointState, Aux]:
        new_y, aux = fn(y, args)
        error = (y**ω - new_y**ω).ω
        scale = (self.atol + self.rtol * ω(new_y).call(jnp.abs)).ω
        new_state = _FixedPointState(self.norm((error**ω / scale**ω).ω))
        return new_y, new_state, aux

    def terminate(
        self,
        fn: Fn[Y, Y, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _FixedPointState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        return state.relative_error < 1, RESULTS.successful

    def postprocess(
        self,
        fn: Fn[Y, Y, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _FixedPointState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        return y, aux, {}


FixedPointIteration.__init__.__doc__ = """**Arguments:**

- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
"""

optimistix/_solver/fixed_point.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, Generic, Optional, Union

import equinox as eqx
import jax
import jax.numpy as jnp
import lineax as lx
from equinox import AbstractVar
from equinox.internal import ω
from jaxtyping import Array, Bool, Int, PyTree, Scalar

from .._custom_types import Args, Aux, DescentState, Fn, Out, SearchState, Y
from .._least_squares import AbstractLeastSquaresSolver
from .._misc import (
    cauchy_termination,
    filter_cond,
    max_norm,
    sum_squares,
    tree_full_like,
    verbose_print,
)
from .._search import (
    AbstractDescent,
    AbstractSearch,
    FunctionInfo,
)
from .._solution import RESULTS
from .learning_rate import LearningRate


def newton_step(
    f_info: Union[
        FunctionInfo.EvalGradHessian,
        FunctionInfo.EvalGradHessianInv,
        FunctionInfo.ResidualJac,
    ],
    linear_solver: lx.AbstractLinearSolver,
) -> tuple[PyTree[Array], RESULTS]:
    """Compute a Newton step.

    For a minimisation problem, this means computing `Hess^{-1} grad`.

    For a least-squares problem, we convert to a minimisation problem via
    `0.5*residuals^2`, which then implies `Hess^{-1} ~ J^T J` (Gauss--Newton
    approximation) and `grad = J^T residuals`.

    Thus `Hess^{-1} grad ~ (J^T J)^{-1} J^T residuals`.   [Equation A]

    Now if `J` is well-posed then this equals `J^{-1} residuals`, which is exactly what
    we compute here.

    And if `J` is ill-posed then [Equation A] is just the normal equations, which should
    almost never be treated directly! (Squares the condition number blahblahblah.) The
    solution of the normal equations matches the pseudoinverse solution `J^{dagger}`
    residuals, which is what we get using an ill-posed-capable linear solver (typically
    QR). So we solve the same linear system as in the well-posed case, we just need to
    set a different linear solver. (Which happens with
    `linear_solver=lx.AutoLinearSolver(well_posed=None)`, which is the recommended
    value.)
    """
    if isinstance(f_info, FunctionInfo.EvalGradHessianInv):
        newton = f_info.hessian_inv.mv(f_info.grad)
        result = RESULTS.successful
    else:
        if isinstance(f_info, FunctionInfo.EvalGradHessian):
            operator = f_info.hessian
            vector = f_info.grad
        elif isinstance(f_info, FunctionInfo.ResidualJac):
            operator = f_info.jac
            vector = f_info.residual
        else:
            raise ValueError(
                "Cannot use a Newton descent with a solver that only evaluates the "
                "gradient, or only the function itself."
            )
        out = lx.linear_solve(operator, vector, linear_solver)
        newton = out.value
        result = RESULTS.promote(out.result)
    return newton, result


class _NewtonDescentState(eqx.Module, Generic[Y]):
    newton: Y
    result: RESULTS


class NewtonDescent(
    AbstractDescent[
        Y,
        Union[
            FunctionInfo.EvalGradHessian,
            FunctionInfo.EvalGradHessianInv,
            FunctionInfo.ResidualJac,
        ],
        _NewtonDescentState,
    ]
):
    """Newton descent direction.

    Given a quadratic bowl `x -> x^T Hess x` -- a local quadratic approximation
    to the target function -- this corresponds to moving in the direction of the bottom
    of the bowl. (Which is *not* the same as steepest descent.)

    This is done by solving a linear system of the form `Hess^{-1} grad`.
    """

    norm: Optional[Callable[[PyTree], Scalar]] = None
    linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=None)

    def init(self, y: Y, f_info_struct: FunctionInfo) -> _NewtonDescentState:
        del f_info_struct
        # Dummy values of the right shape; unused.
        return _NewtonDescentState(y, RESULTS.successful)

    def query(
        self,
        y: Y,
        f_info: Union[
            FunctionInfo.EvalGradHessian,
            FunctionInfo.EvalGradHessianInv,
            FunctionInfo.ResidualJac,
        ],
        state: _NewtonDescentState,
    ) -> _NewtonDescentState:
        del state
        newton, result = newton_step(f_info, self.linear_solver)
        if self.norm is not None:
            newton = (newton**ω / self.norm(newton)).ω
        return _NewtonDescentState(newton, result)

    def step(self, step_size: Scalar, state: _NewtonDescentState) -> tuple[Y, RESULTS]:
        return (-step_size * state.newton**ω).ω, state.result


NewtonDescent.__init__.__doc__ = """**Arguments:**

- `norm`: If passed, then normalise the gradient using this norm. (The returned step
    will have length `step_size` with respect to this norm.) Optimistix includes three
    built-in norms: [`optimistix.max_norm`][], [`optimistix.rms_norm`][], and
    [`optimistix.two_norm`][].
- `linear_solver`: The linear solver used to compute the Newton step.
"""


class _GaussNewtonState(eqx.Module, Generic[Y, Out, Aux, SearchState, DescentState]):
    # Updated every search step
    first_step: Bool[Array, ""]
    y_eval: Y
    search_state: SearchState
    # Updated after each descent step
    f_info: FunctionInfo.ResidualJac
    aux: Aux
    descent_state: DescentState
    # Used for termination
    terminate: Bool[Array, ""]
    result: RESULTS
    # Use for verbose logging
    num_steps: Int[Array, ""]
    num_accepted_steps: Int[Array, ""]
    num_steps_since_acceptance: Int[Array, ""]


def _make_f_info(
    fn: Callable[[Y, Args], tuple[Any, Aux]], y: Y, args: Args, tags: frozenset
) -> tuple[FunctionInfo.ResidualJac, Aux]:
    f_eval, lin_fn, aux_eval = jax.linearize(lambda _y: fn(_y, args), y, has_aux=True)
    jac_eval = lx.FunctionLinearOperator(lin_fn, jax.eval_shape(lambda: y), tags)
    return FunctionInfo.ResidualJac(f_eval, jac_eval), aux_eval


class AbstractGaussNewton(AbstractLeastSquaresSolver[Y, Out, Aux, _GaussNewtonState]):
    """Abstract base class for all Gauss-Newton type methods.

    This includes methods such as [`optimistix.GaussNewton`][],
    [`optimistix.LevenbergMarquardt`][], and [`optimistix.Dogleg`][].

    Subclasses must provide the following abstract attributes, with the following types:

    - `rtol: float`
    - `atol: float`
    - `norm: Callable[[PyTree], Scalar]`
    - `descent: AbstractDescent`
    - `search: AbstractSearch`
    - `verbose: frozenset[str]
    """

    rtol: AbstractVar[float]
    atol: AbstractVar[float]
    norm: AbstractVar[Callable[[PyTree], Scalar]]
    descent: AbstractVar[AbstractDescent[Y, FunctionInfo.ResidualJac, Any]]
    search: AbstractVar[
        AbstractSearch[Y, FunctionInfo.ResidualJac, FunctionInfo.ResidualJac, Any]
    ]
    verbose: AbstractVar[frozenset[str]]

    def init(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _GaussNewtonState:
        f_info_struct, _ = eqx.filter_eval_shape(_make_f_info, fn, y, args, tags)
        f_info = tree_full_like(f_info_struct, 0, allow_static=True)
        return _GaussNewtonState(
            first_step=jnp.array(True),
            y_eval=y,
            search_state=self.search.init(y, f_info_struct),
            f_info=f_info,
            aux=tree_full_like(aux_struct, 0),
            descent_state=self.descent.init(y, f_info_struct),
            terminate=jnp.array(False),
            result=RESULTS.successful,
            num_steps=jnp.array(0),
            num_accepted_steps=jnp.array(0),
            num_steps_since_acceptance=jnp.array(0),
        )

    def step(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _GaussNewtonState,
        tags: frozenset[object],
    ) -> tuple[Y, _GaussNewtonState, Aux]:
        f_eval_info, aux_eval = _make_f_info(fn, state.y_eval, args, tags)
        step_size, accept, search_result, search_state = self.search.step(
            state.first_step,
            y,
            state.y_eval,
            state.f_info,
            f_eval_info,
            state.search_state,
        )
        num_steps = state.num_steps + 1
        num_accepted_steps = state.num_accepted_steps + jnp.where(accept, 1, 0)
        num_steps_since_acceptance = jnp.where(
            accept, 0, state.num_steps_since_acceptance + 1
        )

        def accepted(descent_state):
            descent_state = self.descent.query(state.y_eval, f_eval_info, descent_state)
            y_diff = (state.y_eval**ω - y**ω).ω
            f_diff = (f_eval_info.residual**ω - state.f_info.residual**ω).ω
            terminate = cauchy_termination(
                self.rtol,
                self.atol,
                self.norm,
                state.y_eval,
                y_diff,
                f_eval_info.residual,
                f_diff,
            )
            return state.y_eval, f_eval_info, aux_eval, descent_state, terminate

        def rejected(descent_state):
            return y, state.f_info, state.aux, descent_state, jnp.array(False)

        y, f_info, aux, descent_state, terminate = filter_cond(
            accept, accepted, rejected, state.descent_state
        )

        if len(self.verbose) > 0:
            verbose_step = "step" in self.verbose
            verbose_loss = "loss" in self.verbose
            verbose_accepted = "accepted" in self.verbose
            verbose_step_size = "step_size" in self.verbose
            verbose_y = "y" in self.verbose
            loss_eval = 0.5 * sum_squares(f_eval_info.residual)
            loss = 0.5 * sum_squares(state.f_info.residual)
            verbose_print(
                (verbose_step, "Step", state.num_steps),
                (
                    verbose_step and verbose_accepted,
                    "Accepted steps",
                    state.num_accepted_steps,
                ),
                (
                    verbose_step and verbose_accepted,
                    "Steps since acceptance",
                    state.num_steps_since_acceptance,
                ),
                (verbose_loss, "Loss on this step", loss_eval),
                (
                    verbose_loss and verbose_accepted,
                    "Loss on the last accepted step",
                    loss,
                ),
                (verbose_step_size, "Step size", step_size),
                (verbose_y, "y", state.y_eval),
                (verbose_y and verbose_accepted, "y on the last accepted step", y),
            )

        y_descent, descent_result = self.descent.step(step_size, descent_state)
        y_eval = (y**ω + y_descent**ω).ω
        result = RESULTS.where(
            search_result == RESULTS.successful, descent_result, search_result
        )

        state = _GaussNewtonState(
            first_step=jnp.array(False),
            y_eval=y_eval,
            search_state=search_state,
            f_info=f_info,
            aux=aux,
            descent_state=descent_state,
            terminate=terminate,
            result=result,
            num_steps=num_steps,
            num_accepted_steps=num_accepted_steps,
            num_steps_since_acceptance=num_steps_since_acceptance,
        )
        return y, state, aux

    def terminate(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _GaussNewtonState,
        tags: frozenset[object],
    ):
        return state.terminate, state.result

    def postprocess(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _GaussNewtonState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        return y, aux, {}


class GaussNewton(AbstractGaussNewton[Y, Out, Aux]):
    """Gauss-Newton algorithm, for solving nonlinear least-squares problems.

    Note that regularised approaches like [`optimistix.LevenbergMarquardt`][] are
    usually preferred instead.
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    descent: NewtonDescent[Y]
    search: LearningRate[Y]
    verbose: frozenset[str]

    def __init__(
        self,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree], Scalar] = max_norm,
        linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=None),
        verbose: frozenset[str] = frozenset(),
    ):
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.descent = NewtonDescent(linear_solver=linear_solver)
        self.search = LearningRate(1.0)
        self.verbose = verbose


GaussNewton.__init__.__doc__ = """**Arguments:**

- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
- `linear_solver`: The linear solver used to compute the Newton step.
- `verbose`: Whether to print out extra information about how the solve is proceeding.
    Should be a frozenset of strings, specifying what information to print out. Valid
    entries are `step`, `loss`, `accepted`, `step_size`, `y`. For example
    `verbose=frozenset({"loss", "step_size"})`.
"""

optimistix/_solver/gauss_newton.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, Generic, Optional, Union
from typing_extensions import TypeAlias

import equinox as eqx
import jax
import jax.numpy as jnp
from equinox import AbstractVar
from equinox.internal import ω
from jaxtyping import Array, Bool, PyTree, Scalar

from .._custom_types import Aux, DescentState, Fn, Out, SearchState, Y
from .._minimise import AbstractMinimiser
from .._misc import (
    cauchy_termination,
    filter_cond,
    lin_to_grad,
    max_norm,
    tree_full_like,
)
from .._search import (
    AbstractDescent,
    AbstractSearch,
    FunctionInfo,
)
from .._solution import RESULTS
from .learning_rate import LearningRate


class _SteepestDescentState(eqx.Module, Generic[Y]):
    grad: Y


_FnInfo: TypeAlias = Union[
    FunctionInfo.EvalGrad,
    FunctionInfo.EvalGradHessian,
    FunctionInfo.EvalGradHessianInv,
    FunctionInfo.ResidualJac,
]


class SteepestDescent(AbstractDescent[Y, _FnInfo, _SteepestDescentState]):
    """The descent direction given by locally following the gradient."""

    norm: Optional[Callable[[PyTree], Scalar]] = None

    def init(self, y: Y, f_info_struct: _FnInfo) -> _SteepestDescentState:
        del f_info_struct
        # Dummy; unused
        return _SteepestDescentState(y)

    def query(
        self, y: Y, f_info: _FnInfo, state: _SteepestDescentState
    ) -> _SteepestDescentState:
        if isinstance(
            f_info,
            (
                FunctionInfo.EvalGrad,
                FunctionInfo.EvalGradHessian,
                FunctionInfo.EvalGradHessianInv,
                FunctionInfo.ResidualJac,
            ),
        ):
            grad = f_info.grad
        else:
            raise ValueError(
                "Cannot use `SteepestDescent` with this solver. This is because "
                "`SteepestDescent` requires gradients of the target function, but "
                "this solver does not evaluate such gradients."
            )
        if self.norm is not None:
            grad = (grad**ω / self.norm(grad)).ω
        return _SteepestDescentState(grad)

    def step(
        self, step_size: Scalar, state: _SteepestDescentState
    ) -> tuple[Y, RESULTS]:
        return (-step_size * state.grad**ω).ω, RESULTS.successful


SteepestDescent.__init__.__doc__ = """**Arguments:**

- `norm`: If passed, then normalise the gradient using this norm. (The returned step
    will have length `step_size` with respect to this norm.) Optimistix includes three
    built-in norms: [`optimistix.max_norm`][], [`optimistix.rms_norm`][], and
    [`optimistix.two_norm`][].
"""


class _GradientDescentState(
    eqx.Module, Generic[Y, Out, Aux, SearchState, DescentState]
):
    # Updated every search step
    first_step: Bool[Array, ""]
    y_eval: Y
    search_state: SearchState
    # Updated after each descent step
    f_info: FunctionInfo.EvalGrad
    aux: Aux
    descent_state: DescentState
    # Used for termination
    terminate: Bool[Array, ""]
    result: RESULTS


class AbstractGradientDescent(AbstractMinimiser[Y, Aux, _GradientDescentState]):
    """The gradient descent method for unconstrained minimisation.

    At every step, this algorithm performs a line search along the steepest descent
    direction. You should subclass this to provide it with a particular choice of line
    search. (E.g. [`optimistix.GradientDescent`][] uses a simple learning rate step.)

    Subclasses must provide the following abstract attributes, with the following types:

    - `rtol: float`
    - `atol: float`
    - `norm: Callable[[PyTree], Scalar]`
    - `descent: AbstractDescent`
    - `search: AbstractSearch`
    """

    rtol: AbstractVar[float]
    atol: AbstractVar[float]
    norm: AbstractVar[Callable[[PyTree], Scalar]]
    descent: AbstractVar[AbstractDescent[Y, FunctionInfo.EvalGrad, Any]]
    search: AbstractVar[
        AbstractSearch[Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any]
    ]

    def init(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: jax.ShapeDtypeStruct,
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _GradientDescentState:
        f_info = FunctionInfo.EvalGrad(jnp.zeros(f_struct.shape, f_struct.dtype), y)
        f_info_struct = jax.eval_shape(lambda: f_info)
        return _GradientDescentState(
            first_step=jnp.array(True),
            y_eval=y,
            search_state=self.search.init(y, f_info_struct),
            f_info=f_info,
            aux=tree_full_like(aux_struct, 0),
            descent_state=self.descent.init(y, f_info_struct),
            terminate=jnp.array(False),
            result=RESULTS.successful,
        )

    def step(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _GradientDescentState,
        tags: frozenset[object],
    ) -> tuple[Y, _GradientDescentState, Aux]:
        f_eval, lin_fn, aux_eval = jax.linearize(
            lambda _y: fn(_y, args), state.y_eval, has_aux=True
        )
        step_size, accept, search_result, search_state = self.search.step(
            state.first_step,
            y,
            state.y_eval,
            state.f_info,
            FunctionInfo.Eval(f_eval),
            state.search_state,
        )

        def accepted(descent_state):
            (grad,) = lin_to_grad(lin_fn, y)
            f_eval_info = FunctionInfo.EvalGrad(f_eval, grad)
            descent_state = self.descent.query(state.y_eval, f_eval_info, descent_state)
            y_diff = (state.y_eval**ω - y**ω).ω
            f_diff = (f_eval**ω - state.f_info.f**ω).ω
            terminate = cauchy_termination(
                self.rtol, self.atol, self.norm, state.y_eval, y_diff, f_eval, f_diff
            )
            return state.y_eval, f_eval_info, aux_eval, descent_state, terminate

        def rejected(descent_state):
            return y, state.f_info, state.aux, descent_state, jnp.array(False)

        y, f_info, aux, descent_state, terminate = filter_cond(
            accept, accepted, rejected, state.descent_state
        )

        y_descent, descent_result = self.descent.step(step_size, descent_state)
        y_eval = (y**ω + y_descent**ω).ω
        result = RESULTS.where(
            search_result == RESULTS.successful, descent_result, search_result
        )

        state = _GradientDescentState(
            first_step=jnp.array(False),
            y_eval=y_eval,
            search_state=search_state,
            f_info=f_info,
            aux=aux,
            descent_state=descent_state,
            terminate=terminate,
            result=result,
        )
        return y, state, aux

    def terminate(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _GradientDescentState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        return state.terminate, state.result

    def postprocess(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _GradientDescentState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        return y, aux, {}


class GradientDescent(AbstractGradientDescent[Y, Aux]):
    """Classic gradient descent with a learning rate `learning_rate`."""

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    descent: SteepestDescent[Y]
    search: LearningRate[Y]

    def __init__(
        self,
        learning_rate: float,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree], Scalar] = max_norm,
    ):
        """**Arguments:**

        - `learning_rate`: Specifies a constant learning rate to use at each step.
        - `rtol`: Relative tolerance for terminating the solve.
        - `atol`: Absolute tolerance for terminating the solve.
        - `norm`: The norm used to determine the difference between two iterates in the
            convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
            includes three built-in norms: [`optimistix.max_norm`][],
            [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
        """
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.descent = SteepestDescent()
        self.search = LearningRate(learning_rate)

optimistix/_solver/gradient_methods.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC #
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import cast

import equinox as eqx
import jax.numpy as jnp
from jaxtyping import Array, Bool, Scalar, ScalarLike

from .._custom_types import Y
from .._search import AbstractSearch, FunctionInfo
from .._solution import RESULTS


class LearningRate(AbstractSearch[Y, FunctionInfo, FunctionInfo, None]):
    """Move downhill by taking a step of the fixed size `learning_rate`."""

    learning_rate: ScalarLike = eqx.field(converter=jnp.asarray)

    def init(self, y: Y, f_info_struct: FunctionInfo) -> None:
        return None

    def step(
        self,
        first_step: Bool[Array, ""],
        y: Y,
        y_eval: Y,
        f_info: FunctionInfo,
        f_eval_info: FunctionInfo,
        state: None,
    ) -> tuple[Scalar, Bool[Array, ""], RESULTS, None]:
        del first_step, y, y_eval, f_info, f_eval_info, state
        learning_rate = cast(Array, self.learning_rate)
        return learning_rate, jnp.array(True), RESULTS.successful, None


LearningRate.__init__.__doc__ = """**Arguments:**

- `learning_rate`: The fixed step-size used at each step.
"""

optimistix/_solver/learning_rate.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import cast, Generic, Union

import equinox as eqx
import jax.lax as lax
import jax.numpy as jnp
import jax.tree_util as jtu
import lineax as lx
from equinox.internal import ω
from jaxtyping import Array, Float, PyTree, Scalar, ScalarLike

from .._custom_types import Aux, Out, Y
from .._misc import max_norm, tree_full_like, two_norm
from .._root_find import AbstractRootFinder, root_find
from .._search import AbstractDescent, FunctionInfo
from .._solution import RESULTS
from .gauss_newton import AbstractGaussNewton, newton_step
from .newton_chord import Newton
from .trust_region import ClassicalTrustRegion


class _Damped(eqx.Module):
    operator: lx.AbstractLinearOperator
    damping: Float[Array, ""]

    def __call__(self, y: PyTree[Array]):
        residual = self.operator.mv(y)
        damped = jtu.tree_map(lambda yi: jnp.sqrt(self.damping) * yi, y)
        return residual, damped


def damped_newton_step(
    step_size: Scalar,
    f_info: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
    linear_solver: lx.AbstractLinearSolver,
) -> tuple[PyTree[Array], RESULTS]:
    """Compute a damped Newton step.

    For a minimisation problem, this means solving `(Hess + λI)^{-1} grad`.

    In the (nonlinear) least-squares case, for which the minimisation objective
    is given by `0.5*residual^2`, then we know that `grad=J^T residual`, and we make
    the Gauss--Newton approximation `Hess ~ J^T J`. This reduces the above to
    solving the (linear) least-squares problem
    ```
    [    J   ] [diff]  =  [residual]
    [sqrt(λ)I]         =  [   0    ]
    ```
    This can be seen by observing that the normal equations for the this linear
    least-squares problem is the original linear problem we wanted to solve.
    """

    pred = step_size > jnp.finfo(step_size.dtype).eps
    safe_step_size = jnp.where(pred, step_size, 1)
    lm_param = jnp.where(pred, 1 / safe_step_size, jnp.finfo(step_size).max)
    lm_param = cast(Array, lm_param)
    if isinstance(f_info, FunctionInfo.EvalGradHessian):
        operator = f_info.hessian + lm_param * lx.IdentityLinearOperator(
            f_info.hessian.in_structure()
        )
        vector = f_info.grad
        if lx.is_positive_semidefinite(f_info.hessian):
            operator = lx.TaggedLinearOperator(operator, lx.positive_semidefinite_tag)
    elif isinstance(f_info, FunctionInfo.ResidualJac):
        y_structure = f_info.jac.in_structure()
        operator = lx.FunctionLinearOperator(_Damped(f_info.jac, lm_param), y_structure)
        vector = (f_info.residual, tree_full_like(y_structure, 0))
    else:
        raise ValueError(
            "Damped newton descent cannot be used with a solver that does not "
            "provide (approximate) Hessian information."
        )
    linear_sol = lx.linear_solve(operator, vector, linear_solver, throw=False)
    return linear_sol.value, RESULTS.promote(linear_sol.result)


class _DampedNewtonDescentState(eqx.Module):
    f_info: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac]


class DampedNewtonDescent(
    AbstractDescent[
        Y,
        Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
        _DampedNewtonDescentState,
    ]
):
    """The damped Newton (Levenberg--Marquardt) descent.

    That is: gradient descent is about moving in the direction of `-grad`.
    (Quasi-)Newton descent is about moving in the direction of `-Hess^{-1} grad`. Damped
    Newton interpolates between these two regimes, by moving in the direction of
    `-(Hess + λI)^{-1} grad`.

    The value λ is often referred to as a the "Levenberg--Marquardt" parameter, and in
    version is handled directly, as λ = 1/step_size. Larger step sizes correspond to
    Newton directions; smaller step sizes correspond to gradient directions. (And
    additionally also reduces the step size, hence the term "damping".) This is because
    a line search expects the step to be smaller as the step size decreases.
    """

    # Will probably resolve to either Cholesky (for minimisation problems) or
    # QR (for least-squares problems).
    linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=None)

    def init(
        self,
        y: Y,
        f_info_struct: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
    ) -> _DampedNewtonDescentState:
        del y
        f_info_init = tree_full_like(f_info_struct, 0, allow_static=True)
        return _DampedNewtonDescentState(f_info_init)

    def query(
        self,
        y: Y,
        f_info: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
        state: _DampedNewtonDescentState,
    ) -> _DampedNewtonDescentState:
        del y, state
        return _DampedNewtonDescentState(f_info)

    def step(
        self, step_size: Scalar, state: _DampedNewtonDescentState
    ) -> tuple[Y, RESULTS]:
        sol_value, result = damped_newton_step(
            step_size, state.f_info, self.linear_solver
        )
        y_diff = (-(sol_value**ω)).ω
        return y_diff, result


DampedNewtonDescent.__init__.__doc__ = """**Arguments:**

- `linear_solver`: The linear solver used to compute the Newton step.
"""


class _IndirectDampedNewtonDescentState(eqx.Module, Generic[Y]):
    f_info: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac]
    newton: Y
    newton_norm: Scalar
    result: RESULTS


class IndirectDampedNewtonDescent(
    AbstractDescent[
        Y,
        Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
        _IndirectDampedNewtonDescentState,
    ]
):
    """The indirect damped Newton (Levenberg--Marquardt) trust-region descent.

    If the above line just looks like technical word soup, here's what's going on:

    Gradient descent is about moving in the direction of `-grad`. (Quasi-)Newton descent
    is about moving in the direction of `-Hess^{-1} grad`. Damped Newton interpolates
    between these two regimes, by moving in the direction of
    `-(Hess + λI)^{-1} grad`.

    This can be derived as the dual problem of a trust region method, see Conn, Gould,
    Toint: "Trust-Region Methods" section 7.3. λ is interpreted as a Lagrange
    multiplier. This involves solving a one-dimensional root-finding problem for λ at
    each descent.
    """

    lambda_0: ScalarLike = 1.0
    linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=False)
    # Default tol for `root_finder` because the tol doesn't have to be very strict.
    root_finder: AbstractRootFinder = Newton(rtol=1e-2, atol=1e-2)
    trust_region_norm: Callable[[PyTree], Scalar] = two_norm

    def init(
        self,
        y: Y,
        f_info_struct: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
    ) -> _IndirectDampedNewtonDescentState:
        return _IndirectDampedNewtonDescentState(
            f_info=tree_full_like(f_info_struct, 0, allow_static=True),
            newton=y,
            newton_norm=jnp.array(0.0),
            result=RESULTS.successful,
        )

    def query(
        self,
        y: Y,
        f_info: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
        state: _IndirectDampedNewtonDescentState,
    ) -> _IndirectDampedNewtonDescentState:
        del y
        newton, result = newton_step(f_info, self.linear_solver)
        newton_norm = self.trust_region_norm(newton)
        return _IndirectDampedNewtonDescentState(
            f_info=f_info, newton=newton, newton_norm=newton_norm, result=result
        )

    def step(
        self, step_size: Scalar, state: _IndirectDampedNewtonDescentState
    ) -> tuple[Y, RESULTS]:
        # For trust-region methods like `IndirectDampedNewtonDescent`, the trust-region
        # size directly controls how large a step we take. This is actually somewhat
        # annoying, as a trust region algorithm has no understanding of the scale of a
        # problem unless initialised with a complex initialisation procedure.
        #
        # A simple, heuristic way around this is to scale the trust region `step_size`
        # so that a `step_size` of `1` corresponds to the full length of the Newton
        # step (anything greater than `1` will still accept the Newton step.)
        scaled_step_size = state.newton_norm * step_size

        def comparison_fn(lambda_i: Scalar, _):
            step, _ = damped_newton_step(1 / lambda_i, state.f_info, self.linear_solver)
            return self.trust_region_norm(step) - scaled_step_size

        def reject_newton():
            lambda_out = root_find(
                fn=comparison_fn,
                has_aux=False,
                solver=self.root_finder,
                y0=jnp.asarray(self.lambda_0),
                options=dict(lower=1e-5),
                max_steps=32,
                throw=False,
            ).value
            y_diff, result = damped_newton_step(
                1 / lambda_out, state.f_info, self.linear_solver
            )
            return y_diff, result

        def accept_newton():
            return state.newton, state.result

        # Only do a root-find if we have a small step size, and our Newton step was
        # successful.
        do_root_solve = (state.result == RESULTS.successful) & (step_size < 1)
        neg_y_diff, new_result = lax.cond(do_root_solve, reject_newton, accept_newton)
        return (-(neg_y_diff**ω)).ω, new_result


IndirectDampedNewtonDescent.__init__.__doc__ = """**Arguments:**    

- `lambda_0`: The initial value of the Levenberg--Marquardt parameter used in the root-
    find to hit the trust-region radius. If `IndirectDampedNewtonDescent` is failing,
    this value may need to be increased.
- `linear_solver`: The linear solver used to compute the Newton step.
- `root_finder`: The root finder used to find the Levenberg--Marquardt parameter which
    hits the trust-region radius.
- `trust_region_norm`: The norm used to determine the trust-region shape.
"""


class LevenbergMarquardt(AbstractGaussNewton[Y, Out, Aux]):
    """The Levenberg--Marquardt method.

    This is a classical solver for nonlinear least squares, which works by regularising
    [`optimistix.GaussNewton`][] with a damping factor. This serves to (a) interpolate
    between Gauss--Newton and steepest descent, and (b) limit step size to a local
    region around the current point.

    This is a good algorithm for many least squares problems.
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    descent: DampedNewtonDescent[Y]
    search: ClassicalTrustRegion[Y]
    verbose: frozenset[str]

    def __init__(
        self,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree], Scalar] = max_norm,
        linear_solver: lx.AbstractLinearSolver = lx.QR(),
        verbose: frozenset[str] = frozenset(),
    ):
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.descent = DampedNewtonDescent(linear_solver=linear_solver)
        self.search = ClassicalTrustRegion()
        self.verbose = verbose


LevenbergMarquardt.__init__.__doc__ = """**Arguments:**

- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
- `linear_solver`: The linear solver to use to solve the damped Newton step. Defaults to
    `lineax.QR`.
- `verbose`: Whether to print out extra information about how the solve is proceeding.
    Should be a frozenset of strings, specifying what information to print out. Valid
    entries are `step`, `loss`, `accepted`, `step_size`, `y`. For example
    `verbose=frozenset({"loss", "step_size"})`.
"""


class IndirectLevenbergMarquardt(AbstractGaussNewton[Y, Out, Aux]):
    """The Levenberg--Marquardt method as a true trust-region method.

    This is a variant of [`optimistix.LevenbergMarquardt`][]. The other algorithm works
    by updating the damping factor directly -- this version instead updates a trust
    region, and then fits the damping factor to the size of the trust region.

    Generally speaking [`optimistix.LevenbergMarquardt`][] is preferred, as it performs
    nearly the same algorithm, without the computational overhead of an extra (scalar)
    nonlinear solve.
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    descent: IndirectDampedNewtonDescent[Y]
    search: ClassicalTrustRegion[Y]
    verbose: frozenset[str]

    def __init__(
        self,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree], Scalar] = max_norm,
        lambda_0: ScalarLike = 1.0,
        linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=False),
        root_finder: AbstractRootFinder = Newton(rtol=0.01, atol=0.01),
        verbose: frozenset[str] = frozenset(),
    ):
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.descent = IndirectDampedNewtonDescent(
            lambda_0=lambda_0,
            linear_solver=linear_solver,
            root_finder=root_finder,
        )
        self.search = ClassicalTrustRegion()
        self.verbose = verbose


IndirectLevenbergMarquardt.__init__.__doc__ = """**Arguments:**
    
- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
- `lambda_0`: The initial value of the Levenberg--Marquardt parameter used in the root-
    find to hit the trust-region radius. If `IndirectLevenbergMarquardt` is failing,
    this value may need to be increased.
- `linear_solver`: The linear solver used to compute the Newton step.
- `root_finder`: The root finder used to find the Levenberg--Marquardt parameter which
    hits the trust-region radius.
- `verbose`: Whether to print out extra information about how the solve is proceeding.
    Should be a frozenset of strings, specifying what information to print out. Valid
    entries are `step`, `loss`, `accepted`, `step_size`, `y`. For example
    `verbose=frozenset({"loss", "step_size"})`.
"""

optimistix/_solver/levenberg_marquardt.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools as ft
from collections.abc import Callable
from typing import Any, cast, Generic

import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
import jax.tree_util as jtu
from equinox.internal import ω
from jaxtyping import Array, ArrayLike, Bool, PyTree, Scalar

from .._custom_types import Aux, Fn, Y
from .._minimise import AbstractMinimiser
from .._misc import max_norm, tree_full_like, tree_where
from .._solution import RESULTS


class _NMStats(eqx.Module):
    n_reflect: Scalar
    n_inner_contract: Scalar
    n_outer_contract: Scalar
    n_expand: Scalar
    n_shrink: Scalar


class _NelderMeadState(eqx.Module, Generic[Y, Aux]):
    """
    Information to update and store the simplex of the Nelder Mead update. If
    `dim` is the dimension of the problem, we expect there to be
    `n_vertices = dim + 1` vertices. We expect the leading axis of each leaf
    to be of length `n_vertices`, and the sum of the rest of the axes of all leaves
    together to be `dim`.

    - `simplex`: a PyTree with leading axis of leaves `n_vertices` and sum of the
        rest of the axes of all leaves `dim`.
    - `f_simplex`: a 1-dimensional array of size `n_vertices`.
        The values of the problem function evaluated on each vertex of
        simplex.
    - `best`: A tuple of shape (Scalar, PyTree, Scalar). The tuple contains
        (`f(best_vertex)`, `best_vertex`, index of `best_vertex`) where
        `best_vertex` is the vertex which minimises `f` among all vertices in
        `simplex`.
    - `worst`: A tuple of shape (Scalar, PyTree, Scalar). The tuple contains
        (`f(worst_vertex)`, `worst_vertex`, index of `worst_vertex`) where
        `worst_vertex` is the vertex which maximises `f` among all vertices in
        `simplex`.
    -`second_worst`: A scalar, which is `f(second_worst_vertex)` where
        `second_worst_vertex` is the vertex which maximises `f` among all vertices
        in `simplex` with `worst_vertex` removed.
    - `step`: A scalar. How many steps have been taken so far.
    - `stats`: A `_NMStats` PyTree. This tracks information about the Nelder Mead
        algorithm. Specifically, how many times each of the operations reflect,
        expand, inner contract, outer contract, and shrink are performed.
    - `result`: a [`optimistix.RESULTS`][] object which indicates if we have diverged
        during the course of optimisation.
    - `first_pass`: A bool which indicates if this is the first call to Nelder Mead
        which allows for extra setup. This ultimately exists to save on compilation
        time.
    """

    simplex: PyTree
    f_simplex: Array
    best: tuple[Scalar, Y, Scalar]
    worst: tuple[Scalar, Y, Scalar]
    second_worst_val: Scalar
    stats: _NMStats
    result: RESULTS
    first_pass: Bool[ArrayLike, ""]
    aux: Aux
    step: Scalar


def _update_stats(
    stats,
    reflect: Bool[Array, ""],
    inner_contract: Bool[Array, ""],
    outer_contract: Bool[Array, ""],
    expand: Bool[Array, ""],
    shrink: Bool[Array, ""],
) -> _NMStats:
    return _NMStats(
        stats.n_reflect + jnp.where(reflect, 1, 0),
        stats.n_inner_contract + jnp.where(inner_contract, 1, 0),
        stats.n_outer_contract + jnp.where(outer_contract, 1, 0),
        stats.n_expand + jnp.where(expand, 1, 0),
        stats.n_shrink + jnp.where(shrink, 1, 0),
    )


class NelderMead(AbstractMinimiser[Y, Aux, _NelderMeadState[Y, Aux]]):
    """The Nelder-Mead minimisation algorithm. (Downhill simplex derivative-free
    method.)

    This algorithm is notable in that it only uses function evaluations, and does not
    need gradient evaluations.

    This is usually an "algorithm of last resort". Gradient-based algorithms are usually
    much faster, and be more likely to converge to a minima.

    Comparable to `scipy.optimize.minimize(method="Nelder-Mead")`.
    """

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar] = max_norm
    rdelta: float = 5e-2
    adelta: float = 2.5e-4

    def init(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _NelderMeadState[Y, Aux]:
        aux = tree_full_like(aux_struct, 0)
        y0_simplex = options.get("y0_simplex", False)

        if y0_simplex:
            n_vertices = {x.shape[0] for x in jtu.tree_leaves(y)}
            if len(n_vertices) == 1:
                [n_vertices] = n_vertices
            else:
                raise ValueError(
                    "The PyTree must form a valid simplex. \
                    Got different leading dimension (number of vertices) \
                    for each leaf"
                )
            size = jtu.tree_reduce(
                lambda x, y: x + y, jtu.tree_map(lambda x: x[1:].size, y)
            )
            if n_vertices != size + 1:
                raise ValueError(
                    "The PyTree must form a valid simplex. Got"
                    f"{n_vertices} vertices but dimension {size}."
                )
            simplex = y
        else:
            #
            # The standard approach to creating the initial simplex from a single vector
            # is to add a small constant times each unit vector to the initial vector.
            # The constant is different if the `y` is 0 in the direction of the
            # unit vector. Just because this is standard, does not mean it's well
            # justified. We add rdelta * y[i] + adelta y[i] in the ith unit direction.
            #
            size = jtu.tree_reduce(lambda x, y: x + y, jtu.tree_map(jnp.size, y))
            n_vertices = size + 1
            leaves, treedef = jtu.tree_flatten(y)
            running_size = 0
            new_leaves = []

            for index, leaf in enumerate(leaves):
                # TODO(raderj): This is a difficult bit of Pytree manipulation, and
                # needs to either be simplified if possible or explained, maybe with a
                # follow-along example.
                leaf_size = jnp.size(leaf)
                broadcast_leaves = jnp.repeat(leaf[None, ...], size + 1, axis=0)
                indices = jnp.arange(
                    running_size + 1, running_size + leaf_size + 1, dtype=jnp.int16
                )
                relative_indices = jnp.unravel_index(
                    indices - running_size, shape=leaf.shape
                )
                indices = jnp.unravel_index(indices, shape=broadcast_leaves.shape)
                broadcast_leaves = broadcast_leaves.at[indices].set(
                    broadcast_leaves[indices]
                    + self.adelta
                    + self.rdelta * leaf[relative_indices]
                )
                running_size = running_size + leaf_size
                new_leaves.append(broadcast_leaves)
            simplex = jtu.tree_unflatten(treedef, new_leaves)

        f_simplex = jnp.full(n_vertices, jnp.inf, dtype=f_struct.dtype)
        # Shrink will be called the first time step is called, so remove one from stats
        # at init.
        stats = _NMStats(
            jnp.array(0), jnp.array(0), jnp.array(0), jnp.array(0), jnp.array(-1)
        )

        return _NelderMeadState(
            simplex=simplex,
            f_simplex=f_simplex,
            best=(
                jnp.array(0.0),
                jtu.tree_map(lambda x: x[0], simplex),
                jnp.array(0, dtype=jnp.int32),
            ),
            worst=(
                jnp.array(0.0),
                jtu.tree_map(lambda x: x[0], simplex),
                jnp.array(0, dtype=jnp.int32),
            ),
            second_worst_val=jnp.array(0.0),
            stats=stats,
            result=RESULTS.successful,
            first_pass=jnp.array(True),
            aux=aux,
            step=jnp.array(0),
        )

    def step(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _NelderMeadState[Y, Aux],
        tags: frozenset[object],
    ) -> tuple[Y, _NelderMeadState[Y, Aux], Aux]:
        # TODO(raderj): update to line search api.
        reflect_const = 2
        expand_const = 3
        out_const = 1.5
        in_const = 0.5
        shrink_const = 0.5

        f_best, best, best_index = state.best
        f_worst, worst, worst_index = state.worst
        f_second_worst = state.second_worst_val
        stats = state.stats
        (n_vertices,) = state.f_simplex.shape

        def init_step(state):
            simplex = state.simplex
            f_new_vertex = jnp.array(jnp.inf)
            return (
                state,
                simplex,
                jtu.tree_map(lambda _, x: x[0], y, simplex),
                f_new_vertex,
                state.stats,
                state.aux,
            )

        def main_step(state):
            # TODO(raderj): Calculate the centroid and search dir based upon
            # the prior one.
            search_direction = jtu.tree_map(
                lambda _, x, y: x[...] - y, y, state.simplex, worst
            )
            search_direction = (
                (ω(search_direction) / (n_vertices - 1))
                .call(ft.partial(jnp.sum, axis=0))
                .ω
            )
            reflection = (worst**ω + reflect_const * search_direction**ω).ω

            def eval_new_vertices(vertex_carry, i):
                vertex, (f_vertex, _), stats = vertex_carry

                def internal_eval(f_vertex, stats):
                    expand = f_vertex < f_best
                    inner_contract = f_vertex > f_worst
                    contract = f_vertex > f_second_worst
                    outer_contract = jnp.invert(expand | contract)
                    reflect = (f_vertex > f_best) & (f_vertex < f_second_worst)
                    contract_const = jnp.where(inner_contract, in_const, out_const)
                    #
                    # NOTE: Throughout, we use worst + const * search_direction, not
                    # centroid + const * search direction! The latter
                    # is ubiquitious in Nelder-Mead implementations, but is a
                    # mistake. Nearly every line search assumes:
                    #
                    # a) `delta > 0` for `delta` the line search param.
                    # b) `f(x + delta dir) < f(x)` for suitably small `delta`
                    #
                    # both of these are violated when we start the line search at
                    # centroid (indeed, inner contraction usually requires
                    # multiplication by a negative constant). However, if we start
                    # the search at worst both these issues are avoided.
                    #
                    expanded = (worst**ω + expand_const * search_direction**ω).ω
                    contracted = (worst**ω + contract_const * search_direction**ω).ω
                    new_vertex = tree_where(
                        expand,
                        expanded,
                        vertex,
                    )
                    new_vertex = tree_where(
                        contract,
                        contracted,
                        new_vertex,
                    )
                    stats = _update_stats(
                        stats,
                        reflect,
                        inner_contract,
                        outer_contract,
                        expand,
                        shrink=jnp.array(False),
                    )
                    return new_vertex, stats

                out, stats = lax.cond(
                    i == 1,
                    internal_eval,
                    lambda x, y: (vertex, stats),
                    *(f_vertex, stats),
                )
                return (out, fn(out, args), stats), None

            #
            # We'd like to avoid making two calls to problem.fn in the step to
            # avoid compiling the potentially large problem.fn twice. Instead, we
            # wrap the two evaluations in a single scan and pass different inputs
            # each time.
            #
            # the first iteration of scan will have lax.cond(False) and will return
            # f(reflection) for use in the second iteration which calls internal_eval
            # which uses f(reflection) to determine the next vertex, and returns the
            # vertex along with f(new_vertex) and stats.
            #
            # TODO(raderj): pull out this entire thing into line  api.
            #
            # TODO(raderj): eqxi.scan_trick.
            #
            (new_vertex, (f_new_vertex, aux), stats), _ = lax.scan(
                eval_new_vertices,
                (reflection, (jnp.array(0.0), state.aux), state.stats),
                jnp.arange(2),
            )
            return (state, state.simplex, new_vertex, f_new_vertex, stats, aux)

        state, simplex, new_vertex, f_new_vertex, stats, aux = lax.cond(
            state.first_pass, init_step, main_step, state
        )
        new_best = f_new_vertex < f_best
        best = tree_where(new_best, new_vertex, best)
        f_best = jnp.where(new_best, f_new_vertex, f_best)
        shrink = cast(
            Array, jnp.where(state.first_pass, jnp.array(True), f_new_vertex > f_worst)
        )
        stats = _update_stats(
            stats,
            reflect=jnp.array(False),
            inner_contract=jnp.array(False),
            outer_contract=jnp.array(False),
            expand=jnp.array(False),
            shrink=shrink,
        )

        def shrink_simplex(best, new_vertex, simplex, first_pass):
            # This is just `best + shrink_const * (simplex - best)` but returns
            # a buffer. `best` is passed t
            shrink_simplex = jtu.tree_map(
                lambda b, a: a.at[...].set(b + shrink_const * (a[...] - b)),
                best,
                simplex,
            )
            # if it is the first pass and we just wanted to use shrink_simplex to
            # compute f_simplex, return simplex
            simplex = tree_where(first_pass, simplex, shrink_simplex)
            unwrapped_simplex = jtu.tree_map(lambda _, x: x[...], y, simplex)
            f_simplex, _ = jax.vmap(lambda x: fn(x, args))(unwrapped_simplex)
            return f_simplex, simplex

        def update_simplex(best, new_vertex, simplex, first_pass):
            simplex = jtu.tree_map(
                lambda _, a, b: a.at[worst_index].set(b), best, simplex, new_vertex
            )
            f_simplex = state.f_simplex.at[worst_index].set(f_new_vertex)
            return f_simplex, simplex

        f_new_simplex, new_simplex = lax.cond(
            shrink,
            shrink_simplex,
            update_simplex,
            *(best, new_vertex, state.simplex, state.first_pass),
        )
        #
        # TODO(raderj): only 1 value is updated when not shrinking. This implies
        # that in most cases, rather than do a top_k search in log time, we can
        # just compare f_next_vector and f_best and choose best between those two
        # in constant time. Implement this. A similar thing could likely be done with
        # worst and second worst, with recomputation occurring only when f_new < f_worst
        # but f_new > f_second_worst (otherwise, there will have been a shrink).
        #
        (f_best_neg,), (best_index,) = lax.top_k(-f_new_simplex, 1)
        f_best = -f_best_neg
        (f_worst, f_second_worst), (worst_index, _) = lax.top_k(f_new_simplex, 2)
        _structured_index = lambda a, b: jtu.tree_map(lambda _, x: x[b], y, a)
        best = _structured_index(simplex, best_index)
        worst = _structured_index(simplex, worst_index)
        new_state = _NelderMeadState(
            simplex=new_simplex,
            f_simplex=f_new_simplex,
            best=(f_best, best, best_index),
            worst=(f_worst, worst, worst_index),
            second_worst_val=f_second_worst,
            stats=stats,
            result=state.result,
            first_pass=jnp.array(False),
            aux=aux,
            step=state.step + 1,
        )
        try:
            y0_simplex = options["y0_simplex"]
        except KeyError:
            y0_simplex = False

        if y0_simplex:
            out = simplex
        else:
            out = best

        return out, new_state, aux

    def terminate(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _NelderMeadState[Y, Aux],
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        # TODO(raderj): only check terminate every k
        f_best, best, best_index = state.best
        x_scale = (self.atol + self.rtol * ω(best).call(jnp.abs)).ω
        x_diff = jtu.tree_map(
            lambda _, a, b: jnp.abs(a[...] - b), best, state.simplex, best
        )
        x_converged = self.norm((x_diff**ω / x_scale**ω).ω) < 1
        f_scale = (self.atol + self.rtol * ω(f_best).call(jnp.abs)).ω
        f_diff = (state.f_simplex**ω - f_best**ω).call(jnp.abs).ω
        f_converged = self.norm((f_diff**ω / f_scale**ω).ω) < 1
        #
        # minpack does a further test here where it takes for each unit vector e_i a
        # perturbation "delta" and asserts that f(x + delta e_i) > f(x) and
        # f(x - delta e_i) > f(x). ie. a rough assertion that it is indeed a local
        # minimum. If it fails the algo resets completely. thus process scales as
        # O(dim(y) * T(f)), where T(f) is the cost of evaluating f.
        #
        converged = x_converged & f_converged
        diverged = jnp.any(jnp.invert(jnp.isfinite(f_best)))
        terminate = converged | diverged
        result = RESULTS.where(
            diverged, RESULTS.nonlinear_divergence, RESULTS.successful
        )
        return terminate, result

    def postprocess(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _NelderMeadState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        stats = dict(
            num_reflections=state.stats.n_reflect,
            num_inner_contractions=state.stats.n_inner_contract,
            num_outer_contractions=state.stats.n_outer_contract,
            num_expansions=state.stats.n_expand,
            num_shrinkages=state.stats.n_shrink,
        )
        return y, aux, stats


NelderMead.__init__.__doc__ = """**Arguments:**

- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
- `rdelta`: Nelder-Mead creates an initial simplex by appending a scaled identity 
    matrix to `y`. The `i`th element of this matrix is `rdelta * y_i + adelta`.
    That is, this is the relative size for creating the initial simplex.
- `adelta`: Nelder-Mead creates an initial simplex by appending a scaled identity 
    matrix to `y`. The `i`th element of this matrix is `rdelta * y_i + adelta`.
    That is, this is the absolute size for creating the initial simplex.
"""

optimistix/_solver/nelder_mead.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, Generic, Optional, TYPE_CHECKING

import equinox as eqx
import jax
import jax.lax as lax
import jax.numpy as jnp
import jax.tree_util as jtu
import lineax as lx


if TYPE_CHECKING:
    from typing import ClassVar as AbstractClassVar
else:
    from equinox import AbstractClassVar
from equinox.internal import ω
from jaxtyping import Array, Bool, PyTree, Scalar

from .._custom_types import Aux, Fn, Out, Y
from .._misc import cauchy_termination, max_norm, tree_full_like
from .._root_find import AbstractRootFinder
from .._solution import RESULTS


def _small(diffsize: Scalar) -> Bool[Array, " "]:
    # TODO(kidger): make a more careful choice here -- the existence of this
    # function is pretty ad-hoc.
    resolution = 10 ** (2 - jnp.finfo(diffsize.dtype).precision)
    return diffsize < resolution


def _diverged(rate: Scalar) -> Bool[Array, " "]:
    return jnp.invert(jnp.isfinite(rate)) | (rate > 2)


def _converged(factor: Scalar, tol: float) -> Bool[Array, " "]:
    return (factor > 0) & (factor < tol)


class _NewtonChordState(eqx.Module, Generic[Y]):
    f: PyTree[Array]
    linear_state: Optional[tuple[lx.AbstractLinearOperator, PyTree]]
    diff: Y
    diffsize: Scalar
    diffsize_prev: Scalar
    result: RESULTS
    step: Scalar


class _NewtonChord(AbstractRootFinder[Y, Out, Aux, _NewtonChordState[Y]]):
    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar] = max_norm
    kappa: float = 1e-2
    linear_solver: lx.AbstractLinearSolver = lx.AutoLinearSolver(well_posed=None)
    cauchy_termination: bool = True

    _is_newton: AbstractClassVar[bool]

    def init(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _NewtonChordState[Y]:
        del options, aux_struct
        if self._is_newton:
            linear_state = None
        else:
            # TODO(kidger): evaluate on just the first step, to reduce compile time.
            jac = lx.JacobianLinearOperator(fn, y, args, tags=tags, _has_aux=True)
            jac = lx.linearise(jac)
            init_later_state = self.linear_solver.init(jac, options={})
            init_later_state = lax.stop_gradient(init_later_state)
            linear_state = (jac, init_later_state)
        if self.cauchy_termination:
            f_val = tree_full_like(f_struct, jnp.inf)
        else:
            f_val = None
        return _NewtonChordState(
            f=f_val,
            linear_state=linear_state,
            diff=tree_full_like(y, jnp.inf),
            diffsize=jnp.array(jnp.inf),
            diffsize_prev=jnp.array(1.0),
            result=RESULTS.successful,
            step=jnp.array(0),
        )

    def step(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _NewtonChordState[Y],
        tags: frozenset[object],
    ) -> tuple[Y, _NewtonChordState[Y], Aux]:
        lower = options.get("lower")
        upper = options.get("upper")
        del options
        if self._is_newton:
            fx, lin_fn, aux = jax.linearize(lambda _y: fn(_y, args), y, has_aux=True)
            jac = lx.FunctionLinearOperator(
                lin_fn, jax.eval_shape(lambda: y), tags=tags
            )
            sol = lx.linear_solve(jac, fx, self.linear_solver, throw=False)
        else:
            fx, aux = fn(y, args)
            jac, linear_state = state.linear_state  # pyright: ignore
            linear_state = lax.stop_gradient(linear_state)
            sol = lx.linear_solve(
                jac, fx, self.linear_solver, state=linear_state, throw=False
            )
        diff = sol.value
        new_y = (y**ω - diff**ω).ω
        if lower is not None:
            new_y = jtu.tree_map(lambda a, b: jnp.clip(a, a_min=b), new_y, lower)
        if upper is not None:
            new_y = jtu.tree_map(lambda a, b: jnp.clip(a, a_max=b), new_y, upper)
        if lower is not None or upper is not None:
            diff = (y**ω - new_y**ω).ω
        scale = (self.atol + self.rtol * ω(new_y).call(jnp.abs)).ω
        diffsize = self.norm((diff**ω / scale**ω).ω)
        if self.cauchy_termination:
            f_val = fx
        else:
            f_val = None
        new_state = _NewtonChordState(
            f=f_val,
            linear_state=state.linear_state,
            diff=diff,
            diffsize=diffsize,
            diffsize_prev=state.diffsize,
            result=RESULTS.promote(sol.result),
            step=state.step + 1,
        )
        return new_y, new_state, aux

    def terminate(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _NewtonChordState[Y],
        tags: frozenset[object],
    ):
        del fn, args, options
        if self.cauchy_termination:
            # Compare `f_val` against 0, not against some `f_prev`. This is because
            # we're doing a root-find and know that we're aiming to get close to zero.
            # Note that this does mean that the `rtol` is ignored in f-space, and only
            # `atol` matters.
            terminate = cauchy_termination(
                self.rtol,
                self.atol,
                self.norm,
                y,
                state.diff,
                jtu.tree_map(jnp.zeros_like, state.f),
                state.f,
            )
            terminate_result = RESULTS.successful
        else:
            # TODO(kidger): perform only one iteration when solving a linear system!
            at_least_two = state.step >= 2
            rate = state.diffsize / state.diffsize_prev
            factor = state.diffsize * rate / (1 - rate)
            small = _small(state.diffsize)
            diverged = _diverged(rate)
            converged = _converged(factor, self.kappa)
            terminate = at_least_two & (small | diverged | converged)
            terminate_result = RESULTS.where(
                diverged, RESULTS.nonlinear_divergence, RESULTS.successful
            )
        linsolve_fail = state.result != RESULTS.successful
        result = RESULTS.where(linsolve_fail, state.result, terminate_result)
        terminate = linsolve_fail | terminate
        return terminate, result

    def postprocess(
        self,
        fn: Fn[Y, Out, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _NewtonChordState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        return y, aux, {}


class Newton(_NewtonChord[Y, Out, Aux]):
    """Newton's method for root finding. Also sometimes known as Newton--Raphson.

    Unlike the SciPy implementation of Newton's method, the Optimistix version also
    works for vector-valued (or PyTree-valued) `y`.

    This solver optionally accepts the following `options`:

    - `lower`: The lower bound on the hypercube which contains the root. Should be a
        PyTree of arrays each broadcastable to the corresponding element of `y`. The
        iterates of `y` will be clipped to this hypercube.
    - `upper`: The upper bound on the hypercube which contains the root. Should be a
        PyTree of arrays each broadcastable to the corresponding element of `y`. The
        iterates of `y` will be clipped to this hypercube.
    """

    _is_newton = True


class Chord(_NewtonChord[Y, Out, Aux]):
    """The Chord method of root finding.

    This is equivalent to the Newton method, except that the Jacobian is computed only
    once at the initial point `y0`, and then reused throughout the computation. This is
    a useful way to cheapen the solve, if `y0` is expected to be a good initial guess
    and the target function does not change too rapidly. (For example this is the
    standard technique used in implicit Runge--Kutta methods, when solving differential
    equations.)

    This solver optionally accepts the following `options`:

    - `lower`: The lower bound on the hypercube which contains the root. Should be a
        PyTree of arrays each broadcastable to the corresponding element of `y`. The
        iterates of `y` will be clipped to this hypercube.
    - `upper`: The upper bound on the hypercube which contains the root. Should be a
        PyTree of arrays each broadcastable to the corresponding element of `y`. The
        iterates of `y` will be clipped to this hypercube.
    """

    _is_newton = False


_init_doc = """**Arguments:**

- `rtol`: Relative tolerance for terminating the solve.
- `atol`: Absolute tolerance for terminating the solve.
- `norm`: The norm used to determine the difference between two iterates in the 
    convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
    includes three built-in norms: [`optimistix.max_norm`][],
    [`optimistix.rms_norm`][], and [`optimistix.two_norm`][]. Only used when
    `cauchy_termination=True`.
- `kappa`: A tolerance for early convergence check when `cauchy_termination=False`.
- `linear_solver`: The linear solver used to compute the Newton step.
- `cauchy_termination`: When `True`, use the Cauchy termination condition, that
    two adjacent iterates should have a small difference between them. This is usually
    the standard choice when solving general root finding problems. When `False`, use
    a procedure which attempts to detect slow convergence, and quickly fail the solve
    if so. This is useful when iteratively performing the root-find, refining the
    target problem for those which fail. This comes up when solving differential
    equations with adaptive step sizing and implicit solvers. The exact procedure is as
    described in Section IV.8 of Hairer & Wanner, "Solving Ordinary Differential
    Equations II".
"""

Newton.__init__.__doc__ = _init_doc
Chord.__init__.__doc__ = _init_doc

optimistix/_solver/newton_chord.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, cast, Generic, Union

import equinox as eqx
import jax.numpy as jnp
from equinox.internal import ω
from jaxtyping import Array, PyTree, Scalar

from .._custom_types import Aux, Y
from .._misc import (
    max_norm,
    sum_squares,
    tree_dot,
    tree_full_like,
    tree_where,
)
from .._search import AbstractDescent, AbstractSearch, FunctionInfo
from .._solution import RESULTS
from .backtracking import BacktrackingArmijo
from .gradient_methods import AbstractGradientDescent


def polak_ribiere(grad_vector: Y, grad_prev: Y, y_diff_prev: Y) -> Scalar:
    """The Polak--Ribière formula for β. Used with [`optimistix.NonlinearCG`][] and
    [`optimistix.NonlinearCGDescent`][]."""
    del y_diff_prev
    numerator = tree_dot(grad_vector, (grad_vector**ω - grad_prev**ω).ω)
    denominator = sum_squares(grad_prev)
    # This triggers under two scenarios: (a) at the very start, for which our
    # `grad_prev` is initialised at zero, and (b) at convergence, for which we no longer
    # have a gradient. In either case we set β=0 to revert to just gradient descent.
    pred = denominator > jnp.finfo(denominator.dtype).eps
    safe_denom = jnp.where(pred, denominator, 1)
    out = jnp.where(pred, jnp.clip(numerator / safe_denom, a_min=0), 0)
    return cast(Scalar, out)


def fletcher_reeves(grad: Y, grad_prev: Y, y_diff_prev: Y) -> Scalar:
    """The Fletcher--Reeves formula for β. Used with [`optimistix.NonlinearCG`][] and
    [`optimistix.NonlinearCGDescent`][]."""
    del y_diff_prev
    numerator = sum_squares(grad)
    denominator = sum_squares(grad_prev)
    # Triggers at initialisation and convergence, as above.
    pred = denominator > jnp.finfo(denominator.dtype).eps
    safe_denom = jnp.where(pred, denominator, 1)
    return jnp.where(pred, numerator / safe_denom, 0)


def hestenes_stiefel(grad: Y, grad_prev: Y, y_diff_prev: Y) -> Scalar:
    """The Hestenes--Stiefel formula for β. Used with [`optimistix.NonlinearCG`][] and
    [`optimistix.NonlinearCGDescent`][]."""
    grad_diff = (grad**ω - grad_prev**ω).ω
    numerator = tree_dot(grad, grad_diff)
    denominator = -tree_dot(y_diff_prev, grad_diff)
    # Triggers at initialisation and convergence, as above.
    pred = jnp.abs(denominator) > jnp.finfo(denominator.dtype).eps
    safe_denom = jnp.where(pred, denominator, 1)
    return jnp.where(pred, numerator / safe_denom, 0)


def dai_yuan(grad: Y, grad_prev: Y, y_diff_prev: Y) -> Scalar:
    """The Dai--Yuan formula for β. Used with [`optimistix.NonlinearCG`][] and
    [`optimistix.NonlinearCGDescent`][]."""
    numerator = sum_squares(grad)
    denominator = -tree_dot(y_diff_prev, (grad**ω - grad_prev**ω).ω)
    # Triggers at initialisation and convergence, as above.
    pred = jnp.abs(denominator) > jnp.finfo(denominator.dtype).eps
    safe_denom = jnp.where(pred, denominator, 1)
    return jnp.where(pred, numerator / safe_denom, 0)


class _NonlinearCGDescentState(eqx.Module, Generic[Y]):
    y_diff: Y
    grad: Y


class NonlinearCGDescent(
    AbstractDescent[
        Y,
        Union[
            FunctionInfo.EvalGrad,
            FunctionInfo.EvalGradHessian,
            FunctionInfo.EvalGradHessianInv,
            FunctionInfo.ResidualJac,
        ],
        _NonlinearCGDescentState,
    ]
):
    """The nonlinear conjugate gradient step."""

    method: Callable[[Y, Y, Y], Scalar]

    def init(
        self,
        y: Y,
        f_info_struct: Union[
            FunctionInfo.EvalGrad,
            FunctionInfo.EvalGradHessian,
            FunctionInfo.EvalGradHessianInv,
            FunctionInfo.ResidualJac,
        ],
    ) -> _NonlinearCGDescentState:
        del f_info_struct
        return _NonlinearCGDescentState(
            y_diff=tree_full_like(y, 0),
            grad=tree_full_like(y, 0),
        )

    def query(
        self,
        y: Y,
        f_info: Union[
            FunctionInfo.EvalGrad,
            FunctionInfo.EvalGradHessian,
            FunctionInfo.EvalGradHessianInv,
            FunctionInfo.ResidualJac,
        ],
        state: _NonlinearCGDescentState,
    ) -> _NonlinearCGDescentState:
        if not isinstance(
            f_info,
            (
                FunctionInfo.EvalGrad,
                FunctionInfo.EvalGradHessian,
                FunctionInfo.EvalGradHessianInv,
                FunctionInfo.ResidualJac,
            ),
        ):
            raise ValueError(
                "Cannot use `NonlinearCGDescent` with this solver. This is because "
                "`NonlinearCGDescent` requires gradients of the target function, but "
                "this solver does not evaluate such gradients."
            )
        # On the very first step, we have (from `optim_init`) that `state.y_diff = 0`
        # and `state.grad = 0`. For all methods, this implies that `beta = 0`, so
        # `nonlinear_cg_diretion = neg_grad`, and (as desired) we take just a gradient
        # step on the first step.
        # Furthermore, the same mechanism handles convergence: once
        # `state.{grad, y_diff} = 0`, i.e. our previous step hit a local minima, then
        # on this next step we'll again just use gradient descent, and stop.
        beta = self.method(f_info.grad, state.grad, state.y_diff)
        neg_grad = (-(f_info.grad**ω)).ω
        nonlinear_cg_direction = (neg_grad**ω + beta * state.y_diff**ω).ω
        # Check if this is a descent direction. Use gradient descent if it isn't.
        y_diff = tree_where(
            tree_dot(f_info.grad, nonlinear_cg_direction) < 0,
            nonlinear_cg_direction,
            neg_grad,
        )
        return _NonlinearCGDescentState(y_diff=y_diff, grad=f_info.grad)

    def step(
        self, step_size: Scalar, state: _NonlinearCGDescentState
    ) -> tuple[Y, RESULTS]:
        return (step_size * state.y_diff**ω).ω, RESULTS.successful


NonlinearCGDescent.__init__.__doc__ = """**Arguments:**

- `method`: A callable `method(vector, vector_prev, diff_prev)` describing how to
    calculate the beta parameter of nonlinear CG. Each of these inputs has the meaning
    described above. The "beta parameter" is the sake as can be described as e.g. the
    β_n value
    [on Wikipedia](https://en.wikipedia.org/wiki/Nonlinear_conjugate_gradient_method).
    In practice Optimistix includes four built-in methods:
    [`optimistix.polak_ribiere`][], [`optimistix.fletcher_reeves`][],
    [`optimistix.hestenes_stiefel`][], and [`optimistix.dai_yuan`][].
"""


class NonlinearCG(AbstractGradientDescent[Y, Aux]):
    """The nonlinear conjugate gradient method."""

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    descent: NonlinearCGDescent[Y]
    search: AbstractSearch[Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any]

    def __init__(
        self,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree[Array]], Scalar] = max_norm,
        method: Callable[[Y, Y, Y], Scalar] = polak_ribiere,
        # TODO(raderj): replace the default line search with something better.
        search: AbstractSearch[
            Y, FunctionInfo.EvalGrad, FunctionInfo.Eval, Any
        ] = BacktrackingArmijo(decrease_factor=0.5, slope=0.1),
    ):
        """**Arguments:**

        - `rtol`: Relative tolerance for terminating solve.
        - `atol`: Absolute tolerance for terminating solve.
        - `norm`: The norm used to determine the difference between two iterates in the
            convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
            includes three built-in norms: [`optimistix.max_norm`][],
            [`optimistix.rms_norm`][], and [`optimistix.two_norm`][].
        - `method`: The function which computes `beta` in `NonlinearCG`. Defaults to
            `polak_ribiere`. Optimistix includes four built-in methods:
            [`optimistix.polak_ribiere`][], [`optimistix.fletcher_reeves`][],
            [`optimistix.hestenes_stiefel`][], and [`optimistix.dai_yuan`][], but any
            function `(Y, Y, Y) -> Scalar` will work.
        - `search`: The (line) search to use at each step.
        """
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.descent = NonlinearCGDescent(method=method)
        self.search = search

optimistix/_solver/nonlinear_cg.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections.abc import Callable
from typing import Any, cast

import equinox as eqx
import equinox.internal as eqxi
import jax
import jax.numpy as jnp
from jaxtyping import Array, Bool, Int, PyTree, Scalar

from .._custom_types import Aux, Fn, Y
from .._minimise import AbstractMinimiser
from .._misc import cauchy_termination, max_norm, verbose_print
from .._solution import RESULTS


class _OptaxState(eqx.Module):
    step: Int[Array, ""]
    f: Scalar
    opt_state: Any
    terminate: Bool[Array, ""]


class OptaxMinimiser(AbstractMinimiser[Y, Aux, _OptaxState]):
    """A wrapper to use Optax first-order gradient-based optimisers with
    [`optimistix.minimise`][].
    """

    optim: "optax.GradientTransformation"  # pyright: ignore  # noqa: F821
    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    verbose: frozenset[str]

    def __init__(
        self,
        optim,
        rtol: float,
        atol: float,
        norm: Callable[[PyTree], Scalar] = max_norm,
        verbose: frozenset[str] = frozenset(),
    ):
        """**Arguments:**

        - `optim`: The Optax optimiser to use.
        - `rtol`: Relative tolerance for terminating the solve. Keyword only argument.
        - `atol`: Absolute tolerance for terminating the solve. Keyword only argument.
        - `norm`: The norm used to determine the difference between two iterates in the
            convergence criteria. Should be any function `PyTree -> Scalar`. Optimistix
            includes three built-in norms: [`optimistix.max_norm`][],
            [`optimistix.rms_norm`][], and [`optimistix.two_norm`][]. Keyword only
            argument.
        - `verbose`: Whether to print out extra information about how the solve is
            proceeding. Should be a frozenset of strings, specifying what information to
            print out. Valid entries are `step`, `loss`, `y`. For example
            `verbose=frozenset({"step", "loss"})`.
        """
        # See https://github.com/deepmind/optax/issues/577: Optax has an issue in which
        # it doesn't use pytrees correctly.
        self.optim = eqxi.closure_to_pytree(optim)
        self.rtol = rtol
        self.atol = atol
        self.norm = norm
        self.verbose = verbose

    def init(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        f_struct: PyTree[jax.ShapeDtypeStruct],
        aux_struct: PyTree[jax.ShapeDtypeStruct],
        tags: frozenset[object],
    ) -> _OptaxState:
        del fn, args, options, aux_struct
        opt_state = self.optim.init(y)
        maxval = jnp.array(jnp.finfo(f_struct.dtype).max, f_struct.dtype)
        return _OptaxState(
            step=jnp.array(0), f=maxval, opt_state=opt_state, terminate=jnp.array(False)
        )

    def step(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _OptaxState,
        tags: frozenset[object],
    ) -> tuple[Y, _OptaxState, Aux]:
        del options
        (f, aux), grads = eqx.filter_value_and_grad(fn, has_aux=True)(y, args)
        f = cast(Array, f)
        if len(self.verbose) > 0:
            verbose_print(
                ("step" in self.verbose, "Step", state.step),
                ("loss" in self.verbose, "Loss", f),
                ("y" in self.verbose, "y", y),
            )
        updates, new_opt_state = self.optim.update(grads, state.opt_state)
        new_y = eqx.apply_updates(y, updates)
        terminate = cauchy_termination(
            self.rtol,
            self.atol,
            self.norm,
            y,
            updates,
            f,
            f - state.f,
        )
        new_state = _OptaxState(
            step=state.step + 1, f=f, opt_state=new_opt_state, terminate=terminate
        )
        return new_y, new_state, aux

    def terminate(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        args: PyTree,
        options: dict[str, Any],
        state: _OptaxState,
        tags: frozenset[object],
    ) -> tuple[Bool[Array, ""], RESULTS]:
        del fn, args, options
        return state.terminate, RESULTS.successful

    def postprocess(
        self,
        fn: Fn[Y, Scalar, Aux],
        y: Y,
        aux: Aux,
        args: PyTree,
        options: dict[str, Any],
        state: _OptaxState,
        tags: frozenset[object],
        result: RESULTS,
    ) -> tuple[Y, Aux, dict[str, Any]]:
        return y, aux, {}

optimistix/_solver/optax.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
from typing import TypeVar, Union
from typing_extensions import TypeAlias

import equinox as eqx
import jax.numpy as jnp
from equinox import AbstractVar
from equinox.internal import ω
from jaxtyping import Array, Bool, Scalar, ScalarLike

from .._custom_types import Y
from .._misc import (
    sum_squares,
    tree_dot,
)
from .._search import AbstractSearch, FunctionInfo
from .._solution import RESULTS


class _TrustRegionState(eqx.Module):
    step_size: Scalar


_FnInfo = TypeVar("_FnInfo", bound=FunctionInfo)
_FnEvalInfo: TypeAlias = FunctionInfo


class _AbstractTrustRegion(AbstractSearch[Y, _FnInfo, _FnEvalInfo, _TrustRegionState]):
    """The abstract base class of the trust-region update algorithm.

    Trust region line searches compute the ratio
    `true_reduction/predicted_reduction`, where `true_reduction` is the decrease in `fn`
    between `y` and `new_y`, and `predicted_reduction` is how much we expected the
    function to decrease using an approximation to `fn`.

    The trust-region ratio determines whether to accept or reject a step and the
    next choice of step-size. Specifically:

    - reject the step and decrease stepsize if the ratio is smaller than a
        cutoff `low_cutoff`
    - accept the step and increase the step-size if the ratio is greater than
        another cutoff `high_cutoff` with `low_cutoff < high_cutoff`.
    - else, accept the step and make no change to the step-size.
    """

    high_cutoff: AbstractVar[ScalarLike]
    low_cutoff: AbstractVar[ScalarLike]
    high_constant: AbstractVar[ScalarLike]
    low_constant: AbstractVar[ScalarLike]

    def __post_init__(self):
        # You would not expect `self.low_cutoff` or `self.high_cutoff` to
        # be below zero, but this is technically not incorrect so we don't
        # require it.
        self.low_cutoff, self.high_cutoff = eqx.error_if(  # pyright: ignore
            (self.low_cutoff, self.high_cutoff),
            self.low_cutoff > self.high_cutoff,  # pyright: ignore
            "`low_cutoff` must be below `high_cutoff` in `ClassicalTrustRegion`",
        )
        self.low_constant = eqx.error_if(  # pyright: ignore
            self.low_constant,
            self.low_constant < 0,  # pyright: ignore
            "`low_constant` must be greater than `0` in `ClassicalTrustRegion`",
        )
        self.high_constant = eqx.error_if(  # pyright: ignore
            self.high_constant,
            self.high_constant < 0,  # pyright: ignore
            "`high_constant` must be greater than `0` in `ClassicalTrustRegion`",
        )

    @abc.abstractmethod
    def predict_reduction(self, y_diff: Y, f_info: _FnInfo) -> Scalar:
        ...

    def init(self, y: Y, f_info_struct: _FnInfo) -> _TrustRegionState:
        del f_info_struct
        return _TrustRegionState(step_size=jnp.array(1.0))

    def step(
        self,
        first_step: Bool[Array, ""],
        y: Y,
        y_eval: Y,
        f_info: _FnInfo,
        f_eval_info: _FnEvalInfo,
        state: _TrustRegionState,
    ) -> tuple[Scalar, Bool[Array, ""], RESULTS, _TrustRegionState]:
        y_diff = (y_eval**ω - y**ω).ω
        predicted_reduction = self.predict_reduction(y_diff, f_info)
        # We never actually compute the ratio
        # `true_reduction/predicted_reduction`. Instead, we rewrite the conditions as
        # `true_reduction < const * predicted_reduction` instead, where the inequality
        # flips because we assume `predicted_reduction` is negative.
        # This avoids an expensive division.
        f_min = f_info.as_min()
        f_min_eval = f_eval_info.as_min()
        f_min_diff = f_min_eval - f_min  # This number is probably negative
        accept = f_min_diff <= self.low_cutoff * predicted_reduction
        good = f_min_diff < self.high_cutoff * predicted_reduction
        good = good & (predicted_reduction < 0)

        good = good & jnp.invert(first_step)
        accept = accept | first_step
        mul = jnp.where(good, self.high_constant, 1)
        mul = jnp.where(accept, mul, self.low_constant)
        new_step_size = mul * state.step_size
        new_state = _TrustRegionState(step_size=new_step_size)
        return new_step_size, accept, RESULTS.successful, new_state


class ClassicalTrustRegion(
    _AbstractTrustRegion[
        Y, Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac]
    ]
):
    """The classic trust-region update algorithm which uses a quadratic approximation of
    the objective function to predict reduction.

    Building a quadratic approximation requires an approximation to the Hessian of the
    overall minimisation function. This means that trust region is suitable for use with
    least-squares algorithms (which make the Gauss--Newton approximation
    Hessian~Jac^T J) and for quasi-Newton minimisation algorithms like
    [`optimistix.BFGS`][]. (An error will be raised if you use this with an incompatible
    solver.)
    """

    # This choice of default parameters comes from Gould et al.
    # "Sensitivity of trust region algorithms to their parameters."
    high_cutoff: ScalarLike = 0.99
    low_cutoff: ScalarLike = 0.01
    high_constant: ScalarLike = 3.5
    low_constant: ScalarLike = 0.25

    def predict_reduction(
        self,
        y_diff: Y,
        f_info: Union[FunctionInfo.EvalGradHessian, FunctionInfo.ResidualJac],
    ) -> Scalar:
        """Compute the expected decrease in loss from taking the step `y_diff`.

        The true reduction is
        ```
        fn(y0 + y_diff) - fn(y0)
        ```
        so if `B` is the approximation to the Hessian coming from the quasi-Newton
        method at `y`, and `g` is the gradient of `fn` at `y`, then the predicted
        reduction is
        ```
        g^T y_diff + 1/2 y_diff^T B y_diff
        ```

        **Arguments**:

        - `y_diff`: the proposed step by the descent method.
        - `deriv_info`: the derivative information (on the gradient and Hessian)
            provided by the outer loop.

        **Returns**:

        The expected decrease in loss from moving from `y0` to `y0 + y_diff`.
        """

        if isinstance(f_info, FunctionInfo.EvalGradHessian):
            # Minimisation algorithm. Directly compute the quadratic approximation.
            return tree_dot(
                y_diff,
                (f_info.grad**ω + 0.5 * f_info.hessian.mv(y_diff) ** ω).ω,
            )
        elif isinstance(f_info, FunctionInfo.ResidualJac):
            # Least-squares algorithm. So instead of considering fn (which returns the
            # residuals), instead consider `0.5*fn(y)^2`, and then apply the logic as
            # for minimisation.
            # We get that `g = J^T f0` and `B = J^T J + dJ/dx^T J`.
            # (Here, `f0 = fn(y0)` are the residuals, and `J = dfn/dy(y0)` is the
            # Jacobian of the residuals wrt y.)
            # Then neglect the second term in B (the usual Gauss--Newton approximation)
            # and complete the square.
            # We find that the predicted reduction is
            # `0.5 * ((J y_diff + f0)^T (J y_diff + f0) - f0^T f0)`
            # and this is what is written below.
            #
            # The reason we go through this hassle is because this now involves only
            # a single Jacobian-vector product, rather than the three we would have to
            # make by naively substituting `B = J^T J `and `g = J^T f0` into the general
            # algorithm used for minimisation.
            rtr = sum_squares(f_info.residual)
            jacobian_term = sum_squares(
                (f_info.jac.mv(y_diff) ** ω + f_info.residual**ω).ω
            )
            return 0.5 * (jacobian_term - rtr)
        else:
            raise ValueError(
                "Cannot use `ClassicalTrustRegion` with this solver. This is because "
                "`ClassicalTrustRegion` requires (an approximation to) the Hessian of "
                "the target function, but this solver does not make any estimate of "
                "that information."
            )


# When using a gradient-based method, `LinearTrustRegion` is actually a variant of
# `BacktrackingArmijo`. The linear predicted reduction is the same as the Armijo
# condition. The difference is that unlike standard backtracking,
# `LinearTrustRegion` chooses its next step size based on how well it did in the
# previous iteration.
class LinearTrustRegion(
    _AbstractTrustRegion[
        Y,
        Union[
            FunctionInfo.EvalGrad,
            FunctionInfo.EvalGradHessian,
            FunctionInfo.EvalGradHessianInv,
            FunctionInfo.ResidualJac,
        ],
    ]
):
    """The trust-region update algorithm which uses a linear approximation of
    the objective function to predict reduction.

    Generally speaking you should prefer [`optimistix.ClassicalTrustRegion`][], unless
    you happen to be using a solver (e.g. a non-quasi-Newton minimiser) with which that
    is incompatible.
    """

    # This choice of default parameters comes from Gould et al.
    # "Sensitivity of trust region algorithms to their parameters."
    high_cutoff: ScalarLike = 0.99
    low_cutoff: ScalarLike = 0.01
    high_constant: ScalarLike = 3.5
    low_constant: ScalarLike = 0.25

    def predict_reduction(
        self,
        y_diff: Y,
        f_info: Union[
            FunctionInfo.EvalGrad,
            FunctionInfo.EvalGradHessian,
            FunctionInfo.EvalGradHessianInv,
            FunctionInfo.ResidualJac,
        ],
    ) -> Scalar:
        """Compute the expected decrease in loss from taking the step `y_diff`.

        The true reduction is
        ```
        fn(y0 + y_diff) - fn(y0)
        ```
        so if `g` is the gradient of `fn` at `y`, then the predicted reduction is
        ```
        g^T y_diff
        ```

        **Arguments**:

        - `y_diff`: the proposed step by the descent method.
        - `deriv_info`: the derivative information (on the gradient and Hessian)
            provided by the outer loop.

        **Returns**:

        The expected decrease in loss from moving from `y0` to `y0 + y_diff`.
        """

        if isinstance(
            f_info,
            (
                FunctionInfo.EvalGrad,
                FunctionInfo.EvalGradHessian,
                FunctionInfo.EvalGradHessianInv,
                FunctionInfo.ResidualJac,
            ),
        ):
            return tree_dot(f_info.grad, y_diff)
        else:
            raise ValueError(
                "Cannot use `LinearTrustRegion` with this solver. This is because "
                "`LinearTrustRegion` requires gradients of the target function, but "
                "this solver does not evaluate such gradients."
            )


_init_doc = """In the following, `ratio` refers to the ratio
`true_reduction/predicted_reduction`.

**Arguments**:

- `high_cutoff`: the cutoff such that `ratio > high_cutoff` will accept the step
and increase the step-size on the next iteration.
- `low_cutoff`: the cutoff such that `ratio < low_cutoff` will reject the step
and decrease the step-size on the next iteration.
- `high_constant`: when `ratio > high_cutoff`, multiply the previous step-size by
high_constant`.
- `low_constant`: when `ratio < low_cutoff`, multiply the previous step-size by
low_constant`.
"""

LinearTrustRegion.__init__.__doc__ = _init_doc
ClassicalTrustRegion.__init__.__doc__ = _init_doc

optimistix/_solver/trust_region.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .._ad import implicit_jvp as implicit_jvp

optimistix/internal/__init__.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import dataclasses
import random

import jax
import jax.random as jr
import pytest
from jaxtyping import PRNGKeyArray


jax.config.update("jax_enable_x64", True)


# This offers reproducability -- the initial seed is printed in the repr so we can see
# it when a test fails.
# Note the `eq=False`, which means that `_GetKey `objects have `__eq__` and `__hash__`
# based on object identity.
@dataclasses.dataclass(eq=False)
class _GetKey:
    seed: int
    call: int
    key: PRNGKeyArray

    def __init__(self, seed: int):
        self.seed = seed
        self.call = 0
        self.key = jr.PRNGKey(seed)

    def __call__(self):
        self.call += 1
        return jr.fold_in(self.key, self.call)


@pytest.fixture
def getkey():
    return _GetKey(random.randint(0, 2**31 - 1))

tests/conftest.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools as ft
import operator
import random
from collections.abc import Callable
from typing import Any, TypeVar

import equinox as eqx
import equinox.internal as eqxi
import jax
import jax.flatten_util as jfu
import jax.numpy as jnp
import jax.random as jr
import jax.tree_util as jtu
import lineax as lx
import numpy as np
import optax
from equinox.internal import ω
from jaxtyping import Array, PyTree, Scalar

import optimistix as optx


Y = TypeVar("Y")
Out = TypeVar("Out")
Aux = TypeVar("Aux")


def getkey():
    return jr.PRNGKey(random.randint(0, 2**31 - 1))


def _tree_allclose(x, y, **kwargs):
    if type(x) is not type(y):
        return False
    if isinstance(x, jnp.ndarray):  # pyright: ignore
        if jnp.issubdtype(x.dtype, jnp.inexact):
            return (
                x.shape == y.shape
                and x.dtype == y.dtype
                and jnp.allclose(x, y, **kwargs)
            )
        else:
            return x.shape == y.shape and x.dtype == y.dtype and jnp.all(x == y)
    elif isinstance(x, np.ndarray):
        if np.issubdtype(x.dtype, np.inexact):
            return (
                x.shape == y.shape
                and x.dtype == y.dtype
                and np.allclose(x, y, **kwargs)
            )
        else:
            return x.shape == y.shape and x.dtype == y.dtype and np.all(x == y)
    elif isinstance(x, jax.ShapeDtypeStruct):
        assert x.shape == y.shape and x.dtype == y.dtype
    else:
        return x == y


def tree_allclose(x, y, **kwargs):
    """As `jnp.allclose`, except:
    - It also supports PyTree arguments.
    - It mandates that shapes match as well (no broadcasting)
    """
    same_structure = jtu.tree_structure(x) == jtu.tree_structure(y)
    allclose = ft.partial(_tree_allclose, **kwargs)
    return same_structure and jtu.tree_reduce(
        operator.and_, jtu.tree_map(allclose, x, y), True
    )


def finite_difference_jvp(fn, primals, tangents, eps=None, **kwargs):
    assert jax.config.jax_enable_x64  # pyright: ignore
    out = fn(*primals, **kwargs)
    # Choose ε to trade-off truncation error and floating-point rounding error.
    max_leaves = [jnp.max(jnp.abs(p)) for p in jtu.tree_leaves(primals)] + [1]
    scale = jnp.max(jnp.stack(max_leaves))
    if eps is None:
        ε = np.sqrt(np.finfo(np.float64).eps) * scale
    else:
        # Sometimes we may want to set it manually. finite_difference_jvp is actually
        # pretty inaccurate for nonlinear solves, as these are themselves often only
        # done to a tolerance of 1e-8 or so: the primal pass is already noisy at about
        # the scale of ε.
        ε = eps
    primals_ε = (ω(primals) + ε * ω(tangents)).ω
    out_ε = fn(*primals_ε, **kwargs)
    tangents_out = jtu.tree_map(lambda x, y: (x - y) / ε, out_ε, out)
    # We actually return the perturbed primal.
    # This should still be within all tolerance checks, and means that we have aceesss
    # to both the true primal and the perturbed primal when debugging.
    return out_ε, tangents_out


#
# NOTE: `GN` is shorthand for `gauss_newton`. We want to be sure we test every
# branch of `GN=True` and `GN=False` for all of these solvers.
#
# UNCONSTRAINED MIN
#
# SOLVERS:
# - Gauss-Newton (LM)
# - BFGS
# - GradientDescent
# - Optax
# - NonlinearCG
#
# LINE SEARCHES:
# - BacktrackingArmijo
# - ClassicalTrustRegion
# - LearningRate
#
# DESCENTS:
# - DampedNewtonDescent(GN=...)
# - IndirectDampedNewtonDescent(GN=...)
# - DoglegDescent(GN=...)
# - NewtonDescent(GN=...)
# - Gradient (Not appropriate for GN!)
# - NonlinearCGDescent (Not appropriate for GN!)
#

#
# We define a bunch of helper optimisers which exist to test all code paths, but which
# are not default solvers because they feature methods only advanced users are likely
# to use.
#


class DoglegMax(optx.AbstractGaussNewton[Y, Out, Aux]):
    """Dogleg with trust region shape given by the max norm instead of the two norm."""

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar]
    descent: optx.DoglegDescent[Y]
    search: optx.ClassicalTrustRegion[Y]
    verbose: frozenset[str]

    def __init__(
        self,
        rtol: float,
        atol: float,
    ):
        self.rtol = rtol
        self.atol = atol
        self.norm = optx.max_norm
        self.descent = optx.DoglegDescent(
            linear_solver=lx.AutoLinearSolver(well_posed=False),
            root_finder=optx.Bisection(rtol=0.001, atol=0.001),
            trust_region_norm=optx.max_norm,
        )
        self.search = optx.ClassicalTrustRegion()
        self.verbose = frozenset()


class BFGSDampedNewton(optx.AbstractBFGS):
    """BFGS Hessian + direct Levenberg Marquardt update."""

    rtol: float
    atol: float
    norm: Callable = optx.max_norm
    use_inverse: bool = False
    search: optx.AbstractSearch = optx.ClassicalTrustRegion()
    descent: optx.AbstractDescent = optx.DampedNewtonDescent()


class BFGSIndirectDampedNewton(optx.AbstractBFGS):
    """BFGS Hessian + indirect Levenberg Marquardt update."""

    rtol: float
    atol: float
    norm: Callable = optx.max_norm
    use_inverse: bool = False
    search: optx.AbstractSearch = optx.ClassicalTrustRegion()
    descent: optx.AbstractDescent = optx.IndirectDampedNewtonDescent()


class BFGSDogleg(optx.AbstractBFGS):
    """BFGS Hessian + dogleg update."""

    rtol: float
    atol: float
    norm: Callable = optx.max_norm
    use_inverse: bool = False
    search: optx.AbstractSearch = optx.ClassicalTrustRegion()
    descent: optx.AbstractDescent = optx.DoglegDescent(linear_solver=lx.SVD())


class BFGSBacktracking(optx.AbstractBFGS):
    """Standard BFGS + backtracking line search."""

    rtol: float
    atol: float
    norm: Callable = optx.max_norm
    use_inverse: bool = False
    search: optx.AbstractSearch = optx.BacktrackingArmijo()
    descent: optx.AbstractDescent = optx.NewtonDescent()


class BFGSTrustRegion(optx.AbstractBFGS):
    """Standard BFGS + classical trust region update."""

    rtol: float
    atol: float
    norm: Callable = optx.max_norm
    use_inverse: bool = False
    search: optx.AbstractSearch = optx.LinearTrustRegion()
    descent: optx.AbstractDescent = optx.NewtonDescent()


atol = rtol = 1e-8
_lsqr_only = (
    optx.LevenbergMarquardt(rtol, atol),
    optx.IndirectLevenbergMarquardt(rtol, atol),
    optx.GaussNewton(rtol, atol, linear_solver=lx.AutoLinearSolver(well_posed=False)),
    optx.Dogleg(rtol, atol, linear_solver=lx.AutoLinearSolver(well_posed=False)),
    DoglegMax(rtol, atol),
)


atol = rtol = 1e-8
minimisers = (
    optx.NelderMead(rtol, atol),
    optx.BFGS(rtol, atol, use_inverse=False),
    optx.BFGS(rtol, atol, use_inverse=True),
    BFGSDampedNewton(rtol, atol),
    BFGSIndirectDampedNewton(rtol, atol),
    # Tighter tolerance needed to have bfgs_dogleg pass the JVP test.
    BFGSDogleg(1e-10, 1e-10),
    BFGSBacktracking(rtol, atol, use_inverse=False),
    BFGSBacktracking(rtol, atol, use_inverse=True),
    BFGSTrustRegion(rtol, atol, use_inverse=False),
    BFGSTrustRegion(rtol, atol, use_inverse=True),
    optx.GradientDescent(1.5e-2, rtol, atol),
    # Tighter tolerance needed to have NonlinearCG pass the JVP test.
    optx.NonlinearCG(1e-10, 1e-10),
    optx.OptaxMinimiser(optax.adam(learning_rate=3e-3), rtol=rtol, atol=atol),
)

# the minimisers can handle least squares problems, but the least squares
# solvers cannot handle general minimisation problems.
least_squares_optimisers = _lsqr_only + minimisers
# Remove ones that work, but are just pretty bad!
least_squares_optimisers = [
    x
    for x in least_squares_optimisers
    if not isinstance(x, (optx.GradientDescent, optx.NonlinearCG))
]

#
# MINIMISATION PROBLEMS
#
# Some standard test functions for nonlinear minimisation. Many of these can be found
# on the wikipedia page: https://en.wikipedia.org/wiki/Test_functions_for_optimization
# (Wiki.) The rest can be found in Testing Unconstrained Optimization Software (TUOS)
# by Moré, Garbow, and Hillstrom.
# See also the Fortran package TEST_NLS:
# https://people.math.sc.edu/Burkardt/f_src/test_nls/test_nls.html.
#


def bowl(tree: PyTree[Array], args: Array):
    # Trivial quadratic bowl smoke test for convergence.
    (y, _) = jfu.ravel_pytree(tree)
    matrix = args
    return y.T @ matrix @ y


def diagonal_quadratic_bowl(tree: PyTree[Array], args: PyTree[Array]):
    # A diagonal quadratic bowl smoke test for convergence.
    weight_vector = args
    return (ω(tree).call(jnp.square) * (0.1 + weight_vector**ω)).ω


def rosenbrock(tree: PyTree[Array], args: Scalar):
    # Wiki
    # least squares
    (y, _) = jfu.ravel_pytree(tree)
    const = args
    return (100 * (y[1:] - y[:-1]), const - y[:-1])


def _himmelblau(tree: PyTree[Array], args: PyTree):
    # Wiki
    (y, z) = tree
    const1, const2 = args
    term1 = ((ω(y).call(jnp.square) + z**ω - const1) ** 2).ω
    term2 = ((y**ω + ω(z).call(jnp.square) - const2) ** 2).ω
    return (term1**ω + term2**ω).ω


def matyas(tree: PyTree[Array], args: PyTree):
    # Wiki
    (y, z) = tree
    const1, const2 = args
    term1 = (const1 * (ω(y).call(jnp.square) + ω(z).call(jnp.square))).ω
    term2 = (const2 * y**ω * z**ω).ω
    return (term1**ω - term2**ω).ω


def _eggholder(tree: PyTree[Array], args: PyTree):
    # Wiki
    (y, z) = tree
    # the composition sin . sqrt . abs
    ssa = lambda x: jnp.sin(jnp.sqrt(jnp.abs(x)))
    term1 = (-(z**ω + 47.0) * (0.5 * y**ω + z**ω + 47.0).call(ssa)).ω
    term2 = (-(y**ω) * (y**ω - (z**ω - 47.0)).call(ssa)).ω
    return (term1**ω + term2**ω).ω


def beale(tree: PyTree[Array], args: PyTree):
    # Wiki
    (y, z) = tree
    const1, const2, const3 = args
    term1 = ((const1 - y**ω + y**ω * z**ω) ** 2).ω
    term2 = ((const2 - y**ω + y**ω * ω(z).call(jnp.square)) ** 2).ω
    term3 = ((const3 - y**ω + y**ω * ω(z).call(lambda x: x**3)) ** 2).ω
    return (term1**ω + term2**ω + term3**ω).ω


def variably_dimensioned(tree: PyTree[Array], args: PyTree):
    # TUOS eqn 25
    # least squares
    (y, _) = jfu.ravel_pytree(tree)
    const = args
    increasing = jnp.arange(1, jnp.size(y) + 1)
    term1 = y - const
    term2 = jnp.sum(increasing * (y - const))
    term3 = term2**2
    return (term1, term2, term3)


def trigonometric(tree: PyTree[Array], args: PyTree):
    # TUOS eqn 26
    # least squares
    (y, _) = jfu.ravel_pytree(tree)
    const = args
    sumcos = jnp.sum(jnp.cos(y))
    increasing = jnp.arange(1, jnp.size(y) + 1, dtype=jnp.float32)
    return jnp.size(y) - sumcos + increasing * (const - jnp.cos(y)) - jnp.sin(y)


def simple_nn(model_dynamic: PyTree[Array], args: PyTree):
    # args is the activation functions in
    # the MLP.
    (model_static, data) = args
    model = eqx.combine(model_dynamic, model_static)
    key = jr.PRNGKey(17)
    model_key, data_key = jr.split(key, 2)
    x = jnp.linspace(0, 1, 100)[..., None]
    y = data**2

    def loss(model, x, y):
        pred_y = eqx.filter_vmap(model)(x)
        return (pred_y - y) ** 2

    return loss(model, x, y)


#
# The MLP can be difficult for some of the solvers to optimise. Rather than set
# max_steps to a higher value and iterate for longer, we initialise the MLP
# closer to the minimum by explicitly setting the weights and biases.
#


def get_weights(model):
    layer1, layer2 = model.layers
    return layer1.weight, layer1.bias, layer2.weight, layer2.bias


ffn_init = eqx.nn.MLP(in_size=1, out_size=1, width_size=8, depth=1, key=jr.PRNGKey(17))
weight1 = jnp.array(
    [
        [3.39958394],
        [-0.00648151],
        [-0.57178741],
        [-0.9259611],
        [0.43010087],
        [0.74950293],
        [-0.35932756],
        [1.28393738],
    ]
)
bias1 = jnp.array(
    [
        -0.14702571,
        3.61400359,
        -1.01404942,
        0.3223137,
        2.19991259,
        -0.51860872,
        0.4717813,
        0.9449359,
    ]
)
weight2 = jnp.array(
    [
        [
            0.26476087,
            0.05913492,
            0.40830152,
            0.704758,
            0.64347694,
            0.85807767,
            0.54285774,
            0.05507022,
        ]
    ]
)
bias2 = jnp.array([-2.16578771])
new_ffn_init = eqx.tree_at(get_weights, ffn_init, (weight1, bias1, weight2, bias2))
ffn_dynamic, ffn_static = eqx.partition(new_ffn_init, eqx.is_array)

# _uncoupled_simple default args
diagonal_bowl_init = ({"a": 0.05 * jnp.ones((2, 3, 3))}, (0.05 * jnp.ones(2)))
leaves, treedef = jtu.tree_flatten(diagonal_bowl_init)
key = jr.PRNGKey(17)
diagonal_bowl_args = treedef.unflatten(
    [jr.normal(key, leaf.shape, leaf.dtype) ** 2 for leaf in leaves]
)

# neural net args
ffn_data = jnp.linspace(0, 1, 100)[..., None]
ffn_args = (ffn_static, ffn_data)


least_squares_fn_minima_init_args = (
    (
        diagonal_quadratic_bowl,
        jnp.array(0.0),
        diagonal_bowl_init,
        diagonal_bowl_args,
    ),
    (
        rosenbrock,
        jnp.array(0.0),
        [jnp.array(1.5), jnp.array(1.5)],
        jnp.array(1.0),
    ),
    (
        rosenbrock,
        jnp.array(0.0),
        (1.5 * jnp.ones((2, 4)), {"a": 1.5 * jnp.ones((2, 3, 2))}, ()),
        jnp.array(1.0),
    ),
    (simple_nn, jnp.array(0.0), ffn_dynamic, ffn_args),
    # Skipped for being hard!
    #
    # (
    #     variably_dimensioned,
    #     jnp.array(0.0),
    #     1 - jnp.arange(1, 11) / 10,
    #     jnp.array(1.0),
    # ),
    # (
    #     variably_dimensioned,
    #     jnp.array(0.0),
    #     (1 - jnp.arange(1, 7) / 10, {"a": (1 - jnp.arange(7, 11) / 10)}),
    #     jnp.array(1.0),
    # ),
    # (
    #     trigonometric,
    #     jnp.array(0.0),
    #     jnp.ones(70) / 70,
    #     jnp.array(1.0),
    # ),
    # (
    #     trigonometric,
    #     jnp.array(0.0),
    #     ((jnp.ones(40) / 70, (), {"a": jnp.ones(20) / 70}), jnp.ones(10) / 70),
    #     jnp.array(1.0),
    # ),
)

bowl_init = ({"a": 0.05 * jnp.ones((2, 3, 3))}, (0.05 * jnp.ones(2)))
(flatten_bowl, _) = jfu.ravel_pytree(bowl_init)
key = jr.PRNGKey(17)
matrix = jr.normal(key, (flatten_bowl.size, flatten_bowl.size))
diagonal_bowl_args = matrix.T @ matrix

minimisation_fn_minima_init_args = (
    (bowl, jnp.array(0.0), bowl_init, diagonal_bowl_args),
    (
        _himmelblau,
        jnp.array(0.0),
        [jnp.array(2.0), jnp.array(2.5)],
        (jnp.array(11.0), jnp.array(7.0)),
    ),
    (
        matyas,
        jnp.array(0.0),
        [jnp.array(6.0), jnp.array(6.0)],
        (jnp.array(0.26), jnp.array(0.48)),
    ),
    (
        beale,
        jnp.array(0.0),
        [jnp.array(2.0), jnp.array(0.0)],
        (jnp.array(1.5), jnp.array(2.25), jnp.array(2.625)),
    ),
)

# ROOT FIND/FIXED POINT PROBLEMS
#
# These are mostly derived from common problems arising in
# ordinary and partial differential equations. See
# Diffrax for more on differential equations:
# https://github.com/patrick-kidger/diffrax.
#
# - Simple smoke tests
# - MLP(y)
# - Nonlinear heat equation with Crank Nicholson
# - Implicit midpoint Runge-Kutta in y-space, f-space, and k-space
#


# SETUP
# Helper functions, most of which are single steps of a differential
# equation solve.
def _getsize(y: PyTree[Array]):
    return jtu.tree_reduce(lambda x, y: x + y, jtu.tree_map(jnp.size, y))


def _laplacian(y: PyTree[Array], dx: Scalar):
    (y, unflatten) = jfu.ravel_pytree(y)
    laplacian = jnp.zeros_like(y)
    laplacian = laplacian.at[1:-1].set((y[2:] + y[1:-1] + y[:-2]) / dx)
    return unflatten(y)


def _nonlinear_heat_pde_general(
    y0: PyTree[Array],
    f0: PyTree[Array],
    dx: Scalar,
    t0: Scalar,
    t1: Scalar,
    y: PyTree[Array],
    args: PyTree,
):
    # A single step time `t0` to time `t1` of the Crank Nicolson scheme applied to
    # the nonlinear heat equation:
    # `d y(t, x)/dt = (1 - y(t, x)) Δy(t, x)`
    # where `d/dt` is a partial derivative wrt time and `Δ` is the Laplacian.
    const = args
    stepsize = t1 - t0
    f_val = ((1 - y**ω) * _laplacian(y, dx) ** ω).ω
    return const * (y0**ω + 0.5 * stepsize * (f_val**ω + f0**ω)).ω


# Note that the midpoint methods below assume that `f` is autonomous.


def _midpoint_y_general(
    f: Callable[[PyTree[Array], Any], PyTree[Array]],
    y0: PyTree[Array],
    dt: Scalar,
    y: PyTree[Array],
    args: PyTree,
):
    # Solve an implicit midpoint Runge-Kutta step with fixed point iteration
    # in "y space," ie. the typical representation of a Runge-Kutta method
    f_new = f((0.5 * (y0**ω + y**ω)).ω, args)
    return (y0**ω + dt * f_new**ω).ω


def _midpoint_f_general(
    f: Callable[[PyTree[Array], Any], PyTree[Array]],
    y0: PyTree[Array],
    dt: Scalar,
    y: PyTree[Array],
    args: PyTree,
):
    # Solve an implicit Runge-Kutta step with fixed point iteration
    # in "f-space," ie. we apply `f` to both sides of the Runge-Kutta method
    # and do the fixed-point iteration in this space.
    return f((y0**ω + dt * y**ω).ω, args)


def _midpoint_k_general(
    f: Callable[[PyTree[Array], Any], PyTree[Array]],
    y0: PyTree[Array],
    dt: Scalar,
    y: PyTree[Array],
    args: PyTree,
):
    # Solve an implicit midpoint Runge-Kutta step with fixed point iteration
    # in "k-space," ie. we do the fixed point iterations directly on the `k` terms
    # (see https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods)
    # of the Runge-Kutta method.
    # Note that in practice this just means we are finding `f * dt`
    # instead of `f`.
    return (f((y0**ω + y**ω).ω, args) ** ω * dt).ω


class Robertson(eqx.Module):
    # Take directly from the Diffrax docs, but modified to accept and return
    # a PyTree and have no `t` argument.
    k1: float
    k2: float
    k3: float

    def __call__(self, y, args):
        (y, unflatten) = jfu.ravel_pytree(y)
        f0 = -self.k1 * y[0] + self.k3 * y[1] * y[2]
        f1 = self.k1 * y[0] - self.k2 * y[1] ** 2 - self.k3 * y[1] * y[2]
        f2 = self.k2 * y[1] ** 2
        return unflatten(jnp.stack([f0, f1, f2]))


# PROBLEMS
def _exponential(tree: PyTree[Array], args: PyTree):
    const = args
    return jtu.tree_map(lambda x: jnp.exp(-const * x), tree)


def _cos(tree: PyTree[Array], args: PyTree):
    const = args
    return jtu.tree_map(lambda x: jnp.cos(const * x), tree)


def _nn(tree: PyTree[Array], args: PyTree):
    (y, unflatten) = jfu.ravel_pytree(tree)
    size = _getsize(y)
    weight1, bias1, weight2, bias2 = args
    model = eqx.nn.MLP(
        in_size=size,
        out_size=size,
        width_size=8,
        depth=1,
        activation=jax.nn.softplus,  # more complex than relu
        key=jr.PRNGKey(42),  # doesn't matter! is overwritten immediately
    )
    model = eqx.tree_at(get_weights, model, (weight1, bias1, weight2, bias2))
    return unflatten(0.1 * model(y))


def _nonlinear_heat_pde(y: PyTree[Array], args: PyTree):
    # solve the nonlinear heat equation as above from "0" to "1" in a
    # single step.
    x = jnp.linspace(-1, 1, 100)
    dx = x[1] - x[0]
    y0 = x**2
    f0 = ((1 - y0**ω) * _laplacian(y0, dx) ** ω).ω
    return _nonlinear_heat_pde_general(
        y0, f0, dx, jnp.array(0.0), jnp.array(1.0), y, args
    )


def _midpoint_y_linear(tree: PyTree[Array], args: PyTree):
    (y, unflatten) = jfu.ravel_pytree(tree)
    matrix = args
    f = lambda x, _: matrix @ x
    y0 = ω(y).call(jnp.zeros_like).ω
    dt = jnp.array(1 / (2**4))
    midpoint = _midpoint_y_general(f, y0, dt, y, args)
    return unflatten(midpoint)


def _midpoint_f_linear(tree: PyTree[Array], args: PyTree):
    (y, unflatten) = jfu.ravel_pytree(tree)
    f = lambda x, args: jnp.dot(args, x)
    y0 = ω(y).call(jnp.zeros_like).ω
    dt = jnp.array(1 / (2**4))
    midpoint = _midpoint_f_general(f, y0, dt, y, args)
    if midpoint.size == 1:
        out = midpoint.reshape(1)
    else:
        out = midpoint
    return unflatten(out)


def _midpoint_k_linear(tree: PyTree[Array], args: PyTree):
    (y, unflatten) = jfu.ravel_pytree(tree)
    f = lambda x, args: jnp.dot(args, x)
    y0 = ω(y).call(jnp.zeros_like).ω
    dt = jnp.array(1 / (2**4))
    midpoint = _midpoint_k_general(f, y0, dt, y, args)
    if midpoint.size == 1:
        out = midpoint.reshape(1)
    else:
        out = midpoint
    return unflatten(out)


def _midpoint_y_nonlinear(y: PyTree[Array], args: PyTree):
    size = _getsize(y)
    assert size == 3
    const1, const2, const3 = args
    f = Robertson(const1, const2, const3)
    y0 = ω(y).call(jnp.zeros_like).ω
    dt = jnp.array(1e-8)
    return _midpoint_y_general(f, y0, dt, y, args)


def _midpoint_f_nonlinear(y: PyTree[Array], args: PyTree):
    size = _getsize(y)
    assert size == 3
    const1, const2, const3 = args
    f = Robertson(const1, const2, const3)
    y0 = ω(y).call(jnp.zeros_like).ω
    dt = jnp.array(1e-8)
    return _midpoint_f_general(f, y0, dt, y, args)


def _midpoint_k_nonlinear(y: PyTree[Array], args: PyTree):
    size = _getsize(y)
    assert size == 3
    const1, const2, const3 = args
    f = Robertson(const1, const2, const3)
    y0 = ω(y).call(jnp.zeros_like).ω
    dt = jnp.array(1e-8)
    return _midpoint_k_general(f, y0, dt, y, args)


def trivial(y: Array, args: PyTree):
    return y - 1


# ARGS FOR BISECTION AND FIXED POINT
ones_pytree = ({"a": 0.5 * jnp.ones((3, 2, 4))}, 0.5 * jnp.ones(4))
flat_ones, _ = jfu.ravel_pytree(ones_pytree)
size = flat_ones.size
key = jr.PRNGKey(19)
matrix = jr.normal(key, (size, size))
nn_model = eqx.nn.MLP(
    in_size=size,
    out_size=size,
    width_size=8,
    depth=1,
    activation=jax.nn.softplus,  # more complex than relu
    key=jr.PRNGKey(42),
)
nn_args = get_weights(nn_model)
robertson_args = (jnp.array(0.04), jnp.array(3e7), jnp.array(1e4))
ones_robertson = (jnp.ones(2), {"b": jnp.array(1.0)})
hundred = jnp.ones(100)
single = jnp.array(1.0)
bisection_fn_init_options_args = (
    (
        trivial,
        single,
        {"upper": jnp.array(2.0), "lower": jnp.array(0.5)},
        None,
    ),
    (
        _cos,
        single,
        {"upper": jnp.array(1.0), "lower": jnp.array(0.0)},
        jnp.array(1.0),
    ),
    (
        _exponential,
        single,
        {"upper": jnp.array(1.0), "lower": jnp.array(0.0)},
        jnp.array(1.0),
    ),
    (
        _midpoint_y_linear,
        single,
        {"upper": jnp.array(1.0), "lower": jnp.array(0.0)},
        jnp.array([0.5]),
    ),
    (
        _midpoint_f_linear,
        single,
        {"upper": jnp.array(1.0), "lower": jnp.array(0.0)},
        jnp.array([0.625]),
    ),
    (
        _midpoint_k_linear,
        single,
        {"upper": jnp.array(1.0), "lower": jnp.array(0.0)},
        jnp.array([-0.137]),
    ),
)

fixed_point_fn_init_args = (
    (_cos, ones_pytree, jnp.array(1.0)),
    (_exponential, ones_pytree, jnp.array(1.0)),
    (_nn, ones_pytree, nn_args),
    (_nonlinear_heat_pde, hundred, jnp.array(1.0)),
    (
        _midpoint_y_linear,
        ones_pytree,
        matrix,
    ),
    (
        _midpoint_f_linear,
        ones_pytree,
        matrix,
    ),
    (
        _midpoint_k_linear,
        ones_pytree,
        matrix,
    ),
    (_midpoint_y_nonlinear, ones_robertson, robertson_args),
    (
        _midpoint_f_nonlinear,
        ones_robertson,
        robertson_args,
    ),
    (
        _midpoint_k_nonlinear,
        ones_robertson,
        robertson_args,
    ),
)


# Not really useful enough to be part of the public API, but useful for testing against.
class PiggybackAdjoint(optx.AbstractAdjoint):
    def apply(self, primal_fn, rewrite_fn, inputs, tags):
        del rewrite_fn, tags
        while_loop = ft.partial(eqxi.while_loop, kind="lax")
        return primal_fn(inputs, while_loop)

tests/helpers.py
patrick-kidger/optimistix
import jax.numpy as jnp

import optimistix as optx


def test_fixed_point():
    def fn(y, _):
        return jnp.tanh(y + 1)

    solver = optx.FixedPointIteration(rtol=1e-6, atol=1e-6)
    solver = optx.BestSoFarFixedPoint(solver)
    sol = optx.fixed_point(fn, solver, jnp.array(0.0))
    assert jnp.allclose(sol.value, 0.96118069, rtol=1e-5, atol=1e-5)


def test_root_find():
    def fn(y, _):
        return y - jnp.tanh(y + 1)

    solver = optx.Bisection(rtol=1e-6, atol=1e-6)
    solver = optx.BestSoFarRootFinder(solver)
    options = dict(lower=-1, upper=1)
    sol = optx.root_find(fn, solver, jnp.array(0.0), options=options)
    assert jnp.allclose(sol.value, 0.96118069, rtol=1e-5, atol=1e-5)


def test_least_squares():
    def fn(y, _):
        return y - jnp.tanh(y + 1)

    solver = optx.LevenbergMarquardt(rtol=1e-6, atol=1e-6)
    solver = optx.BestSoFarLeastSquares(solver)
    sol = optx.least_squares(fn, solver, jnp.array(0.0))
    assert jnp.allclose(sol.value, 0.96118069, rtol=1e-5, atol=1e-5)


def test_minimise():
    def fn(y, _):
        return 0.5 * (y - jnp.tanh(y + 1)) ** 2

    solver = optx.BFGS(rtol=1e-6, atol=1e-6)
    solver = optx.BestSoFarMinimiser(solver)
    sol = optx.minimise(fn, solver, jnp.array(0.0))
    assert jnp.allclose(sol.value, 0.96118069, rtol=1e-5, atol=1e-5)

tests/test_best_so_far.py
patrick-kidger/optimistix
import random

import equinox as eqx
import jax.numpy as jnp
import jax.random as jr
import jax.tree_util as jtu
import pytest
from equinox.internal import ω

import optimistix as optx

from .helpers import (
    bisection_fn_init_options_args,
    finite_difference_jvp,
    fixed_point_fn_init_args,
    PiggybackAdjoint,
    tree_allclose,
    trivial,
)


atol = rtol = 1e-6
_fp_solvers = (optx.FixedPointIteration(rtol, atol),)
smoke_aux = (jnp.ones((2, 3)), {"smoke_aux": jnp.ones(2)})


@pytest.mark.parametrize("solver", _fp_solvers)
@pytest.mark.parametrize("_fn, init, args", fixed_point_fn_init_args)
def test_fixed_point(solver, _fn, init, args):
    atol = rtol = 1e-4
    has_aux = random.choice([True, False])
    if has_aux:
        fn = lambda x, args: (_fn(x, args), smoke_aux)
    else:
        fn = _fn
    optx_fp = optx.fixed_point(
        fn, solver, init, has_aux=has_aux, args=args, max_steps=10_000, throw=False
    ).value
    out = fn(optx_fp, args)
    if has_aux:
        f_val, _ = out
    else:
        f_val = out
    assert tree_allclose(optx_fp, f_val, atol=atol, rtol=rtol)


@pytest.mark.parametrize("solver", _fp_solvers)
@pytest.mark.parametrize("_fn, init, args", fixed_point_fn_init_args)
def test_fixed_point_jvp(getkey, solver, _fn, init, args):
    atol = rtol = 1e-3
    has_aux = random.choice([True, False])
    if has_aux:
        fn = lambda x, args: (_fn(x, args), smoke_aux)
    else:
        fn = _fn

    dynamic_args, static_args = eqx.partition(args, eqx.is_array)
    t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), init)
    t_dynamic_args = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), dynamic_args)

    def fixed_point(x, dynamic_args, *, adjoint):
        args = eqx.combine(dynamic_args, static_args)
        return optx.fixed_point(
            fn,
            solver,
            x,
            has_aux=has_aux,
            args=args,
            max_steps=10_000,
            adjoint=adjoint,
            throw=False,
        ).value

    otd = optx.ImplicitAdjoint()
    expected_out, t_expected_out = finite_difference_jvp(
        fixed_point,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=otd,
    )
    out, t_out = eqx.filter_jvp(
        fixed_point,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=otd,
    )
    dto = PiggybackAdjoint()
    expected_out2, t_expected_out2 = finite_difference_jvp(
        fixed_point,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=dto,
    )
    out2, t_out2 = eqx.filter_jvp(
        fixed_point,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=dto,
    )
    assert tree_allclose(expected_out2, expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(out, expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(out2, expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(t_expected_out2, t_expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(t_out, t_expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(t_out2, t_expected_out, atol=atol, rtol=rtol)


@pytest.mark.parametrize(
    "_fn, init, bisection_options, args", bisection_fn_init_options_args
)
def test_bisection(_fn, init, bisection_options, args):
    solver = optx.Bisection(rtol=1e-6, atol=1e-6)
    atol = rtol = 1e-4
    has_aux = random.choice([True, False])

    def root_find_problem(y, args):
        f_val = _fn(y, args)
        return (f_val**ω - y**ω).ω

    if _fn == trivial:
        bisection_problem = _fn
    else:
        bisection_problem = root_find_problem
    if has_aux:
        fn = lambda x, args: (bisection_problem(x, args), smoke_aux)
    else:
        fn = bisection_problem

    optx_fp = optx.root_find(
        fn,
        solver,
        init,
        has_aux=has_aux,
        args=args,
        options=bisection_options,
        max_steps=10_000,
        throw=False,
    ).value
    out = fn(optx_fp, args)
    if has_aux:
        f_val, _ = out
    else:
        f_val = out
    zeros = jtu.tree_map(jnp.zeros_like, f_val)
    assert tree_allclose(f_val, zeros, atol=atol, rtol=rtol)

tests/test_fixed_point.py
patrick-kidger/optimistix
import random

import equinox as eqx
import jax
import jax.numpy as jnp
import jax.random as jr
import jax.tree_util as jtu
import lineax as lx
import pytest

import optimistix as optx

from .helpers import (
    diagonal_quadratic_bowl,
    finite_difference_jvp,
    least_squares_fn_minima_init_args,
    least_squares_optimisers,
    rosenbrock,
    simple_nn,
    tree_allclose,
)


smoke_aux = (jnp.ones((2, 3)), {"smoke_aux": jnp.ones(2)})


@pytest.mark.parametrize("solver", least_squares_optimisers)
@pytest.mark.parametrize("_fn, minimum, init, args", least_squares_fn_minima_init_args)
def test_least_squares(solver, _fn, minimum, init, args):
    atol = rtol = 1e-4
    has_aux = random.choice([True, False])
    if has_aux:
        fn = lambda x, args: (_fn(x, args), smoke_aux)
    else:
        fn = _fn

    optx_argmin = optx.least_squares(
        fn, solver, init, has_aux=has_aux, args=args, max_steps=10_000, throw=False
    ).value
    out = fn(optx_argmin, args)
    if has_aux:
        residual, _ = out
    else:
        residual = out
    optx_min = jtu.tree_reduce(
        lambda x, y: x + y, jtu.tree_map(lambda x: jnp.sum(x**2), residual)
    )
    assert tree_allclose(optx_min, minimum, atol=atol, rtol=rtol)


@pytest.mark.parametrize("solver", least_squares_optimisers)
@pytest.mark.parametrize("_fn, minimum, init, args", least_squares_fn_minima_init_args)
def test_least_squares_jvp(getkey, solver, _fn, minimum, init, args):
    if _fn in (simple_nn, diagonal_quadratic_bowl):
        # These are ridiculously finickity to get references values for the derivatives
        return
    atol = rtol = 1e-2
    has_aux = random.choice([True, False])
    if has_aux:
        fn = lambda x, args: (_fn(x, args), smoke_aux)
    else:
        fn = _fn

    dynamic_args, static_args = eqx.partition(args, eqx.is_array)
    t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), init)
    t_dynamic_args = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), dynamic_args)

    def least_squares(x, dynamic_args, *, adjoint):
        args = eqx.combine(dynamic_args, static_args)
        return optx.least_squares(
            fn,
            solver,
            x,
            has_aux=has_aux,
            args=args,
            max_steps=10_000,
            adjoint=adjoint,
            throw=False,
        ).value

    if _fn is simple_nn:
        otd = optx.ImplicitAdjoint(linear_solver=lx.AutoLinearSolver(well_posed=False))
    else:
        otd = optx.ImplicitAdjoint()
    out, t_out = eqx.filter_jvp(
        least_squares,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=otd,
    )
    if _fn is rosenbrock:
        # Finite difference does a bad job on this one, but we can figure it out
        # analytically.
        assert isinstance(args, jax.Array)
        expected_out = jtu.tree_map(lambda x: jnp.full_like(x, args), init)
        t_expected_out = jtu.tree_map(lambda x: jnp.full_like(x, t_dynamic_args), init)
    else:
        expected_out, t_expected_out = finite_difference_jvp(
            least_squares,
            (init, dynamic_args),
            (t_init, t_dynamic_args),
            adjoint=otd,
        )
    # TODO(kidger): reinstate once we can do jvp-of-custom_vjp. Right now this errors
    #     because of the line searches used internally.
    #
    # dto = PiggybackAdjoint()
    # expected_out2, t_expected_out2 = finite_difference_jvp(
    #     least_squares, (init, dynamic_args), (t_init, t_dynamic_args), adjoint=dto,
    # )
    # out2, t_out2 = eqx.filter_jvp(
    #     least_squares, (init, dynamic_args), (t_init, t_dynamic_args), adjoint=dto,
    # )
    assert tree_allclose(out, expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(t_out, t_expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(expected_out2, expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(out2, expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(t_expected_out2, t_expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(t_out2, t_expected_out, atol=atol, rtol=rtol)

tests/test_least_squares.py
patrick-kidger/optimistix
import functools as ft
import random

import equinox as eqx
import jax.numpy as jnp
import jax.random as jr
import jax.tree_util as jtu
import optax
import pytest

import optimistix as optx

from .helpers import (
    beale,
    bowl,
    finite_difference_jvp,
    matyas,
    minimisation_fn_minima_init_args,
    minimisers,
    tree_allclose,
)


smoke_aux = (jnp.ones((2, 3)), {"smoke_aux": jnp.ones(2)})


@pytest.mark.parametrize("solver", minimisers)
@pytest.mark.parametrize("_fn, minimum, init, args", minimisation_fn_minima_init_args)
def test_minimise(solver, _fn, minimum, init, args):
    if isinstance(solver, optx.GradientDescent):
        max_steps = 100_000
    else:
        max_steps = 10_000
    atol = rtol = 1e-4
    has_aux = random.choice([True, False])
    if has_aux:
        fn = lambda x, args: (_fn(x, args), smoke_aux)
    else:
        fn = _fn
    optx_argmin = optx.minimise(
        fn, solver, init, has_aux=has_aux, args=args, max_steps=max_steps, throw=False
    ).value
    optx_min = _fn(optx_argmin, args)
    assert tree_allclose(optx_min, minimum, atol=atol, rtol=rtol)


@pytest.mark.parametrize("solver", minimisers)
@pytest.mark.parametrize("_fn, minimum, init, args", minimisation_fn_minima_init_args)
def test_minimise_jvp(getkey, solver, _fn, minimum, init, args):
    if isinstance(solver, (optx.GradientDescent, optx.NonlinearCG)):
        max_steps = 100_000
        atol = rtol = 1e-2
    else:
        max_steps = 10_000
        atol = rtol = 1e-3
    has_aux = random.choice([True, False])
    if has_aux:
        fn = lambda x, args: (_fn(x, args), smoke_aux)
    else:
        fn = _fn

    dynamic_args, static_args = eqx.partition(args, eqx.is_array)
    t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), init)
    t_dynamic_args = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), dynamic_args)

    def minimise(x, dynamic_args, *, adjoint):
        args = eqx.combine(dynamic_args, static_args)
        return optx.minimise(
            fn,
            solver,
            x,
            has_aux=has_aux,
            args=args,
            max_steps=max_steps,
            adjoint=adjoint,
            throw=False,
        ).value

    otd = optx.ImplicitAdjoint()
    out, t_out = eqx.filter_jit(ft.partial(eqx.filter_jvp, minimise))(
        (init, dynamic_args), (t_init, t_dynamic_args), adjoint=otd
    )
    if _fn is bowl:
        # Finite difference is very inaccurate on this problem.
        expected_out = t_expected_out = jtu.tree_map(jnp.zeros_like, init)
    elif _fn in (beale, matyas):
        if isinstance(solver, optx.NonlinearCG):
            eps = 1e-3
            atol = rtol = 1e-2  # finite difference does a really bad job on this one
        else:
            eps = 1e-4
        expected_out, t_expected_out = finite_difference_jvp(
            minimise,
            (init, dynamic_args),
            (t_init, t_dynamic_args),
            adjoint=otd,
            eps=eps,
        )
    else:
        expected_out, t_expected_out = finite_difference_jvp(
            minimise, (init, dynamic_args), (t_init, t_dynamic_args), adjoint=otd
        )
    # TODO(kidger): reinstate once we can do jvp-of-custom_vjp. Right now this errors
    #     because of the line searches used internally.
    #
    # dto = PiggybackAdjoint()
    # expected_out2, t_expected_out2 = finite_difference_jvp(
    #     minimise, (init, dynamic_args), (t_init, t_dynamic_args), adjoint=dto,
    # )
    # out2, t_out2 = eqx.filter_jvp(
    #     minimise, (init, dynamic_args), (t_init, t_dynamic_args), adjoint=dto,
    # )
    assert tree_allclose(out, expected_out, atol=atol, rtol=rtol)
    if not isinstance(solver, optx.NelderMead):
        # Nelder-Mead does such a bad job that the finite-difference gradients are
        # noticeably different.
        assert tree_allclose(t_out, t_expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(expected_out2, expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(out2, expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(t_expected_out2, t_expected_out, atol=atol, rtol=rtol)
    # assert tree_allclose(t_out2, t_expected_out, atol=atol, rtol=rtol)


@pytest.mark.parametrize(
    "method",
    [optx.polak_ribiere, optx.fletcher_reeves, optx.hestenes_stiefel, optx.dai_yuan],
)
def test_nonlinear_cg_methods(method):
    solver = optx.NonlinearCG(rtol=1e-10, atol=1e-10, method=method)

    def f(y, _):
        A = jnp.array([[2.0, -1.0], [-1.0, 3.0]])
        b = jnp.array([-100.0, 5.0])
        c = jnp.array(100.0)
        return jnp.einsum("ij,i,j", A, y, y) + jnp.dot(b, y) + c

    # Analytic minimum:
    # 0 = df/dyk
    #   = A_kj y_j + A_ik y_i + b_k
    #   = 2 A_kj y_j + b_k            (A is symmetric)
    # => y = -0.5 A^{-1} b
    #      = [[-0.3, 0.1], [0.1, 0.2]] [-100, 5]
    #      = [29.5, 9]
    y0 = jnp.array([2.0, 3.0])
    sol = optx.minimise(f, solver, y0, max_steps=500)
    assert tree_allclose(sol.value, jnp.array([29.5, 9.0]), rtol=1e-5, atol=1e-5)


def test_optax_recompilation():
    optim1 = optax.chain(
        optax.adam(jnp.array(1e-3)),
        optax.scale_by_schedule(optax.piecewise_constant_schedule(1, {200: 0.1})),
    )
    solver1 = optx.OptaxMinimiser(optim1, rtol=1e-3, atol=1e-3)

    optim2 = optax.chain(
        optax.adam(jnp.array(1e-2)),
        optax.scale_by_schedule(optax.piecewise_constant_schedule(1, {200: 0.1})),
    )
    solver2 = optx.OptaxMinimiser(optim2, rtol=1e-3, atol=1e-3)

    num_called = 0

    def f(x, _):
        nonlocal num_called
        num_called += 1
        return x**2

    optx.minimise(f, solver1, 1.0)
    num_called_so_far = num_called
    optx.minimise(f, solver2, 1.0)
    assert num_called_so_far == num_called

tests/test_minimise.py
patrick-kidger/optimistix
# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import jax
import jax.flatten_util as jfu
import jax.numpy as jnp

import optimistix as optx
import optimistix._misc as optx_misc

from .helpers import tree_allclose


def test_inexact_asarray_no_copy():
    x = jnp.array([1.0])
    assert optx_misc.inexact_asarray(x) is x
    y = jnp.array([1.0, 2.0])
    assert jax.vmap(optx_misc.inexact_asarray)(y) is y


# See JAX issue #15676
def test_inexact_asarray_jvp():
    p, t = jax.jvp(optx_misc.inexact_asarray, (1.0,), (2.0,))
    assert type(p) is not float
    assert type(t) is not float


def test_norm():
    def _square(x):
        return x * jnp.conj(x)

    def _two_norm(x):
        return jnp.sqrt(jnp.sum(_square(jfu.ravel_pytree(x)[0]))).real

    def _rms_norm(x):
        return jnp.sqrt(jnp.mean(_square(jfu.ravel_pytree(x)[0]))).real

    def _max_norm(x):
        return jnp.max(jnp.abs(jfu.ravel_pytree(x)[0]))

    x = [jnp.array(1.0), jnp.arange(4.0).reshape(2, 2)]
    tx = [jnp.array(0.5), jnp.arange(1.0, 5.0).reshape(2, 2)]

    assert jnp.allclose(optx.two_norm(x), _two_norm(x))
    assert jnp.allclose(optx.rms_norm(x), _rms_norm(x))
    assert jnp.allclose(optx.max_norm(x), _max_norm(x))

    two = jax.jvp(optx.two_norm, (x,), (tx,))
    true_two = jax.jvp(_two_norm, (x,), (tx,))
    rms = jax.jvp(optx.rms_norm, (x,), (tx,))
    true_rms = jax.jvp(_rms_norm, (x,), (tx,))
    max = jax.jvp(optx.max_norm, (x,), (tx,))
    true_max = jax.jvp(_max_norm, (x,), (tx,))
    assert tree_allclose(two, true_two)
    assert tree_allclose(rms, true_rms)
    assert tree_allclose(max, true_max)

    zero = [jnp.array(0.0), jnp.zeros((2, 2))]
    two0 = jax.jvp(optx.two_norm, (zero,), (tx,))
    rms0 = jax.jvp(optx.rms_norm, (zero,), (tx,))
    max0 = jax.jvp(optx.max_norm, (zero,), (tx,))
    true0 = (jnp.array(0.0), jnp.array(0.0))
    assert tree_allclose(two0, true0)
    assert tree_allclose(rms0, true0)
    assert tree_allclose(max0[0], true0[0])  # tangent rightly has nonzero value

    x = jnp.array([3 + 1.2j, -0.5 + 4.9j])
    tx = jnp.array([2 - 0.3j, -0.7j])
    two = jax.jvp(optx.two_norm, (x,), (tx,))
    true_two = jax.jvp(_two_norm, (x,), (tx,))
    rms = jax.jvp(optx.rms_norm, (x,), (tx,))
    true_rms = jax.jvp(_rms_norm, (x,), (tx,))
    max = jax.jvp(optx.max_norm, (x,), (tx,))
    true_max = jax.jvp(_max_norm, (x,), (tx,))
    assert two[0].imag == 0
    assert tree_allclose(two, true_two)
    assert rms[0].imag == 0
    assert tree_allclose(rms, true_rms)
    assert max[0].imag == 0
    assert tree_allclose(max, true_max)

tests/test_misc.py
patrick-kidger/optimistix
import random

import equinox as eqx
import jax.numpy as jnp
import jax.random as jr
import jax.tree_util as jtu
import pytest
from equinox.internal import ω

import optimistix as optx

from .helpers import (
    finite_difference_jvp,
    fixed_point_fn_init_args,
    PiggybackAdjoint,
    tree_allclose,
)


atol = rtol = 1e-6
_root_finders = (
    optx.Newton(rtol, atol),
    optx.Chord(rtol, atol),
)
smoke_aux = (jnp.ones((2, 3)), {"smoke_aux": jnp.ones(2)})


@pytest.mark.parametrize("solver", _root_finders)
@pytest.mark.parametrize("_fn, init, args", fixed_point_fn_init_args)
def test_root_find(solver, _fn, init, args):
    atol = rtol = 1e-5
    has_aux = random.choice([True, False])

    def root_find_problem(y, args):
        f_val = _fn(y, args)
        return (f_val**ω - y**ω).ω

    if has_aux:
        fn = lambda x, args: (root_find_problem(x, args), smoke_aux)
    else:
        fn = root_find_problem
    optx_root = optx.root_find(
        fn, solver, init, has_aux=has_aux, args=args, max_steps=10_000, throw=False
    ).value
    out = fn(optx_root, args)
    if has_aux:
        fn_val, _ = out
    else:
        fn_val = out
    zeros = jtu.tree_map(jnp.zeros_like, fn_val)
    assert tree_allclose(fn_val, zeros, atol=atol, rtol=rtol)


@pytest.mark.parametrize("solver", _root_finders)
@pytest.mark.parametrize("_fn, init, args", fixed_point_fn_init_args)
def test_root_find_jvp(getkey, solver, _fn, init, args):
    atol = rtol = 1e-3
    has_aux = random.choice([True, False])

    def root_find_problem(y, args):
        f_val = _fn(y, args)
        return (f_val**ω - y**ω).ω

    if has_aux:
        fn = lambda x, args: (root_find_problem(x, args), smoke_aux)
    else:
        fn = root_find_problem
    dynamic_args, static_args = eqx.partition(args, eqx.is_array)
    t_init = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), init)
    t_dynamic_args = jtu.tree_map(lambda x: jr.normal(getkey(), x.shape), dynamic_args)

    def root_find(x, dynamic_args, *, adjoint):
        args = eqx.combine(dynamic_args, static_args)
        return optx.root_find(
            fn,
            solver,
            x,
            has_aux=has_aux,
            args=args,
            max_steps=10_000,
            adjoint=adjoint,
            throw=False,
        ).value

    otd = optx.ImplicitAdjoint()
    expected_out, t_expected_out = finite_difference_jvp(
        root_find,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=otd,
    )
    out, t_out = eqx.filter_jvp(
        root_find,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=otd,
    )
    dto = PiggybackAdjoint()
    expected_out2, t_expected_out2 = finite_difference_jvp(
        root_find,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=dto,
    )
    out2, t_out2 = eqx.filter_jvp(
        root_find,
        (init, dynamic_args),
        (t_init, t_dynamic_args),
        adjoint=dto,
    )
    assert tree_allclose(expected_out2, expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(out, expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(out2, expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(t_expected_out2, t_expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(t_out, t_expected_out, atol=atol, rtol=rtol)
    assert tree_allclose(t_out2, t_expected_out, atol=atol, rtol=rtol)


def test_bisection_flip():
    for fn in (lambda x, _: x, lambda x, _: -x):
        options = dict(lower=-1, upper=2)
        sol = optx.root_find(
            fn, optx.Bisection(rtol=1e-4, atol=1e-4), 1.0, options=options
        )
        assert jnp.allclose(0, sol.value, atol=1e-3)


@pytest.mark.parametrize(
    "solver", [optx.Newton(rtol=1e-5, atol=1e-5), optx.Chord(rtol=1e-5, atol=1e-5)]
)
def test_newton_bounded(solver):
    y0 = (jnp.array(0.0), jnp.arange(4.0).reshape(2, 2))

    def f(y, _):
        ya, ((yb, yc), (yd, ye)) = y
        foo = ya + yb + yc + yd + ye
        bar = ya - yb + yc - yd + ye
        return foo * bar, ya, yc**2, jnp.tanh(yd), ye + 1

    # Find roots:
    # Clearly ya=0 yc=0 yd=0, ye=-1.
    # Then foo*bar=0 => foo=0 or bar=0.
    # foo=0 => yb=1
    # bar=0 => yb=-1

    if isinstance(solver, optx.Newton):
        tol = 1e-4
    else:
        tol = 1e-2

    lower_bound = (-jnp.inf, jnp.array([[0, -jnp.inf], [-jnp.inf, -jnp.inf]]))
    true_lower_root = (jnp.array(0.0), jnp.array([[1.0, 0.0], [0.0, -1.0]]))
    y0 = (jnp.array(0.1), jnp.array([[1.1, 0.1], [0.1, -1.1]]))
    lower_root = optx.root_find(f, solver, y0, options=dict(lower=lower_bound)).value
    assert tree_allclose(lower_root, true_lower_root, rtol=tol, atol=tol)

    upper_bound = (jnp.inf, jnp.array([[0, jnp.inf], [jnp.inf, jnp.inf]]))
    true_upper_root = (jnp.array(0.0), jnp.array([[-1.0, 0.0], [0.0, -1.0]]))
    y0 = (jnp.array(0.1), jnp.array([[-1.1, 0.1], [0.1, -1.1]]))
    upper_root = optx.root_find(f, solver, y0, options=dict(upper=upper_bound)).value
    assert tree_allclose(upper_root, true_upper_root, rtol=tol, atol=tol)


def test_root_via_min():
    def f(y, _):
        ya, (yb, yc) = y
        return jnp.tanh(ya + 0.1), jnp.tanh(yb - 0.5), jnp.tanh(yc * 2)

    y0 = jnp.array(0.5), jnp.array([-0.3, 0.7])
    sol = optx.root_find(f, optx.BFGS(rtol=1e-8, atol=1e-8), y0)
    assert tree_allclose(sol.value, (jnp.array(-0.1), jnp.array([0.5, 0.0])))


def test_bad_root_via_min():
    def f(y, _):
        ya, (yb, yc) = y
        return 1.0, jnp.tanh(yb - 0.5), jnp.tanh(yc * 2)

    y0 = jnp.array(0.5), jnp.array([-0.3, 0.7])
    sol = optx.root_find(f, optx.BFGS(rtol=1e-8, atol=1e-8), y0, throw=False)
    assert sol.result == optx.RESULTS.nonlinear_max_steps_reached

tests/test_root_find.py
patrick-kidger/optimistix
import jax

import optimistix as optx


def test_minimise():
    @jax.grad
    def f(offset):
        def fn(x, _):
            return x**2 + offset

        solver = optx.GradientDescent(learning_rate=0.1, rtol=0.1, atol=0.1)
        return optx.minimise(fn, solver, 0.0).value

    f(0.0)


def test_least_squares():
    @jax.grad
    def f(offset):
        def fn(x, _):
            return x + offset

        solver = optx.Dogleg(rtol=0.1, atol=0.1)
        return optx.least_squares(fn, solver, 0.0).value

    f(0.0)


def test_root_find():
    @jax.grad
    def f(offset):
        def fn(x, _):
            return x - offset

        solver = optx.Newton(rtol=0.1, atol=0.1)
        return optx.root_find(fn, solver, 0.0).value

    f(0.0)


def test_fixed_point():
    @jax.grad
    def f(offset):
        def fn(x, _):
            return offset

        solver = optx.FixedPointIteration(rtol=0.1, atol=0.1)
        return optx.fixed_point(fn, solver, 0.0).value

    f(0.0)

tests/test_solve.py
patrick-kidger/optimistix