|
| 1 | +# Solving Large Ill-Conditioned Nonlinear Systems with NonlinearSolve.jl |
| 2 | + |
| 3 | +This tutorial is for getting into the extra features of using NonlinearSolve.jl. Solving ill-conditioned nonlinear systems requires specializing the linear solver on properties of the Jacobian in order to cut down on the ``\mathcal{O}(n^3)`` linear solve and the ``\mathcal{O}(n^2)`` back-solves. This tutorial is designed to explain the advanced usage of NonlinearSolve.jl by solving the steady state stiff Brusselator partial differential equation (BRUSS) using NonlinearSolve.jl. |
| 4 | + |
| 5 | +## Definition of the Brusselator Equation |
| 6 | + |
| 7 | +!!! note |
| 8 | + |
| 9 | + Feel free to skip this section: it simply defines the example problem. |
| 10 | + |
| 11 | +The Brusselator PDE is defined as follows: |
| 12 | + |
| 13 | +```math |
| 14 | +\begin{align} |
| 15 | +0 &= 1 + u^2v - 4.4u + \alpha(\frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}) + f(x, y, t)\\ |
| 16 | +0 &= 3.4u - u^2v + \alpha(\frac{\partial^2 v}{\partial x^2} + \frac{\partial^2 v}{\partial y^2}) |
| 17 | +\end{align} |
| 18 | +``` |
| 19 | + |
| 20 | +where |
| 21 | + |
| 22 | +```math |
| 23 | +f(x, y, t) = \begin{cases} |
| 24 | +5 & \quad \text{if } (x-0.3)^2+(y-0.6)^2 ≤ 0.1^2 \text{ and } t ≥ 1.1 \\ |
| 25 | +0 & \quad \text{else} |
| 26 | +\end{cases} |
| 27 | +``` |
| 28 | + |
| 29 | +and the initial conditions are |
| 30 | + |
| 31 | +```math |
| 32 | +\begin{align} |
| 33 | +u(x, y, 0) &= 22\cdot (y(1-y))^{3/2} \\ |
| 34 | +v(x, y, 0) &= 27\cdot (x(1-x))^{3/2} |
| 35 | +\end{align} |
| 36 | +``` |
| 37 | + |
| 38 | +with the periodic boundary condition |
| 39 | + |
| 40 | +```math |
| 41 | +\begin{align} |
| 42 | +u(x+1,y,t) &= u(x,y,t) \\ |
| 43 | +u(x,y+1,t) &= u(x,y,t) |
| 44 | +\end{align} |
| 45 | +``` |
| 46 | + |
| 47 | +To solve this PDE, we will discretize it into a system of ODEs with the finite |
| 48 | +difference method. We discretize `u` and `v` into arrays of the values at each |
| 49 | +time point: `u[i,j] = u(i*dx,j*dy)` for some choice of `dx`/`dy`, and same for |
| 50 | +`v`. Then our ODE is defined with `U[i,j,k] = [u v]`. The second derivative |
| 51 | +operator, the Laplacian, discretizes to become a tridiagonal matrix with |
| 52 | +`[1 -2 1]` and a `1` in the top right and bottom left corners. The nonlinear functions |
| 53 | +are then applied at each point in space (they are broadcast). Use `dx=dy=1/32`. |
| 54 | + |
| 55 | +The resulting `NonlinearProblem` definition is: |
| 56 | + |
| 57 | +```@example ill_conditioned_nlprob |
| 58 | +using NonlinearSolve, LinearAlgebra, SparseArrays, LinearSolve |
| 59 | +
|
| 60 | +const N = 32 |
| 61 | +const xyd_brusselator = range(0, stop = 1, length = N) |
| 62 | +brusselator_f(x, y) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * 5.0 |
| 63 | +limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a |
| 64 | +function brusselator_2d_loop(du, u, p) |
| 65 | + A, B, alpha, dx = p |
| 66 | + alpha = alpha / dx^2 |
| 67 | + @inbounds for I in CartesianIndices((N, N)) |
| 68 | + i, j = Tuple(I) |
| 69 | + x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]] |
| 70 | + ip1, im1, jp1, jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N), |
| 71 | + limit(j - 1, N) |
| 72 | + du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] - |
| 73 | + 4u[i, j, 1]) + |
| 74 | + B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] + |
| 75 | + brusselator_f(x, y) |
| 76 | + du[i, j, 2] = alpha * (u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] - |
| 77 | + 4u[i, j, 2]) + |
| 78 | + A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2] |
| 79 | + end |
| 80 | +end |
| 81 | +p = (3.4, 1.0, 10.0, step(xyd_brusselator)) |
| 82 | +
|
| 83 | +function init_brusselator_2d(xyd) |
| 84 | + N = length(xyd) |
| 85 | + u = zeros(N, N, 2) |
| 86 | + for I in CartesianIndices((N, N)) |
| 87 | + x = xyd[I[1]] |
| 88 | + y = xyd[I[2]] |
| 89 | + u[I, 1] = 22 * (y * (1 - y))^(3 / 2) |
| 90 | + u[I, 2] = 27 * (x * (1 - x))^(3 / 2) |
| 91 | + end |
| 92 | + u |
| 93 | +end |
| 94 | +u0 = init_brusselator_2d(xyd_brusselator) |
| 95 | +prob_brusselator_2d = NonlinearProblem(brusselator_2d_loop, u0, p) |
| 96 | +``` |
| 97 | + |
| 98 | +## Choosing Jacobian Types |
| 99 | + |
| 100 | +When we are solving this nonlinear problem, the Jacobian must be built at many |
| 101 | +iterations, and this can be one of the most |
| 102 | +expensive steps. There are two pieces that must be optimized in order to reach |
| 103 | +maximal efficiency when solving stiff equations: the sparsity pattern and the |
| 104 | +construction of the Jacobian. The construction is filling the matrix |
| 105 | +`J` with values, while the sparsity pattern is what `J` to use. |
| 106 | + |
| 107 | +The sparsity pattern is given by a prototype matrix, the `jac_prototype`, which |
| 108 | +will be copied to be used as `J`. The default is for `J` to be a `Matrix`, |
| 109 | +i.e. a dense matrix. However, if you know the sparsity of your problem, then |
| 110 | +you can pass a different matrix type. For example, a `SparseMatrixCSC` will |
| 111 | +give a sparse matrix. Other sparse matrix types include: |
| 112 | + |
| 113 | + - Bidiagonal |
| 114 | + - Tridiagonal |
| 115 | + - SymTridiagonal |
| 116 | + - BandedMatrix ([BandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BandedMatrices.jl)) |
| 117 | + - BlockBandedMatrix ([BlockBandedMatrices.jl](https://github.com/JuliaLinearAlgebra/BlockBandedMatrices.jl)) |
| 118 | + |
| 119 | +## Declaring a Sparse Jacobian with Automatic Sparsity Detection |
| 120 | + |
| 121 | +Jacobian sparsity is declared by the `jac_prototype` argument in the `NonlinearFunction`. |
| 122 | +Note that you should only do this if the sparsity is high, for example, 0.1% |
| 123 | +of the matrix is non-zeros, otherwise the overhead of sparse matrices can be higher |
| 124 | +than the gains from sparse differentiation! |
| 125 | + |
| 126 | +One of the useful companion tools for NonlinearSolve.jl is |
| 127 | +[Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl). |
| 128 | +This allows for automatic declaration of Jacobian sparsity types. To see this |
| 129 | +in action, we can give an example `du` and `u` and call `jacobian_sparsity` |
| 130 | +on our function with the example arguments, and it will kick out a sparse matrix |
| 131 | +with our pattern, that we can turn into our `jac_prototype`. |
| 132 | + |
| 133 | +```@example ill_conditioned_nlprob |
| 134 | +using Symbolics |
| 135 | +du0 = copy(u0) |
| 136 | +jac_sparsity = Symbolics.jacobian_sparsity((du, u) -> brusselator_2d_loop(du, u, p), |
| 137 | + du0, u0) |
| 138 | +``` |
| 139 | + |
| 140 | +Notice that Julia gives a nice print out of the sparsity pattern. That's neat, and |
| 141 | +would be tedious to build by hand! Now we just pass it to the `NonlinearFunction` |
| 142 | +like as before: |
| 143 | + |
| 144 | +```@example ill_conditioned_nlprob |
| 145 | +ff = NonlinearFunction(brusselator_2d_loop; jac_prototype = float.(jac_sparsity)) |
| 146 | +``` |
| 147 | + |
| 148 | +Build the `NonlinearProblem`: |
| 149 | + |
| 150 | +```@example ill_conditioned_nlprob |
| 151 | +prob_brusselator_2d_sparse = NonlinearProblem(ff, u0, p) |
| 152 | +``` |
| 153 | + |
| 154 | +Now let's see how the version with sparsity compares to the version without: |
| 155 | + |
| 156 | +```@example ill_conditioned_nlprob |
| 157 | +using BenchmarkTools # for @btime |
| 158 | +@btime solve(prob_brusselator_2d, NewtonRaphson()); |
| 159 | +@btime solve(prob_brusselator_2d_sparse, NewtonRaphson()); |
| 160 | +@btime solve(prob_brusselator_2d_sparse, NewtonRaphson(linsolve = KLUFactorization())); |
| 161 | +nothing # hide |
| 162 | +``` |
| 163 | + |
| 164 | +Note that depending on the properties of the sparsity pattern, one may want to try |
| 165 | +alternative linear solvers such as `NewtonRaphson(linsolve = KLUFactorization())` |
| 166 | +or `NewtonRaphson(linsolve = UMFPACKFactorization())` |
| 167 | + |
| 168 | +## Using Jacobian-Free Newton-Krylov |
| 169 | + |
| 170 | +A completely different way to optimize the linear solvers for large sparse |
| 171 | +matrices is to use a Krylov subspace method. This requires choosing a linear |
| 172 | +solver for changing to a Krylov method. To swap the linear solver out, we use |
| 173 | +the `linsolve` command and choose the GMRES linear solver. |
| 174 | + |
| 175 | +```@example ill_conditioned_nlprob |
| 176 | +@btime solve(prob_brusselator_2d, NewtonRaphson(linsolve = KrylovJL_GMRES())); |
| 177 | +nothing # hide |
| 178 | +``` |
| 179 | + |
| 180 | +Notice that this acceleration does not require the definition of a sparsity |
| 181 | +pattern, and can thus be an easier way to scale for large problems. For more |
| 182 | +information on linear solver choices, see the [linear solver documentation](https://docs.sciml.ai/DiffEqDocs/stable/features/linear_nonlinear/#linear_nonlinear). `linsolve` choices are any valid [LinearSolve.jl](https://linearsolve.sciml.ai/dev/) solver. |
| 183 | + |
| 184 | +!!! note |
| 185 | + |
| 186 | + Switching to a Krylov linear solver will automatically change the nonlinear problem solver |
| 187 | + into Jacobian-free mode, dramatically reducing the memory required. This can |
| 188 | + be overridden by adding `concrete_jac=true` to the algorithm. |
| 189 | + |
| 190 | +## Adding a Preconditioner |
| 191 | + |
| 192 | +Any [LinearSolve.jl-compatible preconditioner](https://docs.sciml.ai/LinearSolve/stable/basics/Preconditioners/) |
| 193 | +can be used as a preconditioner in the linear solver interface. To define |
| 194 | +preconditioners, one must define a `precs` function in compatible with nonlinear |
| 195 | +solvers which returns the left and right preconditioners, matrices which |
| 196 | +approximate the inverse of `W = I - gamma*J` used in the solution of the ODE. |
| 197 | +An example of this with using [IncompleteLU.jl](https://github.com/haampie/IncompleteLU.jl) |
| 198 | +is as follows: |
| 199 | + |
| 200 | +```@example ill_conditioned_nlprob |
| 201 | +using IncompleteLU |
| 202 | +function incompletelu(W, du, u, p, t, newW, Plprev, Prprev, solverdata) |
| 203 | + if newW === nothing || newW |
| 204 | + Pl = ilu(W, τ = 50.0) |
| 205 | + else |
| 206 | + Pl = Plprev |
| 207 | + end |
| 208 | + Pl, nothing |
| 209 | +end |
| 210 | +
|
| 211 | +@btime solve(prob_brusselator_2d_sparse, |
| 212 | + NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = incompletelu, |
| 213 | + concrete_jac = true)); |
| 214 | +nothing # hide |
| 215 | +``` |
| 216 | + |
| 217 | +Notice a few things about this preconditioner. This preconditioner uses the |
| 218 | +sparse Jacobian, and thus we set `concrete_jac=true` to tell the algorithm to |
| 219 | +generate the Jacobian (otherwise, a Jacobian-free algorithm is used with GMRES |
| 220 | +by default). Then `newW = true` whenever a new `W` matrix is computed, and |
| 221 | +`newW=nothing` during the startup phase of the solver. Thus, we do a check |
| 222 | +`newW === nothing || newW` and when true, it's only at these points when |
| 223 | +we update the preconditioner, otherwise we just pass on the previous version. |
| 224 | +We use `convert(AbstractMatrix,W)` to get the concrete `W` matrix (matching |
| 225 | +`jac_prototype`, thus `SpraseMatrixCSC`) which we can use in the preconditioner's |
| 226 | +definition. Then we use `IncompleteLU.ilu` on that sparse matrix to generate |
| 227 | +the preconditioner. We return `Pl,nothing` to say that our preconditioner is a |
| 228 | +left preconditioner, and that there is no right preconditioning. |
| 229 | + |
| 230 | +This method thus uses both the Krylov solver and the sparse Jacobian. Not only |
| 231 | +that, it is faster than both implementations! IncompleteLU is fussy in that it |
| 232 | +requires a well-tuned `τ` parameter. Another option is to use |
| 233 | +[AlgebraicMultigrid.jl](https://github.com/JuliaLinearAlgebra/AlgebraicMultigrid.jl) |
| 234 | +which is more automatic. The setup is very similar to before: |
| 235 | + |
| 236 | +```@example ill_conditioned_nlprob |
| 237 | +using AlgebraicMultigrid |
| 238 | +function algebraicmultigrid(W, du, u, p, t, newW, Plprev, Prprev, solverdata) |
| 239 | + if newW === nothing || newW |
| 240 | + Pl = aspreconditioner(ruge_stuben(convert(AbstractMatrix, W))) |
| 241 | + else |
| 242 | + Pl = Plprev |
| 243 | + end |
| 244 | + Pl, nothing |
| 245 | +end |
| 246 | +
|
| 247 | +@btime solve(prob_brusselator_2d_sparse, |
| 248 | + NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid, |
| 249 | + concrete_jac = true)); |
| 250 | +nothing # hide |
| 251 | +``` |
| 252 | + |
| 253 | +or with a Jacobi smoother: |
| 254 | + |
| 255 | +```@example ill_conditioned_nlprob |
| 256 | +function algebraicmultigrid2(W, du, u, p, t, newW, Plprev, Prprev, solverdata) |
| 257 | + if newW === nothing || newW |
| 258 | + A = convert(AbstractMatrix, W) |
| 259 | + Pl = AlgebraicMultigrid.aspreconditioner(AlgebraicMultigrid.ruge_stuben(A, |
| 260 | + presmoother = AlgebraicMultigrid.Jacobi(rand(size(A, |
| 261 | + 1))), |
| 262 | + postsmoother = AlgebraicMultigrid.Jacobi(rand(size(A, |
| 263 | + 1))))) |
| 264 | + else |
| 265 | + Pl = Plprev |
| 266 | + end |
| 267 | + Pl, nothing |
| 268 | +end |
| 269 | +
|
| 270 | +@btime solve(prob_brusselator_2d_sparse, |
| 271 | + NewtonRaphson(linsolve = KrylovJL_GMRES(), precs = algebraicmultigrid2, |
| 272 | + concrete_jac = true)); |
| 273 | +nothing # hide |
| 274 | +``` |
| 275 | + |
| 276 | +For more information on the preconditioner interface, see the |
| 277 | +[linear solver documentation](https://docs.sciml.ai/LinearSolve/stable/basics/Preconditioners/). |
0 commit comments