This is a combination of 7 commits.

Use OrdinaryDiffEq.jl for timestepping Try calling affect inside RHS function Hack for dry bubble, rm copied data Try perf patch Remove update precip from tendencies Apply formatter Bug fix
CliMA · Nov 2, 2021 · 93ece9b · 93ece9b
1 parent cf9d58d
commit 93ece9b
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 77 deletions.
diff --git a/integration_tests/utils/main.jl b/integration_tests/utils/main.jl
@@ -246,36 +246,64 @@ 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
-
-        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)
-
-            # 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
-
-            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)
-
-            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
+
+    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,
+    )
+
+    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
+
+    function affect_filter!(integrator)
+        UnPack.@unpack edmf, grid, state, gm, Case, TS = integrator.p
+        TC.affect_filter!(integrator.u, edmf, grid, state, gm, Case, TS)
     end
+
+    callback_io = ODE.DiscreteCallback(condition_io, affect_io!)
+    callback_filters = ODE.DiscreteCallback(condition_every_iter, affect_filter!)
+
+    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_filters, callback_io),
+        progress = true,
+        progress_message = (dt, u, p, t) -> t,
+    )
+
     TC.close_files(sim.Stats) # #removeVarsHack
     return
 end

diff --git a/src/Fields.jl b/src/Fields.jl
@@ -89,3 +89,16 @@ function FieldFromNamedTuple(space, nt::NamedTuple)
     cmv(z) = nt
     return cmv.(CC.Fields.coordinate_field(space))
 end
+
+# https://github.com/CliMA/ClimaCore.jl/issues/275
+transform_broadcasted(bc::Base.Broadcast.Broadcasted{CC.Fields.FieldVectorStyle}, symb, axes) =
+    Base.Broadcast.Broadcasted(bc.f, map(arg -> transform_broadcasted(arg, symb, axes), bc.args), axes)
+transform_broadcasted(fv::CC.Fields.FieldVector, symb, axes) = parent(getproperty(fv, symb))
+transform_broadcasted(x, symb, axes) = x
+@inline function Base.copyto!(dest::CC.Fields.FieldVector, bc::Base.Broadcast.Broadcasted{CC.Fields.FieldVectorStyle})
+    for symb in propertynames(dest)
+        p = parent(getproperty(dest, symb))
+        Base.copyto!(p, transform_broadcasted(bc, symb, axes(p)))
+    end
+    return dest
+end
diff --git a/src/Turbulence_PrognosticTKE.jl b/src/Turbulence_PrognosticTKE.jl
@@ -230,8 +230,36 @@ function compute_diffusive_fluxes(
     return
 end
 
+function affect_filter!(u, edmf, grid, state, gm, Case, TS)
+    # TODO: figure out why this filter kills the DryBubble results if called at t = 0.
+    if Case.casename == "DryBubble" && !(TS.t > 0)
+        return nothing
+    end
+    parent(state.prog) .= parent(u)
+    prog_en = center_prog_environment(state)
+    up = edmf.UpdVar
+    ###
+    ### 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.ρatke[k] = max(prog_en.ρatke[k], 0.0)
+        prog_en.ρaHvar[k] = max(prog_en.ρaHvar[k], 0.0)
+        prog_en.ρaQTvar[k] = max(prog_en.ρaQTvar[k], 0.0)
+        prog_en.ρaHQTcov[k] = max(prog_en.ρaHQTcov[k], -sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k]))
+        prog_en.ρaHQTcov[k] = min(prog_en.ρaHQTcov[k], sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k]))
+    end
+    parent(u) .= parent(state.prog)
+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
+    affect_filter!(Y, edmf, grid, state, gm, Case, TS)
+    parent(state.prog) .= parent(Y)
 
     gm = gm
     up = edmf.UpdVar
@@ -277,53 +305,7 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca
     compute_en_tendencies!(edmf, grid, state, param_set, TS, :QTvar, :ρaQTvar, n_updrafts)
     compute_en_tendencies!(edmf, grid, state, param_set, TS, :HQTcov, :ρaHQTcov, n_updrafts)
 
-    ###
-    ### update (to be removed)
-    ###
-
-    Δt = TS.dt
-    prog_up = center_prog_updrafts(state)
-    prog_up_f = face_prog_updrafts(state)
-    prog_pr = center_prog_precipitation(state)
-    tendencies_pr = center_tendencies_precipitation(state)
-
-    @inbounds for k in real_center_indices(grid)
-        prog_en.ρatke[k] += TS.dt * tendencies_en.ρatke[k]
-        prog_en.ρaHvar[k] += TS.dt * tendencies_en.ρaHvar[k]
-        prog_en.ρaQTvar[k] += TS.dt * tendencies_en.ρaQTvar[k]
-        prog_en.ρaHQTcov[k] += TS.dt * tendencies_en.ρaHQTcov[k]
-        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
-        @inbounds for i in 1:(up.n_updrafts)
-            prog_up[i].ρarea[k] += Δt * tendencies_up[i].ρarea[k]
-            prog_up[i].ρaθ_liq_ice[k] += Δt * tendencies_up[i].ρaθ_liq_ice[k]
-            prog_up[i].ρaq_tot[k] += Δt * tendencies_up[i].ρaq_tot[k]
-        end
-        if has_precip
-            prog_pr.qr[k] += tendencies_pr.qr[k] * TS.dt
-        end
-    end
-
-    @inbounds for k in real_face_indices(grid)
-        @inbounds for i in 1:(up.n_updrafts)
-            prog_up_f[i].ρaw[k] += Δt * tendencies_up_f[i].ρaw[k]
-        end
-    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.ρatke[k] = max(prog_en.ρatke[k], 0.0)
-        prog_en.ρaHvar[k] = max(prog_en.ρaHvar[k], 0.0)
-        prog_en.ρaQTvar[k] = max(prog_en.ρaQTvar[k], 0.0)
-        prog_en.ρaHQTcov[k] = max(prog_en.ρaHQTcov[k], -sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k]))
-        prog_en.ρaHQTcov[k] = min(prog_en.ρaHQTcov[k], sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k]))
-    end
+    parent(dY) .= parent(state.tendencies)
 
     return
 end