NonlinearSolve.jl: High-Performance and Robust Solvers for Systems of Nonlinear Equations in Julia

Table of Links

Abstract and 1. Introduction

2. Mathematical Description and 2.1. Numerical Algorithms for Nonlinear Equations

2.2. Globalization Strategies

2.3. Sensitivity Analysis

2.4. Matrix Coloring & Sparse Automatic Differentiation

3. Special Capabilities

3.1. Composable Building Blocks

3.2. Smart PolyAlgortihm Defaults

3.3. Non-Allocating Static Algorithms inside GPU Kernels

3.4. Automatic Sparsity Exploitation

3.5. Generalized Jacobian-Free Nonlinear Solvers using Krylov Methods

4. Results and 4.1. Robustness on 23 Test Problems

4.2. Initializing the Doyle-Fuller-Newman (DFN) Battery Model

4.3. Large Ill-Conditioned Nonlinear Brusselator System

5. Conclusion and References

Abstract. Efficiently solving nonlinear equations underpins numerous scientific and engineering disciplines, yet scaling these solutions for complex system models remains a challenge. This paper presents NonlinearSolve.jl – a suite of high-performance open-source nonlinear equation solvers implemented natively in the Julia programming language. NonlinearSolve.jl distinguishes itself by offering a unified API that accommodates a diverse range of solver specifications alongside features such as automatic algorithm selection based on runtime analysis, support for GPU-accelerated computation through static array kernels, and the utilization of sparse automatic differentiation and Jacobian-free Krylov methods for large-scale problem-solving. Through rigorous comparison with established tools such as Sundials and MINPACK, NonlinearSolve.jl demonstrates unparalleled robustness and efficiency, achieving significant advancements in solving benchmark problems and challenging real-world applications. The capabilities of NonlinearSolve.jl unlock new potentials in modeling and simulation across various domains, making it a valuable addition to the computational toolkit of researchers and practitioners alike.

1. Introduction.

Nonlinear equations are ubiquitous across the sciences, governing the behavior of complex systems in domains ranging from physics to biology to machine learning. Solving steady state (or fixed-point) equations, implicit Ordinary Differential Equation (ODE) [1], semi-explicit Differential Algebraic Equations (DAEs), boundary value differential equations [2, 3, 4, 5], and memory-efficient deep learning [6, 7, 8, 9] all boil down to solving nonlinear equations. Despite their prevalence in real-world problems, finding solutions to nonlinear equations at scale remains a major challenge.

This paper presents NonlinearSolve.jl - an open-source, high-performance nonlinear equation-solving framework implemented in the Julia Programming Language [10]. NonlinearSolve.jl provides a flexible, easy-to-use interface for specifying and solving general nonlinear equations. Under the hood, it employs various state-of-the-art techniques such as automatic sparsity exploitation, fast automatic differentiation, and jacobian-free methods to solve systems at scale reliably. NonlinearSolve.jl is part of the Scientific Machine Learning (SciML) package ecosystem of Julia[1] and is, in turn, extended by various packages to solve domain-specific problems.

The key capabilities and contributions of NonlinearSolve.jl include:

• Unified API for Rapid Experimentation with Nonlinear Solver Options: Users can seamlessly switch between solver algorithms (including the same algorithm from different software packages), line search methods, trust region schemes, automatic-differentiation backends, linear solvers (including Krylov methods), and sparsity detection algorithms.

• Smart Polyalgorithmic Solver Selection for Robustness: Automated runtime checks select appropriate internal solver methods and parameters, balancing speed and reliability.

• High-performance Algorithms for Small Problems: Specialized non-allocating kernels solve small systems extremely efficiently – outperforming existing software.

• Automatic Sparsity Exploitation for Large Systems: Approximation techniques automatically detect sparsity patterns in system Jacobians. Colored sparse matrix computations coupled with sparse linear solvers accelerate both derivative calculations and overall solve times.

• Jacobian-Free Krylov Methods: Matrix-free linear solvers avoid explicit Jacobian materialization, reducing memory overhead and enabling the solution of otherwise computationally infeasible systems.

• Composable Blocks for Building Solvers: All solvers in NonlinearSolve.jl are built using a composition of fundamental building blocks – Descent Algorithm, Linear Solver, Jacobian Computation Algorithm, and Globalization Strategy.

• Symbolic Tearing of Nonlinear Systems: ModelingToolkit.jl [11] automatically performs tearing of nonlinear systems for NonlinearSolve.jl and reduces the computational complexity of solving those systems.

We demonstrate the capabilities of NonlinearSolve.jl with several numerical experiments. The software reliably solves 23 nonlinear test problems [Subsection 4.1], outperforming the open-source Sundials toolkit and reference MINPACK implementation. Furthermore, we show NonlinearSolve.jl’s applicability to large real-world problems by benchmarking it on applications like initializing a DAE battery model [Subsection 4.2] and steady-state Partial Differential Equation (PDE) solving [Subsection 4.3]. NonlinearSolve.jl enables fast yet robust nonlinear equation solving on problems where existing tools break down. Its efficiency, flexibility, and seamless GPU support unlock new modeling and simulation capabilities across application domains.

1.1. Comparison to Existing Software. Most existing nonlinear solvers exhibit flexibility, performance, and scalability limitations when applied to real-world nonlinear problems. For example, Sundials KINSOL [12, 13] is specialized for implicit ODE solvers, and tricks like Jacobian reuse often fail for simple test problems [Subsection 4.1]. MINPACK [14, 15] forms the base for popular packages like SciPy [16] and MATLAB fsolve [17]; however, this lacks the support of tools like Automatic Differentiation (AD), sparsity detection, and Krylov methods, is often more than 100 times slower than our solvers [Figure 10] and fails for numerically challenging problems [Subsection 4.2]. PETSc SNES [18] is a specialized set of solvers for large systems in a distributed setup; however, they lack the support for sparse AD and instead can compute sparse Jacobians exclusively via colored finite differencing. Optimistix [19] solves systems of nonlinear equations for machine learning-focused applications and, as such, lacks support for any form of sparsity handling that is uncommon in their target domain. We draw much inspiration from pre-existing Julia Packages like NLsolve.jl and NLSolvers.jl [20]; however, they had similar problems of scalability and lacked specialized linear solvers as evident from their slower relative performance [Section 4]. In contrast to existing software, NonlinearSolve.jl is designed to be robust, performant, and scalable in general while having controls to be customized to specific applications

1.2. Domain Specific Extensions of NonlinearSolve.jl. In addition to the core framework for solving nonlinear systems, NonlinearSolve.jl is extended by several packages to solve domain-specific problems. These extensions include:

1. SteadyStateDiffEq.jl: Enables computing the steady state of dynamical systems. Catalyst.jl’s[2] tutorials[3] demonstrate how this is used to determine the steady states of a dimerization reaction network.

   
   

2. DeepEquilibriumNetworks.jl: Utilizes the sensitivity analysis capabilities of NonlinearSolve.jl to train deep equilibrium networks [Subsection 2.3].

   
   

3. DiffEqCallbacks.jl: Provides ManifoldCallback for solving semi-explicit ODEs on arbitrary manifolds. Our tutorials[4] demonstrate how NonlinearSolve.jl can be used to conserve energy and angular momentum in the Kepler Problem [22].

   
   

4. OrdinaryDiffEq.jl: Uses NonlinearSolve.jl to initialize DAE solves. In Subsection 4.2, we solve the DAE initialization problem for a battery model using NonlinearSolve.jl.

[1] https://sciml.ai/

[2] Catalyst.jl [21] is a package for modeling chemical reaction networks.

[3] https://docs.sciml.ai/Catalyst/dev/catalyst_applications/nonlinear_solve/#nonlinear_solve

[4] https://docs.sciml.ai/DiffEqDocs/stable/examples/kepler_problem/#The-Kepler-Proble