DOCSEARCHbyAlgolia
Submit Docs

optimistix

repository·main·Indexed 20 days ago

https://github.com/patrick-kidger/optimistix

A JAX and Equinox-based library for nonlinear solvers, including root finding, minimisation, fixed points, and least squares. It provides a modular ecosystem with specialized solvers like Newton, BFGS, and Levenberg-Marquardt, and supports autodifferentiation of nonlinear solves via the implicit function theorem using optimistix.ImplicitAdjoint. The library is designed for extensibility through abstract base classes for minimisers, root finders, and descent directions.

Tokens
18.8K
Snippets
42
Records
95
Agent score
69%

What's inside optimistix

  1. Minimisation vs Least-Squares in Optimistix

    main

    Optimistix handles both standard minimisation and least-squares problems using the same underlying abstractions.

    Minimisation

    Target: Minimize $f: \mathbb{R}^n \to \mathbb{R}$. Evaluated quantities typically include:

    • Function values $f(y)$
    • Gradients $\nabla f(y)$
    • Hessian approximations $\nabla^2 f(y)$

    Least-Squares

    Target: Minimize $0.5 \sum_i r(y)_i^2$, where $r: \mathbb{R}^n \to \mathbb{R}^m$ is a residual function.

    • Function evaluations: $f(y)$ is computed directly from residuals.
    • Gradients: $\nabla f(y) = r(y) \nabla r(y)$, computed efficiently via vector-Jacobian products.
    • Hessians: Approximated via the Gauss–Newton approximation: $\nabla^2 f(y) \approx (\nabla r(y))^T (\nabla r(y))$.

  2. Define custom searches, descent directions, and adjoints

    main

    Optimistix allows for deep customization of the optimization process through several specialized ABCs:

    • Searches, Trust Regions, and Learning Rates: Subclass optimistix.AbstractSearch to define custom search logic.
    • Descent Directions: Subclass optimistix.AbstractDescent to define custom directions such as steepest descent, Newton steps, or Levenberg--Marquardt damped steps.
    • Adjoints: Subclass optimistix.AbstractAdjoint to implement custom autodifferentiation strategies.

  3. Use FunctionInfo to inspect function properties

    main
    The optimistix.FunctionInfo class and its specialized subclasses provide metadata and evaluation capabilities for functions used within Optimistix. These classes allow you to inspect properties like as_min (indicating if a function is treated as a minimization objective) and provide structured ways to evaluate functions, their gradients, Hessians, and residuals.

  4. Use trust region methods with ClassicalTrustRegion or LinearTrustRegion

    main

    Optimistix provides trust region methods to constrain the step size within a region where a local model (usually quadratic) is considered reliable.

    • optimistix.ClassicalTrustRegion: Implements a standard trust region approach.
    • optimistix.LinearTrustRegion: Implements a trust region approach based on a linear model.

  5. Understanding accepted and rejected steps

    main

    In Optimistix, the optimization loop is flattened, meaning there is no distinction between the steps of the overall optimizer and the steps within a sub-routine like a line search. This allows for performance optimizations like reusing expensive computations (e.g., solving a linear system for a Newton descent) across multiple search steps.

    • Accepted steps: Steps where the descent is re-evaluated. For example, in a backtracking search, a step is accepted if it satisfies the Armijo condition.
    • Rejected steps: Steps where the descent direction is kept fixed while the search evaluates different scalar values (e.g., different step lengths in a line search).

    Searches decide which steps are accepted or rejected.

  7. Compare Optimistix with Optax, JAXopt, and jax.scipy.optimize.minimize

    main

    Optimistix vs Optax

    Optax is dedicated to first-order gradient-based minimization (useful for neural networks). Optimistix has a broader scope, including root-finding and second-order methods like optimistix.LevenbergMarquardt or optimistix.Newton. For first-order gradient methods, you should use optimistix.OptaxMinimiser to pair Optimistix with Optax.

    Optimistix vs JAXopt

    Use Optimistix if you need:

    • Faster compilation and execution.
    • Solvers not in JAXopt (e.g., optimistix.Newton for root-finding).
    • Better integration with the Equinox ecosystem.
    • Flexibility to mix-and-match different optimizers.

    Use JAXopt if you need:

    • To optimize functions that cannot be jax.jit'd.
    • Constrained optimization (beyond hypercube constraints) or quadratic programming.

    Optimistix vs jax.scipy.optimize.minimize

    jax.scipy.optimize.minimize is being removed from JAX. Optimistix provides optimistix.compat.minimize as a drop-in replacement.

  8. Use AbstractSearch for custom search strategies

    main
    The optimistix.AbstractSearch class serves as the base interface for implementing custom search strategies in Optimistix. If you need to define a custom way to determine step sizes or update parameters during an optimization process, you should subclass AbstractSearch and implement its required interface.

  10. Use fixed-point solvers in Optimistix

    main

    The optimistix.fixed_point module provides tools to find a fixed point x such that f(x) = x.

    Optimistix is highly flexible regarding the solver argument. While it provides dedicated fixed-point solvers, you can also pass in:

    • Any Root Finder: Optimistix will automatically rewrite the problem to find the root of f(x) - x = 0.
    • Any Least Squares Solver: Optimistix will rewrite the problem to minimize sum((f(x) - x)^2).
    • Any Minimiser: Optimistix will rewrite the problem to minimize the same squared difference objective.

  11. How searches and descents work together

    main

    Optimistix uses a generalized approach to optimization by decoupling the determination of step size from the determination of the update direction. This is achieved through two primary abstractions: Searches and Descents.

    Searches

    Searches consume information about the function (values, gradients, Hessians, etc.) and produce a single scalar value. This scalar represents:

    • The distance along a line search.
    • The radius of a trust region.
    • The value of a learning rate.

    Descents

    Descents consume the function information along with the scalar value produced by a search to compute the actual update to apply to the parameters. Examples include:

    • -scalar * gradient (Steepest Descent)
    • -scalar * (Hessian⁻¹ gradient) (Newton/Gauss-Newton)
    • The distance along a dogleg path (Dogleg)
    • (Jacobianᵀ Jacobian + scalar * I)⁻¹ gradient (Levenberg–Marquardt)

    Combining Abstractions

    By mixing and matching these, you can construct various algorithms:

    • Gradient Descent: Fixed learning rate + Steepest descent.
    • Levenberg–Marquardt: Trust region algorithm + Damped Newton descent.
    • BFGS: Backtracking Armijo line search + Newton descent.

  12. Use the Solution abstraction

    main
    In optimistix, a Solution object represents the result of an optimization or integration process. It encapsulates the final state and the metadata associated with the solver's execution. While the provided documentation snippet is a high-level representation, users typically interact with Solution objects returned by solver functions to access the computed results.