Allow grouping of input files into output files

README.md

@@ -20,6 +20,30 @@ If there are ancillary fields that should only be present in only some of the ou
 
 Example of these variables are the `latitude_longitude` found in atmosphere files or the `uarea`, `tmask`, `tarea`, `VGRDb`, `VGRDi`, `VGRDs` variables from ice files.
 
+## Grouping Input Files
+
+Input file paths can be grouped together using the `--input-group-regex INPUT_GROUP_REGEX` command line option.
+This can be used for example to group together 12 monthly input files into one yearly output file per variable.
+This option uses a regex with a named capture group "wild" to identify the portion of the files that varies across the group.
+For example, to group togther these two sets of monthly files into yearly output files:
+```
+aiihca.pa-234501_mon.nc  aiihca.pe-234501_dai.nc
+aiihca.pa-234502_mon.nc  aiihca.pe-234502_dai.nc
+aiihca.pa-234503_mon.nc  aiihca.pe-234503_dai.nc
+aiihca.pa-234504_mon.nc  aiihca.pe-234504_dai.nc
+aiihca.pa-234505_mon.nc  aiihca.pe-234505_dai.nc
+aiihca.pa-234506_mon.nc  aiihca.pe-234506_dai.nc
+aiihca.pa-234507_mon.nc  aiihca.pe-234507_dai.nc
+aiihca.pa-234508_mon.nc  aiihca.pe-234508_dai.nc
+aiihca.pa-234509_mon.nc  aiihca.pe-234509_dai.nc
+aiihca.pa-234510_mon.nc  aiihca.pe-234510_dai.nc
+aiihca.pa-234511_mon.nc  aiihca.pe-234511_dai.nc
+aiihca.pa-234512_mon.nc  aiihca.pe-234512_dai.nc
+```
+use the regex `aiihca\.p[ae]-\d{4}(?P<wild>\d{2})_(mon|dai)\.nc` which uses the `wild` capture group to match the month.
+
+Any file paths that don't match the regex will be processed individually.
+
 ## Config File
 
 The `-c`/`--command-line-file` option can be used to supply a filepath to a file that contains command line options.
@@ -73,13 +97,19 @@ options:
                         access-esm1p6.{component}.{dimensions}.{field}.{freq}.{time_cell_method}.{datestamp}.nc
                         splitnc will attempt to deduce all the components of the filename. If this option is not given
                         {field}_{original_filename} will be used.
+                        Note: This option also enables special preprocessing for ESM1.6 files.
   --fix-cell-methods    Correct cell_methods by adding 'time: point' to cell_methods for variables that have 'time' but
                         not 'time_bnds' and no other 'time' cell_methods.
   --file-freq FILE_FREQ
                         Specify the frequency of the files (not the data), e.g. if each file contains a month of data
                         then the file-frequency is '1mon'. Used to determine the resolution of the timestamp for ESM1.6
                         filenames. Follows the ACCESS frequency vocabulary (e.g. '1yr', '1mon', '1day', '1hr'), any
                         unrecognised frequency will use the full timestamp. Defaults to '1yr'.
+  --input-group-regex INPUT_GROUP_REGEX
+                        Specify a regex that will be used to group a subset of the input into a single set. E.g. group together 12
+                        input monthly files to a single year of output. Use a named capture group "wild" to specify the portion of
+                        the filename that varies. E.g. to group monthly files with this pattern - "aiihca.pa-YYYYMM_mon.nc" - use
+                        the regex \"aiihca\.pa-\d{4}(?P<wild>\d{2})_mon\.nc\".
   --output-dir OUTPUT_DIR
                         Output directory for the processed files. If not given output files will be placed in the same
                         directory as the original file.

pyproject.toml

@@ -17,6 +17,7 @@ requires-python = ">=3.9"
 dependencies = [
     "netcdf4",
     "xarray",
+    "dask",
 ]
 
 [project.optional-dependencies]

src/splitnc/esm1p6.py

@@ -131,7 +131,8 @@ def _build_datestamp(ds, field_name, file_freq):
 
     # Calculate the middle point
     first, last = time_arr.min(), time_arr.max()
-    datestamp_dt = (first + (last - first) / 2).dt
+    datestamp_dt = (first + (last - first) / 2).compute().dt
+        # Need to .compute when using open_mfdataset
 
     return "." + datestamp_dt.strftime(fmt).data.flatten()[0]
 
@@ -156,3 +157,17 @@ def build_esm1p6_filename(ds, field_name, input_filepath, esm1p6_filename=False,
         raise
 
     return template.format(**d)
+
+
+def preprocess_esm1p6_files(ds):
+    """
+    Prepare ESM1.6 files before grouping.
+    - Round off surface_altitude in atmos files to remove variation due to numerical issues
+    """
+    if 'surface_altitude' in ds:
+        # Surface altitude is measured in meters
+        # surface_altitude in different files can vary down around 10^-10
+        # Round it off to 4 decimal places
+        ds['surface_altitude'] = ds['surface_altitude'].round(4)
+
+    return ds

src/splitnc/splitnc.py

@@ -10,7 +10,7 @@
 
 import xarray as xr
 
-from splitnc.esm1p6 import build_esm1p6_filename
+from splitnc.esm1p6 import build_esm1p6_filename, preprocess_esm1p6_files
 
 
 def determine_field_vars(ds):
@@ -272,7 +272,67 @@ def build_filename(ds, field_name, input_filepath, esm1p6_filename=False, file_f
         return f"{field_name}_{input_filepath.name}"
 
 
-def process_file(filepath, **kwargs):
+def group_filepaths(filepaths, group_regex):
+    r"""
+    Group together files from the list of filepaths that match the given regex
+    with only the portion in the capture group "wild" varying.
+
+    E.g. if files follow the patterns
+    - aiihca.pa-YYYYMM_mon.nc and
+    - aiihca.pe-YYYYMM_dai.nc
+    use "aiihca\.p[ae]-\d{4}(?P<wild>\d{2})_(mon|dai)\.nc" to group together
+    months for each year and freq. Grouped filepaths will be returned as a list
+
+    Any filepath that doesn't match the regex will be returned alone, i.e. in a
+    group of length 1.
+
+    Returns a list of lists of filepath strings
+    """
+    grouped_filepaths = []
+    while len(filepaths) > 0:
+        f = filepaths[0]
+        if m:=re.search(group_regex, f):
+            # We need to know which indices the "wild" group has in the filepath
+            wild_span = m.span("wild")
+
+            # Replace the wild match in the orginal string with a wild regex
+            # Use double {{ }} to escape them in f-strings
+            group_regx = re.compile(m.string[:wild_span[0]] + f".{{{len(m['wild'])}}}" + m.string[wild_span[1]:])
+
+            # Get the filepaths that match the regex and remove them from the filepaths list
+            group_list = [fp for fp in filepaths if group_regx.search(fp)]
+            filepaths = [fp for fp in filepaths if not group_regx.search(fp)]
+        else:
+            # If the regex doesn't match the group regex treat it as a solo group
+            group_list = [filepaths.pop(0)]
+
+        grouped_filepaths.append(group_list)
+
+    return grouped_filepaths
+
+
+def process_files(**kwargs):
+    # Prepare the filepath list
+    filepaths_list = kwargs.pop("filepaths")
+    if input_group_regex:=kwargs['input_group_regex']:
+        logging.debug(f"Grouping filepaths according to regex: {input_group_regex}")
+
+        # Group files together according to the input_file_date_regex
+        filepaths_list = group_filepaths(filepaths_list, input_group_regex)
+    else:
+        # Treat every filepath as a size 1 group
+        filepaths_list = [[f] for f in filepaths_list]
+
+    logging.debug("Filepaths groups as follows:\n" + "\n".join(
+        [f"{i}: {filepaths}" for i, filepaths in enumerate(filepaths_list)]
+    ))
+
+    # Process each filepath group
+    for filepaths in filepaths_list:
+        process_filegroup(filepaths, **kwargs)
+
+
+def process_filegroup(filepaths, **kwargs):
     # Define default kwargs and update them with kwargs
     kwargs = {
         "excluded_vars": [],
@@ -287,12 +347,51 @@ def process_file(filepath, **kwargs):
         "overwrite": False,
     } | kwargs
 
-    logging.debug(f"Processing {filepath}")
-    filepath = Path(filepath)
+    logging.debug(f"Processing {filepaths}")
+
+    filepaths = [Path(f) for f in filepaths]
+
+    # xarray drops .encoding when using open_mfdataset with more than one file
+    # So save the encodings when loading and reapply
+    encoding_map = {}
+    def save_encoding(ds):
+        for v in ds.variables:
+            enc = ds[v].encoding
+
+            # Remove some keys from the encoding as these will not always match
+            # and aren't important here
+            keys_to_delete = ['source', 'chunksizes', 'preferred_chunks', 'original_shape']
+            for del_key in keys_to_delete:
+                try:
+                    del enc[del_key]
+                except KeyError:
+                    # If the key isn't there do nothing
+                    pass
+
+            if v in encoding_map and encoding_map[v] != enc:
+                raise ValueError(f"Encodings for {v} doesn't match across all files: {enc}")
+
+            encoding_map[v] = enc
+
+        return ds
+
+    def preprocess(ds):
+        ds = save_encoding(ds)
+
+        if kwargs['use_esm1p6_filenames']:
+            ds = preprocess_esm1p6_files(ds)
+
+        return ds
 
     # Use cftime to suppress warnings
     decoder = xr.coders.CFDatetimeCoder(time_unit='us')
-    with xr.open_dataset(filepath, decode_times=decoder) as ds:
+    with xr.open_mfdataset(filepaths, decode_times=decoder, combine="nested", 
+        compat="no_conflicts", join="outer", preprocess=preprocess) as ds:
+        # Reapply the saved encodings if they're missing
+        for v in ds.variables:
+            if not ds[v].encoding:
+                ds[v].encoding = encoding_map[v]
+
         # Resolve any regex in the excluded_vars list
         if excluded_vars:=kwargs["excluded_vars"]:
             excluded_vars = match_regex_list(excluded_vars, ds.variables)
@@ -326,6 +425,9 @@ def process_file(filepath, **kwargs):
             rename_dict = {}
         logging.debug(f"Rename dict is {rename_dict}")
 
+        if len(field_vars) == 0:
+            logging.warning(f"No field variables to process for {filepaths}")
+
         for v in field_vars:
             # Get the list of vars to keep for this field
             logging.debug(f"Determining dependent variables for field variable {v}")
@@ -381,16 +483,17 @@ def process_file(filepath, **kwargs):
             if kwargs["fix_cell_methods"]:
                 fix_cell_methods(ds_v, v)
 
+            # Output path construction assumes the first path can be used
             if output_dir:=kwargs["output_dir"]:
                 output_dir = Path(output_dir)
             else:
-                output_dir = filepath.parent
+                output_dir = filepaths[0].parent
 
             # Build the output filepath
             filename = build_filename(
                 ds=ds_v,
                 field_name=v,
-                input_filepath=filepath,
+                input_filepath=filepaths[0],
                 esm1p6_filename=kwargs["use_esm1p6_filenames"],
                 file_freq=kwargs["file_freq"],
             )
@@ -489,7 +592,8 @@ def globbable_string_list(string_list):
         help="Use the ESM1.6 filename pattern for the output files: "
         "access-esm1p6.{component}.{dimensions}.{field}.{freq}.{time_cell_method}.{datestamp}.nc"
         " splitnc will attempt to deduce all the components of the filename. "
-        "If this option is not given {field}_{original_filename} will be used."
+        "If this option is not given {field}_{original_filename} will be used. "
+        "Note: This option also enables special preprocessing for ESM1.6 files."
     )
     parser.add_argument(
         "--fix-cell-methods",
@@ -508,6 +612,15 @@ def globbable_string_list(string_list):
         "'1hr'), any unrecognised frequency will use the full timestamp. "
         "Defaults to '1yr'."
     )
+    parser.add_argument(
+        "--input-group-regex",
+        help="Specify a regex that will be used to group a subset of the input "
+        "into a single set. E.g. group together 12 input monthly files to "
+        "a single year of output. Use a named capture group \"wild\" to "
+        "specify the portion of the filename that varies. E.g. to group monthly "
+        "files with this pattern - \"aiihca.pa-YYYYMM_mon.nc\" - use the regex "
+        r"\"aiihca\.pa-\d{4}(?P<wild>\d{2})_mon\.nc\"."
+    )
     parser.add_argument(
         "--output-dir",
         help="Output directory for the processed files. If not given output "
@@ -541,9 +654,10 @@ def globbable_string_list(string_list):
     args = parser.parse_args(args=cmdline_args)
 
     # File paths may need flattened since glob was used
-    args.filepaths = [
+    # Sort the list to ensure repeatable behaviour
+    args.filepaths = sorted([
         filepath for glob_list in args.filepaths for filepath in glob_list
-    ]
+    ])
 
     # If the command line yaml was supplied use the contents instead of argv
     if args.command_line_file:
@@ -572,8 +686,7 @@ def main():
         logging.error("No files to process.")
         raise ValueError("No files to process.")
 
-    for f in args.filepaths:
-        process_file(f, **vars(args))
+    process_files(**vars(args))
 
 
 if __name__ == "__main__":