Add some CUDA compatibility

Author: Sbozzolo

ext/InterpolationsRegridderExt.jl

@@ -94,14 +94,20 @@ function Regridders.regrid(regridder::InterpolationsRegridder, data, dimensions)
         # horizontal separately
         @warn "Ignoring boundary conditions, implementing Periodic, Flat, Flat"
 
-        return map(regridder.coordinates) do coord
-            ClimaComms.allowscalar(
-                interpolation_3d_z,
-                device,
-                data,
-                dimensions_FT...,
-                totuple(coord)...,
-            )
+        adapted_data = Adapt.adapt(ClimaComms.array_type(regridder.target_space), data)
+        xs, ys, zs = dimensions_FT
+        adapted_xs = Adapt.adapt(ClimaComms.array_type(regridder.target_space), xs)
+        adapted_ys = Adapt.adapt(ClimaComms.array_type(regridder.target_space), ys)
+        adapted_zs = Adapt.adapt(ClimaComms.array_type(regridder.target_space), zs)
+
+        return ClimaComms.allowscalar(ClimaComms.device(regridder.target_space)) do
+            map(regridder.coordinates) do coord
+                interpolation_3d_z(
+                    adapted_data,
+                    adapted_xs, adapted_ys, adapted_zs,
+                    totuple(coord)...,
+                )
+            end
         end
     else
         # Make a linear spline
@@ -157,7 +163,8 @@ function interpolation_3d_z(data, xs, ys, zs, target_x, target_y, target_z)
     # end
 
     # Find nearest neighbors
-    target_x = mod(target_x, maximum(xs) - minimum(xs))
+    x_period = xs[end] - xs[begin]
+    target_x = mod(target_x, x_period)
 
     x_index = searchsortedfirst(xs, target_x)
     y_index = searchsortedfirst(ys, target_y)
@@ -174,15 +181,16 @@ function interpolation_3d_z(data, xs, ys, zs, target_x, target_y, target_z)
     end
 
     # 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)
+
+    z00 = @view zs[x0_index, y0_index, :]
+    z01 = @view zs[x0_index, y1_index, :]
+    z10 = @view zs[x1_index, y0_index, :]
+    z11 = @view zs[x1_index, y1_index, :]
+
+    f00 = linear_interp_z(view(data,x0_index, y0_index, :), z00, target_z)
+    f01 = linear_interp_z(view(data,x0_index, y1_index, :), z01, target_z)
+    f10 = linear_interp_z(view(data,x1_index, y0_index, :), z10, target_z)
+    f11 = linear_interp_z(view(data,x1_index, y1_index, :), z11, target_z)
 
     # Bilinear interpolation in x-y plane
     val = bilinear_interp(