#!/usr/bin/env python3 """Concatenate indexed PEM NetCDF outputs along the Time dimension. Builds a cumulative time axis from consecutive files and writes one merged dataset for downstream diagnostics. """ ## @file concat_time_netcdf.py # @author JB Clement # @date 16/10/2023 import os import re import sys import glob import readline import argparse import xarray as xr import dask def complete_path(text, state): matches = glob.glob(text + '*') return matches[state] if state < len(matches) else None readline.set_completer_delims(' \t\n;') readline.set_completer(complete_path) readline.parse_and_bind("tab: complete") def parse_args(): parser = argparse.ArgumentParser( description="Concatenate multiple NetCDF files along the Time dimension" ) parser.add_argument( "--folder", type=str, help="Path to the directory containing the NetCDF files" ) parser.add_argument( "--basename", type=str, help="Base name of the files, e.g., 'diagevo' for files like diagevo1.nc" ) parser.add_argument( "--start", type=int, help="Starting index of the files to include" ) parser.add_argument( "--end", type=int, help="Ending index of the files to include (inclusive)" ) parser.add_argument( "--output", type=str, help="Output filename for the concatenated NetCDF (default: merged.nc)" ) parser.add_argument( "--dt", type=float, help="Time spacing [planetary year] assumed for files holding a single " "record. Without it, such a file is assumed to span one time step " "of the spacing found in the previous files." ) parser.add_argument( "--chunk", type=int, default=100, help="Number of time records read at once (default: 100). Lower it if " "the merge still exceeds the memory limit, raise it to go faster. " "Ignored when dask is unavailable." ) parser.add_argument( "--batch", action="store_true", help="Never prompt: take the default for every option left unset" ) return parser.parse_args() def prompt_with_default(prompt_text, default, cast_fn=None, batch=False): if batch: return default prompt = f"{prompt_text} [press Enter for default {default}]: " while True: try: user_input = input(prompt) except KeyboardInterrupt: print("\nInterrupted.") sys.exit(1) except EOFError: print(f"\nNo input available for '{prompt_text}'. " f"Pass it on the command line or use --batch.") sys.exit(1) if not user_input.strip(): return default try: return cast_fn(user_input) if cast_fn else user_input except ValueError: print(f"Invalid value. Expecting {cast_fn.__name__}. Please try again.") def find_index_range(folder, basename): pattern = os.path.join(folder, f"{basename}*.nc") files = glob.glob(pattern) indices = [] for f in files: name = os.path.basename(f) m = re.match(fr"{re.escape(basename)}(\d+)\.nc$", name) if m: indices.append(int(m.group(1))) if not indices: raise FileNotFoundError(f"No files matching {basename}*.nc found in {folder}") return min(indices), max(indices) def check_grid(ds, fpath, ref_sizes, ref_path): """Abort if the dimensions of 'ds' other than Time differ from the reference ones.""" for dim, size in ds.sizes.items(): if dim == 'Time': continue if dim not in ref_sizes: print(f"Grid mismatch: '{fpath}' has dimension '{dim}', absent from '{ref_path}'.") sys.exit(1) if ref_sizes[dim] != size: print(f"Grid mismatch: dimension '{dim}' is {size} in '{fpath}' " f"but {ref_sizes[dim]} in '{ref_path}'.") sys.exit(1) for dim in ref_sizes: if dim != 'Time' and dim not in ds.sizes: print(f"Grid mismatch: dimension '{dim}' of '{ref_path}' is absent from '{fpath}'.") sys.exit(1) def main(): args = parse_args() # Folder and basename: prompt only for what was not given on the command line folder = args.folder if args.folder is not None else prompt_with_default( "Enter the folder path containing NetCDF files", "diags", batch=args.batch ) basename = args.basename if args.basename is not None else prompt_with_default( "Enter the base filename", "diagevo", batch=args.batch ) # Determine available index range min_idx, max_idx = find_index_range(folder, basename) print(f"Found files from index {min_idx} to {max_idx}.") # Prompt for start/end with discovered defaults start = args.start if args.start is not None else prompt_with_default( "Enter the starting file index", min_idx, cast_fn=int, batch=args.batch ) end = args.end if args.end is not None else prompt_with_default( "Enter the ending file index", max_idx, cast_fn=int, batch=args.batch ) # Validate range if start < min_idx or end > max_idx or start > end: print(f"Invalid range: must be between {min_idx} and {max_idx}, and start <= end.") sys.exit(1) # Output filename output = args.output if args.output is not None else prompt_with_default( "Enter the output filename (including .nc)", "merged.nc", batch=args.batch ) # Build and verify file list file_list = [ os.path.join(folder, f"{basename}{i}.nc") for i in range(start, end + 1) ] for fpath in file_list: if not os.path.isfile(fpath): raise FileNotFoundError(f"File not found: {fpath}") # Opening files with chunks and writing the results slab by slab to save memory open_kwargs = {"decode_times": False} open_kwargs["chunks"] = {"Time": args.chunk} # Offset Time values to make them cumulative datasets = [] time_offset = 0. last_dt = args.dt ref_sizes, ref_path = None, None for fpath in file_list: ds = xr.open_dataset(fpath, **open_kwargs) if 'Time' not in ds.coords: raise ValueError(f"'Time' coordinate not found in {fpath}") # All the files must share the same grid: 'compat="override"' below skips that check if ref_sizes is None: ref_sizes, ref_path = dict(ds.sizes), fpath else: check_grid(ds, fpath, ref_sizes, ref_path) time_vals = ds['Time'].values # 'assign_coords' drops the attributes of the coordinate, so restore them time_attrs = ds['Time'].attrs ds = ds.assign_coords(Time=time_vals + time_offset) ds['Time'].attrs = time_attrs datasets.append(ds) # The PEM writes 'n_yr_run' before incrementing it, so a file holding n records # spans n time steps and the next file must start one step after its last record if len(time_vals) > 1: last_dt = float(time_vals[1] - time_vals[0]) elif last_dt is None: last_dt = 1. print(f"Warning: '{fpath}' holds a single record and no --dt was given: " f"assuming it spans {last_dt} planetary year.") time_offset += last_dt*len(time_vals) # Concatenate. 'data_vars="minimal"' keeps the time-independent grid variables # ('ap', 'bp', 'soildepth', 'cell_area') free of a spurious Time dimension merged_ds = xr.concat( datasets, dim="Time", data_vars="minimal", coords="minimal", compat="override", combine_attrs="override", ) # Optionally decode CF conventions after loading try: merged_ds = xr.decode_cf(merged_ds) except Exception as e: print(f"Warning: CF decoding failed: {e}\nProceeding with raw time values.") # Inspect and save try: tmax = merged_ds.Time.max().values print(f"Final time value: {tmax:.0f}") except Exception: print("Time variables not decoded correctly.") # Coordinates must not get the default fill value of NaN encoding = {name: {"_FillValue": None} for name in merged_ds.coords} print(f"Writing {output}...") merged_ds.to_netcdf(output, encoding=encoding) print(f"Merged dataset written to {output}") if __name__ == "__main__": main()