From b693eb3a14fed9580098f7e202b92c369dfab094 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Sat, 30 Sep 2023 12:02:23 -0700 Subject: [PATCH] Use and factor out post_implicit cb --- src/time_stepper/hc_ars343.jl | 45 ++++++++++++++-------------------- src/time_stepper/imex_ark.jl | 33 ++++++++++++++----------- src/time_stepper/imex_ssprk.jl | 35 +++++++++++++++----------- 3 files changed, 58 insertions(+), 55 deletions(-) diff --git a/src/time_stepper/hc_ars343.jl b/src/time_stepper/hc_ars343.jl index 81f53e8658c..738d8612eb9 100644 --- a/src/time_stepper/hc_ars343.jl +++ b/src/time_stepper/hc_ars343.jl @@ -39,25 +39,25 @@ function CTS.step_u!( lim!(U, p, t_exp, u) @. U += dt * a_exp[i, 1] * T_exp[1] dss!(U, p, t_exp) - # post_explicit!(U, p, t_exp) @. temp = U # used in closures + post_explicit!(U, p, t_exp) let i = i t_imp = t + dt * c_imp[i] implicit_equation_residual! = (residual, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!(residual, Ui, p, t_imp) @. residual = temp + dt * a_imp[i, i] * residual - Ui end implicit_equation_jacobian! = (jacobian, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp) end - call_post_implicit! = Ui -> begin - # post_implicit!(Ui, p, t_imp) - end + call_post_implicit! = + Ui -> begin + @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i]) + post_implicit!(Ui, p, t_imp) + end CTS.solve_newton!( newtons_method, newtons_method_cache, @@ -68,9 +68,6 @@ function CTS.step_u!( ) end - @. T_imp[i] = (U - temp) / (dt * a_imp[i, i]) - - post_explicit!(U, p, t_exp) T_lim!(T_lim[i], U, p, t_exp) T_exp!(T_exp[i], U, p, t_exp) @@ -84,25 +81,25 @@ function CTS.step_u!( @. U += dt * a_exp[i, 2] * T_exp[2] @. U += dt * a_imp[i, 2] * T_imp[2] dss!(U, p, t_exp) - # post_explicit!(U, p, t_exp) @. temp = U # used in closures + post_explicit!(U, p, t_exp) let i = i t_imp = t + dt * c_imp[i] implicit_equation_residual! = (residual, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!(residual, Ui, p, t_imp) @. residual = temp + dt * a_imp[i, i] * residual - Ui end implicit_equation_jacobian! = (jacobian, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp) end - call_post_implicit! = Ui -> begin - # post_implicit!(Ui, p, t_imp) - end + call_post_implicit! = + Ui -> begin + @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i]) + post_implicit!(Ui, p, t_imp) + end CTS.solve_newton!( newtons_method, newtons_method_cache, @@ -113,9 +110,6 @@ function CTS.step_u!( ) end - @. T_imp[i] = (U - temp) / (dt * a_imp[i, i]) - - post_explicit!(U, p, t_exp) T_lim!(T_lim[i], U, p, t_exp) T_exp!(T_exp[i], U, p, t_exp) i = 4 @@ -131,25 +125,25 @@ function CTS.step_u!( @. U += dt * a_imp[i, 2] * T_imp[2] @. U += dt * a_imp[i, 3] * T_imp[3] dss!(U, p, t_exp) - # post_explicit!(U, p, t_exp) @. temp = U # used in closures + post_explicit!(U, p, t_exp) let i = i t_imp = t + dt * c_imp[i] implicit_equation_residual! = (residual, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!(residual, Ui, p, t_imp) @. residual = temp + dt * a_imp[i, i] * residual - Ui end implicit_equation_jacobian! = (jacobian, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp) end - call_post_implicit! = Ui -> begin - # post_implicit!(Ui, p, t_imp) - end + call_post_implicit! = + Ui -> begin + @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i]) + post_implicit!(Ui, p, t_imp) + end CTS.solve_newton!( newtons_method, newtons_method_cache, @@ -160,9 +154,6 @@ function CTS.step_u!( ) end - @. T_imp[i] = (U - temp) / (dt * a_imp[i, i]) - - post_explicit!(U, p, t_exp) T_lim!(T_lim[i], U, p, t_exp) T_exp!(T_exp[i], U, p, t_exp) diff --git a/src/time_stepper/imex_ark.jl b/src/time_stepper/imex_ark.jl index bef37373df9..0f230821516 100644 --- a/src/time_stepper/imex_ark.jl +++ b/src/time_stepper/imex_ark.jl @@ -50,30 +50,42 @@ function CTS.step_u!( end dss!(U, p, t_exp) + if iszero(a_imp[i, i]) + post_explicit!(U, p, t_exp) + end if !iszero(a_imp[i, i]) # Implicit solve @assert !isnothing(newtons_method) @. temp = U + post_explicit!(U, p, t_exp) # TODO: can/should we remove these closures? implicit_equation_residual! = (residual, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!(residual, Ui, p, t_imp) @. residual = temp + dt * a_imp[i, i] * residual - Ui end implicit_equation_jacobian! = (jacobian, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp) end + call_post_implicit! = + Ui -> begin + if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && !iszero(a_imp[i, i]) + # If T_imp[i] is being treated implicitly, ensure that it + # exactly satisfies the implicit equation. + @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i]) + end + post_implicit!(Ui, p, t_imp) + end CTS.solve_newton!( newtons_method, newtons_method_cache, U, implicit_equation_residual!, implicit_equation_jacobian!, + call_post_implicit!, ) end @@ -81,21 +93,14 @@ function CTS.step_u!( # give the same results for redundant columns (as long as the implicit # tendency only acts in the vertical direction). - if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i]) - if iszero(a_imp[i, i]) - # If its coefficient is 0, T_imp[i] is effectively being - # treated explicitly. - post_implicit!(Ui, p, t_imp) - T_imp!(T_imp[i], U, p, t_imp) - else - # If T_imp[i] is being treated implicitly, ensure that it - # exactly satisfies the implicit equation. - @. T_imp[i] = (U - temp) / (dt * a_imp[i, i]) - end + if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && + iszero(a_imp[i, i]) + # If its coefficient is 0, T_imp[i] is effectively being + # treated explicitly. + T_imp!(T_imp[i], U, p, t_imp) end if !all(iszero, a_exp[:, i]) || !iszero(b_exp[i]) - post_explicit!(U, p, t_exp) T_lim!(T_lim[i], U, p, t_exp) T_exp!(T_exp[i], U, p, t_exp) end diff --git a/src/time_stepper/imex_ssprk.jl b/src/time_stepper/imex_ssprk.jl index 520730ca654..10fc7921541 100644 --- a/src/time_stepper/imex_ssprk.jl +++ b/src/time_stepper/imex_ssprk.jl @@ -50,22 +50,36 @@ function CTS.step_u!( iszero(a_imp[i, j]) && continue @. U += dt * a_imp[i, j] * T_imp[j] end + if iszero(a_imp[i, i]) + post_explicit!(U, p, t_exp) + end if !iszero(a_imp[i, i]) # Implicit solve @assert !isnothing(newtons_method) @. temp = U + post_explicit!(U, p, t_exp) # TODO: can/should we remove these closures? implicit_equation_residual! = (residual, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!(residual, Ui, p, t_imp) @. residual = temp + dt * a_imp[i, i] * residual - Ui end implicit_equation_jacobian! = (jacobian, Ui) -> begin - post_implicit!(Ui, p, t_imp) T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp) end + + call_post_implicit! = + Ui -> begin + if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && + !iszero(a_imp[i, i]) + # If T_imp[i] is being treated implicitly, ensure that it + # exactly satisfies the implicit equation. + @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i]) + end + post_implicit!(Ui, p, t_imp) + end + CTS.solve_newton!( newtons_method, newtons_method_cache, @@ -79,21 +93,14 @@ function CTS.step_u!( # give the same results for redundant columns (as long as the implicit # tendency only acts in the vertical direction). - if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i]) - if iszero(a_imp[i, i]) - # If its coefficient is 0, T_imp[i] is effectively being - # treated explicitly. - post_implicit!(U, p, t_imp) - T_imp!(T_imp[i], U, p, t_imp) - else - # If T_imp[i] is being treated implicitly, ensure that it - # exactly satisfies the implicit equation. - @. T_imp[i] = (U - temp) / (dt * a_imp[i, i]) - end + if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && + iszero(a_imp[i, i]) + # If its coefficient is 0, T_imp[i] is effectively being + # treated explicitly. + T_imp!(T_imp[i], U, p, t_imp) end if !iszero(β[i]) - post_explicit!(U, p, t_exp) T_lim!(T_lim, U, p, t_exp) T_exp!(T_exp, U, p, t_exp) end