consider 4-d variables for plot climatology

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

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: