PETSc SNES Example 2

This implements src/snes/examples/tutorials/ex2.c from PETSc and examples/SNES_ex2.jl from PETSc.jl using automatic sparsity detection and automatic differentiation using NonlinearSolve.jl.

This solves the equations sequentially. Newton method to solve u'' + u^{2} = f, sequentially.

import NonlinearSolve as NLS
import PETSc
import LinearAlgebra
import SparseConnectivityTracer
import BenchmarkTools: @benchmark

u0 = fill(0.5, 128)

function form_residual!(resid, x, _)
    n = length(x)
    xp = LinRange(0.0, 1.0, n)
    F = 6xp .+ (xp .+ 1e-12) .^ 6

    dx = 1 / (n - 1)
    resid[1] = x[1]
    for i in 2:(n - 1)
        resid[i] = (x[i - 1] - 2x[i] + x[i + 1]) / dx^2 + x[i] * x[i] - F[i]
    end
    resid[n] = x[n] - 1

    return
end

form_residual! (generic function with 1 method)

To use automatic sparsity detection, we need to specify sparsity keyword argument to NonlinearFunction. See Automatic Sparsity Detection for more details.

nlfunc_dense = NLS.NonlinearFunction(form_residual!)
nlfunc_sparse = NLS.NonlinearFunction(
    form_residual!; sparsity = SparseConnectivityTracer.TracerSparsityDetector())

nlprob_dense = NLS.NonlinearProblem(nlfunc_dense, u0)
nlprob_sparse = NLS.NonlinearProblem(nlfunc_sparse, u0)

NonlinearProblem with uType Vector{Float64}. In-place: true
u0: 128-element Vector{Float64}:
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 ⋮
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5
 0.5

Now we can solve the problem using PETScSNES or with one of the native NonlinearSolve.jl solvers.

sol_dense_nr = NLS.solve(nlprob_dense, NLS.NewtonRaphson(); abstol = 1e-8)
sol_dense_snes = NLS.solve(nlprob_dense, NLS.PETScSNES(); abstol = 1e-8)
sol_dense_nr .- sol_dense_snes

128-element Vector{Float64}:
 1.4458254803788363e-19
 1.9232094825645703e-17
 3.831977053378455e-17
 5.740681096721295e-17
 7.649046326885234e-17
 9.557919776817525e-17
 1.1466793226749816e-16
 1.337634430303991e-16
 1.5281829621183185e-16
 1.71954464556201e-16
 ⋮
 1.1102230246251565e-16
 1.1102230246251565e-16
 1.1102230246251565e-16
 0.0
 1.1102230246251565e-16
 0.0
 0.0
 0.0
 0.0

sol_sparse_nr = NLS.solve(nlprob_sparse, NLS.NewtonRaphson(); abstol = 1e-8)
sol_sparse_snes = NLS.solve(nlprob_sparse, NLS.PETScSNES(); abstol = 1e-8)
sol_sparse_nr .- sol_sparse_snes

128-element Vector{Float64}:
 -1.1350517561031018e-43
  1.3316628480258483e-17
  2.6633256960516966e-17
  3.995115599019633e-17
  5.326820798692844e-17
  6.659034218134408e-17
  7.990570011218168e-17
  9.321428177944124e-17
  1.0657707355532509e-16
  1.1985855016827252e-16
  ⋮
  4.440892098500626e-16
  4.440892098500626e-16
  3.3306690738754696e-16
  2.220446049250313e-16
  2.220446049250313e-16
  1.1102230246251565e-16
  1.1102230246251565e-16
  0.0
  0.0

As expected the solutions are the same (upto floating point error). Now let's compare the runtimes.

Runtimes

Dense Jacobian

@benchmark NLS.solve($(nlprob_dense), $(NLS.NewtonRaphson()); abstol = 1e-8)

BenchmarkTools.Trial: 273 samples with 1 evaluation per sample.
 Range (min … max):  17.486 ms … 71.970 ms  ┊ GC (min … max): 0.00% … 72.20%
 Time  (median):     18.166 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   18.333 ms ±  3.271 ms  ┊ GC (mean ± σ):  1.04% ±  4.37%

     ▃▂▁                  ▆█▇▅▁ ▂   ▁▂                         
  ▄▄█████▇█▄▄▁▁▁▁▁▁▄▁▁▁▁▁▆███████▇▇▆██▇█▆▄▄▁▄▆▄▆▄▁▄▁▁▄▁▁▁▁▄▁▄ ▆
  17.5 ms      Histogram: log(frequency) by time        19 ms <

 Memory estimate: 1.32 MiB, allocs estimate: 1122.

@benchmark NLS.solve($(nlprob_dense), $(NLS.PETScSNES()); abstol = 1e-8)

BenchmarkTools.Trial: 598 samples with 1 evaluation per sample.
 Range (min … max):  8.019 ms …  11.524 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     8.144 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   8.370 ms ± 563.622 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▃▇█▆▃▃▃▁                                                     
  █████████▆▄▆▆▆▆▇▄▁▇▆▅▄▆▁▄▆▄▅▆▆▆▅▄▆▆▅▅▄▄▅▄▆▄▁▄▆▄▁▅▄▆▆▅▄▄▁▁▅▄ ▇
  8.02 ms      Histogram: log(frequency) by time      10.5 ms <

 Memory estimate: 229.66 KiB, allocs estimate: 452.

Sparse Jacobian

@benchmark NLS.solve($(nlprob_sparse), $(NLS.NewtonRaphson()); abstol = 1e-8)

BenchmarkTools.Trial: 4754 samples with 1 evaluation per sample.
 Range (min … max):  948.973 μs … 66.934 ms  ┊ GC (min … max): 0.00% … 76.84%
 Time  (median):     990.654 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):     1.051 ms ±  1.522 ms  ┊ GC (mean ± σ):  3.64% ±  2.61%

    ▅██▇▄▃  ▂▃▄▃▂▁                                              
  ▃████████▇██████▇▅▅▃▃▃▂▂▂▂▂▂▂▂▁▁▂▁▂▁▂▁▁▂▁▁▁▁▁▁▁▂▁▂▁▁▁▁▁▁▁▁▁▁ ▃
  949 μs          Histogram: frequency by time          1.2 ms <

 Memory estimate: 253.53 KiB, allocs estimate: 2685.

@benchmark NLS.solve($(nlprob_sparse), $(NLS.PETScSNES()); abstol = 1e-8)

BenchmarkTools.Trial: 6761 samples with 1 evaluation per sample.
 Range (min … max):  601.732 μs … 475.020 ms  ┊ GC (min … max): 0.00% … 29.63%
 Time  (median):     646.193 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   738.992 μs ±   5.769 ms  ┊ GC (mean ± σ):  2.82% ±  0.36%

   ▃█▄                        ▁                                  
  ▄████▇█▇▇▇▇▆▅▅▄▃▃▃▃▂▂▂▂▂▂▂▅███▇▇▆▄▄▄▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▂▁▂▂▂▂▂ ▃
  602 μs           Histogram: frequency by time          863 μs <

 Memory estimate: 185.41 KiB, allocs estimate: 2392.