From c9f0a1496b1e6968e816727653cdee3af3a096ce Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Mon, 31 Mar 2025 14:44:11 -0700 Subject: [PATCH 01/10] 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 --- docs/src/matrix_fields.md | 4 + src/MatrixFields/MatrixFields.jl | 1 + src/MatrixFields/field_name.jl | 18 ++ src/MatrixFields/field_name_dict.jl | 140 +++++++++++- test/MatrixFields/field_names.jl | 17 +- test/MatrixFields/scalar_fieldmatrix.jl | 273 ++++++++++++++++++++++++ 6 files changed, 441 insertions(+), 12 deletions(-) create mode 100644 test/MatrixFields/scalar_fieldmatrix.jl diff --git a/docs/src/matrix_fields.md b/docs/src/matrix_fields.md index 4c89aa765d..650035101d 100644 --- a/docs/src/matrix_fields.md +++ b/docs/src/matrix_fields.md @@ -89,6 +89,9 @@ preconditioner_cache check_preconditioner lazy_or_concrete_preconditioner apply_preconditioner +get_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 ``` diff --git a/src/MatrixFields/MatrixFields.jl b/src/MatrixFields/MatrixFields.jl index e2acdc67c1..43853e0096 100644 --- a/src/MatrixFields/MatrixFields.jl +++ b/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 diff --git a/src/MatrixFields/field_name.jl b/src/MatrixFields/field_name.jl index bd05f844c1..3b66470292 100644 --- a/src/MatrixFields/field_name.jl +++ b/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) && diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 2d98bb3ff3..d6aebb60e5 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/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)) diff --git a/test/MatrixFields/field_names.jl b/test/MatrixFields/field_names.jl index 0253dd57bc..6966124e60 100644 --- a/test/MatrixFields/field_names.jl +++ b/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)], ) diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl new file mode 100644 index 0000000000..b303ac6d08 --- /dev/null +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -0,0 +1,273 @@ +using Test +using JET + +import ClimaCore: Geometry, Domains, Meshes, Spaces, Fields, MatrixFields +import ClimaCore.Utilities: half +import ClimaCore.RecursiveApply: ⊠ +import ClimaComms +import LinearAlgebra: I, norm, ldiv!, mul! +import ClimaCore.MatrixFields: @name +ClimaComms.@import_required_backends +include("matrix_field_test_utils.jl") + +@testset "get_field_first_index_offset" begin + struct singleton{T} + x::T + end + struct two_fields{T1, T2} + x::T1 + y::T2 + end + function test_get_field_first_index_offset( + name, + ::Type{T}, + ::Type{S}, + expected_offset, + ) where {T, S} + @test_all MatrixFields.get_field_first_index_offset(name, T, S) == + expected_offset + end + test_get_field_first_index_offset( + @name(x), + Float64, + singleton{singleton{singleton{singleton{Float64}}}}, + 0, + ) + test_get_field_first_index_offset( + @name(x.x.x.x), + Float64, + singleton{singleton{singleton{singleton{Float64}}}}, + 0, + ) + test_get_field_first_index_offset( + @name(y.x), + Float64, + two_fields{two_fields{Float64, Float64}, two_fields{Float64, Float64}}, + 2, + ) + test_get_field_first_index_offset( + @name(y.y), + Float64, + two_fields{ + two_fields{Float64, Float64}, + two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + }, + 3, + ) + test_get_field_first_index_offset( + @name(y.y), + Float32, + two_fields{two_fields{Float64, Float64}, two_fields{Float64, Float64}}, + 6, + ) + test_get_field_first_index_offset( + @name(y.y.x), + Float64, + two_fields{ + two_fields{Float64, Float64}, + two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + }, + 3, + ) + test_get_field_first_index_offset( + @name(y.y.y.x), + Float64, + two_fields{ + two_fields{Float64, Float64}, + two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + }, + 4, + ) +end + +@testset "broadcasted_get_field_type" begin + struct singleton{T} + x::T + end + struct two_fields{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{Float64}}}}, + singleton{singleton{singleton{Float64}}}, + ) + test_broadcasted_get_field_type( + @name(x.x.x), + singleton{singleton{singleton{singleton{Float64}}}}, + singleton{Float64}, + ) + test_broadcasted_get_field_type( + @name(y.x), + two_fields{ + two_fields{Float64, Float64}, + two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + }, + Float64, + ) + test_broadcasted_get_field_type( + @name(y.y.y), + two_fields{ + two_fields{Float64, Float64}, + two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + }, + singleton{Float64}, + ) +end + +@testset "fieldmatrix to scalar fieldmatrix unit tests" begin + FT = Float64 + center_space, face_space = test_spaces(FT) + surface_space = Spaces.level(face_space, half) + seed!(1) + ᶜvec = random_field(FT, center_space) + ᶠvec = random_field(FT, face_space) + sfc_vec = random_field(FT, surface_space) + λ = 10 + ᶜᶜmat1 = random_field(DiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) + ᶜᶠmat2 = random_field(BidiagonalMatrixRow{FT}, center_space) ./ λ + ᶠᶜmat2 = random_field(BidiagonalMatrixRow{FT}, face_space) ./ λ + ᶜᶜmat3 = random_field(TridiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) + ᶠᶠmat3 = random_field(TridiagonalMatrixRow{FT}, face_space) ./ λ .+ (I,) + + e¹² = Geometry.Covariant12Vector(1, 2) + e³ = Geometry.Covariant3Vector(1) + e₃ = Geometry.Contravariant3Vector(1) + + ρχ_unit = (; ρq_tot = 1, ρq_liq = 1, ρq_ice = 1, ρq_rai = 1, ρq_sno = 1) + ρaχ_unit = + (; ρaq_tot = 1, ρaq_liq = 1, ρaq_ice = 1, ρaq_rai = 1, ρaq_sno = 1) + + dry_center_gs_unit = (; ρ = 1, ρe_tot = 1, uₕ = e¹²) + 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₃',) + ᶜᶜmat3_scalar_uₕ = ᶜᶜmat3 .* (e¹²,) + ᶜᶜmat1_scalar_uₕ = ᶜᶜmat1 .* (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) + + A = 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₃)) => ᶠᶠ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)) => ᶜᶜmat3_ρχ_scalar, + (@name(c.uₕ), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3_scalar_uₕ, + (@name(c.sgsʲs.:(1).ρa), @name(c.uₕ)) => ᶜᶜ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₃)) => ᶜᶠ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₃)) => ᶜᶠ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₃), + ) + + scalar_keys = MatrixFields.get_scalar_keys(A) + @test all( + f -> f isa MatrixFields.UniformScaling || eltype(eltype(f)) <: FT, + getindex.(Ref(A), scalar_keys), + ) + @test (@allocated MatrixFields.get_scalar_keys(A)) == 0 + @test_opt MatrixFields.get_scalar_keys(A) + foreach(scalar_keys) do key + @test_opt A[key] + @test (@allocated A[key]) == 0 + end + + A_scalar = MatrixFields.scalar_fieldmatrix(A) + @test all( + f -> f isa MatrixFields.UniformScaling || eltype(eltype(f)) <: FT, + getindex.(Ref(A_scalar), scalar_keys), + ) + @test all(scalar_keys) do key + entry = A_scalar[key] + entry isa MatrixFields.UniformScaling && return true + key′, entry′ = + filter(pair -> MatrixFields.is_child_value(key, pair[1]), pairs(A))[1] + if key != key′ + parent(entry) isa SubArray || + parent(entry) === parent(entry′) || + return false + end + parent(parent(entry)) === parent(entry′) || return false + return true + end + function scalar_fieldmatrix_wrapper(field_matrix_of_tensors) + MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors) + return + end + scalar_fieldmatrix_wrapper(A) + @test (@allocated scalar_fieldmatrix_wrapper(A)) == 0 + @test_opt MatrixFields.scalar_fieldmatrix(A) + + A_copy = deepcopy(A) + @test all(scalar_keys) do key + entry_copy = A_copy[key] + entry_original = A[key] + entry_flattened = A_scalar[key] + entry_flattened isa MatrixFields.UniformScaling && return true + entry_original == entry_copy || return false + entry_original === entry_flattened || return false + entry_flattened .*= 2 + entry_original == 2 .* entry_copy || return false + return true + end +end From ea5f42161400e0e7821f99f641e84a8d7e22abea Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Tue, 22 Apr 2025 17:02:00 -0700 Subject: [PATCH 02/10] Make suggested changes to tests and field_name_dict.jl --- src/MatrixFields/field_name_dict.jl | 56 ++++---- test/MatrixFields/field_matrix_solvers.jl | 66 ++------- test/MatrixFields/matrix_field_test_utils.jl | 85 ++++++++++- test/MatrixFields/scalar_fieldmatrix.jl | 143 ++++++------------- 4 files changed, 174 insertions(+), 176 deletions(-) diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index d6aebb60e5..1a8f3a5a16 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/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) diff --git a/test/MatrixFields/field_matrix_solvers.jl b/test/MatrixFields/field_matrix_solvers.jl index 524d5b8b9e..d6ea0c7f40 100644 --- a/test/MatrixFields/field_matrix_solvers.jl +++ b/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, ) diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl index 9e65be8830..878603c490 100644 --- a/test/MatrixFields/matrix_field_test_utils.jl +++ b/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} diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl index b303ac6d08..e6e5e63000 100644 --- a/test/MatrixFields/scalar_fieldmatrix.jl +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -3,18 +3,16 @@ using JET import ClimaCore: Geometry, Domains, Meshes, Spaces, Fields, MatrixFields import ClimaCore.Utilities: half -import ClimaCore.RecursiveApply: ⊠ import ClimaComms -import LinearAlgebra: I, norm, ldiv!, mul! import ClimaCore.MatrixFields: @name ClimaComms.@import_required_backends include("matrix_field_test_utils.jl") @testset "get_field_first_index_offset" begin - struct singleton{T} + struct Singleton{T} x::T end - struct two_fields{T1, T2} + struct Two_fields{T1, T2} x::T1 y::T2 end @@ -30,61 +28,61 @@ include("matrix_field_test_utils.jl") test_get_field_first_index_offset( @name(x), Float64, - singleton{singleton{singleton{singleton{Float64}}}}, + Singleton{Singleton{Singleton{Singleton{Float64}}}}, 0, ) test_get_field_first_index_offset( @name(x.x.x.x), Float64, - singleton{singleton{singleton{singleton{Float64}}}}, + Singleton{Singleton{Singleton{Singleton{Float64}}}}, 0, ) test_get_field_first_index_offset( @name(y.x), Float64, - two_fields{two_fields{Float64, Float64}, two_fields{Float64, Float64}}, + Two_fields{Two_fields{Float64, Float64}, Two_fields{Float64, Float64}}, 2, ) test_get_field_first_index_offset( @name(y.y), Float64, - two_fields{ - two_fields{Float64, Float64}, - two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + Two_fields{ + Two_fields{Float64, Float64}, + Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, }, 3, ) test_get_field_first_index_offset( @name(y.y), Float32, - two_fields{two_fields{Float64, Float64}, two_fields{Float64, Float64}}, + Two_fields{Two_fields{Float64, Float64}, Two_fields{Float64, Float64}}, 6, ) test_get_field_first_index_offset( @name(y.y.x), Float64, - two_fields{ - two_fields{Float64, Float64}, - two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + Two_fields{ + Two_fields{Float64, Float64}, + Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, }, 3, ) test_get_field_first_index_offset( @name(y.y.y.x), Float64, - two_fields{ - two_fields{Float64, Float64}, - two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + Two_fields{ + Two_fields{Float64, Float64}, + Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, }, 4, ) end @testset "broadcasted_get_field_type" begin - struct singleton{T} + struct Singleton{T} x::T end - struct two_fields{T1, T2} + struct Two_fields{T1, T2} x::T1 y::T2 end @@ -98,29 +96,29 @@ end end test_broadcasted_get_field_type( @name(x), - singleton{singleton{singleton{singleton{Float64}}}}, - singleton{singleton{singleton{Float64}}}, + Singleton{Singleton{Singleton{Singleton{Float64}}}}, + Singleton{Singleton{Singleton{Float64}}}, ) test_broadcasted_get_field_type( @name(x.x.x), - singleton{singleton{singleton{singleton{Float64}}}}, - singleton{Float64}, + Singleton{Singleton{Singleton{Singleton{Float64}}}}, + Singleton{Float64}, ) test_broadcasted_get_field_type( @name(y.x), - two_fields{ - two_fields{Float64, Float64}, - two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + Two_fields{ + Two_fields{Float64, Float64}, + Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, }, Float64, ) test_broadcasted_get_field_type( @name(y.y.y), - two_fields{ - two_fields{Float64, Float64}, - two_fields{Float64, two_fields{Float64, singleton{Float64}}}, + Two_fields{ + Two_fields{Float64, Float64}, + Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, }, - singleton{Float64}, + Singleton{Float64}, ) end @@ -129,9 +127,6 @@ end center_space, face_space = test_spaces(FT) surface_space = Spaces.level(face_space, half) seed!(1) - ᶜvec = random_field(FT, center_space) - ᶠvec = random_field(FT, face_space) - sfc_vec = random_field(FT, surface_space) λ = 10 ᶜᶜmat1 = random_field(DiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) ᶜᶠmat2 = random_field(BidiagonalMatrixRow{FT}, center_space) ./ λ @@ -147,80 +142,30 @@ end ρaχ_unit = (; ρaq_tot = 1, ρaq_liq = 1, ρaq_ice = 1, ρaq_rai = 1, ρaq_sno = 1) - dry_center_gs_unit = (; ρ = 1, ρe_tot = 1, uₕ = e¹²) - 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₃',) - ᶜᶜmat3_scalar_uₕ = ᶜᶜmat3 .* (e¹²,) - ᶜᶜmat1_scalar_uₕ = ᶜᶜmat1 .* (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) - - A = 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₃)) => ᶠᶠ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)) => ᶜᶜmat3_ρχ_scalar, - (@name(c.uₕ), @name(c.sgsʲs.:(1).ρa)) => ᶜᶜmat3_scalar_uₕ, - (@name(c.sgsʲs.:(1).ρa), @name(c.uₕ)) => ᶜᶜ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₃)) => ᶜᶠ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₃)) => ᶜᶠ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₃), + 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, ) scalar_keys = MatrixFields.get_scalar_keys(A) + @test length(scalar_keys) > length(keys(A)) @test all( f -> f isa MatrixFields.UniformScaling || eltype(eltype(f)) <: FT, getindex.(Ref(A), scalar_keys), From 5fa5ed5b328bbecce90e8f5dda3b935411772837 Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Thu, 24 Apr 2025 11:34:32 -0700 Subject: [PATCH 03/10] Revert unrolled_findfirst --- src/MatrixFields/field_name_dict.jl | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 1a8f3a5a16..1640d5d030 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -294,7 +294,7 @@ function get_field_first_index_offset( child_type = fieldtype(S, child_name) remaining_field_chain = drop_first(name) field_index = - UnrolledUtilities.unrolled_findfirst(isequal(child_name), fieldnames(S)) + 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 @@ -313,10 +313,7 @@ entries in the `FieldMatrix` `dict`. """ function get_scalar_keys(dict::FieldMatrix) keys_tuple = unrolled_flatmap(keys(dict).values) do key - entry = values(dict)[UnrolledUtilities.unrolled_findfirst( - isequal(key), - keys(dict).values, - )] + entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] entry = entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry unrolled_map( From 01957aa6894056c30d81c505ba8c46243ed7be27 Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Mon, 28 Apr 2025 11:43:17 -0700 Subject: [PATCH 04/10] Clean up field matrix tests and add support for DiagonalMatrixRows --- src/MatrixFields/field_name_dict.jl | 32 ++++-- test/MatrixFields/field_matrix_solvers.jl | 22 +--- test/MatrixFields/matrix_field_test_utils.jl | 67 +++++++++---- test/MatrixFields/scalar_fieldmatrix.jl | 100 ++++--------------- 4 files changed, 92 insertions(+), 129 deletions(-) diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 1640d5d030..d42982bf80 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -152,18 +152,29 @@ 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::ScalingFieldMatrixEntry, +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, name_pair::FieldNamePair, key_error, -) = +) if name_pair[1] == name_pair[2] entry + elseif name_pair[2] == @name() && has_field(entry, name_pair[1]) + DiagonalMatrixRow(get_field(entry, name_pair[1])) elseif is_overlapping_name(name_pair[1], name_pair[2]) throw(key_error) else zero(entry) end +end function get_internal_entry( entry::ColumnwiseBandMatrixField, name_pair::FieldNamePair, @@ -317,12 +328,15 @@ function get_scalar_keys(dict::FieldMatrix) entry = entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry unrolled_map( - filtered_names( - field -> - (field isa UniformScaling) || - eltype(field) == eltype(parent(field)), - entry, - ), + filtered_names(entry) do field + if field isa UniformScaling + return true + elseif field isa Fields.Field + return eltype(field) == eltype(eltype(field)) + else + return eltype(field) == typeof(field) + end + end, ) do name (append_internal_name(key[1], name), key[2]) end diff --git a/test/MatrixFields/field_matrix_solvers.jl b/test/MatrixFields/field_matrix_solvers.jl index d6ea0c7f40..4f65861fd5 100644 --- a/test/MatrixFields/field_matrix_solvers.jl +++ b/test/MatrixFields/field_matrix_solvers.jl @@ -459,7 +459,10 @@ end ), b = b_moist_dycore_diagnostic_edmf, ) - + ( + A_moist_dycore_prognostic_edmf_prognostic_surface, + b_moist_dycore_prognostic_edmf_prognostic_surface, + ) = dycore_prognostic_EDMF_FieldMatrix(FT) test_field_matrix_solver(; test_name = "similar solve to ClimaAtmos's moist dycore + prognostic \ EDMF + prognostic surface temperature with implicit \ @@ -473,22 +476,7 @@ end n_iters = 6, ), ), - 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, - ), + A = A_moist_dycore_prognostic_edmf_prognostic_surface, b = b_moist_dycore_prognostic_edmf_prognostic_surface, ) end diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl index 878603c490..989abaf8ab 100644 --- a/test/MatrixFields/matrix_field_test_utils.jl +++ b/test/MatrixFields/matrix_field_test_utils.jl @@ -140,29 +140,53 @@ 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, -) +# also returns corresponding FieldVector +function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} + seed!(1) # For reproducibility with random fields + center_space, face_space = test_spaces(FT) + surface_space = Spaces.level(face_space, half) + surface_space = Spaces.level(face_space, half) + sfc_vec = random_field(FT, surface_space) + ᶜvec = random_field(FT, center_space) + ᶠvec = random_field(FT, face_space) + seed!(1) + λ = 10 + ᶜᶜmat1 = random_field(DiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) + ᶜᶠmat2 = random_field(BidiagonalMatrixRow{FT}, center_space) ./ λ + ᶠᶜmat2 = random_field(BidiagonalMatrixRow{FT}, face_space) ./ λ + ᶜᶜmat3 = random_field(TridiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) + ᶠᶠmat3 = random_field(TridiagonalMatrixRow{FT}, face_space) ./ λ .+ (I,) + + e¹² = Geometry.Covariant12Vector(1, 1) + e³ = Geometry.Covariant3Vector(1) + e₃ = Geometry.Contravariant3Vector(1) + + ρχ_unit = (; ρq_tot = 1, ρq_liq = 1, ρq_ice = 1, ρq_rai = 1, ρq_sno = 1) + ρaχ_unit = + (; ρaq_tot = 1, ρaq_liq = 1, ρaq_ice = 1, ρaq_rai = 1, ρaq_sno = 1) - ᶜᶜmat3_uₕ_scalar = ᶜᶜmat3 .* (e¹²,) + + ᶠᶜmat2_u₃_scalar = ᶠᶜmat2 .* (e³,) + ᶜᶠmat2_scalar_u₃ = ᶜᶠmat2 .* (e₃',) + ᶠᶠmat3_u₃_u₃ = ᶠᶠmat3 .* (e³ * e₃',) + ᶜᶠmat2_ρχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit ⊠ e₃')), ᶜᶠmat2) + ᶜᶜmat3_uₕ_scalar = + DiagonalMatrixRow(Geometry.Covariant12Vector(FT(1), FT(1))) ᶜᶠ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( + + dry_center_gs_unit = (; ρ = 1, ρe_tot = 1, uₕ = e¹²) + center_gs_unit = (; dry_center_gs_unit..., ρatke = 1, ρχ = ρχ_unit) + center_sgsʲ_unit = (; ρa = 1, ρae_tot = 1, ρaχ = ρaχ_unit) + + b = Fields.FieldVector(; + sfc = sfc_vec .* ((; T = 1),), + c = ᶜvec .* ((; center_gs_unit..., sgsʲs = (center_sgsʲ_unit,)),), + f = ᶠvec .* ((; u₃ = e³, sgsʲs = ((; u₃ = e³),)),), + ) + A = MatrixFields.FieldMatrix( # GS-GS blocks: (@name(sfc), @name(sfc)) => I, (@name(c.ρ), @name(c.ρ)) => I, @@ -193,8 +217,10 @@ function dycore_prognostic_EDMF_FieldMatrix(; (@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.ρ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), @@ -216,6 +242,7 @@ function dycore_prognostic_EDMF_FieldMatrix(; (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => deepcopy(ᶠᶠmat3_u₃_u₃), ) + return A, b end # Generate extruded finite difference spaces for testing. Include topography # when possible. diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl index e6e5e63000..c2ef87ef15 100644 --- a/test/MatrixFields/scalar_fieldmatrix.jl +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -124,95 +124,29 @@ end @testset "fieldmatrix to scalar fieldmatrix unit tests" begin FT = Float64 - center_space, face_space = test_spaces(FT) - surface_space = Spaces.level(face_space, half) - seed!(1) - λ = 10 - ᶜᶜmat1 = random_field(DiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) - ᶜᶠmat2 = random_field(BidiagonalMatrixRow{FT}, center_space) ./ λ - ᶠᶜmat2 = random_field(BidiagonalMatrixRow{FT}, face_space) ./ λ - ᶜᶜmat3 = random_field(TridiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) - ᶠᶠmat3 = random_field(TridiagonalMatrixRow{FT}, face_space) ./ λ .+ (I,) - - e¹² = Geometry.Covariant12Vector(1, 2) - e³ = Geometry.Covariant3Vector(1) - e₃ = Geometry.Contravariant3Vector(1) - - ρχ_unit = (; ρq_tot = 1, ρq_liq = 1, ρq_ice = 1, ρq_rai = 1, ρq_sno = 1) - ρaχ_unit = - (; ρaq_tot = 1, ρaq_liq = 1, ρaq_ice = 1, ρaq_rai = 1, ρaq_sno = 1) - - - ᶠᶜmat2_u₃_scalar = ᶠᶜmat2 .* (e³,) - ᶜᶠmat2_scalar_u₃ = ᶜᶠmat2 .* (e₃',) - ᶠᶠmat3_u₃_u₃ = ᶠᶠmat3 .* (e³ * e₃',) - ᶜᶠmat2_ρχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit ⊠ e₃')), ᶜᶠmat2) - 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, - ) - - scalar_keys = MatrixFields.get_scalar_keys(A) - @test length(scalar_keys) > length(keys(A)) - @test all( - f -> f isa MatrixFields.UniformScaling || eltype(eltype(f)) <: FT, - getindex.(Ref(A), scalar_keys), - ) - @test (@allocated MatrixFields.get_scalar_keys(A)) == 0 - @test_opt MatrixFields.get_scalar_keys(A) - foreach(scalar_keys) do key - @test_opt A[key] - @test (@allocated A[key]) == 0 + A, _ = dycore_prognostic_EDMF_FieldMatrix(FT) + for key in MatrixFields.get_scalar_keys(A) + @test_all A[key] isa MatrixFields.ColumnwiseBandMatrixField ? + eltype(eltype(A[key])) == eltype(parent(A[key])) : + eltype(eltype(A[key])) == eltype(A[key]) end - - A_scalar = MatrixFields.scalar_fieldmatrix(A) @test all( - f -> f isa MatrixFields.UniformScaling || eltype(eltype(f)) <: FT, - getindex.(Ref(A_scalar), scalar_keys), + entry -> + entry isa MatrixFields.UniformScaling || + eltype(eltype(entry)) <: FT, + MatrixFields.scalar_fieldmatrix(A).entries, ) - @test all(scalar_keys) do key - entry = A_scalar[key] - entry isa MatrixFields.UniformScaling && return true - key′, entry′ = - filter(pair -> MatrixFields.is_child_value(key, pair[1]), pairs(A))[1] - if key != key′ - parent(entry) isa SubArray || - parent(entry) === parent(entry′) || - return false - end - parent(parent(entry)) === parent(entry′) || return false - return true + test_get(A, entry, key) = A[key] == entry + for (key, entry) in MatrixFields.scalar_fieldmatrix(A) + @test test_get(A, entry, key) + @test (@allocated test_get(A, entry, key)) == 0 + @test_opt test_get(A, entry, key) end function scalar_fieldmatrix_wrapper(field_matrix_of_tensors) - MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors) - return + A_scalar = MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors) + return true end - scalar_fieldmatrix_wrapper(A) + scalar_fieldmatrix_wrapper(A) # compile the wrapper function @test (@allocated scalar_fieldmatrix_wrapper(A)) == 0 @test_opt MatrixFields.scalar_fieldmatrix(A) - - A_copy = deepcopy(A) - @test all(scalar_keys) do key - entry_copy = A_copy[key] - entry_original = A[key] - entry_flattened = A_scalar[key] - entry_flattened isa MatrixFields.UniformScaling && return true - entry_original == entry_copy || return false - entry_original === entry_flattened || return false - entry_flattened .*= 2 - entry_original == 2 .* entry_copy || return false - return true - end end From 63ac1322626ec0cb1f0f46d7bca77beae7dd3c7a Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Mon, 28 Apr 2025 12:51:19 -0700 Subject: [PATCH 05/10] CamelCase struct name --- test/MatrixFields/scalar_fieldmatrix.jl | 38 ++++++++++++------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl index c2ef87ef15..12d0110096 100644 --- a/test/MatrixFields/scalar_fieldmatrix.jl +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -12,7 +12,7 @@ include("matrix_field_test_utils.jl") struct Singleton{T} x::T end - struct Two_fields{T1, T2} + struct TwoFields{T1, T2} x::T1 y::T2 end @@ -40,39 +40,39 @@ include("matrix_field_test_utils.jl") test_get_field_first_index_offset( @name(y.x), Float64, - Two_fields{Two_fields{Float64, Float64}, Two_fields{Float64, Float64}}, + TwoFields{TwoFields{Float64, Float64}, TwoFields{Float64, Float64}}, 2, ) test_get_field_first_index_offset( @name(y.y), Float64, - Two_fields{ - Two_fields{Float64, Float64}, - Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, + TwoFields{ + TwoFields{Float64, Float64}, + TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, }, 3, ) test_get_field_first_index_offset( @name(y.y), Float32, - Two_fields{Two_fields{Float64, Float64}, Two_fields{Float64, Float64}}, + TwoFields{TwoFields{Float64, Float64}, TwoFields{Float64, Float64}}, 6, ) test_get_field_first_index_offset( @name(y.y.x), Float64, - Two_fields{ - Two_fields{Float64, Float64}, - Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, + TwoFields{ + TwoFields{Float64, Float64}, + TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, }, 3, ) test_get_field_first_index_offset( @name(y.y.y.x), Float64, - Two_fields{ - Two_fields{Float64, Float64}, - Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, + TwoFields{ + TwoFields{Float64, Float64}, + TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, }, 4, ) @@ -82,7 +82,7 @@ end struct Singleton{T} x::T end - struct Two_fields{T1, T2} + struct TwoFields{T1, T2} x::T1 y::T2 end @@ -106,17 +106,17 @@ end ) test_broadcasted_get_field_type( @name(y.x), - Two_fields{ - Two_fields{Float64, Float64}, - Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, + TwoFields{ + TwoFields{Float64, Float64}, + TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, }, Float64, ) test_broadcasted_get_field_type( @name(y.y.y), - Two_fields{ - Two_fields{Float64, Float64}, - Two_fields{Float64, Two_fields{Float64, Singleton{Float64}}}, + TwoFields{ + TwoFields{Float64, Float64}, + TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, }, Singleton{Float64}, ) From 8eada4541b7b6743353d6fa52f73cdc0a5f07825 Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Mon, 28 Apr 2025 15:33:36 -0700 Subject: [PATCH 06/10] Clean up tests and get_scalar_keys --- src/MatrixFields/field_name_dict.jl | 20 +++--- test/MatrixFields/matrix_field_test_utils.jl | 75 +++++++++----------- test/MatrixFields/scalar_fieldmatrix.jl | 56 ++++++++------- 3 files changed, 71 insertions(+), 80 deletions(-) diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index d42982bf80..9875e0e9c6 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -327,17 +327,15 @@ function get_scalar_keys(dict::FieldMatrix) entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] entry = entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry - unrolled_map( - filtered_names(entry) do field - if field isa UniformScaling - return true - elseif field isa Fields.Field - return eltype(field) == eltype(eltype(field)) - else - return eltype(field) == typeof(field) - end - end, - ) do name + unrolled_map(filtered_names(entry) do field + if field isa UniformScaling + true + elseif field isa Fields.Field + eltype(field) == eltype(eltype(field)) + else + eltype(field) == typeof(field) + end + end) do name (append_internal_name(key[1], name), key[2]) end end diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl index 989abaf8ab..579073bc33 100644 --- a/test/MatrixFields/matrix_field_test_utils.jl +++ b/test/MatrixFields/matrix_field_test_utils.jl @@ -190,57 +190,48 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} # 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₃), + (@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)) => 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₃), + (@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₃)) => - deepcopy(ᶜᶠmat2_scalar_u₃), + ᶜᶠ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₃), + ᶜᶠ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)) => - deepcopy(ᶠᶜmat2_u₃_scalar), + ᶠᶜ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₃), + ᶠᶜmat2_u₃_scalar, + (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => ᶠᶠmat3_u₃_u₃, ) return A, b end diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl index 12d0110096..3f030d7482 100644 --- a/test/MatrixFields/scalar_fieldmatrix.jl +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -9,6 +9,7 @@ ClimaComms.@import_required_backends include("matrix_field_test_utils.jl") @testset "get_field_first_index_offset" begin + FT = Float64 struct Singleton{T} x::T end @@ -27,58 +28,59 @@ include("matrix_field_test_utils.jl") end test_get_field_first_index_offset( @name(x), - Float64, - Singleton{Singleton{Singleton{Singleton{Float64}}}}, + FT, + Singleton{Singleton{Singleton{Singleton{FT}}}}, 0, ) test_get_field_first_index_offset( @name(x.x.x.x), - Float64, - Singleton{Singleton{Singleton{Singleton{Float64}}}}, + FT, + Singleton{Singleton{Singleton{Singleton{FT}}}}, 0, ) test_get_field_first_index_offset( @name(y.x), - Float64, - TwoFields{TwoFields{Float64, Float64}, TwoFields{Float64, Float64}}, + FT, + TwoFields{TwoFields{FT, FT}, TwoFields{FT, FT}}, 2, ) test_get_field_first_index_offset( @name(y.y), - Float64, + FT, TwoFields{ - TwoFields{Float64, Float64}, - TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, + TwoFields{FT, FT}, + TwoFields{FT, TwoFields{FT, Singleton{FT}}}, }, 3, ) test_get_field_first_index_offset( @name(y.y), Float32, - TwoFields{TwoFields{Float64, Float64}, TwoFields{Float64, Float64}}, + TwoFields{TwoFields{FT, FT}, TwoFields{FT, FT}}, 6, ) test_get_field_first_index_offset( @name(y.y.x), - Float64, + FT, TwoFields{ - TwoFields{Float64, Float64}, - TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, + TwoFields{FT, FT}, + TwoFields{FT, TwoFields{FT, Singleton{FT}}}, }, 3, ) test_get_field_first_index_offset( @name(y.y.y.x), - Float64, + FT, TwoFields{ - TwoFields{Float64, Float64}, - TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, + TwoFields{FT, FT}, + TwoFields{FT, TwoFields{FT, Singleton{FT}}}, }, 4, ) end @testset "broadcasted_get_field_type" begin + FT = Float64 struct Singleton{T} x::T end @@ -96,29 +98,29 @@ end end test_broadcasted_get_field_type( @name(x), - Singleton{Singleton{Singleton{Singleton{Float64}}}}, - Singleton{Singleton{Singleton{Float64}}}, + Singleton{Singleton{Singleton{Singleton{FT}}}}, + Singleton{Singleton{Singleton{FT}}}, ) test_broadcasted_get_field_type( @name(x.x.x), - Singleton{Singleton{Singleton{Singleton{Float64}}}}, - Singleton{Float64}, + Singleton{Singleton{Singleton{Singleton{FT}}}}, + Singleton{FT}, ) test_broadcasted_get_field_type( @name(y.x), TwoFields{ - TwoFields{Float64, Float64}, - TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, + TwoFields{FT, FT}, + TwoFields{FT, TwoFields{FT, Singleton{FT}}}, }, - Float64, + FT, ) test_broadcasted_get_field_type( @name(y.y.y), TwoFields{ - TwoFields{Float64, Float64}, - TwoFields{Float64, TwoFields{Float64, Singleton{Float64}}}, + TwoFields{FT, FT}, + TwoFields{FT, TwoFields{FT, Singleton{FT}}}, }, - Singleton{Float64}, + Singleton{FT}, ) end @@ -136,7 +138,7 @@ end eltype(eltype(entry)) <: FT, MatrixFields.scalar_fieldmatrix(A).entries, ) - test_get(A, entry, key) = A[key] == entry + test_get(A, entry, key) = A[key] === entry for (key, entry) in MatrixFields.scalar_fieldmatrix(A) @test test_get(A, entry, key) @test (@allocated test_get(A, entry, key)) == 0 From 9af2119cdf74358ba13d47c6138a0153514b0017 Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Fri, 2 May 2025 17:10:50 -0700 Subject: [PATCH 07/10] wip backup --- src/MatrixFields/field_name.jl | 4 + src/MatrixFields/field_name_dict.jl | 113 ++++++++++++++++--- test/MatrixFields/matrix_field_test_utils.jl | 6 +- 3 files changed, 107 insertions(+), 16 deletions(-) diff --git a/src/MatrixFields/field_name.jl b/src/MatrixFields/field_name.jl index 3b66470292..42d36769b6 100644 --- a/src/MatrixFields/field_name.jl +++ b/src/MatrixFields/field_name.jl @@ -50,6 +50,10 @@ 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} = last(name_chain) +drop_last(::FieldName{name_chain}) where {name_chain} = + FieldName(name_chain[1:end-1]...) + has_field(x, ::FieldName{()}) = true has_field(x, name::FieldName) = extract_first(name) in propertynames(x) && diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 9875e0e9c6..6889f31b0e 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -185,10 +185,9 @@ function get_internal_entry( T = eltype(eltype(entry)) if name_pair == (@name(), @name()) entry - elseif name_pair[1] == name_pair[2] + elseif name_pair[1] == name_pair[2] && !broadcasted_has_field(T, name_pair[1]) # multiplication case 3 or 4, first argument - @assert T <: Geometry.SingleValue && - !broadcasted_has_field(T, name_pair[1]) + @assert T <: Geometry.SingleValue entry elseif name_pair[2] == @name() && broadcasted_has_field(T, name_pair[1]) # multiplication case 2 or 4, second argument @@ -236,6 +235,48 @@ function get_internal_entry( end end # Note: This assumes that the entry is in a FieldMatrixBroadcasted. end + elseif broadcasted_has_field(T, name_pair[1]) && broadcasted_has_field(T, name_pair[2]) + # this should only be the case when both independent and dependent var are axisvectors + @assert T <: Geometry.SingleValue && !(T <: Number) + @assert drop_last(name_pair[1]) == drop_last(name_pair[2]) == @name(components.data) + row_index = extract_last(name_pair[1]) + col_index = extract_last(name_pair[2]) + (n_rows, n_cols) = map(length, axes(T)) + @assert row_index <= n_rows && col_index <= n_cols + flattened_index = n_rows * (col_index - 1) + row_index + band_element_size = div(sizeof(T), sizeof(eltype(T))) + 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 = band_matrix_row_type( + outer_diagonals(T_band)..., + eltype(T), + ) + field_dim_size = DataLayouts.ncomponents(Fields.field_values(entry)) + band_element_size = div(sizeof(T), sizeof(eltype(T))) + @assert band_element_size == n_rows * n_cols + parent_indices = DataLayouts.to_data_specific_field( + singleton_datalayout, + ( + :, + :, + flattened_index:band_element_size:field_dim_size, + :, + :, + ), + ) + # Main.@infiltrate + 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 throw(key_error) end @@ -322,26 +363,70 @@ end Returns a `FieldMatrixKeys` object that contains the keys of all the scalar entries in the `FieldMatrix` `dict`. """ -function get_scalar_keys(dict::FieldMatrix) +function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) keys_tuple = unrolled_flatmap(keys(dict).values) do key entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] - entry = - entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry - unrolled_map(filtered_names(entry) do field - if field isa UniformScaling - true - elseif field isa Fields.Field - eltype(field) == eltype(eltype(field)) + if entry isa UniformScaling # uniformscalings can only contain numbers + (key,) + elseif entry isa ColumnwiseBandMatrixField + first_band = entry.entries.:(1) + target_eltype = eltype(parent(first_band)) + if eltype(first_band) == target_eltype + (key,) else - eltype(field) == typeof(field) + dependent_var = get_field(Y, key[1]) + independent_var = get_field(Y, key[2]) + dependent_type = eltype(dependent_var) + independent_type = eltype(independent_var) + # @Main.infiltrate + @assert dependent_type <: Geometry.SingleValue || + independent_type <: Geometry.SingleValue || + "cannot get scalar keys for key $key" + + # figure out if we need to drill into key[1] or key[2], or both + # @show key + unrolled_flatmap(filtered_names(x -> eltype(x) <: target_eltype, dependent_var)) do dependent_name + unrolled_map(filtered_names(x -> eltype(x) <: target_eltype, independent_var)) do independent_name + (append_internal_name(key[1], dependent_name), append_internal_name(key[2], independent_name)) + end + # @Main.infiltrate + # key + end + # (key,) end - end) do name - (append_internal_name(key[1], name), key[2]) + else + # TODO: Fix me + (key,) end + + # entry = + # entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry + # unrolled_map(filtered_names(entry) do field + # if field isa UniformScaling + # true + # elseif field isa Fields.Field + # eltype(field) == eltype(eltype(field)) + # else + # eltype(field) == typeof(field) + # end + # end) do name + # (append_internal_name(key[1], name), key[2]) + # end end return FieldMatrixKeys(keys_tuple) end +function new_get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) + scalar_field_vector_keys = MatrixFields.filtered_names(Y) do field + field isa Fields.Field && eltype(field) == eltype(parent(field)) + end + map(keys(dict).values) do key + first_key_is_scalar = unrolled_any(isequal(key[1]), scalar_field_vector_keys) + second_key_is_scalar = unrolled_any(isequal(key[2]), scalar_field_vector_keys) + @assert first_key_is_scalar || second_key_is_scalar "$key" + end +end + """ scalar_fieldmatrix(field_matrix::FieldMatrix) diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl index 579073bc33..69ed05e449 100644 --- a/test/MatrixFields/matrix_field_test_utils.jl +++ b/test/MatrixFields/matrix_field_test_utils.jl @@ -156,7 +156,7 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} ᶠᶜmat2 = random_field(BidiagonalMatrixRow{FT}, face_space) ./ λ ᶜᶜmat3 = random_field(TridiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) ᶠᶠmat3 = random_field(TridiagonalMatrixRow{FT}, face_space) ./ λ .+ (I,) - + # Geometry.Covariant123Vector(1, 2, 3) * Geometry.Covariant12Vector(1, 2)' e¹² = Geometry.Covariant12Vector(1, 1) e³ = Geometry.Covariant3Vector(1) e₃ = Geometry.Contravariant3Vector(1) @@ -172,6 +172,7 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} ᶜᶠmat2_ρχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit ⊠ e₃')), ᶜᶠmat2) ᶜᶜmat3_uₕ_scalar = DiagonalMatrixRow(Geometry.Covariant12Vector(FT(1), FT(1))) + ᶜᶜmat3_uₕ_uₕ = ᶜᶜmat3 .* (e¹² * 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) @@ -189,11 +190,12 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} A = MatrixFields.FieldMatrix( # GS-GS blocks: (@name(sfc), @name(sfc)) => I, + (@name(sfc), @name(c.uₕ)) => ᶜᶜmat3 .* (Geometry.Covariant123Vector(1, 2, 3) * Geometry.Covariant12Vector(1, 2)',), (@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.uₕ), @name(c.uₕ)) => ᶜᶜmat3_uₕ_uₕ, (@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₃, From e1c1c223f056b09b85011da1800019667397c654 Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Thu, 8 May 2025 15:34:54 -0700 Subject: [PATCH 08/10] Minimal working with allocs --- src/MatrixFields/field_name.jl | 6 +- src/MatrixFields/field_name_dict.jl | 273 +++++++++++++------ test/MatrixFields/matrix_field_test_utils.jl | 13 +- test/MatrixFields/scalar_fieldmatrix.jl | 18 +- 4 files changed, 215 insertions(+), 95 deletions(-) diff --git a/src/MatrixFields/field_name.jl b/src/MatrixFields/field_name.jl index 42d36769b6..d3f0549da0 100644 --- a/src/MatrixFields/field_name.jl +++ b/src/MatrixFields/field_name.jl @@ -50,9 +50,9 @@ 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} = last(name_chain) -drop_last(::FieldName{name_chain}) where {name_chain} = - FieldName(name_chain[1:end-1]...) +extract_last(::FieldName{name_chain}) where {name_chain} = name_chain[length(name_chain)] +# drop_last(::FieldName{name_chain}) where {name_chain} = +# FieldName(name_chain[1:(end - 1)]...) has_field(x, ::FieldName{()}) = true has_field(x, name::FieldName) = diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 6889f31b0e..9e860ff079 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -165,10 +165,14 @@ function get_internal_entry( name_pair::FieldNamePair, key_error, ) + # TODO: Add other cases if name_pair[1] == name_pair[2] entry - elseif name_pair[2] == @name() && has_field(entry, name_pair[1]) - DiagonalMatrixRow(get_field(entry, name_pair[1])) + elseif name_pair[2] == @name() && + broadcasted_has_field(eltype(entry), name_pair[1]) + DiagonalMatrixRow( + broadcasted_get_field(entry.entries.:(1), name_pair[1]), + ) elseif is_overlapping_name(name_pair[1], name_pair[2]) throw(key_error) else @@ -185,13 +189,32 @@ function get_internal_entry( T = eltype(eltype(entry)) if name_pair == (@name(), @name()) entry - elseif name_pair[1] == name_pair[2] && !broadcasted_has_field(T, name_pair[1]) - # multiplication case 3 or 4, first argument - @assert T <: Geometry.SingleValue + elseif name_pair[1] == name_pair[2] && + !broadcasted_has_field(T, name_pair[1]) + # @show "aa" + # # multiplication case 3 or 4, first argument + @assert T <: Number 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]) + elseif name_pair[1] == @name() || name_pair[2] == @name() + + target_chain = if name_pair[1] == @name() + if broadcasted_has_field(T, name_pair[2]) + # this case should be dscalar/dvec with T isa vec + name_pair[2] + else + # this should be dscalar/dvec with T isa adjoint + append_internal_name(@name(parent), name_pair[2]) + end + else + if broadcasted_has_field(T, name_pair[1]) + # this case should be dtuple/dscalar or dvec/dscalar with T isa vec + name_pair[1] + else + # this should be dvec/dscalar with T isa adjoint + append_internal_name(@name(parent), name_pair[1]) + end + end + target_field_eltype = broadcasted_get_field_type(T, target_chain) if target_field_eltype == eltype(parent(entry)) T_band = eltype(entry) singleton_datalayout = @@ -203,7 +226,7 @@ function get_internal_entry( ) field_dim_size = DataLayouts.ncomponents(Fields.field_values(entry)) scalar_field_offset = get_field_first_index_offset( - name_pair[1], + target_chain, target_field_eltype, T, ) @@ -231,57 +254,131 @@ function get_internal_entry( else Base.broadcasted(entry) do matrix_row map(matrix_row) do matrix_row_entry - broadcasted_get_field(matrix_row_entry, name_pair[1]) + broadcasted_get_field(matrix_row_entry, target_chain) end - end # Note: This assumes that the entry is in a FieldMatrixBroadcasted. + end + end + elseif name_pair[2] != @name() && name_pair[1] != @name() + # this should only be the case with dvec/dvec or dNTuple/dvec + if T <: Geometry.SingleValue + # @assert drop_last(name_pair[1]) == + # drop_last(name_pair[2]) == + # @name(components.data) + row_index = extract_last(name_pair[1]) + col_index = extract_last(name_pair[2]) + (n_rows, n_cols) = map(length, axes(T)) + @assert row_index <= n_rows && col_index <= n_cols + flattened_index = n_rows * (col_index - 1) + row_index + elseif eltype(T) <: Geometry.SingleValue #TODO: nested tuples? + # @assert drop_last(name_pair[2]) == @name(components.data) + modified_first_name = + broadcasted_has_field(T, name_pair[1]) ? name_pair[1] : + append_internal_name(@name(parent), name_pair[1]) + flattened_index = + get_field_first_index_offset( + name_pair[1], + broadcasted_get_field_type(T, name_pair[1]), + T, + ) + extract_last(name_pair[2]) + else + error("Cannot get entry for key $name_pair") end - elseif broadcasted_has_field(T, name_pair[1]) && broadcasted_has_field(T, name_pair[2]) - # this should only be the case when both independent and dependent var are axisvectors - @assert T <: Geometry.SingleValue && !(T <: Number) - @assert drop_last(name_pair[1]) == drop_last(name_pair[2]) == @name(components.data) - row_index = extract_last(name_pair[1]) - col_index = extract_last(name_pair[2]) - (n_rows, n_cols) = map(length, axes(T)) - @assert row_index <= n_rows && col_index <= n_cols - flattened_index = n_rows * (col_index - 1) + row_index band_element_size = div(sizeof(T), sizeof(eltype(T))) T_band = eltype(entry) - singleton_datalayout = - DataLayouts.singleton(Fields.field_values(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 = band_matrix_row_type( - outer_diagonals(T_band)..., - eltype(T), - ) + scalar_band_type = + band_matrix_row_type(outer_diagonals(T_band)..., eltype(eltype(T))) field_dim_size = DataLayouts.ncomponents(Fields.field_values(entry)) band_element_size = div(sizeof(T), sizeof(eltype(T))) - @assert band_element_size == n_rows * n_cols parent_indices = DataLayouts.to_data_specific_field( - singleton_datalayout, - ( - :, - :, - flattened_index:band_element_size:field_dim_size, - :, - :, - ), - ) - # Main.@infiltrate + singleton_datalayout, + (:, :, flattened_index: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)) + + 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 throw(key_error) end end +function get_scalar_keys(dict::FieldMatrix) + keys_tuple = unrolled_flatmap(keys(dict).values) do key + entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] + entry = + entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry + unrolled_map(filtered_names(entry) do field + if field isa UniformScaling + true + elseif field isa Fields.Field + eltype(field) == eltype(eltype(field)) + else + eltype(field) == typeof(field) + end + end) do name + (append_internal_name(key[1], name), key[2]) + end + end + return FieldMatrixKeys(keys_tuple) +end +# function combine_name_pair(name_pair::Tuple{FieldName, FieldName{()}}, T) +# end + +# function combine_name_pair(name_pair::Tuple{FieldName{()}, FieldName}, ::Type{T}) where {T} +# T <: NamedTuple && error("Cannot return ") +# # @assert eltype() +# end + +# function combine_name_pair(name_pair::FieldNamePair, T) +# end + +# function foobarbaz(combined_name_chain, T, entry, target_field_eltype) +# band_element_size = div(sizeof(T), sizeof(eltype(T))) +# 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 = band_matrix_row_type( +# outer_diagonals(T_band)..., +# eltype(T), +# ) +# field_dim_size = DataLayouts.ncomponents(Fields.field_values(entry)) +# band_element_size = div(sizeof(T), sizeof(eltype(target_field_eltype))) +# first_index = get_field_first_index_offset( +# combined_name_chain, +# target_field_eltype, +# T, +# ) +# parent_indices = DataLayouts.to_data_specific_field( +# singleton_datalayout, +# ( +# :, +# :, +# first_index:band_element_size:field_dim_size, +# :, +# :, +# ), +# ) +# target_data = view(parent(entry), parent_indices...) +# values = DataLayouts.union_all(singleton_datalayout){ +# target_data, +# Base.tail( +# DataLayouts.type_params(Fields.field_values(entry)), +# )..., +# }( +# scalar_data, +# ) +# Fields.Field(values, axes(entry)) +# 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) @@ -368,10 +465,15 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] if entry isa UniformScaling # uniformscalings can only contain numbers (key,) - elseif entry isa ColumnwiseBandMatrixField + elseif entry isa ColumnwiseBandMatrixField || + entry isa DiagonalMatrixRow first_band = entry.entries.:(1) target_eltype = eltype(parent(first_band)) - if eltype(first_band) == target_eltype + if entry isa ColumnwiseBandMatrixField && + eltype(first_band) <: target_eltype + (key,) + elseif entry isa DiagonalMatrixRow && + typeof(first_band) <: target_eltype (key,) else dependent_var = get_field(Y, key[1]) @@ -380,52 +482,55 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) independent_type = eltype(independent_var) # @Main.infiltrate @assert dependent_type <: Geometry.SingleValue || - independent_type <: Geometry.SingleValue || - "cannot get scalar keys for key $key" + independent_type <: Geometry.SingleValue || + "cannot get scalar keys for key $key" # figure out if we need to drill into key[1] or key[2], or both # @show key - unrolled_flatmap(filtered_names(x -> eltype(x) <: target_eltype, dependent_var)) do dependent_name - unrolled_map(filtered_names(x -> eltype(x) <: target_eltype, independent_var)) do independent_name - (append_internal_name(key[1], dependent_name), append_internal_name(key[2], independent_name)) + unrolled_flatmap( + filtered_names( + x -> eltype(x) <: target_eltype, + dependent_var, + ), + ) do dependent_name + unrolled_map( + filtered_names( + x -> eltype(x) <: target_eltype, + independent_var, + ), + ) do independent_name + ( + append_internal_name(key[1], dependent_name), + append_internal_name(key[2], independent_name), + ) end # @Main.infiltrate # key end # (key,) end + # elseif entry isa DiagonalMatrixRow + # target_eltype = eltype(parent(get_field(Y, key[1]))) + # # TODO: unify target_eltype + # (key,) else - # TODO: Fix me - (key,) + error("Cannot get scalar keys for key $key") end - # entry = - # entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry - # unrolled_map(filtered_names(entry) do field - # if field isa UniformScaling - # true - # elseif field isa Fields.Field - # eltype(field) == eltype(eltype(field)) - # else - # eltype(field) == typeof(field) - # end - # end) do name - # (append_internal_name(key[1], name), key[2]) - # end end return FieldMatrixKeys(keys_tuple) end -function new_get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) - scalar_field_vector_keys = MatrixFields.filtered_names(Y) do field - field isa Fields.Field && eltype(field) == eltype(parent(field)) - end - map(keys(dict).values) do key - first_key_is_scalar = unrolled_any(isequal(key[1]), scalar_field_vector_keys) - second_key_is_scalar = unrolled_any(isequal(key[2]), scalar_field_vector_keys) - @assert first_key_is_scalar || second_key_is_scalar "$key" - end -end +# function new_get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) +# scalar_field_vector_keys = MatrixFields.filtered_names(Y) do field +# field isa Fields.Field && eltype(field) == eltype(parent(field)) +# end +# map(keys(dict).values) do key +# first_key_is_scalar = unrolled_any(isequal(key[1]), scalar_field_vector_keys) +# second_key_is_scalar = unrolled_any(isequal(key[2]), scalar_field_vector_keys) +# @assert first_key_is_scalar || second_key_is_scalar "$key" +# end +# end """ scalar_fieldmatrix(field_matrix::FieldMatrix) @@ -467,6 +572,14 @@ function scalar_fieldmatrix(field_matrix::FieldMatrix) return FieldNameDict(scalar_keys, entries) end +function scalar_fieldmatrix(field_matrix::FieldMatrix, Y::Fields.FieldVector) + scalar_keys = get_scalar_keys(field_matrix, Y) + 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)) @@ -776,8 +889,8 @@ function Base.Broadcast.broadcasted( ) product_value = scaling_value(entry1) * scaling_value(entry2) product_value isa Number ? - UniformScaling(product_value) : - DiagonalMatrixRow(product_value) + (UniformScaling(product_value),) : + (DiagonalMatrixRow(product_value),) elseif entry1 isa ScalingFieldMatrixEntry Base.Broadcast.broadcasted(*, (scaling_value(entry1),), entry2) elseif entry2 isa ScalingFieldMatrixEntry diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl index 69ed05e449..33d2808cad 100644 --- a/test/MatrixFields/matrix_field_test_utils.jl +++ b/test/MatrixFields/matrix_field_test_utils.jl @@ -149,7 +149,6 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} sfc_vec = random_field(FT, surface_space) ᶜvec = random_field(FT, center_space) ᶠvec = random_field(FT, face_space) - seed!(1) λ = 10 ᶜᶜmat1 = random_field(DiagonalMatrixRow{FT}, center_space) ./ λ .+ (I,) ᶜᶠmat2 = random_field(BidiagonalMatrixRow{FT}, center_space) ./ λ @@ -158,6 +157,7 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} ᶠᶠmat3 = random_field(TridiagonalMatrixRow{FT}, face_space) ./ λ .+ (I,) # Geometry.Covariant123Vector(1, 2, 3) * Geometry.Covariant12Vector(1, 2)' e¹² = Geometry.Covariant12Vector(1, 1) + e₁₂ = Geometry.Contravariant12Vector(1, 1) e³ = Geometry.Covariant3Vector(1) e₃ = Geometry.Contravariant3Vector(1) @@ -172,7 +172,14 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} ᶜᶠmat2_ρχ_u₃ = map(Base.Fix1(map, Base.Fix2(⊠, ρχ_unit ⊠ e₃')), ᶜᶠmat2) ᶜᶜmat3_uₕ_scalar = DiagonalMatrixRow(Geometry.Covariant12Vector(FT(1), FT(1))) - ᶜᶜmat3_uₕ_uₕ = ᶜᶜmat3 .* (e¹² * e¹²',) + # ᶜᶜmat3_uₕ_uₕ = ᶜᶜmat3 .* (e¹² * e₁₂',) + ᶜᶜmat3_uₕ_uₕ = + ᶜᶜmat3 .* ( + Geometry.Covariant12Vector(1, 0) * + Geometry.Contravariant12Vector(1, 0)' + + Geometry.Covariant12Vector(0, 1) * + Geometry.Contravariant12Vector(0, 1)', + ) ᶜᶠ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) @@ -190,7 +197,7 @@ function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} A = MatrixFields.FieldMatrix( # GS-GS blocks: (@name(sfc), @name(sfc)) => I, - (@name(sfc), @name(c.uₕ)) => ᶜᶜmat3 .* (Geometry.Covariant123Vector(1, 2, 3) * Geometry.Covariant12Vector(1, 2)',), + # (@name(sfc), @name(c.uₕ)) => ᶜᶜmat3 .* (Geometry.Covariant123Vector(1, 2, 3) * Geometry.Covariant12Vector(1, 2)',), (@name(c.ρ), @name(c.ρ)) => I, (@name(c.ρe_tot), @name(c.ρe_tot)) => ᶜᶜmat3, (@name(c.ρatke), @name(c.ρatke)) => ᶜᶜmat3, diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl index 3f030d7482..9f8bb7263d 100644 --- a/test/MatrixFields/scalar_fieldmatrix.jl +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -126,8 +126,8 @@ end @testset "fieldmatrix to scalar fieldmatrix unit tests" begin FT = Float64 - A, _ = dycore_prognostic_EDMF_FieldMatrix(FT) - for key in MatrixFields.get_scalar_keys(A) + A, b = dycore_prognostic_EDMF_FieldMatrix(FT) + for key in MatrixFields.get_scalar_keys(A, b) @test_all A[key] isa MatrixFields.ColumnwiseBandMatrixField ? eltype(eltype(A[key])) == eltype(parent(A[key])) : eltype(eltype(A[key])) == eltype(A[key]) @@ -136,19 +136,19 @@ end entry -> entry isa MatrixFields.UniformScaling || eltype(eltype(entry)) <: FT, - MatrixFields.scalar_fieldmatrix(A).entries, + MatrixFields.scalar_fieldmatrix(A, b).entries, ) test_get(A, entry, key) = A[key] === entry - for (key, entry) in MatrixFields.scalar_fieldmatrix(A) + for (key, entry) in MatrixFields.scalar_fieldmatrix(A, b) @test test_get(A, entry, key) @test (@allocated test_get(A, entry, key)) == 0 @test_opt test_get(A, entry, key) end - function scalar_fieldmatrix_wrapper(field_matrix_of_tensors) - A_scalar = MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors) + function scalar_fieldmatrix_wrapper(field_matrix_of_tensors, b) + A_scalar = MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors, b) return true end - scalar_fieldmatrix_wrapper(A) # compile the wrapper function - @test (@allocated scalar_fieldmatrix_wrapper(A)) == 0 - @test_opt MatrixFields.scalar_fieldmatrix(A) + scalar_fieldmatrix_wrapper(A, b) # compile the wrapper function + @test (@allocated scalar_fieldmatrix_wrapper(A, b)) == 0 + @test_opt MatrixFields.scalar_fieldmatrix(A, b) end From d519f329750511b4205df5f14c8cf7d1756cdcfb Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Tue, 13 May 2025 11:15:24 -0700 Subject: [PATCH 09/10] WIP1 --- src/MatrixFields/field_name.jl | 3 +- src/MatrixFields/field_name_dict.jl | 267 ++++++++++++------- test/MatrixFields/matrix_field_test_utils.jl | 10 +- test/MatrixFields/scalar_fieldmatrix.jl | 32 ++- 4 files changed, 194 insertions(+), 118 deletions(-) diff --git a/src/MatrixFields/field_name.jl b/src/MatrixFields/field_name.jl index d3f0549da0..39de0e976b 100644 --- a/src/MatrixFields/field_name.jl +++ b/src/MatrixFields/field_name.jl @@ -50,7 +50,8 @@ 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)] +extract_last(::FieldName{name_chain}) where {name_chain} = + name_chain[length(name_chain)] # drop_last(::FieldName{name_chain}) where {name_chain} = # FieldName(name_chain[1:(end - 1)]...) diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 9e860ff079..6f9db8717f 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -191,9 +191,13 @@ function get_internal_entry( entry elseif name_pair[1] == name_pair[2] && !broadcasted_has_field(T, name_pair[1]) - # @show "aa" + # @show name_pair # # multiplication case 3 or 4, first argument - @assert T <: Number + # if T <: Number + # entry + # else + # throw(key_error) + # end entry elseif name_pair[1] == @name() || name_pair[2] == @name() @@ -203,6 +207,9 @@ function get_internal_entry( name_pair[2] else # this should be dscalar/dvec with T isa adjoint + if !(T <: Adjoint) + throw(key_error) + end append_internal_name(@name(parent), name_pair[2]) end else @@ -211,6 +218,9 @@ function get_internal_entry( name_pair[1] else # this should be dvec/dscalar with T isa adjoint + if !(T <: Adjoint) + throw(key_error) + end append_internal_name(@name(parent), name_pair[1]) end end @@ -281,7 +291,13 @@ function get_internal_entry( T, ) + extract_last(name_pair[2]) else - error("Cannot get entry for key $name_pair") + # error("Cannot get entry for key $name_pair") + throw(key_error) + # Base.broadcasted(entry) do matrix_row + # map(matrix_row) do matrix_row_entry + # broadcasted_get_field(matrix_row_entry, target_chain) + # end + # end end band_element_size = div(sizeof(T), sizeof(eltype(T))) T_band = eltype(entry) @@ -310,73 +326,24 @@ function get_internal_entry( end end -function get_scalar_keys(dict::FieldMatrix) - keys_tuple = unrolled_flatmap(keys(dict).values) do key - entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] - entry = - entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry - unrolled_map(filtered_names(entry) do field - if field isa UniformScaling - true - elseif field isa Fields.Field - eltype(field) == eltype(eltype(field)) - else - eltype(field) == typeof(field) - end - end) do name - (append_internal_name(key[1], name), key[2]) - end - end - return FieldMatrixKeys(keys_tuple) -end -# function combine_name_pair(name_pair::Tuple{FieldName, FieldName{()}}, T) -# end - -# function combine_name_pair(name_pair::Tuple{FieldName{()}, FieldName}, ::Type{T}) where {T} -# T <: NamedTuple && error("Cannot return ") -# # @assert eltype() -# end - -# function combine_name_pair(name_pair::FieldNamePair, T) -# end - -# function foobarbaz(combined_name_chain, T, entry, target_field_eltype) -# band_element_size = div(sizeof(T), sizeof(eltype(T))) -# 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 = band_matrix_row_type( -# outer_diagonals(T_band)..., -# eltype(T), -# ) -# field_dim_size = DataLayouts.ncomponents(Fields.field_values(entry)) -# band_element_size = div(sizeof(T), sizeof(eltype(target_field_eltype))) -# first_index = get_field_first_index_offset( -# combined_name_chain, -# target_field_eltype, -# T, -# ) -# parent_indices = DataLayouts.to_data_specific_field( -# singleton_datalayout, -# ( -# :, -# :, -# first_index:band_element_size:field_dim_size, -# :, -# :, -# ), -# ) -# target_data = view(parent(entry), parent_indices...) -# values = DataLayouts.union_all(singleton_datalayout){ -# target_data, -# Base.tail( -# DataLayouts.type_params(Fields.field_values(entry)), -# )..., -# }( -# scalar_data, -# ) -# Fields.Field(values, axes(entry)) +# function get_scalar_keys(dict::FieldMatrix) +# keys_tuple = unrolled_flatmap(keys(dict).values) do key +# entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] +# entry = +# entry isa ColumnwiseBandMatrixField ? entry.entries.:(1) : entry +# unrolled_map(filtered_names(entry) do field +# if field isa UniformScaling +# true +# elseif field isa Fields.Field +# eltype(field) == eltype(eltype(field)) +# else +# eltype(field) == typeof(field) +# end +# end) do name +# (append_internal_name(key[1], name), key[2]) +# end +# end +# return FieldMatrixKeys(keys_tuple) # end # Similar behavior to indexing an array with a slice. @@ -454,6 +421,24 @@ if hasfield(Method, :recursion_relation) end end +function bypass_adjoint(field::Fields.Field) + if eltype(field) <: Adjoint + return (eltype(field.parent), true) + else + return (eltype(field), false) + end +end + +function bypass_adjoint(d::T) where {T} + if T <: Adjoint + return (typeof(d.parent), true) + else + return (typeof(d), false) + end +end + + + """ get_scalar_keys(dict::FieldMatrix) @@ -461,6 +446,7 @@ Returns a `FieldMatrixKeys` object that contains the keys of all the scalar entries in the `FieldMatrix` `dict`. """ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) + target_eltype = eltype(Y) keys_tuple = unrolled_flatmap(keys(dict).values) do key entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] if entry isa UniformScaling # uniformscalings can only contain numbers @@ -468,44 +454,131 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) elseif entry isa ColumnwiseBandMatrixField || entry isa DiagonalMatrixRow first_band = entry.entries.:(1) - target_eltype = eltype(parent(first_band)) - if entry isa ColumnwiseBandMatrixField && - eltype(first_band) <: target_eltype - (key,) - elseif entry isa DiagonalMatrixRow && - typeof(first_band) <: target_eltype + # if entry isa DiagonalMatrixRow, then first_band is not a field + get_type = entry isa ColumnwiseBandMatrixField ? eltype : typeof + + if get_type(first_band) <: target_eltype (key,) else + dependent_var = get_field(Y, key[1]) independent_var = get_field(Y, key[2]) dependent_type = eltype(dependent_var) independent_type = eltype(independent_var) - # @Main.infiltrate + @assert dependent_type <: Geometry.SingleValue || independent_type <: Geometry.SingleValue || "cannot get scalar keys for key $key" - # figure out if we need to drill into key[1] or key[2], or both # @show key - unrolled_flatmap( - filtered_names( - x -> eltype(x) <: target_eltype, - dependent_var, - ), - ) do dependent_name + # TODO: Make the cth axis check for all combinations + if independent_type <: Geometry.AxisTensor && + dependent_type <: Geometry.AxisTensor + (unwrapped_row_eltype, is_adjoint) = + bypass_adjoint(first_band) + @assert unwrapped_row_eltype <: Geometry.Axis2Tensor + axis_components = + map(length, axes(unwrapped_row_eltype)) unrolled_map( - filtered_names( - x -> eltype(x) <: target_eltype, - independent_var, + MatrixFields.unrolled_product( + 1:axis_components[1], + 1:axis_components[2], ), - ) do independent_name + ) do (component1, component2) ( - append_internal_name(key[1], dependent_name), - append_internal_name(key[2], independent_name), + append_internal_name( + key[1], + append_internal_name( + @name(components.data), + MatrixFields.FieldName(component1), + ), + ), + append_internal_name( + key[2], + append_internal_name( + @name(components.data), + MatrixFields.FieldName(component2), + ), + ), ) end - # @Main.infiltrate - # key + elseif ( + dependent_type <: Number && + independent_type <: Geometry.AxisTensor + ) || ( + independent_type <: Number && + dependent_type <: Geometry.AxisTensor + ) + (unwrapped_row_eltype, is_adjoint) = + bypass_adjoint(first_band) + @assert unwrapped_row_eltype <: Geometry.AxisVector + ncomponents = length(axes(unwrapped_row_eltype)[1]) + unrolled_map(1:ncomponents) do component + if dependent_type <: Geometry.AxisTensor + ( + append_internal_name( + key[1], + append_internal_name( + @name(components.data), + MatrixFields.FieldName(component), + ), + ), + key[2], + ) + else + ( + key[1], + append_internal_name( + key[2], + append_internal_name( + @name(components.data), + MatrixFields.FieldName(component), + ), + ), + ) + end + end + + elseif dependent_type <: Union{NamedTuple, Tuple} && + independent_type <: Geometry.AxisVector + # @assert unwrapped_row_eltype <: Geometry.AxisVector + unrolled_flatmap( + filtered_names( + x -> get_type(x) <: Geometry.SingleValue, + first_band, + ), + ) do dependent_name + (unwrapped_row_eltype, is_adjoint) = bypass_adjoint( + get_field(first_band, dependent_name), + ) + ncomponents = length(axes(unwrapped_row_eltype)[1]) + unrolled_map(1:ncomponents) do component + ( + append_internal_name(key[1], dependent_name), + append_internal_name( + key[2], + append_internal_name( + @name(components.data), + MatrixFields.FieldName(component), + ), + ), + ) + end + end + # + elseif dependent_type <: Union{NamedTuple, Tuple} && + independent_type <: Number + unrolled_map( + filtered_names( + x -> eltype(x) <: target_eltype, + dependent_var, + ), + ) do dependent_name + + (append_internal_name(key[1], dependent_name), key[2]) + + end + end # (key,) end @@ -521,16 +594,6 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) return FieldMatrixKeys(keys_tuple) end -# function new_get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) -# scalar_field_vector_keys = MatrixFields.filtered_names(Y) do field -# field isa Fields.Field && eltype(field) == eltype(parent(field)) -# end -# map(keys(dict).values) do key -# first_key_is_scalar = unrolled_any(isequal(key[1]), scalar_field_vector_keys) -# second_key_is_scalar = unrolled_any(isequal(key[2]), scalar_field_vector_keys) -# @assert first_key_is_scalar || second_key_is_scalar "$key" -# end -# end """ scalar_fieldmatrix(field_matrix::FieldMatrix) diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl index 33d2808cad..d1b14cbf37 100644 --- a/test/MatrixFields/matrix_field_test_utils.jl +++ b/test/MatrixFields/matrix_field_test_utils.jl @@ -141,9 +141,15 @@ 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 # also returns corresponding FieldVector -function dycore_prognostic_EDMF_FieldMatrix(::Type{FT}) where {FT} +function dycore_prognostic_EDMF_FieldMatrix( + ::Type{FT}, + center_space = nothing, + face_space = nothing, +) where {FT} seed!(1) # For reproducibility with random fields - center_space, face_space = test_spaces(FT) + if isnothing(center_space) || isnothing(face_space) + center_space, face_space = test_spaces(FT) + end surface_space = Spaces.level(face_space, half) surface_space = Spaces.level(face_space, half) sfc_vec = random_field(FT, surface_space) diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl index 9f8bb7263d..07ddedfdc6 100644 --- a/test/MatrixFields/scalar_fieldmatrix.jl +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -1,7 +1,8 @@ using Test using JET -import ClimaCore: Geometry, Domains, Meshes, Spaces, Fields, MatrixFields +import ClimaCore: + Geometry, Domains, Meshes, Spaces, Fields, MatrixFields, CommonSpaces import ClimaCore.Utilities: half import ClimaComms import ClimaCore.MatrixFields: @name @@ -128,9 +129,11 @@ end FT = Float64 A, b = dycore_prognostic_EDMF_FieldMatrix(FT) for key in MatrixFields.get_scalar_keys(A, b) - @test_all A[key] isa MatrixFields.ColumnwiseBandMatrixField ? - eltype(eltype(A[key])) == eltype(parent(A[key])) : - eltype(eltype(A[key])) == eltype(A[key]) + # @show key + # TODO: test_all + @test A[key] isa MatrixFields.ColumnwiseBandMatrixField ? + eltype(eltype(A[key])) == eltype(parent(A[key])) : + eltype(eltype(A[key])) == eltype(A[key]) end @test all( entry -> @@ -141,14 +144,17 @@ end test_get(A, entry, key) = A[key] === entry for (key, entry) in MatrixFields.scalar_fieldmatrix(A, b) @test test_get(A, entry, key) - @test (@allocated test_get(A, entry, key)) == 0 - @test_opt test_get(A, entry, key) + # @test (@allocated test_get(A, entry, key)) == 0 + # @test_opt test_get(A, entry, key) end - function scalar_fieldmatrix_wrapper(field_matrix_of_tensors, b) - A_scalar = MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors, b) - return true - end - scalar_fieldmatrix_wrapper(A, b) # compile the wrapper function - @test (@allocated scalar_fieldmatrix_wrapper(A, b)) == 0 - @test_opt MatrixFields.scalar_fieldmatrix(A, b) + # function scalar_fieldmatrix_wrapper(field_matrix_of_tensors, b) + # A_scalar = MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors, b) + # return nothing + # end + # scalar_fieldmatrix_wrapper(A, b) + # scalar_fieldmatrix_wrapper(A, b) # compile the wrapper function + # al = @allocated scalar_fieldmatrix_wrapper(A, b) + # @show al + # @test (@allocated scalar_fieldmatrix_wrapper(A, b)) == 0 + # @test_opt MatrixFields.scalar_fieldmatrix(A, b) end From 62ceaa8fe1ea41cdb09ad05d5161b1e88a962ab6 Mon Sep 17 00:00:00 2001 From: imreddyTeja Date: Thu, 15 May 2025 17:31:29 -0700 Subject: [PATCH 10/10] WIP more allocs fix --- src/MatrixFields/field_name_dict.jl | 39 +++++++++++++------------ test/MatrixFields/scalar_fieldmatrix.jl | 27 ++++++++--------- 2 files changed, 32 insertions(+), 34 deletions(-) diff --git a/src/MatrixFields/field_name_dict.jl b/src/MatrixFields/field_name_dict.jl index 6f9db8717f..eb5a3f31eb 100644 --- a/src/MatrixFields/field_name_dict.jl +++ b/src/MatrixFields/field_name_dict.jl @@ -187,10 +187,15 @@ function get_internal_entry( # Ensure compatibility with RecursiveApply (i.e., with rmul). # See note above matrix_product_keys in field_name_set.jl for more details. T = eltype(eltype(entry)) + # @show !broadcasted_has_field(T, name_pair[1]) + # @show name_pair + # @show T if name_pair == (@name(), @name()) entry - elseif name_pair[1] == name_pair[2] && - !broadcasted_has_field(T, name_pair[1]) + elseif name_pair[1] == name_pair[2] && !broadcasted_has_field(T, name_pair[1]) && + !(T <: Geometry.SingleValue) #&& + #!broadcasted_has_field(T, name_pair[1]) + # @assert !(T <: Geometry.SingleValue) # @show name_pair # # multiplication case 3 or 4, first argument # if T <: Number @@ -198,7 +203,9 @@ function get_internal_entry( # else # throw(key_error) # end - entry + # @show "aaa" + # entry + nothing elseif name_pair[1] == @name() || name_pair[2] == @name() target_chain = if name_pair[1] == @name() @@ -423,17 +430,17 @@ end function bypass_adjoint(field::Fields.Field) if eltype(field) <: Adjoint - return (eltype(field.parent), true) + return eltype(field.parent) else - return (eltype(field), false) + return eltype(field) end end function bypass_adjoint(d::T) where {T} if T <: Adjoint - return (typeof(d.parent), true) + return typeof(d.parent) else - return (typeof(d), false) + return typeof(d) end end @@ -446,8 +453,8 @@ Returns a `FieldMatrixKeys` object that contains the keys of all the scalar entries in the `FieldMatrix` `dict`. """ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) - target_eltype = eltype(Y) keys_tuple = unrolled_flatmap(keys(dict).values) do key + target_eltype = eltype(Y) entry = dict[unrolled_filter(isequal(key), keys(dict).values)[1]] if entry isa UniformScaling # uniformscalings can only contain numbers (key,) @@ -474,7 +481,7 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) # TODO: Make the cth axis check for all combinations if independent_type <: Geometry.AxisTensor && dependent_type <: Geometry.AxisTensor - (unwrapped_row_eltype, is_adjoint) = + unwrapped_row_eltype = bypass_adjoint(first_band) @assert unwrapped_row_eltype <: Geometry.Axis2Tensor axis_components = @@ -509,7 +516,7 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) independent_type <: Number && dependent_type <: Geometry.AxisTensor ) - (unwrapped_row_eltype, is_adjoint) = + unwrapped_row_eltype = bypass_adjoint(first_band) @assert unwrapped_row_eltype <: Geometry.AxisVector ncomponents = length(axes(unwrapped_row_eltype)[1]) @@ -548,11 +555,11 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) first_band, ), ) do dependent_name - (unwrapped_row_eltype, is_adjoint) = bypass_adjoint( + unwrapped_row_eltyped = bypass_adjoint( get_field(first_band, dependent_name), ) - ncomponents = length(axes(unwrapped_row_eltype)[1]) - unrolled_map(1:ncomponents) do component + ncomponents1 = length(axes(unwrapped_row_eltyped)[1]) + unrolled_map(1:ncomponents1) do component ( append_internal_name(key[1], dependent_name), append_internal_name( @@ -565,7 +572,6 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) ) end end - # elseif dependent_type <: Union{NamedTuple, Tuple} && independent_type <: Number unrolled_map( @@ -578,14 +584,9 @@ function get_scalar_keys(dict::FieldMatrix, Y::Fields.FieldVector) (append_internal_name(key[1], dependent_name), key[2]) end - end # (key,) end - # elseif entry isa DiagonalMatrixRow - # target_eltype = eltype(parent(get_field(Y, key[1]))) - # # TODO: unify target_eltype - # (key,) else error("Cannot get scalar keys for key $key") end diff --git a/test/MatrixFields/scalar_fieldmatrix.jl b/test/MatrixFields/scalar_fieldmatrix.jl index 07ddedfdc6..489ddf202e 100644 --- a/test/MatrixFields/scalar_fieldmatrix.jl +++ b/test/MatrixFields/scalar_fieldmatrix.jl @@ -129,9 +129,7 @@ end FT = Float64 A, b = dycore_prognostic_EDMF_FieldMatrix(FT) for key in MatrixFields.get_scalar_keys(A, b) - # @show key - # TODO: test_all - @test A[key] isa MatrixFields.ColumnwiseBandMatrixField ? + @test_all A[key] isa MatrixFields.ColumnwiseBandMatrixField ? eltype(eltype(A[key])) == eltype(parent(A[key])) : eltype(eltype(A[key])) == eltype(A[key]) end @@ -144,17 +142,16 @@ end test_get(A, entry, key) = A[key] === entry for (key, entry) in MatrixFields.scalar_fieldmatrix(A, b) @test test_get(A, entry, key) - # @test (@allocated test_get(A, entry, key)) == 0 - # @test_opt test_get(A, entry, key) + @test (@allocated test_get(A, entry, key)) == 0 + @test_opt test_get(A, entry, key) end - # function scalar_fieldmatrix_wrapper(field_matrix_of_tensors, b) - # A_scalar = MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors, b) - # return nothing - # end - # scalar_fieldmatrix_wrapper(A, b) - # scalar_fieldmatrix_wrapper(A, b) # compile the wrapper function - # al = @allocated scalar_fieldmatrix_wrapper(A, b) - # @show al - # @test (@allocated scalar_fieldmatrix_wrapper(A, b)) == 0 - # @test_opt MatrixFields.scalar_fieldmatrix(A, b) + + function scalar_fieldmatrix_wrapper(field_matrix_of_tensors, b) + A_scalar = MatrixFields.scalar_fieldmatrix(field_matrix_of_tensors, b) + return nothing + end + scalar_fieldmatrix_wrapper(A, b) + scalar_fieldmatrix_wrapper(A, b) # compile the wrapper function + @test (@allocated scalar_fieldmatrix_wrapper(A, b)) == 0 + @test_opt MatrixFields.scalar_fieldmatrix(A, b) end