Try to fix FieldVector inference

CliMA · Sep 8, 2023 · f006d70 · f006d70
1 parent 4205c01
commit f006d70
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 20 deletions.
diff --git a/src/TurbulenceConvection_deprecated/types.jl b/src/TurbulenceConvection_deprecated/types.jl
@@ -212,21 +212,22 @@ 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
 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,
-    )
+    nt_prog = (; cent = prog_cent_column, face = prog_face_column)
+    nt_aux = (; cent = aux_cent_column, face = aux_face_column)
+    nt_tends = (; cent = tends_cent_column, face = tends_face_column)
+    FT = Spaces.undertype(axes(prog.c))
+    prog_column = Fields.FieldVector{FT}(map(Fields.wrap, nt_prog))
+    aux_column = Fields.FieldVector{FT}(map(Fields.wrap, nt_aux))
+    tends_column = CC.Fields.FieldVector{FT}(map(Fields.wrap, nt_tends))
     surface_conditions = CC.column(p.sfc_conditions, colidx)[]
     return State(
         prog_column,
@@ -243,10 +244,11 @@ function tc_column_state(prog, p, tendencies::Nothing, colidx, t)
     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)
+    nt_prog = (; cent = prog_cent_column, face = prog_face_column)
+    nt_aux = (; cent = aux_cent_column, face = aux_face_column)
+    FT = Spaces.undertype(axes(prog.c))
+    prog_column = Fields.FieldVector{FT}(map(Fields.wrap, nt_prog))
+    aux_column = Fields.FieldVector{FT}(map(Fields.wrap, nt_aux))
     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,15 @@ 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])
+    FT = Spaces.undertype(axes(Y.c))
+    ntY = (; c = Y.c.turbconv[colidx], f = Y.f.turbconv[colidx])
+    ntYt = (; c = Yₜ.c.turbconv[colidx], f = Yₜ.f.turbconv[colidx])
+    ntY_filtered =
+        (; c = Y_filtered.c.turbconv[colidx], f = Y_filtered.f.turbconv[colidx])
+
+    Yₜ_turbconv = Fields.FieldVector{FT}(map(Fields.wrap, ntYt))
+    Y_filtered_turbconv = Fields.FieldVector{FT}(map(Fields.wrap, ntY_filtered))
+    Y_turbconv = Fields.FieldVector{FT}(map(Fields.wrap, ntY))
     Yₜ_turbconv .+= (Y_filtered_turbconv .- Y_turbconv) ./ Δt
     return nothing
 end