Merge #1504 #1515

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. 1515: remove distributed TempestRemap tests r=charleskawczynski a=simonbyrne This isn't how we should do distributed remapping, so i'll just disable them. Avoids #1513. - [x] Code follows the [style guidelines](https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/) OR N/A. - [x] Unit tests are included OR N/A. - [x] Code is exercised in an integration test OR N/A. - [x] Documentation has been added/updated OR N/A. Co-authored-by: Charles Kawczynski <[email protected]> Co-authored-by: Simon Byrne <[email protected]>
CliMA · Oct 23, 2023 · ebdfa12 · ebdfa12
3 parents 9c22047 + f144a9a + 273ca9b
commit ebdfa12
Show file tree

Hide file tree

Showing 3 changed files with 158 additions and 39 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -681,32 +681,6 @@ steps:
         agents:
           slurm_ntasks: 3
 
-  - group: "Unit: lib/ClimaCoreTempestRemap"
-    steps:
-
-      - label: "Unit: lib/TempestRemap MPI test"
-        key: "write_exodus_mpi"
-        command: "srun julia --color=yes --project=lib/ClimaCoreTempestRemap lib/ClimaCoreTempestRemap/test/mpi_tests/exodus.jl"
-        timeout_in_minutes: 5
-        env:
-          CLIMACOMMS_CONTEXT: "MPI"
-        retry:
-          automatic: true
-        agents:
-          slurm_nodes: 3
-          slurm_tasks_per_node: 1
-
-      - label: "Unit: lib/TempestRemap distributed remapping test"
-        key: "online_remap_mpi"
-        command: "srun julia --color=yes --project=lib/ClimaCoreTempestRemap lib/ClimaCoreTempestRemap/test/mpi_tests/online_remap.jl"
-        timeout_in_minutes: 5
-        env:
-          CLIMACOMMS_CONTEXT: "MPI"
-          CI_OUTPUT_DIR: "output/online_remap"
-        agents:
-          slurm_nodes: 3
-          slurm_tasks_per_node: 1
-
   # TODO: improve performance label: [inference, allocs, flops, latency, benchmark, boundscheck]
   # TODO: use perf env for perf jobs, or merge test/perf envs
   - group: "Perf: Geometry"

diff --git 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
@@ -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