Skip to content

Commit

Permalink
Use OrdinaryDiffEq.jl for timestepping
Browse files Browse the repository at this point in the history
  • Loading branch information
@charleskawczynski
charleskawczynski committed Oct 30, 2021
1 parent ffc24ab commit c88579d
Show file tree
Hide file tree
Showing 2 changed files with 116 additions and 76 deletions.
133 changes: 109 additions & 24 deletions integration_tests/utils/main.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -242,36 +242,121 @@ function run(sim::Simulation1d)
grid = sim.grid
state = sim.state
TC.open_files(sim.Stats) # #removeVarsHack
while sim.TS.t <= sim.TS.t_max
TC.update(sim.Turb, grid, state, sim.GMV, sim.Case, sim.TS)
TC.update(sim.TS)

if mod(iter, 100) == 0
progress = sim.TS.t / sim.TS.t_max
@show progress
end
t_span = (0.0, sim.TS.t_max)
params = (;
edmf = sim.Turb,
grid = grid,
state = state,
gm = sim.GMV,
Case = sim.Case,
TS = sim.TS,
Stats = sim.Stats,
io_nt = sim.io_nt,
)

if mod(round(Int, sim.TS.t), round(Int, sim.Stats.frequency)) == 0
# TODO: is this the best location to call diagnostics?
TC.compute_diagnostics!(sim.Turb, sim.GMV, grid, state, sim.Case, sim.TS)
condition_io(u, t, integrator) = mod(round(Int, t), round(Int, sim.Stats.frequency)) == 0
condition_every_iter(u, t, integrator) = true

function affect_io!(integrator)
UnPack.@unpack edmf, io_nt, grid, state, Case, gm, TS, Stats = integrator.p
parent(state.prog) .= integrator.u
param_set = TC.parameter_set(gm)
# TODO: is this the best location to call diagnostics?
TC.compute_diagnostics!(edmf, gm, grid, state, Case, TS)

# TODO: remove `vars` hack that avoids
# https://github.com/Alexander-Barth/NCDatasets.jl/issues/135
# opening/closing files every step should be okay. #removeVarsHack
# TurbulenceConvection.io(sim) # #removeVarsHack
TC.write_simulation_time(Stats, TS.t) # #removeVarsHack

TC.io(io_nt.aux, Stats)
TC.io(io_nt.diagnostics, Stats)
TC.io(io_nt.prog, Stats)
TC.io(io_nt.tendencies, Stats)

TC.io(gm, grid, state, Stats) # #removeVarsHack
TC.io(Case, grid, state, Stats) # #removeVarsHack
TC.io(edmf, grid, state, Stats, TS, param_set) # #removeVarsHack
end

# TODO: remove `vars` hack that avoids
# https://github.com/Alexander-Barth/NCDatasets.jl/issues/135
# opening/closing files every step should be okay. #removeVarsHack
# TurbulenceConvection.io(sim) # #removeVarsHack
TC.write_simulation_time(sim.Stats, sim.TS.t) # #removeVarsHack
function affect_filter!(integrator)
UnPack.@unpack edmf, grid, state, gm, Case, TS = integrator.p
parent(state.prog) .= parent(integrator.u)
prog_en = TC.center_prog_environment(state)
up = edmf.UpdVar
###
### Filters
###
TC.set_edmf_surface_bc(edmf, grid, state, up, Case.Sur)
TC.filter_updraft_vars(edmf, grid, state, gm)

@inbounds for k in TC.real_center_indices(grid)
prog_en.tke[k] = max(prog_en.tke[k], 0.0)
prog_en.Hvar[k] = max(prog_en.Hvar[k], 0.0)
prog_en.QTvar[k] = max(prog_en.QTvar[k], 0.0)
prog_en.HQTcov[k] = max(prog_en.HQTcov[k], -sqrt(prog_en.Hvar[k] * prog_en.QTvar[k]))
prog_en.HQTcov[k] = min(prog_en.HQTcov[k], sqrt(prog_en.Hvar[k] * prog_en.QTvar[k]))
end
parent(integrator.u) .= parent(state.prog)
end

TC.io(sim.io_nt.aux, sim.Stats)
TC.io(sim.io_nt.diagnostics, sim.Stats)
TC.io(sim.io_nt.prog, sim.Stats)
TC.io(sim.io_nt.tendencies, sim.Stats)
# TODO: this should not be handled via callbacks
function affect_implicit!(integrator)
UnPack.@unpack edmf, grid, state, gm, TS = integrator.p
parent(state.prog) .= parent(integrator.u)
implicit_eqs = edmf.implicit_eqs
# Matrix is the same for all variables that use the same eddy diffusivity, we can construct once and reuse
param_set = TC.parameter_set(gm)
up = edmf.UpdVar
prog_en = TC.center_prog_environment(state)

common_args = (
grid,
param_set,
state,
TS,
up.n_updrafts,
edmf.minimum_area,
edmf.pressure_plume_spacing,
edmf.frac_turb_entr,
edmf.entr_sc,
edmf.mixing_length,
)

implicit_eqs.A_TKE .= TC.construct_tridiag_diffusion_en(common_args..., true)
implicit_eqs.A_Hvar .= TC.construct_tridiag_diffusion_en(common_args..., false)
implicit_eqs.A_QTvar .= TC.construct_tridiag_diffusion_en(common_args..., false)
implicit_eqs.A_HQTcov .= TC.construct_tridiag_diffusion_en(common_args..., false)

implicit_eqs.b_TKE .= TC.en_diffusion_tendencies(grid, state, TS, :tke, up.n_updrafts)
implicit_eqs.b_Hvar .= TC.en_diffusion_tendencies(grid, state, TS, :Hvar, up.n_updrafts)
implicit_eqs.b_QTvar .= TC.en_diffusion_tendencies(grid, state, TS, :QTvar, up.n_updrafts)
implicit_eqs.b_HQTcov .= TC.en_diffusion_tendencies(grid, state, TS, :HQTcov, up.n_updrafts)
parent(prog_en.tke) .= implicit_eqs.A_TKE \ implicit_eqs.b_TKE
parent(prog_en.Hvar) .= implicit_eqs.A_Hvar \ implicit_eqs.b_Hvar
parent(prog_en.QTvar) .= implicit_eqs.A_QTvar \ implicit_eqs.b_QTvar
parent(prog_en.HQTcov) .= implicit_eqs.A_HQTcov \ implicit_eqs.b_HQTcov

parent(integrator.u) .= parent(state.prog)

TC.io(sim.GMV, grid, state, sim.Stats) # #removeVarsHack
TC.io(sim.Case, grid, state, sim.Stats) # #removeVarsHack
TC.io(sim.Turb, grid, state, sim.Stats, sim.TS, sim.param_set) # #removeVarsHack
end
iter += 1
end
callback_io = ODE.DiscreteCallback(condition_io, affect_io!)
callback_filters = ODE.DiscreteCallback(condition_every_iter, affect_filter!)
callback_implicit_solve = ODE.DiscreteCallback(condition_every_iter, affect_implicit!)

prob = ODE.ODEProblem(TC.step!, state.prog, t_span, params; dt = sim.TS.dt)
sol = ODE.solve(
prob,
ODE.Euler(),
reltol = 1e-12,
abstol = 1e-12,
callback = ODE.CallbackSet(callback_implicit_solve, callback_filters, callback_io),
progress = true,
progress_message = (dt, u, p, t) -> t,
)

TC.close_files(sim.Stats) # #removeVarsHack
return
end
Expand Down
59 changes: 7 additions & 52 deletions src/Turbulence_PrognosticTKE.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -233,7 +233,12 @@ function compute_diffusive_fluxes(
end

# Perform the update of the scheme
function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Case::CasesBase, TS::TimeStepping)
function step!(dY, Y, params, t)
UnPack.@unpack edmf, grid, state, gm, Case, TS = params
TS.t = t
parent(state.prog) .= parent(Y)
cprog = deepcopy(Y.cent)
fprog = deepcopy(Y.face)

gm = gm
up = edmf.UpdVar
Expand Down Expand Up @@ -275,64 +280,14 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca
# compute tendencies
compute_gm_tendencies!(edmf, grid, state, Case, gm, TS)
compute_updraft_tendencies(edmf, grid, state, gm)
# ----------- TODO: move to compute_tendencies
implicit_eqs = edmf.implicit_eqs
# Matrix is the same for all variables that use the same eddy diffusivity, we can construct once and reuse

common_args = (
grid,
param_set,
state,
TS,
up.n_updrafts,
edmf.minimum_area,
edmf.pressure_plume_spacing,
edmf.frac_turb_entr,
edmf.entr_sc,
edmf.mixing_length,
)

implicit_eqs.A_TKE .= construct_tridiag_diffusion_en(common_args..., true)
implicit_eqs.A_Hvar .= construct_tridiag_diffusion_en(common_args..., false)
implicit_eqs.A_QTvar .= construct_tridiag_diffusion_en(common_args..., false)
implicit_eqs.A_HQTcov .= construct_tridiag_diffusion_en(common_args..., false)

implicit_eqs.b_TKE .= en_diffusion_tendencies(grid, state, TS, :tke, n_updrafts)
implicit_eqs.b_Hvar .= en_diffusion_tendencies(grid, state, TS, :Hvar, n_updrafts)
implicit_eqs.b_QTvar .= en_diffusion_tendencies(grid, state, TS, :QTvar, n_updrafts)
implicit_eqs.b_HQTcov .= en_diffusion_tendencies(grid, state, TS, :HQTcov, n_updrafts)
# -----------

###
### update
###
update_updraft(edmf, grid, state, gm, TS)
if edmf.Precip.precipitation_model == "clima_1m"
update_precipitation(edmf.Precip, grid, state, up_thermo, en_thermo, edmf.PrecipPhys, TS)
end
parent(prog_en.tke) .= implicit_eqs.A_TKE \ implicit_eqs.b_TKE
parent(prog_en.Hvar) .= implicit_eqs.A_Hvar \ implicit_eqs.b_Hvar
parent(prog_en.QTvar) .= implicit_eqs.A_QTvar \ implicit_eqs.b_QTvar
parent(prog_en.HQTcov) .= implicit_eqs.A_HQTcov \ implicit_eqs.b_HQTcov
@inbounds for k in real_center_indices(grid)
prog_gm.u[k] += tendencies_gm.u[k] * TS.dt
prog_gm.v[k] += tendencies_gm.v[k] * TS.dt
prog_gm.θ_liq_ice[k] += tendencies_gm.θ_liq_ice[k] * TS.dt
prog_gm.q_tot[k] += tendencies_gm.q_tot[k] * TS.dt
end

###
### Filters
###
set_edmf_surface_bc(edmf, grid, state, up, Case.Sur)
filter_updraft_vars(edmf, grid, state, gm)
@inbounds for k in real_center_indices(grid)
prog_en.tke[k] = max(prog_en.tke[k], 0.0)
prog_en.Hvar[k] = max(prog_en.Hvar[k], 0.0)
prog_en.QTvar[k] = max(prog_en.QTvar[k], 0.0)
prog_en.HQTcov[k] = max(prog_en.HQTcov[k], -sqrt(prog_en.Hvar[k] * prog_en.QTvar[k]))
prog_en.HQTcov[k] = min(prog_en.HQTcov[k], sqrt(prog_en.Hvar[k] * prog_en.QTvar[k]))
end
parent(dY) .= parent(state.tendencies)

return
end
Expand Down

0 comments on commit c88579d

Please sign in to comment.