Add scalar_fieldmatrix

Add a function to convert a FieldMatrix where each matrix entry
has an eltype of some struct into a FieldMatrix where each entry
has an eltype of a scalar.

Add additional tests for scalar_matrixfields

Use @test_all in tests

Author: imreddyTeja

docs/src/matrix_fields.md

@@ -89,6 +89,9 @@ preconditioner_cache
 check_preconditioner
 lazy_or_concrete_preconditioner
 apply_preconditioner
+scalar_keys
+get_field_first_index_offset
+broadcasted_get_field_type
 ```
 
 ## Utilities
@@ -98,4 +101,5 @@ column_field2array
 column_field2array_view
 field2arrays
 field2arrays_view
+scalar_fieldmatrix
 ```

src/MatrixFields/MatrixFields.jl

@@ -58,6 +58,7 @@ import ..Utilities: PlusHalf, half
 import ..RecursiveApply:
     rmap, rmaptype, rpromote_type, rzero, rconvert, radd, rsub, rmul, rdiv
 import ..RecursiveApply: ⊠, ⊞, ⊟
+import ..DataLayouts
 import ..DataLayouts: AbstractData
 import ..DataLayouts: vindex
 import ..Geometry

src/MatrixFields/field_name.jl

@@ -59,6 +59,24 @@ get_field(x, ::FieldName{()}) = x
 get_field(x, name::FieldName) =
     get_field(getproperty(x, extract_first(name)), drop_first(name))
 
+"""
+    broadcasted_get_field_type(::Type{X}, name::FieldName)
+
+Returns the type of the field accessed by `name` in the type `X`.
+"""
+broadcasted_get_field_type(::Type{X}, ::FieldName{()}) where {X} = X
+broadcasted_get_field_type(::Type{X}, name::FieldName) where {X} =
+    broadcasted_get_field_type(
+        fieldtype(X, extract_first(name)),
+        drop_first(name),
+    )
+if hasfield(Method, :recursion_relation)
+    dont_limit = (args...) -> true
+    for m in methods(broadcasted_get_field_type)
+        m.recursion_relation = dont_limit
+    end
+end
+
 broadcasted_has_field(::Type{X}, ::FieldName{()}) where {X} = true
 broadcasted_has_field(::Type{X}, name::FieldName) where {X} =
     extract_first(name) in fieldnames(X) &&

src/MatrixFields/field_name_dict.jl

@@ -181,11 +181,51 @@ function get_internal_entry(
         entry
     elseif name_pair[2] == @name() && broadcasted_has_field(T, name_pair[1])
         # multiplication case 2 or 4, second argument
-        Base.broadcasted(entry) do matrix_row
-            map(matrix_row) do matrix_row_entry
-                broadcasted_get_field(matrix_row_entry, name_pair[1])
-            end
-        end # Note: This assumes that the entry is in a FieldMatrixBroadcasted.
+        target_field_eltype = broadcasted_get_field_type(T, name_pair[1])
+        if target_field_eltype <: Number
+            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)),
+            }
+            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)))
+            parent_indices = DataLayouts.to_data_specific_field(
+                singleton_datalayout,
+                (
+                    :,
+                    :,
+                    (1 + scalar_field_offset):band_element_size:field_dim_size,
+                    :,
+                    :,
+                ),
+            )
+            scalar_data = view(parent(entry), parent_indices...)
+            values = DataLayouts.union_all(singleton_datalayout){
+                scalar_band_type,
+                Base.tail(
+                    DataLayouts.type_params(Fields.field_values(entry)),
+                )...,
+            }(
+                scalar_data,
+            )
+            Fields.Field(values, axes(entry))
+        else
+            Base.broadcasted(entry) do matrix_row
+                map(matrix_row) do matrix_row_entry
+                    broadcasted_get_field(matrix_row_entry, name_pair[1])
+                end
+            end # Note: This assumes that the entry is in a FieldMatrixBroadcasted.
+        end
     else
         throw(key_error)
     end
@@ -237,6 +277,96 @@ function Base.one(matrix::FieldMatrix)
     return FieldNameDict(inferred_diagonal_keys, entries)
 end
 
+"""
+    get_field_first_index_offset(name::FieldName, ::Type{T}, ::Type{S})
+
+Returns the offset of the the field with name `name` in an object of type `S`
+in multiples of `sizeof(T)`.
+"""
+function get_field_first_index_offset(
+    name::FieldName,
+    ::Type{T},
+    ::Type{S},
+) where {T, S}
+    if name == @name()
+        return 0
+    end
+    child_name = extract_first(name)
+    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]
+    return DataLayouts.fieldtypeoffset(T, S, field_index) +
+           get_field_first_index_offset(remaining_field_chain, T, child_type)
+end
+if hasfield(Method, :recursion_relation)
+    dont_limit = (args...) -> true
+    for m in methods(get_field_first_index_offset)
+        m.recursion_relation = dont_limit
+    end
+end
+
+"""
+    get_scalar_keys(dict::FieldMatrix)
+
+Returns a `FieldMatrixKeys` object that contains the keys of all the scalar
+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 =
+            entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry
+        unrolled_map(
+            filtered_child_names(
+                field -> eltype(field) <: Number,
+                entry,
+                @name()
+            ),
+        ) do name
+            (append_internal_name(key[1], name), key[2])
+        end
+    end
+    return FieldMatrixKeys(keys_tuple)
+end
+
+"""
+    scalar_fieldmatrix(field_matrix::FieldMatrix)
+
+Constructs a `FieldNameDict` where the keys and entries are views
+of the entries of `field_matrix`, which corresponding to the
+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)
+A = MatrixFields.FieldMatrix(
+    (@name(biz), @name(baz)) => mat1,
+    (@name(bip), @name(bop)) => mat2,
+)
+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))
+```
+"""
+function scalar_fieldmatrix(field_matrix::FieldMatrix)
+    scalar_keys = get_scalar_keys(field_matrix)
+    entries = unrolled_map(scalar_keys.values) do key
+        field_matrix[key]
+    end
+    return FieldNameDict(scalar_keys, entries)
+end
+
 replace_name_tree(dict::FieldNameDict, name_tree) =
     FieldNameDict(replace_name_tree(keys(dict), name_tree), values(dict))
 

test/MatrixFields/field_names.jl

@@ -770,9 +770,9 @@ end
         (@name(a), @name(a)) => -I_CT3XC3,
     )
 
-    for (vector, matrix, I_foo, I_a) in (
-        (vector_of_scalars, matrix_of_scalars, I, I),
-        (vector_of_vectors, matrix_of_tensors, I_C12XCT12, I_CT3XC3),
+    for (vector, matrix, I_foo, I_a, is_scalar_test) in (
+        (vector_of_scalars, matrix_of_scalars, I, I, true),
+        (vector_of_vectors, matrix_of_tensors, I_C12XCT12, I_CT3XC3, false),
     )
         @test_all MatrixFields.field_vector_view(vector) ==
                   MatrixFields.FieldVectorView(
@@ -842,10 +842,13 @@ end
         @test_all matrix[@name(a.c), @name(a.b)] == zero(I_a)
         @test_all matrix[@name(foo._value), @name(foo._value)] ==
                   matrix[@name(foo), @name(foo)]
-
-        @test_all matrix[@name(foo._value), @name(a.b)] isa
-                  Base.AbstractBroadcasted
-        @test Base.materialize(matrix[@name(foo._value), @name(a.b)]) == map(
+        entry = matrix[@name(foo._value), @name(a.b)]
+        @test_all entry isa (
+            is_scalar_test ? MatrixFields.ColumnwiseBandMatrixField :
+            Base.AbstractBroadcasted
+        )
+        entry = is_scalar_test ? entry : Base.materialize(entry)
+        @test entry == map(
             row -> map(foo -> foo.value, row),
             matrix[@name(foo), @name(a.b)],
         )