diff --git a/src/Operators/integrals.jl b/src/Operators/integrals.jl index f919ee9548..2c1333733e 100644 --- a/src/Operators/integrals.jl +++ b/src/Operators/integrals.jl @@ -1,6 +1,8 @@ -import ..RecursiveApply: rzero, ⊠, ⊞ +import ..RecursiveApply: rzero, ⊠, ⊞, radd, rmul import ..Utilities import ..DataLayouts +import RootSolvers +import ClimaComms """ column_integral_definite!(∫field::Field, ᶜfield::Field) @@ -86,6 +88,21 @@ end Sets `ᶠ∫field```(z) = \\int_0^z\\,```ᶜfield```(z')\\,dz'``. The input `ᶜfield` must lie on a cell-center space, and the output `ᶠ∫field` must lie on the corresponding cell-face space. + + column_integral_indefinite!( + f::Function, + ᶠ∫field::Fields.ColumnField, + ϕ₀ = 0, + average = (ϕ⁻, ϕ⁺) -> (ϕ⁻ + ϕ⁺) / 2, + ) + +The indefinite integral `ᶠ∫field = ϕ(z) = ∫ f(ϕ,z) dz` given: + +- `f` the integral integrand function (which may be a function) +- `ᶠ∫field` the resulting (scalar) field `ϕ(z)` +- `ϕ₀` (optional) the boundary condition +- `average` (optional) a function to compute the cell center + average between two cell faces (`ϕ⁻, ϕ⁺`). """ column_integral_indefinite!(ᶠ∫field::Fields.Field, ᶜfield::Fields.Field) = column_integral_indefinite!(ClimaComms.device(ᶠ∫field), ᶠ∫field, ᶜfield) @@ -155,6 +172,94 @@ function _column_integral_indefinite!( end end +dual_space(space::Spaces.FaceExtrudedFiniteDifferenceSpace) = + Spaces.CenterExtrudedFiniteDifferenceSpace(space) +dual_space(space::Spaces.CenterExtrudedFiniteDifferenceSpace) = + Spaces.FaceExtrudedFiniteDifferenceSpace(space) + +dual_space(space::Spaces.FaceFiniteDifferenceSpace) = + Spaces.CenterFiniteDifferenceSpace(space) +dual_space(space::Spaces.CenterFiniteDifferenceSpace) = + Spaces.FaceFiniteDifferenceSpace(space) + +# First, dispatch on device: +column_integral_indefinite!( + f::Function, + ᶠ∫field::Fields.Field, + ϕ₀ = zero(eltype(ᶠ∫field)), + average = (ϕ⁻, ϕ⁺) -> (ϕ⁻ + ϕ⁺) / 2, +) = column_integral_indefinite!( + f, + ClimaComms.device(ᶠ∫field), + ᶠ∫field, + ϕ₀, + average, +) + +##### +##### CPU +##### + +column_integral_indefinite!( + f::Function, + ::ClimaComms.AbstractCPUDevice, + ᶠ∫field, + args..., +) = column_integral_indefinite_cpu!(f, ᶠ∫field, args...) + +column_integral_indefinite!( + f::Function, + ::ClimaComms.AbstractCPUDevice, + ᶠ∫field::Fields.FaceExtrudedFiniteDifferenceField, + args..., +) = + Fields.bycolumn(axes(ᶠ∫field)) do colidx + column_integral_indefinite_cpu!(f, ᶠ∫field[colidx], args...) + end + +#= +Function-based signature, solve for ϕ: +``` +∂ϕ/∂z = f(ϕ,z) +(ᶠϕ^{k+1}-ᶠϕ^{k})/ᶜΔz = ᶜf(ᶜϕ̄,ᶜz) +ᶜϕ̄ = (ϕ^{k+1}+ϕ^{k})/2 +(ᶠϕ^{k+1}-ᶠϕ^{k})/ᶜΔz = ᶜf((ᶠϕ^{k+1}+ᶠϕ^{k})/2,ᶜz) +root equation: (_ϕ-ϕ^{k})/Δz = f((_ϕ+ϕ^{k})/2,ᶜz) +``` +=# +function column_integral_indefinite_cpu!( + f::Function, + ᶠ∫field::Fields.ColumnField, + ϕ₀ = zero(eltype(ᶠ∫field)), + average = (ϕ⁻, ϕ⁺) -> (ϕ⁻ + ϕ⁺) / 2, +) + cspace = dual_space(axes(ᶠ∫field)) + ᶜzfield = Fields.coordinate_field(cspace) + face_space = axes(ᶠ∫field) + first_level = Operators.left_idx(face_space) + last_level = Operators.right_idx(face_space) + ᶜΔzfield = Fields.Δz_field(ᶜzfield) + @inbounds Fields.level(ᶠ∫field, first_level)[] = ϕ₀ + ϕ₁ = ϕ₀ + @inbounds for level in (first_level + 1):last_level + ᶜz = Fields.level(ᶜzfield.z, level - half)[] + ᶜΔz = Fields.level(ᶜΔzfield, level - half)[] + ϕ₀ = ϕ₁ + root_eq(_x) = (_x - ϕ₀) / ᶜΔz - f(average(_x, ϕ₀), ᶜz) + sol = RootSolvers.find_zero( + root_eq, + RootSolvers.NewtonsMethodAD(ϕ₀), + RootSolvers.CompactSolution(), + ) + ϕ₁ = sol.root + f₁ = f(average(ϕ₀, ϕ₁), ᶜz) + ᶜintegrand_lev = f₁ + @inbounds Fields.level(ᶠ∫field, level)[] = + radd(Fields.level(ᶠ∫field, level - 1)[], rmul(ᶜintegrand_lev, ᶜΔz)) + end + return nothing +end + """ column_mapreduce!(fn, op, reduced_field::Field, fields::Field...) diff --git a/test/Operators/integrals.jl b/test/Operators/integrals.jl index eb6ed00600..e7d8c71a33 100644 --- a/test/Operators/integrals.jl +++ b/test/Operators/integrals.jl @@ -76,8 +76,44 @@ function test_column_integral_indefinite!(center_space, alloc_lim) test_allocs(alloc_lim, allocs, "test_column_integral_indefinite!") end -function test_column_mapreduce!(space, alloc_lim; broken = false) - broken && return +function test_column_integral_indefinite_fn!(center_space, alloc_lim) + face_space = center_to_face_space(center_space) + ᶜz = Fields.coordinate_field(center_space).z + ᶠz = Fields.coordinate_field(face_space).z + FT = Spaces.undertype(center_space) + + ᶜu = Dict() + ᶠ∫u_ref = Dict() + ᶠ∫u_test = Dict() + + + # ᶜu = map(z -> (; one = one(z), powers = (z, z^2, z^3)), ᶜz) + # ᶠ∫u_ref = map(z -> (; one = z, powers = (z^2 / 2, z^3 / 3, z^4 / 4)), ᶠz) + # ᶠ∫u_test = similar(ᶠ∫u_ref) + + for (i, fn) in enumerate(((ϕ, z) -> z, (ϕ, z) -> z^2, (ϕ, z) -> z^3)) + ᶜu = ᶜz .^ i + ᶠ∫u_ref = ᶠz .^ (i + 1) ./ (i + 1) + ᶠ∫u_test = similar(ᶠ∫u_ref) + + column_integral_indefinite!(fn, ᶠ∫u_test) + ref_array = + parent(Fields.level(ᶠ∫u_ref, Operators.right_idx(face_space))) + test_array = + parent(Fields.level(ᶠ∫u_test, Operators.right_idx(face_space))) + max_relative_error = + maximum(@. abs((ref_array - test_array) / ref_array)) + @test max_relative_error <= 0.006 # Less than 0.6% error at the top level. + + # @test_opt column_integral_indefinite!(fn, ᶠ∫u_test) + + allocs = @allocated column_integral_indefinite!(fn, ᶠ∫u_test) + test_allocs(alloc_lim, allocs, "test_column_integral_indefinite_fn!") + end + +end + +function test_column_mapreduce!(space, alloc_lim) z_field = Fields.coordinate_field(space).z z_top_field = Fields.level(z_field, Operators.right_idx(space)) sin_field = @. sin(pi * z_field / z_top_field) @@ -137,7 +173,7 @@ end i_lim[(1, Float64)] = 1936 i_lim[(2, Float64)] = 4720 - i_lim[(3, Float64)] = 2368 + i_lim[(3, Float64)] = 2544 i_lim[(4, Float64)] = 8176 lim = Dict() @@ -173,26 +209,30 @@ end TU.CenterExtrudedFiniteDifferenceSpace(FT; context), i_lim[(4, FT)], ) + broken || test_column_integral_indefinite_fn!( + TU.ColumnCenterFiniteDifferenceSpace(FT; context), + i_lim[(3, FT)], + ) + broken || test_column_integral_indefinite_fn!( + TU.CenterExtrudedFiniteDifferenceSpace(FT; context), + i_lim[(4, FT)], + ) - test_column_mapreduce!( + broken || test_column_mapreduce!( TU.ColumnCenterFiniteDifferenceSpace(FT; context), lim[(1, FT)], - broken = broken, ) - test_column_mapreduce!( + broken || test_column_mapreduce!( TU.ColumnFaceFiniteDifferenceSpace(FT; context), lim[(2, FT)], - broken = broken, ) - test_column_mapreduce!( + broken || test_column_mapreduce!( TU.CenterExtrudedFiniteDifferenceSpace(FT; context), - lim[(3, FT)]; - broken = broken, + lim[(3, FT)], ) - test_column_mapreduce!( + broken || test_column_mapreduce!( TU.FaceExtrudedFiniteDifferenceSpace(FT; context), - lim[(4, FT)]; - broken = broken, + lim[(4, FT)], ) end end