Support InMemory(N) with split timeseries (#5607)

* fix split path

* fix

Author: simone-silvestri

src/OutputReaders/field_time_series.jl

@@ -909,23 +909,15 @@ function FieldTimeSeries(file::JLD2.JLDFile, name::String;
     Nt = time_indices_length(backend, times)
     @apply_regionally data = new_data(eltype(grid), grid, loc, indices, Nt)
 
-    # For OnDisk backend with split files, store a SplitFilePath so getindex
-    # can find the correct part file for each time index.
-    fts_path = if !isnothing(Nparts) && backend isa OnDisk
-        SplitFilePath(part_paths, cumsum(iterations_per_part))
-    else
-        path
-    end
+    fts_path = isnothing(Nparts) ? path : SplitFilePath(part_paths, cumsum(iterations_per_part))
 
     time_series = FieldTimeSeries{LX, LY, LZ}(data, grid, backend, boundary_conditions, indices,
                                               times, fts_path, name, time_indexing, reader_kw)
 
     if isnothing(Nparts)
         set!(time_series, path, name)
     else
-        for path in part_paths
-            set!(time_series, path, name; warn_missing_data=false)
-        end
+        set!(time_series, fts_path, name)
     end
 
     return time_series

src/OutputReaders/set_field_time_series.jl

@@ -40,6 +40,20 @@ function set!(fts::InMemoryFTS, path::String, args::String...; kwargs...)
     end
 end
 
+function set!(fts::InMemoryFTS, sfp::SplitFilePath, name::String=fts.name)
+    Ntotal = last(sfp.cumulative_length)
+    needed_paths = String[]
+    for n in time_indices(fts)
+        (n < 1 || n > Ntotal) && continue
+        part_path, _ = file_and_local_index(sfp, n)
+        part_path ∉ needed_paths && push!(needed_paths, part_path)
+    end
+    for part_path in needed_paths
+        set!(fts, part_path, name; warn_missing_data=false)
+    end
+    return nothing
+end
+
 # Convenience method with default path
 function set!(fts::InMemoryFTS; kwargs...)
     return set!(fts, fts.path; kwargs...)
@@ -141,6 +155,7 @@ function initialize_file!(file, name, fts)
 end
 
 set!(fts::OnDiskFTS, path::String, name::String; kwargs...) = nothing
+set!(fts::OnDiskFTS, ::SplitFilePath, name::String=fts.name) = nothing
 
 function set!(fts::InMemoryFTS, f::Function)
     cpu_times = on_architecture(CPU(), fts.times)

test/test_output_readers.jl

@@ -2,7 +2,7 @@ include("dependencies_for_runtests.jl")
 
 using Oceananigans.Units: Time
 using Oceananigans.Fields: indices, interpolate!
-using Oceananigans.OutputReaders: Cyclical, Clamp, Linear
+using Oceananigans.OutputReaders: Cyclical, Clamp, Linear, SplitFilePath
 
 using Random
 using NCDatasets
@@ -325,6 +325,9 @@ function test_field_time_series_split_files(arch)
     model = NonhydrostaticModel(grid, tracers=:c)
     simulation = Simulation(model, Δt=1, stop_time=10)
 
+    set_tracer_to_iteration!(sim) = fill!(parent(sim.model.tracers.c), sim.model.clock.iteration)
+    add_callback!(simulation, set_tracer_to_iteration!, IterationInterval(1))
+
     simulation.output_writers[:fields] = JLD2Writer(model, model.tracers;
                                                      filename = "split_test",
                                                      dir = dir,
@@ -341,6 +344,7 @@ function test_field_time_series_split_files(arch)
     @test length(fts_mem.times) == 11
     @test fts_mem[1] isa Field
     @test fts_mem[11] isa Field
+    @test fts_mem.path isa SplitFilePath
 
     # Test OnDisk backend with split files
     fts_disk = FieldTimeSeries(abs_path, "c"; backend=OnDisk(), architecture=arch)
@@ -351,6 +355,12 @@ function test_field_time_series_split_files(arch)
         @test fts_disk[n] isa Field
     end
 
+    fts_partly = FieldTimeSeries(abs_path, "c"; backend=InMemory(2), architecture=arch)
+    @test fts_partly.path isa SplitFilePath
+    for n in 1:length(fts_partly.times)
+        @test Array(interior(fts_partly[n])) == Array(interior(fts_mem[n]))
+    end
+
     rm(dir, recursive=true, force=true)
     return nothing
 end