From 368ac55327b679ba927abeabcd486e7baf6401e0 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Tue, 29 Oct 2024 14:57:52 -0700 Subject: [PATCH] work in progress [skip ci] --- NEWS.md | 4 + docs/make.jl | 1 + docs/src/remapping.md | 119 ++ src/Remapping/distributed_remapping.jl | 168 +-- test/Remapping/distributed_remapping.jl | 1353 ++++++++++++----------- 5 files changed, 914 insertions(+), 731 deletions(-) create mode 100644 docs/src/remapping.md diff --git a/NEWS.md b/NEWS.md index ce8a0b9412..a1cb4d2067 100644 --- a/NEWS.md +++ b/NEWS.md @@ -6,6 +6,10 @@ main - Fixed world-age issue on Julia 1.11 issue [Julia#54780](https://github.com/JuliaLang/julia/issues/54780), PR [#2034](https://github.com/CliMA/ClimaCore.jl/pull/2034). +### ![][badge-✨feature/enhancement] Various improvements to `Remapper` + + + v0.14.19 ------- diff --git a/docs/make.jl b/docs/make.jl index a2496aed68..66bed38583 100644 --- 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 new file mode 100644 index 0000000000..168064afd2 --- /dev/null +++ b/docs/src/remapping.md @@ -0,0 +1,119 @@ +# Remapping to regular grids + +`ClimaCore` horizontal domains are spectral elements. + + +### Non-conservative remapping + +# Remapper + +## Constructors + +```julia +Remapper(space, target_hcoords, target_zcoords, buffer_length = 1) +Remapper(space; target_hcoords, target_zcoords, buffer_length = 1) +Remapper(space, target_hcoords; buffer_length = 1) +Remapper(space, target_zcoords; buffer_length = 1) +``` + +Return a `Remapper` responsible for interpolating any `Field` defined on the given `space` to the Cartesian product of `target_hcoords` with `target_zcoords`. + +`target_zcoords` can be `nothing` for interpolation on horizontal spaces. Similarly, `target_hcoords` can be `nothing` for interpolation on vertical spaces. + +The `Remapper` is designed to not be tied to any particular `Field`. You can use the same `Remapper` for any `Field` as long as they are all defined on the same `topology`. + +`Remapper` is the main argument to the `interpolate` function. + +### Keyword arguments + +`buffer_length` is size of the internal buffer in the Remapper 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`. + +## Interpolation + +```julia +interpolate(remapper::Remapper, fields) +interpolate!(dest, remapper::Remapper, fields) +``` + +Interpolate the given `field`(s) as prescribed by `remapper`. + +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`. + +This call mutates the internal (private) state of the `remapper`. + +Horizontally, interpolation is performed with the barycentric formula in [Berrut2004], equation (3.2). Vertical interpolation is linear except in the boundary elements where it is 0th order. + +`interpolate!` writes the output to the given `dest`ination. `dest` is expected to be defined on the root process and to be `nothing` for the other processes. + +**Note:** `interpolate` allocates new arrays and has some internal type-instability, `interpolate!` is non-allocating and type-stable. + +When using `interpolate!`, the `dest`ination has to be the same array type as the device in use (e.g., `CuArray` for CUDA runs). + +### 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, :, :, :] +``` + +## Convenience Interpolation + +```julia +interpolate(field::ClimaCore.Fields; + hresolution = 180, + resolution = 50, + target_hcoords = default_target_hcoords(space; hresolution), + target_zcoords = default_target_vcoords(space; vresolution) + ) +``` + +Interpolate the given fields on the Cartesian product of `target_hcoords` with `target_zcoords` (if not empty). + +Coordinates have to be `ClimaCore.Geometry.Points`. + +**Note:** do not use this method when performance is important. Instead, define a `Remapper` and call `interpolate(remapper, fields)`. Different `Field`s defined on the same `Space` can share a `Remapper`, so that interpolation can be optimized. + +### 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 simply + +```julia +julia> interpolate(field, hcoords, zcoords) +``` + +Coordinates can be specified: + +```julia +julia> longpts = range(-180.0, 180.0, 21) +julia> latpts = range(-80.0, 80.0, 21) +julia> zpts = range(0.0, 1000.0, 21) + +julia> hcoords = [Geometry.LatLongPoint(lat, long) for long in longpts, lat in latpts] +julia> zcoords = [Geometry.ZPoint(z) for z in zpts] + +julia> interpolate(field, hcoords, zcoords) +``` +``` + +### Conservative remapping with `TempestRemap` + +This section hasn't been written yet. You can help by writing it. diff --git a/src/Remapping/distributed_remapping.jl b/src/Remapping/distributed_remapping.jl index d7516c7c55..fc03e26d29 100644 --- a/src/Remapping/distributed_remapping.jl +++ b/src/Remapping/distributed_remapping.jl @@ -885,6 +885,9 @@ function interpolate(remapper::Remapper, fields) if !isa_vertical_space # For spaces with an horizontal component, reshape the output so that it is a nice grid. _apply_mpi_bitmask!(remapper, num_fields) + else + # For purely vertical spaces, just move to _interpolated_values + remapper._interpolated_values .= remapper._local_interpolated_values end # Finally, we have to send all the _interpolated_values to root and sum them up to @@ -899,12 +902,83 @@ function interpolate(remapper::Remapper, fields) interpolated_values end +# dest has to be allowed to be nothing because interpolation happens only on the root +# process +function interpolate!( + dest::Union{Nothing, <:AbstractArray}, + remapper::Remapper, + fields, +) + only_one_field = fields isa Fields.Field + if only_one_field + fields = [fields] + end + isa_vertical_space = remapper.space isa Spaces.FiniteDifferenceSpace + + if !isnothing(dest) + # !isnothing(dest) means that this is the root process, in this case, the size have + # to match (ignoring the buffer_length) + dest_size = only_one_field ? size(dest) : size(dest)[1:(end - 1)] + + dest_size == size(remapper._interpolated_values)[1:(end - 1)] || error( + "Destination array is not compatible with remapper (size mismatch)", + ) + + expected_array_type = + ClimaComms.array_type(ClimaComms.device(remapper.comms_ctx)) + + found_type = nameof(typeof(dest)) + + dest isa expected_array_type || + error("dest is a $found_type, expected $expected_array_type") + end + index_field_begin, index_field_end = + 1, min(length(fields), remapper.buffer_length) + + while true + num_fields = 1 + index_field_end - index_field_begin + + # Reset interpolated_values. This is needed because we collect distributed results + # with a + reduction. + _reset_interpolated_values!(remapper) + # Perform the interpolations (horizontal and vertical) + _set_interpolated_values!( + remapper, + view(fields, index_field_begin:index_field_end), + ) + + if !isa_vertical_space + # For spaces with an horizontal component, reshape the output so that it is a nice grid. + _apply_mpi_bitmask!(remapper, num_fields) + else + # For purely vertical spaces, just move to _interpolated_values + remapper._interpolated_values .= remapper._local_interpolated_values + end + + # Finally, we have to send all the _interpolated_values to root and sum them up to + # obtain the final answer. + _collect_interpolated_values!( + dest, + remapper, + index_field_begin, + index_field_end; + only_one_field, + ) + + index_field_end != length(fields) || break + index_field_begin = index_field_begin + remapper.buffer_length + index_field_end = + min(length(fields), index_field_end + remapper.buffer_length) + end + return nothing +end + """ interpolate(field::ClimaCore.Fields; hresolution = 180, resolution = 50, - target_hcoords = get_target_hcoords(space; hresolution), - target_zcoords = get_target_cords(space; vresolution) + target_hcoords = default_target_hcoords(space; hresolution), + target_zcoords = default_target_vcoords(space; vresolution) ) Interpolate the given fields on the Cartesian product of `target_hcoords` with @@ -943,8 +1017,8 @@ function interpolate( field::Fields.Field; vresolution = 50, hresolution = 100, - target_hcoords = get_target_hcoords(axes(field); hresolution), - target_zcoords = get_target_zcoords(axes(field); vresolution), + target_hcoords = default_target_hcoords(axes(field); hresolution), + target_zcoords = default_target_zcoords(axes(field); vresolution), ) return interpolate(field, axes(field); hresolution, vresolution) end @@ -954,29 +1028,29 @@ function interpolate(field::Fields.Field, target_hcoords, target_zcoords) return interpolate(remapper, field) end -function get_target_hcoords(space::Spaces.AbstractSpace; hresolution) - return get_target_hcoords(Spaces.horizontal_space(space); hresolution) +function default_target_hcoords(space::Spaces.AbstractSpace; hresolution) + return default_target_hcoords(Spaces.horizontal_space(space); hresolution) end -function get_target_hcoords( +function default_target_hcoords( space::Spaces.SpectralElementSpace2D; hresolution = 180, ) topology = Spaces.topology(space) mesh = topology.mesh domain = Meshes.domain(mesh) - PT1 = typeof(domain.interval1.coord_min) - PT2 = typeof(domain.interval2.coord_min) + PointType1 = typeof(domain.interval1.coord_min) + PointType2 = typeof(domain.interval2.coord_min) x1min = Geometry.component(domain.interval1.coord_min, 1) x2min = Geometry.component(domain.interval2.coord_min, 1) x1max = Geometry.component(domain.interval1.coord_max, 1) x2max = Geometry.component(domain.interval2.coord_max, 1) - x1 = map(PT1, range(x1min, x1max; length = hresolution)) - x2 = map(PT2, range(x2min, x2max; length = hresolution)) + x1 = map(PointType1, range(x1min, x1max; length = hresolution)) + x2 = map(PointType2, range(x2min, x2max; length = hresolution)) return Base.Iterators.product((x1, x2)) end -function get_target_hcoords(space::Spaces.SpectralElementSpace1D; hresolution = 180) +function default_target_hcoords(space::Spaces.SpectralElementSpace1D; hresolution = 180) topology = Spaces.topology(space) mesh = topology.mesh domain = Meshes.domain(mesh) @@ -986,76 +1060,8 @@ function get_target_hcoords(space::Spaces.SpectralElementSpace1D; hresolution = return PointType.(range(x1min, x1max; length = hresolution)) end -function get_target_zcoords(space; vresolution = 50) +function default_target_zcoords(space; vresolution = 50) return Geometry.ZPoint.( range(z_min(space), z_max(space); length = vresolution) ) end - -# dest has to be allowed to be nothing because interpolation happens only on the root -# process -function interpolate!( - dest::Union{Nothing, <:AbstractArray}, - remapper::Remapper, - fields, -) - only_one_field = fields isa Fields.Field - if only_one_field - fields = [fields] - end - isa_vertical_space = remapper.space isa Spaces.FiniteDifferenceSpace - - if !isnothing(dest) - # !isnothing(dest) means that this is the root process, in this case, the size have - # to match (ignoring the buffer_length) - dest_size = only_one_field ? size(dest) : size(dest)[1:(end - 1)] - - dest_size == size(remapper._interpolated_values)[1:(end - 1)] || error( - "Destination array is not compatible with remapper (size mismatch)", - ) - - expected_array_type = - ClimaComms.array_type(ClimaComms.device(remapper.comms_ctx)) - - found_type = nameof(typeof(dest)) - - dest isa expected_array_type || - error("dest is a $found_type, expected $expected_array_type") - end - index_field_begin, index_field_end = - 1, min(length(fields), remapper.buffer_length) - - while true - num_fields = 1 + index_field_end - index_field_begin - - # Reset interpolated_values. This is needed because we collect distributed results - # with a + reduction. - _reset_interpolated_values!(remapper) - # Perform the interpolations (horizontal and vertical) - _set_interpolated_values!( - remapper, - view(fields, index_field_begin:index_field_end), - ) - - if !isa_vertical_space - # For spaces with an horizontal component, reshape the output so that it is a nice grid. - _apply_mpi_bitmask!(remapper, num_fields) - end - - # Finally, we have to send all the _interpolated_values to root and sum them up to - # obtain the final answer. - _collect_interpolated_values!( - dest, - remapper, - index_field_begin, - index_field_end; - only_one_field, - ) - - index_field_end != length(fields) || break - index_field_begin = index_field_begin + remapper.buffer_length - index_field_end = - min(length(fields), index_field_end + remapper.buffer_length) - end - return nothing -end diff --git a/test/Remapping/distributed_remapping.jl b/test/Remapping/distributed_remapping.jl index 6eccde622a..74f716d50f 100644 --- a/test/Remapping/distributed_remapping.jl +++ b/test/Remapping/distributed_remapping.jl @@ -31,638 +31,638 @@ atexit() do global_logger(prev_logger) end -@testset "Utils" begin - # batched_ranges(num_fields, buffer_length) - @test Remapping.batched_ranges(1, 1) == [1:1] - @test Remapping.batched_ranges(1, 2) == [1:1] - @test Remapping.batched_ranges(2, 2) == [1:2] - @test Remapping.batched_ranges(3, 2) == [1:2, 3:3] -end - -on_gpu = device isa ClimaComms.CUDADevice -broken = false - -if !on_gpu - @testset "2D extruded" begin - vertdomain = Domains.IntervalDomain( - Geometry.ZPoint(0.0), - Geometry.ZPoint(1000.0); - boundary_names = (: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 = 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( - hv_center_space, - hcoords, - zcoords, - buffer_length = 2, - ) - - interp_x = Remapping.interpolate(remapper, coords.x) - interp_x2 = Remapping.interpolate(coords.x, hcoords, zcoords) - if ClimaComms.iamroot(context) - @test Array(interp_x) ≈ [x for x in xpts, z in zpts] - @test Array(interp_x2) ≈ [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 Array(interp_z[:, 2:(end - 1)]) ≈ expected_z[:, 2:(end - 1)] - @test Array(interp_z[:, 1]) ≈ - [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts] - @test Array(interp_z[:, end]) ≈ - [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts] - end - - # Remapping two fields - interp_xx = Remapping.interpolate(remapper, [coords.x, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ interp_xx[:, :, 1] - @test interp_x ≈ interp_xx[:, :, 2] - end - - # Remapping three fields (more than the buffer length) - interp_xxx = - Remapping.interpolate(remapper, [coords.x, coords.x, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ interp_xxx[:, :, 1] - @test interp_x ≈ interp_xxx[:, :, 2] - @test interp_x ≈ interp_xxx[:, :, 3] - end - - # Remapping in-place one field - dest = ArrayType(zeros(21, 21)) - Remapping.interpolate!(dest, remapper, coords.x) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest - end - - # Two fields - dest = ArrayType(zeros(21, 21, 2)) - Remapping.interpolate!(dest, remapper, [coords.x, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, 1] - @test interp_x ≈ dest[:, :, 2] - end - - # Three fields (more than buffer length) - dest = ArrayType(zeros(21, 21, 3)) - Remapping.interpolate!(dest, remapper, [coords.x, coords.x, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, 1] - @test interp_x ≈ dest[:, :, 2] - @test interp_x ≈ dest[:, :, 3] - end - end -end - -@testset "3D box" begin - vertdomain = Domains.IntervalDomain( - Geometry.ZPoint(0.0), - Geometry.ZPoint(1000.0); - boundary_names = (:bottom, :top), - ) - - vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) - verttopo = Topologies.IntervalTopology( - ClimaComms.SingletonCommsContext(device), - vertmesh, - ) - vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) - - horzdomain = Domains.RectangleDomain( - Geometry.XPoint(-500.0) .. Geometry.XPoint(500.0), - Geometry.YPoint(-500.0) .. Geometry.YPoint(500.0), - x1periodic = true, - x2periodic = true, - ) - - quad = Quadratures.GLL{4}() - horzmesh = Meshes.RectilinearMesh(horzdomain, 10, 10) - horztopology = Topologies.Topology2D( - ClimaComms.SingletonCommsContext(device), - horzmesh, - ) - horzspace = Spaces.SpectralElementSpace2D(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) - ypts = range(-500.0, 500.0, length = 21) - zpts = range(0.0, 1000.0, length = 21) - hcoords = [Geometry.XYPoint(x, y) for x in xpts, y in ypts] - zcoords = [Geometry.ZPoint(z) for z in zpts] - - remapper = - Remapping.Remapper(hv_center_space, hcoords, zcoords, buffer_length = 2) - - interp_x = Remapping.interpolate(remapper, coords.x) - interp_x2 = Remapping.interpolate(coords.x, hcoords, zcoords) - if ClimaComms.iamroot(context) - @test Array(interp_x) ≈ [x for x in xpts, y in ypts, z in zpts] - @test Array(interp_x2) ≈ [x for x in xpts, y in ypts, z in zpts] - end - - interp_y = Remapping.interpolate(remapper, coords.y) - if ClimaComms.iamroot(context) - @test Array(interp_y) ≈ [y for x in xpts, y in ypts, z in zpts] - end - - interp_z = Remapping.interpolate(remapper, coords.z) - expected_z = [z for x in xpts, y in ypts, z in zpts] - if ClimaComms.iamroot(context) - @test Array(interp_z[:, :, 2:(end - 1)]) ≈ expected_z[:, :, 2:(end - 1)] - @test Array(interp_z[:, :, 1]) ≈ - [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts, y in ypts] - @test Array(interp_z[:, :, end]) ≈ - [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts, y in ypts] - end - - # Remapping two fields - interp_xy = Remapping.interpolate(remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_x ≈ interp_xy[:, :, :, 1] - @test interp_y ≈ interp_xy[:, :, :, 2] - end - # Remapping three fields (more than the buffer length) - interp_xyx = Remapping.interpolate(remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ interp_xyx[:, :, :, 1] - @test interp_y ≈ interp_xyx[:, :, :, 2] - @test interp_x ≈ interp_xyx[:, :, :, 3] - end - - # Remapping in-place one field - dest = ArrayType(zeros(21, 21, 21)) - Remapping.interpolate!(dest, remapper, coords.x) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest - end - - # Two fields - dest = ArrayType(zeros(21, 21, 21, 2)) - Remapping.interpolate!(dest, remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, :, 1] - @test interp_y ≈ dest[:, :, :, 2] - end - - # Three fields (more than buffer length) - dest = ArrayType(zeros(21, 21, 21, 3)) - Remapping.interpolate!(dest, remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, :, 1] - @test interp_y ≈ dest[:, :, :, 2] - @test interp_x ≈ dest[:, :, :, 3] - end - - - # Horizontal space - horiz_space = Spaces.horizontal_space(hv_center_space) - horiz_remapper = Remapping.Remapper(horiz_space, hcoords, buffer_length = 2) - - coords = Fields.coordinate_field(horiz_space) - - interp_x = Remapping.interpolate(horiz_remapper, coords.x) - # Only root has the final result - if ClimaComms.iamroot(context) - @test Array(interp_x) ≈ [x for x in xpts, y in ypts] - end - - interp_y = Remapping.interpolate(horiz_remapper, coords.y) - if ClimaComms.iamroot(context) - @test Array(interp_y) ≈ [y for x in xpts, y in ypts] - end - - # Two fields - interp_xy = Remapping.interpolate(horiz_remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_xy[:, :, 1] ≈ interp_x - @test interp_xy[:, :, 2] ≈ interp_y - end - - # Three fields - interp_xyx = - Remapping.interpolate(horiz_remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_xyx[:, :, 1] ≈ interp_x - @test interp_xyx[:, :, 2] ≈ interp_y - @test interp_xyx[:, :, 3] ≈ interp_x - end - - # Remapping in-place one field - # - # We have to change remapper for GPU to make sure it works for when have have only one - # field - dest = ArrayType(zeros(21, 21)) - Remapping.interpolate!(dest, horiz_remapper, coords.x) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest - end - - - # Two fields - dest = ArrayType(zeros(21, 21, 2)) - Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, 1] - @test interp_y ≈ dest[:, :, 2] - end - - # Three fields (more than buffer length) - dest = ArrayType(zeros(21, 21, 3)) - Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, 1] - @test interp_y ≈ dest[:, :, 2] - @test interp_x ≈ dest[:, :, 3] - end - -end - - -@testset "3D box - space filling curve" begin - vertdomain = Domains.IntervalDomain( - Geometry.ZPoint(0.0), - Geometry.ZPoint(1000.0); - boundary_names = (:bottom, :top), - ) - - vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) - verttopo = Topologies.IntervalTopology( - ClimaComms.SingletonCommsContext(device), - vertmesh, - ) - vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) - - horzdomain = Domains.RectangleDomain( - Geometry.XPoint(-500.0) .. Geometry.XPoint(500.0), - Geometry.YPoint(-500.0) .. Geometry.YPoint(500.0), - x1periodic = true, - x2periodic = true, - ) - - quad = Quadratures.GLL{4}() - horzmesh = Meshes.RectilinearMesh(horzdomain, 10, 10) - horztopology = Topologies.Topology2D( - ClimaComms.SingletonCommsContext(device), - horzmesh, - Topologies.spacefillingcurve(horzmesh), - ) - horzspace = Spaces.SpectralElementSpace2D(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) - ypts = range(-500.0, 500.0, length = 21) - zpts = range(0.0, 1000.0, length = 21) - hcoords = [Geometry.XYPoint(x, y) for x in xpts, y in ypts] - zcoords = [Geometry.ZPoint(z) for z in zpts] - - remapper = - Remapping.Remapper(hv_center_space, hcoords, zcoords, buffer_length = 2) - - interp_x = Remapping.interpolate(remapper, coords.x) - interp_x2 = Remapping.interpolate(coords.x, hcoords, zcoords) - if ClimaComms.iamroot(context) - @test Array(interp_x) ≈ [x for x in xpts, y in ypts, z in zpts] - @test Array(interp_x2) ≈ [x for x in xpts, y in ypts, z in zpts] - end - - interp_y = Remapping.interpolate(remapper, coords.y) - if ClimaComms.iamroot(context) - @test Array(interp_y) ≈ [y for x in xpts, y in ypts, z in zpts] - end - - interp_z = Remapping.interpolate(remapper, coords.z) - expected_z = [z for x in xpts, y in ypts, z in zpts] - if ClimaComms.iamroot(context) - @test Array(interp_z[:, :, 2:(end - 1)]) ≈ expected_z[:, :, 2:(end - 1)] - @test Array(interp_z[:, :, 1]) ≈ - [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts, y in ypts] - @test Array(interp_z[:, :, end]) ≈ - [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts, y in ypts] - end - - # Remapping two fields - interp_xy = Remapping.interpolate(remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_x ≈ interp_xy[:, :, :, 1] - @test interp_y ≈ interp_xy[:, :, :, 2] - end - # Remapping three fields (more than the buffer length) - interp_xyx = Remapping.interpolate(remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ interp_xyx[:, :, :, 1] - @test interp_y ≈ interp_xyx[:, :, :, 2] - @test interp_x ≈ interp_xyx[:, :, :, 3] - end - - # Remapping in-place one field - dest = ArrayType(zeros(21, 21, 21)) - Remapping.interpolate!(dest, remapper, coords.x) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest - end - - # Two fields - dest = ArrayType(zeros(21, 21, 21, 2)) - Remapping.interpolate!(dest, remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, :, 1] - @test interp_y ≈ dest[:, :, :, 2] - end - - # Three fields (more than buffer length) - dest = ArrayType(zeros(21, 21, 21, 3)) - Remapping.interpolate!(dest, remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, :, 1] - @test interp_y ≈ dest[:, :, :, 2] - @test interp_x ≈ dest[:, :, :, 3] - end - - - # Horizontal space - horiz_space = Spaces.horizontal_space(hv_center_space) - horiz_remapper = Remapping.Remapper(horiz_space, hcoords, buffer_length = 2) - - coords = Fields.coordinate_field(horiz_space) - - interp_x = Remapping.interpolate(horiz_remapper, coords.x) - # Only root has the final result - if ClimaComms.iamroot(context) - @test Array(interp_x) ≈ [x for x in xpts, y in ypts] - end - - interp_y = Remapping.interpolate(horiz_remapper, coords.y) - if ClimaComms.iamroot(context) - @test Array(interp_y) ≈ [y for x in xpts, y in ypts] - end - - # Two fields - interp_xy = Remapping.interpolate(horiz_remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_xy[:, :, 1] ≈ interp_x - @test interp_xy[:, :, 2] ≈ interp_y - end - - # Three fields - interp_xyx = - Remapping.interpolate(horiz_remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_xyx[:, :, 1] ≈ interp_x - @test interp_xyx[:, :, 2] ≈ interp_y - @test interp_xyx[:, :, 3] ≈ interp_x - end - - # Remapping in-place one field - dest = ArrayType(zeros(21, 21)) - Remapping.interpolate!(dest, horiz_remapper, coords.x) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest - end - - - # Two fields - dest = ArrayType(zeros(21, 21, 2)) - Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, 1] - @test interp_y ≈ dest[:, :, 2] - end - - # Three fields (more than buffer length) - dest = ArrayType(zeros(21, 21, 3)) - Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y, coords.x]) - if ClimaComms.iamroot(context) - @test interp_x ≈ dest[:, :, 1] - @test interp_y ≈ dest[:, :, 2] - @test interp_x ≈ dest[:, :, 3] - end -end - -@testset "3D sphere" begin - vertdomain = Domains.IntervalDomain( - Geometry.ZPoint(0.0), - Geometry.ZPoint(1000.0); - boundary_names = (:bottom, :top), - ) - - vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) - verttopo = Topologies.IntervalTopology( - ClimaComms.SingletonCommsContext(ClimaComms.device()), - vertmesh, - ) - vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) - - horzdomain = Domains.SphereDomain(1e6) - - quad = Quadratures.GLL{4}() - horzmesh = Meshes.EquiangularCubedSphere(horzdomain, 6) - horztopology = Topologies.Topology2D(context, horzmesh) - horzspace = Spaces.SpectralElementSpace2D(horztopology, quad) - - hv_center_space = - Spaces.ExtrudedFiniteDifferenceSpace(horzspace, vert_center_space) - - longpts = range(-120.0, 120.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] - - remapper = - Remapping.Remapper(hv_center_space, hcoords, zcoords, buffer_length = 2) - - coords = Fields.coordinate_field(hv_center_space) - - interp_sin_long = Remapping.interpolate(remapper, sind.(coords.long)) - interp_sin_long2 = - Remapping.interpolate(sind.(coords.long), hcoords, zcoords) - # Only root has the final result - if ClimaComms.iamroot(context) - @test Array(interp_sin_long) ≈ - [sind(x) for x in longpts, y in latpts, z in zpts] rtol = 0.01 - @test Array(interp_sin_long2) ≈ - [sind(x) for x in longpts, y in latpts, z in zpts] rtol = 0.01 - end - - interp_sin_lat = Remapping.interpolate(remapper, sind.(coords.lat)) - if ClimaComms.iamroot(context) - @test Array(interp_sin_lat) ≈ - [sind(y) for x in longpts, y in latpts, z in zpts] rtol = 0.01 - end - - interp_z = Remapping.interpolate(remapper, coords.z) - expected_z = [z for x in longpts, y in latpts, z in zpts] - if ClimaComms.iamroot(context) - @test Array(interp_z[:, :, 2:(end - 1)]) ≈ expected_z[:, :, 2:(end - 1)] - @test Array(interp_z[:, :, 1]) ≈ - [1000.0 * (0 / 30 + 1 / 30) / 2 for x in longpts, y in latpts] - @test Array(interp_z[:, :, end]) ≈ - [1000.0 * (29 / 30 + 30 / 30) / 2 for x in longpts, y in latpts] - end - - # Remapping two fields - interp_long_lat = - Remapping.interpolate(remapper, [sind.(coords.long), sind.(coords.lat)]) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ interp_long_lat[:, :, :, 1] - @test interp_sin_lat ≈ interp_long_lat[:, :, :, 2] - end - # Remapping three fields (more than the buffer length) - interp_long_lat_long = Remapping.interpolate( - remapper, - [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], - ) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ interp_long_lat_long[:, :, :, 1] - @test interp_sin_lat ≈ interp_long_lat_long[:, :, :, 2] - @test interp_sin_long ≈ interp_long_lat_long[:, :, :, 3] - end - - # Remapping in-place one field - dest = ArrayType(zeros(21, 21, 21)) - Remapping.interpolate!(dest, remapper, sind.(coords.long)) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ dest - end - - # Two fields - dest = ArrayType(zeros(21, 21, 21, 2)) - Remapping.interpolate!( - dest, - remapper, - [sind.(coords.long), sind.(coords.lat)], - ) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ dest[:, :, :, 1] - @test interp_sin_lat ≈ dest[:, :, :, 2] - end - - # Three fields (more than buffer length) - dest = ArrayType(zeros(21, 21, 21, 3)) - Remapping.interpolate!( - dest, - remapper, - [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], - ) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ dest[:, :, :, 1] - @test interp_sin_lat ≈ dest[:, :, :, 2] - @test interp_sin_long ≈ dest[:, :, :, 3] - end - - # Horizontal space - horiz_space = Spaces.horizontal_space(hv_center_space) - horiz_remapper = Remapping.Remapper(horiz_space, hcoords, buffer_length = 2) - - coords = Fields.coordinate_field(horiz_space) - - interp_sin_long = Remapping.interpolate(horiz_remapper, sind.(coords.long)) - # Only root has the final result - if ClimaComms.iamroot(context) - @test Array(interp_sin_long) ≈ [sind(x) for x in longpts, y in latpts] rtol = 0.01 - end - - interp_sin_lat = Remapping.interpolate(horiz_remapper, sind.(coords.lat)) - if ClimaComms.iamroot(context) - @test Array(interp_sin_lat) ≈ [sind(y) for x in longpts, y in latpts] rtol = 0.01 - end - - # Two fields - interp_sin_long_lat = Remapping.interpolate( - horiz_remapper, - [sind.(coords.long), sind.(coords.lat)], - ) - if ClimaComms.iamroot(context) - @test interp_sin_long_lat[:, :, 1] ≈ interp_sin_long - @test interp_sin_long_lat[:, :, 2] ≈ interp_sin_lat - end - - # Three fields - interp_sin_long_lat_long = Remapping.interpolate( - horiz_remapper, - [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], - ) - if ClimaComms.iamroot(context) - @test interp_sin_long_lat_long[:, :, 1] ≈ interp_sin_long - @test interp_sin_long_lat_long[:, :, 2] ≈ interp_sin_lat - @test interp_sin_long_lat_long[:, :, 3] ≈ interp_sin_long - end - - # Remapping in-place one field - dest = ArrayType(zeros(21, 21)) - Remapping.interpolate!(dest, horiz_remapper, sind.(coords.long)) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ dest - end - - # Two fields - dest = ArrayType(zeros(21, 21, 2)) - Remapping.interpolate!( - dest, - horiz_remapper, - [sind.(coords.long), sind.(coords.lat)], - ) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ dest[:, :, 1] - @test interp_sin_lat ≈ dest[:, :, 2] - end - - # Three fields (more than buffer length) - dest = ArrayType(zeros(21, 21, 3)) - Remapping.interpolate!( - dest, - horiz_remapper, - [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], - ) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ dest[:, :, 1] - @test interp_sin_lat ≈ dest[:, :, 2] - @test interp_sin_long ≈ dest[:, :, 3] - end -end - -@testset "Purely vertical space" begin +# @testset "Utils" begin +# # batched_ranges(num_fields, buffer_length) +# @test Remapping.batched_ranges(1, 1) == [1:1] +# @test Remapping.batched_ranges(1, 2) == [1:1] +# @test Remapping.batched_ranges(2, 2) == [1:2] +# @test Remapping.batched_ranges(3, 2) == [1:2, 3:3] +# end + +# on_gpu = device isa ClimaComms.CUDADevice +# broken = false + +# if !on_gpu +# @testset "2D extruded" begin +# vertdomain = Domains.IntervalDomain( +# Geometry.ZPoint(0.0), +# Geometry.ZPoint(1000.0); +# boundary_names = (: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 = 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( +# hv_center_space, +# hcoords, +# zcoords, +# buffer_length = 2, +# ) + +# interp_x = Remapping.interpolate(remapper, coords.x) +# interp_x2 = Remapping.interpolate(coords.x, hcoords, zcoords) +# if ClimaComms.iamroot(context) +# @test Array(interp_x) ≈ [x for x in xpts, z in zpts] +# @test Array(interp_x2) ≈ [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 Array(interp_z[:, 2:(end - 1)]) ≈ expected_z[:, 2:(end - 1)] +# @test Array(interp_z[:, 1]) ≈ +# [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts] +# @test Array(interp_z[:, end]) ≈ +# [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts] +# end + +# # Remapping two fields +# interp_xx = Remapping.interpolate(remapper, [coords.x, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ interp_xx[:, :, 1] +# @test interp_x ≈ interp_xx[:, :, 2] +# end + +# # Remapping three fields (more than the buffer length) +# interp_xxx = +# Remapping.interpolate(remapper, [coords.x, coords.x, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ interp_xxx[:, :, 1] +# @test interp_x ≈ interp_xxx[:, :, 2] +# @test interp_x ≈ interp_xxx[:, :, 3] +# end + +# # Remapping in-place one field +# dest = ArrayType(zeros(21, 21)) +# Remapping.interpolate!(dest, remapper, coords.x) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest +# end + +# # Two fields +# dest = ArrayType(zeros(21, 21, 2)) +# Remapping.interpolate!(dest, remapper, [coords.x, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, 1] +# @test interp_x ≈ dest[:, :, 2] +# end + +# # Three fields (more than buffer length) +# dest = ArrayType(zeros(21, 21, 3)) +# Remapping.interpolate!(dest, remapper, [coords.x, coords.x, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, 1] +# @test interp_x ≈ dest[:, :, 2] +# @test interp_x ≈ dest[:, :, 3] +# end +# end +# end + +# @testset "3D box" begin +# vertdomain = Domains.IntervalDomain( +# Geometry.ZPoint(0.0), +# Geometry.ZPoint(1000.0); +# boundary_names = (:bottom, :top), +# ) + +# vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) +# verttopo = Topologies.IntervalTopology( +# ClimaComms.SingletonCommsContext(device), +# vertmesh, +# ) +# vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) + +# horzdomain = Domains.RectangleDomain( +# Geometry.XPoint(-500.0) .. Geometry.XPoint(500.0), +# Geometry.YPoint(-500.0) .. Geometry.YPoint(500.0), +# x1periodic = true, +# x2periodic = true, +# ) + +# quad = Quadratures.GLL{4}() +# horzmesh = Meshes.RectilinearMesh(horzdomain, 10, 10) +# horztopology = Topologies.Topology2D( +# ClimaComms.SingletonCommsContext(device), +# horzmesh, +# ) +# horzspace = Spaces.SpectralElementSpace2D(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) +# ypts = range(-500.0, 500.0, length = 21) +# zpts = range(0.0, 1000.0, length = 21) +# hcoords = [Geometry.XYPoint(x, y) for x in xpts, y in ypts] +# zcoords = [Geometry.ZPoint(z) for z in zpts] + +# remapper = +# Remapping.Remapper(hv_center_space, hcoords, zcoords, buffer_length = 2) + +# interp_x = Remapping.interpolate(remapper, coords.x) +# interp_x2 = Remapping.interpolate(coords.x, hcoords, zcoords) +# if ClimaComms.iamroot(context) +# @test Array(interp_x) ≈ [x for x in xpts, y in ypts, z in zpts] +# @test Array(interp_x2) ≈ [x for x in xpts, y in ypts, z in zpts] +# end + +# interp_y = Remapping.interpolate(remapper, coords.y) +# if ClimaComms.iamroot(context) +# @test Array(interp_y) ≈ [y for x in xpts, y in ypts, z in zpts] +# end + +# interp_z = Remapping.interpolate(remapper, coords.z) +# expected_z = [z for x in xpts, y in ypts, z in zpts] +# if ClimaComms.iamroot(context) +# @test Array(interp_z[:, :, 2:(end - 1)]) ≈ expected_z[:, :, 2:(end - 1)] +# @test Array(interp_z[:, :, 1]) ≈ +# [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts, y in ypts] +# @test Array(interp_z[:, :, end]) ≈ +# [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts, y in ypts] +# end + +# # Remapping two fields +# interp_xy = Remapping.interpolate(remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ interp_xy[:, :, :, 1] +# @test interp_y ≈ interp_xy[:, :, :, 2] +# end +# # Remapping three fields (more than the buffer length) +# interp_xyx = Remapping.interpolate(remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ interp_xyx[:, :, :, 1] +# @test interp_y ≈ interp_xyx[:, :, :, 2] +# @test interp_x ≈ interp_xyx[:, :, :, 3] +# end + +# # Remapping in-place one field +# dest = ArrayType(zeros(21, 21, 21)) +# Remapping.interpolate!(dest, remapper, coords.x) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest +# end + +# # Two fields +# dest = ArrayType(zeros(21, 21, 21, 2)) +# Remapping.interpolate!(dest, remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, :, 1] +# @test interp_y ≈ dest[:, :, :, 2] +# end + +# # Three fields (more than buffer length) +# dest = ArrayType(zeros(21, 21, 21, 3)) +# Remapping.interpolate!(dest, remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, :, 1] +# @test interp_y ≈ dest[:, :, :, 2] +# @test interp_x ≈ dest[:, :, :, 3] +# end + + +# # Horizontal space +# horiz_space = Spaces.horizontal_space(hv_center_space) +# horiz_remapper = Remapping.Remapper(horiz_space, hcoords, buffer_length = 2) + +# coords = Fields.coordinate_field(horiz_space) + +# interp_x = Remapping.interpolate(horiz_remapper, coords.x) +# # Only root has the final result +# if ClimaComms.iamroot(context) +# @test Array(interp_x) ≈ [x for x in xpts, y in ypts] +# end + +# interp_y = Remapping.interpolate(horiz_remapper, coords.y) +# if ClimaComms.iamroot(context) +# @test Array(interp_y) ≈ [y for x in xpts, y in ypts] +# end + +# # Two fields +# interp_xy = Remapping.interpolate(horiz_remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_xy[:, :, 1] ≈ interp_x +# @test interp_xy[:, :, 2] ≈ interp_y +# end + +# # Three fields +# interp_xyx = +# Remapping.interpolate(horiz_remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_xyx[:, :, 1] ≈ interp_x +# @test interp_xyx[:, :, 2] ≈ interp_y +# @test interp_xyx[:, :, 3] ≈ interp_x +# end + +# # Remapping in-place one field +# # +# # We have to change remapper for GPU to make sure it works for when have have only one +# # field +# dest = ArrayType(zeros(21, 21)) +# Remapping.interpolate!(dest, horiz_remapper, coords.x) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest +# end + + +# # Two fields +# dest = ArrayType(zeros(21, 21, 2)) +# Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, 1] +# @test interp_y ≈ dest[:, :, 2] +# end + +# # Three fields (more than buffer length) +# dest = ArrayType(zeros(21, 21, 3)) +# Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, 1] +# @test interp_y ≈ dest[:, :, 2] +# @test interp_x ≈ dest[:, :, 3] +# end + +# end + + +# @testset "3D box - space filling curve" begin +# vertdomain = Domains.IntervalDomain( +# Geometry.ZPoint(0.0), +# Geometry.ZPoint(1000.0); +# boundary_names = (:bottom, :top), +# ) + +# vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) +# verttopo = Topologies.IntervalTopology( +# ClimaComms.SingletonCommsContext(device), +# vertmesh, +# ) +# vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) + +# horzdomain = Domains.RectangleDomain( +# Geometry.XPoint(-500.0) .. Geometry.XPoint(500.0), +# Geometry.YPoint(-500.0) .. Geometry.YPoint(500.0), +# x1periodic = true, +# x2periodic = true, +# ) + +# quad = Quadratures.GLL{4}() +# horzmesh = Meshes.RectilinearMesh(horzdomain, 10, 10) +# horztopology = Topologies.Topology2D( +# ClimaComms.SingletonCommsContext(device), +# horzmesh, +# Topologies.spacefillingcurve(horzmesh), +# ) +# horzspace = Spaces.SpectralElementSpace2D(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) +# ypts = range(-500.0, 500.0, length = 21) +# zpts = range(0.0, 1000.0, length = 21) +# hcoords = [Geometry.XYPoint(x, y) for x in xpts, y in ypts] +# zcoords = [Geometry.ZPoint(z) for z in zpts] + +# remapper = +# Remapping.Remapper(hv_center_space, hcoords, zcoords, buffer_length = 2) + +# interp_x = Remapping.interpolate(remapper, coords.x) +# interp_x2 = Remapping.interpolate(coords.x, hcoords, zcoords) +# if ClimaComms.iamroot(context) +# @test Array(interp_x) ≈ [x for x in xpts, y in ypts, z in zpts] +# @test Array(interp_x2) ≈ [x for x in xpts, y in ypts, z in zpts] +# end + +# interp_y = Remapping.interpolate(remapper, coords.y) +# if ClimaComms.iamroot(context) +# @test Array(interp_y) ≈ [y for x in xpts, y in ypts, z in zpts] +# end + +# interp_z = Remapping.interpolate(remapper, coords.z) +# expected_z = [z for x in xpts, y in ypts, z in zpts] +# if ClimaComms.iamroot(context) +# @test Array(interp_z[:, :, 2:(end - 1)]) ≈ expected_z[:, :, 2:(end - 1)] +# @test Array(interp_z[:, :, 1]) ≈ +# [1000.0 * (0 / 30 + 1 / 30) / 2 for x in xpts, y in ypts] +# @test Array(interp_z[:, :, end]) ≈ +# [1000.0 * (29 / 30 + 30 / 30) / 2 for x in xpts, y in ypts] +# end + +# # Remapping two fields +# interp_xy = Remapping.interpolate(remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ interp_xy[:, :, :, 1] +# @test interp_y ≈ interp_xy[:, :, :, 2] +# end +# # Remapping three fields (more than the buffer length) +# interp_xyx = Remapping.interpolate(remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ interp_xyx[:, :, :, 1] +# @test interp_y ≈ interp_xyx[:, :, :, 2] +# @test interp_x ≈ interp_xyx[:, :, :, 3] +# end + +# # Remapping in-place one field +# dest = ArrayType(zeros(21, 21, 21)) +# Remapping.interpolate!(dest, remapper, coords.x) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest +# end + +# # Two fields +# dest = ArrayType(zeros(21, 21, 21, 2)) +# Remapping.interpolate!(dest, remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, :, 1] +# @test interp_y ≈ dest[:, :, :, 2] +# end + +# # Three fields (more than buffer length) +# dest = ArrayType(zeros(21, 21, 21, 3)) +# Remapping.interpolate!(dest, remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, :, 1] +# @test interp_y ≈ dest[:, :, :, 2] +# @test interp_x ≈ dest[:, :, :, 3] +# end + + +# # Horizontal space +# horiz_space = Spaces.horizontal_space(hv_center_space) +# horiz_remapper = Remapping.Remapper(horiz_space, hcoords, buffer_length = 2) + +# coords = Fields.coordinate_field(horiz_space) + +# interp_x = Remapping.interpolate(horiz_remapper, coords.x) +# # Only root has the final result +# if ClimaComms.iamroot(context) +# @test Array(interp_x) ≈ [x for x in xpts, y in ypts] +# end + +# interp_y = Remapping.interpolate(horiz_remapper, coords.y) +# if ClimaComms.iamroot(context) +# @test Array(interp_y) ≈ [y for x in xpts, y in ypts] +# end + +# # Two fields +# interp_xy = Remapping.interpolate(horiz_remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_xy[:, :, 1] ≈ interp_x +# @test interp_xy[:, :, 2] ≈ interp_y +# end + +# # Three fields +# interp_xyx = +# Remapping.interpolate(horiz_remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_xyx[:, :, 1] ≈ interp_x +# @test interp_xyx[:, :, 2] ≈ interp_y +# @test interp_xyx[:, :, 3] ≈ interp_x +# end + +# # Remapping in-place one field +# dest = ArrayType(zeros(21, 21)) +# Remapping.interpolate!(dest, horiz_remapper, coords.x) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest +# end + + +# # Two fields +# dest = ArrayType(zeros(21, 21, 2)) +# Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, 1] +# @test interp_y ≈ dest[:, :, 2] +# end + +# # Three fields (more than buffer length) +# dest = ArrayType(zeros(21, 21, 3)) +# Remapping.interpolate!(dest, horiz_remapper, [coords.x, coords.y, coords.x]) +# if ClimaComms.iamroot(context) +# @test interp_x ≈ dest[:, :, 1] +# @test interp_y ≈ dest[:, :, 2] +# @test interp_x ≈ dest[:, :, 3] +# end +# end + +# @testset "3D sphere" begin +# vertdomain = Domains.IntervalDomain( +# Geometry.ZPoint(0.0), +# Geometry.ZPoint(1000.0); +# boundary_names = (:bottom, :top), +# ) + +# vertmesh = Meshes.IntervalMesh(vertdomain, nelems = 30) +# verttopo = Topologies.IntervalTopology( +# ClimaComms.SingletonCommsContext(ClimaComms.device()), +# vertmesh, +# ) +# vert_center_space = Spaces.CenterFiniteDifferenceSpace(verttopo) + +# horzdomain = Domains.SphereDomain(1e6) + +# quad = Quadratures.GLL{4}() +# horzmesh = Meshes.EquiangularCubedSphere(horzdomain, 6) +# horztopology = Topologies.Topology2D(context, horzmesh) +# horzspace = Spaces.SpectralElementSpace2D(horztopology, quad) + +# hv_center_space = +# Spaces.ExtrudedFiniteDifferenceSpace(horzspace, vert_center_space) + +# longpts = range(-120.0, 120.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] + +# remapper = +# Remapping.Remapper(hv_center_space, hcoords, zcoords, buffer_length = 2) + +# coords = Fields.coordinate_field(hv_center_space) + +# interp_sin_long = Remapping.interpolate(remapper, sind.(coords.long)) +# interp_sin_long2 = +# Remapping.interpolate(sind.(coords.long), hcoords, zcoords) +# # Only root has the final result +# if ClimaComms.iamroot(context) +# @test Array(interp_sin_long) ≈ +# [sind(x) for x in longpts, y in latpts, z in zpts] rtol = 0.01 +# @test Array(interp_sin_long2) ≈ +# [sind(x) for x in longpts, y in latpts, z in zpts] rtol = 0.01 +# end + +# interp_sin_lat = Remapping.interpolate(remapper, sind.(coords.lat)) +# if ClimaComms.iamroot(context) +# @test Array(interp_sin_lat) ≈ +# [sind(y) for x in longpts, y in latpts, z in zpts] rtol = 0.01 +# end + +# interp_z = Remapping.interpolate(remapper, coords.z) +# expected_z = [z for x in longpts, y in latpts, z in zpts] +# if ClimaComms.iamroot(context) +# @test Array(interp_z[:, :, 2:(end - 1)]) ≈ expected_z[:, :, 2:(end - 1)] +# @test Array(interp_z[:, :, 1]) ≈ +# [1000.0 * (0 / 30 + 1 / 30) / 2 for x in longpts, y in latpts] +# @test Array(interp_z[:, :, end]) ≈ +# [1000.0 * (29 / 30 + 30 / 30) / 2 for x in longpts, y in latpts] +# end + +# # Remapping two fields +# interp_long_lat = +# Remapping.interpolate(remapper, [sind.(coords.long), sind.(coords.lat)]) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ interp_long_lat[:, :, :, 1] +# @test interp_sin_lat ≈ interp_long_lat[:, :, :, 2] +# end +# # Remapping three fields (more than the buffer length) +# interp_long_lat_long = Remapping.interpolate( +# remapper, +# [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], +# ) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ interp_long_lat_long[:, :, :, 1] +# @test interp_sin_lat ≈ interp_long_lat_long[:, :, :, 2] +# @test interp_sin_long ≈ interp_long_lat_long[:, :, :, 3] +# end + +# # Remapping in-place one field +# dest = ArrayType(zeros(21, 21, 21)) +# Remapping.interpolate!(dest, remapper, sind.(coords.long)) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ dest +# end + +# # Two fields +# dest = ArrayType(zeros(21, 21, 21, 2)) +# Remapping.interpolate!( +# dest, +# remapper, +# [sind.(coords.long), sind.(coords.lat)], +# ) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ dest[:, :, :, 1] +# @test interp_sin_lat ≈ dest[:, :, :, 2] +# end + +# # Three fields (more than buffer length) +# dest = ArrayType(zeros(21, 21, 21, 3)) +# Remapping.interpolate!( +# dest, +# remapper, +# [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], +# ) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ dest[:, :, :, 1] +# @test interp_sin_lat ≈ dest[:, :, :, 2] +# @test interp_sin_long ≈ dest[:, :, :, 3] +# end + +# # Horizontal space +# horiz_space = Spaces.horizontal_space(hv_center_space) +# horiz_remapper = Remapping.Remapper(horiz_space, hcoords, buffer_length = 2) + +# coords = Fields.coordinate_field(horiz_space) + +# interp_sin_long = Remapping.interpolate(horiz_remapper, sind.(coords.long)) +# # Only root has the final result +# if ClimaComms.iamroot(context) +# @test Array(interp_sin_long) ≈ [sind(x) for x in longpts, y in latpts] rtol = 0.01 +# end + +# interp_sin_lat = Remapping.interpolate(horiz_remapper, sind.(coords.lat)) +# if ClimaComms.iamroot(context) +# @test Array(interp_sin_lat) ≈ [sind(y) for x in longpts, y in latpts] rtol = 0.01 +# end + +# # Two fields +# interp_sin_long_lat = Remapping.interpolate( +# horiz_remapper, +# [sind.(coords.long), sind.(coords.lat)], +# ) +# if ClimaComms.iamroot(context) +# @test interp_sin_long_lat[:, :, 1] ≈ interp_sin_long +# @test interp_sin_long_lat[:, :, 2] ≈ interp_sin_lat +# end + +# # Three fields +# interp_sin_long_lat_long = Remapping.interpolate( +# horiz_remapper, +# [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], +# ) +# if ClimaComms.iamroot(context) +# @test interp_sin_long_lat_long[:, :, 1] ≈ interp_sin_long +# @test interp_sin_long_lat_long[:, :, 2] ≈ interp_sin_lat +# @test interp_sin_long_lat_long[:, :, 3] ≈ interp_sin_long +# end + +# # Remapping in-place one field +# dest = ArrayType(zeros(21, 21)) +# Remapping.interpolate!(dest, horiz_remapper, sind.(coords.long)) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ dest +# end + +# # Two fields +# dest = ArrayType(zeros(21, 21, 2)) +# Remapping.interpolate!( +# dest, +# horiz_remapper, +# [sind.(coords.long), sind.(coords.lat)], +# ) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ dest[:, :, 1] +# @test interp_sin_lat ≈ dest[:, :, 2] +# end + +# # Three fields (more than buffer length) +# dest = ArrayType(zeros(21, 21, 3)) +# Remapping.interpolate!( +# dest, +# horiz_remapper, +# [sind.(coords.long), sind.(coords.lat), sind.(coords.long)], +# ) +# if ClimaComms.iamroot(context) +# @test interp_sin_long ≈ dest[:, :, 1] +# @test interp_sin_lat ≈ dest[:, :, 2] +# @test interp_sin_long ≈ dest[:, :, 3] +# end +# end + +# @testset "Purely vertical space" begin vertdomain = Domains.IntervalDomain( Geometry.ZPoint(0.0), Geometry.ZPoint(1000.0); @@ -675,26 +675,79 @@ end vertmesh, ) cspace = Spaces.CenterFiniteDifferenceSpace(verttopo) - space = Spaces.FaceFiniteDifferenceSpace(cspace) + fspace = Spaces.FaceFiniteDifferenceSpace(cspace) zpts = range(0.0, 1000.0, 21) zcoords = [Geometry.ZPoint(z) for z in zpts] - remapper = - Remapping.Remapper(space; target_zcoords = zcoords, buffer_length = 2) - - coords = Fields.coordinate_field(space) - - interp_z = Remapping.interpolate(remapper, coords.z) - expected_z = zpts - if ClimaComms.iamroot(context) - @test interp_z == collect(expected_z) - end - - # Remapping two fields - interp = Remapping.interpolate(remapper, [cos.(coords.z), sin.(coords.z)]) - if ClimaComms.iamroot(context) - @test interp_sin_long ≈ interp_long_lat[:, :, :, 1] - @test interp_sin_lat ≈ interp_long_lat[:, :, :, 2] - end + # Center space + cremapper = + Remapping.Remapper(cspace; target_zcoords = zcoords, buffer_length = 2) + ccoords = Fields.coordinate_field(cspace) + cinterp_z = Remapping.interpolate(cremapper, ccoords.z) + cexpected_z = copy(zpts) + # if ClimaComms.iamroot(context) + # @test Array(cinterp_z) ≈ cexpected_z + # end + + # Face space + fremapper = + Remapping.Remapper(fspace; target_zcoords = zcoords, buffer_length = 2) + fcoords = Fields.coordinate_field(fspace) + finterp_z = Remapping.interpolate(fremapper, fcoords.z) + fexpected_z = copy(zpts) + if ClimaComms.iamroot(context) + @test Array(finterp_z) ≈ fexpected_z + end + + + + # # Remapping two fields + # interp = Remapping.interpolate(remapper, [cos.(coords.z), sin.(coords.z)]) + # if ClimaComms.iamroot(context) + # @test interp_sin_long ≈ interp_long_lat[:, :, :, 1] + # @test interp_sin_lat ≈ interp_long_lat[:, :, :, 2] + # end + + # # Remapping three fields (more than the buffer length) + # interp_xxx = + # Remapping.interpolate(remapper, [coords.x, coords.x, coords.x]) + # if ClimaComms.iamroot(context) + # @test interp_x ≈ interp_xxx[:, :, 1] + # @test interp_x ≈ interp_xxx[:, :, 2] + # @test interp_x ≈ interp_xxx[:, :, 3] + # end + + # # Remapping in-place one field + # dest = ArrayType(zeros(21, 21)) + # Remapping.interpolate!(dest, horiz_remapper, sind.(coords.long)) + # if ClimaComms.iamroot(context) + # @test interp_sin_long ≈ dest + # end + + # # Two fields + # dest = ArrayType(zeros(21, 21, 2)) + # Remapping.interpolate!( + # dest, + # horiz_remapper, + # [sind.(coords.long), sind.(coords.lat)], + # ) + + # # Three fields (more than buffer length) + # dest = ArrayType(zeros(21, 21, 3)) + # Remapping.interpolate!(dest, remapper, [coords.x, coords.x, coords.x]) + # if ClimaComms.iamroot(context) + # @test interp_x ≈ dest[:, :, 1] + # @test interp_x ≈ dest[:, :, 2] + # @test interp_x ≈ dest[:, :, 3] + # end +# end + +@testset "Convenience interpolate" begin + # 3D Sphere space + # 3D box space + # 2D slice space + # Purely horizontal 2D space + # Purely horizontal 1D space + # Purely vertical space end