diff --git a/plot_mpas_netcdf.py b/plot_mpas_netcdf.py index 8ff09a2..d51f1f1 100644 --- a/plot_mpas_netcdf.py +++ b/plot_mpas_netcdf.py @@ -12,7 +12,8 @@ import signal import traceback import time -from datetime import datetime +import re +from datetime import datetime, timedelta import gc, psutil proc = psutil.Process(os.getpid()) @@ -88,17 +89,25 @@ def open_ux_subset(gridfile, datafiles, vars_to_keep): logger.error("For MPAS this is usually a history file or an init.nc file") raise e - keep_set = set(vars_to_keep) | {"xtime"} # always keep xtime - + keep_set = set(vars_to_keep) | {"xtime"} # always keep xtime if present xr_ds_list = [] - for f in datafiles: + for i, f in enumerate(datafiles): logger.debug(f"Opening dataset file {f}\nMemory usage:{proc.memory_info().rss/1024**2} MB") ds = xr.open_dataset(f, decode_cf=False, chunks={}) # lazy missing = [v for v in vars_to_keep if v not in ds.variables] if missing: raise KeyError(f"{f} missing required variables: {missing}") + # Normalize time to xtime if not already present + if "xtime" not in ds.variables: + timestring = _extract_timestring(ds, f, i) + ds["xtime"] = xr.DataArray( + [[c for c in timestring.ljust(64)[:64]]], + dims=["Time", "StrLen"] + ) + logger.debug(f"Synthesized xtime={timestring} for {f}") + available_keep = [v for v in keep_set if v in ds.variables] xr_ds_list.append(ds[available_keep]) @@ -453,6 +462,41 @@ def deep_merge(dict1, dict2): return result +def _extract_timestring(ds: xr.Dataset, filepath: str, file_index: int = 0) -> str: + """ + Extract a '%Y-%m-%d_%H:%M:%S' timestring from a single raw xarray Dataset. + Tries: CF Time variable, filename parsing, dummy fallback. + Only called from open_ux_subset() when xtime is absent. + file_index is used to generate unique dummy timestamps when all else fails. + """ + # CF-compliant Time variable with units attribute + if "Time" in ds.variables and "units" in ds["Time"].attrs: + units = ds["Time"].attrs["units"] + match = re.match(r'(\w+) since (\d{4}-\d{2}-\d{2}[T _]\d{2}:\d{2}:\d{2})', units) + if match: + unit_type = match.group(1).lower() + base_time = datetime.strptime( + match.group(2).replace("T", " ").replace("_", " "), "%Y-%m-%d %H:%M:%S" + ) + multipliers = {"seconds": 1, "minutes": 60, "hours": 3600, "days": 86400} + if unit_type in multipliers: + offset_secs = float(ds["Time"].isel(Time=0).values) * multipliers[unit_type] + valid_time = base_time + timedelta(seconds=offset_secs) + logger.debug(f"Decoded time from CF Time variable with units: {units}") + return valid_time.strftime("%Y-%m-%d_%H:%M:%S") + + # Parse time from filename (e.g. diag.2025-06-07_00.00.00.nc) + match = re.search(r'(\d{4}-\d{2}-\d{2}_\d{2})\.(\d{2})\.(\d{2})', os.path.basename(filepath)) + if match: + timestring = f"{match.group(1)}:{match.group(2)}:{match.group(3)}" + logger.info(f"No time variable found; extracted time from filename: {timestring}") + return timestring + + dummy = datetime(1900, 1, 1) + timedelta(hours=file_index) + logger.warning(f"Could not determine valid time for {filepath}; using dummy time value {dummy.strftime('%Y-%m-%d_%H:%M:%S')}") + return dummy.strftime("%Y-%m-%d_%H:%M:%S") + + def setup_args(config_d: dict,uxds: ux.UxDataset): """ Sets up the argument list for plotit to allow for parallelization with Python starmap @@ -484,32 +528,14 @@ def setup_args(config_d: dict,uxds: ux.UxDataset): else: raise TypeError(f"Invalid level {vardict['lev']} specified for variable {var}") - # Extract time strings - # If multiple timesteps in a dataset, loop over times if "Time" in uxds[var].dims: - times=[] - for i in range(uxds.sizes["Time"]): - logger.debug(f"Plotting time step {i}") - if "xtime" in uxds: - times.append("".join(uxds["xtime"].isel(Time=i).values.astype(str))) - else: - logger.warning(f"'xtime' variable not found in input file, using dummy time value") - times.append("1900-01-01_00:00:00") for lev in levels: - i=0 - for timestring in times: + for i in range(uxds.sizes["Time"]): + timestring = "".join(uxds["xtime"].isel(Time=i).values.astype(str)).strip() args.append( (config_d,uxds,var,lev,i,timestring) ) - i+=1 - else: logger.debug(f"{var} has no time dimension") - if "xtime" in uxds: - logger.debug("Using first xtime value in file") - timestring="".join(uxds["xtime"].isel(Time=0).values.astype(str)) - else: - logger.warning(f"'xtime' variable not found in input file, using dummy time value") - timestring="1900-01-01_00:00:00" - + timestring = "".join(uxds["xtime"].isel(Time=0).values.astype(str)).strip() for lev in levels: args.append( (config_d,uxds,var,lev,-1,timestring) )