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 .+ 1.0e-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 = 1.0e-8)
sol_dense_snes = NLS.solve(nlprob_dense, NLS.PETScSNES(); abstol = 1.0e-8)
sol_dense_nr .- sol_dense_snes

128-element Vector{Float64}:
 -1.864697855818134e-18
  1.2347061020361197e-11
  2.4685383368459304e-11
  3.7006144222378416e-11
  4.930045000034076e-11
  6.155934255050585e-11
  7.377380744331714e-11
  8.593478259175634e-11
  9.803316708759113e-11
  1.10059830606829e-10
  ⋮
  6.50459686113436e-11
  5.686628945511529e-11
  4.870381875576868e-11
  4.055777935718652e-11
  3.242639490252941e-11
  2.4307444945748102e-11
  1.619882006309581e-11
  8.097522652406042e-12
  0.0

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

128-element Vector{Float64}:
 -1.1350517561031018e-43
  1.2346984563567488e-11
  2.468522858928182e-11
  3.700591112229946e-11
  4.9300138583595506e-11
  6.155895281370616e-11
  7.377333937968675e-11
  8.593423618096646e-11
  9.803254237029935e-11
  1.1005912750172014e-10
  ⋮
  6.50441922545042e-11
  5.686473514288082e-11
  4.870259751044159e-11
  4.0556669134161893e-11
  3.242539570180725e-11
  2.430688983423579e-11
  1.619837597388596e-11
  8.097300607801117e-12
  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 = 1.0e-8)

BenchmarkTools.Trial: 382 samples with 1 evaluation per sample.
 Range (min … max):  12.843 ms … 43.865 ms  ┊ GC (min … max): 0.00% … 68.73%
 Time  (median):     13.012 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   13.114 ms ±  1.586 ms  ┊ GC (mean ± σ):  0.60% ±  3.52%

   ▄█▇▆▁     ▁▅▁▂          ▁                                   
  ██████▄▃▅▄▄██████▇▆▄▅▁▃▄▆█▇▄▆▆▃▄▂▃▂▃▁▁▁▁▁▃▁▂▁▁▃▂▁▂▁▁▁▁▁▂▁▁▂ ▃
  12.8 ms         Histogram: frequency by time        13.6 ms <

 Memory estimate: 701.36 KiB, allocs estimate: 603.

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

BenchmarkTools.Trial: 607 samples with 1 evaluation per sample.
 Range (min … max):  8.030 ms …  10.180 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     8.230 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   8.239 ms ± 240.045 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

   ▁     ▁█▂                                                   
  ▄██▆▃▂▂███▅▃▃▃▂▂▂▂▂▃▂▂▁▂▁▁▂▂▁▁▂▁▁▁▁▁▁▂▁▁▂▂▂▁▁▁▁▂▁▁▁▁▁▂▁▁▁▁▂ ▃
  8.03 ms         Histogram: frequency by time        9.47 ms <

 Memory estimate: 227.96 KiB, allocs estimate: 471.

Sparse Jacobian

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

BenchmarkTools.Trial: 5731 samples with 1 evaluation per sample.
 Range (min … max):  810.783 μs …  21.307 ms  ┊ GC (min … max): 0.00% … 70.66%
 Time  (median):     827.804 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   873.644 μs ± 806.908 μs  ┊ GC (mean ± σ):  3.59% ±  3.72%

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

 Memory estimate: 253.55 KiB, allocs estimate: 2749.

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

BenchmarkTools.Trial: 6318 samples with 1 evaluation per sample.
 Range (min … max):  603.473 μs … 417.665 ms  ┊ GC (min … max): 0.00% … 31.88%
 Time  (median):     653.207 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   790.878 μs ±   7.039 ms  ┊ GC (mean ± σ):  4.68% ±  0.53%

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

 Memory estimate: 183.75 KiB, allocs estimate: 2410.