fix implicit tensor rep tests

Author: imreddyTeja

src/MatrixFields/field_name.jl

@@ -50,9 +50,6 @@ extract_first(::FieldName{name_chain}) where {name_chain} = first(name_chain)
 drop_first(::FieldName{name_chain}) where {name_chain} =
     FieldName(Base.tail(name_chain)...)
 
-extract_last(::FieldName{name_chain}) where {name_chain} =
-    name_chain[length(name_chain)]
-
 has_field(x, ::FieldName{()}) = true
 has_field(x, name::FieldName) =
     extract_first(name) in propertynames(x) &&

src/MatrixFields/field_name_dict.jl

@@ -152,62 +152,42 @@ get_internal_key(child_name_pair::FieldNamePair, name_pair::FieldNamePair) = (
 get_internal_entry(entry, name::FieldName, key_error) = get_field(entry, name)
 get_internal_entry(entry, name_pair::FieldNamePair, key_error) =
     name_pair == (@name(), @name()) ? entry : throw(key_error)
-get_internal_entry(entry::UniformScaling, name_pair::FieldNamePair, key_error) =
-    if name_pair[1] == name_pair[2]
-        entry
-    elseif is_overlapping_name(name_pair[1], name_pair[2])
-        throw(key_error)
-    else
-        zero(entry)
-    end
 function get_internal_entry(
-    entry::DiagonalMatrixRow{T},
+    entry::ScalingFieldMatrixEntry,
     name_pair::FieldNamePair,
     key_error,
-) where {T}
+)
+    T = eltype(entry)
     if name_pair == (@name(), @name())
         return entry
-    elseif T <: Geometry.Axis2Tensor
-        (name_pair[1] == @name() || name_pair[2] == @name()) && throw(key_error) # Cannot slice a 2D tensor
-        row_index = extract_first(name_pair[1])
-        col_index = extract_first(name_pair[2])
-        if (!(row_index isa Number) || !(col_index isa Number))
-            name_pair[1] == name_pair[2] && return entry # multiplication case 3 or 4, first argument
-            is_overlapping_name(name_pair[1], name_pair[2]) ||
-                return zero(entry)
-            throw(key_error)
-        end
-        (n_rows, n_cols) = map(length, axes(scaling_value(entry)))
-        @assert row_index <= n_rows && col_index <= n_cols
-        return DiagonalMatrixRow(
-            broadcasted_get_field(
-                scaling_value(entry).components.data,
-                FieldName(n_rows * (col_index - 1) + row_index),
-            ),
+    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)),
         )
-    elseif T <: Geometry.AxisVectorOrAdj
-        return get_internal_entry(
-            DiagonalMatrixRow(
-                getfield(
-                    getfield(
-                        T <: Geometry.AxisVector ? scaling_value(entry) :
-                        scaling_value(entry).parent,
-                        :components,
-                    ),
-                    :data,
-                ),
-            ),
-            name_pair,
-            key_error,
+        col_index = extract_first(
+            extract_internal_name(name_pair[2], @name(components.data)),
         )
-    elseif name_pair[1] == name_pair[2] &&
-           !broadcasted_has_field(eltype(entry), name_pair[1])
-        return entry
+        return DiagonalMatrixRow(scaling_value(entry)[row_index, col_index])
     else
-        non_empty_chain = name_pair[1] != @name() ? name_pair[1] : name_pair[2]
-        return DiagonalMatrixRow(
-            broadcasted_get_field(scaling_value(entry), non_empty_chain),
-        )
+        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
     end
 end
 function get_internal_entry(
@@ -216,13 +196,28 @@ function get_internal_entry(
     key_error,
 )
     name_pair == (@name(), @name()) && return entry
-    (index_offset, target_type) = field_offset_and_type(
-        name_pair,
-        eltype(parent(entry)),
-        eltype(eltype(entry)),
-        key_error,
-    )
-    target_type <: eltype(eltype(entry)) && return entry # multiplication case 3 or 4, first argument
+    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) =
+        field_offset_and_type(name_pair, T, S, key_error)
     if target_type <: eltype(parent(entry))
         band_element_size =
             DataLayouts.typesize(eltype(parent(entry)), eltype(eltype(entry)))
@@ -254,7 +249,6 @@ function get_internal_entry(
     end
 end
 
-
 # Similar behavior to indexing an array with a slice.
 function Base.getindex(dict::FieldNameDict, new_keys::FieldNameSet)
     common_keys = intersect(keys(dict), new_keys)
@@ -320,76 +314,34 @@ function field_offset_and_type(
     key_error,
 ) where {S, T}
     name_pair == (@name(), @name()) && return (0, S) # base case
-    if S <: Geometry.Axis2Tensor{T}
-        (name_pair[1] == @name() || name_pair[2] == @name()) && throw(key_error) # Cannot slice a 2D tensor
-        row_index = extract_first(name_pair[1])
-        col_index = extract_first(name_pair[2])
-        if (!(row_index isa Number) || !(col_index isa Number))
-            name_pair[1] == name_pair[2] && return (0, S) # multiplication case 3 or 4, first argument
-            throw(key_error)
-        end
-        (n_rows, n_cols) = map(length, axes(S))
-        @assert row_index <= n_rows && col_index <= n_cols
-        return (n_rows * (col_index - 1) + row_index - 1, T)
-    elseif (
-        name_pair[1] == name_pair[2] && !broadcasted_has_field(S, name_pair[1])
-    )
-        return (0, S)
-    elseif name_pair[1] == @name()
-        return field_offset_and_type(name_pair[2], T, S, key_error)
-    elseif name_pair[2] == @name()
-        return field_offset_and_type(name_pair[1], T, S, key_error)
-    else
-        child_name = extract_first(name_pair[1])
-        (child_name in fieldnames(S)) || throw(key_error)
-        child_type = fieldtype(S, child_name)
-        remaining_field_chain = (drop_first(name_pair[1]), name_pair[2])
-        field_index = unrolled_filter(
-            i -> fieldname(S, i) == child_name,
-            1:fieldcount(S),
-        )[1]
-        (remaining_offset, end_type) = field_offset_and_type(
-            remaining_field_chain,
-            T,
-            child_type,
-            key_error,
+    if S <: Geometry.Axis2Tensor{T} # 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)),
         )
-        return (
-            DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset,
-            end_type,
-        )
-    end
-end
-
-"""
-    field_offset_and_type(name::FieldName, ::Type{T}, ::Type{S}, key_error)
-
-Returns the offset of the the field with name `name` in an object of type `S`
-in multiples of `sizeof(T)` and the type of the field with name `name` in an object of type `S`
-in multiples of `sizeof(T)`
-"""
-function field_offset_and_type(
-    name::FieldName,
-    ::Type{T},
-    ::Type{S},
-    key_error,
-) where {T, S}
-    name == @name() && return (0, S) # base case
-    if S <: Geometry.AdjointAxisVector
-        return field_offset_and_type(name, T, fieldtype(S, :parent), key_error)
-    elseif S <: Geometry.AxisVector
-        (remaining_offset, end_type) = field_offset_and_type(
-            name,
-            T,
-            fieldtype(fieldtype(S, :components), :data),
-            key_error,
+        col_index = extract_first(
+            extract_internal_name(name_pair[2], @name(components.data)),
         )
-        return (remaining_offset, end_type)
+        ((row_index isa Number) && (col_index isa Number)) || throw(key_error) # slicing not supported
+        (n_rows, n_cols) = map(length, axes(S))
+        (row_index <= n_rows && col_index <= n_cols) || throw(key_error)
+        return (n_rows * (col_index - 1) + row_index - 1, T)
     else
-        child_name = extract_first(name)
+        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_type = fieldtype(S, child_name)
-        remaining_field_chain = drop_first(name)
         field_index = unrolled_filter(
             i -> fieldname(S, i) == child_name,
             1:fieldcount(S),
@@ -422,7 +374,6 @@ when indexing `dict`.
 """
 function get_scalar_keys(dict::FieldMatrix, ::Type{FT}) where {FT}
     keys_tuple = unrolled_flatmap(keys(dict).values) do outer_key
-        # target_eltype = eltype(Y)
         unrolled_map(get_scalar_keys(eltype(dict[outer_key]), FT)) do inner_key
             (
                 append_internal_name(outer_key[1], inner_key[1]),
@@ -444,24 +395,33 @@ function get_scalar_keys(::Type{T}, ::Type{FT}) where {T, FT}
         return ((@name(), @name()),)
     elseif T <: BandMatrixRow
         return get_scalar_keys(eltype(T), FT)
-    elseif T <: Geometry.AxisVector
-        return unrolled_map(1:length(axes(T)[1])) do component
-            (FieldName(component), @name())
+    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
+                append_internal_name.(
+                    Ref(@name(components.data)),
+                    (FieldName(row_component), FieldName(col_component)),
+                )
+            end
         end
     elseif T <: Geometry.AdjointAxisVector
         return unrolled_map(
-            1:length(axes(fieldtype(T, :parent))[1]),
-        ) do component
-            (@name(), FieldName(component))
+            get_scalar_keys(fieldtype(T, :parent), FT),
+        ) do inner_key
+            (inner_key[2], inner_key[1]) # assumes that adjoints only appear with d/dvec
         end
-    elseif T <: Geometry.Axis2Tensor
-        unrolled_flatmap(1:length(axes(T)[1])) do row_component
-            unrolled_map(1:length(axes(T)[2])) do col_component
-                (FieldName(row_component), FieldName(col_component))
-            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),
+        ) do inner_key
+            (
+                append_internal_name(@name(components), inner_key[1]),
+                inner_key[2],
+            )
         end
     else
-        unrolled_flatmap(fieldnames(T)) do inner_field
+        return unrolled_flatmap(fieldnames(T)) do inner_field
             unrolled_map(
                 get_scalar_keys(fieldtype(T, inner_field), FT),
             ) do inner_key

test/MatrixFields/scalar_fieldmatrix.jl

@@ -30,31 +30,31 @@ include("matrix_field_test_utils.jl")
                   (expected_offset, E)
     end
     test_field_offset_and_type(
-        @name(x),
+        (@name(x), @name()),
         FT,
         Singleton{Singleton{Singleton{Singleton{FT}}}},
         0,
         Singleton{Singleton{Singleton{FT}}},
         KeyError(@name(x.x.x.x)),
     )
     test_field_offset_and_type(
-        @name(x.x.x.x),
+        (@name(), @name(x.x.x.x)),
         FT,
         Singleton{Singleton{Singleton{Singleton{FT}}}},
         0,
         FT,
         KeyError(@name(x.x.x.x)),
     )
     test_field_offset_and_type(
-        @name(y.x),
+        (@name(), @name(y.x)),
         FT,
         TwoFields{TwoFields{FT, FT}, TwoFields{FT, FT}},
         2,
         FT,
         KeyError(@name(y.x)),
     )
     test_field_offset_and_type(
-        @name(y.y),
+        (@name(y), @name(y)),
         FT,
         TwoFields{
             TwoFields{FT, FT},
@@ -64,14 +64,6 @@ include("matrix_field_test_utils.jl")
         TwoFields{FT, Singleton{FT}},
         KeyError(@name(y.y.x)),
     )
-    test_field_offset_and_type(
-        @name(y.y),
-        Float32,
-        TwoFields{TwoFields{FT, FT}, TwoFields{FT, FT}},
-        6,
-        FT,
-        KeyError(@name(y.y.x)),
-    )
     test_field_offset_and_type(
         (@name(y.y), @name(x)),
         FT,