Expose support for multiple files to DataHandler

Sbozzolo · Sbozzolo · commit 4229def06652 · 2024-11-28T08:59:54.000-08:00
This commit surfaces the new capability introduced recently in
NCFileReaders to read multiple files as if they were one.

In this commit, I only add support for reading one variable split across
multiple files, but extending to composing variables should be
straightforward (when such capability will be required).
diff --git a/NEWS.md b/NEWS.md
@@ -4,13 +4,45 @@ ClimaUtilities.jl Release Notes
 main
 ------
 
-#### Support reading time data across multiple files. PR [#127](https://github.com/CliMA/ClimaUtilities.jl/pull/127)
+v0.1.20
+------
+
+#### Support reading time data across multiple files. PRs [#127](https://github.com/CliMA/ClimaUtilities.jl/pull/127), [#132](https://github.com/CliMA/ClimaUtilities.jl/pull/132)
 
 `NCFileReader`s can now read multiple files at the same time. The files have to
 contain temporal data for the given variable and they are aggregated along the
 time dimension. To use this feature, just pass a vector of file paths to the
 constructor.
 
+This capability is also available to `DataHandler`s and `TimeVaryingInput`s. To
+use this feature, just pass the list of files that contain your variable of
+interested, for example
+```julia
+timevaryinginput = TimeVaryingInputs.TimeVaryingInput(["era5_1980.nc", "era5_1981.nc"],
+                                                      "u",
+                                                      target_space,
+                                                      start_date = Dates.DateTime(1980, 1, 1),
+                                                      regridder_type = :InterpolationsRegridder)
+```
+You can also compose variables
+```julia
+timevaryinginput = TimeVaryingInputs.TimeVaryingInput(["era5_1980.nc", "era5_1981.nc", "era5_1982.nc"],
+                                                      ["u", "v"],
+                                                      target_space,
+                                                      start_date = Dates.DateTime(1980, 1, 1),
+                                                      regridder_type = :InterpolationsRegridder,
+                                                      compose_function = (x, y) -> sqrt(x^2 + y^2))
+```
+
+When you compose variables, pay attention that `TimeVaryingInput` implements
+some heuristics to disambiguate the case where the passed list of files is split
+along the time or the variable dimension. You can always pass a list of lists to
+be explicit in your intentions. Read the
+[documentation](https://clima.github.io/ClimaUtilities.jl/dev/datahandling.html#Heuristics-to-do-what-you-mean)
+to learn more about this.
+
+This capability is only available for the `InterpolationsRegridder`.
+
 v0.1.19
 ------
 
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ClimaUtilities"
 uuid = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
 authors = ["Gabriele Bozzola <gbozzola@caltech.edu>", "Julia Sloan <jsloan@caltech.edu>"]
-version = "0.1.19"
+version = "0.1.20"
 
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
diff --git a/docs/src/datahandling.md b/docs/src/datahandling.md
@@ -9,17 +9,18 @@ This is no trivial task. Among the challenges:
 - IO can be very expensive,
 - CPU/GPU communication can be a bottleneck.
 
-The `DataHandling` takes the divide and conquer approach: the various core tasks
+The `DataHandling` takes the divide-and-conquer approach: the various core tasks
 and features and split into other independent modules (chiefly
 [`FileReaders`](@ref), and [`Regridders`](@ref)). Such modules can be developed,
 tested, and extended independently (as long as they maintain a consistent
 interface). For instance, if need arises, the `DataHandler` can be used (almost)
 directly to process files with a different format from NetCDF.
 
 The key struct in `DataHandling` is the `DataHandler`. The `DataHandler`
-contains one or more `FileReader`(s), a `Regridder`, and other metadata necessary to perform
-its operations (e.g., target `ClimaCore.Space`). The `DataHandler` can be used
-for static or temporal data, and exposes the following key functions:
+contains one or more `FileReader`(s), a `Regridder`, and other metadata
+necessary to perform its operations (e.g., target `ClimaCore.Space`). The
+`DataHandler` can be used for static or temporal data, and exposes the following
+key functions:
 - `regridded_snapshot(time)`: to obtain the regridded field at the given `time`.
                     `time` has to be available in the data.
 - `available_times` (`available_dates`): to list all the `times` (`dates`) over
@@ -66,6 +67,50 @@ are composed.
 Composing multiple input variables is currently only supported with the
 `InterpolationsRegridder`, not with `TempestRegridder`.
 
+Sometimes, the time development of a variable is split across multiple NetCDF
+files. `DataHandler` knows how to combine them and treat multiple files as if
+they were a single one. To use this feature, just pass the list of NetCDF files
+(while the file don't have to be sorted, it is good practice to pass them sorted
+in ascending order by time). 
+
+#### Heuristics to do-what-you-mean
+
+`DataHandler` tries to interpret the files provided and identify if they are
+split across variables or along the time dimension. The heuristics implement are
+the following:
+- When just a file is passed, it is assumed that it contains everything 
+- When multiple files are passed, `DataHandler` will assume that the files are
+  split along variables if the number of files is the same the number of
+  variables, otherwise, it will assume that each file contains all the variables
+  for a portion of the total time.
+- When the above assumption is incorrect, you can pass a list of list of files
+  that fully specifies variables and times.
+
+For example, 
+```julia
+data_handler = DataHandling.DataHandler(
+    ["era1980.nc", "era1981.nc"],
+    ["lai_hv", "lai_lv"],
+    target_space;
+    compose_function = (x, y) -> x + y,
+)
+```
+
+In this case, `DataHandler` will incorrectly assume that `lai_hv` is contained
+in `era1980.nc`, and `lai_lv` is in `era1980.nc`. Instead, construct the
+`data_handler` by passing a list of lists 
+```julia
+files = ["era1980.nc", "era1981.nc"] 
+data_handler = DataHandling.DataHandler(
+    [files, files],
+    ["lai_hv", "lai_lv"],
+    target_space;
+    compose_function = (x, y) -> x + y,
+)
+```
+where each element of the list is the collection of files that contain the time
+evolution of that variable.
+
 ## Example: Linear interpolation of a single data variable
 
 As an example, let us implement a simple linear interpolation for a variable `u`
@@ -108,6 +153,11 @@ function linear_interpolation(data_handler, time)
 end
 ```
 
+If, for example, the data was split across multiple files named `era5_1980.nc`,
+`era5_1981.nc`, ... (e.g., each file containing one year), we could directly
+pass the list to the constructor for `DataHandler` (instead of just passing one
+file path), which knows how to combine them.
+
 ### Example appendix: Using multiple input data variables
 
 Suppose that the input NetCDF file `era5_example.nc` contains two variables `u`
diff --git a/docs/src/inputs.md b/docs/src/inputs.md
@@ -76,7 +76,7 @@ compose_function = (x, y) -> x + y
 # Define pre-processing function to convert units of input
 unit_conversion_func = (data) -> 1000 * data
 
-data_handler = TimeVaryingInputs.TimeVaryingInput("era5_example.nc",
+timevaryinginput = TimeVaryingInputs.TimeVaryingInput("era5_example.nc",
                                         ["u", "v"],
                                         target_space,
                                         start_date = Dates.DateTime(2000, 1, 1),
@@ -88,6 +88,29 @@ data_handler = TimeVaryingInputs.TimeVaryingInput("era5_example.nc",
 The same arguments (excluding `start_date`) could be passed to a
 `SpaceVaryingInput` to compose multiple input variables with that type.
 
+#### Example: Data split across multiple NetCDF files
+
+Often, large datasets come chunked, meaning that the data is split across
+multiple files with each file containing only a subset of the time interval.
+`TimeVaryingInput`s know to combine data across multiple files as it were
+provided in a single file. To do use this feature, just pass the list of file
+paths. While it is not required for the files to be in order, it is good
+practice to pass them in ascending order by time.
+
+For example:
+```julia
+timevaryinginput = TimeVaryingInputs.TimeVaryingInput(["era5_1980.nc", "era5_1981.nc"],
+                                                       "u",
+                                                       target_space,
+                                                       start_date = Dates.DateTime(1980, 1, 1),
+                                                       regridder_type = :InterpolationsRegridder
+                                                       )
+```
+
+This capability is only available for the `InterpolationsRegridder`.
+
+Read more about this feature in the page about [`DataHandler`](@ref).
+
 ### Extrapolation boundary conditions
 
 `TimeVaryingInput`s can have multiple boundary conditions for extrapolation. By
diff --git a/ext/DataHandlingExt.jl b/ext/DataHandlingExt.jl
@@ -93,9 +93,40 @@ struct DataHandler{
     preallocated_read_data::PR
 end
 
+
+"""
+    _check_file_paths_varnames(file_paths, varnames, regridder_type, compose_function)
+
+Check consistency of `file_paths`, `varnames`, `regridder_type`, and `compose_function` for
+our current `DataHandler`.
+"""
+function _check_file_paths_varnames(
+    file_paths,
+    varnames,
+    regridder_type,
+    compose_function,
+)
+    # Verify that the number of file paths and variable names are consistent
+    if length(varnames) == 1
+        # Multiple files are not not supported by TempestRegridder
+        (length(file_paths) > 1 && regridder_type == :TempestRegridder) &&
+            error("TempestRegridder does not support multiple input files")
+    else
+        # We have multiple variables
+        # This is not supported by TempestRegridder
+        regridder_type == :TempestRegridder &&
+            error("TempestRegridder does not support multiple input variables")
+
+        # We need a compose_function when passed multiple variables
+        compose_function == identity && error(
+            "`compose_function` must be specified when using multiple input variables",
+        )
+    end
+end
+
 """
-    DataHandler(file_paths::Union{AbstractString, AbstractArray{<:AbstractString}},
-                varnames::Union{AbstractString, AbstractArray{<:AbstractString}},
+    DataHandler(file_paths,
+                varnames,
                 target_space::ClimaCore.Spaces.AbstractSpace;
                 start_date::Dates.DateTime = Dates.DateTime(1979, 1, 1),
                 regridder_type = nothing,
@@ -104,8 +135,14 @@ end
                 file_reader_kwargs = ())
 
 Create a `DataHandler` to read `varnames` from `file_paths` and remap them to `target_space`.
-`file_paths` may contain either one path for all variables or one path for each variable.
-In the latter case, the entries of `file_paths` and `varnames` are expected to match based on position.
+
+This function supports reading across multiple files and composing variables that are in
+different files.
+
+
+`file_paths` may contain either one path for all variables or one path for each variable. In
+the latter case, the entries of `file_paths` and `varnames` are expected to match based on
+position.
 
 The DataHandler maintains an LRU cache of Fields that were previously computed.
 
@@ -114,7 +151,21 @@ Creating this object results in the file being accessed (to preallocate some mem
 Positional arguments
 =====================
 
-- `file_paths`: Paths of the NetCDF file(s) that contain the input data.
+- `file_paths`: Paths of the NetCDF file(s) that contain the input data. `file_paths` should
+  be as "do-what-I-mean" as possible, meaning that it should behave as you expect.
+
+  To be specific, there are three options for `file_paths`:
+  - It is a string that points to a single NetCDF file.
+  - It is a list that points to multiple NetCDF files. In this case, we support two modes:
+    1. if `varnames` is a vector with the number of entries as `file_paths`, we assume that
+       each file contains a different variable.
+    2. otherwise, we assume that each file contains all the variables and is temporal chunk.
+  - It is a list of lists of paths to NetCDF files, where the inner list identifies temporal
+    chunks of a given variable, and the outer list identifies different variables
+    (supporting the mode where different variables live in different files and their time
+    development is split across multiple files). In other words, `file_paths[i]` is the list
+    of files that define the temporal evolution of `varnames[i]`
+
 - `varnames`: Names of the datasets in the NetCDF that have to be read and processed.
 - `target_space`: Space where the simulation is run, where the data has to be regridded to.
 
@@ -138,13 +189,16 @@ everything more type stable.)
                       It can be a NamedTuple, or a Dictionary that maps Symbols to values.
 - `file_reader_kwargs`: Additional keywords to be passed to the constructor of the file reader.
                         It can be a NamedTuple, or a Dictionary that maps Symbols to values.
-- `compose_function`: Function to combine multiple input variables into a single data variable.
-                    The default, to be used in the case of one input variable, is the identity.
-                    Note that the order of `varnames` must match the argument order of `compose_function`.
+- `compose_function`: Function to combine multiple input variables into a single data
+                      variable. The default, to be used in the case of one input variable,
+                      is the identity. The compose function has to take N arguments, where
+                      N is the number of variables in `varnames`, and return a scalar.
+                      The order of the arguments in `compose_function` has to match the order
+                      of `varnames`. This function will be broadcasted to data read from file.
 """
 function DataHandling.DataHandler(
-    file_paths::Union{AbstractString, AbstractArray{<:AbstractString}},
-    varnames::Union{AbstractString, AbstractArray{<:AbstractString}},
+    file_paths,
+    varnames,
     target_space::ClimaCore.Spaces.AbstractSpace;
     start_date::Union{Dates.DateTime, Dates.Date} = Dates.DateTime(1979, 1, 1),
     regridder_type = nothing,
@@ -182,38 +236,51 @@ function DataHandling.DataHandler(
         varnames = [varnames]
     end
 
-    # Verify that the number of file paths and variable names match
-    (length(file_paths) > 1 && length(file_paths) != length(varnames)) && error(
-        "Number of file paths ($(length(file_paths))) and variable names ($(length(varnames))) do not match.",
-    )
-
-    # Verify that `compose_function` is specified when using multiple input variables
-    (length(varnames) > 1 && compose_function == identity) && error(
-        "`compose_function` must be specified when using multiple input variables",
-    )
-
-    # Verify that `compose_function` is identity when using a single input variable
-    (length(varnames) == 1 && compose_function != identity) && error(
-        "`compose_function` must be identity when using a single input variable",
-    )
-
-    # TempestRegridder does not support multiple input variables
-    (length(varnames) > 1 && regridder_type == :TempestRegridder) &&
-        error("TempestRegridder does not support multiple input variables")
-
     # Determine which regridder to use if not already specified
     regridder_type =
         isnothing(regridder_type) ? Regridders.default_regridder_type() :
         regridder_type
 
-    # Construct a file reader, which deals with ingesting data and is possibly buffered/cached, for each input file
-    file_readers = Dict{String, AbstractFileReader}()
-    all_vars_in_same_file = length(file_paths) == 1
-    for (i, varname) in enumerate(varnames)
-        file_path = all_vars_in_same_file ? first(file_paths) : file_paths[i]
-        file_readers[varname] =
-            NCFileReader(file_path, varname; file_reader_kwargs...)
+    _check_file_paths_varnames(
+        file_paths,
+        varnames,
+        regridder_type,
+        compose_function,
+    )
+
+    # We have to deal with the case with have 1 FileReader (with possibly multiple files),
+    # or with N FileReaders (for when variables are split across files, and with possibly
+    # multiple files). To accommodate all these cases, we cast everything into the format
+    # where we have a list of lists, where the outer list is along variable names, and the
+    # inner list is along times. This is the most general input we expect from this
+    # constructor.
+
+    is_file_paths_list_of_lists = !(first(file_paths) isa AbstractString)
+
+    if !is_file_paths_list_of_lists
+        # If is_file_paths_list_of_lists not already a list of lists, we have two options:
+        # 1. file_paths identifies the temporal development of the variables
+        # 2. file_paths identifies different variables
+
+        # We use as heuristic that when the number of files provided is the same as the
+        # number of variables, that means that the files include different variables
+        if length(file_paths) == length(varnames)
+            # One file per variable
+            file_paths = [[f] for f in file_paths]
+        else
+            # Every file per every variable
+            file_paths = [copy(file_paths) for _ in varnames]
+        end
     end
+    # Now, we have a list of lists, where file_paths[i] is the list of files that define the
+    # temporal evolution of varnames[i]
+
+    # Construct the file readers, which deals with ingesting data and is possibly
+    # buffered/cached, for each variable
+    file_readers = Dict(
+        varname => NCFileReader(paths, varname; file_reader_kwargs...) for
+        (paths, varname) in zip(file_paths, varnames)
+    )
 
     # Verify that the spatial dimensions are the same for each variable
     @assert length(
@@ -248,9 +315,11 @@ function DataHandling.DataHandler(
             regridder_kwargs = merge((; regrid_dir), regridder_kwargs)
         end
 
-        # Note: using one arbitrary element of `varnames` and of `file_paths` assumes
-        #  that all input variables will use the same regridding
-        regridder_args = (target_space, first(varnames), first(file_paths))
+        # Note: using one arbitrary element of `varnames` and of `file_paths`
+        # assumes that all input variables will use the same regridding (there
+        # are two firsts in file_paths because we now have a list of lists)
+        regridder_args =
+            (target_space, first(varnames), first(first(file_paths)))
     elseif regridder_type == :InterpolationsRegridder
         regridder_args = (target_space,)
     end
diff --git a/test/data_handling.jl b/test/data_handling.jl
