From 1a0e9efdc3cfbdae24a1116f21c0cc3be5087f4e Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Thu, 7 Nov 2024 17:03:03 +1100
Subject: [PATCH] Added continuation

---
 src/Main.jl       | 23 +++++++++++++++++++++-
 src/parameters.jl | 49 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/Main.jl b/src/Main.jl
index 39f40a0..a6b3efc 100644
--- a/src/Main.jl
+++ b/src/Main.jl
@@ -317,6 +317,8 @@ function _fe_operator(U,V,params)
   mfs = _multi_field_style(params)
   if has_transient(params)
     _ode_fe_operator(mfs,U,V,params)
+  elseif has_continuation(params)
+    _continuation_fe_operator(mfs,U,V,params)
   else
     _fe_operator(mfs,U,V,params)
   end
@@ -339,7 +341,7 @@ function _fe_operator(::BlockMultiFieldStyle,U,V,params)
   return FEOperator(res,jac,U,V,assem)
 end
 
-function _ode_fe_operator(::ConsecutiveMultiFieldStyle,U,V,params)
+function _ode_fe_operator(mfs,U,V,params)
   k = params[:fespaces][:k]
   res, jac, jac_t = weak_form(params,k)
   Tm = params[:solver][:matrix_type]
@@ -348,6 +350,25 @@ function _ode_fe_operator(::ConsecutiveMultiFieldStyle,U,V,params)
   return TransientFEOperator(res,(jac,jac_t),U,V,assembler=assem)
 end
 
+function _continuation_fe_operator(mfs,U,V,params)
+  niter  = params[:continuation][:niter]
+  alphas = params[:continuation][:alphas]
+  nsteps = length(niter)
+
+  ops = map(alphas) do α
+    p = duplicate_params(params)
+    p[:fluid][:α] = α
+    _fe_operator(mfs,U,V,p)
+  end
+
+  op = _fe_operator(mfs,U,V,params)
+  for step in nsteps:-1:1
+    op = ContinuationFEOperator(ops[step],op,niter[step])
+  end
+
+  return op
+end
+
 # Sub-triangulations
 
 const DiscreteModelTypes = Union{Gridap.DiscreteModel,GridapDistributed.DistributedDiscreteModel}
diff --git a/src/parameters.jl b/src/parameters.jl
index 76e3f5f..ef5edff 100644
--- a/src/parameters.jl
+++ b/src/parameters.jl
@@ -98,6 +98,7 @@ function add_default_params(_params)
     :multigrid => false,
     :check_valid => false,
     :transient => false,
+    :continuation => false,
     :x0 => false,
   )
   _check_mandatory(_params,mandatory,"")
@@ -111,6 +112,7 @@ function add_default_params(_params)
     :multigrid => nothing,
     :check_valid => true,
     :transient => nothing,
+    :continuation => nothing,
     :x0 => :zero,
   )
   params = _add_optional(_params,mandatory,optional,_params,"")
@@ -127,12 +129,35 @@ function add_default_params(_params)
   if !isnothing(params[:transient])
     params[:transient] = params_transient(params)
   end
+  if !isnothing(params[:continuation])
+    params[:continuation] = params_continuation(params)
+  end
   params
 end
 
 default_ptimer(model) = PTimer(DebugArray(LinearIndices((1,))))
 default_ptimer(model::GridapDistributed.DistributedDiscreteModel) = PTimer(get_parts(model))
 
+"""
+Duplicates all paramaters in `params`, such that 
+  - isbits objects (e.g. numbers and lighweight structs) are deep-copied. 
+  - heavier objects are copied by reference.
+
+This is useful to create a new set of parameters where 
+constants can be modified without affecting the original ones.
+"""
+function duplicate_params(params)
+  new_params = Dict{Symbol,Any}()
+  for (key,value) in params
+    if isa(value,Dict)
+      new_params[key] = duplicate_params(value)
+    else
+      new_params[key] = value
+    end
+  end
+  return new_params
+end
+
 """
 Valid keys for `params[:solver]` are the following:
 
@@ -669,3 +694,27 @@ function default_solver_params(::Val{:forward})
     :solver => :forward,
   )
 end
+
+"""
+Continuation parameters
+
+Valid keys for the dictionaries in `params[:continuation]` are the following.
+
+# Mandatory keys
+
+- `:niter`: Number of nonlinear iterations per continuation step. 
+- `:alphas`: Array of continuation parameters.
+
+"""
+function params_continuation(params::Dict{Symbol,Any})
+  mandatory = Dict{Symbol,Any}(
+    :niter => true,
+    :alphas => true,
+  )
+  optional = Dict{Symbol,Any}()
+  continuation = _check_mandatory_and_add_optional_weak(params[:continuation],mandatory,optional,params,"[:continuation]")
+  @assert length(continuation[:niter]) == length(continuation[:alphas])
+  return continuation
+end
+
+has_continuation(params) = haskey(params,:continuation) && !isnothing(params[:continuation])