diff --git a/pcmdi_metrics/io/xcdat_openxml.py b/pcmdi_metrics/io/xcdat_openxml.py index 5c3b2208c..423cf3dd2 100644 --- a/pcmdi_metrics/io/xcdat_openxml.py +++ b/pcmdi_metrics/io/xcdat_openxml.py @@ -114,7 +114,7 @@ def fix_noncompliant_attr(ds: xr.Dataset) -> xr.Dataset: def _xcdat_openxml( - xmlfile: str, data_var: str = None, decode_times: bool = True + xmlfile: str, data_var: str = None, decode_times: bool = True, chunks=None ) -> xr.Dataset: """Open input file (xml generated by cdscan) @@ -142,11 +142,11 @@ def _xcdat_openxml( if len(ncfile_list) > 1: ds = xc.open_mfdataset( - ncfile_list, data_var=data_var, decode_times=decode_times + ncfile_list, data_var=data_var, decode_times=decode_times, chunks=chunks ) else: ds = xc.open_dataset( - ncfile_list[0], data_var=data_var, decode_times=decode_times + ncfile_list[0], data_var=data_var, decode_times=decode_times, chunks=chunks ) return ds diff --git a/pcmdi_metrics/mean_climate/lib/calculate_climatology.py b/pcmdi_metrics/mean_climate/lib/calculate_climatology.py index e49e87c00..3266c62a5 100644 --- a/pcmdi_metrics/mean_climate/lib/calculate_climatology.py +++ b/pcmdi_metrics/mean_climate/lib/calculate_climatology.py @@ -76,8 +76,8 @@ def calculate_climatology( # Subset given time period d = d.sel( time=slice( - start_yr_str + "-" + start_mo_str + "-" + start_da_str, - end_yr_str + "-" + end_mo_str + "-" + end_da_str, + f"{start_yr_str}-{start_mo_str}-{start_da_str}", + f"{end_yr_str}-{end_mo_str}-{end_da_str}", ) ) @@ -113,34 +113,61 @@ def calculate_climatology( # Save to netcdf file if periodinname is None: addf = ( - "." - + start_yr_str - + start_mo_str - + "-" - + end_yr_str - + end_mo_str - + "." - + s - + "." - + ver - + ".nc" + f".{start_yr_str}{start_mo_str}-{end_yr_str}{end_mo_str}.{s}.{ver}.nc" ) if periodinname is not None: addf = "." + s + "." + ver + ".nc" if outfilename is not None: out = os.path.join(outdir, outfilename) + out_season = out.replace(".nc", addf) - print("output file is", out_season) d_clim_dict[s].to_netcdf( out_season ) # global attributes are automatically saved as well + print("output file:", out_season) + + # Plot climatology if plot and s == "AC": - plot_climatology( - d_ac, - var, - season_to_plot="all", - output_filename=out_season.replace(".nc", ".png"), - ) + # Check if variable is 4D + if is_4d_variable(d_ac, var): + # Plot 3 levels (hPa) for 4D variables for quick check + levels_to_plot = [200, 500, 850] + else: + levels_to_plot = [None] + + # Plot climatology for each level + for level in levels_to_plot: + output_fig_path = out_season.replace(".nc", ".png") + if level is not None: + if var in output_fig_path: + output_fig_path = os.path.join( + outdir, + output_fig_path.split("/")[-1].replace( + var, f"{var}-{level}" + ), + ) + else: + output_fig_path = output_fig_path.replace( + ".png", f"-{level}.png" + ) + + # plot climatology for each level + plot_climatology( + d_ac, + var, + level=level, + season_to_plot="all", + output_filename=output_fig_path, + period=f"{start_yr_str}-{end_yr_str}", + ) + + print("output figure:", output_fig_path) + + +def is_4d_variable(ds, data_var): + da = ds[data_var] + print("data_var, da.shape:", data_var, da.shape) + return len(da.shape) == 4 diff --git a/pcmdi_metrics/mean_climate/lib/plot_clim_maps.py b/pcmdi_metrics/mean_climate/lib/plot_clim_maps.py index c6fb3c72b..b0ce5942c 100644 --- a/pcmdi_metrics/mean_climate/lib/plot_clim_maps.py +++ b/pcmdi_metrics/mean_climate/lib/plot_clim_maps.py @@ -890,6 +890,12 @@ def _load_variable_setting( "colormap": cc.cm.rainbow, "colormap_diff": "jet", }, + 500: { + "levels": np.linspace(-45, 5, 21), + "levels_diff": None, + "colormap": cc.cm.rainbow, + "colormap_diff": "RdBu_r", + }, 850: { "levels": np.arange(-35, 40, 5), "levels_diff": [-15, -10, -5, -2, -1, -0.5, 0, 0.5, 1, 2, 5, 10, 15], @@ -936,6 +942,12 @@ def _load_variable_setting( "colormap": "PiYG_r", "colormap_diff": "RdBu_r", }, + 500: { + "levels": np.arange(-40, 45, 5), + "levels_diff": np.linspace(-20, 20, 21), + "colormap": "PiYG_r", + "colormap_diff": "RdBu_r", + }, 850: { "levels": [ -25, @@ -967,6 +979,12 @@ def _load_variable_setting( "colormap": "PiYG_r", "colormap_diff": "RdBu_r", }, + 500: { + "levels": np.linspace(-10, 10, 11), + "levels_diff": np.linspace(-5, 5, 6), + "colormap": "PiYG_r", + "colormap_diff": "RdBu_r", + }, 850: { "levels": np.linspace(-10, 10, 11), "levels_diff": np.linspace(-5, 5, 6), @@ -984,30 +1002,39 @@ def _load_variable_setting( }, } - # Check if the variable and level exist in the settings - - in_dict = False + # Initialize + levels = None + levels_diff = None + cmap = None + cmap_diff = None + cmap_ext = None + cmap_ext_diff = None + # Check if the variable and level exist in the settings if data_var in var_setting_dict: if level in var_setting_dict[data_var]: settings = var_setting_dict[data_var][level] - levels = settings["levels"] - levels_diff = settings["levels_diff"] - cmap = _get_colormap(settings["colormap"]) - cmap_diff = _get_colormap(settings["colormap_diff"]) + levels = settings.get("levels", None) + levels_diff = settings.get("levels_diff", None) + cmap = _get_colormap(settings.get("colormap", None)) + cmap_diff = _get_colormap(settings.get("colormap_diff", None)) cmap_ext = settings.get("colormap_ext", "both") cmap_ext_diff = "both" - in_dict = True # Use default settings if not found - if not in_dict: - vmin = float(ds[data_var].min()) - vmax = float(ds[data_var].max()) + vmin = float(ds[data_var].min()) + vmax = float(ds[data_var].max()) + if levels is None: levels = np.linspace(vmin, vmax, 21) + if levels_diff is None: levels_diff = np.linspace(vmin / 2.0, vmax / 2.0, 21) + if cmap is None: cmap = plt.get_cmap("jet") + if cmap_diff is None: cmap_diff = plt.get_cmap("RdBu_r") + if cmap_ext is None: cmap_ext = "both" + if cmap_ext_diff is None: cmap_ext_diff = "both" if diff: