DatasetRestoring only via Metadata (#486)

navidcy · web-flow · commit cdbe1699145b · 2025-04-28T03:42:23.000+02:00
* DatasetRestoring only from metadata

* better show method for DatasetRestoring and Metadata

* typed arg

* fix test

* better DatasetRestoring show method

* better DatasetRestoring show method

* more robust Dataset cycling boundaries test

* fix start_idx bug + add convenience metadata constructor

* test the Metadata convenience constructor

* disambiguate Metadata constructors

* validate kwargs

* fix docstring

* it's isnothing

* fix test

* fix Metadata constructor
diff --git a/src/DataWrangling/metadata.jl b/src/DataWrangling/metadata.jl
@@ -17,23 +17,37 @@ end
 
 Metadata holding a specific dataset information.
 
-Arguments
-=========
+Argument
+========
 - `variable_name`: a symbol representing the name of the variable (for example, `:temperature`,
                    `:salinity`, `:u_velocity`, etc)
 
 Keyword Arguments
 =================
-- `dataset`: The dataset of the dataset. Supported datasets are `ECCO2Monthly()`, `ECCO2Daily()`,
-             `ECCO4Monthly()`, `EN4Monthly(), `RepeatYearJRA55()`, or `MultiYearJRA55()`.
-- `dates`: The dates of the dataset, in a `AbstractCFDateTime` format. Note this can either be a range
-           or a vector of dates, representing a time-series. For a single date, use [`Metadatum`](@ref).
+- `dataset`: Supported datasets are `ECCO2Monthly()`, `ECCO2Daily()`, `ECCO4Monthly()`, `EN4Monthly(),
+             `RepeatYearJRA55()`, or `MultiYearJRA55()`.
+- `dates`: The dates of the dataset (`Dates.AbstractDateTime` or `CFTime.AbstractCFDateTime`).
+           Note this can either be a range or a vector of dates, representing a time-series.
+           For a single date, use [`Metadatum`](@ref).
+- `start_date`: If `dates = nothing`, we can prescribe the first date of metadata as a date
+                (`Dates.AbstractDateTime` or `CFTime.AbstractCFDateTime`). `start_date` should lie
+                within the date range of the dataset. Default: nothing.
+- `end_date`: If `dates = nothing`, we can prescribe the last date of metadata as a date
+              (`Dates.AbstractDateTime` or `CFTime.AbstractCFDateTime`). `end_date` should lie
+                within the date range of the dataset. Default: nothing.
 - `dir`: The directory where the dataset is stored.
 """
 function Metadata(variable_name;
                   dataset,
-                  dates=all_dates(dataset, variable_name)[1:1],
-                  dir=default_download_directory(dataset))
+                  dates = all_dates(dataset, variable_name),
+                  dir = default_download_directory(dataset),
+                  start_date = nothing,
+                  end_date = nothing)
+
+    if !isnothing(start_date) && !isnothing(end_date)
+        @info "Slicing date range within $start_date and $end_date"
+        dates = compute_native_date_range(dates, start_date, end_date)
+    end
 
     return Metadata(variable_name, dataset, dates, dir)
 end
@@ -54,7 +68,9 @@ function Metadatum(variable_name;
                    date=first_date(dataset, variable_name),
                    dir=default_download_directory(dataset))
 
-    # TODO: validate that `date` is actually a single date?
+    date isa Union{CFTime.AbstractCFDateTime, Dates.AbstractDateTime} ||
+        throw(ArgumentError("date must be Union{Dates.AbstractDateTime, CFTime.AbstractCFDateTime}"))
+
     return Metadata(variable_name, dataset, date, dir)
 end
 
@@ -69,7 +85,7 @@ Base.show(io::IO, metadata::Metadata) =
     "├── name: $(metadata.name)", '\n',
     "├── dataset: $(metadata.dataset)", '\n',
     "├── dates: $(metadata.dates)", '\n',
-    "└── data directory: $(metadata.dir)")
+    "└── dir: $(metadata.dir)")
 
 # Treat Metadata as an array to allow iteration over the dates.
 Base.length(metadata::Metadata) = length(metadata.dates)
@@ -162,7 +178,7 @@ metadata_filename(metadata) = [metadata_filename(metadatum) for metadatum in met
 """
     compute_native_date_range(native_dates, start_date, end_date)
 
-Compute the range of dates that fall within the specified start and end date.
+Compute the range of `native_dates` that fall within the specified `start_date` and `end_date`.
 """
 function compute_native_date_range(native_dates, start_date, end_date)
     if last(native_dates) < end_date
@@ -175,7 +191,7 @@ function compute_native_date_range(native_dates, start_date, end_date)
 
     start_idx = findfirst(x -> x ≥ start_date, native_dates)
     end_idx   = findfirst(x -> x ≥ end_date,   native_dates)
-    start_idx = start_idx > 1 ? start_idx - 1 : start_idx
+    start_idx = (start_idx > 1 && native_dates[start_idx] > start_date) ? start_idx - 1 : start_idx
     end_idx   = isnothing(end_idx) ? length(native_dates) : end_idx
 
     return native_dates[start_idx:end_idx]
diff --git a/src/DataWrangling/restoring.jl b/src/DataWrangling/restoring.jl
@@ -221,15 +221,12 @@ end
 end
 
 """
-    DatasetRestoring(variable_name::Symbol, [ arch_or_grid = CPU(), ];
-                     dataset,
-                     start_date = first_date(dataset, variable_name),
-                     end_date = last_date(dataset, variable_name),
-                     time_indices_in_memory = 2,
-                     time_indexing = Cyclical(),
+    DatasetRestoring(metadata::Metadata,
+                     arch_or_grid = CPU();
+                     rate,
                      mask = 1,
-                     rate = 1,
-                     dir = default_download_directory(dataset),
+                     time_indices_in_memory = 2, # Not more than this if we want to use GPU!
+                     time_indexing = Cyclical(),
                      inpainting = NearestNeighborInpainting(Inf),
                      cache_inpainted_data = true)
 
@@ -247,7 +244,7 @@ from the dataset of choice to the simulation grid and time.
 Arguments
 =========
 
-- `variable_name`: The name of the variable to restore. Choices include:
+- `metadata`: The medatada for a dataset variable to restore. Choices for variables include:
   * `:temperature`,
   * `:salinity`,
   * `:u_velocity`,
@@ -260,54 +257,25 @@ Arguments
                   `arch_or_grid = CPU()` or `arch_or_grid = GPU()`, data is interpolated
                   on-the-fly when the forcing tendency is computed. Default: CPU().
 
-!!! info "Providing `Metadata` instead of `variable_name`"
-    Note that `Metadata` may be provided as the first argument instead of `variable_name`.
-    In this case the `dataset`, `start_date`, and `end_date` kwargs (described below)
-    cannot be provided since they are inferred from `Metadata`.
-
 Keyword Arguments
 =================
 
-- `dataset`: The dataset; required keyword argument if `variable_name` argument is provided.
-
-- `start_date`: The starting date to use for the dataset. Default: `first_date(dataset, variable_name)`.
+- `rate`: The restoring rate, i.e., the inverse of the restoring timescale (in s⁻¹).
 
-- `end_date`: The ending date to use for the dataset. Default: `end_date(dataset, variable_name)`.
+- `mask`: The mask value. Can be a function of `(x, y, z, time)`, an array, or a number.
 
 - `time_indices_in_memory`: The number of time indices to keep in memory. The number is chosen based on
                             a trade-off between increased performance (more indices in memory) and reduced
                             memory footprint (fewer indices in memory). Default: 2.
 
-- `time_indexing`: The time indexing scheme for the field time series.
-
-- `mask`: The mask value. Can be a function of `(x, y, z, time)`, an array, or a number.
-
-- `rate`: The restoring rate, i.e., the inverse of the restoring timescale (in s⁻¹).
-
-- `dir`: The directory where the native data is located. If the data does not exist it will
-         be automatically downloaded. Default: `default_download_directory(dataset)`.
+- `time_indexing`: The time indexing scheme for the field time series. Default: `Cyclical()`.
 
 - `inpainting`: inpainting algorithm, see [`inpaint_mask!`](@ref). Default: `NearestNeighborInpainting(Inf)`.
 
 - `cache_inpainted_data`: If `true`, the data is cached to disk after inpainting for later retrieving.
                           Default: `true`.
 """
-function DatasetRestoring(variable_name::Symbol,
-                          arch_or_grid = CPU();
-                          dataset,
-                          dir = default_download_directory(dataset),
-                          start_date = first_date(dataset, variable_name),
-                          end_date = last_date(dataset, variable_name),
-                          kw...)
-
-    native_dates = all_dates(dataset, variable_name)
-    dates = compute_native_date_range(native_dates, start_date, end_date)
-    metadata = Metadata(variable_name, dataset, dates, dir)
-
-    return DatasetRestoring(metadata, arch_or_grid; kw...)
-end
-
-function DatasetRestoring(metadata,
+function DatasetRestoring(metadata::Metadata,
                           arch_or_grid = CPU();
                           rate,
                           mask = 1,
@@ -316,6 +284,8 @@ function DatasetRestoring(metadata,
                           inpainting = NearestNeighborInpainting(Inf),
                           cache_inpainted_data = true)
 
+    download_dataset(metadata)
+
     fts = FieldTimeSeries(metadata, arch_or_grid;
                           time_indices_in_memory,
                           time_indexing,
@@ -334,13 +304,17 @@ function DatasetRestoring(metadata,
     return DatasetRestoring(fts, maybe_native_grid, mask, field_name, rate)
 end
 
-function Base.show(io::IO, p::DatasetRestoring)
+function Base.show(io::IO, dsr::DatasetRestoring)
     print(io, "DatasetRestoring:", '\n',
-              "├── restored variable: ", summary(p.variable_name), '\n',
-              "├── restoring dataset: ", summary(p.field_time_series.backend.metadata), '\n',
-              "├── restoring rate: ", p.rate, '\n',
-              "├── mask: ", summary(p.mask), '\n',
-              "└── grid: ", summary(p.native_grid))
+              "├── variable_name: ", summary(dsr.variable_name), '\n',
+              "├── rate: ", dsr.rate, '\n',
+              "├── field_time_series: ", summary(dsr.field_time_series), '\n',
+              "│    ├── dataset: ", summary(dsr.field_time_series.backend.metadata.dataset), '\n',
+              "│    ├── dates: ", dsr.field_time_series.backend.metadata.dates, '\n',
+              "│    ├── time_indexing: ", summary(dsr.field_time_series.time_indexing), '\n',
+              "│    └── dir: ", dsr.field_time_series.backend.metadata.dir, '\n',
+              "├── mask: ", summary(dsr.mask), '\n',
+              "└── native_grid: ", summary(dsr.native_grid))
 end
 
 regularize_forcing(forcing::DatasetRestoring, field, field_name, model_field_names) = forcing
diff --git a/test/test_ecco4_en4.jl b/test/test_ecco4_en4.jl
@@ -173,7 +173,8 @@ for arch in test_architectures, dataset in test_datasets
                                                z = (z₁, 0))
 
         for name in (:temperature, :salinity)
-            var_restoring = DatasetRestoring(name, arch; dataset, start_date, end_date, mask, inpainting, rate=1/1000)
+            metadata = Metadata(name; dates, dataset)
+            var_restoring = DatasetRestoring(metadata, arch; mask, inpainting, rate=1/1000)
 
             fill!(var_restoring.field_time_series[1], 1.0)
             fill!(var_restoring.field_time_series[2], 1.0)
@@ -210,7 +211,8 @@ for arch in test_architectures, dataset in test_datasets
             true
         end
 
-        forcing_T = DatasetRestoring(:temperature, arch;  dataset, start_date, end_date, inpainting, rate=1/1000)
+        Tmetadata = Metadata(:temperature; dates, dataset)
+        forcing_T = DatasetRestoring(Tmetadata, arch; inpainting, rate=1/1000)
 
         ocean = ocean_simulation(grid; forcing = (; T = forcing_T), verbose=false)
 
@@ -233,30 +235,44 @@ for arch in test_architectures, dataset in test_datasets
         end_date = DateTime(1993, 5, 1)
         dates = start_date : Month(1) : end_date
 
-        T_restoring = DatasetRestoring(:temperature, arch; dataset, start_date, end_date, inpainting, rate=1/1000)
+        time_indices_in_memory = 2
 
-        times = native_times(T_restoring.field_time_series.backend.metadata)
-        ocean = ocean_simulation(grid, forcing = (; T = T_restoring))
+        Tmetadata1 = Metadata(:temperature; dates, dataset)
+        Tmetadata2 = Metadata(:temperature; start_date, end_date, dataset)
 
-        ocean.model.clock.time = times[3] + 2 * Units.days
-        update_state!(ocean.model)
+        for Tmetadata in (Tmetadata1, Tmetadata2)
+            T_restoring = DatasetRestoring(Tmetadata, arch; time_indices_in_memory, inpainting, rate=1/1000)
 
-        @test T_restoring.field_time_series.backend.start == 3
+            times = native_times(T_restoring.field_time_series.backend.metadata)
+            ocean = ocean_simulation(grid, forcing = (; T = T_restoring))
 
-        # Compile
-        time_step!(ocean)
+            # start a bit after time_index
+            time_index = 3
+            ocean.model.clock.time = times[time_index] + 2 * Units.days
+            update_state!(ocean.model)
 
-        # Try stepping out of the dataset bounds
-        ocean.model.clock.time = last(times) + 2 * Units.days
+            @test time_indices(T_restoring.field_time_series) ==
+                Tuple(range(time_index, length=time_indices_in_memory))
 
-        update_state!(ocean.model)
+            @test T_restoring.field_time_series.backend.start == time_index
 
-        @test begin
+            # Compile
             time_step!(ocean)
-            true
-        end
 
-        # The backend has cycled to the end
-        @test time_indices(T_restoring.field_time_series) == (6, 1)
+            # Try stepping out of the dataset bounds
+            # start a bit after last time_index
+            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) ==
+                mod1.(Tuple(range(length(times), length=time_indices_in_memory)), length(times))
+        end
     end
 end
diff --git a/test/test_jra55.jl b/test/test_jra55.jl
@@ -142,7 +142,7 @@ using ClimaOcean.OceanSeaIceModels: PrescribedAtmosphere
         backend = JRA55NetCDFBackend(10)
         Ta = JRA55FieldTimeSeries(:temperature; dataset, start_date, end_date, backend)
 
-        @test Second(end_date - start_date).value ≈ Ta.times[end-1] - Ta.times[1]
+        @test Second(end_date - start_date).value ≈ Ta.times[end] - Ta.times[1]
 
         # Test we can access all the data
         for t in eachindex(Ta.times)
