From 47079aa740eab3328f3da43734f6a8eb4747d4b9 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Sun, 8 Sep 2024 22:15:36 -0600 Subject: [PATCH] Fix FieldTimeSeries interpolation for single column grids with x, y locations (#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 --- src/Fields/interpolate.jl | 3 +- src/Grids/grid_generation.jl | 5 +- src/Grids/grid_utils.jl | 5 +- src/Grids/rectilinear_grid.jl | 1 - .../field_time_series_indexing.jl | 22 +++++---- src/Utils/schedules.jl | 1 + test/test_output_readers.jl | 49 ++++++++++++++----- 7 files changed, 57 insertions(+), 29 deletions(-) diff --git a/src/Fields/interpolate.jl b/src/Fields/interpolate.jl index 202ab16462..6acd5b7ad9 100644 --- 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 index 1d30922ba6..b0ba1ad7dc 100644 --- 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 index 2d09ba2f29..b1c94204f1 100644 --- 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 index 8b5c71b4ac..bc6c36a694 100644 --- 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 index a7a19faad7..6a4683f9d0 100644 --- 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₂ = interpolating_time_indices(time_indexing, times, at_time) + ñ, 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 index fdd632f173..49302cfe7b 100644 --- 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 index 1f5db696c5..b33923a6c9 100644 --- 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)