Extend Remapper to FiniteDifferenceSpaces, update interpolate interface

NEWS.md

@@ -4,20 +4,28 @@ ClimaCore.jl Release Notes
 main
 -------
 
+### ![][badge-✨feature/enhancement] Various improvements to `Remapper` [2060](https://github.com/CliMA/ClimaCore.jl/pull/2060)
+
+The `ClimaCore.Remapping` module received two improvements. First, `Remapper` is
+now compatible with purely vertical `Space`s (performing a linear
+interpolation), making it compatible with column setups. Second, a new set of
+simplified interpolation functions are provided.
+
+Now, interpolating a `Field` `field` is as easy as
+```julia
+import ClimaCore.Remapping: interpolate
+output_array = interpolate(field)
+```
+The target coordinates are automatically determined, but can also be customized.
+Refer to the [documentation](https://clima.github.io/ClimaCore.jl/dev/remapping/)
+for more information.
+
 v0.14.21
 --------
 
  - Support for new TVD limiters were added, PR [1662]
    (https://github.com/CliMA/ClimaCore.jl/pull/1662).
 
- - We've added new convenience constructors for spaces PR [2082](https://github.com/CliMA/ClimaCore.jl/pull/2082). Here are links to the new constructors:
-   - [ExtrudedCubedSphereSpace]()
-   - [CubedSphereSpace]()
-   - [ColumnSpace]()
-   - [Box3DSpace]()
-   - [SliceXZSpace]()
-   - [RectangleXYSpace]()
-
 ### ![][badge-🐛bugfix] Bug fixes
 
 - Fixed writing/reading purely vertical spaces. PR [2102](https://github.com/CliMA/ClimaCore.jl/pull/2102)
@@ -30,6 +38,14 @@ v0.14.21
 face-centered (and viceversa). These functions only work for vertical and
 extruded spaces.
 
+- We've added new convenience constructors for spaces PR [2082](https://github.com/CliMA/ClimaCore.jl/pull/2082). Here are links to the new constructors:
+  - [ExtrudedCubedSphereSpace](https://clima.github.io/ClimaCore.jl/dev/api/#ClimaCore.CommonSpaces.ExtrudedCubedSphereSpace)
+  - [CubedSphereSpace](https://clima.github.io/ClimaCore.jl/dev/api/#ClimaCore.CommonSpaces.CubedSphereSpace)
+  - [ColumnSpace](https://clima.github.io/ClimaCore.jl/dev/api/#ClimaCore.CommonSpaces.ColumnSpace)
+  - [Box3DSpace](https://clima.github.io/ClimaCore.jl/dev/api/#ClimaCore.CommonSpaces.Box3DSpace)
+  - [SliceXZSpace](https://clima.github.io/ClimaCore.jl/dev/api/#ClimaCore.CommonSpaces.SliceXZSpace)
+  - [RectangleXYSpace](https://clima.github.io/ClimaCore.jl/dev/api/#ClimaCore.CommonSpaces.RectangleXYSpace)
+
 v0.14.20
 --------
 

docs/make.jl

@@ -77,6 +77,7 @@ withenv("GKSwstype" => "nul") do
             "Installation and How-to Guides" => "installation_instructions.md",
             "Geometry" => "geometry.md",
             "Operators" => "operators.md",
+            "Remapping" => "remapping.md",
             "MatrixFields" => "matrix_fields.md",
             "API" => "api.md",
             "Developer docs" => ["Performance tips" => "performance_tips.md"],

docs/src/remapping.md

@@ -0,0 +1,153 @@
+# Remapping to regular grids
+
+`ClimaCore` horizontal domains are spectral elements. Points are not distributed
+uniformly within an element, and elements are also not necessarily organized in
+a simple way. For these reasons, remapping to regular grids becomes a
+fundamental operations when inspecting the simulation output. In this section,
+we describe the remappers currently available in `ClimaCore`.
+
+Broadly speaking, we can classify remappers in two categories: conservative, and
+non-conservative. Conservative remappers preserve areas (and masses) when
+interpolating from the spectral grid to Cartesian ones. Conservative remappers
+are non-local operations (meaning that they require communication between
+different elements) and are more expensive, so they are typically reserved to
+operations where physical conservation is important (e.g., exchange between
+component models in a coupled simulation). On the other hand, non-conservative
+remappers are local to an element and faster to evaluate, which makes them
+suitable to operations like diagnostics and plotting, where having perfect
+physical conservation is not as important.
+
+## Non-conservative remapping
+
+Non-conservative remappers are fast and do not require communication, but they
+are not as accurate as conservative remappers, especially with large elements
+with sharp gradients. These remappers are better suited for diagnostics and
+plots.
+
+The main non-conservative remapper currently implemented utilizes a Lagrange
+interpolation with the barycentric formula in [Berrut2004], equation (3.2), for
+the horizontal interpolation. Vertical interpolation is linear except in the
+boundary elements where it is 0th order.
+
+### Quick start
+
+Assuming you have a `ClimaCore` `Field` with name `field`, the simplest way to
+interpolate onto a uniform grid is with
+```julia
+julia> import ClimaCore.Remapping
+julia> Remapping.interpolate(field)
+```
+
+This will return an `Array` (or a `CuArray`) with the `field` interpolated on
+some uniform grid that is automatically determined based on the `Space` of the
+given `field`. To obtain such coordinates, you can call the
+`Remapping.default_target_hcoords` and `Remapping.default_target_zcoords`
+functions. These functions return an `Array` with the coordinates over which
+interpolation will occur. These arrays are of type `Geometry.Point`s.
+
+`ClimaCore.Remapping.interpolate` allocates new output arrays. As such, it is
+not suitable for performance-critical applications.
+`ClimaCore.Remapping.interpolate!` performs interpolation in-place. When using
+the in-place version`, the `dest`ination has to have the same array type as the
+device in use (e.g., `CuArray` for CUDA runs) and has to be `nothing` for
+non-root processes. For performance-critical applications, it is preferable to a
+`ClimaCore.Remapping.Remapper` and use it directly (see next Section).
+
+#### Example
+
+Given `field`, a `Field` defined on a cubed sphere.
+
+By default, a target uniform grid is chosen (with resolution `hresolution` and
+`vresolution`), so remapping is
+```julia
+interpolated_array = interpolate(field, hcoords, zcoords)
+```
+Coordinates can be specified:
+
+```julia
+longpts = range(-180.0, 180.0, 21)
+latpts = range(-80.0, 80.0, 21)
+zpts = range(0.0, 1000.0, 21)
+
+hcoords = [Geometry.LatLongPoint(lat, long) for long in longpts, lat in latpts]
+zcoords = [Geometry.ZPoint(z) for z in zpts]
+
+interpolated_array = interpolate(field, hcoords, zcoords)
+```
+The output is defined on the Cartesian product of `hcoords` with `zcoords`.
+
+If the default target coordinates are being used, it is possible to broadcast
+`ClimaCore.Geometry.components` to extract them as a vector of tuples (and then
+broadcast `getindex` to extract the respective coordinates as vectors).
+
+### The `Remapper` object
+
+A `Remapping.Remapper` is an object that is tied to a specified `Space` and can
+interpolate scalar `Field`s defined on that space onto a predefined target grid.
+The grid does not have to be regular, but it has to be defined as a Cartesian
+product between some horizontal and vertical coordinates (meaning, for each
+horizontal point, there is a fixed column of vertical coordinates).
+
+Let us create our first remapper, assuming we have `space` defined on the surface of the sphere
+```julia
+import ClimaCore.Geometry: LatLongPoint, ZPoint
+import ClimaCore.Remapping: Remapper
+
+hcoords = [Geometry.LatLongPoint(lat, long) for long in -180.:180., lat in -90.:90.]
+remapper = Remapper(space, target_hcoords)
+```
+This `remapper` object knows can interpolate `Field`s defined on `space` with the same `interpolate` and `interpolate!` functions.
+```julia
+import ClimaCore.Fields: coordinate_field
+import ClimaCore.Remapping: interpolate, interpolate!
+
+example_field = coordinate_field(space)
+interpolated_array = interpolate(remapper, example_field)
+
+# Interpolate in place
+interpolate!(interpolated_array, remapper, example_field)
+```
+
+Multiple fields defined on the same space can be interpolate at the same time
+```julia
+example_field2 = cosd.(example_field)
+interpolated_arrays = interpolate(remapper, [example_field, example_field2])
+```
+
+When interpolating multiple fields, greater performance can be achieved by
+creating the `Remapper` with a larger internal buffer to store intermediate
+values for interpolation. Effectively, this controls how many fields can be
+remapped simultaneously in `interpolate`. When more fields than `buffer_length`
+are passed, the remapper will batch the work in sizes of `buffer_length`. The optimal number of fields passed is the `buffer_length` of the `remapper`. If
+more fields are passed, the `remapper` will batch work with size up to its
+`buffer_length`.
+
+#### Example
+
+Given `field1`,`field2`, two `Field` defined on a cubed sphere.
+
+```julia
+longpts = range(-180.0, 180.0, 21)
+latpts = range(-80.0, 80.0, 21)
+zpts = range(0.0, 1000.0, 21)
+
+hcoords = [Geometry.LatLongPoint(lat, long) for long in longpts, lat in latpts]
+zcoords = [Geometry.ZPoint(z) for z in zpts]
+
+space = axes(field1)
+
+remapper = Remapper(space, hcoords, zcoords)
+
+int1 = interpolate(remapper, field1)
+int2 = interpolate(remapper, field2)
+
+# Or
+int12 = interpolate(remapper, [field1, field2])
+# With int1 = int12[1, :, :, :]
+```
+
+## Conservative remapping with `TempestRemap`
+
+This section hasn't been written yet. You can help by writing it.
+
+TODO: finish writing this section.

ext/cuda/remapping_distributed.jl

@@ -17,8 +17,14 @@ function _set_interpolated_values_device!(
 )
     # FIXME: Avoid allocation of tuple
     field_values = tuple(map(f -> Fields.field_values(f), fields)...)
-    nblocks, _ = size(interpolation_matrix[1])
-    nthreads = length(vert_interpolation_weights)
+
+    purely_vertical_space = isnothing(interpolation_matrix)
+    num_horizontal_points =
+        purely_vertical_space ? 1 : prod(size(local_horiz_indices))
+    num_points = num_horizontal_points * length(vert_interpolation_weights)
+    max_threads = 256
+    nthreads = min(num_points, max_threads)
+    nblocks = cld(num_points, nthreads)
     args = (
         out,
         interpolation_matrix,
@@ -131,6 +137,41 @@ function set_interpolated_values_kernel!(
     return nothing
 end
 
+# GPU, vertical case
+function set_interpolated_values_kernel!(
+    out::AbstractArray,
+    ::Nothing,
+    ::Nothing,
+    vert_interpolation_weights,
+    vert_bounding_indices,
+    field_values,
+)
+    # TODO: Check the memory access pattern. This was not optimized and likely inefficient!
+    num_fields = length(field_values)
+    num_vert = length(vert_bounding_indices)
+
+    vindex = (blockIdx().y - Int32(1)) * blockDim().y + threadIdx().y
+    findex = (blockIdx().z - Int32(1)) * blockDim().z + threadIdx().z
+
+    totalThreadsY = gridDim().y * blockDim().y
+    totalThreadsZ = gridDim().z * blockDim().z
+
+    CI = CartesianIndex
+    for j in vindex:totalThreadsY:num_vert
+        v_lo, v_hi = vert_bounding_indices[j]
+        A, B = vert_interpolation_weights[j]
+        for k in findex:totalThreadsZ:num_fields
+            if j ≤ num_vert && k ≤ num_fields
+                out[j, k] = (
+                    A * field_values[k][CI(1, 1, 1, v_lo, 1)] +
+                    B * field_values[k][CI(1, 1, 1, v_hi, 1)]
+                )
+            end
+        end
+    end
+    return nothing
+end
+
 function _set_interpolated_values_device!(
     out::AbstractArray,
     fields::AbstractArray{<:Fields.Field},

src/Remapping/Remapping.jl

@@ -5,6 +5,7 @@ using LinearAlgebra, StaticArrays
 import ClimaComms
 import ..DataLayouts,
     ..Geometry,
+    ..Domains,
     ..Spaces,
     ..Grids,
     ..Topologies,