WIPP1

Author: imreddyTeja

src/MatrixFields/field_name_dict.jl

@@ -163,31 +163,52 @@ function get_internal_entry(
     elseif T <: Geometry.Axis2Tensor &&
            all(n -> is_child_name(n, @name(components.data)), name_pair)
         # two indices needed to index into a 2d tensor (one can be Colon())
-        row_index = extract_first(
-            extract_internal_name(name_pair[1], @name(components.data)),
+        internal_row_name =
+            extract_internal_name(name_pair[1], @name(components.data))
+        internal_col_name =
+            extract_internal_name(name_pair[2], @name(components.data))
+        row_index = extract_first(internal_row_name)
+        col_index = extract_first(internal_col_name)
+        return get_internal_entry(
+            DiagonalMatrixRow(scaling_value(entry)[row_index, col_index]),
+            (drop_first(internal_row_name), drop_first(internal_col_name)),
+            key_error,
         )
-        col_index = extract_first(
-            extract_internal_name(name_pair[2], @name(components.data)),
+    elseif T <: Geometry.AdjointAxisVector
+        return get_internal_entry(
+            DiagonalMatrixRow(getfield(scaling_value(entry), :parent)),
+            name_pair,
+            key_error,
         )
-        return DiagonalMatrixRow(scaling_value(entry)[row_index, col_index])
     else
-        combined_chain = append_internal_name(name_pair[1], name_pair[2])
-        modified_chain =
-            T <: Geometry.AdjointAxisVector && # indexing adjoint of axis vector
-            extract_first(combined_chain) != :parent ? # is the same as its parent
-            extract_internal_name(combined_chain, @name(components.data)) :
-            combined_chain
-        if !broadcasted_has_field(T, modified_chain) # implicit tensor structure case
-            T <: Geometry.SingleValue || throw(key_error) # optimization only works with scalars
-            hasfield(T, extract_first(combined_chain)) && throw(key_error)
-            name_pair[1] == name_pair[2] && return entry # multiplication case 3, 4 first arg
-            is_overlapping_name(name_pair[1], name_pair[2]) && throw(key_error)
-            return zero(entry) # off diagonal
-        else
-            return DiagonalMatrixRow(
-                broadcasted_get_field(scaling_value(entry), modified_chain),
-            ) # multiplication case 2, 4 second arg
-        end
+        child_name, remaining_chain =
+            if name_pair[1] != @name() &&
+               extract_first(name_pair[1]) in fieldnames(T)
+                @inline (
+                    extract_first(name_pair[1]),
+                    (drop_first(name_pair[1]), name_pair[2]),
+                )
+            elseif name_pair[2] != @name() &&
+                   extract_first(name_pair[2]) in fieldnames(T)
+                @inline (
+                    extract_first(name_pair[2]),
+                    (name_pair[1], drop_first(name_pair[2])),
+                )
+            elseif !any(isequal(@name()), name_pair) # implicit tensor structure
+                return get_internal_entry(
+                    extract_first(name_pair[1]) == extract_first(name_pair[2]) ?
+                    entry : zero(entry),
+                    (drop_first(name_pair[1]), drop_first(name_pair[2])),
+                    key_error,
+                )
+            else
+                throw(key_error)
+            end
+        return get_internal_entry(
+            DiagonalMatrixRow(getfield(scaling_value(entry), child_name)),
+            remaining_chain,
+            key_error,
+        )
     end
 end
 function get_internal_entry(
@@ -198,27 +219,9 @@ function get_internal_entry(
     name_pair == (@name(), @name()) && return entry
     S = eltype(eltype(entry))
     T = eltype(parent(entry))
-    first_name = extract_first(append_internal_name(name_pair...))
-    if !hasfield(S, first_name) &&
-       !(hasfield(S, :parent) && hasfield(fieldtype(S, :parent), first_name)) &&
-       S <: Geometry.SingleValue
-        if name_pair[1] == name_pair[2]
-            return entry # multiplication case 3 or 4, first arg
-        elseif is_overlapping_name(name_pair[1], name_pair[2])
-            throw(key_error)
-        else
-            throw(key_error)
-            return Base.broadcasted(entry) do matrix_row # off diagonal
-                map(matrix_row) do matrix_row_entry
-                    zero(S)
-                end
-            end
-        end
-    end
-    # multiplication case 2 or 4, second arg
-    (index_offset, target_type) =
+    (start_offset, target_type, apply_zero) =
         field_offset_and_type(name_pair, T, S, key_error)
-    if target_type <: eltype(parent(entry))
+    if target_type <: eltype(parent(entry)) && !apply_zero
         band_element_size =
             DataLayouts.typesize(eltype(parent(entry)), eltype(eltype(entry)))
         singleton_datalayout = DataLayouts.singleton(Fields.field_values(entry))
@@ -227,7 +230,7 @@ function get_internal_entry(
         field_dim_size = DataLayouts.ncomponents(Fields.field_values(entry))
         parent_indices = DataLayouts.to_data_specific_field(
             singleton_datalayout,
-            (:, :, (index_offset + 1):band_element_size:field_dim_size, :, :),
+            (:, :, (start_offset + 1):band_element_size:field_dim_size, :, :),
         )
         scalar_data = view(parent(entry), parent_indices...)
         values = DataLayouts.union_all(singleton_datalayout){
@@ -237,6 +240,12 @@ function get_internal_entry(
             scalar_data,
         )
         return Fields.Field(values, axes(entry))
+    elseif apply_zero
+        return Base.broadcasted(entry) do matrix_row
+            map(matrix_row) do matrix_row_entry
+               zero(target_type)
+            end
+        end
     else
         return Base.broadcasted(entry) do matrix_row
             map(matrix_row) do matrix_row_entry
@@ -248,6 +257,12 @@ function get_internal_entry(
         end
     end
 end
+if hasfield(Method, :recursion_relation)
+    dont_limit = (args...) -> true
+    for m in methods(get_internal_entry)
+        m.recursion_relation = dont_limit
+    end
+end
 
 # Similar behavior to indexing an array with a slice.
 function Base.getindex(dict::FieldNameDict, new_keys::FieldNameSet)
@@ -313,40 +328,68 @@ function field_offset_and_type(
     ::Type{S},
     key_error,
 ) where {S, T}
-    name_pair == (@name(), @name()) && return (0, S) # base case
-    if S <: Geometry.Axis2Tensor{T} # special case to calculate index
+    name_pair == (@name(), @name()) && return (0, S, false) # base case
+    if S <: Geometry.Axis2Tensor # special case to calculate index
         (name_pair[1] == @name() || name_pair[2] == @name()) && throw(key_error)
-        row_index = extract_first(
-            extract_internal_name(name_pair[1], @name(components.data)),
-        )
-        col_index = extract_first(
-            extract_internal_name(name_pair[2], @name(components.data)),
-        )
+        internal_row_name =
+            extract_internal_name(name_pair[1], @name(components.data))
+        internal_col_name =
+            extract_internal_name(name_pair[2], @name(components.data))
+        row_index = extract_first(internal_row_name)
+        col_index = extract_first(internal_col_name)
         ((row_index isa Number) && (col_index isa Number)) || throw(key_error) # slicing not supported
         (n_rows, n_cols) = map(length, axes(S))
+        (remaining_offset, end_type, apply_zero) = field_offset_and_type(
+            (drop_first(internal_row_name), drop_first(internal_col_name)),
+            T,
+            eltype(S),
+            key_error,
+        )
         (row_index <= n_rows && col_index <= n_cols) || throw(key_error)
-        return (n_rows * (col_index - 1) + row_index - 1, T)
+        return (
+            (n_rows * (col_index - 1) + row_index - 1) + remaining_offset,
+            end_type,
+            apply_zero,
+        )
+    elseif S <: Geometry.AdjointAxisVector
+        return field_offset_and_type(name_pair, T, fieldtype(S, 1), key_error)
     else
-        child_name, remaining_field_chain = if name_pair[1] != @name()
-            extract_first(name_pair[1]), (drop_first(name_pair[1]), name_pair[2])
-        else
-            extract_first(name_pair[2]), (@name(), drop_first(name_pair[2]))
-        end
-        # indexing adjoint of axis vector is the same as indexing its parent
-        (S <: Geometry.AdjointAxisVector && child_name != :parent) &&
-            return field_offset_and_type(
-                name_pair,
-                T,
-                fieldtype(S, 1),
-                key_error,
-            )
-        (child_name in fieldnames(S)) || throw(key_error)
+        child_name, remaining_field_chain =
+            if name_pair[1] != @name() &&
+               extract_first(name_pair[1]) in fieldnames(S)
+                @inline (
+                    extract_first(name_pair[1]),
+                    (drop_first(name_pair[1]), name_pair[2]),
+                )
+            elseif name_pair[2] != @name() &&
+                   extract_first(name_pair[2]) in fieldnames(S)
+                @inline (
+                    extract_first(name_pair[2]),
+                    (name_pair[1], drop_first(name_pair[2])),
+                )
+            elseif !any(isequal(@name()), name_pair) # implicit tensor structure
+                (remaining_offset, end_type, apply_zero) =
+                    field_offset_and_type(
+                        (drop_first(name_pair[1]), drop_first(name_pair[2])),
+                        T,
+                        fieldtype(S, 1),
+                        key_error,
+                    )
+                return (
+                    remaining_offset,
+                    end_type,
+                    extract_first(name_pair[1]) == extract_first(name_pair[2]) ?
+                    apply_zero : true,
+                )
+            else
+                throw(key_error)
+            end
         child_type = fieldtype(S, child_name)
         field_index = unrolled_filter(
             i -> fieldname(S, i) == child_name,
             1:fieldcount(S),
         )[1]
-        (remaining_offset, end_type) = field_offset_and_type(
+        (remaining_offset, end_type, apply_zero) = field_offset_and_type(
             remaining_field_chain,
             T,
             child_type,
@@ -355,7 +398,9 @@ function field_offset_and_type(
         return (
             DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset,
             end_type,
+            apply_zero,
         )
+
     end
 end
 if hasfield(Method, :recursion_relation)
@@ -366,16 +411,16 @@ if hasfield(Method, :recursion_relation)
 end
 
 """
-    get_scalar_keys(dict::FieldMatrix, FT)
+    get_scalar_keys(dict::FieldMatrix)
 
 Returns a `FieldMatrixKeys` object that contains the keys that result in
-a `ScalingFieldMatrixEntry{FT}` or a `ColumnwiseBandMatrixField` with bands of eltype `FT`
+a `ScalingFieldMatrixEntry{<:Number}` or a `ColumnwiseBandMatrixField` with bands of eltype `< :Number`
 when indexing `dict`.
 """
-function get_scalar_keys(dict::FieldMatrix, ::Type{FT}) where {FT}
+function get_scalar_keys(dict::FieldMatrix)
     keys_tuple = unrolled_flatmap(keys(dict).values) do outer_key
-        unrolled_map(get_scalar_keys(eltype(dict[outer_key]), FT)) do inner_key
-            (
+        @inline unrolled_map(get_scalar_keys(eltype(dict[outer_key]))) do inner_key
+            @inline (
                 append_internal_name(outer_key[1], inner_key[1]),
                 append_internal_name(outer_key[2], inner_key[2]),
             )
@@ -385,16 +430,16 @@ function get_scalar_keys(dict::FieldMatrix, ::Type{FT}) where {FT}
 end
 
 """
-    get_scalar_keys(T::Type, FT::Type)
+    get_scalar_keys(T::Type)
 
 Returns a tuple of `FieldNamePair` objects that correspond to any children
 of `T` that are of type `FT`.
 """
-function get_scalar_keys(::Type{T}, ::Type{FT}) where {T, FT}
-    if T <: FT || T <: Bool # identity has eltype Bool
+function get_scalar_keys(::Type{T}) where {T}
+    if T <: Number # TODO: is this tight enough of a Type? what about complex and duals and plushalfs
         return ((@name(), @name()),)
     elseif T <: BandMatrixRow
-        return get_scalar_keys(eltype(T), FT)
+        return get_scalar_keys(eltype(T))
     elseif T <: Geometry.Axis2Tensor
         return unrolled_flatmap(1:length(axes(T)[1])) do row_component
             unrolled_map(1:length(axes(T)[2])) do col_component
@@ -405,28 +450,35 @@ function get_scalar_keys(::Type{T}, ::Type{FT}) where {T, FT}
             end
         end
     elseif T <: Geometry.AdjointAxisVector
-        return unrolled_map(
-            get_scalar_keys(fieldtype(T, :parent), FT),
-        ) do inner_key
+        return unrolled_map(get_scalar_keys(fieldtype(T, :parent))) do inner_key
             (inner_key[2], inner_key[1]) # assumes that adjoints only appear with d/dvec
         end
     elseif T <: Geometry.AxisVector # special case to avoid recursing into the axis field
         # TODO: this should be able to be combined with the else case, but it causes runtime dispatch
         return unrolled_map(
-            get_scalar_keys(fieldtype(T, :components), FT),
+            get_scalar_keys(fieldtype(T, :components)),
         ) do inner_key
             (
                 append_internal_name(@name(components), inner_key[1]),
                 inner_key[2],
             )
         end
+        # return unrolled_flatmap((:components,)) do inner_field
+        #     @inline unrolled_map(
+        #         # get_scalar_keys(fieldtype(T, inner_field)),
+        #         get_scalar_keys(fieldtype(T, :components)),
+        #     ) do inner_key
+        #         @inline (
+        #             append_internal_name(FieldName(inner_field), inner_key[1]),
+        #             inner_key[2],
+        #         )
+        #     end
+        # end
     else
-        return unrolled_flatmap(fieldnames(T)) do inner_field
-            unrolled_map(
-                get_scalar_keys(fieldtype(T, inner_field), FT),
-            ) do inner_key
+        return unrolled_flatmap(fieldnames(T)) do inner_name
+            unrolled_map(get_scalar_keys(fieldtype(T, inner_name))) do inner_key
                 (
-                    append_internal_name(FieldName(inner_field), inner_key[1]),
+                    append_internal_name(FieldName(inner_name), inner_key[1]),
                     inner_key[2],
                 )
             end
@@ -442,7 +494,7 @@ end
 
 
 """
-    scalar_fieldmatrix(field_matrix::FieldMatrix, FT)
+    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
@@ -464,7 +516,7 @@ A = MatrixFields.FieldMatrix(
     (@name(c.uₕ), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3_uₕ_scalar,
 )
 
-A_scalar = MatrixFields.scalar_fieldmatrix(A, Float64)
+A_scalar = MatrixFields.scalar_fieldmatrix(A)
 keys(A_scalar)
 # Output:
 # (@name(c.ρχ.ρq_liq), @name(f.u₃.:(1)))
@@ -473,8 +525,8 @@ keys(A_scalar)
 # (@name(c.uₕ.:(2)), @name(c.sgsʲs.:(1).ρa))
 ```
 """
-function scalar_fieldmatrix(field_matrix::FieldMatrix, ::Type{FT}) where {FT}
-    scalar_keys = get_scalar_keys(field_matrix, FT)
+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