Add qbo wavelet plot support

e3sm_diags/plot/qbo_plot.py

@@ -16,10 +16,11 @@
 logger = custom_logger(__name__)
 
 PANEL_CFG = [
-    (0.075, 0.70, 0.6, 0.225),
-    (0.075, 0.425, 0.6, 0.225),
-    (0.725, 0.425, 0.2, 0.5),
-    (0.075, 0.075, 0.85, 0.275),
+    (0.075, 0.75, 0.6, 0.175),
+    (0.075, 0.525, 0.6, 0.175),
+    (0.725, 0.525, 0.2, 0.4),
+    (0.075, 0.285, 0.85, 0.175),
+    (0.075, 0.04, 0.85, 0.175),
 ]
 
 LABEL_SIZE = 14
@@ -51,7 +52,7 @@ class ZAxis(TypedDict):
 
 
 def plot(parameter: QboParameter, test_dict, ref_dict):
-    fig = plt.figure(figsize=(14, 14))
+    fig = plt.figure(figsize=(14, 18))
 
     test_z_axis = xc.get_dim_coords(test_dict["qbo"], axis="Z")
     ref_z_axis = xc.get_dim_coords(ref_dict["qbo"], axis="Z")
@@ -154,7 +155,7 @@ def plot(parameter: QboParameter, test_dict, ref_dict):
 
     # Panel 3 (Bottom)
     x = dict(
-        axis_range=[0, 50],
+        axis_range=[5, 50],
         axis_scale="linear",
         label="Period (months)",
         data=test_dict["period_new"],
@@ -174,6 +175,28 @@ def plot(parameter: QboParameter, test_dict, ref_dict):
     title = "QBO Spectral Density (Eq. 18-22 hPa zonal winds)"
     _add_color_map(3, fig, "line", title, x, y)
 
+    # Panel 4 (Bottom/Bottom)
+    x = dict(
+        axis_range=[5, 50],
+        axis_scale="linear",
+        data=test_dict["wave_period"],
+        data_label=test_dict["name"],
+        data2=ref_dict["wave_period"],
+        data2_label=ref_dict["name"],
+        label="Period (months)",
+    )
+    y = dict(
+        axis_range=[-1, 105],
+        axis_scale="linear",
+        data=test_dict["wavelet"],
+        data_label=None,
+        data2=ref_dict["wavelet"],
+        data2_label=None,
+        label="Variance (" + "m\u00b2/s\u00b2" + ")",
+    )
+    title = "QBO Wavelet (Eq. 18-22 hPa zonal winds)"
+    _add_color_map(4, fig, "line", title, x, y)
+
     plt.tight_layout()
 
     # Figure title
@@ -225,6 +248,29 @@ def _add_color_map(
             fontsize=LABEL_SIZE,
         )
 
+    if subplot_num == 3 or subplot_num == 4:
+        # Find the index of the wavelet maximum value
+        test_ymax_idx = list(y["data"]).index(max(y["data"]))
+        ref_ymax_idx = list(y["data2"]).index(max(y["data2"]))  # type: ignore
+
+        # Use the index to get the period value for peak of spectra
+        test_y_max_xval = list(x["data"])[test_ymax_idx]
+        ref_y_max_xval = list(x["data2"])[ref_ymax_idx]  # type: ignore
+
+        # Plot vertical lines for period peaks
+        ax.axvline(
+            x=test_y_max_xval,
+            ymax=max(y["data"]) / y["axis_range"][1],
+            color="k",
+            linestyle="-",
+        )
+        ax.axvline(
+            x=ref_y_max_xval,
+            ymax=max(y["data2"]) / y["axis_range"][1],  # type: ignore
+            color="r",
+            linestyle="--",
+        )
+
     ax.set_title(title, size=LABEL_SIZE, weight="demi")
     ax.set_xlabel(x["label"], size=LABEL_SIZE)
     ax.set_ylabel(y["label"], size=LABEL_SIZE)