Optimistix is a JAX library for nonlinear solvers, including root finding, minimisation, fixed points, and least squares. It features highly modular optimisers, interoperable solvers (e.g., converting root find problems to least squares), PyTree-based state management, fast compilation and runtimes, and deep integration with the JAX ecosystem for features like autodiff, autoparallelism, and GPU/TPU support. As of version 0.1.0, it requires Python 3.11+ and is under active, rapid development with frequent updates.
pip install optimistixVerified import paths — ran on the pinned version, not inferred.
This quickstart demonstrates finding a fixed point for an implicit Euler step of an ODE. It uses `optimistix.fixed_point` with a `Newton` solver. The `fn` defines the function for which the fixed point is sought, taking `y` and `args` and returning the next `y` value. The solution object `sol` contains the `value` of the fixed point and a `result` code indicating success or failure.
Update `verbose` usage: `verbose=True` to print everything, `verbose=False` (default) to print nothing, or provide a custom `callable` for fine-grained control, instead of `verbose=frozenset({'loss'})`.To handle errors programmatically without raising an exception, pass `throw=False` to the top-level solve function (e.g., `optx.fixed_point(..., throw=False)`). The `sol.result` attribute can then be inspected for success/failure codes (0 for success).
If the solution is suboptimal, consider improving the initial guess (`y0`), trying different solvers (consult 'How to choose a solver' in the docs), or reformulating the problem (e.g., fitting parts of a time series incrementally).
Use `optimistix.compat.minimize` as a drop-in replacement or migrate to native Optimistix APIs like `optx.minimise` for more control and JAX ecosystem compatibility.
If the absolute best value across all steps is critical, users may need to implement custom logic to store and compare intermediate values during a stepped solve, rather than relying solely on the final `sol.value`.
Check inputs to the problem for `NaN` or `inf` values. If solving a linear least-squares problem, pass `solver=AutoLinearSolver(well_posed=False)`. If the problem is inherently ill-conditioned, consider a more robust solver or re-parametrisation. Placing `jax.debug.print` or `jax.debug.breakpoint` can help diagnose the issue.
Increase the `max_steps` argument in the solve function (e.g., `optx.fixed_point(..., max_steps=N)`). Verify that the problem actually has a solution. Loosen `rtol` (relative tolerance) or `atol` (absolute tolerance) if appropriate for the application.
This often points to a poorly scaled problem, a bad initial guess (`y0`), or an unsuitable solver. Try different initial guesses, consider scaling your problem variables, or switch to a more robust solver for 'messier' problems (e.g., `OptaxMinimiser` for minimisation, `LevenbergMarquardt` or `Dogleg` for root-finding/least-squares).
Verify the mathematical properties of the function being solved. If you expect a root or fixed point but the solver fails, it may be converging to a local minimum of the squared residual instead of zero. For problems where a root is not guaranteed, consider using a minimisation algorithm on the squared residual `f(y)^2` instead of a root finder directly.