fix broken tests

Author: imreddyTeja

docs/src/matrix_fields.md

@@ -90,9 +90,7 @@ check_preconditioner
 lazy_or_concrete_preconditioner
 apply_preconditioner
 get_scalar_keys
-get_field_first_index_offset
-broadcasted_get_field_type
-inner_type_ignore_adjoint
+field_offset_and_type
 ```
 
 ## Utilities
@@ -160,15 +158,16 @@ If the key `(@name(name1), @name(name2))` corresponds to an entry, then
 `(@name(name1.foo.bar.buz), @name(name2.biz.bop.fud))`.
 
 Currently, internal values cannot be extracted in all situations. Extracting interal values
-works when:
+works when indexing an object of type `eltype(entry)` with the
+second key of the internal key pair appended to the first results in a scalar.
+If the internal keys index to a non-scalar `Field`, a broadcasted object is returned.
 
-- The second name in the internal key is empty, and the first name in the internal key accesses internal values for the type of element contained in each row of the entry. This does not work when the element type of each row is a 2d tensor.
+When the entry is a `Field` of `Axis2Tensor`s, and both internal names are numbers that would index
+an `Axis2Tensor` with the same axis.
 
-- The first name in the internal key is empty, and the type of element contained in each row of the entry is an `AxisVector` or the adjoint of an `AxisVector`. In this case, the second name must access inernal values for the type of `AxisVector` contained in each row.
+This does not work when the internal keys index to a `Field` of sliced tensors.
 
-- The element type of each row in the entry is a 2d tensor, and the internal key is of the form `(@name(components.data.:(1)), @name(components.data.:(2)))`, but possibly with different numbers to index into the 2d tensor
-
-- The element type of each row in the entry is some number of nested `Tuple`s and `NamedTuple`s, and the first name in the internal key accesses an `AxisVector` or the adjoint of an `AxisVector` from the outer `Tuple`/`NamedTuple`, and the second name in the inernal key accesses a component of the `AxisVector`
+Extracting internal values from a `DiagonalMatrixRow` works in all cases, except when
 
 If the `FieldMatrix` represents a Jacobian, then extracting internal values works when an entry represents:
 

src/MatrixFields/field_name.jl

@@ -62,18 +62,6 @@ 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),
-    )
-
 broadcasted_has_field(::Type{X}, ::FieldName{()}) where {X} = true
 broadcasted_has_field(::Type{X}, name::FieldName) where {X} =
     extract_first(name) in fieldnames(X) &&
@@ -214,7 +202,4 @@ if hasfield(Method, :recursion_relation)
     for m in methods(get_subtree_at_name)
         m.recursion_relation = dont_limit
     end
-    for m in methods(broadcasted_get_field_type)
-        m.recursion_relation = dont_limit
-    end
 end

src/MatrixFields/field_name_dict.jl

@@ -168,10 +168,15 @@ function get_internal_entry(
     if name_pair == (@name(), @name())
         return entry
     elseif T <: Geometry.Axis2Tensor
-        (name_pair[1] == @name() || name_pair[2] == @name()) &&
-            error("Cannot slice a 2D tensor")
+        (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(
@@ -215,7 +220,9 @@ function get_internal_entry(
         name_pair,
         eltype(parent(entry)),
         eltype(eltype(entry)),
+        key_error,
     )
+    target_type <: eltype(eltype(entry)) && return entry # multiplication case 3 or 4, first argument
     if target_type <: eltype(parent(entry))
         band_element_size =
             DataLayouts.typesize(eltype(parent(entry)), eltype(eltype(entry)))
@@ -295,7 +302,7 @@ function Base.one(matrix::FieldMatrix)
 end
 
 """
-    field_offset_and_type(name_pair::FieldNamePair, ::Type{T}, ::Type{S})
+    field_offset_and_type(name_pair::FieldNamePair, ::Type{T}, ::Type{S}, key_error)
 
 Returns the offset of the field with name `name_pair` in an object of type `S` in
 multiples of `sizeof(T)` and the type of the field with name `name_pair`.
@@ -310,13 +317,17 @@ function field_offset_and_type(
     name_pair::FieldNamePair,
     ::Type{T},
     ::Type{S},
+    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()) &&
-            error("Cannot slice a 2D tensor")
+        (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)
@@ -325,19 +336,24 @@ function field_offset_and_type(
     )
         return (0, S)
     elseif name_pair[1] == @name()
-        return field_offset_and_type(name_pair[2], T, S)
+        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)
+        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)
+        (remaining_offset, end_type) = field_offset_and_type(
+            remaining_field_chain,
+            T,
+            child_type,
+            key_error,
+        )
         return (
             DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset,
             end_type,
@@ -346,7 +362,7 @@ function field_offset_and_type(
 end
 
 """
-    field_offset_and_type(name::FieldName, ::Type{T}, ::Type{S})
+    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`
@@ -356,27 +372,34 @@ 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))
+        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,
         )
         return (remaining_offset, end_type)
     else
         child_name = extract_first(name)
+        (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),
         )[1]
-        (remaining_offset, end_type) =
-            field_offset_and_type(remaining_field_chain, T, child_type)
+        (remaining_offset, end_type) = field_offset_and_type(
+            remaining_field_chain,
+            T,
+            child_type,
+            key_error,
+        )
         return (
             DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset,
             end_type,

test/MatrixFields/field_names.jl

@@ -834,12 +834,8 @@ end
 
         @test_throws KeyError matrix[@name(a), @name(a.c)]
         @test_throws KeyError matrix[@name(a.c), @name(a)]
-        @test_throws AssertionError matrix[@name(foo), @name(foo._value)]
-        if is_scalar_test
-            @test_throws KeyError matrix[@name(foo._value), @name(foo)]
-        else
-            @test_throws AssertionError matrix[@name(foo._value), @name(foo)]
-        end
+        @test_throws KeyError matrix[@name(foo), @name(foo._value)]
+        @test_throws KeyError matrix[@name(foo._value), @name(foo)]
 
         @test_all matrix[@name(a), @name(a)] == -I_a
         @test_all matrix[@name(a.c), @name(a.c)] == -I_a

test/MatrixFields/scalar_fieldmatrix.jl

@@ -24,8 +24,9 @@ include("matrix_field_test_utils.jl")
         ::Type{S},
         expected_offset,
         ::Type{E},
+        key_error,
     ) where {T, S, E}
-        @test_all MatrixFields.field_offset_and_type(name, T, S) ==
+        @test_all MatrixFields.field_offset_and_type(name, T, S, key_error) ==
                   (expected_offset, E)
     end
     test_field_offset_and_type(
@@ -34,20 +35,23 @@ include("matrix_field_test_utils.jl")
         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),
         FT,
         Singleton{Singleton{Singleton{Singleton{FT}}}},
         0,
         FT,
+        KeyError(@name(x.x.x.x)),
     )
     test_field_offset_and_type(
         @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),
@@ -58,78 +62,37 @@ include("matrix_field_test_utils.jl")
         },
         3,
         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.x),
+        (@name(y.y), @name(x)),
         FT,
         TwoFields{
             TwoFields{FT, FT},
             TwoFields{FT, TwoFields{FT, Singleton{FT}}},
         },
         3,
         FT,
+        KeyError(@name(y.y.x.x)),
     )
     test_field_offset_and_type(
-        @name(y.y.y.x),
+        (@name(y.y.y), @name(y.x)),
         FT,
         TwoFields{
             TwoFields{FT, FT},
             TwoFields{FT, TwoFields{FT, Singleton{FT}}},
         },
         4,
         FT,
-    )
-end
-
-@testset "broadcasted_get_field_type" begin
-    FT = Float64
-    struct Singleton{T}
-        x::T
-    end
-    struct TwoFields{T1, T2}
-        x::T1
-        y::T2
-    end
-    function test_broadcasted_get_field_type(
-        name,
-        ::Type{T},
-        expected_type,
-    ) where {T}
-        @test_all MatrixFields.broadcasted_get_field_type(T, name) ==
-                  expected_type
-    end
-    test_broadcasted_get_field_type(
-        @name(x),
-        Singleton{Singleton{Singleton{Singleton{FT}}}},
-        Singleton{Singleton{Singleton{FT}}},
-    )
-    test_broadcasted_get_field_type(
-        @name(x.x.x),
-        Singleton{Singleton{Singleton{Singleton{FT}}}},
-        Singleton{FT},
-    )
-    test_broadcasted_get_field_type(
-        @name(y.x),
-        TwoFields{
-            TwoFields{FT, FT},
-            TwoFields{FT, TwoFields{FT, Singleton{FT}}},
-        },
-        FT,
-    )
-    test_broadcasted_get_field_type(
-        @name(y.y.y),
-        TwoFields{
-            TwoFields{FT, FT},
-            TwoFields{FT, TwoFields{FT, Singleton{FT}}},
-        },
-        Singleton{FT},
+        KeyError(@name(y.y.y.x.x)),
     )
 end