Skip to content

Commit

Permalink
Use and factor out post_implicit cb
Browse files Browse the repository at this point in the history
  • Loading branch information
@charleskawczynski
charleskawczynski committed Sep 30, 2023
1 parent 3fe38f6 commit b693eb3
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 55 deletions.
45 changes: 18 additions & 27 deletions src/time_stepper/hc_ars343.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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)

Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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,
Expand All @@ -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)

Expand Down
33 changes: 19 additions & 14 deletions src/time_stepper/imex_ark.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,52 +50,57 @@ 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

# We do not need to DSS U again because the implicit solve should
# 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
Expand Down
35 changes: 21 additions & 14 deletions src/time_stepper/imex_ssprk.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down

0 comments on commit b693eb3

Please sign in to comment.