CDAT Migration Phase 2: Refactor diurnal_cycle set

auxiliary_tools/cdat_regression_testing/666-diurnal-cycle/run.cfg

@@ -0,0 +1,10 @@
+[#]
+sets = ["diurnal_cycle"]
+case_id = "TRMM-3B43v-7_3hr"
+variables = ["PRECT"]
+ref_name = "TRMM-3B43v-7_3hr"
+# seasons = ["ANN", "DJF", "MAM", "JJA", "SON"]
+seasons = ["ANN"]
+# regions = ["CONUS", "20S20N", "W_Pacific", "Amazon","global","50S50N"]
+regions = ["20S20N"]
+reference_name = "TRMM-3B43v-7"

auxiliary_tools/cdat_regression_testing/666-diurnal-cycle/run_script.py

@@ -0,0 +1,12 @@
+# %%
+# python -m auxiliary_tools.cdat_regression_testing.666-diurnal-cycle.run_script
+from auxiliary_tools.cdat_regression_testing.base_run_script import run_set
+
+SET_NAME = "diurnal_cycle"
+SET_DIR = "666-diurnal-cycle"
+CFG_PATH: str | None = None
+# CFG_PATH: str | None = "/global/u2/v/vo13/E3SM-Project/e3sm_diags/auxiliary_tools/cdat_regression_testing/666-diurnal-cycle/run.cfg"
+MULTIPROCESSING = True
+
+# %%
+run_set(SET_NAME, SET_DIR, CFG_PATH, MULTIPROCESSING)

auxiliary_tools/cdat_regression_testing/utils.py

@@ -1,4 +1,5 @@
 import math
+import os
 from typing import List
 
 import pandas as pd
@@ -163,7 +164,7 @@ def get_num_metrics_above_diff_thres(
     )
 
 
-def get_image_diffs(actual_path: str, expected_path: str):
+def get_image_diffs(actual_path: str, expected_path: str) -> str | None:
     """Get the diffs between two images.
 
     This function is useful for comparing two datasets that can't be compared
@@ -183,10 +184,28 @@ def get_image_diffs(actual_path: str, expected_path: str):
     expected_png = Image.open(expected_path).convert("RGB")
 
     diff = ImageChops.difference(actual_png, expected_png)
-
     draw = ImageDraw.Draw(diff)
-    (left, upper, right, lower) = diff.getbbox()
-    draw.rectangle(((left, upper), (right, lower)), outline="red")
 
-    diff_path = actual_path.replace("actual", "diff")
+    try:
+        (left, upper, right, lower) = diff.getbbox()
+    except TypeError as e:
+        if "cannot unpack non-iterable NoneType object" in str(e):
+            print("     * Plots are identical")
+
+        return None
+    else:
+        draw.rectangle(((left, upper), (right, lower)), outline="red")
+
+    # Create the diff directory.
+    split_actual_path = actual_path.split("/")
+    split_actual_path[-2] = f"{split_actual_path[-2]}_diff"
+    actual_dir_path = ("/").join(split_actual_path[0:-1])
+    os.makedirs(actual_dir_path, exist_ok=True)
+
+    # Save the png file to the diff directory.
+    diff_path = ("/").join(split_actual_path)
     diff.save(diff_path)
+
+    print(f"     * Difference path {diff_path}")
+
+    return diff_path

e3sm_diags/driver/diurnal_cycle_driver.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
-import os
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
-import cdms2
+import xarray as xr
 
-import e3sm_diags
-from e3sm_diags.driver import utils
+from e3sm_diags.driver.utils.dataset_xr import Dataset
+from e3sm_diags.driver.utils.diurnal_cycle_xr import composite_diurnal_cycle
+from e3sm_diags.driver.utils.io import _get_output_filename_filepath
+from e3sm_diags.driver.utils.regrid import _apply_land_sea_mask, _subset_on_region
 from e3sm_diags.logger import custom_logger
-from e3sm_diags.plot import plot
+from e3sm_diags.plot.diurnal_cycle_plot import plot as plot_func
 
 logger = custom_logger(__name__)
 
@@ -22,94 +23,109 @@ def run_diag(parameter: DiurnalCycleParameter) -> DiurnalCycleParameter:
     ref_name = getattr(parameter, "ref_name", "")
     regions = parameter.regions
 
-    test_data = utils.dataset.Dataset(parameter, test=True)
-    ref_data = utils.dataset.Dataset(parameter, ref=True)
-
-    for season in seasons:
-        # Get the name of the data, appended with the years averaged.
-        parameter.test_name_yrs = utils.general.get_name_and_yrs(
-            parameter, test_data, season
-        )
-        parameter.ref_name_yrs = utils.general.get_name_and_yrs(
-            parameter, ref_data, season
-        )
-
-        # Get land/ocean fraction for masking.
-        try:
-            land_frac = test_data.get_climo_variable("LANDFRAC", season)
-            ocean_frac = test_data.get_climo_variable("OCNFRAC", season)
-        except Exception:
-            mask_path = os.path.join(
-                e3sm_diags.INSTALL_PATH, "acme_ne30_ocean_land_mask.nc"
-            )
-            with cdms2.open(mask_path) as f:
-                land_frac = f("LANDFRAC")
-                ocean_frac = f("OCNFRAC")
-
-        for var in variables:
-            logger.info("Variable: {}".format(var))
-            test = test_data.get_climo_variable(var, season)
-            ref = ref_data.get_climo_variable(var, season)
-
-            parameter.var_id = var
-            parameter.viewer_descr[var] = (
-                test.long_name
-                if hasattr(test, "long_name")
-                else "No long_name attr in test data."
-            )
+    test_ds = Dataset(parameter, data_type="test")
+    ref_ds = Dataset(parameter, data_type="ref")
+
+    for var_key in variables:
+        logger.info("Variable: {}".format(var_key))
+        parameter.var_id = var_key
+
+        for season in seasons:
+            parameter._set_name_yrs_attrs(test_ds, ref_ds, season)
+
+            ds_land_sea_mask: xr.Dataset = test_ds._get_land_sea_mask(season)
+
+            ds_test = test_ds.get_climo_dataset(var_key, season)
+            ds_ref = ref_ds.get_climo_dataset(var_key, season)
 
             for region in regions:
-                test_domain = utils.general.select_region(
-                    region, test, land_frac, ocean_frac, parameter
-                )
-                ref_domain = utils.general.select_region(
-                    region, ref, land_frac, ocean_frac, parameter
+                if "land" in region or "ocean" in region:
+                    test_domain = _apply_land_sea_mask(
+                        ds_test,
+                        ds_land_sea_mask,
+                        var_key,
+                        region,  # type: ignore
+                        parameter.regrid_tool,
+                        parameter.regrid_method,
+                    )
+
+                    ref_domain = _apply_land_sea_mask(
+                        ds_ref,
+                        ds_land_sea_mask,
+                        var_key,
+                        region,  # type: ignore
+                        parameter.regrid_tool,
+                        parameter.regrid_method,
+                    )
+                else:
+                    test_domain = ds_test.copy()
+                    ref_domain = ds_ref.copy()
+
+                test_domain = _subset_on_region(test_domain, var_key, region)
+                ref_domain = _subset_on_region(ref_domain, var_key, region)
+
+                parameter.viewer_descr[var_key] = ds_test[var_key].attrs.get(
+                    "long_name", "No long_name attr in test data."
                 )
-
-                parameter.output_file = "-".join([ref_name, var, season, region])
+                parameter.output_file = "-".join([ref_name, var_key, season, region])
                 parameter.main_title = str(
-                    " ".join([var, "Diurnal Cycle ", season, region])
+                    " ".join([var_key, "Diurnal Cycle ", season, region])
                 )
 
                 (
                     test_cmean,
                     test_amplitude,
                     test_maxtime,
-                ) = utils.diurnal_cycle.composite_diurnal_cycle(test_domain, season)
+                ) = composite_diurnal_cycle(
+                    test_domain, var_key, season
+                )  # type: ignore
                 (
                     ref_cmean,
                     ref_amplitude,
                     ref_maxtime,
-                ) = utils.diurnal_cycle.composite_diurnal_cycle(ref_domain, season)
+                ) = composite_diurnal_cycle(
+                    ref_domain, var_key, season
+                )  # type: ignore
+
                 parameter.var_region = region
-                plot(
-                    parameter.current_set,
+
+                plot_func(
                     test_maxtime,
                     test_amplitude,
                     ref_maxtime,
                     ref_amplitude,
                     parameter,
                 )
-                utils.general.save_ncfiles(
-                    parameter.current_set,
-                    test_cmean,
-                    ref_cmean,
-                    None,
-                    parameter,
-                )
-                utils.general.save_ncfiles(
-                    parameter.current_set,
-                    test_amplitude,
-                    ref_amplitude,
-                    None,
-                    parameter,
+
+                ds_out_test = xr.Dataset(
+                    data_vars={
+                        test_cmean.name: test_cmean,
+                        test_amplitude.name: test_amplitude,
+                        test_maxtime.name: test_maxtime,
+                    }
                 )
-                utils.general.save_ncfiles(
-                    parameter.current_set,
-                    test_maxtime,
-                    ref_maxtime,
-                    None,
-                    parameter,
+                ds_out_ref = xr.Dataset(
+                    data_vars={
+                        ref_cmean.name: ref_cmean,
+                        ref_amplitude.name: ref_amplitude,
+                        ref_maxtime.name: ref_maxtime,
+                    }
                 )
 
+                _write_vars_to_netcdf(parameter, var_key, ds_out_test, "test")
+                _write_vars_to_netcdf(parameter, var_key, ds_out_ref, "ref")
+
     return parameter
+
+
+def _write_vars_to_netcdf(
+    parameter: DiurnalCycleParameter,
+    var_key: str,
+    ds: xr.Dataset,
+    data_type: Literal["test", "ref"],
+):
+    _, filepath = _get_output_filename_filepath(parameter, data_type)
+
+    ds.to_netcdf(filepath)
+
+    logger.info(f"'{var_key}' {data_type} variable output saved in: {filepath}")