Add and test function-based indefinite integrals

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,19 @@ 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,
+    )
+
+Populates `ᶠ∫field` with the indefinite integral defined by:
+```
+∂ϕ/∂z = f(ϕ,z).        i.e., `ᶠ∫field = ϕ(z)`
+```
+Note that the boundary value `ϕ₀` is optionally supplied.
 """
 column_integral_indefinite!(ᶠ∫field::Fields.Field, ᶜfield::Fields.Field) =
     column_integral_indefinite!(ClimaComms.device(ᶠ∫field), ᶠ∫field, ᶜfield)
@@ -155,6 +170,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)
@@ -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