Various improvements to Remapper

This commit: - extends the remapper to work with purely vertical spaces - adds documentation about remapping - adds user-facing functions to directly remap a field Co-authored-by: charleskawczynski <[email protected]>
CliMA · Dec 10, 2024 · 2ad3e18 · 2ad3e18
1 parent 340603b
commit 2ad3e18
Show file tree

Hide file tree

Showing 10 changed files with 854 additions and 83 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -4,6 +4,24 @@ 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.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 +30,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)

diff --git a/docs/make.jl b/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"],

diff --git a/docs/src/remapping.md b/docs/src/remapping.md
@@ -0,0 +1,159 @@
+# 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 not 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.
+
+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
+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 (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`](@ref) 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)
+```
+
+!!! note
+
+    `interpolate` allocates new arrays and has some internal
+     type-instability, `interpolate!` is non-allocating and type-stable. 
+     When using `interpolate!`, the destination has to be defined on the 
+     root process and has to be `nothing` for the other MPI processes.
+
+
+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.
diff --git a/ext/cuda/remapping_distributed.jl b/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},

diff --git a/src/Remapping/Remapping.jl b/src/Remapping/Remapping.jl
@@ -5,6 +5,7 @@ using LinearAlgebra, StaticArrays
 import ClimaComms
 import ..DataLayouts,
     ..Geometry,
+    ..Domains,
     ..Spaces,
     ..Grids,
     ..Topologies,