Update REF ENSO recipes for CMIP7 (#4470)

Co-authored-by: Romain Beucher <romain.beucher@anu.edu.au>

Author: flicj191

esmvaltool/diag_scripts/enso_metrics/climatology_diagnostic1.py

@@ -47,7 +47,7 @@ def plot_level1(input_data, cfg, provenance):
             cube,
         )
 
-        if dataset["project"] == "CMIP6":
+        if dataset["project"].startswith("CMIP"):
             qplt.plot(cube, label=dataset["dataset"])
             model_data = cube.data
         else:
@@ -201,15 +201,16 @@ def main(cfg):
     )
     # for each select obs and iterate others, obs last
     for grp, var_attr in variable_groups.items():
-        logger.info("%s : %d, %s", grp, len(var_attr), pformat(var_attr))
+        datasets = [attr["dataset"] for attr in var_attr]
+        logger.info("%s : %d, %s", grp, len(var_attr), datasets)
         pairs = [var_attr[-1]]  # obs to list
-        prov = provenance_record(grp, list(cfg["input_data"].keys()))
+        prov = provenance_record(grp, [attr["filename"] for attr in var_attr])
         for metadata in var_attr:
             logger.info("iterate though datasets\n %s", pformat(metadata))
-            if metadata["project"] == "CMIP6":
+            if metadata["project"].startswith("CMIP"):
                 pairs.append(metadata)
                 fig, filename, rmse = plot_level1(pairs, cfg, prov)
-
+                logger.info("metric %s", rmse)
                 save_figure(
                     "_".join(filename),
                     prov,

esmvaltool/diag_scripts/enso_metrics/climatology_diagnosticlevel2.py

@@ -26,6 +26,7 @@ def plotmaps_level2(input_data, grp):
     """Create map plots for pair of input data."""
     fig = plt.figure(figsize=(18, 6))
     proj = ccrs.Orthographic(central_longitude=210.0)
+    cubes = []
     for plt_pos, dataset in enumerate(input_data, start=121):
         logger.info(
             "dataset: %s - %s",
@@ -34,6 +35,9 @@ def plotmaps_level2(input_data, grp):
         )
         cube, cbar_label = load_seacycle_stdev(dataset)
 
+        # return cubes to save
+        cubes.append(cube)
+
         ax1 = plt.subplot(plt_pos, projection=proj)
         ax1.add_feature(cfeature.LAND, facecolor="gray")
         ax1.coastlines()
@@ -61,7 +65,7 @@ def plotmaps_level2(input_data, grp):
     cbar = fig.colorbar(cf1, cax=cax, orientation="horizontal", extend="both")
     cbar.set_label(cbar_label)
 
-    return fig
+    return fig, cubes
 
 
 def load_seacycle_stdev(dataset):
@@ -123,16 +127,15 @@ def provenance_record(var_grp, ancestor_files):
     return record
 
 
-def save_plotdata(plotdata, group, pairs, cfg):
-    """Save both obs and model plotted data."""
-    for i, cube in enumerate(plotdata):
-        data_prov = provenance_record(group, [pairs[i]["filename"]])
-        datafile = [
-            pairs[i]["dataset"],
-            pairs[i]["short_name"],
-            pairs[i]["preprocessor"],
-        ]
-        save_data("_".join(datafile), data_prov, cfg, cube)
+def save_plotdata(plotdata, group, pairs, cfg, i):
+    """Save both obs and model plotted data, i: pairs index."""
+    data_prov = provenance_record(group, [pairs[i]["filename"]])
+    datafile = [
+        pairs[i]["dataset"],
+        pairs[i]["short_name"],
+        pairs[i]["preprocessor"],
+    ]
+    save_data("_".join(datafile), data_prov, cfg, plotdata[i])
 
 
 def main(cfg):
@@ -153,9 +156,10 @@ def main(cfg):
             # create pairs, add obs first to list
             pairs = [var_attr[-1]]
             logger.info("iterate though datasets\n %s", pformat(metadata))
-            if metadata["project"] == "CMIP6":
+            if metadata["project"].startswith("CMIP"):
                 pairs.append(metadata)
-                fig = plotmaps_level2(pairs, grp)
+                fig, data_cubes = plotmaps_level2(pairs, grp)
+                save_plotdata(data_cubes, grp, pairs, cfg, 1)
                 filename = "_".join(
                     [
                         metadata["dataset"],
@@ -170,6 +174,8 @@ def main(cfg):
                     figure=fig,
                     dpi=300,
                 )
+        # save obs data for each group at end so not repeated for each model
+        save_plotdata(data_cubes, grp, pairs, cfg, 0)
 
 
 if __name__ == "__main__":

esmvaltool/diag_scripts/enso_metrics/climatology_diagnosticlevel3.py

@@ -48,6 +48,7 @@ def plotmaps_level3(input_data, itcz=False):
         "ts": np.arange(20, 31, 1),
         "tauu": np.arange(-80, 90, 10),
     }
+    cubes = []
     for plt_pos, dataset in enumerate(input_data, start=121):
         logger.info(
             "dataset: %s - %s",
@@ -56,6 +57,8 @@ def plotmaps_level3(input_data, itcz=False):
         )
         cube, cbar_label, x_label = load_seacycle_stat(dataset, itcz)
 
+        # return cubes to save
+        cubes.append(cube)
         ax1 = plt.subplot(plt_pos)
         cf1 = iplt.contourf(
             cube,
@@ -78,7 +81,7 @@ def plotmaps_level3(input_data, itcz=False):
     cbar = fig.colorbar(cf1, cax=cax, orientation="horizontal", extend="both")
     cbar.set_label(cbar_label)
 
-    return fig
+    return fig, cubes
 
 
 def load_seacycle_stat(dataset, itcz=False):
@@ -169,16 +172,11 @@ def provenance_record(var_grp, ancestor_files):
     return record
 
 
-def save_plotdata(plotdata, group, pairs, cfg):
-    """Save both obs and model plotted data."""
-    for i, cube in enumerate(plotdata):
-        data_prov = provenance_record(group, [pairs[i]["filename"]])
-        datafile = [
-            pairs[i]["dataset"],
-            pairs[i]["short_name"],
-            pairs[i]["preprocessor"],
-        ]
-        save_data("_".join(datafile), data_prov, cfg, cube)
+def save_plotdata(plotdata, group, pairs, cfg, i):
+    """Save both obs and model plotted data, i: pairs index."""
+    data_prov = provenance_record(group, [pairs[i]["filename"]])
+    datafile = [pairs[i]["dataset"], pairs[i]["short_name"], group, "level3"]
+    save_data("_".join(datafile), data_prov, cfg, plotdata[i])
 
 
 def main(cfg):
@@ -193,16 +191,17 @@ def main(cfg):
     )
     # for each select obs and iterate others, obs last
     for grp, var_attr in variable_groups.items():
-        logger.info("%s : %d, %s", grp, len(var_attr), pformat(var_attr))
+        datasets = [attr["dataset"] for attr in var_attr]
+        logger.info("%s : %d, %s", grp, len(var_attr), datasets)
         if grp in ["pr_seacycle", "sst_seacycle", "tauu_seacycle"]:
             for metadata in var_attr:
                 # create pairs, add obs first to list
                 pairs = [var_attr[-1]]
                 logger.info("iterate though datasets\n %s", pformat(metadata))
-                if metadata["project"] == "CMIP6":
+                if metadata["project"].startswith("CMIP"):
                     pairs.append(metadata)
-                    fig = plotmaps_level3(pairs, itcz=False)
-                    # save_plotdata(data_cubes, grp, pairs, cfg)
+                    fig, data_cubes = plotmaps_level3(pairs, itcz=False)
+                    save_plotdata(data_cubes, grp, pairs, cfg, 1)
                     filename = "_".join(
                         [
                             metadata["dataset"],
@@ -224,7 +223,8 @@ def main(cfg):
                     if grp == "pr_seacycle":
                         # replace pr with doubleITCZ
                         grp_itcz = "doubleITCZ_seacycle"
-                        fig = plotmaps_level3(pairs, itcz=True)
+                        fig, itcz_cubes = plotmaps_level3(pairs, itcz=True)
+                        save_plotdata(itcz_cubes, grp_itcz, pairs, cfg, 1)
                         filename = "_".join(
                             [
                                 metadata["dataset"],
@@ -243,6 +243,12 @@ def main(cfg):
                             figure=fig,
                             dpi=300,
                         )
+                        obs_itcz = itcz_cubes[0]
+            # save obs data for each group at end so not repeated for each model
+            save_plotdata(data_cubes, grp, pairs, cfg, 0)
+            if obs_itcz:
+                save_plotdata(itcz_cubes, grp_itcz, pairs, cfg, 0)
+                obs_itcz = None
 
 
 if __name__ == "__main__":

esmvaltool/diag_scripts/enso_metrics/enso_diag1metrics.py

@@ -22,6 +22,7 @@
     get_diagnostic_filename,
     group_metadata,
     run_diagnostic,
+    save_data,
     save_figure,
     select_metadata,
 )
@@ -135,8 +136,9 @@ def sst_regressed(n34_cube):
 
         # Select the data for the current year and append it to n34_selected
         year_enso = iris.Constraint(
-            time=lambda cell, enso_epoch=enso_epoch: cell.point.year
-            in enso_epoch,
+            time=lambda cell, enso_epoch=enso_epoch: (
+                cell.point.year in enso_epoch
+            ),
         )
         cube_2 = n34_cube.extract(year_enso)
         n34_selected.append(cube_2.data.data)
@@ -234,7 +236,7 @@ def compute_enso_metrics(input_pair, dt_ls, var_group, metric):
             "ENSO lifecycle",
             dt_ls,
         )
-
+        # save - require months
         data_to_save.append(data_to_cube(obs, months, metric))
         data_to_save.append(data_to_cube(model, months, metric))
 
@@ -275,8 +277,6 @@ def compute_enso_metrics(input_pair, dt_ls, var_group, metric):
             "ENSO seasonality",
             dt_ls,
         )
-        data_to_save.append(data_to_cube(data_values[0], None, metric))
-        data_to_save.append(data_to_cube(data_values[1], None, metric))
 
     elif metric == "13asymmetry":
         model_skew = skew(input_pair[1][var_group[0]].data, axis=0)
@@ -291,8 +291,6 @@ def compute_enso_metrics(input_pair, dt_ls, var_group, metric):
             "ENSO skewness",
             dt_ls,
         )
-        data_to_save.append(data_to_cube(obs_skew, None, metric))
-        data_to_save.append(data_to_cube(model_skew, None, metric))
 
     elif metric == "14duration":
         model = sst_regressed(input_pair[1][var_group[0]])
@@ -313,15 +311,14 @@ def compute_enso_metrics(input_pair, dt_ls, var_group, metric):
             "ENSO duration",
             dt_ls,
         )
-        data_to_save.append(data_to_cube(data_values[0], None, metric))
-        data_to_save.append(data_to_cube(data_values[1], None, metric))
+
     elif metric == "15diversity":
         for datas in input_pair:  # obs 0, mod 1
             events = enso_events(datas[var_group[0]])
             results_lon = diversity(datas[var_group[1]], events)
             results_lon["enso"] = results_lon["nino"] + results_lon["nina"]
             data_values.append(iqr(results_lon["enso"]))
-        # compute(obs, mod)
+
         val = compute(data_values[0], data_values[1])
         fig = plot_level1(
             data_values,
@@ -330,6 +327,8 @@ def compute_enso_metrics(input_pair, dt_ls, var_group, metric):
             "ENSO diversity",
             dt_ls,
         )
+
+    if metric not in ["09pattern", "10lifecycle"]:
         data_to_save.append(data_to_cube(data_values[0], None, metric))
         data_to_save.append(data_to_cube(data_values[1], None, metric))
 
@@ -458,10 +457,27 @@ def group_obs_models(obs, models, metric, var_preproc, cfg):
                 figure=output[1],
                 dpi=300,
             )
+
+            # save data_cubes output[2]
+            save_plotdata(output[2], metric, [obs, attributes], cfg)
+
         # clear value,fig
         output = None
 
 
+def save_plotdata(plotdata, metric, pairs, cfg):
+    """Save both obs and model plotted data."""
+    for i, cube in enumerate(plotdata):
+        files = [attr["filename"] for attr in pairs[i]]
+        data_prov = get_provenance_record(metric, files)
+        datafile = [
+            pairs[i][0]["dataset"],
+            pairs[i][0]["short_name"],
+            metric,
+        ]
+        save_data("_".join(datafile), data_prov, cfg, cube)
+
+
 def get_provenance_record(metric, ancestor_files):
     """Create a provenance record describing the diagnostic data and plot."""
     caption = {
@@ -549,11 +565,21 @@ def main(cfg):
                 variable_group=var_prep,
                 project="OBS6",
             )
+            obs += select_metadata(
+                input_data,
+                variable_group=var_prep,
+                project="obs4MIPs",
+            )
             models += select_metadata(
                 input_data,
                 variable_group=var_prep,
                 project="CMIP6",
             )
+            models += select_metadata(
+                input_data,
+                variable_group=var_prep,
+                project="CMIP7",
+            )
 
         group_obs_models(obs, models, metric, var_preproc, cfg)