Make suggested changes to tests and field_name_dict.jl

Author: imreddyTeja

src/MatrixFields/field_name_dict.jl

@@ -182,23 +182,22 @@ function get_internal_entry(
     elseif name_pair[2] == @name() && broadcasted_has_field(T, name_pair[1])
         # multiplication case 2 or 4, second argument
         target_field_eltype = broadcasted_get_field_type(T, name_pair[1])
-        if target_field_eltype <: Number
+        if target_field_eltype == eltype(parent(entry))
             T_band = eltype(entry)
             singleton_datalayout =
                 DataLayouts.singleton(Fields.field_values(entry))
             # BandMatrixRow with same lowest diagonal and bandwidth as `entry`, with a scalar eltype
-            scalar_band_type = BandMatrixRow{
-                T_band.parameters[1],
-                T_band.parameters[2],
-                eltype(parent(entry)),
-            }
+            scalar_band_type = band_matrix_row_type(
+                outer_diagonals(T_band)...,
+                target_field_eltype,
+            )
             field_dim_size = DataLayouts.ncomponents(Fields.field_values(entry))
             scalar_field_offset = get_field_first_index_offset(
                 name_pair[1],
                 target_field_eltype,
                 T,
             )
-            band_element_size = Int(div(sizeof(T), sizeof(target_field_eltype)))
+            band_element_size = div(sizeof(T), sizeof(target_field_eltype))
             parent_indices = DataLayouts.to_data_specific_field(
                 singleton_datalayout,
                 (
@@ -295,7 +294,7 @@ function get_field_first_index_offset(
     child_type = fieldtype(S, child_name)
     remaining_field_chain = drop_first(name)
     field_index =
-        unrolled_filter(i -> fieldname(S, i) == child_name, 1:fieldcount(S))[1]
+        UnrolledUtilities.unrolled_findfirst(isequal(child_name), fieldnames(S))
     return DataLayouts.fieldtypeoffset(T, S, field_index) +
            get_field_first_index_offset(remaining_field_chain, T, child_type)
 end
@@ -314,14 +313,18 @@ entries in the `FieldMatrix` `dict`.
 """
 function get_scalar_keys(dict::FieldMatrix)
     keys_tuple = unrolled_flatmap(keys(dict).values) do key
-        _, entry = unrolled_filter(pair -> key == pair[1], pairs(dict))[1]
+        entry = values(dict)[UnrolledUtilities.unrolled_findfirst(
+            isequal(key),
+            keys(dict).values,
+        )]
         entry =
             entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry
         unrolled_map(
-            filtered_child_names(
-                field -> eltype(field) <: Number,
+            filtered_names(
+                field ->
+                    (field isa UniformScaling) ||
+                        eltype(field) == eltype(parent(field)),
                 entry,
-                @name()
             ),
         ) do name
             (append_internal_name(key[1], name), key[2])
@@ -339,24 +342,27 @@ scalar components of entries of `field_matrix`.
 
 # Example usage
 ```julia
-struct foo{T1, T2}
-    a::T
-    b::T2
-end
-mat1 = fill(DiagonalMatrixRow(ClimaCore.Geometry.Covariant12Vector(1.0, 2.0)), space)
-mat2 = fill(DiagonalMatrixRow(foo(foo(1.0, 2.0), 3.0)), space)
+e¹² = Geometry.Covariant12Vector(1.6, 0.7)
+e₃ = Geometry.Contravariant3Vector(1.0)
+ᶜᶜmat3 = fill(TridiagonalMatrixRow(2.0, 3.2, 2.1), center_space)
+ᶜᶠmat2 = fill(BidiagonalMatrixRow(4.3, 1.7), center_space)
+ᶜᶜmat3_uₕ_scalar = ᶜᶜmat3 .* (e¹²,)
+ρχ_unit = (;ρq_liq = 1.0, ρq_ice = 1.0)
+ᶜᶠmat2_ρχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit ⊠ e₃')), ᶜᶠmat2)
+
+
 A = MatrixFields.FieldMatrix(
-    (@name(biz), @name(baz)) => mat1,
-    (@name(bip), @name(bop)) => mat2,
+    (@name(c.ρχ), @name(f.u₃)) => ᶜᶠmat2_ρχ_u₃,
+    (@name(c.uₕ), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3_uₕ_scalar,
 )
+
 A_scalar = MatrixFields.scalar_fieldmatrix(A)
 keys(A_scalar)
 # Output:
-# (@name(biz.components.data.:(1)), @name(baz))
-# (@name(biz.components.data.:(2)), @name(baz))
-# (@name(bip.a.a), @name(bop))
-# (@name(bip.a.b), @name(bop))
-# (@name(bip.b), @name(bop))
+# (@name(c.ρχ.ρq_liq.parent.components.data.:(1)), @name(f.u₃))
+# (@name(c.ρχ.ρq_ice.parent.components.data.:(1)), @name(f.u₃))
+# (@name(c.uₕ.components.data.:(1)), @name(c.sgsʲs.:(1).ρa))
+# (@name(c.uₕ.components.data.:(2)), @name(c.sgsʲs.:(1).ρa))
 ```
 """
 function scalar_fieldmatrix(field_matrix::FieldMatrix)

test/MatrixFields/field_matrix_solvers.jl

@@ -376,15 +376,10 @@ end
     center_gs_unit = (; dry_center_gs_unit..., ρatke = 1, ρχ = ρχ_unit)
     center_sgsʲ_unit = (; ρa = 1, ρae_tot = 1, ρaχ = ρaχ_unit)
 
-    ᶜᶜmat3_uₕ_scalar = ᶜᶜmat3 .* (e¹²,)
     ᶠᶜmat2_u₃_scalar = ᶠᶜmat2 .* (e³,)
     ᶜᶠmat2_scalar_u₃ = ᶜᶠmat2 .* (e₃',)
-    ᶜᶠmat2_uₕ_u₃ = ᶜᶠmat2 .* (e¹² * e₃',)
     ᶠᶠmat3_u₃_u₃ = ᶠᶠmat3 .* (e³ * e₃',)
-    ᶜᶜmat3_ρχ_scalar = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit)), ᶜᶜmat3)
-    ᶜᶜmat3_ρaχ_scalar = map(Base.Fix1(map, Base.Fix2(⊠, ρaχ_unit)), ᶜᶜmat3)
     ᶜᶠmat2_ρχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit ⊠ e₃')), ᶜᶠmat2)
-    ᶜᶠmat2_ρaχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρaχ_unit ⊠ e₃')), ᶜᶠmat2)
     # We need to use Fix1 and Fix2 instead of defining anonymous functions in
     # order for the result of map to be inferrable.
 
@@ -478,52 +473,21 @@ end
                 n_iters = 6,
             ),
         ),
-        A = MatrixFields.FieldMatrix(
-            # GS-GS blocks:
-            (@name(sfc), @name(sfc)) => I,
-            (@name(c.ρ), @name(c.ρ)) => I,
-            (@name(c.ρe_tot), @name(c.ρe_tot)) => ᶜᶜmat3,
-            (@name(c.ρatke), @name(c.ρatke)) => ᶜᶜmat3,
-            (@name(c.ρχ), @name(c.ρχ)) => ᶜᶜmat3,
-            (@name(c.uₕ), @name(c.uₕ)) => ᶜᶜmat3,
-            (@name(c.ρ), @name(f.u₃)) => ᶜᶠmat2_scalar_u₃,
-            (@name(c.ρe_tot), @name(f.u₃)) => ᶜᶠmat2_scalar_u₃,
-            (@name(c.ρatke), @name(f.u₃)) => ᶜᶠmat2_scalar_u₃,
-            (@name(c.ρχ), @name(f.u₃)) => ᶜᶠmat2_ρχ_u₃,
-            (@name(f.u₃), @name(c.ρ)) => ᶠᶜmat2_u₃_scalar,
-            (@name(f.u₃), @name(c.ρe_tot)) => ᶠᶜmat2_u₃_scalar,
-            (@name(f.u₃), @name(f.u₃)) => ᶠᶠmat3_u₃_u₃,
-            # GS-SGS blocks:
-            (@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρae_tot)) => ᶜᶜmat3,
-            (@name(c.ρχ.ρq_tot), @name(c.sgsʲs.:(1).ρaχ.ρaq_tot)) => ᶜᶜmat3,
-            (@name(c.ρχ.ρq_liq), @name(c.sgsʲs.:(1).ρaχ.ρaq_liq)) => ᶜᶜmat3,
-            (@name(c.ρχ.ρq_ice), @name(c.sgsʲs.:(1).ρaχ.ρaq_ice)) => ᶜᶜmat3,
-            (@name(c.ρχ.ρq_rai), @name(c.sgsʲs.:(1).ρaχ.ρaq_rai)) => ᶜᶜmat3,
-            (@name(c.ρχ.ρq_sno), @name(c.sgsʲs.:(1).ρaχ.ρaq_sno)) => ᶜᶜmat3,
-            (@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3,
-            (@name(c.ρatke), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3,
-            (@name(c.ρχ), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3_ρχ_scalar,
-            (@name(c.uₕ), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3_uₕ_scalar,
-            (@name(c.ρe_tot), @name(f.sgsʲs.:(1).u₃)) => ᶜᶠmat2_scalar_u₃,
-            (@name(c.ρatke), @name(f.sgsʲs.:(1).u₃)) => ᶜᶠmat2_scalar_u₃,
-            (@name(c.ρχ), @name(f.sgsʲs.:(1).u₃)) => ᶜᶠmat2_ρχ_u₃,
-            (@name(c.uₕ), @name(f.sgsʲs.:(1).u₃)) => ᶜᶠmat2_uₕ_u₃,
-            (@name(f.u₃), @name(c.sgsʲs.:(1).ρa)) => ᶠᶜmat2_u₃_scalar,
-            (@name(f.u₃), @name(f.sgsʲs.:(1).u₃)) => ᶠᶠmat3_u₃_u₃,
-            # SGS-SGS blocks:
-            (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).ρa)) => I,
-            (@name(c.sgsʲs.:(1).ρae_tot), @name(c.sgsʲs.:(1).ρae_tot)) => I,
-            (@name(c.sgsʲs.:(1).ρaχ), @name(c.sgsʲs.:(1).ρaχ)) => I,
-            (@name(c.sgsʲs.:(1).ρa), @name(f.sgsʲs.:(1).u₃)) =>
-                ᶜᶠmat2_scalar_u₃,
-            (@name(c.sgsʲs.:(1).ρae_tot), @name(f.sgsʲs.:(1).u₃)) =>
-                ᶜᶠmat2_scalar_u₃,
-            (@name(c.sgsʲs.:(1).ρaχ), @name(f.sgsʲs.:(1).u₃)) => ᶜᶠmat2_ρaχ_u₃,
-            (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).ρa)) =>
-                ᶠᶜmat2_u₃_scalar,
-            (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).ρae_tot)) =>
-                ᶠᶜmat2_u₃_scalar,
-            (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => ᶠᶠmat3_u₃_u₃,
+        A = dycore_prognostic_EDMF_FieldMatrix(;
+            ᶜᶜmat1,
+            ᶜᶠmat2,
+            ᶠᶜmat2,
+            ᶜᶜmat3,
+            ᶠᶠmat3,
+            e¹²,
+            e³,
+            e₃,
+            ρχ_unit,
+            ρaχ_unit,
+            ᶜᶠmat2_ρχ_u₃,
+            ᶠᶠmat3_u₃_u₃,
+            ᶜᶠmat2_scalar_u₃,
+            ᶠᶜmat2_u₃_scalar,
         ),
         b = b_moist_dycore_prognostic_edmf_prognostic_surface,
     )

test/MatrixFields/matrix_field_test_utils.jl

@@ -21,6 +21,10 @@ import ClimaCore:
     Operators,
     Quadratures
 using ClimaCore.MatrixFields
+import ClimaCore.Utilities: half
+import ClimaCore.RecursiveApply: ⊠
+import LinearAlgebra: I, norm, ldiv!, mul!
+import ClimaCore.MatrixFields: @name
 
 # Test that an expression is true and that it is also type-stable.
 macro test_all(expression)
@@ -32,7 +36,7 @@ macro test_all(expression)
     end
 end
 
-# Compute the minimum time (in seconds) required to run an expression after it 
+# Compute the minimum time (in seconds) required to run an expression after it
 # has been compiled. This macro is used instead of @benchmark from
 # BenchmarkTools.jl because the latter is extremely slow (it appears to keep
 # triggering recompilations and allocating a lot of memory in the process).
@@ -134,6 +138,85 @@ function test_field_broadcast(;
     end
 end
 
+# Create a field matrix for a similar solve to ClimaAtmos's moist dycore + prognostic,
+# EDMF + prognostic surface temperature with implicit acoustic waves and SGS fluxes
+function dycore_prognostic_EDMF_FieldMatrix(;
+    ᶜᶜmat1,
+    ᶜᶠmat2,
+    ᶠᶜmat2,
+    ᶜᶜmat3,
+    ᶠᶠmat3,
+    e¹²,
+    e³,
+    e₃,
+    ρχ_unit,
+    ρaχ_unit,
+    ᶜᶠmat2_ρχ_u₃,
+    ᶠᶠmat3_u₃_u₃,
+    ᶜᶠmat2_scalar_u₃,
+    ᶠᶜmat2_u₃_scalar,
+)
+
+    ᶜᶜmat3_uₕ_scalar = ᶜᶜmat3 .* (e¹²,)
+    ᶜᶠmat2_uₕ_u₃ = ᶜᶠmat2 .* (e¹² * e₃',)
+    ᶜᶜmat3_ρχ_scalar = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit)), ᶜᶜmat3)
+    ᶜᶜmat3_ρaχ_scalar = map(Base.Fix1(map, Base.Fix2(⊠, ρaχ_unit)), ᶜᶜmat3)
+    ᶜᶠmat2_ρaχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρaχ_unit ⊠ e₃')), ᶜᶠmat2)
+    return MatrixFields.FieldMatrix(
+        # GS-GS blocks:
+        (@name(sfc), @name(sfc)) => I,
+        (@name(c.ρ), @name(c.ρ)) => I,
+        (@name(c.ρe_tot), @name(c.ρe_tot)) => deepcopy(ᶜᶜmat3),
+        (@name(c.ρatke), @name(c.ρatke)) => deepcopy(ᶜᶜmat3),
+        (@name(c.ρχ), @name(c.ρχ)) => deepcopy(ᶜᶜmat3),
+        (@name(c.uₕ), @name(c.uₕ)) => deepcopy(ᶜᶜmat3),
+        (@name(c.ρ), @name(f.u₃)) => deepcopy(ᶜᶠmat2_scalar_u₃),
+        (@name(c.ρe_tot), @name(f.u₃)) => deepcopy(ᶜᶠmat2_scalar_u₃),
+        (@name(c.ρatke), @name(f.u₃)) => deepcopy(ᶜᶠmat2_scalar_u₃),
+        (@name(c.ρχ), @name(f.u₃)) => deepcopy(ᶜᶠmat2_ρχ_u₃),
+        (@name(f.u₃), @name(c.ρ)) => deepcopy(ᶠᶜmat2_u₃_scalar),
+        (@name(f.u₃), @name(c.ρe_tot)) => deepcopy(ᶠᶜmat2_u₃_scalar),
+        (@name(f.u₃), @name(f.u₃)) => deepcopy(ᶠᶠmat3_u₃_u₃),
+        # GS-SGS blocks:
+        (@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρae_tot)) => deepcopy(ᶜᶜmat3),
+        (@name(c.ρχ.ρq_tot), @name(c.sgsʲs.:(1).ρaχ.ρaq_tot)) =>
+            deepcopy(ᶜᶜmat3),
+        (@name(c.ρχ.ρq_liq), @name(c.sgsʲs.:(1).ρaχ.ρaq_liq)) =>
+            deepcopy(ᶜᶜmat3),
+        (@name(c.ρχ.ρq_ice), @name(c.sgsʲs.:(1).ρaχ.ρaq_ice)) =>
+            deepcopy(ᶜᶜmat3),
+        (@name(c.ρχ.ρq_rai), @name(c.sgsʲs.:(1).ρaχ.ρaq_rai)) =>
+            deepcopy(ᶜᶜmat3),
+        (@name(c.ρχ.ρq_sno), @name(c.sgsʲs.:(1).ρaχ.ρaq_sno)) =>
+            deepcopy(ᶜᶜmat3),
+        (@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρa)) => deepcopy(ᶜᶜmat3),
+        (@name(c.ρatke), @name(c.sgsʲs.:(1).ρa)) => deepcopy(ᶜᶜmat3),
+        (@name(c.ρχ), @name(c.sgsʲs.:(1).ρa)) => deepcopy(ᶜᶜmat3_ρχ_scalar),
+        (@name(c.uₕ), @name(c.sgsʲs.:(1).ρa)) => deepcopy(ᶜᶜmat3_uₕ_scalar),
+        (@name(c.ρe_tot), @name(f.sgsʲs.:(1).u₃)) => deepcopy(ᶜᶠmat2_scalar_u₃),
+        (@name(c.ρatke), @name(f.sgsʲs.:(1).u₃)) => deepcopy(ᶜᶠmat2_scalar_u₃),
+        (@name(c.ρχ), @name(f.sgsʲs.:(1).u₃)) => deepcopy(ᶜᶠmat2_ρχ_u₃),
+        (@name(c.uₕ), @name(f.sgsʲs.:(1).u₃)) => deepcopy(ᶜᶠmat2_uₕ_u₃),
+        (@name(f.u₃), @name(c.sgsʲs.:(1).ρa)) => deepcopy(ᶠᶜmat2_u₃_scalar),
+        (@name(f.u₃), @name(f.sgsʲs.:(1).u₃)) => deepcopy(ᶠᶠmat3_u₃_u₃),
+        # SGS-SGS blocks:
+        (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).ρa)) => I,
+        (@name(c.sgsʲs.:(1).ρae_tot), @name(c.sgsʲs.:(1).ρae_tot)) => I,
+        (@name(c.sgsʲs.:(1).ρaχ), @name(c.sgsʲs.:(1).ρaχ)) => I,
+        (@name(c.sgsʲs.:(1).ρa), @name(f.sgsʲs.:(1).u₃)) =>
+            deepcopy(ᶜᶠmat2_scalar_u₃),
+        (@name(c.sgsʲs.:(1).ρae_tot), @name(f.sgsʲs.:(1).u₃)) =>
+            deepcopy(ᶜᶠmat2_scalar_u₃),
+        (@name(c.sgsʲs.:(1).ρaχ), @name(f.sgsʲs.:(1).u₃)) =>
+            deepcopy(ᶜᶠmat2_ρaχ_u₃),
+        (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).ρa)) =>
+            deepcopy(ᶠᶜmat2_u₃_scalar),
+        (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).ρae_tot)) =>
+            deepcopy(ᶠᶜmat2_u₃_scalar),
+        (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) =>
+            deepcopy(ᶠᶠmat3_u₃_u₃),
+    )
+end
 # Generate extruded finite difference spaces for testing. Include topography
 # when possible.
 function test_spaces(::Type{FT}) where {FT}