Fix type instabilities

Author: charleskawczynski

src/utils/variable_manipulations.jl

@@ -159,28 +159,14 @@ Converts every variable of the form `ρχ` in the grid-scale state `gs` into the
 specific variable `χ` by dividing it by `ρ`. All other variables in `gs` are
 omitted from the result.
 """
-specific_gs(gs) =
-    NamedTuple{specific_gs_names(gs)}(map(name -> gs.:($name) / gs.ρ, relevant_gs_names(gs)))
-
-"""
-    relevant_gs_names(gs)
-
-Returns relevant grid-scale state `gs` names for determining specific variables.
-"""
-@generated relevant_gs_names(::Type{GS}) where {GS} =
-    filter(name -> has_prefix(name, :ρ) && name != :ρ, Base._nt_names(GS))
-
-@inline relevant_gs_names(gs) = relevant_gs_names(typeof(gs))
-
-"""
-    specific_gs_names(gs)
-
-Returns relevant specific grid-scale state `gs` names.
-"""
-@generated specific_gs_names(::Type{GS}) where {GS} =
-    map(name -> remove_prefix(name, :ρ), relevant_gs_names(GS))
-
-@inline specific_gs_names(gs) = specific_gs_names(typeof(gs))
+@generated function specific_gs(gs)
+    gs_names = Base._nt_names(gs)
+    relevant_gs_names =
+        filter(name -> has_prefix(name, :ρ) && name != :ρ, gs_names)
+    specific_gs_names = map(name -> remove_prefix(name, :ρ), relevant_gs_names)
+    specific_gs_values = map(name -> :(gs.$name / gs.ρ), relevant_gs_names)
+    return :(NamedTuple{$specific_gs_names}(($(specific_gs_values...),)))
+end
 
 """
     all_specific_gs(gs)
@@ -199,12 +185,14 @@ Arguments:
 Returns:
 - A new `NamedTuple` containing only the specific quantities (e.g., `:q_tot`, `:e_tot`).
 """
-all_specific_gs(gs) =
-    NamedTuple{specific_gs_names(eltype(gs))}(
-        UU.unrolled_map(relevant_gs_names(eltype(gs))) do name
-            lazy.(specific.(getproperty(gs, name), gs.ρ))
-        end
-    )
+@generated function all_specific_gs(gs)
+    gs_names = Base._nt_names(gs)
+    relevant_gs_names =
+        filter(name -> has_prefix(name, :ρ) && name != :ρ, gs_names)
+    specific_gs_names = map(name -> remove_prefix(name, :ρ), relevant_gs_names)
+    specific_gs_values = map(name -> :(lazy.(specific.(gs.$name, gs.ρ))), relevant_gs_names)
+    return :(NamedTuple{$specific_gs_names}(($(specific_gs_values...),)))
+end
 
 """
     specific_sgs(sgs, gs, turbconv_model)

test/utilities.jl

@@ -5,7 +5,6 @@ import Dates
 using Random
 Random.seed!(1234)
 import ClimaAtmos as CA
-import LazyBroadcast: lazy
 using NCDatasets
 
 include("test_helpers.jl")
@@ -34,21 +33,6 @@ end
     @test CA.isdivisible(Dates.Month(1), Dates.Hour(1))
 end
 
-@testset "Variable manipulations" begin
-    (; helem, cent_space, face_space) = get_cartesian_spaces()
-    ccoords, fcoords = get_coords(cent_space, face_space)
-    FT = eltype(ccoords.x)
-    nt = (;ρ=FT(2),ρa=FT(4), ρe_tot=FT(4))
-    @test CA.relevant_gs_names(typeof(nt)) == (:ρa, :ρe_tot)
-    @test CA.specific_gs_names(nt) == (:a, :e_tot)
-    @test CA.specific_gs(nt) == (a = 2.0, e_tot = 2.0)
-    f = Fields.Field(typeof(nt), cent_space)
-    @test CA.all_specific_gs(f) == (;
-        a = lazy.(CA.specific.(f.ρa, f.ρ)),
-        e_tot = lazy.(CA.specific.(f.ρe_tot, f.ρ)),
-    )
-end
-
 @testset "kinetic_energy (c.f. analytical function)" begin
     # Test kinetic energy function for staggered grids
     # given an analytical expression for the velocity profiles