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}:
 -2.323470425923449e-18
  1.2346990982277283e-11
  2.4685243751171575e-11
  3.7005935025569234e-11
  4.930017122485716e-11
  6.155899419634783e-11
  7.37733895104847e-11
  8.593429508024948e-11
  9.803261002451491e-11
  1.100592039379733e-10
  ⋮
  6.504441429910912e-11
  5.686484616518328e-11
  4.870259751044159e-11
  4.0556780156464356e-11
  3.242550672410971e-11
  2.430688983423579e-11
  1.619848699618842e-11
  8.097300607801117e-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}:
 -2.0178697886277366e-42
  1.234698080051774e-11
  2.4685221063394083e-11
  3.7005899831350276e-11
  4.9300123526737835e-11
  6.15589339960222e-11
  7.37733168011765e-11
  8.593420983485367e-11
  9.803251222947895e-11
  1.1005909362040225e-10
  ⋮
  6.504352612068942e-11
  5.6864180031368505e-11
  4.870215342123174e-11
  4.0556336067254506e-11
  3.242517365720232e-11
  2.4306667789630865e-11
  1.6198264951583496e-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: 324 samples with 1 evaluation per sample.
 Range (min … max):  15.170 ms … 48.531 ms  ┊ GC (min … max): 0.00% … 67.61%
 Time  (median):     15.372 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):   15.478 ms ±  1.847 ms  ┊ GC (mean ± σ):  0.65% ±  3.76%

   ▂▅▂▂▄             ▂            ▄█▂▄                         
  ▆█████▇▅▄▄▃▃▃▃▃▄▆█▆█▅▅▇▄▃▃▃▄▃▁▃▇█████▄▇▆▅▄▄▃▃▄▄▃▁▁▁▁▁▃▁▃▁▃▃ ▄
  15.2 ms         Histogram: frequency by time        15.7 ms <

 Memory estimate: 701.36 KiB, allocs estimate: 603.

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

BenchmarkTools.Trial: 487 samples with 1 evaluation per sample.
 Range (min … max):  10.044 ms …  12.677 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     10.248 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   10.280 ms ± 175.915 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

         ▁ ▂▁█▅▆▃▁▂▁                                            
  ▃▃▃▃▃▃▆███████████▇▆▅▅▄▅▄▃▃▃▂▂▃▂▃▂▃▂▃▁▁▁▁▂▁▂▁▁▁▁▁▂▂▁▁▁▁▁▂▁▁▂ ▃
  10 ms           Histogram: frequency by time           11 ms <

 Memory estimate: 227.96 KiB, allocs estimate: 471.

Sparse Jacobian

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

BenchmarkTools.Trial: 4908 samples with 1 evaluation per sample.
 Range (min … max):  954.628 μs …  20.626 ms  ┊ GC (min … max): 0.00% … 69.26%
 Time  (median):     979.413 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):     1.018 ms ± 692.581 μs  ┊ GC (mean ± σ):  2.53% ±  3.51%

             ▄██▆▅▂                                              
  ▂▂▂▂▂▂▂▂▂▄████████▇▅▄▄▃▃▃▃▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂▂▂▂▂▂▂ ▃
  955 μs           Histogram: frequency by time         1.06 ms <

 Memory estimate: 253.55 KiB, allocs estimate: 2749.

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

BenchmarkTools.Trial: 5527 samples with 1 evaluation per sample.
 Range (min … max):  772.979 μs … 309.056 ms  ┊ GC (min … max): 0.00% … 26.19%
 Time  (median):     831.609 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   904.042 μs ±   4.146 ms  ┊ GC (mean ± σ):  1.62% ±  0.35%

         ▅█▇▅▆▆▄▃▂           ▂▄▃▁                                
  ▂▃▅▄▄▄▆█████████▆▇▅▄▄▄▄▄▃▄▇████▆▅▃▃▃▂▂▂▃▃▂▂▂▂▁▁▁▁▂▁▁▁▁▁▁▁▁▁▂▁ ▃
  773 μs           Histogram: frequency by time         1.01 ms <

 Memory estimate: 183.75 KiB, allocs estimate: 2410.