Merge #1297

1297: Try improving inference r=charleskawczynski a=charleskawczynski This PR is a step towards improving inference Co-authored-by: Charles Kawczynski <[email protected]>
CliMA · Sep 19, 2022 · 13affe0 · 13affe0 · ilopezgp · Sep 19, 2022
2 parents 920f81f + 345b787
commit 13affe0
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 32 deletions.
diff --git a/driver/callbacks.jl b/driver/callbacks.jl
@@ -12,13 +12,12 @@ end
 condition_every_iter(u, t, integrator) = true
 
 function affect_io!(integrator)
-    UnPack.@unpack edmf, calibrate_io, precip_model, aux, io_nt, diagnostics, surf_params, param_set, Stats, skip_io =
-        integrator.p
-    skip_io && return nothing
-    t = integrator.t
-    prog = integrator.u
+    integrator.p.skip_io && return nothing
 
-    CC.Fields.bycolumn(axes(prog.cent)) do colidx
+    CC.Fields.bycolumn(axes(integrator.u.cent)) do colidx
+        (; edmf, calibrate_io, precip_model, aux, io_nt, diagnostics, surf_params, param_set, Stats) = integrator.p
+        t = integrator.t
+        prog = integrator.u
         stats = Stats[colidx]
         # TODO: remove `vars` hack that avoids
         # https://github.com/Alexander-Barth/NCDatasets.jl/issues/135
@@ -68,22 +67,23 @@ function affect_io!(integrator)
 
         surf = get_surface(surf_params, grid, state, t, param_set)
         io(surf, surf_params, grid, state, stats, t)
+        nothing
     end
 
     ODE.u_modified!(integrator, false) # We're legitamately not mutating `u` (the state vector)
 end
 
 function affect_filter!(integrator)
-    UnPack.@unpack edmf, param_set, aux, case, surf_params = integrator.p
-    t = integrator.t
-    prog = integrator.u
-    prog = integrator.u
+    CC.Fields.bycolumn(axes(integrator.u.cent)) do colidx
+        (; edmf, param_set, aux, case, surf_params) = integrator.p
+        t = integrator.t
+        prog = integrator.u
 
-    CC.Fields.bycolumn(axes(prog.cent)) do colidx
         state = TC.column_prog_aux(prog, aux, colidx)
         grid = TC.Grid(state)
         surf = get_surface(surf_params, grid, state, t, param_set)
         TC.affect_filter!(edmf, grid, state, param_set, surf, t)
+        nothing
     end
 
     # We're lying to OrdinaryDiffEq.jl, in order to avoid
@@ -139,27 +139,26 @@ function compute_dt_max(state::TC.State, edmf::TC.EDMFModel, dt_max::FT, CFL_lim
     return dt_max
 end
 
-function dt_max!(integrator)
-    UnPack.@unpack edmf, aux, TS = integrator.p
-    prog = integrator.u
-
-    dt_max = TS.dt_max # initialize dt_max
-    CC.Fields.bycolumn(axes(prog.cent)) do colidx
+dt_max!(integrator) = dt_max!(integrator, integrator.p.TS.dt_max)
+function dt_max!(integrator, dt_max)
+    CC.Fields.bycolumn(axes(integrator.u.cent)) do colidx
+        (; edmf, aux, TS) = integrator.p
+        prog = integrator.u
         state = TC.column_prog_aux(prog, aux, colidx)
         dt_max = compute_dt_max(state, edmf, dt_max, TS.cfl_limit)
     end
     to_float(f) = f isa ForwardDiff.Dual ? ForwardDiff.value(f) : f
+    (; TS) = integrator.p
     TS.dt_max_edmf = to_float(dt_max)
 
     ODE.u_modified!(integrator, false)
 end
 
 function monitor_cfl!(integrator)
-    UnPack.@unpack edmf, aux, TS = integrator.p
-    prog = integrator.u
-    frac_tol = 0.05
-
-    CC.Fields.bycolumn(axes(prog.cent)) do colidx
+    CC.Fields.bycolumn(axes(integrator.u.cent)) do colidx
+        (; edmf, aux, TS) = integrator.p
+        frac_tol = 0.05
+        prog = integrator.u
         state = TC.column_prog_aux(prog, aux, colidx)
         grid = TC.Grid(state)
         prog_gm = TC.center_prog_grid_mean(state)
@@ -193,6 +192,7 @@ function monitor_cfl!(integrator)
                 )
             end
         end
+        nothing
     end
 
     ODE.u_modified!(integrator, false)

diff --git a/driver/dycore.jl b/driver/dycore.jl
@@ -210,16 +210,14 @@ end
 
 # Compute the sum of tendencies for the scheme
 function ∑tendencies!(tendencies::FV, prog::FV, params::NT, t::Real) where {NT, FV <: CC.Fields.FieldVector}
-    UnPack.@unpack edmf, precip_model, param_set, case = params
-    UnPack.@unpack surf_params, radiation, forcing, aux, TS = params
+    CC.Fields.bycolumn(axes(prog.cent)) do colidx
+        UnPack.@unpack edmf, precip_model, param_set, case = params
+        UnPack.@unpack surf_params, radiation, forcing, aux, TS = params
 
-    thermo_params = TCP.thermodynamics_params(param_set)
-    tends_face = tendencies.face
-    tends_cent = tendencies.cent
-    parent(tends_face) .= 0
-    parent(tends_cent) .= 0
+        thermo_params = TCP.thermodynamics_params(param_set)
 
-    CC.Fields.bycolumn(axes(prog.cent)) do colidx
+        parent(tendencies.face[colidx]) .= 0
+        parent(tendencies.cent[colidx]) .= 0
         state = TC.column_state(prog, aux, tendencies, colidx)
         grid = TC.Grid(state)
 
@@ -250,6 +248,7 @@ function ∑tendencies!(tendencies::FV, prog::FV, params::NT, t::Real) where {NT
 
         TC.compute_turbconv_tendencies!(edmf, grid, state, param_set, surf, Δt)
         compute_gm_tendencies!(edmf, grid, state, surf, radiation, forcing, param_set)
+        nothing
     end
 
     return nothing

diff --git a/src/types.jl b/src/types.jl
@@ -771,10 +771,16 @@ function Base.summary(io::IO, edmf::EDMFModel)
 end
 
 
-struct State{P, A, T}
+struct State{P, A, T, G}
     prog::P
     aux::A
     tendencies::T
+    grid::G
+end
+
+function State(prog::P, aux::A, tendencies::T) where {P, A, T}
+    grid = Grid(axes(prog.cent))
+    return State{P, A, T, typeof(grid)}(prog, aux, tendencies, grid)
 end
 
 """
@@ -826,7 +832,7 @@ function column_diagnostics(diagnostics, colidx)
 end
 
 
-Grid(state::State) = Grid(axes(state.prog.cent))
+Grid(state::State) = state.grid
 
 float_type(state::State) = eltype(state.prog)
 # float_type(field::CC.Fields.Field) = CC.Spaces.undertype(axes(field))