Merge #1504

1504: Add and test function-based indefinite integrals r=charleskawczynski a=charleskawczynski

This PR adds an indefinite integral implementation that uses a function-based input argument. Concretely, given `ϕ(z)`, this interface allows us to solve for `∂ϕ/∂z = f(ϕ,z)`. The existing field-based implementation cannot be used in this situation, because `f` is a function of `ϕ`, which is what we're computing. So, this implementation uses RootSolvers to find `ϕ` at the next level in order to provide a somewhat generic interface.

Unfortunately, one limitation of this is that `ϕ` must be a scalar field, and it cannot be, e.g., a `NamedTuple`. We could try to lift this assumption, but that would require adjustments to RootSolvers.

Closes #1492.

Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

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...)
 

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