Various improvements to Remapper

[skip ci]

Co-authored-by: charleskawczynski <[email protected]>

NEWS.md

@@ -4,6 +4,11 @@ ClimaCore.jl Release Notes
 main
 -------
 
+### ![][badge-✨feature/enhancement] Various improvements to `Remapper`
+
+
+v0.14.20
+---------
  - 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]()
@@ -12,8 +17,6 @@ main
    - [SliceXZSpace]()
    - [RectangleXYSpace]()
 
-v0.14.20
-
  - We've added new convenience constructors for grids PR [1848](https://github.com/CliMA/ClimaCore.jl/pull/1848). Here are links to the new constructors:
    - [ExtrudedCubedSphereGrid](https://github.com/CliMA/ClimaCore.jl/blob/cbb193042fac3b4bef33251fbc0f232427bfe506/src/CommonGrids/CommonGrids.jl#L85-L144)
    - [CubedSphereGrid](https://github.com/CliMA/ClimaCore.jl/blob/cbb193042fac3b4bef33251fbc0f232427bfe506/src/CommonGrids/CommonGrids.jl#L200-L235)

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,183 @@
+# 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 going
+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 n t require communication, but they
+are not as accurate as conservative remappers, especially with large elements
+with sharp gradients. These remappers are optimally suited for diagnostics and
+plots.
+
+# `Remapper`
+
+The main 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 into 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
+definition of `field`. To obtain such coordinates, you can call the
+[`Remapping.default_target_hcoords`](@ref) and
+[`Remapping.default_target_zcoords`](@ref) functions. These functions return an
+Array with the coordinates over which interpolation will occur. These arrays are
+of `Geometry.Point`s.
+
+[`Remapping.interpolate`](@ref) allocates new output arrays. As such, it is not
+suitable for performance-critical applications. [`Remapping.interpolate!`](@ref)
+performs interpolation in-place. When using the in-place version`, the
+`dest`ination has to be the same array type as the device in use (e.g.,
+`CuArray` for CUDA runs).
+
+For performance-critical applications, it is preferable to a
+[`Remapping.Remapper`](@ref) and use it directly. A [`Remapping.Remapper`](@ref)
+is an object that is tied to a specified `Space` and can interpolate scalar
+`Field`s defined on that space on a defined target grid. The grid does not need
+to be regular, but it has to be defined as a Cartesian product between some
+horizontal and vertical coordinates. 
+
+For example, let us create a remapper to
+
+
+```julia
+Remapper(space, target_hcoords, target_zcoords, buffer_length = 1)
+Remapper(space; target_hcoords, target_zcoords, buffer_length = 1)
+Remapper(space, target_hcoords; buffer_length = 1)
+Remapper(space, target_zcoords; buffer_length = 1)
+```
+
+Return a `Remapper` responsible for interpolating any `Field` defined on the
+given `space` to the Cartesian product of `target_hcoords` with
+`target_zcoords`.
+
+`target_zcoords` can be `nothing` for interpolation on horizontal spaces.
+Similarly, `target_hcoords` can be `nothing` for interpolation on vertical
+spaces.
+
+Multiple `fields` can be interpolated at the same time as long as they share the
+underlying space.
+
+`Remapper` is the main argument to the `interpolate` function.
+
+### Keyword arguments
+
+`buffer_length` is size of the internal buffer in the Remapper 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`.
+
+## Interpolation
+
+```julia
+interpolate(remapper::Remapper, fields)
+interpolate!(dest, remapper::Remapper, fields)
+```
+
+Interpolate the given `field`(s) as prescribed by `remapper`.
+
+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`.
+
+This call mutates the internal (private) state of the `remapper`.
+
+`interpolate!` writes the output to the given `dest`ination. `dest` is expected
+to be defined on the root process and to be `nothing` for the other processes.
+
+**Note:** `interpolate` allocates new arrays and has some internal
+type-instability, `interpolate!` is non-allocating and type-stable.
+
+### 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, :, :, :]
+```
+
+## Convenience Interpolation
+
+```julia
+interpolate(field::ClimaCore.Fields;
+           hresolution = 180,
+           resolution = 50,
+           target_hcoords = default_target_hcoords(space; hresolution),
+           target_zcoords = default_target_vcoords(space; vresolution)
+           )
+```
+
+Interpolate the given fields on the Cartesian product of `target_hcoords` with `target_zcoords` (if not empty).
+
+Coordinates have to be `ClimaCore.Geometry.Points`.
+
+**Note:** do not use this method when performance is important. Instead, define a `Remapper` and call `interpolate(remapper, fields)`. Different `Field`s defined on the same `Space` can share a `Remapper`, so that interpolation can be optimized.
+
+### 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 simply
+
+```julia
+julia> interpolate(field, hcoords, zcoords)
+```
+
+Coordinates can be specified:
+
+```julia
+julia> longpts = range(-180.0, 180.0, 21)
+julia> latpts = range(-80.0, 80.0, 21)
+julia> zpts = range(0.0, 1000.0, 21)
+
+julia> hcoords = [Geometry.LatLongPoint(lat, long) for long in longpts, lat in latpts]
+julia> zcoords = [Geometry.ZPoint(z) for z in zpts]
+
+julia> interpolate(field, hcoords, zcoords)
+```
+```
+
+### Conservative remapping with `TempestRemap`
+
+This section hasn't been written yet. You can help by writing it.

ext/cuda/remapping_distributed.jl

@@ -131,6 +131,40 @@ 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,
+    ::Nothing,
+)
+    # TODO: Check the memory access pattern. This was not optimized and likely inefficient!
+    num_fields = length(field_values)
+
+    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,