Added ContinuationFEOperators

src/NonlinearSolvers/ContinuationFEOperators.jl

@@ -0,0 +1,140 @@
+
+"""
+    mutable struct ContinuationSwitch{A}
+      callback :: Function
+      caches   :: A
+      switched :: Bool
+    end
+
+    ContinuationSwitch(callback::Function, caches)
+
+Switch that implements the switching logic for a ContinuationFEOperator.
+
+The callback function must provide the following signatures:
+
+- Reset: `callback(u::FEFunction,::Nothing,cache) -> (switch::Bool, new_cache)`
+- Update: `callback(u::FEFunction,b::AbstractVector,cache) -> (switch::Bool, new_cache)`
+
+where `u` is the current solution (as a FEFunction) and `b` is the current residual. The first 
+signature is called by `allocate_residual` and is typically used to reset the switch for a new 
+nonlinear solve. The second signature is called by `residual!` and is used to update the switch
+based on the current residual.
+
+The cache is (potentially) mutated between each call, and can hold any extra information 
+needed for the continuation logic.
+"""
+mutable struct ContinuationSwitch{A}
+  callback :: Function
+  caches   :: A
+  switched :: Bool
+
+  function ContinuationSwitch(callback::Function, caches)
+    A = typeof(caches)
+    new{A}(callback, caches, false)
+  end
+end
+
+has_switched(s::ContinuationSwitch) = s.switched
+
+switch!(s::ContinuationSwitch, u) = switch!(s, u, nothing)
+
+function switch!(s::ContinuationSwitch, u, b)
+  has_switched(s) && return s.switched
+  s.switched, s.caches = s.callback(u, b, s.caches)
+  if has_switched(s)
+    @debug "ContinuationFEOperator: Switching operators!"
+  end
+  return s.switched
+end
+
+"""
+  ContinuationSwitch(niter::Int)
+
+Switch that will change operators after `niter` iterations.
+"""
+function ContinuationSwitch(niter::Int)
+  caches = (;it = -1, niter = niter)
+  callback(u,::Nothing,c) = (false, (;it = -1, niter = c.niter))
+  callback(u,b,c) = (c.it+1 >= c.niter, (;it = c.it+1, niter = c.niter))
+  return ContinuationSwitch(callback, caches)
+end
+
+"""
+    struct ContinuationFEOperator <: FESpaces.FEOperator
+      op1    :: FEOperator
+      op2    :: FEOperator
+      switch :: ContinuationSwitch
+      reuse_caches :: Bool
+    end
+
+FEOperator implementing the continuation method for nonlinear solvers. It switches between 
+its two operators when the switch is triggered.
+
+Continuation between more that two operators can be achieved by daisy-chaining two or more 
+ContinuationFEOperators.
+
+If `reuse_caches` is `true`, the Jacobian of the first operator is reused for the second 
+operator. This is only possible if the sparsity pattern of the Jacobian does not change.
+"""
+struct ContinuationFEOperator{A,B,C} <: FESpaces.FEOperator
+  op1 :: A
+  op2 :: B
+  switch :: C
+  reuse_caches :: Bool
+
+  function ContinuationFEOperator(
+    op1::FEOperator,
+    op2::FEOperator,
+    switch::ContinuationSwitch;
+    reuse_caches::Bool = true
+  )
+    A, B, C = typeof(op1), typeof(op2), typeof(switch)
+    new{A,B,C}(op1, op2, switch, reuse_caches)
+  end
+end
+
+has_switched(op::ContinuationFEOperator) = has_switched(op.switch)
+
+"""
+  ContinuationFEOperator(op1::FEOperator, op2::FEOperator, niter::Int; reuse_caches::Bool = true)
+
+ContinuationFEOperator that switches between `op1` and `op2` after `niter` iterations.
+"""
+function ContinuationFEOperator(
+  op1::FEOperator, op2::FEOperator, niter::Int; reuse_caches::Bool = true
+)
+  switch = ContinuationSwitch(niter)
+  return ContinuationFEOperator(op1, op2, switch; reuse_caches)
+end
+
+# FEOperator API
+
+function FESpaces.get_test(op::ContinuationFEOperator)
+  ifelse(!has_switched(op), get_test(op.op1), get_test(op.op2))
+end
+
+function FESpaces.get_trial(op::ContinuationFEOperator)
+  ifelse(!has_switched(op), get_trial(op.op1), get_trial(op.op2))
+end
+
+function Algebra.allocate_residual(op::ContinuationFEOperator,u)
+  switch!(op.switch, u, nothing)
+  ifelse(!has_switched(op), allocate_residual(op.op1, u), allocate_residual(op.op2, u))
+end
+
+function Algebra.residual!(b::AbstractVector,op::ContinuationFEOperator,u)
+  switch!(op.switch, u, b)
+  ifelse(!has_switched(op), residual!(b, op.op1, u), residual!(b, op.op2, u))
+end
+
+function Algebra.allocate_jacobian(op::ContinuationFEOperator,u)
+  ifelse(!has_switched(op), allocate_jacobian(op.op1, u), allocate_jacobian(op.op2, u))
+end
+
+function Algebra.jacobian!(A::AbstractMatrix,op::ContinuationFEOperator,u)
+  if op.reuse_caches
+    ifelse(!has_switched(op), jacobian!(A, op.op1, u), jacobian!(A, op.op2, u))
+  else
+    ifelse(!has_switched(op), jacobian(op.op1, u), jacobian(op.op2, u))
+  end
+end

src/NonlinearSolvers/NewtonRaphsonSolver.jl

@@ -19,6 +19,8 @@ function NewtonSolver(ls;maxiter=100,atol=1e-12,rtol=1.e-6,verbose=0,name="Newto
   return NewtonSolver(ls,log)
 end
 
+AbstractTrees.children(s::NewtonSolver) = [s.ls]
+
 struct NewtonCache
   A::AbstractMatrix
   b::AbstractVector

src/NonlinearSolvers/NonlinearSolvers.jl

@@ -1,4 +1,5 @@
 module NonlinearSolvers
+  using AbstractTrees
   using LinearAlgebra
   using SparseArrays
   using SparseMatricesCSR
@@ -13,6 +14,9 @@ module NonlinearSolvers
   using GridapSolvers.SolverInterfaces
   using GridapSolvers.MultilevelTools
 
+  include("ContinuationFEOperators.jl")
+  export ContinuationFEOperator
+
   include("NewtonRaphsonSolver.jl")
   export NewtonSolver
 

test/NonlinearSolvers/NonlinearSolversTests.jl

@@ -13,7 +13,9 @@ using Gridap.Algebra
 using GridapSolvers
 using GridapSolvers.NonlinearSolvers, GridapSolvers.LinearSolvers
 
-function main(ranks,model,solver)
+const ModelTypes = Union{<:DiscreteModel,<:GridapDistributed.DistributedDiscreteModel}
+
+function main(ranks,model::ModelTypes,solver::Symbol)
   u_sol(x) = sum(x)
 
   order = 1
@@ -40,6 +42,9 @@ function main(ranks,model,solver)
     nls = NLsolveNonlinearSolver(ls; show_trace=true, method=:anderson, linesearch=BackTracking(), iterations=40, m=10)
   elseif solver == :newton
     nls = NewtonSolver(ls,maxiter=20,verbose=true)
+  elseif solver == :newton_continuation
+    op = ContinuationFEOperator(op,op,2;reuse_caches=true)
+    nls = NewtonSolver(ls,maxiter=20,verbose=true)
   else
     @error "Unknown solver"
   end
@@ -52,15 +57,15 @@ function main(ranks,model,solver)
 end
 
 # Serial
-function main(solver)
+function main(solver::Symbol)
   model = CartesianDiscreteModel((0,1,0,1),(8,8))
   ranks = DebugArray([1])
   main(ranks,model,solver)
 end
 
 # Distributed
-function main(distribute,solver)
-  ranks = distribute(LinearIndices((4,)))
+function main(distribute,np,solver::Symbol)
+  ranks = distribute(LinearIndices((prod(np),)))
   model = CartesianDiscreteModel((0,1,0,1),(8,8))
   main(ranks,model,solver)
 end

test/NonlinearSolvers/mpi/NonlinearSolversTests.jl

@@ -0,0 +1,10 @@
+module NonlinearSolversTestsMPI
+using MPI, PartitionedArrays
+include("../NonlinearSolversTests.jl")
+
+with_mpi() do distribute
+  NonlinearSolversTests.main(distribute,4,:newton)
+  NonlinearSolversTests.main(distribute,4,:newton_continuation)
+end
+
+end

test/NonlinearSolvers/mpi/runtests.jl

@@ -0,0 +1,20 @@
+using Test
+using MPI
+using GridapSolvers
+
+function run_tests(testdir)
+  istest(f) = endswith(f, ".jl") && !(f=="runtests.jl")
+  testfiles = sort(filter(istest, readdir(testdir)))
+  @time @testset "$f" for f in testfiles
+    MPI.mpiexec() do cmd
+      np = 4
+      cmd = `$cmd -n $(np) --oversubscribe $(Base.julia_cmd()) --project=. $(joinpath(testdir, f))`
+      @show cmd
+      run(cmd)
+      @test true
+    end
+  end
+end
+
+# MPI tests
+run_tests(@__DIR__)

test/NonlinearSolvers/seq/runtests.jl

@@ -1,15 +1,29 @@
 
 using Test
+using PartitionedArrays
 
 include("../NonlinearSolversTests.jl")
 
-@testset "NonlinearSolvers" begin
+@testset "NonlinearSolvers - Serial" begin
   @testset "NLSolvers" begin
     NonlinearSolversTests.main(:nlsolvers_newton)
     NonlinearSolversTests.main(:nlsolvers_trust_region)
     NonlinearSolversTests.main(:nlsolvers_anderson)
   end
   @testset "Newton" begin
     NonlinearSolversTests.main(:newton)
+    NonlinearSolversTests.main(:newton_continuation)
+  end
+end
+
+@testset "NonlinearSolvers - Sequential" begin
+  @testset "NLSolvers" begin
+    NonlinearSolversTests.main(DebugArray,4,:nlsolvers_newton)
+    NonlinearSolversTests.main(DebugArray,4,:nlsolvers_trust_region)
+    NonlinearSolversTests.main(DebugArray,4,:nlsolvers_anderson)
+  end
+  @testset "Newton" begin
+    NonlinearSolversTests.main(DebugArray,4,:newton)
+    NonlinearSolversTests.main(DebugArray,4,:newton_continuation)
   end
 end