From 90752686fe9be25320bbd3a572d6f0821caccac1 Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Fri, 20 Sep 2024 11:59:13 -0700 Subject: [PATCH] extend dss functions for FieldVectors --- NEWS.md | 3 +++ src/Fields/Fields.jl | 2 +- src/Fields/fieldvector.jl | 35 ++++++++++++++++++++++++++++++ test/Fields/field_opt.jl | 45 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 84 insertions(+), 1 deletion(-) diff --git a/NEWS.md b/NEWS.md index 004b157664..bca8ccb1a3 100644 --- a/NEWS.md +++ b/NEWS.md @@ -4,6 +4,9 @@ ClimaCore.jl Release Notes main ------- +- Extended `create_dss_buffer` and `weighted_dss!` for `FieldVector`s, rather than +just `Field`s. PR [#2000](https://github.com/CliMA/ClimaCore.jl/pull/2000). + v0.14.16 ------- diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index c8c0ab7baf..171daf699e 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -374,7 +374,7 @@ function interpcoord(elemrange, x::Real) end """ - Spaces.weighted_dss!(f::Field[, ghost_buffer = Spaces.create_dss_buffer(field)]) + Spaces.weighted_dss!(f::Field, dss_buffer = Spaces.create_dss_buffer(field)) Apply weighted direct stiffness summation (DSS) to `f`. This operates in-place (i.e. it modifies the `f`). `ghost_buffer` contains the necessary information diff --git a/src/Fields/fieldvector.jl b/src/Fields/fieldvector.jl index 158f317b80..1dc3ee30f8 100644 --- a/src/Fields/fieldvector.jl +++ b/src/Fields/fieldvector.jl @@ -184,6 +184,41 @@ end return dest end +""" + Spaces.create_dss_buffer(fv::FieldVector) + +Create a NamedTuple of buffers for communicating neighbour information of +each Field in `fv`. In this NamedTuple, the name of each field is mapped +to the buffer. +""" +function Spaces.create_dss_buffer(fv::FieldVector) + NamedTuple{propertynames(fv)}( + map( + key -> Spaces.create_dss_buffer(getproperty(fv, key)), + propertynames(fv), + ), + ) +end + +""" + Spaces.weighted_dss!(fv::FieldVector, dss_buffer = Spaces.create_dss_buffer(fv)) + +Apply weighted direct stiffness summation (DSS) to each field in `fv`. +If a `dss_buffer` object is not provided, a buffer will be created for each +field in `fv`. +Note that using the `Pair` interface here parallelizes the `weighted_dss!` calls. +""" +function Spaces.weighted_dss!( + fv::FieldVector, + dss_buffer = Spaces.create_dss_buffer(fv), +) + pairs = map(propertynames(fv)) do key + Pair(getproperty(fv, key), getproperty(dss_buffer, key)) + end + Spaces.weighted_dss!(pairs...) +end + + # Recursively call transform_bc_args() on broadcast arguments in a way that is statically reducible by the optimizer # see Base.Broadcast.preprocess_args @inline transform_bc_args(args::Tuple, inds...) = ( diff --git a/test/Fields/field_opt.jl b/test/Fields/field_opt.jl index 24955cdbb5..b57c4f3206 100644 --- a/test/Fields/field_opt.jl +++ b/test/Fields/field_opt.jl @@ -393,4 +393,49 @@ using JET @test_opt ifelsekernel!(S, ρ) end +@testset "dss of FieldVectors" begin + function field_vec(center_space, face_space) + Y = Fields.FieldVector( + c = map(Fields.coordinate_field(center_space)) do coord + FT = Spaces.undertype(center_space) + (; + ρ = FT(coord.lat + coord.long), + uₕ = Geometry.Covariant12Vector( + FT(coord.lat), + FT(coord.long), + ), + ) + end, + f = map(Fields.coordinate_field(face_space)) do coord + FT = Spaces.undertype(face_space) + (; w = Geometry.Covariant3Vector(FT(coord.lat + coord.long))) + end, + ) + return Y + end + + fv = field_vec(toy_sphere(Float64)...) + + c_copy = copy(getproperty(fv, :c)) + f_copy = copy(getproperty(fv, :f)) + + # Test that dss_buffer is created and has the correct keys + dss_buffer = Spaces.create_dss_buffer(fv) + @test haskey(dss_buffer, :c) + @test haskey(dss_buffer, :f) + + # Test weighted_dss! with and without preallocated buffer + Spaces.weighted_dss!(fv, dss_buffer) + @test getproperty(fv, :c) ≈ Spaces.weighted_dss!(c_copy) + @test getproperty(fv, :f) ≈ Spaces.weighted_dss!(f_copy) + + fv = field_vec(toy_sphere(Float64)...) + c_copy = copy(getproperty(fv, :c)) + f_copy = copy(getproperty(fv, :f)) + + Spaces.weighted_dss!(fv) + @test getproperty(fv, :c) ≈ Spaces.weighted_dss!(c_copy) + @test getproperty(fv, :f) ≈ Spaces.weighted_dss!(f_copy) +end + nothing