Merge #552

552: Change update interface, rename to step r=charleskawczynski a=charleskawczynski This is a prep PR for #473. This PR - renames `update` to `step!` - Changes the `step!` interface to be compatible with OrdinaryDiffEq.jl - Moves the `State` definition into TC.jl `src/`, so that we can pass its properties into a local struct and continue using the dycore API Co-authored-by: Charles Kawczynski <[email protected]>
CliMA · Nov 10, 2021 · bd1bf97 · bd1bf97
2 parents 55edc69 + ccbb4ff
commit bd1bf97
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 19 deletions.
diff --git a/integration_tests/utils/main.jl b/integration_tests/utils/main.jl
@@ -21,12 +21,6 @@ include("parameter_set.jl")
 include("Cases.jl")
 import .Cases
 
-struct State{P, A, T}
-    prog::P
-    aux::A
-    tendencies::T
-end
-
 struct Simulation1d
     io_nt::NamedTuple
     grid
@@ -310,7 +304,7 @@ function Simulation1d(namelist)
     aux = CC.Fields.FieldVector(cent = aux_cent_fields, face = aux_face_fields)
     diagnostics = CC.Fields.FieldVector(cent = diagnostic_cent_fields, face = diagnostic_face_fields)
 
-    state = State(prog, aux, tendencies)
+    state = TC.State(prog, aux, tendencies)
 
     TC.compute_ref_state!(state, grid, param_set; ref_params...)
 
@@ -347,14 +341,19 @@ function run(sim::Simulation1d)
     iter = 0
     grid = sim.grid
     state = sim.state
+    tendencies = state.tendencies
+    prog = state.prog
+    aux = state.aux
+    TS = sim.TS
     diagnostics = sim.diagnostics
     sim.skip_io || 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)
+    params = (; edmf = sim.Turb, grid = grid, gm = sim.GMV, case = sim.Case, TS = TS, aux = aux)
+    while TS.t <= TS.t_max
+        TC.step!(tendencies, prog, params, TS.t)
+        TC.update(TS)
 
         if mod(iter, 100) == 0
-            progress = sim.TS.t / sim.TS.t_max
+            progress = TS.t / TS.t_max
             @show progress
         end
 
@@ -366,14 +365,14 @@ function run(sim::Simulation1d)
             # 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.write_simulation_time(sim.Stats, TS.t) # #removeVarsHack
 
             TC.io(sim.io_nt.diagnostics, sim.Stats)
             TC.io(sim.io_nt.aux, 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
+            TC.io(sim.Turb, grid, state, sim.Stats, TS, sim.param_set) # #removeVarsHack
         end
         iter += 1
     end

diff --git a/perf/common.jl b/perf/common.jl
@@ -13,8 +13,14 @@ TurbulenceConvection.io(sim::Simulation1d) = nothing
 update_n(sim, N::Int) = update_n(sim, Val(N))
 
 function update_n(sim, ::Val{N}) where {N}
+    grid = sim.grid
+    TS = sim.TS
+    prog = sim.state.prog
+    aux = sim.state.aux
+    tendencies = sim.state.tendencies
+    params = (; edmf = sim.Turb, grid = grid, gm = sim.GMV, case = sim.Case, TS = TS, aux = aux)
     for i in 1:N
-        TC.update(sim.Turb, sim.grid, sim.state, sim.GMV, sim.Case, sim.TS)
+        TC.step!(tendencies, prog, params, TS.t)
     end
     return nothing
 end

diff --git a/src/Turbulence_PrognosticTKE.jl b/src/Turbulence_PrognosticTKE.jl
@@ -252,7 +252,10 @@ 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!(tendencies, prog, params, t)
+    UnPack.@unpack edmf, grid, gm, case, aux, TS = params
+
+    state = State(prog, aux, tendencies)
 
     gm = gm
     up = edmf.UpdVar
@@ -267,8 +270,8 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca
     # Some of these methods should probably live in `compute_tendencies`, when written, but we'll
     # treat them as auxiliary variables for now, until we disentangle the tendency computations.
 
-    compute_updraft_surface_bc(edmf, grid, state, Case)
-    update_aux!(edmf, gm, grid, state, Case, param_set, TS)
+    compute_updraft_surface_bc(edmf, grid, state, case)
+    update_aux!(edmf, gm, grid, state, case, param_set, TS)
 
     parent(state.tendencies) .= 0
     tendencies_gm = center_tendencies_grid_mean(state)
@@ -284,7 +287,7 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca
     end
 
     # compute tendencies
-    compute_gm_tendencies!(edmf, grid, state, Case, gm, TS)
+    compute_gm_tendencies!(edmf, grid, state, case, gm, TS)
     compute_updraft_tendencies(edmf, grid, state, gm)
 
     compute_en_tendencies!(edmf, grid, state, param_set, TS, Val(:tke), Val(:ρatke))
@@ -331,7 +334,7 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca
     ###
     ### Filters
     ###
-    set_edmf_surface_bc(edmf, grid, state, up, Case.Sur)
+    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)

diff --git a/src/types.jl b/src/types.jl
@@ -565,3 +565,9 @@ mutable struct EDMF_PrognosticTKE{N_up, A1}
     end
 end
 parameter_set(obj) = obj.param_set
+
+struct State{P, A, T}
+    prog::P
+    aux::A
+    tendencies::T
+end