Update tropical_subseasonal set for eamxx output (#982)

chengzhuzhang · tomvothecoder · web-flow · commit 436456117022 · 2025-05-27T15:45:49.000-07:00
* update tropical_subseasonal_for_eamxx

* use sub_monthly flag

* support U_at_850hPa with missing value

* add example pre-processing and run script

* fix pre-commit issues

* Remove commented out code in example run script

---------

Co-authored-by: tomvothecoder &lt;tomvothecoder@gmail.com&gt;
diff --git a/e3sm_diags/derivations/derivations.py b/e3sm_diags/derivations/derivations.py
@@ -881,6 +881,9 @@
         ("wind_speed_10m",): rename,  # EAMxx
         ("si10",): rename,
     },
+    "U850": {
+        ("U_at_850hPa",): rename,  # EAMxx
+    },
     "QREFHT": {
         ("QREFHT",): lambda q: convert_units(q, target_units="g/kg"),
         ("qv_2m",): lambda q: convert_units(q, target_units="g/kg"),  # EAMxx
diff --git a/e3sm_diags/driver/tropical_subseasonal_driver.py b/e3sm_diags/driver/tropical_subseasonal_driver.py
@@ -9,16 +9,14 @@
 
 from __future__ import annotations
 
-import glob
-from typing import TYPE_CHECKING  # , Optional
+from typing import TYPE_CHECKING, Tuple
 
 import numpy as np
 import xarray as xr
 
 from e3sm_diags.driver.utils import zwf_functions as wf
 from e3sm_diags.driver.utils.climo_xr import ClimoFreq
 from e3sm_diags.driver.utils.dataset_xr import Dataset
-from e3sm_diags.driver.utils.general import pad_year
 from e3sm_diags.logger import _setup_child_logger
 from e3sm_diags.plot.tropical_subseasonal_plot import plot
 
@@ -47,35 +45,28 @@ def run_diag(parameter: TropicalSubseasonalParameter) -> TropicalSubseasonalPara
     ref_data = Dataset(parameter, data_type="ref")
 
     for variable in parameter.variables:
-        test, test_start, test_end = calculate_spectrum(
-            parameter.test_data_path,
-            variable,
-            parameter.test_start_yr,
-            parameter.test_end_yr,
-        )
+        # Get test dataset
+        test_ds = test_data.get_time_series_dataset(variable, single_point=True)
+        test_spectrum, test_start, test_end = calculate_spectrum(test_ds, variable)
+
+        # Update parameters with actual time range
         parameter.test_start_yr = test_start
         parameter.test_end_yr = test_end
         parameter.test_name_yrs = test_data.get_name_yrs_attr(season)
+
+        # Get reference dataset
         if run_type == "model_vs_model":
-            ref, ref_start, ref_end = calculate_spectrum(
-                parameter.reference_data_path,
-                variable,
-                parameter.ref_start_yr,
-                parameter.ref_end_yr,
-            )
+            ref_ds = ref_data.get_time_series_dataset(variable, single_point=True)
+            ref_spectrum, ref_start, ref_end = calculate_spectrum(ref_ds, variable)
         elif run_type == "model_vs_obs":
             # TODO use pre-calculated spectral power
             # if parameter.ref_start_yr == "":
             #    parameter.ref_name_yrs = parameter.reference_name
             #    # read precalculated data.
             # else:
-            ref_data_path = f"{parameter.reference_data_path}/{parameter.ref_name}"
-            ref, ref_start, ref_end = calculate_spectrum(
-                ref_data_path,
-                variable,
-                parameter.ref_start_yr,
-                parameter.ref_end_yr,
-            )
+            ref_ds = ref_data.get_time_series_dataset(variable, single_point=True)
+            ref_spectrum, ref_start, ref_end = calculate_spectrum(ref_ds, variable)
+
         parameter.ref_start_yr = ref_start
         parameter.ref_end_yr = ref_end
         parameter.ref_name_yrs = ref_data.get_name_yrs_attr(season)
@@ -84,16 +75,24 @@ def run_diag(parameter: TropicalSubseasonalParameter) -> TropicalSubseasonalPara
         for diff_name in ["raw_sym", "raw_asy", "norm_sym", "norm_asy", "background"]:
             diff = (
                 100
-                * (test[f"spec_{diff_name}"] - ref[f"spec_{diff_name}"])
-                / ref[f"spec_{diff_name}"]
+                * (
+                    test_spectrum[f"spec_{diff_name}"]
+                    - ref_spectrum[f"spec_{diff_name}"]
+                )
+                / ref_spectrum[f"spec_{diff_name}"]
             )
             diff.name = f"spec_{diff_name}"
-            diff.attrs.update(test[f"spec_{diff_name}"].attrs)
+            diff.attrs.update(test_spectrum[f"spec_{diff_name}"].attrs)
 
             parameter.spec_type = diff_name
             parameter.output_file = f"{parameter.var_id}_{parameter.spec_type}_15N-15S"
             parameter.diff_title = "percent difference"
-            plot(parameter, test[f"spec_{diff_name}"], ref[f"spec_{diff_name}"], diff)
+            plot(
+                parameter,
+                test_spectrum[f"spec_{diff_name}"],
+                ref_spectrum[f"spec_{diff_name}"],
+                diff,
+            )
 
             if "norm" in diff_name:
                 parameter.spec_type = f"{diff_name}_zoom"
@@ -102,16 +101,33 @@ def run_diag(parameter: TropicalSubseasonalParameter) -> TropicalSubseasonalPara
                 )
                 plot(
                     parameter,
-                    test[f"spec_{diff_name}"],
-                    ref[f"spec_{diff_name}"],
+                    test_spectrum[f"spec_{diff_name}"],
+                    ref_spectrum[f"spec_{diff_name}"],
                     diff,
                     do_zoom=True,
                 )
 
     return parameter
 
 
-def calculate_spectrum(path, variable, start_year, end_year):
+def calculate_spectrum(ds: xr.Dataset, variable: str) -> Tuple[xr.Dataset, str, str]:
+    """Calculate wavenumber-frequency power spectra for a variable.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset containing the variable of interest
+    variable : str
+        Name of the variable to analyze
+
+    Returns
+    -------
+    Tuple[xr.Dataset, str, str]
+        Tuple containing:
+        - Dataset with spectral power components
+        - Start year (as string)
+        - End year (as string)
+    """
     # latitude bounds for analysis
     latBound = (-15, 15)
     # SAMPLES PER DAY
@@ -131,42 +147,13 @@ def calculate_spectrum(path, variable, start_year, end_year):
         "dosymmetries": True,
         "rmvLowFrq": True,
     }
-    # TODO the time subsetting and variable derivation should be replaced during
-    # cdat revamp.
-
-    start_year_str = pad_year(start_year)
-    end_year_str = pad_year(end_year)
-    try:
-        var = xr.open_mfdataset(glob.glob(f"{path}/{variable}_*.nc")).sel(
-            lat=slice(-15, 15),
-            time=slice(f"{start_year_str}-01-01", f"{end_year_str}-12-31"),
-        )[variable]
-        actual_start = var.time.dt.year.values[0]
-        actual_end = var.time.dt.year.values[-1]
-    except OSError:
-        logger.info(
-            f"No files to open for {variable} within {start_year} and {end_year} from {path}."
-        )
-        raise
-    # Unit conversion
-    if var.name == "PRECT":
-        if var.attrs["units"] == "m/s" or var.attrs["units"] == "m s{-1}":
-            logger.info(
-                "\nBEFORE unit conversion: Max/min of data: "
-                + str(var.values.max())
-                + "   "
-                + str(var.values.min())
-            )
-            var.values = (
-                var.values * 1000.0 * 86400.0
-            )  # convert m/s to mm/d, do not alter metadata (yet)
-            var.attrs["units"] = "mm/d"  # adjust metadata to reflect change in units
-            logger.info(
-                "\nAFTER unit conversion: Max/min of data: "
-                + str(var.values.max())
-                + "   "
-                + str(var.values.min())
-            )
+
+    # Extract variable data from dataset and subset to tropical latitudes
+    var = ds[variable].sel(lat=slice(-15, 15))
+
+    # Get actual time range from the data
+    actual_start = var.time.dt.year.values[0]
+    actual_end = var.time.dt.year.values[-1]
 
     # Wavenumber Frequency Analysis
     spec_all = wf_analysis(var, **opt)
@@ -220,6 +207,13 @@ def wf_analysis(x, **kwargs):
     # OPTIONAL kwargs:
     # segsize, noverlap, spd, latitude_bounds (tuple: (south, north)), dosymmetries, rmvLowFrq
 
+    # Interpolate missing values along longitude before spectral analysis
+    if np.any(np.isnan(x)):
+        logger.info(
+            "Interpolating missing values along longitude before spectral analysis"
+        )
+        x = x.interpolate_na(dim="lon", method="linear", fill_value="extrapolate")
+
     z2 = wf.spacetime_power(x, **kwargs)
     z2avg = z2.mean(dim="component")
     z2.loc[{"frequency": 0}] = np.nan  # get rid of spurious power at \nu = 0 (mean)
diff --git a/examples/e3sm_diags_for_eamxx/run_e3sm_diags_eamxx256_climo.py b/examples/e3sm_diags_for_eamxx/run_e3sm_diags_eamxx256_climo.py
@@ -0,0 +1,54 @@
+import os
+from e3sm_diags.parameter.core_parameter import CoreParameter
+from e3sm_diags.parameter.diurnal_cycle_parameter import DiurnalCycleParameter
+from e3sm_diags.parameter.tropical_subseasonal_parameter import (
+    TropicalSubseasonalParameter,
+)
+from e3sm_diags.run import runner
+
+param = CoreParameter()
+
+param.reference_data_path = (
+    "/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/climatology"
+)
+
+test_base_path = "/pscratch/sd/c/chengzhu/ne256pg2_ne256pg2.F20TR-SCREAMv1.rainfrac1.spanc1000.auto2700.acc150.n0128"
+ref_climo = "/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/climatology/"
+ref_ts = "/global/cfs/cdirs/e3sm/diagnostics/observations/Atm/time-series"
+param.test_data_path = f"{test_base_path}/rgr/climo"
+param.test_name = "1ma_ne30pg2.AVERAGE.nmonths_x1"
+param.seasons = ["ANN", "DJF", "MAM", "JJA", "SON"]
+# param.save_netcdf = True
+
+prefix = "/global/cfs/cdirs/e3sm/www/zhang40/tests/eamxx"
+param.results_dir = os.path.join(prefix, "eamxx_ne256_0520_trop")
+params = [param]
+
+trop_param = TropicalSubseasonalParameter()
+trop_param.test_data_path = f"{test_base_path}/rgr/ts_daily"
+trop_param.short_test_name = "Daily_3hi_ne30pg2"
+trop_param.test_start_yr = 1996
+trop_param.test_end_yr = 2001
+
+# Obs
+trop_param.reference_data_path = ref_ts
+trop_param.ref_start_yr = 2001
+trop_param.ref_end_yr = 2010
+
+params.append(trop_param)
+
+dc_param = DiurnalCycleParameter()
+dc_param.test_data_path = f"{test_base_path}/rgr/climo_diurnal_3hrly"
+dc_param.test_name = "3hi_ne30pg2.INSTANT.nhours_x3"
+dc_param.short_test_name = "3hi_ne30pg2"
+# Plotting diurnal cycle amplitude on different scales. Default is True
+dc_param.normalize_test_amp = False
+
+# Obs
+dc_param.reference_data_path = ref_climo
+
+params.append(dc_param)
+
+runner.sets_to_run = ["tropical_subseasonal"]
+
+runner.run_diags(params)
diff --git a/examples/e3sm_diags_for_eamxx/test_eamxx_ne256.sh b/examples/e3sm_diags_for_eamxx/test_eamxx_ne256.sh
@@ -0,0 +1,45 @@
+#!/bin/bash                                 
+
+source /global/common/software/e3sm/anaconda_envs/load_latest_e3sm_unified_pm-cpu.sh
+#
+drc_in=/global/cfs/cdirs/e3sm/beydoun/ne256pg2_ne256pg2.F20TR-SCREAMv1.rainfrac1.spanc1000.auto2700.acc150.n0128/run
+drc_out=/pscratch/sd/c/chengzhu/ne256pg2_ne256pg2.F20TR-SCREAMv1.rainfrac1.spanc1000.auto2700.acc150.n0128
+map_file=/global/cfs/projectdirs/e3sm/zender/maps/map_ne30pg2_to_cmip6_180x360_traave.20231201.nc
+#Monthly averaged: 1ma_ne30pg2.AVERAGE.nmonths_x1.
+#3hourly instantaneous: 3hi_ne30pg2.INSTANT.nhours_x3.
+#1hourly instantaneous: 1hi_ne30pg2.INSTANT.nhours_x1.
+start=1996
+end=2001
+
+echo "Generating Climatology files"
+mkdir -p $drc_out
+drc_rgr=${drc_out}/rgr/climo
+drc=${drc_out}/native/climo
+echo $drc_out
+cd ${drc_in};eval ls 1ma_ne30pg2.AVERAGE.nmonths_x1.*{${start}..${end}}*.nc | ncclimo -P eamxx -p serial --fml_nm=1ma_ne30pg2.AVERAGE.nmonths_x1 --yr_srt=${start} --yr_end=${end} --drc_out=$drc -O $drc_rgr --map=${map_file}
+
+echo "Generating Diurnal Cycle Climo files"
+# Following drc_in included 3hi_ne30pg2.INSTANT.nhours_x3.*nc but with time_bnds
+drc_in=/global/cfs/cdirs/e3sm/chengzhu/ne256pg2_ne256pg2.F20TR-SCREAMv1.rainfrac1.spanc1000.auto2700.acc150.n0128/run
+
+drc_rgr=${drc_out}/rgr/climo_diurnal_3hrly
+drc=${drc_out}/native/climo_diurnal_3hrly
+echo ${drc_in}
+
+cd ${drc_in};eval ls 3hi_ne30pg2.INSTANT.nhours_x3.*{${start}..${end}}*-10800.nc | ncclimo -P eamxx --clm_md=hfc --caseid=3hi_ne30pg2.INSTANT.nhours_x3 -v precip_liq_surf_mass_flux,precip_ice_surf_mass_flux --yr_srt=${start} --yr_end=${end} --drc_out=${drc} -O $drc_rgr --map=${map_file} #--tpd=8
+##
+echo "Generating per-variable time-series from 3hourly instantaneous output."
+drc_rgr=${drc_out}/rgr/ts_3hrly
+drc=${drc_out}/native/ts_3hrly
+
+cd ${drc_in};eval ls 3hi_ne30pg2.INSTANT.nhours_x3.*{${start}..${end}}*-10800.nc | ncclimo -P eamxx --clm_md=hfs --caseid=3hi_ne30pg2.INSTANT.nhours_x3 --var=precip_liq_surf_mass_flux,precip_ice_surf_mass_flux,U_at_850hPa,LW_flux_up_at_model_top --yr_srt=$start --yr_end=$end --drc_out=${drc} -O $drc_rgr --map=${map_file} --split # --tpd=8
+
+echo "Average 3hrly to daily"
+drc_rgr=${drc_out}/rgr/ts_daily
+mkdir -p $drc_rgr
+cd ${drc_out}/rgr/ts_3hrly
+for i in *.nc; do ncra --mro -O -d time,0,,8,8 "$i" "${drc_rgr}/$i";done
+
+exit
+
+
