Use OrdinaryDiffEq.jl for timestepping, rearrange update_aux!, and remove buoyancy extrapolation

integration_tests/utils/generate_namelist.jl

@@ -139,7 +139,6 @@ function default_namelist(case_name::String)
 
     namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["updraft_number"] = 1
     namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["entrainment"] = "moisture_deficit"
-    namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["extrapolate_buoyancy"] = true
     namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["use_local_micro"] = true
     namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["constant_area"] = false
     namelist_defaults["turbulence"]["EDMF_PrognosticTKE"]["calculate_tke"] = true

integration_tests/utils/main.jl

@@ -232,36 +232,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

src/Turbulence_PrognosticTKE.jl

@@ -382,7 +382,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
@@ -424,64 +429,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.Rain.rain_model == "clima_1m"
         update_rain(edmf.Rain, grid, state, up_thermo, en_thermo, edmf.RainPhys, 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

src/types.jl

@@ -499,7 +499,6 @@ mutable struct EDMF_PrognosticTKE{A1, A2, IE}
     n_updrafts::Int
     drag_sign::Int
     asp_label
-    extrapolate_buoyancy::Bool
     surface_area::Float64
     max_area::Float64
     aspect_ratio::Float64
@@ -587,9 +586,6 @@ mutable struct EDMF_PrognosticTKE{A1, A2, IE}
             default = "const",
         )
 
-        extrapolate_buoyancy =
-            parse_namelist(namelist, "turbulence", "EDMF_PrognosticTKE", "extrapolate_buoyancy"; default = true)
-
         # Get values from namelist
         # set defaults at some point?
         surface_area = namelist["turbulence"]["EDMF_PrognosticTKE"]["surface_area"]
@@ -745,7 +741,6 @@ mutable struct EDMF_PrognosticTKE{A1, A2, IE}
             n_updrafts,
             drag_sign,
             asp_label,
-            extrapolate_buoyancy,
             surface_area,
             max_area,
             aspect_ratio,

src/update_aux.jl

@@ -84,14 +84,15 @@ function update_aux!(edmf, gm, grid, state, Case, param_set, TS)
     #####
     ##### diagnose_GMV_moments
     #####
-    #! format: off
+
     get_GMV_CoVar(edmf, grid, state, :Hvar, :θ_liq_ice)
     get_GMV_CoVar(edmf, grid, state, :QTvar, :q_tot)
     get_GMV_CoVar(edmf, grid, state, :HQTcov, :θ_liq_ice, :q_tot)
     GMV_third_m(edmf, grid, state, :Hvar, :θ_liq_ice, :H_third_m)
     GMV_third_m(edmf, grid, state, :QTvar, :q_tot, :QT_third_m)
     GMV_third_m(edmf, grid, state, :tke, :w, :W_third_m)
-    #! format: on
+
+    satadjust(gm, grid, state)
 
     #####
     ##### decompose_environment
@@ -115,6 +116,26 @@ function update_aux!(edmf, gm, grid, state, Case, param_set, TS)
     end
 
     @inbounds for k in real_center_indices(grid)
+
+        #####
+        ##### updraft buoyancy
+        #####
+
+        @inbounds for i in 1:(up.n_updrafts)
+            if aux_up[i].area[k] > 0.0
+                ts_up = thermo_state_pθq(param_set, p0_c[k], aux_up[i].θ_liq_ice[k], aux_up[i].q_tot[k])
+                aux_up[i].q_liq[k] = TD.liquid_specific_humidity(ts_up)
+                aux_up[i].q_ice[k] = TD.ice_specific_humidity(ts_up)
+                aux_up[i].T[k] = TD.air_temperature(ts_up)
+                ρ = TD.air_density(ts_up)
+                aux_up[i].buoy[k] = buoyancy_c(param_set, ρ0_c[k], ρ)
+                aux_up[i].RH[k] = TD.relative_humidity(ts_up)
+            else
+                aux_up[i].buoy[k] = aux_en.buoy[k]
+                aux_up[i].RH[k] = aux_en.RH[k]
+            end
+        end
+
         a_bulk_c = aux_tc.bulk.area[k]
         aux_tc.bulk.q_tot[k] = 0
         aux_tc.bulk.q_liq[k] = 0
@@ -167,36 +188,9 @@ function update_aux!(edmf, gm, grid, state, Case, param_set, TS)
         aux_en.RH[k] = TD.relative_humidity(ts_en)
 
         #####
-        ##### buoyancy
+        ##### grid mean buoyancy
         #####
 
-        @inbounds for i in 1:(up.n_updrafts)
-            if aux_up[i].area[k] > 0.0
-                ts_up = thermo_state_pθq(param_set, p0_c[k], aux_up[i].θ_liq_ice[k], aux_up[i].q_tot[k])
-                aux_up[i].q_liq[k] = TD.liquid_specific_humidity(ts_up)
-                aux_up[i].q_ice[k] = TD.ice_specific_humidity(ts_up)
-                aux_up[i].T[k] = TD.air_temperature(ts_up)
-                ρ = TD.air_density(ts_up)
-                aux_up[i].buoy[k] = buoyancy_c(param_set, ρ0_c[k], ρ)
-                aux_up[i].RH[k] = TD.relative_humidity(ts_up)
-            elseif k > kc_surf
-                if aux_up[i].area[k - 1] > 0.0 && edmf.extrapolate_buoyancy
-                    qt = aux_up[i].q_tot[k - 1]
-                    h = aux_up[i].θ_liq_ice[k - 1]
-                    ts_up = thermo_state_pθq(param_set, p0_c[k], h, qt)
-                    ρ = TD.air_density(ts_up)
-                    aux_up[i].buoy[k] = buoyancy_c(param_set, ρ0_c[k], ρ)
-                    aux_up[i].RH[k] = TD.relative_humidity(ts_up)
-                else
-                    aux_up[i].buoy[k] = aux_en.buoy[k]
-                    aux_up[i].RH[k] = aux_en.RH[k]
-                end
-            else
-                aux_up[i].buoy[k] = aux_en.buoy[k]
-                aux_up[i].RH[k] = aux_en.RH[k]
-            end
-        end
-
         aux_gm.buoy[k] = (1.0 - aux_tc.bulk.area[k]) * aux_en.buoy[k]
         @inbounds for i in 1:(up.n_updrafts)
             aux_gm.buoy[k] += aux_up[i].area[k] * aux_up[i].buoy[k]
@@ -225,7 +219,7 @@ function update_aux!(edmf, gm, grid, state, Case, param_set, TS)
     @inbounds for k in real_center_indices(grid)
         aux_gm.q_liq[k] = (a_up_bulk[k] * aux_tc.bulk.q_liq[k] + (1 - a_up_bulk[k]) * aux_en.q_liq[k])
         aux_gm.q_ice[k] = (a_up_bulk[k] * aux_tc.bulk.q_ice[k] + (1 - a_up_bulk[k]) * aux_en.q_ice[k])
-        aux_gm.T[k] = (a_up_bulk[k] * aux_tc.bulk.T[k] + (1 - a_up_bulk[k]) * aux_en.T[k])
+        # aux_gm.T[k] = (a_up_bulk[k] * aux_tc.bulk.T[k] + (1 - a_up_bulk[k]) * aux_en.T[k])
         aux_gm.buoy[k] = (a_up_bulk[k] * aux_tc.bulk.buoy[k] + (1 - a_up_bulk[k]) * aux_en.buoy[k])
     end
     compute_pressure_plume_spacing(edmf, param_set)