From 3db2d25d13ff8e051c11f0ccac038677823e9fe2 Mon Sep 17 00:00:00 2001
From: Simone Silvestri <silvestri.simone0@gmail.com>
Date: Wed, 11 Dec 2024 10:04:06 +0100
Subject: [PATCH] Bugfix `set!`ting ECCO FTS with cycling boundaries (#290)

* add test

* cycling test

* fix ecco cycling

* no need for a memory index

* correct tests

* Import update_state!

* timesteppers

* correct tests

---------

Co-authored-by: Gregory L. Wagner <wagner.greg@gmail.com>
---
 src/DataWrangling/ECCO/ECCO_restoring.jl |  4 +-
 test/test_ecco.jl                        | 51 +++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/src/DataWrangling/ECCO/ECCO_restoring.jl b/src/DataWrangling/ECCO/ECCO_restoring.jl
index 7400d0f7..5f235985 100644
--- a/src/DataWrangling/ECCO/ECCO_restoring.jl
+++ b/src/DataWrangling/ECCO/ECCO_restoring.jl
@@ -61,12 +61,10 @@ cache_inpainted_data(::ECCONetCDFBackend{native, cache_data}) where {native, cac
 
 function set!(fts::ECCONetCDFFTS) 
     backend = fts.backend
-    start   = backend.start
     inpainting = backend.inpainting
     cache_data = cache_inpainted_data(backend)
-    len = backend.length
 
-    for t in start:start+len-1
+    for t in time_indices(fts)
         # Set each element of the time-series to the associated file
         metadatum = @inbounds backend.metadata[t] 
         set!(fts[t], metadatum; inpainting, cache_inpainted_data=cache_data)
diff --git a/test/test_ecco.jl b/test/test_ecco.jl
index 9a3bb9c3..5b70ab6f 100644
--- a/test/test_ecco.jl
+++ b/test/test_ecco.jl
@@ -7,7 +7,11 @@ using ClimaOcean
 using ClimaOcean.ECCO
 using ClimaOcean.ECCO: ECCO_field, metadata_path, ECCO_times
 using ClimaOcean.DataWrangling: NearestNeighborInpainting
+
 using Oceananigans.Grids: topology
+using Oceananigans.OutputReaders: time_indices
+using Oceananigans.TimeSteppers: update_state!
+using Oceananigans.Units
 
 using CUDA: @allowscalar
 
@@ -174,7 +178,8 @@ end
 
 @testset "Setting temperature and salinity to ECCO" begin
     for arch in test_architectures
-        grid = LatitudeLongitudeGrid(size=(10, 10, 10),
+        grid = LatitudeLongitudeGrid(arch; 
+                                     size=(10, 10, 10),
                                      latitude=(-60, -40),
                                      longitude=(10, 15),
                                      z=(-200, 0),
@@ -185,3 +190,47 @@ end
         set!(ocean.model, T=ECCOMetadata(:temperature, date), S=ECCOMetadata(:salinity, date))
     end
 end
+
+@testset "ECCO dataset cycling boundaries" begin
+    for arch in test_architectures
+        grid = LatitudeLongitudeGrid(arch;
+                                     size=(10, 10, 10),
+                                     latitude=(-60, -40),
+                                     longitude=(10, 15),
+                                     z=(-200, 0),
+                                     halo = (7, 7, 7))
+
+        start_date = DateTimeProlepticGregorian(1993, 1, 1)
+        end_date = DateTimeProlepticGregorian(1993, 5, 1)
+        dates = start_date : Month(1) : end_date
+
+        t_restoring = ECCORestoring(arch, :temperature;
+                                    dates,
+                                    rate = 1 / 1000.0,
+                                    inpainting)
+
+        times = ECCO_times(t_restoring.field_time_series.backend.metadata)
+        ocean = ocean_simulation(grid, forcing = (; T = t_restoring))
+
+        ocean.model.clock.time = times[3] + 2 * Units.days
+        update_state!(ocean.model)
+
+        @test t_restoring.field_time_series.backend.start == 3
+
+        # Compile
+        time_step!(ocean)
+
+        # Try stepping out of the ECCO dataset bounds
+        ocean.model.clock.time = last(times) + 2 * Units.days
+
+        update_state!(ocean.model)
+        
+        @test begin
+            time_step!(ocean)
+            true
+        end
+
+        # The backend has cycled to the end
+        @test time_indices(t_restoring.field_time_series) == (5, 1)
+    end
+end