Try to fix FieldVector inference

Skip calling Fields.wrap
CliMA · Sep 8, 2023 · b5543c2 · b5543c2
1 parent 4205c01
commit b5543c2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 30 deletions.
diff --git a/src/TurbulenceConvection_deprecated/types.jl b/src/TurbulenceConvection_deprecated/types.jl
@@ -212,21 +212,23 @@ float_type(state::State) = eltype(state.prog)
 # float_type(field::CC.Fields.Field) = CC.Spaces.undertype(axes(field))
 float_type(field::CC.Fields.Field) = eltype(parent(field))
 
+import ClimaCore.Fields as Fields
+import ClimaCore.Spaces as Spaces
+
+
+Base.@propagate_inbounds function field_vector_column(
+    fv::FieldVector{T},
+    colidx::Fields.ColumnIndex,
+) where {T}
+    values = map(x -> x[colidx], Fields._values(fv))
+    return Fields.FieldVector{T, typeof(values)}(values)
+end
+
 function tc_column_state(prog, p, tendencies, colidx, t)
-    prog_cent_column = CC.column(prog.c, colidx)
-    prog_face_column = CC.column(prog.f, colidx)
-    aux_cent_column = CC.column(p.edmf_cache.aux.cent, colidx)
-    aux_face_column = CC.column(p.edmf_cache.aux.face, colidx)
-    tends_cent_column = CC.column(tendencies.c, colidx)
-    tends_face_column = CC.column(tendencies.f, colidx)
-    prog_column =
-        CC.Fields.FieldVector(cent = prog_cent_column, face = prog_face_column)
-    aux_column =
-        CC.Fields.FieldVector(cent = aux_cent_column, face = aux_face_column)
-    tends_column = CC.Fields.FieldVector(
-        cent = tends_cent_column,
-        face = tends_face_column,
-    )
+    FT = Spaces.undertype(axes(prog.c))
+    prog_column = field_vector_column(prog, colidx)
+    aux_column = field_vector_column(p.edmf_cache.aux, colidx)
+    tends_column = field_vector_column(tendencies, colidx)
     surface_conditions = CC.column(p.sfc_conditions, colidx)[]
     return State(
         prog_column,
@@ -239,14 +241,8 @@ function tc_column_state(prog, p, tendencies, colidx, t)
 end
 
 function tc_column_state(prog, p, tendencies::Nothing, colidx, t)
-    prog_cent_column = CC.column(prog.c, colidx)
-    prog_face_column = CC.column(prog.f, colidx)
-    aux_cent_column = CC.column(p.edmf_cache.aux.cent, colidx)
-    aux_face_column = CC.column(p.edmf_cache.aux.face, colidx)
-    prog_column =
-        CC.Fields.FieldVector(cent = prog_cent_column, face = prog_face_column)
-    aux_column =
-        CC.Fields.FieldVector(cent = aux_cent_column, face = aux_face_column)
+    prog_column = field_vector_column(prog, colidx)
+    aux_column = field_vector_column(p.edmf_cache.aux, colidx)
     tends_column = nothing
     surface_conditions = CC.column(p.sfc_conditions, colidx)[]
     return State(

diff --git a/src/dycore_equations_deprecated/sgs_flux_tendencies.jl b/src/dycore_equations_deprecated/sgs_flux_tendencies.jl
@@ -237,14 +237,9 @@ function explicit_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, ::TC.EDMFModel)
     # Note: This "filter relaxation tendency" can be scaled down if needed, but
     # it must be present in order to prevent Y and Y_filtered from diverging
     # during each timestep.
-    Yₜ_turbconv =
-        Fields.FieldVector(c = Yₜ.c.turbconv[colidx], f = Yₜ.f.turbconv[colidx])
-    Y_filtered_turbconv = Fields.FieldVector(
-        c = Y_filtered.c.turbconv[colidx],
-        f = Y_filtered.f.turbconv[colidx],
-    )
-    Y_turbconv =
-        Fields.FieldVector(c = Y.c.turbconv[colidx], f = Y.f.turbconv[colidx])
+    Yₜ_turbconv = field_vector_column(Yₜ, colidx)
+    Y_filtered_turbconv = field_vector_column(Y_filtered, colidx)
+    Y_turbconv = field_vector_column(Y, colidx)
     Yₜ_turbconv .+= (Y_filtered_turbconv .- Y_turbconv) ./ Δt
     return nothing
 end