Merge #1497

1497: Add support for topography in `distributed_remapping` r=simonbyrne a=Sbozzolo This PR adds a new mode to the `interpolate` function for distributed remapping, `interpolate_topography`. When `interpolate_topography` is `true`, the target `zcoords` are assumed to be with respect to the mean sea level as opposed to with respect to the surface. `zcoords` that are below the surface are filled with `NaN`s. Closes #1488 Co-authored-by: Gabriele Bozzola <[email protected]>
CliMA · Oct 18, 2023 · 5fb3b4b · 5fb3b4b
2 parents 32be46e + 8cdd86b
commit 5fb3b4b
Show file tree

Hide file tree

Showing 5 changed files with 312 additions and 15 deletions.
diff --git a/src/Hypsography/Hypsography.jl b/src/Hypsography/Hypsography.jl
@@ -23,10 +23,10 @@ end
 
 Locate vertical levels using an exponential function between the surface field and the top
 of the domain, using the method of [Schar2002](@cite). This method is modified
-such no warping is applied above some user defined parameter 0 ≤ ηₕ < 1.0, where the lower and upper 
-bounds represent the domain bottom and top respectively. `s` governs the decay rate. 
-If the decay-scale is poorly specified (i.e., `s * zₜ` is lower than the maximum 
-surface elevation), a warning is thrown and `s` is adjusted such that it `szₜ > maximum(z_surface)`. 
+such no warping is applied above some user defined parameter 0 ≤ ηₕ < 1.0, where the lower and upper
+bounds represent the domain bottom and top respectively. `s` governs the decay rate.
+If the decay-scale is poorly specified (i.e., `s * zₜ` is lower than the maximum
+surface elevation), a warning is thrown and `s` is adjusted such that it `szₜ > maximum(z_surface)`.
 """
 struct SLEVEAdaption{F <: Fields.Field, FT <: Real} <: HypsographyAdaption
     # Union can be removed once deprecation removed.
@@ -201,15 +201,15 @@ end
 
 Option for 2nd order diffusive smoothing of generated terrain.
 Mutate (smooth) a given elevation profile `f` before assigning the surface
-elevation to the `HypsographyAdaption` type. A spectral second-order diffusion 
+elevation to the `HypsographyAdaption` type. A spectral second-order diffusion
 operator is applied with forward-Euler updates to generate
 profiles for each new iteration. Steps to generate smoothed terrain (
-represented as a ClimaCore Field) are as follows: 
+represented as a ClimaCore Field) are as follows:
 - Compute discrete elevation profile f
 - Compute diffuse_surface_elevation!(f, κ, iter). f is mutated.
 - Define `Hypsography.LinearAdaption(f)`
 - Define `ExtrudedFiniteDifferenceSpace` with new surface elevation.
-Default diffusion parameters are appropriate for spherical arrangements. 
+Default diffusion parameters are appropriate for spherical arrangements.
 For `zmax-zsfc` == 𝒪(10^4), κ == 𝒪(10^8), dt == 𝒪(10⁻¹).
 """
 function diffuse_surface_elevation!(

diff --git a/src/Remapping/Remapping.jl b/src/Remapping/Remapping.jl
@@ -4,7 +4,13 @@ using LinearAlgebra, StaticArrays
 
 import ClimaComms
 import ..DataLayouts,
-    ..Geometry, ..Spaces, ..Topologies, ..Meshes, ..Operators, ..Fields
+    ..Geometry,
+    ..Spaces,
+    ..Topologies,
+    ..Meshes,
+    ..Operators,
+    ..Fields,
+    ..Hypsography
 
 using ..RecursiveApply
 

diff --git a/src/Remapping/distributed_remapping.jl b/src/Remapping/distributed_remapping.jl
@@ -124,7 +124,6 @@ The `Remapper` is designed to not be tied to any particular `Field`. You can use
 `Remapper` for any `Field` as long as they are all defined on the same `topology`.
 
 `Remapper` is the main argument to the `interpolate` function.
-
 """
 function Remapper(target_hcoords, target_zcoords, space)
 
@@ -192,10 +191,19 @@ end
 
 
 """
-   interpolate(remapper, field)
+   interpolate(remapper, field; physical_z = false)
 
 Interpolate the given `field` as prescribed by `remapper`.
 
+
+Keyword arguments
+==================
+
+- `physical_z`: When `true`, interpolate to the physical z coordinates, taking into account
+                hypsography. If `false` (the default), interpolation will be based on the
+                reference z coordinate, i.e. will correspond to constant model levels.
+                `NaN`s are returned for values that are below the surface.
+
 Example
 ========
 
@@ -218,7 +226,11 @@ int2 = interpolate(remapper, field2)
 
 ```
 """
-function interpolate(remapper::Remapper, field::T) where {T <: Fields.Field}
+function interpolate(
+    remapper::Remapper,
+    field::T;
+    physical_z = false,
+) where {T <: Fields.Field}
 
     axes(field) == remapper.space ||
         error("Field is defined on a different space than remapper")
@@ -257,7 +269,8 @@ function interpolate(remapper::Remapper, field::T) where {T <: Fields.Field}
                     field,
                     remapper.target_zcoords,
                     weights,
-                    gidx,
+                    gidx;
+                    physical_z,
                 ) for (gidx, weights) in
                 zip(remapper.local_indices, remapper.interpolation_coeffs)
             )'

diff --git a/src/Remapping/interpolate_array.jl b/src/Remapping/interpolate_array.jl
@@ -32,6 +32,24 @@ function interpolate_slab(
     return x
 end
 
+"""
+    interpolate_slab_level(
+                           field::Fields.Field,
+                           h::Integer,
+                           Is::Tuple,
+                           zcoord,
+                           )
+
+Vertically interpolate the given `field` on `zcoord`.
+
+The field is linearly interpolated across two neighboring vertical elements.
+
+For centered-valued fields, if `zcoord` is in the top (bottom) half of a top (bottom)
+element in a column, no interpolation is performed and the value at the cell center is
+returned. Effectively, this means that the interpolation is first-order accurate across the
+column, but zeroth-order accurate close to the boundaries.
+
+"""
 function interpolate_slab_level(
     field::Fields.Field,
     h::Integer,
@@ -197,12 +215,56 @@ and global index in the topology.
 
 The coefficients `weights` are computed with `Spaces.Quadratures.interpolation_matrix`.
 See also `interpolate_array`.
+
+Keyword arguments
+==================
+
+- `physical_z`: When `true`, the given `zpts` are interpreted as "from mean sea
+                            level" (instead of "from surface") and hypsography is
+                            interpolated. `NaN`s are returned for values that are below the
+                            surface.
+
 """
 function interpolate_column(
     field::Fields.ExtrudedFiniteDifferenceField,
     zpts,
     weights,
-    gidx,
+    gidx;
+    physical_z = false,
 )
-    return [interpolate_slab_level(field, gidx, weights, z) for z in zpts]
+    space = axes(field)
+    FT = Spaces.undertype(space)
+
+    # When we don't have hypsography, there is no notion of "interpolating hypsography". In
+    # this case, the reference vertical points coincide with the physical ones. Setting
+    # physical_z = false ensures that zpts_ref = zpts
+    if space.hypsography isa Spaces.Flat
+        physical_z = false
+    end
+
+    # If we physical_z, we have to move the z coordinates from physical to
+    # reference ones.
+    if physical_z
+        # We are hardcoding the transformation from Hypsography.LinearAdaption
+        space.hypsography isa Hypsography.LinearAdaption ||
+            error("Cannot interpolate $(space.hypsography) hypsography")
+
+        z_surface = interpolate_slab(
+            space.hypsography.surface,
+            Fields.SlabIndex(nothing, gidx),
+            weights,
+        )
+        z_top = Spaces.vertical_topology(space).mesh.domain.coord_max.z
+        zpts_ref = [
+            Geometry.ZPoint((z.z - z_surface) / (1 - z_surface / z_top)) for
+            z in zpts
+        ]
+    else
+        zpts_ref = zpts
+    end
+
+    return [
+        z.z >= 0 ? interpolate_slab_level(field, gidx, weights, z) : FT(NaN) for
+        z in zpts_ref
+    ]
 end