From 6eeb21c7241bed5b5f5881e42c8eb85f3c7c1627 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Tue, 9 Nov 2021 16:53:54 -0800
Subject: [PATCH] Change update interface, rename to step

---
 integration_tests/utils/main.jl | 25 ++++++++++++-------------
 src/Turbulence_PrognosticTKE.jl | 13 ++++++++-----
 src/types.jl                    |  6 ++++++
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/integration_tests/utils/main.jl b/integration_tests/utils/main.jl
index 39ace1185e..fecf45fa6e 100644
--- 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/src/Turbulence_PrognosticTKE.jl b/src/Turbulence_PrognosticTKE.jl
index e56d7ae171..eafc16d9d0 100755
--- 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
index 5370085011..1758836a63 100644
--- 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