diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 955e6dd9b2..d152dbb517 100755 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -237,12 +237,25 @@ steps: command: "srun julia --color=yes --check-bounds=yes --project=test test/Remapping/distributed_remapping.jl" env: CLIMACOMMS_CONTEXT: "MPI" + CLIMACOMMS_DEVICE: "CPU" agents: slurm_ntasks: 2 - label: "Unit: distributed remapping (1 process)" key: distributed_remapping_1proc command: "julia --color=yes --check-bounds=yes --project=test test/Remapping/distributed_remapping.jl" + env: + CLIMACOMMS_DEVICE: "CPU" + + - label: "Unit: distributed remapping with CUDA" + key: distributed_remapping_gpu + command: "srun julia --color=yes --check-bounds=yes --project=test test/Remapping/distributed_remapping.jl" + env: + CLIMACOMMS_CONTEXT: "MPI" + CLIMACOMMS_DEVICE: "CUDA" + agents: + slurm_ntasks: 2 + slurm_gpus: 1 - label: "Unit: distributed gather" key: unit_distributed_gather4 diff --git a/src/Remapping/Remapping.jl b/src/Remapping/Remapping.jl index 77ff1060e7..107775ff81 100644 --- a/src/Remapping/Remapping.jl +++ b/src/Remapping/Remapping.jl @@ -13,6 +13,7 @@ import ..DataLayouts, ..Hypsography using ..RecursiveApply +using CUDA include("interpolate_array.jl") include("distributed_remapping.jl") diff --git a/src/Remapping/distributed_remapping.jl b/src/Remapping/distributed_remapping.jl index 5818453d41..7617730fd7 100644 --- a/src/Remapping/distributed_remapping.jl +++ b/src/Remapping/distributed_remapping.jl @@ -239,10 +239,17 @@ function interpolate( if length(remapper.target_zcoords) == 0 out_local_array = zeros(FT, size(remapper.local_target_hcoords_bitmask)) - interpolated_values = [ - interpolate_slab(field, Fields.SlabIndex(nothing, gidx), weights) for (gidx, weights) in - zip(remapper.local_indices, remapper.interpolation_coeffs) - ] + + interpolated_values = zeros(FT, length(remapper.local_indices)) + slab_indices = + [Fields.SlabIndex(nothing, gidx) for gidx in remapper.local_indices] + + interpolate_slab!( + interpolated_values, + field, + slab_indices, + remapper.interpolation_coeffs, + ) # out_local_array[remapper.local_target_hcoords_bitmask] returns a view on space we # want to write on diff --git a/src/Remapping/interpolate_array.jl b/src/Remapping/interpolate_array.jl index 6a2467b30f..1c4fb8700d 100644 --- a/src/Remapping/interpolate_array.jl +++ b/src/Remapping/interpolate_array.jl @@ -1,35 +1,209 @@ -function interpolate_slab( +""" + interpolate_slab!(output_array, field, slab_indices, weights) + +Interpolate horizontal field on the given `slab_indices` using the given interpolation +`weights`. + +`interpolate_slab!` interpolates several values at a fixed `z` coordinate. For this reason, +it requires several slab indices and weights. + +""" +interpolate_slab!(output_array, field::Fields.Field, slab_indices, weights) = + interpolate_slab!( + output_array, + field::Fields.Field, + slab_indices, + weights, + ClimaComms.device(field), + ) + + +# CPU kernel for 3D configurations +function interpolate_slab!( + output_array, field::Fields.Field, - slabidx::Fields.SlabIndex, - (I1,)::Tuple{<:AbstractArray}, -) + slab_indices, + weights::AbstractArray{Tuple{A, A}}, + device::ClimaComms.AbstractCPUDevice, +) where {A} space = axes(field) - x = zero(eltype(field)) - QS = Spaces.quadrature_style(space) - Nq = Spaces.Quadratures.degrees_of_freedom(QS) + FT = Spaces.undertype(space) + + for index in 1:length(output_array) + (I1, I2) = weights[index] + Nq1, Nq2 = length(I1), length(I2) + + output_array[index] = zero(FT) - for i in 1:Nq - ij = CartesianIndex((i,)) - x += I1[i] * Operators.get_node(space, field, ij, slabidx) + for j in 1:Nq2, i in 1:Nq1 + ij = CartesianIndex((i, j)) + output_array[index] += + I1[i] * + I2[j] * + Operators.get_node(space, field, ij, slab_indices[index]) + end end - return x end -function interpolate_slab( +# CPU kernel for 2D configurations +function interpolate_slab!( + output_array, field::Fields.Field, - slabidx::Fields.SlabIndex, - (I1, I2)::Tuple{<:AbstractArray, <:AbstractArray}, + slab_indices, + weights::AbstractArray{Tuple{A}}, + device::ClimaComms.AbstractCPUDevice, +) where {A} + space = axes(field) + FT = Spaces.undertype(space) + + for index in 1:length(output_array) + (I1,) = weights[index] + Nq = length(I1) + + output_array[index] = zero(FT) + + for i in 1:Nq + ij = CartesianIndex((i,)) + output_array[index] += + I1[i] * + Operators.get_node(space, field, ij, slab_indices[index]) + end + end +end + +# GPU +function interpolate_slab!( + output_array, + field::Fields.Field, + slab_indices, + weights, + device::ClimaComms.CUDADevice, ) space = axes(field) - x = zero(eltype(field)) - QS = Spaces.quadrature_style(space) - Nq = Spaces.Quadratures.degrees_of_freedom(QS) + FT = Spaces.undertype(space) + + output_cuarray = CuArray(zeros(FT, length(output_array))) + cuweights = CuArray(weights) + cuslab_indices = CuArray(slab_indices) + + nitems = length(output_array) + nthreads, nblocks = Spaces._configure_threadblock(nitems) + + @cuda threads = (nthreads) blocks = (nblocks) interpolate_slab_kernel!( + output_cuarray, + field, + cuslab_indices, + cuweights, + ) + + output_array .= Array(output_cuarray) +end + +# GPU kernel for 3D configurations +function interpolate_slab_kernel!( + output_array, + field, + slab_indices, + weights::AbstractArray{Tuple{A, A}}, +) where {A} + index = threadIdx().x + (blockIdx().x - 1) * blockDim().x + space = axes(field) + FT = Spaces.undertype(space) + + if index <= length(output_array) + I1, I2 = weights[index] + Nq1, Nq2 = length(I1), length(I2) - for j in 1:Nq, i in 1:Nq - ij = CartesianIndex((i, j)) - x += I1[i] * I2[j] * Operators.get_node(space, field, ij, slabidx) + output_array[index] = zero(FT) + + for j in 1:Nq2, i in 1:Nq1 + ij = CartesianIndex((i, j)) + output_array[index] += + I1[i] * + I2[j] * + Operators.get_node(space, field, ij, slab_indices[index]) + end + end + return nothing +end + +# GPU kernel for 2D configurations +function interpolate_slab_kernel!( + output_array, + field, + slab_indices, + weights::AbstractArray{Tuple{A}}, +) where {A} + index = threadIdx().x + (blockIdx().x - 1) * blockDim().x + space = axes(field) + FT = Spaces.undertype(space) + + if index <= length(output_array) + I1, = weights[index] + Nq = length(I1) + + output_array[index] = zero(FT) + + for i in 1:Nq + ij = CartesianIndex((i)) + output_array[index] += + I1[i] * + Operators.get_node(space, field, ij, slab_indices[index]) + end + end + return nothing +end + +""" + vertical_indices_ref_coordinate(space, zcoord) + +Return the vertical indices of the elements below and above `zcoord`. + +Return also the correct reference coordinate `zcoord` for vertical interpolation. +""" +function vertical_indices end + +function vertical_indices_ref_coordinate( + space::Spaces.FaceExtrudedFiniteDifferenceSpace, + zcoord, +) + vert_topology = Spaces.vertical_topology(space) + vert_mesh = vert_topology.mesh + + velem = Meshes.containing_element(vert_mesh, zcoord) + ξ3, = Meshes.reference_coordinates(vert_mesh, velem, zcoord) + v_lo, v_hi = velem - half, velem + half + return v_lo, v_hi, ξ3 +end + +function vertical_indices_ref_coordinate( + space::Spaces.CenterExtrudedFiniteDifferenceSpace, + zcoord, +) + vert_topology = Spaces.vertical_topology(space) + vert_mesh = vert_topology.mesh + Nz = Spaces.nlevels(space) + + velem = Meshes.containing_element(vert_mesh, zcoord) + ξ3, = Meshes.reference_coordinates(vert_mesh, velem, zcoord) + if ξ3 < 0 + if Topologies.isperiodic(Spaces.vertical_topology(space)) + v_lo = mod1(velem - 1, Nz) + else + v_lo = max(velem - 1, 1) + end + v_hi = velem + ξ3 = ξ3 + 1 + else + v_lo = velem + if Topologies.isperiodic(Spaces.vertical_topology(space)) + v_hi = mod1(velem + 1, Nz) + else + v_hi = min(velem + 1, Nz) + end + ξ3 = ξ3 - 1 end - return x + return v_lo, v_hi, ξ3 end """ @@ -37,10 +211,13 @@ end field::Fields.Field, h::Integer, Is::Tuple, - zcoord, + zpts; + fill_value = eltype(field)(NaN) ) -Vertically interpolate the given `field` on `zcoord`. +Vertically interpolate the given `field` on `zpts`. + +`interpolate_slab_level!` interpolates several values at a fixed horizontal coordinate. The field is linearly interpolated across two neighboring vertical elements. @@ -49,47 +226,192 @@ element in a column, no interpolation is performed and the value at the cell cen returned. Effectively, this means that the interpolation is first-order accurate across the column, but zeroth-order accurate close to the boundaries. +Return `fill_value` when the vertical coordinate is negative. + """ -function interpolate_slab_level( +function interpolate_slab_level!( + output_array, field::Fields.Field, h::Integer, Is::Tuple, - zcoord, + vertical_indices_ref_coordinates, +) + device = ClimaComms.device(field) + + interpolate_slab_level!( + output_array, + field, + vertical_indices_ref_coordinates, + h, + Is, + device, + ) +end + +# CPU kernel for 3D configurations +function interpolate_slab_level!( + output_array, + field::Fields.Field, + vidx_ref_coordinates, + h::Integer, + (I1, I2)::Tuple{<:AbstractArray, <:AbstractArray}, + device::ClimaComms.AbstractCPUDevice, ) space = axes(field) - vert_topology = Spaces.vertical_topology(space) - vert_mesh = vert_topology.mesh - Nz = Spaces.nlevels(space) + FT = Spaces.undertype(space) + Nq1, Nq2 = length(I1), length(I2) + + for index in 1:length(vidx_ref_coordinates) + v_lo, v_hi, ξ3 = vidx_ref_coordinates[index] + + f_lo = zero(FT) + f_hi = zero(FT) + + for j in 1:Nq2, i in 1:Nq1 + ij = CartesianIndex((i, j)) + f_lo += + I1[i] * + I2[j] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_lo, h)) + f_hi += + I1[i] * + I2[j] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_hi, h)) + end - velem = Meshes.containing_element(vert_mesh, zcoord) - ξ3, = Meshes.reference_coordinates(vert_mesh, velem, zcoord) - if space isa Spaces.FaceExtrudedFiniteDifferenceSpace - v_lo = velem - half - v_hi = velem + half - elseif space isa Spaces.CenterExtrudedFiniteDifferenceSpace - if ξ3 < 0 - if Topologies.isperiodic(Spaces.vertical_topology(space)) - v_lo = mod1(velem - 1, Nz) - else - v_lo = max(velem - 1, 1) - end - v_hi = velem - ξ3 = ξ3 + 1 - else - v_lo = velem - if Topologies.isperiodic(Spaces.vertical_topology(space)) - v_hi = mod1(velem + 1, Nz) - else - v_hi = min(velem + 1, Nz) - end - ξ3 = ξ3 - 1 + output_array[index] = ((1 - ξ3) * f_lo + (1 + ξ3) * f_hi) / 2 + end +end + +# CPU kernel for 2D configurations +function interpolate_slab_level!( + output_array, + field::Fields.Field, + vidx_ref_coordinates, + h::Integer, + (I1,)::Tuple{<:AbstractArray}, + device::ClimaComms.AbstractCPUDevice, +) + space = axes(field) + FT = Spaces.undertype(space) + Nq = length(I1) + + for index in 1:length(vidx_ref_coordinates) + v_lo, v_hi, ξ3 = vidx_ref_coordinates[index] + + f_lo = zero(FT) + f_hi = zero(FT) + + for i in 1:Nq + ij = CartesianIndex((i,)) + f_lo += + I1[i] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_lo, h)) + f_hi += + I1[i] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_hi, h)) end + output_array[index] = ((1 - ξ3) * f_lo + (1 + ξ3) * f_hi) / 2 end - f_lo = interpolate_slab(field, Fields.SlabIndex(v_lo, h), Is) - f_hi = interpolate_slab(field, Fields.SlabIndex(v_hi, h), Is) - return ((1 - ξ3) * f_lo + (1 + ξ3) * f_hi) / 2 end +# GPU +function interpolate_slab_level!( + output_array, + field::Fields.Field, + vidx_ref_coordinates, + h::Integer, + Is::Tuple, + device::ClimaComms.CUDADevice, +) + cuvidx_ref_coordinates = CuArray(vidx_ref_coordinates) + + output_cuarray = CuArray( + zeros(Spaces.undertype(axes(field)), length(vidx_ref_coordinates)), + ) + + nitems = length(zpts) + nthreads, nblocks = Spaces._configure_threadblock(nitems) + @cuda threads = (nthreads) blocks = (nblocks) interpolate_slab_level_kernel!( + output_cuarray, + field, + cuvidx_ref_coordinates, + h, + Is, + ) + output_array .= Array(output_cuarray) +end + +# GPU kernel for 3D configurations +function interpolate_slab_level_kernel!( + output_array, + field, + vidx_ref_coordinates, + h, + (I1, I2)::Tuple{<:AbstractArray, <:AbstractArray}, +) + index = threadIdx().x + (blockIdx().x - 1) * blockDim().x + space = axes(field) + FT = Spaces.undertype(space) + Nq1, Nq2 = length(I1), length(I2) + + if index <= length(vidx_ref_coordinates) + v_lo, v_hi, ξ3 = vidx_ref_coordinates[index] + + f_lo = zero(FT) + f_hi = zero(FT) + + for j in 1:Nq2, i in 1:Nq1 + ij = CartesianIndex((i, j)) + f_lo += + I1[i] * + I2[j] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_lo, h)) + f_hi += + I1[i] * + I2[j] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_hi, h)) + end + output_array[index] = ((1 - ξ3) * f_lo + (1 + ξ3) * f_hi) / 2 + end + return nothing +end + +# GPU kernel for 2D configurations +function interpolate_slab_level_kernel!( + output_array, + field, + vidx_ref_coordinates, + h, + (I1,)::Tuple{<:AbstractArray}, +) + index = threadIdx().x + (blockIdx().x - 1) * blockDim().x + space = axes(field) + FT = Spaces.undertype(space) + Nq = length(I1) + + if index <= length(vidx_ref_coordinates) + v_lo, v_hi, ξ3 = vidx_ref_coordinates[index] + + f_lo = zero(FT) + f_hi = zero(FT) + + for i in 1:Nq + ij = CartesianIndex((i,)) + f_lo += + I1[i] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_lo, h)) + f_hi += + I1[i] * + Operators.get_node(space, field, ij, Fields.SlabIndex(v_hi, h)) + end + + output_array[index] = ((1 - ξ3) * f_lo + (1 + ξ3) * f_hi) / 2 + end + return nothing +end + + """ interpolate_array(field, xpts, ypts) interpolate_array(field, xpts, ypts, zpts) @@ -128,6 +450,10 @@ function interpolate_array( array = zeros(T, length(xpts), length(zpts)) FT = Spaces.undertype(space) + + vertical_indices_ref_coordinates = + [vertical_indices_ref_coordinate(space, zcoord) for zcoord in zpts] + for (ix, xcoord) in enumerate(xpts) hcoord = xcoord helem = Meshes.containing_element(horz_mesh, hcoord) @@ -136,9 +462,13 @@ function interpolate_array( weights = interpolation_weights(horz_mesh, hcoord, quad_points) h = helem - for (iz, zcoord) in enumerate(zpts) - array[ix, iz] = interpolate_slab_level(field, h, weights, zcoord) - end + interpolate_slab_level!( + view(array, ix, :), + field, + h, + weights, + vertical_indices_ref_coordinates, + ) end return array end @@ -159,6 +489,10 @@ function interpolate_array( array = zeros(T, length(xpts), length(ypts), length(zpts)) FT = Spaces.undertype(space) + + vertical_indices_ref_coordinates = + [vertical_indices_ref_coordinate(space, zcoord) for zcoord in zpts] + for (iy, ycoord) in enumerate(ypts), (ix, xcoord) in enumerate(xpts) hcoord = Geometry.product_coordinates(xcoord, ycoord) helem = Meshes.containing_element(horz_mesh, hcoord) @@ -168,10 +502,13 @@ function interpolate_array( gidx = horz_topology.orderindex[helem] h = gidx - for (iz, zcoord) in enumerate(zpts) - array[ix, iy, iz] = - interpolate_slab_level(field, h, weights, zcoord) - end + interpolate_slab_level!( + view(array, ix, iy, :), + field, + h, + weights, + vertical_indices_ref_coordinates, + ) end return array end @@ -228,12 +565,12 @@ Keyword arguments function interpolate_column( field::Fields.ExtrudedFiniteDifferenceField, zpts, - weights, + Is, gidx; physical_z = false, + fill_value = Spaces.undertype(axes(field))(NaN), ) 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 @@ -242,29 +579,62 @@ function interpolate_column( physical_z = false end - # If we physical_z, we have to move the z coordinates from physical to - # reference ones. + output_array = zeros(Spaces.undertype(space), length(zpts)) + + # If we have physical_z, we have to move the z coordinates from physical to reference + # ones. We also have to deal with the fact that the output_array is not going to be + # fully obtained through interpolation, and the head of it will be simply a collection + # of `fill_values` (up to the surface). When we have topography, we compute the + # reference z and take the positive ones (above the surface). Then, we the top of + # `output_array` with `fill_value` and the rest with interpolated ones. To achieve this, + # we have to make sure that we are passing around views of the same array (as opposed to + # copied of it). 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( + FT = Spaces.undertype(axes(field)) + + # interpolate_slab! takes a vector + z_surface = [zero(FT)] + + interpolate_slab!( + z_surface, space.hypsography.surface, - Fields.SlabIndex(nothing, gidx), - weights, + [Fields.SlabIndex(nothing, gidx)], + [Is], ) + z_surface = z_surface[] 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 + z in zpts if z.z > z_surface ] + + # When zpts = zpts_ref, all the points are above the surface + num_points_below_surface = length(zpts) - length(zpts_ref) + + fill!( + (@view output_array[1:(1 + num_points_below_surface)]), + fill_value, + ) else zpts_ref = zpts + num_points_below_surface = 0 end - return [ - z.z >= 0 ? interpolate_slab_level(field, gidx, weights, z) : FT(NaN) for - z in zpts_ref - ] + vertical_indices_ref_coordinates = + [vertical_indices_ref_coordinate(space, zcoord) for zcoord in zpts_ref] + + interpolate_slab_level!( + (@view output_array[(1 + num_points_below_surface):end]), + field, + gidx, + Is, + vertical_indices_ref_coordinates, + ) + + return output_array end diff --git a/test/Remapping/distributed_remapping.jl b/test/Remapping/distributed_remapping.jl index 52b82fa657..92dff5db08 100644 --- a/test/Remapping/distributed_remapping.jl +++ b/test/Remapping/distributed_remapping.jl @@ -15,7 +15,7 @@ import ClimaCore: using ClimaComms const context = ClimaComms.context() const pid, nprocs = ClimaComms.init(context) -const device = ClimaComms.CPUSingleThreaded() +const device = ClimaComms.device() # log output only from root process logger_stream = ClimaComms.iamroot(context) ? stderr : devnull @@ -24,93 +24,98 @@ atexit() do global_logger(prev_logger) end -@testset "2D extruded" begin - vertdomain = Domains.IntervalDomain( - Geometry.ZPoint(0.0), - Geometry.ZPoint(1000.0); - boundary_tags = (:bottom, :top), - ) - - vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) - verttopo = Topologies.IntervalTopology( - ClimaComms.SingletonCommsContext(device), - vertmesh, - ) - vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) - - horzdomain = Domains.IntervalDomain( - Geometry.XPoint(-500.0) .. Geometry.XPoint(500.0), - periodic = true, - ) - - quad = Spaces.Quadratures.GLL{4}() - horzmesh = Meshes.IntervalMesh(horzdomain, nelems = 10) - horztopology = Topologies.IntervalTopology( - ClimaComms.SingletonCommsContext(device), - horzmesh, - ) - horzspace = Spaces.SpectralElementSpace1D(horztopology, quad) - - hv_center_space = - Spaces.ExtrudedFiniteDifferenceSpace(horzspace, vert_center_space) - - coords = Fields.coordinate_field(hv_center_space) - - xpts = range(-500.0, 500.0, length = 21) - zpts = range(0.0, 1000.0, length = 21) - hcoords = [Geometry.XPoint(x) for x in xpts] - zcoords = [Geometry.ZPoint(z) for z in zpts] - - remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space) - - interp_x = Remapping.interpolate(remapper, coords.x) - if ClimaComms.iamroot(context) - @test interp_x ≈ [x for x in xpts, z in zpts] - end - - interp_z = Remapping.interpolate(remapper, coords.z) - expected_z = [z for x in xpts, z in zpts] - if ClimaComms.iamroot(context) - @test interp_z[:, 2:(end - 1)] ≈ expected_z[:, 2:(end - 1)] - @test interp_z[:, 1] ≈ [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts] - @test interp_z[:, end] ≈ - [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts] - end - - # With hypsography - z_surface = cosd.(Fields.coordinate_field(horzspace).x) .+ 20 - hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( - horzspace, - vert_center_space, - Hypsography.LinearAdaption(z_surface), - ) - - coords = Fields.coordinate_field(hv_center_space) - - remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space;) - - function ref_z(x, z) - z_surface = cosd(x) + 20 - z_top = vertdomain.coord_max.z - return (z - z_surface) / (1 - z_surface / z_top) - end - - val_or_NaN(x, z, val) = ref_z(x, z) >= 0 ? val : NaN - - interp_x = Remapping.interpolate(remapper, coords.x; physical_z = true) - if ClimaComms.iamroot(context) - @test interp_x ≈ [val_or_NaN(x, z, x) for x in xpts, z in zpts] nans = true - end - interp_z = Remapping.interpolate(remapper, coords.z; physical_z = true) - expected_z = [val_or_NaN(x, z, z) for x in xpts, z in zpts] - if ClimaComms.iamroot(context) - @test interp_z[:, 2:(end - 1)] ≈ expected_z[:, 2:(end - 1)] nans = true - - z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 - @test interp_z[:, 1] ≈ [val_or_NaN(x, z_bottom, z_bottom) for x in xpts] nans = true - - z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 - @test interp_z[:, end] ≈ [val_or_NaN(x, z_top, z_top) for x in xpts] rtol = 0.02 nans = true +if !(device isa ClimaComms.CUDADevice) + @testset "2D extruded" begin + vertdomain = Domains.IntervalDomain( + Geometry.ZPoint(0.0), + Geometry.ZPoint(1000.0); + boundary_tags = (:bottom, :top), + ) + + vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) + verttopo = Topologies.IntervalTopology( + ClimaComms.SingletonCommsContext(device), + vertmesh, + ) + vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) + + horzdomain = Domains.IntervalDomain( + Geometry.XPoint(-500.0) .. Geometry.XPoint(500.0), + periodic = true, + ) + + quad = Spaces.Quadratures.GLL{4}() + horzmesh = Meshes.IntervalMesh(horzdomain, nelems = 10) + horztopology = Topologies.IntervalTopology( + ClimaComms.SingletonCommsContext(device), + horzmesh, + ) + horzspace = Spaces.SpectralElementSpace1D(horztopology, quad) + + hv_center_space = + Spaces.ExtrudedFiniteDifferenceSpace(horzspace, vert_center_space) + + coords = Fields.coordinate_field(hv_center_space) + + xpts = range(-500.0, 500.0, length = 21) + zpts = range(0.0, 1000.0, length = 21) + hcoords = [Geometry.XPoint(x) for x in xpts] + zcoords = [Geometry.ZPoint(z) for z in zpts] + + remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space) + + interp_x = Remapping.interpolate(remapper, coords.x) + if ClimaComms.iamroot(context) + @test interp_x ≈ [x for x in xpts, z in zpts] + end + + interp_z = Remapping.interpolate(remapper, coords.z) + expected_z = [z for x in xpts, z in zpts] + if ClimaComms.iamroot(context) + @test interp_z[:, 2:(end - 1)] ≈ expected_z[:, 2:(end - 1)] + @test interp_z[:, 1] ≈ + [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts] + @test interp_z[:, end] ≈ + [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts] + end + + # With hypsography + z_surface = cosd.(Fields.coordinate_field(horzspace).x) .+ 20 + hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( + horzspace, + vert_center_space, + Hypsography.LinearAdaption(z_surface), + ) + + coords = Fields.coordinate_field(hv_center_space) + + remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space;) + + function ref_z(x, z) + z_surface = cosd(x) + 20 + z_top = vertdomain.coord_max.z + return (z - z_surface) / (1 - z_surface / z_top) + end + + val_or_NaN(x, z, val) = ref_z(x, z) >= 0 ? val : NaN + + interp_x = Remapping.interpolate(remapper, coords.x; physical_z = true) + if ClimaComms.iamroot(context) + @test interp_x ≈ [val_or_NaN(x, z, x) for x in xpts, z in zpts] nans = true + end + interp_z = Remapping.interpolate(remapper, coords.z; physical_z = true) + expected_z = [val_or_NaN(x, z, z) for x in xpts, z in zpts] + if ClimaComms.iamroot(context) + @test interp_z[:, 2:(end - 1)] ≈ expected_z[:, 2:(end - 1)] nans = + true + + z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 + @test interp_z[:, 1] ≈ + [val_or_NaN(x, z_bottom, z_bottom) for x in xpts] nans = true + + z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 + @test interp_z[:, end] ≈ [val_or_NaN(x, z_top, z_top) for x in xpts] rtol = 0.02 nans = true + end end end @@ -192,58 +197,60 @@ end @test interp_y ≈ [y for x in xpts, y in ypts] end - - # With hypsography - z_surface = - cosd.(Fields.coordinate_field(horzspace).x) .+ - cosd.(Fields.coordinate_field(horzspace).y) .+ 20 - hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( - horzspace, - vert_center_space, - Hypsography.LinearAdaption(z_surface), - ) - - coords = Fields.coordinate_field(hv_center_space) - - remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space;) - - function ref_z(x, y, z) - z_surface = cosd(x) + cosd(y) + 20 - z_top = vertdomain.coord_max.z - return (z - z_surface) / (1 - z_surface / z_top) - end - - val_or_NaN(x, y, z, val) = ref_z(x, y, z) >= 0 ? val : NaN - - interp_x = Remapping.interpolate(remapper, coords.x; physical_z = true) - if ClimaComms.iamroot(context) - @test interp_x ≈ - [val_or_NaN(x, y, z, x) for x in xpts, y in ypts, z in zpts] nans = - true - end - - interp_y = Remapping.interpolate(remapper, coords.y; physical_z = true) - if ClimaComms.iamroot(context) - @test interp_y ≈ - [val_or_NaN(x, y, z, y) for x in xpts, y in ypts, z in zpts] nans = - true - end - - interp_z = Remapping.interpolate(remapper, coords.z; physical_z = true) - expected_z = [val_or_NaN(x, y, z, z) for x in xpts, y in ypts, z in zpts] - if ClimaComms.iamroot(context) - @test interp_z[:, :, 2:(end - 1)] ≈ expected_z[:, :, 2:(end - 1)] nans = - true - - z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 - @test interp_z[:, :, 1] ≈ - [val_or_NaN(x, y, z_bottom, z_bottom) for x in xpts, y in ypts] nans = - true - - z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 - @test interp_z[:, :, end] ≈ - [val_or_NaN(x, y, z_top, z_top) for x in xpts, y in ypts] rtol = - 0.02 nans = true + if !(device isa ClimaComms.CUDADevice) + # With hypsography + z_surface = + cosd.(Fields.coordinate_field(horzspace).x) .+ + cosd.(Fields.coordinate_field(horzspace).y) .+ 20 + hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( + horzspace, + vert_center_space, + Hypsography.LinearAdaption(z_surface), + ) + + coords = Fields.coordinate_field(hv_center_space) + + remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space;) + + function ref_z(x, y, z) + z_surface = cosd(x) + cosd(y) + 20 + z_top = vertdomain.coord_max.z + return (z - z_surface) / (1 - z_surface / z_top) + end + + val_or_NaN(x, y, z, val) = ref_z(x, y, z) >= 0 ? val : NaN + + interp_x = Remapping.interpolate(remapper, coords.x; physical_z = true) + if ClimaComms.iamroot(context) + @test interp_x ≈ + [val_or_NaN(x, y, z, x) for x in xpts, y in ypts, z in zpts] nans = + true + end + + interp_y = Remapping.interpolate(remapper, coords.y; physical_z = true) + if ClimaComms.iamroot(context) + @test interp_y ≈ + [val_or_NaN(x, y, z, y) for x in xpts, y in ypts, z in zpts] nans = + true + end + + interp_z = Remapping.interpolate(remapper, coords.z; physical_z = true) + expected_z = + [val_or_NaN(x, y, z, z) for x in xpts, y in ypts, z in zpts] + if ClimaComms.iamroot(context) + @test interp_z[:, :, 2:(end - 1)] ≈ expected_z[:, :, 2:(end - 1)] nans = + true + + z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 + @test interp_z[:, :, 1] ≈ [ + val_or_NaN(x, y, z_bottom, z_bottom) for x in xpts, y in ypts + ] nans = true + + z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 + @test interp_z[:, :, end] ≈ + [val_or_NaN(x, y, z_top, z_top) for x in xpts, y in ypts] rtol = + 0.02 nans = true + end end end @@ -328,57 +335,60 @@ end @test interp_y ≈ [y for x in xpts, y in ypts] end - # With hypsography - z_surface = - cosd.(Fields.coordinate_field(horzspace).x) .+ - cosd.(Fields.coordinate_field(horzspace).y) .+ 20 - hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( - horzspace, - vert_center_space, - Hypsography.LinearAdaption(z_surface), - ) - - coords = Fields.coordinate_field(hv_center_space) - - remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space;) - - function ref_z(x, y, z) - z_surface = cosd(x) + cosd(y) + 20 - z_top = vertdomain.coord_max.z - return (z - z_surface) / (1 - z_surface / z_top) - end - - val_or_NaN(x, y, z, val) = ref_z(x, y, z) >= 0 ? val : NaN - - interp_x = Remapping.interpolate(remapper, coords.x; physical_z = true) - if ClimaComms.iamroot(context) - @test interp_x ≈ - [val_or_NaN(x, y, z, x) for x in xpts, y in ypts, z in zpts] nans = - true - end - - interp_y = Remapping.interpolate(remapper, coords.y; physical_z = true) - if ClimaComms.iamroot(context) - @test interp_y ≈ - [val_or_NaN(x, y, z, y) for x in xpts, y in ypts, z in zpts] nans = - true - end - - interp_z = Remapping.interpolate(remapper, coords.z; physical_z = true) - expected_z = [val_or_NaN(x, y, z, z) for x in xpts, y in ypts, z in zpts] - if ClimaComms.iamroot(context) - @test interp_z[:, :, 2:(end - 1)] ≈ expected_z[:, :, 2:(end - 1)] nans = - true - - z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 - @test interp_z[:, :, 1] ≈ - [val_or_NaN(x, y, z_bottom, z_bottom) for x in xpts, y in ypts] nans = - true - - z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 - @test interp_z[:, :, end] ≈ - [val_or_NaN(x, y, z_top, z_top) for x in xpts, y in ypts] rtol = - 0.02 nans = true + if !(device isa ClimaComms.CUDADevice) + # With hypsography + z_surface = + cosd.(Fields.coordinate_field(horzspace).x) .+ + cosd.(Fields.coordinate_field(horzspace).y) .+ 20 + hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( + horzspace, + vert_center_space, + Hypsography.LinearAdaption(z_surface), + ) + + coords = Fields.coordinate_field(hv_center_space) + + remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space;) + + function ref_z(x, y, z) + z_surface = cosd(x) + cosd(y) + 20 + z_top = vertdomain.coord_max.z + return (z - z_surface) / (1 - z_surface / z_top) + end + + val_or_NaN(x, y, z, val) = ref_z(x, y, z) >= 0 ? val : NaN + + interp_x = Remapping.interpolate(remapper, coords.x; physical_z = true) + if ClimaComms.iamroot(context) + @test interp_x ≈ + [val_or_NaN(x, y, z, x) for x in xpts, y in ypts, z in zpts] nans = + true + end + + interp_y = Remapping.interpolate(remapper, coords.y; physical_z = true) + if ClimaComms.iamroot(context) + @test interp_y ≈ + [val_or_NaN(x, y, z, y) for x in xpts, y in ypts, z in zpts] nans = + true + end + + interp_z = Remapping.interpolate(remapper, coords.z; physical_z = true) + expected_z = + [val_or_NaN(x, y, z, z) for x in xpts, y in ypts, z in zpts] + if ClimaComms.iamroot(context) + @test interp_z[:, :, 2:(end - 1)] ≈ expected_z[:, :, 2:(end - 1)] nans = + true + + z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 + @test interp_z[:, :, 1] ≈ [ + val_or_NaN(x, y, z_bottom, z_bottom) for x in xpts, y in ypts + ] nans = true + + z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 + @test interp_z[:, :, end] ≈ + [val_or_NaN(x, y, z_top, z_top) for x in xpts, y in ypts] rtol = + 0.02 nans = true + end end end @@ -457,68 +467,77 @@ end @test interp_sin_lat ≈ [sind(y) for x in longpts, y in latpts] rtol = 0.01 end - # With hypsography - # - # When we have hypsography, we should allow for larger error because elements close to - # the surface are not reconstructed well - z_surface = - cosd.(Fields.coordinate_field(horzspace).lat) .+ - cosd.(Fields.coordinate_field(horzspace).long) .+ 20 - hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( - horzspace, - vert_center_space, - Hypsography.LinearAdaption(z_surface), - ) - - coords = Fields.coordinate_field(hv_center_space) - - remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space) - - function ref_z(x, y, z) - z_surface = cosd(x) + cosd(y) + 20 - z_top = vertdomain.coord_max.z - return (z - z_surface) / (1 - z_surface / z_top) - end - - val_or_NaN(x, y, z, val) = ref_z(x, y, z) >= 0 ? val : NaN - - interp_sin_long = - Remapping.interpolate(remapper, sind.(coords.long); physical_z = true) - - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ [ - val_or_NaN(x, y, z, sind(x)) for x in longpts, y in latpts, - z in zpts - ] rtol = 0.06 nans = true - end - - interp_sin_lat = - Remapping.interpolate(remapper, sind.(coords.lat); physical_z = true) - if ClimaComms.iamroot(context) - # We have to add atol because we have an element that is identically zero - @test interp_sin_lat ≈ [ - val_or_NaN(x, y, z, sind(y)) for x in longpts, y in latpts, - z in zpts - ] rtol = 0.01 nans = true atol = 1e-15 - end - - interp_z = Remapping.interpolate(remapper, coords.z, physical_z = true) - - FT = Spaces.undertype(hv_center_space) - - expected_z = [z for x in longpts, y in latpts, z in zpts] - - if ClimaComms.iamroot(context) - @test interp_z[:, :, 2:(end - 1)] ≈ expected_z[:, :, 2:(end - 1)] - - z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 - @test interp_z[:, :, 1] ≈ [ - val_or_NaN(x, y, z_bottom, z_bottom) for x in longpts, y in latpts - ] nans = true - - z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 - @test interp_z[:, :, end] ≈ - [val_or_NaN(x, y, z_top, z_top) for x in longpts, y in latpts] nans = - true rtol = 0.02 + if !(device isa ClimaComms.CUDADevice) + # With hypsography + # + # When we have hypsography, we should allow for larger error because elements close to + # the surface are not reconstructed well + z_surface = + cosd.(Fields.coordinate_field(horzspace).lat) .+ + cosd.(Fields.coordinate_field(horzspace).long) .+ 20 + hv_center_space = Spaces.ExtrudedFiniteDifferenceSpace( + horzspace, + vert_center_space, + Hypsography.LinearAdaption(z_surface), + ) + + coords = Fields.coordinate_field(hv_center_space) + + remapper = Remapping.Remapper(hcoords, zcoords, hv_center_space) + + function ref_z(x, y, z) + z_surface = cosd(x) + cosd(y) + 20 + z_top = vertdomain.coord_max.z + return (z - z_surface) / (1 - z_surface / z_top) + end + + val_or_NaN(x, y, z, val) = ref_z(x, y, z) >= 0 ? val : NaN + + interp_sin_long = Remapping.interpolate( + remapper, + sind.(coords.long); + physical_z = true, + ) + + if ClimaComms.iamroot(context) + @test interp_sin_long ≈ [ + val_or_NaN(x, y, z, sind(x)) for x in longpts, y in latpts, + z in zpts + ] rtol = 0.06 nans = true + end + + interp_sin_lat = Remapping.interpolate( + remapper, + sind.(coords.lat); + physical_z = true, + ) + if ClimaComms.iamroot(context) + # We have to add atol because we have an element that is identically zero + @test interp_sin_lat ≈ [ + val_or_NaN(x, y, z, sind(y)) for x in longpts, y in latpts, + z in zpts + ] rtol = 0.01 nans = true atol = 1e-15 + end + + interp_z = Remapping.interpolate(remapper, coords.z, physical_z = true) + + FT = Spaces.undertype(hv_center_space) + + expected_z = [z for x in longpts, y in latpts, z in zpts] + + if ClimaComms.iamroot(context) + @test interp_z[:, :, 2:(end - 1)] ≈ expected_z[:, :, 2:(end - 1)] + + z_bottom = 1000.0 * (0 / 30 + 1 / 30) / 2 + @test interp_z[:, :, 1] ≈ [ + val_or_NaN(x, y, z_bottom, z_bottom) for x in longpts, + y in latpts + ] nans = true + + z_top = 1000.0 * (29 / 30 + 30 / 30) / 2 + @test interp_z[:, :, end] ≈ + [val_or_NaN(x, y, z_top, z_top) for x in longpts, y in latpts] nans = + true rtol = 0.02 + end end end