Add support for interpolation with 3D altitude

Sbozzolo · Sbozzolo · commit 921d9d139ace · 2024-12-16T17:59:00.000-08:00
This commit adds support for interpolating datasets of the form
(lon, lat, z), with (lon, lat) being 1D variables, and z being z(lon,
lat) (a 3D variable). This is accomplished by splitting interpolation
into two: first, we perform linear interpolation along the vertical
direction, and the bilinear in the horizontal.

More specifically, I interpolate the nodal points around the target
point onto the same altitude. Then, I perform bilinear interpolation
with the interpolated values.

This code is not polished and does not run on GPU. I am not very happy
about how it looks, but it does the job
diff --git a/ext/InterpolationsRegridderExt.jl b/ext/InterpolationsRegridderExt.jl
@@ -80,21 +80,172 @@ Regrid the given data as defined on the given dimensions to the `target_space` i
 This function is allocating.
 """
 function Regridders.regrid(regridder::InterpolationsRegridder, data, dimensions)
+    # TODO: There is room for improvement in this function...
+
     FT = ClimaCore.Spaces.undertype(regridder.target_space)
     dimensions_FT = map(d -> FT.(d), dimensions)
 
-    # Make a linear spline
-    itp = Intp.extrapolate(
-        Intp.interpolate(dimensions_FT, FT.(data), Intp.Gridded(Intp.Linear())),
-        regridder.extrapolation_bc,
+    coordinates = ClimaCore.Fields.coordinate_field(regridder.target_space)
+
+    has_3d_z = length(size(last(dimensions))) == 3
+    if eltype(coordinates) <: ClimaCore.Geometry.LatLongZPoint && has_3d_z
+        # If we have 3D altitudes, we do linear in the vertical and bilinear
+        # horizontal separately
+
+        # TODO: Support other boundary conditions
+        any(x -> x != Intp.Throw(), regridder.extrapolation_bc) && error(
+            "Only Throw() boundary condition currently implemented $(regridder.extrapolation_bc)",
+        )
+        return map(regridder.coordinates) do coord
+            interpolation_3d_z(data, dimensions_FT..., totuple(coord)...)
+        end
+    else
+        # Make a linear spline
+        itp = Intp.extrapolate(
+            Intp.interpolate(
+                dimensions_FT,
+                FT.(data),
+                Intp.Gridded(Intp.Linear()),
+            ),
+            regridder.extrapolation_bc,
+        )
+
+        # Move it to GPU (if needed)
+        gpuitp = Adapt.adapt(ClimaComms.array_type(regridder.target_space), itp)
+
+        return map(regridder.coordinates) do coord
+            gpuitp(totuple(coord)...)
+        end
+    end
+end
+
+"""
+    interpolation_3d_z(data, xs, ys, zs, target_x, target_y, target_z)
+
+Perform bilinear + vertical interpolation on a 3D dataset.
+
+This function first performs linear interpolation along the z-axis at the four
+corners of the cell containing the target (x, y) point. Then, it performs
+bilinear interpolation in the x-y plane using the z-interpolated values.
+
+# Arguments
+- `data`: A 3D array of data values.
+- `xs`: A vector of x-coordinates corresponding to the first dimension of `data`.
+- `ys`: A vector of y-coordinates corresponding to the second dimension of `data`.
+- `zs`: A 3D array of z-coordinates. `zs[i, j, :]` provides the z-coordinates for the data point `data[i, j, :]`.
+- `target_x`: The x-coordinate of the target point.
+- `target_y`: The y-coordinate of the target point.
+- `target_z`: The z-coordinate of the target point.
+"""
+function interpolation_3d_z(data, xs, ys, zs, target_x, target_y, target_z)
+    # Check boundaries
+    if target_x < xs[begin] || target_x > xs[end]
+        error(
+            "target_x is out of bounds: $(target_x) not in [$(xs[1]), $(xs[end])]",
+        )
+    end
+    if target_y < ys[begin] || target_y > ys[end]
+        error(
+            "target_y is out of bounds: $(target_y) not in [$(ys[1]), $(ys[end])]",
+        )
+    end
+
+    # Find nearest neighbors
+    x_index = searchsortedfirst(xs, target_x)
+    y_index = searchsortedfirst(ys, target_y)
+
+    x0_index = x_index == 1 ? x_index : x_index - 1
+    x1_index = x0_index + 1
+    y0_index = y_index == 1 ? y_index : y_index - 1
+    y1_index = y0_index + 1
+
+    # Interpolate in z-direction
+    z00 = zs[x0_index, y0_index, :]
+    z01 = zs[x0_index, y1_index, :]
+    z10 = zs[x1_index, y0_index, :]
+    z11 = zs[x1_index, y1_index, :]
+
+    f00 = linear_interp_z(data[x0_index, y0_index, :], z00, target_z)
+    f01 = linear_interp_z(data[x0_index, y1_index, :], z01, target_z)
+    f10 = linear_interp_z(data[x1_index, y0_index, :], z10, target_z)
+    f11 = linear_interp_z(data[x1_index, y1_index, :], z11, target_z)
+
+    # Bilinear interpolation in x-y plane
+    return bilinear_interp(
+        f00,
+        f01,
+        f10,
+        f11,
+        xs[x0_index],
+        xs[x1_index],
+        ys[y0_index],
+        ys[y1_index],
+        target_x,
+        target_y,
     )
+end
+
+"""
+    linear_interp_z(f, z, target_z)
+
+Perform linear interpolation along the z-axis.
 
-    # Move it to GPU (if needed)
-    gpuitp = Adapt.adapt(ClimaComms.array_type(regridder.target_space), itp)
+# Arguments
+- `f`: A vector of function values corresponding to the z-coordinates in `z`.
+- `z`: A vector of z-coordinates.
+- `target_z`: The z-coordinate at which to interpolate.
+
+# Returns
+The linearly interpolated value at `target_z`.
+"""
+function linear_interp_z(f, z, target_z)
+    if target_z < z[begin] || target_z > z[end]
+        error(
+            "target_z is out of bounds: $(target_z) not in [$(z[1]), $(z[end])]",
+        )
+    end
 
-    return map(regridder.coordinates) do coord
-        gpuitp(totuple(coord)...)
+    index = searchsortedfirst(z, target_z)
+    # Handle edge cases for index
+    if index == 1
+        z0 = z[index]
+        z1 = z[index + 1]
+        f0 = f[index]
+        f1 = f[index + 1]
+    else
+        z0 = z[index - 1]
+        z1 = z[index]
+        f0 = f[index - 1]
+        f1 = f[index]
     end
+
+    return f0 + (target_z - z0) / (z1 - z0) * (f1 - f0)
+end
+
+"""
+    bilinear_interp(f00, f01, f10, f11, x0, x1, y0, y1, target_x, target_y)
+
+Perform bilinear interpolation on a 2D plane.
+
+# Arguments
+- `f00`: Function value at (x0, y0).
+- `f01`: Function value at (x0, y1).
+- `f10`: Function value at (x1, y0).
+- `f11`: Function value at (x1, y1).
+- `x0`: x-coordinate of the first corner.
+- `x1`: x-coordinate of the second corner.
+- `y0`: y-coordinate of the first corner.
+- `y1`: y-coordinate of the second corner.
+- `target_x`: The x-coordinate of the target point.
+- `target_y`: The y-coordinate of the target point.
+"""
+function bilinear_interp(f00, f01, f10, f11, x0, x1, y0, y1, target_x, target_y)
+    return (
+        (x1 - target_x) * (y1 - target_y) / ((x1 - x0) * (y1 - y0)) * f00 +
+        (x1 - target_x) * (target_y - y0) / ((x1 - x0) * (y1 - y0)) * f01 +
+        (target_x - x0) * (y1 - target_y) / ((x1 - x0) * (y1 - y0)) * f10 +
+        (target_x - x0) * (target_y - y0) / ((x1 - x0) * (y1 - y0)) * f11
+    )
 end
 
 end
diff --git a/test/interpolations_regridder.jl b/test/interpolations_regridder.jl
@@ -0,0 +1,107 @@
+using Test
+import ClimaUtilities
+import Interpolations
+import ClimaUtilities: Regridders
+import ClimaComms
+import ClimaCore
+
+linear_interp_z = Base.get_extension(ClimaUtilities, :ClimaUtilitiesClimaCoreInterpolationsExt).InterpolationsRegridderExt.linear_interp_z
+bilinear_interp = Base.get_extension(ClimaUtilities, :ClimaUtilitiesClimaCoreInterpolationsExt).InterpolationsRegridderExt.bilinear_interp
+interpolation_3d_z = Base.get_extension(ClimaUtilities, :ClimaUtilitiesClimaCoreInterpolationsExt).InterpolationsRegridderExt.interpolation_3d_z
+
+const context = ClimaComms.context()
+ClimaComms.init(context)
+
+include("TestTools.jl")
+
+@testset "Interpolation Tests" begin
+    @testset "linear_interp_z" begin
+        f = [1.0, 3.0, 5.0]
+        z = [10.0, 20.0, 30.0]
+        @test linear_interp_z(f, z, 15.0) ≈ 2.0
+        @test linear_interp_z(f, z, 25.0) ≈ 4.0
+        @test linear_interp_z(f, z, 10.0) ≈ 1.0
+        @test linear_interp_z(f, z, 30.0) ≈ 5.0
+
+        # Out of bounds
+        f = [1.0, 3.0]
+        z = [10.0, 20.0]
+        @test_throws ErrorException linear_interp_z(f, z, 5.0)
+        @test_throws ErrorException linear_interp_z(f, z, 25.0)
+
+
+        # One point
+        f = [2.5]
+        z = [15.0]
+        @test_throws ErrorException linear_interp_z(f, z, 10.0)
+        @test_throws ErrorException linear_interp_z(f, z, 20.0)
+
+        # Non uniform spacing
+        f = [2.0, 4.0, 8.0]
+        z = [1.0, 3.0, 7.0]
+        @test linear_interp_z(f,z, 1.0) ≈ 2.0
+        @test linear_interp_z(f,z, 3.0) ≈ 4.0
+        @test linear_interp_z(f,z, 7.0) ≈ 8.0
+        @test linear_interp_z(f,z, 2.0) ≈ 3.0
+        @test linear_interp_z(f,z, 5.0) ≈ 6.0
+    end
+
+    @testset "interpolation_3d_z" begin
+        # Test cases for the main 3D interpolation function
+
+        # Create some sample data
+        xs = [1.0, 2.0, 3.0]
+        ys = [4.0, 5.0, 6.0]
+        zs = zeros(3, 3, 8)
+        for k in 1:8
+            for j in 1:3
+                for i in 1:3
+                    zs[i, j, k] = i + j + k
+                end
+            end
+        end
+        data = reshape(1:(3*3*8), 3, 3, 8)
+
+        # Exact point
+        @test interpolation_3d_z(data, xs, ys, zs, xs[1], ys[1], zs[1, 1, 4]) ≈ data[1, 1, 4]
+
+        # Interpolated point
+        @test interpolation_3d_z(data, xs, ys, zs, 2.5, 5.5, 7.5) ≈ 20.5
+
+        # Out of bounds
+        @test_throws ErrorException interpolation_3d_z(data, xs, ys, zs, 0.5, 4.5, 3.5)
+        @test_throws ErrorException interpolation_3d_z(data, xs, ys, zs, 2.5, 4.5, 4.5)
+    end
+
+    @testset "Regrid" begin
+
+        lon, lat, z =
+            collect(-180.0:1:180), collect(-90.0:1:90), collect(0.0:1.0:100.0)
+        size3D = (361, 181, 101)
+        data_z3D = zeros(size3D)
+
+        for i in 1:length(lon)
+            for j in 1:length(lat)
+                data_z3D[i, j, :] .= z
+            end
+        end
+        dimensions3D = (lon, lat, data_z3D)
+
+        FT = Float64
+        spaces = make_spherical_space(FT; context)
+        hv_center_space = spaces.hybrid
+        extrapolation_bc = (
+            Interpolations.Throw(),
+            Interpolations.Throw(),
+            Interpolations.Throw(),
+        )
+        reg_hv = Regridders.InterpolationsRegridder(
+            hv_center_space;
+            extrapolation_bc,
+        )
+        regridded_z = Regridders.regrid(reg_hv, data_z3D, dimensions3D)
+        @test maximum(ClimaCore.Fields.level(regridded_z, 2)) ≈ 0.15
+        @test minimum(ClimaCore.Fields.level(regridded_z, 2)) ≈ 0.15
+
+    end
+end
