Merge #101

101: Clean up update_j and KrylovMethod r=charleskawczynski a=dennisYatunin



Co-authored-by: Dennis Yatunin <[email protected]>

Author: bors[bot]

docs/src/newtons_method.md

@@ -44,8 +44,9 @@ MultipleConditions
 UpdateSignalHandler
 UpdateEvery
 UpdateEveryN
+UpdateEveryDt
 UpdateSignal
-NewStep
+NewTimeStep
 NewNewtonSolve
 NewNewtonIteration
 ```

src/solvers/imex_ark.jl

@@ -220,12 +220,15 @@ function cache(
             i -> Symbol(:f, χ, :_, i) => similar(u),
             filter(i -> save_tendency(i, a), i_range(a)),
         )
+    γs = unique(filter(!iszero, diag(as[2])))
+    γ = length(γs) == 1 ? γs[1] : nothing
     u = prob.u0
     Uis = map(
         i -> Symbol(:U, i) => similar(u),
         filter(i -> !(i in u_alias_is(as[1], as[2])), i_range(as[1])[1:end - 1])
     )
     _cache = NamedTuple((
+        :γ => γ,
         :U_temp => similar(u),
         Uis...,
         f_cache(:exp, as[1], typeof(prob.f.f2))...,
@@ -253,6 +256,14 @@ struct ImplicitErrorJacobian{W, P, T}
     t::T
     Δt::T
 end
+struct FirstImplicitErrorJacobian{W, U, P, T, Γ}
+    Wfact!::W
+    u::U
+    p::P
+    t::T
+    Δt::T
+    γ::Γ
+end
 
 (implicit_error::ImplicitError)(f, u) =
     implicit_error(f, u, implicit_error.ode_f!)
@@ -266,6 +277,14 @@ function ((; û, p, t, Δt)::ImplicitError)(f, u, ode_f!)
     f .= û .+ Δt .* f .- u
 end
 ((; Wfact!, p, t, Δt)::ImplicitErrorJacobian)(j, u) = Wfact!(j, u, p, Δt, t)
+function ((; Wfact!, u, p, t, Δt, γ)::FirstImplicitErrorJacobian)(j)
+    isnothing(γ) &&
+        error(
+            "Cannot compute implicit error Jacobian for timestep becasue a_imp \
+             does not have a unique value of γ. Try using a different tableau."
+        )
+    Wfact!(j, u, p, Δt * typeof(Δt)(γ), t)
+end
 
 function step_u_expr(
     ::Type{<:IMEXARKCache{as, cs}},
@@ -296,13 +315,11 @@ function step_u_expr(
         (; f1, f2) = f;
         (; newtons_method) = alg;
         (; _cache, newtons_method_cache) = cache;
-        isnothing(f1.Wfact) || run!(
-            newtons_method.update_j,
-            newtons_method_cache.update_j_cache,
-            NewStep(),
-            ImplicitErrorJacobian(f1.Wfact, p, t, dt * $(FT(as[1][end, end]))),
-            newtons_method_cache.j,
-            u,
+        isnothing(f1.Wfact) || update!(
+            newtons_method,
+            newtons_method_cache,
+            NewTimeStep(t),
+            FirstImplicitErrorJacobian(f1.Wfact, u, p, t, dt, _cache.γ),
         );
     )
 
@@ -436,12 +453,15 @@ function not_generated_cache(
             filter(i -> save_tendency(i, a), i_range(a)),
         )
 
+    γs = unique(filter(!iszero, diag(as[2])))
+    γ = length(γs) == 1 ? γs[1] : nothing
     u = prob.u0
     Uis = map(
         i -> Symbol(:U, i) => similar(u),
         filter(i -> !(i in u_alias_is(as[1], as[2])), i_range(as[1])[1:end - 1])
     )
     _cache = NamedTuple((
+        :γ => γ,
         :U_temp => similar(u),
         Uis...,
         f_cache(:exp, as[1], typeof(prob.f.f2))...,
@@ -481,13 +501,11 @@ function not_generated_step_u!(integrator, cache::IMEXARKCache{as, cs}) where {a
         f_types = (typeof(f2), typeof(f1))
         (; u_alias_is_, first_i_s, new_js_s, js_to_save_s, has_implicit_step_s, save_tendency_s, old_js_s) = _cache
 
-        isnothing(f1.Wfact) || run!(
-            newtons_method.update_j,
-            newtons_method_cache.update_j_cache,
-            NewStep(),
-            ImplicitErrorJacobian(f1.Wfact, p, t, dt * FT(as[1][end, end])),
-            newtons_method_cache.j,
-            u,
+        isnothing(f1.Wfact) || update!(
+            newtons_method,
+            newtons_method_cache,
+            NewTimeStep(t),
+            FirstImplicitErrorJacobian(f1.Wfact, u, p, t, dt, _cache.γ),
         )
 
         function Δu_broadcast(i, j, χ, a, f_type, first_i_)

src/solvers/newtons_method.jl

@@ -325,9 +325,9 @@ end
 
 """
     KrylovMethod(;
+        type = Val(Krylov.GmresSolver),
         jacobian_free_jvp = nothing,
         forcing_term = ConstantForcing(0),
-        type = Val(Krylov.GmresSolver),
         args = (20,),
         kwargs = (;),
         solve_kwargs = (;),
@@ -347,13 +347,14 @@ where `x_prototype` is `similar` to `x` (and also to `Δx` and `f`).
 
 This is primarily a wrapper for a `Krylov.KrylovSolver` from `Krylov.jl`. In
 `allocate_cache`, the solver is constructed with
-`solver = type(l, l, args..., Krylov.ktypeof(x_prototype); kwargs...)` (note
-that `type` must be passed through in a `Val` struct), where
+`solver = type(l, l, args..., Krylov.ktypeof(x_prototype); kwargs...)`, where
 `l = length(x_prototype)` and `Krylov.ktypeof(x_prototype)` is a subtype of
 `DenseVector` that can be used to store `x_prototype`. By default, the solver
 is a `Krylov.GmresSolver` with a Krylov subspace size of 20 (the default Krylov
 subspace size for this solver in `Krylov.jl`). In `run!`, the solver is run with
 `Krylov.solve!(solver, opj, f; M, ldiv, atol, rtol, verbose, solve_kwargs...)`.
+The solver's type can be changed by specifying a different value for `type`,
+though this value has to be wrapped in a `Val` to avoid runtime compilation.
 
 In the call to `Krylov.solve!`, `opj` is a `LinearOperator` that represents
 `j(x[n])`, which the solver uses by evaluating `mul!(jΔx, opj, Δx)`. If a
@@ -388,7 +389,7 @@ each iteration of the Krylov method. If a debugger is specified, it is run
 before the call to `Kyrlov.solve!`.
 """
 Base.@kwdef struct KrylovMethod{
-    T <: Val,
+    T <: Val{<:Krylov.KrylovSolver},
     J <: Union{Nothing, JacobianFreeJVP},
     F <: ForcingTerm,
     A <: Tuple,
@@ -412,7 +413,6 @@ solver_type(::KrylovMethod{Val{T}}) where {T} = T
 function allocate_cache(alg::KrylovMethod, x_prototype)
     (; jacobian_free_jvp, forcing_term, args, kwargs, debugger) = alg
     type = solver_type(alg)
-    @assert type isa Type{<:Krylov.KrylovSolver}
     l = length(x_prototype)
     return (;
         jacobian_free_jvp_cache = isnothing(jacobian_free_jvp) ? nothing :
@@ -466,7 +466,7 @@ end
 """
     NewtonsMethod(;
         max_iters = 1,
-        update_j = UpdateEvery(NewNewtonIteration()),
+        update_j = UpdateEvery(NewNewtonIteration),
         krylov_method = nothing,
         convergence_checker = nothing,
         verbose = false,
@@ -512,11 +512,23 @@ for its preconditioners, so, since the value computed with `j!` is used as a
 preconditioner in Krylov methods with a Jacobian-free JVP, using such a Krylov
 method requires specifying a `j_prototype` that can be passed to `ldiv!`.
 
-If `j(x)` changes sufficiently slowly, `update_j` can be changed from
-`UpdateEvery(NewNewtonIteration())` to some other `UpdateSignalHandler` in order
-to make the approximation `j(x[n]) ≈ j(x₀)`, where `x₀` is a previous value of
-`x[n]` (this could even be a value from a previous `run!` of Newton's method).
-When Newton's method uses this approximation, it is called the "chord method".
+If `j(x)` changes sufficiently slowly, `update_j` may be changed from
+`UpdateEvery(NewNewtonIteration)` to some other `UpdateSignalHandler` that
+gets triggered less frequently, such as `UpdateEvery(NewNewtonSolve)`. This
+can be used to make the approximation `j(x[n]) ≈ j(x₀)`, where `x₀` is a
+previous value of `x[n]` (possibly even a value from a previous `run!` of
+Newton's method). When Newton's method uses such an approximation, it is called
+the "chord method".
+
+In addition, `update_j` can be set to an `UpdateSignalHandler` that gets
+triggered by signals that originate outside of Newton's method, such as
+`UpdateEvery(NewTimeStep)`. It is possible to send any signal for updating `j`
+to Newton's method while it is not running by calling
+`update!(::NewtonsMethod, cache, ::UpdateSignal, j!)`, where in this case
+`j!(j)` is a function that sets `j` in-place without any dependence on `x`
+(since `x` is not necessarily defined while Newton's method is not running, this
+version of `j!` does not take `x` as an argument). This can be used to make the
+approximation `j(x[n]) ≈ j₀`, where `j₀` can have an arbitrary value.
 
 If a convergence checker is provided, it gets used to determine whether to stop
 iterating on iteration `n` based on the value `x[n]` and its error `Δx[n]`;
@@ -534,7 +546,7 @@ Base.@kwdef struct NewtonsMethod{
     C <: Union{Nothing, ConvergenceChecker},
 }
     max_iters::Int = 1
-    update_j::U = UpdateEvery(NewNewtonIteration())
+    update_j::U = UpdateEvery(NewNewtonIteration)
     krylov_method::K = nothing
     convergence_checker::C = nothing
     verbose::Bool = false
@@ -547,7 +559,7 @@ function allocate_cache(alg::NewtonsMethod, x_prototype, j_prototype = nothing)
         (isnothing(krylov_method) || isnothing(krylov_method.jacobian_free_jvp))
     )
     return (;
-        update_j_cache = allocate_cache(update_j),
+        update_j_cache = allocate_cache(update_j, eltype(x_prototype)),
         krylov_method_cache = isnothing(krylov_method) ? nothing :
             allocate_cache(krylov_method, x_prototype),
         convergence_checker_cache = isnothing(convergence_checker) ? nothing :
@@ -596,3 +608,9 @@ function run!(alg::NewtonsMethod, cache, x, f!, j! = nothing)
         end
     end
 end
+
+function update!(alg::NewtonsMethod, cache, signal::UpdateSignal, j!)
+    (; update_j) = alg
+    (; update_j_cache, j) = cache
+    isnothing(j) || run!(update_j, update_j_cache, signal, j!, j)
+end

src/solvers/update_signal_handler.jl

@@ -1,5 +1,5 @@
-export UpdateSignal, NewStep, NewNewtonSolve, NewNewtonIteration
-export UpdateSignalHandler, UpdateEvery, UpdateEveryN
+export UpdateSignal, NewTimeStep, NewNewtonSolve, NewNewtonIteration
+export UpdateSignalHandler, UpdateEvery, UpdateEveryN, UpdateEveryDt
 
 """
     UpdateSignal
@@ -10,11 +10,13 @@ operation is performed.
 abstract type UpdateSignal end
 
 """
-    NewStep()
+    NewTimeStep(t)
 
-The signal for a new time step.
+The signal for a new time step at time `t`.
 """
-struct NewStep <: UpdateSignal end
+struct NewTimeStep{T} <: UpdateSignal
+    t::T
+end
 
 """
     NewNewtonSolve()
@@ -37,62 +39,78 @@ struct NewNewtonIteration <: UpdateSignal end
 Updates a value upon receiving an appropriate `UpdateSignal`. This is done by
 calling `run!(::UpdateSignalHandler, cache, ::UpdateSignal, f!, args...)`, where
 `f!` is function such that `f!(args...)` modifies the desired value in-place.
-The `cache` can be obtained with `allocate_cache(::UpdateSignalHandler)`.
+The `cache` can be obtained with `allocate_cache(::UpdateSignalHandler, FT)`,
+where `FT` is the floating-point type of the integrator.
 """
 abstract type UpdateSignalHandler end
 
 """
-    UpdateEvery(update_signal)
+    UpdateEvery(update_signal_type)
 
-An `UpdateSignalHandler` that executes the update every time it is `run!` with
-`update_signal`.
+An `UpdateSignalHandler` that performs the update whenever it is `run!` with an
+`UpdateSignal` of type `update_signal_type`.
 """
-struct UpdateEvery{U <: UpdateSignal} <: UpdateSignalHandler
-    update_signal::U
-end
+struct UpdateEvery{U <: UpdateSignal} <: UpdateSignalHandler end
+UpdateEvery(::Type{U}) where {U} = UpdateEvery{U}()
 
-allocate_cache(::UpdateSignalHandler) = (;)
-
-function run!(alg::UpdateEvery{U}, cache, ::U, f!, args...) where {
-    U <: UpdateSignal,
-}
-    f!(args...)
-    return true
-end
+run!(alg::UpdateEvery{U}, cache, ::U, f!, args...) where {U} = f!(args...)
 
 """
-    UpdateEveryN(update_signal, n, reset_n_signal = nothing)
+    UpdateEveryN(n, update_signal_type, reset_signal_type = Nothing)
 
-An `UpdateSignalHandler` that executes the update every `n`-th time it is `run!`
-with `update_signal`. If `reset_n_signal` is specified, then the value of `n` is
-reset to 0 every time the signal handler is `run!` with `reset_n_signal`.
+An `UpdateSignalHandler` that performs the update every `n`-th time it is `run!`
+with an `UpdateSignal` of type `update_signal_type`. If `reset_signal_type` is
+specified, then the counter (which gets incremented from 0 to `n` and then gets
+reset to 0 when it is time to perform another update) is reset to 0 whenever the
+signal handler is `run!` with an `UpdateSignal` of type `reset_signal_type`.
 """
 struct UpdateEveryN{U <: UpdateSignal, R <: Union{Nothing, UpdateSignal}} <:
     UpdateSignalHandler
-    update_signal::U
     n::Int
-    reset_n_signal::R
 end
-UpdateEveryN(update_signal, n, reset_n_signal = nothing) =
-    UpdateEveryN(update_signal, n, reset_n_signal)
+UpdateEveryN(n, ::Type{U}, ::Type{R} = Nothing) where {U, R} =
+    UpdateEveryN{U, R}(n)
 
-allocate_cache(::UpdateEveryN) = (; n = Ref(0))
+allocate_cache(::UpdateEveryN, _) = (; counter = Ref(0))
 
-function run!(alg::UpdateEveryN{U}, cache, ::U, f!, args...) where {
-    U <: UpdateSignal,
-}
-    cache.n[] += 1
-    if cache.n[] == alg.n
+function run!(alg::UpdateEveryN{U}, cache, ::U, f!, args...) where {U}
+    (; n) = alg
+    (; counter) = cache
+    if counter[] == 0
         f!(args...)
-        cache.n[] = 0
-        return true
     end
-    return false
+    counter[] += 1
+    if counter[] == n
+        counter[] = 0
+    end
+end
+function run!(alg::UpdateEveryN{<:Any, R}, cache, ::R, f!, args...) where {R}
+    (; counter) = cache
+    counter[] = 0
 end
-function run!(alg::UpdateEveryN{U, R}, cache, ::R, f!, args...) where {
-    U,
-    R <: UpdateSignal,
-}
-    cache.n[] = 0
-    return false
+
+"""
+    UpdateEveryDt(dt)
+
+An `UpdateSignalHandler` that performs the update whenever it is `run!` with an
+`UpdateSignal` of type `NewTimeStep` and the difference between the current time
+and the previous update time is no less than `dt`.
+"""
+struct UpdateEveryDt{T} <: UpdateSignalHandler
+    dt::T
+end
+
+# TODO: This assumes that typeof(t) == FT, which might not always be correct.
+allocate_cache(alg::UpdateEveryDt, ::Type{FT}) where {FT} =
+    (; is_first_t = Ref(true), prev_update_t = Ref{FT}())
+
+function run!(alg::UpdateEveryDt, cache, signal::NewTimeStep, f!, args...)
+    (; dt) = alg
+    (; is_first_t, prev_update_t) = cache
+    (; t) = signal
+    if is_first_t[] || abs(t - prev_update_t[]) >= dt
+        f!(args...)
+        is_first_t[] = false
+        prev_update_t[] = t
+    end
 end