make metrics work on regional subset

Author: chengzhuzhang

auxiliary_tools/debug/968-native-grid-vis/run_lat_lon_native.py

@@ -15,37 +15,37 @@
 # Create parameter objects for 3 different runs
 params = []
 
-# (1) First test configuration
-param1 = LatLonNativeParameter()
-param1.results_dir = f"/lcrc/group/e3sm/public_html/diagnostic_output/{username}/tests/lat_lon_native_test_1"
-param1.test_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
-param1.test_name = "v3.LR.amip_0101"
-param1.short_test_name = "v3.LR.amip_0101"
-param1.reference_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
-param1.ref_name = "v3.HR.test4"
-param1.short_ref_name = "v3.HR.test4"
-param1.seasons = ["DJF"]
-param1.test_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne30pg2.nc"
-param1.ref_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne120pg2.nc"
-param1.case_id = "model_vs_model"
-param1.run_type = "model_vs_model"
-params.append(param1)
-
-# (2) Second test configuration
-param2 = LatLonNativeParameter()
-param2.results_dir = f"/lcrc/group/e3sm/public_html/diagnostic_output/{username}/tests/lat_lon_native_test_2"
-param2.test_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
-param2.test_file = "v3.LR.amip_0101_DJF_climo.nc"
-param2.short_test_name = "v3.LR.amip_0101"
-param2.reference_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
-param2.ref_file = "v3.HR.test4_DJF_climo.nc"
-param2.short_ref_name = "v3.HR.test4"
-param2.seasons = ["DJF"]
-param2.test_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne30pg2.nc"
-param2.ref_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne120pg2.nc"
-param2.case_id = "model_vs_model"
-param2.run_type = "model_vs_model"
-params.append(param2)
+## (1) First test configuration
+#param1 = LatLonNativeParameter()
+#param1.results_dir = f"/lcrc/group/e3sm/public_html/diagnostic_output/{username}/tests/lat_lon_native_test_1"
+#param1.test_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param1.test_name = "v3.LR.amip_0101"
+#param1.short_test_name = "v3.LR.amip_0101"
+#param1.reference_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param1.ref_name = "v3.HR.test4"
+#param1.short_ref_name = "v3.HR.test4"
+#param1.seasons = ["DJF"]
+#param1.test_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne30pg2.nc"
+#param1.ref_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne120pg2.nc"
+#param1.case_id = "model_vs_model"
+#param1.run_type = "model_vs_model"
+#params.append(param1)
+#
+## (2) Second test configuration
+#param2 = LatLonNativeParameter()
+#param2.results_dir = f"/lcrc/group/e3sm/public_html/diagnostic_output/{username}/tests/lat_lon_native_test_2"
+#param2.test_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param2.test_file = "v3.LR.amip_0101_DJF_climo.nc"
+#param2.short_test_name = "v3.LR.amip_0101"
+#param2.reference_data_path = "/lcrc/group/e3sm/public_html/e3sm_diags_test_data/native_grid"
+#param2.ref_file = "v3.HR.test4_DJF_climo.nc"
+#param2.short_ref_name = "v3.HR.test4"
+#param2.seasons = ["DJF"]
+#param2.test_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne30pg2.nc"
+#param2.ref_grid_file = "/lcrc/group/e3sm/diagnostics/grids/ne120pg2.nc"
+#param2.case_id = "model_vs_model"
+#param2.run_type = "model_vs_model"
+#params.append(param2)
 
 # (3) Third test configuration
 param3 = LatLonNativeParameter()

e3sm_diags/driver/lat_lon_native_driver.py

@@ -5,6 +5,7 @@
 
 import uxarray as ux
 
+from e3sm_diags.derivations.default_regions_xr import REGION_SPECS
 from e3sm_diags.driver import METRICS_DEFAULT_VALUE
 from e3sm_diags.driver.utils.dataset_native import NativeDataset
 from e3sm_diags.driver.utils.type_annotations import MetricsDict
@@ -223,14 +224,25 @@ def _run_diags_2d(
                 _compute_diff_between_grids(uxds_test, uxds_ref, var_key)
             )
 
-            # Create metrics dictionary using remapped datasets (following lat_lon_driver pattern)
+            # Apply regional subsetting to all datasets before metrics calculation
+            uxds_test_subset = _apply_regional_subsetting(uxds_test, var_key, region)
+            uxds_ref_subset = _apply_regional_subsetting(uxds_ref, var_key, region)
+            uxds_test_remapped_subset = _apply_regional_subsetting(
+                uxds_test_remapped, var_key, region
+            )
+            uxds_ref_remapped_subset = _apply_regional_subsetting(
+                uxds_ref_remapped, var_key, region
+            )
+            uxds_diff_subset = _apply_regional_subsetting(uxds_diff, var_key, region)
+
+            # Create metrics dictionary using regionally subsetted datasets
             metrics_dict = _create_metrics_dict(
                 var_key,
-                uxds_test,
-                uxds_ref,
-                uxds_test_remapped,
-                uxds_ref_remapped,
-                uxds_diff,
+                uxds_test_subset,
+                uxds_ref_subset,
+                uxds_test_remapped_subset,
+                uxds_ref_remapped_subset,
+                uxds_diff_subset,
             )
 
             # Store metrics in parameter for plot function to access
@@ -240,7 +252,7 @@ def _run_diags_2d(
                 var_key, season, region, ref_name, ilev=None
             )
 
-            # Call plot function directly (pass region parameter)
+            # Call plot function with original datasets for visualization
             plot_func(
                 parameter,
                 var_key,
@@ -252,10 +264,13 @@ def _run_diags_2d(
         else:
             logger.info(f"Processing {var_key} for region {region} (model-only)")
 
-            # Create metrics dictionary for model-only run
+            # Apply regional subsetting to test dataset before metrics calculation
+            uxds_test_subset = _apply_regional_subsetting(uxds_test, var_key, region)
+
+            # Create metrics dictionary for model-only run using regionally subsetted dataset
             metrics_dict = _create_metrics_dict(
                 var_key,
-                uxds_test,
+                uxds_test_subset,
                 None,  # No reference dataset
                 None,  # No remapped test dataset (not needed for model-only)
                 None,  # No remapped reference dataset
@@ -269,7 +284,7 @@ def _run_diags_2d(
                 var_key, season, region, ref_name, ilev=None
             )
 
-            # Call plot function directly (pass region parameter)
+            # Call plot function with original dataset for visualization
             plot_func(
                 parameter,
                 var_key,
@@ -588,8 +603,8 @@ def _create_metrics_dict(
     var_test = uxds_test[var_key]
     metrics_dict: MetricsDict = {
         "test": {
-            "min": var_test.min().item(),
-            "max": var_test.max().item(),
+            "min": [var_test.min().item()],
+            "max": [var_test.max().item()],
             "mean": [var_test.weighted_mean().item()],
             "std": METRICS_DEFAULT_VALUE,  # Not implemented yet for native grids
         },
@@ -603,8 +618,8 @@ def _create_metrics_dict(
     if uxds_ref is not None and var_key in uxds_ref:
         var_ref = uxds_ref[var_key]
         metrics_dict["ref"] = {
-            "min": var_ref.min().item(),
-            "max": var_ref.max().item(),
+            "min": [var_ref.min().item()],
+            "max": [var_ref.max().item()],
             "mean": [var_ref.weighted_mean().item()],
             "std": METRICS_DEFAULT_VALUE,  # Not implemented yet for native grids
         }
@@ -613,8 +628,8 @@ def _create_metrics_dict(
     if uxds_test_remapped is not None and var_key in uxds_test_remapped:
         var_test_remapped = uxds_test_remapped[var_key]
         metrics_dict["test_regrid"] = {
-            "min": var_test_remapped.min().item(),
-            "max": var_test_remapped.max().item(),
+            "min": [var_test_remapped.min().item()],
+            "max": [var_test_remapped.max().item()],
             "mean": [var_test_remapped.weighted_mean().item()],
             "std": METRICS_DEFAULT_VALUE,  # Not implemented yet for native grids
         }
@@ -623,8 +638,8 @@ def _create_metrics_dict(
     if uxds_ref_remapped is not None and var_key in uxds_ref_remapped:
         var_ref_remapped = uxds_ref_remapped[var_key]
         metrics_dict["ref_regrid"] = {
-            "min": var_ref_remapped.min().item(),
-            "max": var_ref_remapped.max().item(),
+            "min": [var_ref_remapped.min().item()],
+            "max": [var_ref_remapped.max().item()],
             "mean": [var_ref_remapped.weighted_mean().item()],
             "std": METRICS_DEFAULT_VALUE,  # Not implemented yet for native grids
         }
@@ -653,8 +668,8 @@ def _create_metrics_dict(
     if uxds_diff is not None and var_key in uxds_diff:
         var_diff = uxds_diff[var_key]
         metrics_dict["diff"] = {
-            "min": var_diff.min().item(),
-            "max": var_diff.max().item(),
+            "min": [var_diff.min().item()],
+            "max": [var_diff.max().item()],
             "mean": [var_diff.weighted_mean().item()],
             "std": METRICS_DEFAULT_VALUE,  # Not implemented yet for native grids
         }
@@ -683,3 +698,72 @@ def _set_default_metric_values(metrics_dict: MetricsDict) -> MetricsDict:
         }
 
     return metrics_dict
+
+
+def _apply_regional_subsetting(
+    uxds: ux.UxDataset | None, var_key: str, region: str
+) -> ux.UxDataset | None:
+    """Apply regional subsetting to a uxarray dataset based on region specification.
+
+    This function follows the same pattern as the regional subsetting in
+    lat_lon_native_plot.py but moves it to the driver for consistency.
+
+    Parameters
+    ----------
+    uxds : ux.UxDataset or None
+        The uxarray dataset to subset.
+    var_key : str
+        The variable key to subset.
+    region : str
+        The region specification (e.g., "global", "CONUS", etc.).
+
+    Returns
+    -------
+    ux.UxDataset or None
+        The regionally subsetted dataset, or None if input was None.
+    """
+    if uxds is None:
+        return uxds
+
+    # Get region specs (same logic as in plot function)
+    region_specs = REGION_SPECS.get(region, None)
+
+    if region_specs is None:
+        # Unknown region, return original dataset
+        logger.warning(
+            f"Region '{region}' not found in REGION_SPECS. Using global dataset."
+        )
+        return uxds
+
+    # Get bounds (same logic as in plot function)
+    lat_bounds = region_specs.get("lat", (-90, 90))  # type: ignore
+    lon_bounds = region_specs.get("lon", (0, 360))  # type: ignore
+    is_global_domain = lat_bounds == (-90, 90) and lon_bounds == (0, 360)
+
+    if is_global_domain:
+        # Global domain, no subsetting needed
+        return uxds
+
+    try:
+        # Check if target variable exists
+        if var_key not in uxds.data_vars:
+            logger.warning(
+                f"Variable '{var_key}' not found in dataset. Available vars: {list(uxds.data_vars)}"
+            )
+            return uxds
+
+        # Apply subsetting to the specific variable
+        var_subset = uxds[var_key].subset.bounding_box(lon_bounds, lat_bounds)
+
+        # Create new dataset from subsetted variable
+        uxds_subset = var_subset.to_dataset()
+        uxds_subset.attrs.update(uxds.attrs)
+        uxds_subset[var_key].attrs.update(uxds[var_key].attrs)
+        return uxds_subset
+
+    except Exception as e:
+        logger.warning(
+            f"Failed to apply regional subsetting for region '{region}': {e}"
+        )
+        logger.warning("Using global dataset instead.")
+        return uxds

e3sm_diags/plot/lat_lon_native_plot.py

@@ -109,53 +109,28 @@ def plot(  # noqa: C901
 
     logger.info(f"Region: {region}, lat_bounds: {lat_bounds}, lon_bounds: {lon_bounds}")
 
-    # Extract metrics directly from the uxarray dataset
+    # Extract metrics from parameter.metrics_dict (calculated in driver with regional subsetting)
     if uxds_test is not None and var_key in uxds_test:
-        # ------------------------------------------------------------
-        # FIXME: Metrics extraction for test, ref, and diff datasets are duplicated,
-        # extract to a helper function.
-        # ------------------------------------------------------------
-        if is_global_domain:
-            test_min = uxds_test[var_key].min().item()
-            test_max = uxds_test[var_key].max().item()
-            # For native grid, use weighted mean
-            test_mean = uxds_test[var_key].weighted_mean().item()
-        else:
-            test_subset = uxds_test[var_key].subset.bounding_box(
-                lon_bounds,
-                lat_bounds,
-            )
-            test_min = test_subset.min().item()
-            test_max = test_subset.max().item()
-            # For native grid, use weighted mean
-            test_mean = test_subset.weighted_mean().item()
-
         units = uxds_test[var_key].attrs.get("units", "")
-        # ------------------------------------------------------------
+
+        # Get test metrics from parameter.metrics_dict
+        try:
+            test_min = parameter.metrics_dict["test_regrid"]["min"][0]  # type: ignore
+            test_max = parameter.metrics_dict["test_regrid"]["max"][0]  # type: ignore
+            test_mean = parameter.metrics_dict["test_regrid"]["mean"][0]  # type: ignore
+        except (KeyError, IndexError, TypeError) as e:
+            logger.warning(
+                f"Failed to get test metrics from metrics_dict: {e}, using NaN"
+            )
+            test_min = test_max = test_mean = float("nan")
     else:
         # This should not happen since we check earlier, but just in case
         logger.error(f"Missing test data for variable {var_key} in native grid dataset")
-
         return
 
     # Extract metrics for reference data if available
     ref_min = ref_max = ref_mean = diff_min = diff_max = diff_mean = None
     if has_reference and uxds_ref is not None:
-        # ------------------------------------------------------------
-        # FIXME: Metrics extraction for test, ref, and diff datasets are duplicated,
-        # extract to a helper function.
-        # ------------------------------------------------------------
-        if is_global_domain:
-            ref_min = uxds_ref[var_key].min().item()
-            ref_max = uxds_ref[var_key].max().item()
-            ref_mean = uxds_ref[var_key].weighted_mean().item()
-        else:
-            ref_subset = uxds_ref[var_key].subset.bounding_box(lon_bounds, lat_bounds)
-            ref_min = ref_subset.min().item()
-            ref_max = ref_subset.max().item()
-            ref_mean = ref_subset.weighted_mean().item()
-        # ------------------------------------------------------------
-
         ref_units = uxds_ref[var_key].attrs.get("units", "")
 
         # Check if units match between test and reference
@@ -164,25 +139,28 @@ def plot(  # noqa: C901
                 f"Units mismatch between test ({units}) and reference ({ref_units})"
             )
 
-        # Calculate approximate metrics for difference if not already available
+        # Get reference metrics from parameter.metrics_dict
+        try:
+            ref_min = parameter.metrics_dict["ref"]["min"][0]  # type: ignore
+            ref_max = parameter.metrics_dict["ref"]["max"][0]  # type: ignore
+            ref_mean = parameter.metrics_dict["ref"]["mean"][0]  # type: ignore
+        except (KeyError, IndexError, TypeError):
+            logger.warning(
+                "Failed to get reference metrics from metrics_dict, using NaN"
+            )
+            ref_min = ref_max = ref_mean = float("nan")
+
+        # Get difference metrics from parameter.metrics_dict
         if uxds_diff is not None and var_key in uxds_diff:
-            # ------------------------------------------------------------
-            # FIXME: Metrics extraction for test, ref, and diff datasets are duplicated,
-            # extract to a helper function.
-            # ------------------------------------------------------------
-            if is_global_domain:
-                diff_min = uxds_diff[var_key].min().item()
-                diff_max = uxds_diff[var_key].max().item()
-                diff_mean = uxds_diff[var_key].weighted_mean().item()
-            else:
-                diff_subset = uxds_diff[var_key].subset.bounding_box(
-                    lon_bounds,
-                    lat_bounds,
+            try:
+                diff_min = parameter.metrics_dict["diff"]["min"][0]  # type: ignore
+                diff_max = parameter.metrics_dict["diff"]["max"][0]  # type: ignore
+                diff_mean = parameter.metrics_dict["diff"]["mean"][0]  # type: ignore
+            except (KeyError, IndexError, TypeError):
+                logger.warning(
+                    "Failed to get diff metrics from metrics_dict, using NaN"
                 )
-                diff_min = diff_subset.min().item()
-                diff_max = diff_subset.max().item()
-                diff_mean = diff_subset.weighted_mean().item()
-            # ------------------------------------------------------------
+                diff_min = diff_max = diff_mean = float("nan")
 
     # Create panels following the lat_lon_plot layout
     # Panel 1: Test data (always created)