Migrate CA column ops to ClimaCore

.buildkite/pipeline.yml

@@ -502,6 +502,18 @@ steps:
           slurm_mem: 20GB
           slurm_gpus: 1
 
+      - label: "Unit: Integrals (CPU)"
+        key: "cpu_integrals"
+        command:
+          - "julia --color=yes --check-bounds=yes --project=test test/Operators/integrals.jl"
+
+      - label: "Unit: Integrals (GPU)"
+        key: "gpu_integrals"
+        command:
+          - "julia --color=yes --check-bounds=yes --project=test test/Operators/integrals.jl"
+        agents:
+          slurm_gpus: 1
+
       - label: "Unit: Thomas Algorithm"
         key: "cpu_thomas_algorithm"
         command:

docs/src/operators.md

@@ -97,6 +97,14 @@ SetDivergence
 Extrapolate
 ```
 
+## Integrals
+
+```@docs
+column_integral_definite!
+column_indefinite_integral!
+column_mapreduce!
+```
+
 ## Internal APIs
 
 ```@docs

src/Operators/integrals.jl

@@ -1,86 +1,258 @@
-#####
-##### Column integrals (definite)
-#####
+import ..RecursiveApply: rzero, ⊠, ⊞
 
 """
-    column_integral_definite!(col∫field::Field, field::Field)
+    column_integral_definite!(∫field::Field, ᶜfield::Field)
 
-The definite vertical column integral, `col∫field`, of field `field`.
+Sets `∫field```{}= \\int_0^{z_{max}}\\,```ᶜfield```(z)\\,dz``, where ``z_{max}``
+is the value of `z` at the top of the domain. The input `ᶜfield` must lie on a
+cell-center space, and the output `∫field` must lie on the corresponding
+horizontal space.
 """
-column_integral_definite!(col∫field::Fields.Field, field::Fields.Field) =
-    column_integral_definite!(ClimaComms.device(axes(field)), col∫field, field)
+column_integral_definite!(∫field::Fields.Field, ᶜfield::Fields.Field) =
+    column_integral_definite!(ClimaComms.device(axes(ᶜfield)), ∫field, ᶜfield)
 
 function column_integral_definite!(
     ::ClimaComms.CUDADevice,
-    col∫field::Fields.Field,
-    field::Fields.Field,
+    ∫field::Fields.Field,
+    ᶜfield::Fields.Field,
 )
-    Ni, Nj, _, _, Nh = size(Fields.field_values(col∫field))
+    Ni, Nj, _, _, Nh = size(Fields.field_values(∫field))
     nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh)
     @cuda threads = nthreads blocks = nblocks column_integral_definite_kernel!(
-        col∫field,
-        field,
+        ∫field,
+        ᶜfield,
     )
 end
 
 function column_integral_definite_kernel!(
-    col∫field::Fields.SpectralElementField,
-    field::Fields.ExtrudedFiniteDifferenceField,
+    ∫field::Fields.SpectralElementField,
+    ᶜfield::Fields.CenterExtrudedFiniteDifferenceField,
 )
     idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x
-    Ni, Nj, _, _, Nh = size(Fields.field_values(field))
+    Ni, Nj, _, _, Nh = size(Fields.field_values(ᶜfield))
     if idx <= Ni * Nj * Nh
         i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx)
-        colfield = Spaces.column(field, i, j, h)
-        _column_integral_definite!(Spaces.column(col∫field, i, j, h), colfield)
+        ∫field_column = Spaces.column(∫field, i, j, h)
+        ᶜfield_column = Spaces.column(ᶜfield, i, j, h)
+        _column_integral_definite!(∫field_column, ᶜfield_column)
     end
     return nothing
 end
 
 column_integral_definite_kernel!(
-    col∫field::Fields.PointField,
-    field::Fields.FiniteDifferenceField,
-) = _column_integral_definite!(col∫field, field)
+    ∫field::Fields.PointField,
+    ᶜfield::Fields.CenterFiniteDifferenceField,
+) = _column_integral_definite!(∫field, ᶜfield)
 
-function column_integral_definite!(
+column_integral_definite!(
     ::ClimaComms.AbstractCPUDevice,
-    col∫field::Fields.SpectralElementField,
-    field::Fields.ExtrudedFiniteDifferenceField,
-)
-    Fields.bycolumn(axes(field)) do colidx
-        _column_integral_definite!(col∫field[colidx], field[colidx])
-        nothing
+    ∫field::Fields.SpectralElementField,
+    ᶜfield::Fields.CenterExtrudedFiniteDifferenceField,
+) =
+    Fields.bycolumn(axes(ᶜfield)) do colidx
+        _column_integral_definite!(∫field[colidx], ᶜfield[colidx])
     end
-    return nothing
-end
 
 column_integral_definite!(
     ::ClimaComms.AbstractCPUDevice,
-    col∫field::Fields.PointField,
-    field::Fields.FiniteDifferenceField,
-) = _column_integral_definite!(col∫field, field)
+    ∫field::Fields.PointField,
+    ᶜfield::Fields.CenterFiniteDifferenceField,
+) = _column_integral_definite!(∫field, ᶜfield)
 
 function _column_integral_definite!(
-    col∫field::Fields.PointField,
-    field::Fields.ColumnField,
+    ∫field::Fields.PointField,
+    ᶜfield::Fields.ColumnField,
+)
+    space = axes(ᶜfield)
+    Δz = Fields.Δz_field(space)
+    first_level = Operators.left_idx(space)
+    last_level = Operators.right_idx(space)
+    ∫field[] = rzero(eltype(∫field))
+    @inbounds for level in first_level:last_level
+        ∫field[] =
+            ∫field[] ⊞ Fields.level(ᶜfield, level)[] ⊠ Fields.level(Δz, level)[]
+    end
+end
+
+"""
+    column_indefinite_integral!(ᶠ∫field::Field, ᶜfield::Field)
+
+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_indefinite_integral!(ᶠ∫field::Fields.Field, ᶜfield::Fields.Field) =
+    column_indefinite_integral!(ClimaComms.device(ᶠ∫field), ᶠ∫field, ᶜfield)
+
+function column_indefinite_integral!(
+    ::ClimaComms.CUDADevice,
+    ᶠ∫field::Fields.Field,
+    ᶜfield::Fields.Field,
 )
-    space = axes(field)
-    Δz_field = Fields.Δz_field(space)
-    Nv = Spaces.nlevels(space)
-
-    col∫field[] = 0
-    @inbounds for idx in 1:Nv
-        col∫field[] +=
-            reduction_getindex(field, idx) * reduction_getindex(Δz_field, idx)
+    Ni, Nj, _, _, Nh = size(Fields.field_values(ᶠ∫field))
+    nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh)
+    @cuda threads = nthreads blocks = nblocks column_indefinite_integral_kernel!(
+        ᶠ∫field,
+        ᶜfield,
+    )
+end
+
+function column_indefinite_integral_kernel!(
+    ᶠ∫field::Fields.FaceExtrudedFiniteDifferenceField,
+    ᶜfield::Fields.CenterExtrudedFiniteDifferenceField,
+)
+    idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+    Ni, Nj, _, _, Nh = size(Fields.field_values(ᶜfield))
+    if idx <= Ni * Nj * Nh
+        i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx)
+        ᶠ∫field_column = Spaces.column(ᶠ∫field, i, j, h)
+        ᶜfield_column = Spaces.column(ᶜfield, i, j, h)
+        _column_indefinite_integral!(ᶠ∫field_column, ᶜfield_column)
     end
     return nothing
 end
 
-reduction_getindex(column_field, index) = @inbounds getidx(
-    axes(column_field),
-    column_field,
-    Interior(),
-    index - 1 + left_idx(axes(column_field)),
+column_indefinite_integral_kernel!(
+    ᶠ∫field::Fields.FaceFiniteDifferenceField,
+    ᶜfield::Fields.CenterFiniteDifferenceField,
+) = _column_indefinite_integral!(ᶠ∫field, ᶜfield)
+
+column_indefinite_integral!(
+    ::ClimaComms.AbstractCPUDevice,
+    ᶠ∫field::Fields.FaceExtrudedFiniteDifferenceField,
+    ᶜfield::Fields.CenterExtrudedFiniteDifferenceField,
+) =
+    Fields.bycolumn(axes(ᶜfield)) do colidx
+        _column_indefinite_integral!(ᶠ∫field[colidx], ᶜfield[colidx])
+    end
+
+column_indefinite_integral!(
+    ::ClimaComms.AbstractCPUDevice,
+    ᶠ∫field::Fields.FaceFiniteDifferenceField,
+    ᶜfield::Fields.CenterFiniteDifferenceField,
+) = _column_indefinite_integral!(ᶠ∫field, ᶜfield)
+
+function _column_indefinite_integral!(
+    ᶠ∫field::Fields.ColumnField,
+    ᶜfield::Fields.ColumnField,
+)
+    face_space = axes(ᶠ∫field)
+    ᶜΔz = Fields.Δz_field(ᶜfield)
+    first_level = Operators.left_idx(face_space)
+    last_level = Operators.right_idx(face_space)
+    @inbounds Fields.level(ᶠ∫field, first_level)[] = rzero(eltype(ᶜfield))
+    @inbounds for level in (first_level + 1):last_level
+        Fields.level(ᶠ∫field, level)[] =
+            Fields.level(ᶠ∫field, level - 1)[] ⊞
+            Fields.level(ᶜfield, level - half)[] ⊠
+            Fields.level(ᶜΔz, level - half)[]
+    end
+end
+
+"""
+    column_mapreduce!(fn, op, reduced_field::Field, fields::Field...)
+
+Applies mapreduce along the vertical direction. The input `fields` must all lie
+on the same space, and the output `reduced_field` must lie on the corresponding
+horizontal space. The function `fn` is mapped over every input, and the function
+`op` is used to reduce the outputs of `fn`.
+"""
+column_mapreduce!(
+    fn::F,
+    op::O,
+    reduced_field::Fields.Field,
+    fields::Fields.Field...,
+) where {F, O} = column_mapreduce_device!(
+    ClimaComms.device(reduced_field),
+    fn,
+    op,
+    reduced_field,
+    fields...,
 )
 
-# TODO: add support for indefinite integrals
+function column_mapreduce_device!(
+    ::ClimaComms.CUDADevice,
+    fn::F,
+    op::O,
+    reduced_field::Fields.Field,
+    fields::Fields.Field...,
+) where {F, O}
+    Ni, Nj, _, _, Nh = size(Fields.field_values(reduced_field))
+    nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh)
+    @cuda threads = nthreads blocks = nblocks column_mapreduce_kernel!(
+        fn,
+        op,
+        reduced_field,
+        fields...,
+    )
+end
+
+function column_mapreduce_kernel!(
+    fn::F,
+    op::O,
+    reduced_field::Fields.SpectralElementField,
+    fields::Fields.ExtrudedFiniteDifferenceField...,
+) where {F, O}
+    idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+    Ni, Nj, _, _, Nh = size(Fields.field_values(field))
+    if idx <= Ni * Nj * Nh
+        i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx)
+        reduced_field_column = Spaces.column(reduced_field, i, j, h)
+        field_columns = map(field -> Spaces.column(field, i, j, h), fields)
+        _column_mapreduce!(fn, op, reduced_field_column, field_columns...)
+    end
+    return nothing
+end
+
+column_mapreduce_kernel!(
+    fn::F,
+    op::O,
+    reduced_field::Fields.PointField,
+    fields::Fields.FiniteDifferenceField...,
+) where {F, O} = _column_mapreduce!(fn, op, reduced_field, fields...)
+
+column_mapreduce_device!(
+    ::ClimaComms.AbstractCPUDevice,
+    fn::F,
+    op::O,
+    reduced_field::Fields.SpectralElementField,
+    fields::Fields.ExtrudedFiniteDifferenceField...,
+) where {F, O} =
+    Fields.bycolumn(axes(reduced_field)) do colidx
+        column_fields = map(field -> field[colidx], fields)
+        _column_mapreduce!(fn, op, reduced_field[colidx], column_fields...)
+    end
+
+column_mapreduce_device!(
+    ::ClimaComms.AbstractCPUDevice,
+    fn::F,
+    op::O,
+    reduced_field::Fields.PointField,
+    fields::Fields.FiniteDifferenceField...,
+) where {F, O} = _column_mapreduce!(fn, op, reduced_field, fields...)
+
+function _column_mapreduce!(
+    fn::F,
+    op::O,
+    reduced_field::Fields.PointField,
+    fields::Fields.ColumnField...,
+) where {F, O}
+    space = axes(first(fields))
+    for field in Base.rest(fields)
+        axes(field) === space ||
+            error("All inputs to column_mapreduce must lie on the same space")
+    end
+    first_level = Operators.left_idx(space)
+    last_level = Operators.right_idx(space)
+    # TODO: This code is allocating memory. In particular, even if we comment
+    # out the rest of this function, the first line alone allocates memory.
+    # This problem is not fixed by replacing map with ntuple or unrolled_map.
+    first_level_values =
+        map(field -> (@inbounds Fields.level(field, first_level)[]), fields)
+    reduced_field[] = fn(first_level_values...)
+    for level in (first_level + 1):last_level
+        values = map(field -> (@inbounds Fields.level(field, level)[]), fields)
+        reduced_field[] = op(reduced_field[], fn(values...))
+    end
+    return nothing
+end