Merge #962

962: Add support for definite column integrals r=simonbyrne a=charleskawczynski

This PR adds `column_integral_definite!`, a non-allocating definite integral for columns. The idea is that it will replace the existing (allocating) definite integral [here](https://github.com/CliMA/ClimaAtmos.jl/blob/44cee55def2a51433c5dcdcdae010b7777cd741b/examples/hybrid/sphere/baroclinic_wave_utilities.jl#L167-L186). However, this implementation has several advantages:
 - It's compatible with `bycolumn`, so the computation can occur in parallel
 - It's allocation-free
 - It's _tested_ on single column and sphere configurations

Looking at the flame graph for allocations (using the Julia 1.8 Allocs module with `sample_rate=1`), this is responsible for most of the remaining allocations in ClimaAtmos:
<img width="1993" alt="Screen Shot 2022-09-26 at 7 57 05 AM" src="https://user-images.githubusercontent.com/1880641/192353757-03101c41-2c0b-4ccb-a8b9-43faa78680f8.png">


The interface for the added function captures two cases:

```julia
function column_integral_definite!(col∫field::Field, field::Field)
    bycolumn(axes(field)) do colidx
        column_integral_definite!(col∫field[colidx], field[colidx])
        nothing
    end
    return nothing
end
```
and
```julia
function column_integral_definite!(
    col∫field::PointField,
    field::ColumnField,
)
    `@inbounds` col∫field[] = column_integral_definite(field)
    return nothing
end
```

A step towards closing #943, #748, [CA#686](CliMA/ClimaAtmos.jl#686).

## A note on an alternative approach
Ideally, we would write this function as `column_integral_definite!(fn, args...)` where we might be able to write a broadcast statement like:

```julia
`@.` f2D = column_integral_definite(f3D) do z
     f3D.x * z^2
end
```

however, this would require us to define broadcasting between planar and 3D domains, which is not trivial (or maybe not possible) because of ambiguities. The ambiguities arise because (2D, 3D) broadcasting may want different things in different cases, for example:

 - (f2D, f3D) -> f3D: mul full 3D field by planar surface value
 - (f2D, f3D) -> f2D: perform reduction over z-coordinate to yield 2D field

The situation is similar when thinking about what happens when we make views. For example,

```julia
Fields.bycolumn(axes(f3D)) do colidx
     `@.` f2D[colidx] = column_integral_definite(f3D[colidx]) do z
         f3D[colidx].x * z^2
      end
end
```

Now, we have to define how `DataF` data layouts broadcast with `VF`. Again, we have two cases:

 - (f0D, f1D) -> f1D: mul full 1D field by 0D field
 - (f0D, f1D) -> f0D: perform reduction over z-coordinate to yield 0D field

My vote/preference is to support the first cases (which is partially supported already) and write custom functions (e.g., reductions) that operate on single fields for the second case.

Co-authored-by: Charles Kawczynski <[email protected]>
Co-authored-by: Simon Byrne <[email protected]>

src/Fields/Fields.jl

@@ -256,6 +256,14 @@ local_geometry_field(space::AbstractSpace) =
 local_geometry_field(field::Field) = local_geometry_field(axes(field))
 
 
+"""
+    dz_field(field::Field)
+    dz_field(space::AbstractSpace)
+
+A `Field` containing the `Δz` values on the same space as the given field.
+"""
+dz_field(field::Field) = dz_field(axes(field))
+dz_field(space::AbstractSpace) = Field(Spaces.dz_data(space), space)
 
 include("broadcast.jl")
 include("mapreduce.jl")

src/Operators/Operators.jl

@@ -21,5 +21,6 @@ include("stencilcoefs.jl")
 include("operator2stencil.jl")
 include("pointwisestencil.jl")
 include("remapping.jl")
+include("integrals.jl")
 
 end # module

src/Operators/integrals.jl

@@ -0,0 +1,50 @@
+#####
+##### Column integrals (definite)
+#####
+
+"""
+    column_integral_definite(field::Field)
+
+The definite vertical column integral of field `field`.
+"""
+column_integral_definite(field::Fields.CenterExtrudedFiniteDifferenceField) =
+    column_integral_definite(field, 1)
+column_integral_definite(field::Fields.FaceExtrudedFiniteDifferenceField) =
+    column_integral_definite(field, PlusHalf(1))
+
+column_integral_definite(field::Fields.Field, one_index::Union{Int, PlusHalf}) =
+    column_integral_definite(similar(level(field, one_index)), field)
+
+"""
+    column_integral_definite!(col∫field::Field, field::Field)
+
+The definite vertical column integral, `col∫field`, of field `field`.
+"""
+function column_integral_definite!(col∫field::Fields.Field, field::Fields.Field)
+    Fields.bycolumn(axes(field)) do colidx
+        column_integral_definite!(col∫field[colidx], field[colidx])
+        nothing
+    end
+    return nothing
+end
+
+function column_integral_definite!(
+    col∫field::Fields.PointField,
+    field::Fields.ColumnField,
+)
+    @inbounds col∫field[] = column_integral_definite(field)
+    return nothing
+end
+
+function column_integral_definite(field::Fields.ColumnField)
+    field_data = Fields.field_values(field)
+    Δz_data = Spaces.dz_data(axes(field))
+    Nv = size(Δz_data, 5)
+    ∫field = zero(eltype(field))
+    @inbounds for j in 1:Nv
+        ∫field += field_data[j] * Δz_data[j]
+    end
+    return ∫field
+end
+
+# TODO: add support for indefinite integrals

src/Spaces/finitedifference.jl

@@ -153,6 +153,27 @@ local_geometry_data(space::CenterFiniteDifferenceSpace) =
 local_geometry_data(space::FaceFiniteDifferenceSpace) =
     space.face_local_geometry
 
+
+"""
+    z_component(::Type{<:Goemetry.AbstractPoint})
+
+The index of the z-component of an abstract point
+in an `AxisTensor`.
+"""
+z_component(::Type{<:Geometry.LatLongZPoint}) = 9
+z_component(::Type{<:Geometry.ZPoint}) = 1
+
+"""
+    dz_data(space::AbstractSpace)
+
+A DataLayout containing the `Δz` on a given space `space`.
+"""
+function dz_data(space::AbstractSpace)
+    lg = local_geometry_data(space)
+    data_layout_type = eltype(lg.coordinates)
+    return getproperty(lg.∂x∂ξ.components.data, z_component(data_layout_type))
+end
+
 function left_boundary_name(space::FiniteDifferenceSpace)
     boundaries = Topologies.boundaries(Spaces.topology(space))
     propertynames(boundaries)[1]

test/Fields/field.jl

@@ -1,4 +1,5 @@
 using Test
+using OrderedCollections
 using StaticArrays, IntervalSets
 import ClimaCore
 import ClimaCore.Utilities: PlusHalf
@@ -17,6 +18,21 @@ end
 
 include(joinpath(@__DIR__, "util_spaces.jl"))
 
+fc_index(
+    i,
+    ::Union{
+        Spaces.FaceExtrudedFiniteDifferenceSpace,
+        Spaces.FaceFiniteDifferenceSpace,
+    },
+) = PlusHalf(i)
+fc_index(
+    i,
+    ::Union{
+        Spaces.CenterExtrudedFiniteDifferenceSpace,
+        Spaces.CenterFiniteDifferenceSpace,
+    },
+) = i
+
 function spectral_space_2D(; n1 = 1, n2 = 1, Nij = 4)
     domain = Domains.RectangleDomain(
         Geometry.XPoint(-3.0) .. Geometry.XPoint(5.0),
@@ -309,21 +325,6 @@ end
     @test isapprox(Fields.field_values(add_field)[], FT(2π))
 end
 
-fc_index(
-    i,
-    ::Union{
-        Spaces.FaceExtrudedFiniteDifferenceSpace,
-        Spaces.FaceFiniteDifferenceSpace,
-    },
-) = PlusHalf(i)
-fc_index(
-    i,
-    ::Union{
-        Spaces.CenterExtrudedFiniteDifferenceSpace,
-        Spaces.CenterFiniteDifferenceSpace,
-    },
-) = i
-
 @testset "Level" begin
     FT = Float64
     for space in all_spaces(FT)
@@ -517,3 +518,113 @@ end
     C .= Ref(zero(eltype(C)))
     @test all(==(0.0), parent(C))
 end
+function integrate_bycolumn!(∫y, Y)
+    Fields.bycolumn(axes(Y.y)) do colidx
+        Operators.column_integral_definite!(∫y[colidx], Y.y[colidx])
+        nothing
+    end
+end
+
+"""
+    convergence_rate(err, Δh)
+
+Estimate convergence rate given vectors `err` and `Δh`
+
+    err = C Δh^p+ H.O.T
+    err_k ≈ C Δh_k^p
+    err_k/err_m ≈ Δh_k^p/Δh_m^p
+    log(err_k/err_m) ≈ log((Δh_k/Δh_m)^p)
+    log(err_k/err_m) ≈ p*log(Δh_k/Δh_m)
+    log(err_k/err_m)/log(Δh_k/Δh_m) ≈ p
+
+"""
+convergence_rate(err, Δh) =
+    [log(err[i] / err[i - 1]) / log(Δh[i] / Δh[i - 1]) for i in 2:length(Δh)]
+
+@testset "Definite column integrals bycolumn" begin
+    FT = Float64
+    results = OrderedCollections.OrderedDict()
+    ∫y_analytic = 1 - cos(1) - (0 - cos(0))
+    function col_field_copy(y)
+        col_copy = similar(y[Fields.ColumnIndex((1, 1), 1)])
+        return Fields.Field(Fields.field_values(col_copy), axes(col_copy))
+    end
+    for zelem in (2^2, 2^3, 2^4, 2^5)
+        for space in all_spaces(FT; zelem)
+            # Filter out spaces without z coordinates:
+            :z in propertynames(Spaces.coordinates_data(space)) || continue
+            # Skip spaces incompatible with Fields.bycolumn:
+            (
+                space isa Spaces.ExtrudedFiniteDifferenceSpace ||
+                space isa Spaces.SpectralElementSpace1D ||
+                space isa Spaces.SpectralElementSpace2D
+            ) || continue
+
+            Y = FieldFromNamedTuple(space, (; y = FT(1)))
+            zcf = Fields.coordinate_field(Y.y).z
+            Δz = Fields.dz_field(axes(zcf))
+            Δz_col = Δz[Fields.ColumnIndex((1, 1), 1)]
+            Δz_1 = parent(Δz_col)[1]
+            key = (space, zelem)
+            if !haskey(results, key)
+                results[key] =
+                    Dict(:maxerr => 0, :Δz_1 => FT(0), :∫y => [], :y => [])
+            end
+            ∫y = Spaces.level(similar(Y.y), fc_index(1, space))
+            ∫y .= 0
+            y = Y.y
+            @. y .= 1 + sin(zcf)
+            # Compute max error against analytic solution
+            maxerr = 0
+            Fields.bycolumn(axes(Y.y)) do colidx
+                Operators.column_integral_definite!(∫y[colidx], y[colidx])
+                maxerr = max(
+                    maxerr,
+                    maximum(abs.(parent(∫y[colidx]) .- ∫y_analytic)),
+                )
+                nothing
+            end
+            results[key][:Δz_1] = Δz_1
+            results[key][:maxerr] = maxerr
+            push!(results[key][:∫y], ∫y)
+            push!(results[key][:y], y)
+            nothing
+        end
+    end
+    maxerrs = map(key -> results[key][:maxerr], collect(keys(results)))
+    Δzs_1 = map(key -> results[key][:Δz_1], collect(keys(results)))
+    cr = convergence_rate(maxerrs, Δzs_1)
+    @test 2 < sum(abs.(cr)) / length(cr) < 2.01
+end
+
+ClimaCore.enable_threading() = false # launching threads allocates
+@testset "Allocation tests for integrals" begin
+    FT = Float64
+    for space in all_spaces(FT)
+        # Filter out spaces without z coordinates:
+        :z in propertynames(Spaces.coordinates_data(space)) || continue
+        Y = FieldFromNamedTuple(space, (; y = FT(1)))
+        zcf = Fields.coordinate_field(Y.y).z
+        ∫y = Spaces.level(similar(Y.y), fc_index(1, space))
+        ∫y .= 0
+        y = Y.y
+        @. y .= 1 + sin(zcf)
+        # Implicit bycolumn
+        Operators.column_integral_definite!(∫y, y) # compile first
+        p = @allocated Operators.column_integral_definite!(∫y, y)
+        @test p == 0
+        # Skip spaces incompatible with Fields.bycolumn:
+        (
+            space isa Spaces.ExtrudedFiniteDifferenceSpace ||
+            space isa Spaces.SpectralElementSpace1D ||
+            space isa Spaces.SpectralElementSpace2D
+        ) || continue
+        # Explicit bycolumn
+        integrate_bycolumn!(∫y, Y) # compile first
+        p = @allocated integrate_bycolumn!(∫y, Y)
+        @test p == 0
+        nothing
+    end
+    nothing
+end
+ClimaCore.enable_threading() = false

test/Fields/util_spaces.jl

@@ -3,7 +3,7 @@ import ClimaCore as CC
 #=
 Return a vector of toy spaces for testing
 =#
-function all_spaces(::Type{FT}) where {FT}
+function all_spaces(::Type{FT}; zelem = 10) where {FT}
 
     # 1d domain space
     domain = CC.Domains.IntervalDomain(
@@ -58,7 +58,6 @@ function all_spaces(::Type{FT}) where {FT}
     radius = FT(128)
     zlim = (0, 1)
     helem = 4
-    zelem = 10
     Nq = 4
 
     vertdomain = CC.Domains.IntervalDomain(