Fix FieldTimeSeries interpolation for single column grids with x, y l…

…ocations (#3766) * Fix issue showing single column grids with x, y location * Fix FieldTimeSeries interpolation for SCM grids * Bump version * Add test * Import _node into FieldTimeSeries indexing * Relax test stingency --------- Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>
CliMA · Sep 9, 2024 · 47079aa · 47079aa · navidcy · Sep 9, 2024
1 parent 6d1d503
commit 47079aa
Show file tree

Hide file tree

Showing 7 changed files with 57 additions and 29 deletions.
diff --git a/src/Fields/interpolate.jl b/src/Fields/interpolate.jl
@@ -183,7 +183,6 @@ end
 @inline function _fractional_indices((x, y), grid, ℓx, ℓy, ::Nothing)
     ii = fractional_x_index(x, (ℓx, ℓy, nothing), grid)
     jj = fractional_y_index(y, (ℓx, ℓy, nothing), grid)
-
     return (ii, jj, nothing)
 end
 
@@ -324,7 +323,7 @@ end
 @inline flatten_node(x, ::Nothing) = flatten_node(x)
 
 @inline flatten_node(x) = tuple(x)
-@inline flatten_node(::Nothing) = tuple(x)
+@inline flatten_node(::Nothing) = tuple()
 
 @kernel function _interpolate!(to_field, to_grid, to_location,
                                from_field, from_grid, from_location)

diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl
@@ -123,8 +123,9 @@ end
 generate_coordinate(FT, ::Flat, N, H, c::Number, coordinate_name, arch) =
     FT(1), range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1)
 
-generate_coordinate(FT, ::Flat, N, H, c::Tuple{Number, Number}, coordinate_name, arch) =
-    FT(1), c, c, FT(1), FT(1)
+# What's the use case for this?
+# generate_coordinate(FT, ::Flat, N, H, c::Tuple{Number, Number}, coordinate_name, arch) =
+#     FT(1), c, c, FT(1), FT(1)
 
 generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch) =
     FT(1), nothing, nothing, FT(1), FT(1)

diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl
@@ -175,7 +175,7 @@ constant grid spacing `Δ`, and interior extent `L`.
 
 # Grid domains
 @inline domain(topo, N, ξ) = CUDA.@allowscalar ξ[1], ξ[N+1]
-@inline domain(::Flat, N, ξ::AbstractArray) = CUDA.@allowscalar ξ[1], ξ[1]
+@inline domain(::Flat, N, ξ::AbstractArray) = ξ[1]
 @inline domain(::Flat, N, ξ::Number) = ξ
 @inline domain(::Flat, N, ::Nothing) = nothing
 
@@ -347,7 +347,8 @@ function domain_summary(topo, name, (left, right))
                   topo isa Bounded ? "Bounded  " :
                   topo isa FullyConnected ? "FullyConnected " :
                   topo isa LeftConnected ? "LeftConnected  " :
-                  "RightConnected "
+                  topo isa RightConnected ? "RightConnected  " :
+                  error("Unexpected topology $topo together with the domain end points ($left, $right)")
 
     return string(topo_string, name, " ∈ [",
                   prettysummary(left), ", ",

diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl
@@ -66,7 +66,6 @@ stretched_dimensions(::YZRegularRG) = tuple(1)
 stretched_dimensions(::XZRegularRG) = tuple(2)
 stretched_dimensions(::XYRegularRG) = tuple(3)
 
-
 """
     RectilinearGrid([architecture = CPU(), FT = Float64];
                     size,

diff --git a/src/OutputReaders/field_time_series_indexing.jl b/src/OutputReaders/field_time_series_indexing.jl
@@ -1,4 +1,5 @@
-using Oceananigans.Fields: interpolator, _interpolate, fractional_indices
+using Oceananigans.Grids: _node
+using Oceananigans.Fields: interpolator, _interpolate, fractional_indices, flatten_node
 using Oceananigans.Architectures: architecture
 
 import Oceananigans.Fields: interpolate
@@ -154,27 +155,28 @@ end
 ##### Linear time- and space-interpolation of a FTS
 #####
 
-@inline function interpolate(at_node, at_time_index::Time, from_fts::FlavorOfFTS, from_loc, from_grid)
+@inline function interpolate(to_node, to_time_index::Time, from_fts::FlavorOfFTS, from_loc, from_grid)
     data = from_fts.data
     times = from_fts.times
     backend = from_fts.backend
     time_indexing = from_fts.time_indexing
-    return interpolate(at_node, at_time_index, data, from_loc, from_grid, times, backend, time_indexing)
+    return interpolate(to_node, to_time_index, data, from_loc, from_grid, times, backend, time_indexing)
 end
 
-@inline function interpolate(at_node, at_time_index::Time, data::OffsetArray,
+@inline function interpolate(to_node, to_time_index::Time, data::OffsetArray,
                              from_loc, from_grid, times, backend, time_indexing)
 
-    at_time = at_time_index.time
+    to_time = to_time_index.time
 
     # Build space interpolators
-    ii, jj, kk = fractional_indices(at_node, from_grid, from_loc...)
+    to_node = flatten_node(to_node...)
+    ii, jj, kk = fractional_indices(to_node, from_grid, from_loc...)
 
     ix = interpolator(ii)
     iy = interpolator(jj)
     iz = interpolator(kk)
 
-    ñ, n₁, n₂ = interpolating_time_indices(time_indexing, times, at_time)
+    ñ, n₁, n₂ = interpolating_time_indices(time_indexing, times, to_time)
 
     Nt = length(times)
     m₁ = memory_index(backend, time_indexing, Nt, n₁)
@@ -216,10 +218,10 @@ end
     # 4D index, cool!
     i, j, k, n = @index(Global, NTuple)
 
-    target_node = node(i, j, k, target_grid, target_location...)
-    at_time     = @inbounds Time(target_times[n])
+    target_node = _node(i, j, k, target_grid, target_location...)
+    to_time     = @inbounds Time(target_times[n])
 
-    @inbounds target_fts[i, j, k, n] = interpolate(target_node, at_time,
+    @inbounds target_fts[i, j, k, n] = interpolate(target_node, to_time,
                                                    source_fts, source_location, source_grid)
 end
 

diff --git a/src/Utils/schedules.jl b/src/Utils/schedules.jl
@@ -290,3 +290,4 @@ Base.summary(schedule::SpecifiedTimes) = string("SpecifiedTimes(", specified_tim
 Base.summary(schedule::ConsecutiveIterations) = string("ConsecutiveIterations(",
                                                        summary(schedule.parent), ", ",
                                                        schedule.consecutive_iterations, ")")
+
diff --git a/test/test_output_readers.jl b/test/test_output_readers.jl
@@ -36,16 +36,20 @@ function generate_some_interesting_simulation_data(Nx, Ny, Nz; architecture=CPU(
 
     fields_to_output = merge(model.velocities, model.tracers, computed_fields)
 
+    filepath3d = "test_3d_output_with_halos.jld2"
+    filepath2d = "test_2d_output_with_halos.jld2"
+    filepath1d = "test_1d_output_with_halos.jld2"
+
     simulation.output_writers[:jld2_3d_with_halos] =
         JLD2OutputWriter(model, fields_to_output,
-                         filename = "test_3d_output_with_halos.jld2",
+                         filename = filepath3d,
                          with_halos = true,
                          schedule = TimeInterval(30seconds),
                          overwrite_existing = true)
 
     simulation.output_writers[:jld2_2d_with_halos] =
         JLD2OutputWriter(model, fields_to_output,
-                         filename = "test_2d_output_with_halos.jld2",
+                         filename = filepath2d,
                          indices = (:, :, grid.Nz),
                          with_halos = true,
                          schedule = TimeInterval(30seconds),
@@ -55,33 +59,28 @@ function generate_some_interesting_simulation_data(Nx, Ny, Nz; architecture=CPU(
 
     simulation.output_writers[:jld2_1d_with_halos] =
         JLD2OutputWriter(model, profiles,
-                         filename = "test_1d_output_with_halos.jld2",
+                         filename = filepath1d,
                          with_halos = true,
                          schedule = TimeInterval(30seconds),
                          overwrite_existing = true)
 
     run!(simulation)
 
-    return nothing
+    return filepath1d, filepath2d, filepath3d
 end
 
 @testset "OutputReaders" begin
     @info "Testing output readers..."
 
-    Nx, Ny, Nz = 16, 10, 5
-    generate_some_interesting_simulation_data(Nx, Ny, Nz)
     Nt = 5
-
-    filepath3d = "test_3d_output_with_halos.jld2"
-    filepath2d = "test_2d_output_with_halos.jld2"
-    filepath1d = "test_1d_output_with_halos.jld2"
+    Nx, Ny, Nz = 16, 10, 5
+    filepath1d, filepath2d, filepath3d = generate_some_interesting_simulation_data(Nx, Ny, Nz)
 
     for arch in archs
         @testset "FieldTimeSeries{InMemory} [$(typeof(arch))]" begin
             @info "  Testing FieldTimeSeries{InMemory} [$(typeof(arch))]..."
 
-            ## 3D Fields
-
+            # 3D Fields
             u3 = FieldTimeSeries(filepath3d, "u", architecture=arch)
             v3 = FieldTimeSeries(filepath3d, "v", architecture=arch)
             w3 = FieldTimeSeries(filepath3d, "w", architecture=arch)
@@ -129,6 +128,32 @@ end
             interpolate!(u3i, u3)
             @test all(interior(u3i) .≈ interior(u3))
 
+            # Interpolation to a _located_ single column grid
+            grid3 = RectilinearGrid(arch, size=(3, 3, 3), x=(0.5, 3.5), y=(0.5, 3.5), z=(0.5, 3.5),
+                                    topology = (Periodic, Periodic, Bounded))
+
+            grid1 = RectilinearGrid(arch, size=3, x=1.3, y=2.7, z=(0.5, 3.5),
+                                    topology=(Flat, Flat, Bounded))
+
+            times = [1, 2]
+            c3 = FieldTimeSeries{Center, Center, Center}(grid3, times)
+            c1 = FieldTimeSeries{Center, Center, Center}(grid1, times)
+
+            for n in 1:length(times)
+                tn = times[n]
+                c₀(x, y, z) = (x + y + z) * tn
+                set!(c3[n], c₀)
+            end
+
+            interpolate!(c1, c3)
+
+            # Convert to CPU for testing
+            c11 = interior(c1[1], 1, 1, :) |> Array
+            c12 = interior(c1[2], 1, 1, :) |> Array
+
+            @test c11 ≈ [5.0, 6.0, 7.0]
+            @test c12 ≈ [10.0, 12.0, 14.0]
+
             ## 2D sliced Fields
 
             u2 = FieldTimeSeries(filepath2d, "u", architecture=arch)