JJB tropical subseasonal diags - Wheeler and Kiladis

.pre-commit-config.yaml

@@ -1,4 +1,4 @@
-exclude: "docs|node_modules|migrations|.git|.tox|examples|analysis_data_preprocess|auxiliary_tools|conda/meta.yaml"
+exclude: "docs|node_modules|migrations|.git|.tox|examples|analysis_data_preprocess|auxiliary_tools|conda/meta.yaml|e3sm_diags/driver/utils/zwf_functions.py"
 default_stages: [commit]
 fail_fast: true
 

README.md

@@ -50,6 +50,8 @@ E3SM Diags is modeled after the National Center for Atmospheric Research (NCAR)
 ### New Feature added during v2 development
 | Feature name <br />(set name) | Brief Introduction                                                                                                                                                                                     | Developers Contributors*                                                                                                                    | Released version |
 |--- | ---------------- | --- | --- |
+| Wavenumber frequency analysis (tropical_subseasonal) 		| The wavenumber frequency analysis (Wheeler-Kiladis Diagram) for Tropical subseasonal analysis 													 | Jim Benedict, Brian Medeiros, Jill Zhang, Tom Vo     | 2.11.0
+
 | T5050 diagnostic /Mixed phase partition (mp_partition)        | Temperature at which cloud top is 50% ice, 50% liquid, following McCoy et al. (2015).                                                                                                                  | Yuying Zhang, Jill Zhang, Jiwen Fan, Yunpeng Shan                                                                                                                  | 2.9.0            |
 | ARM diagnostics v3 (arm_diags)                               | Enhanced ARM diagnostics with new aerosol-cloud-interaction and aerosol activation metrics                                                                                                             | Xiaojian Zheng, Jill Zhang, Cheng Tao, Shaocheng Xie                                                                                                                              | 2.9.0            |
 | Aerosol budget diagnostics (aerosol_budget)                  | Global mean burdens, source/sink budgets and lifetimes for aerosol species.                                                                                                                            | Jill Zhang, Susannah Burrows, Naser Mahfouz, Tom Vo, Kai Zhang, Taufiq Hassan Mozumder, Xue Zheng, Ziming Ke, Yan Feng, Manish Shrivastava, Hailong Wang, Mingxuan Wu | 2.8.0            |

auxiliary_tools/tropical_subseasonal_diags/run_tropical_subseasonal.py

@@ -0,0 +1,26 @@
+import os
+from e3sm_diags.parameter.tropical_subseasonal_parameter import TropicalSubseasonalParameter
+from e3sm_diags.run import runner
+
+param = TropicalSubseasonalParameter()
+
+param.reference_data_path = '/global/cfs/cdirs/e3sm/e3sm_diags/obs_for_e3sm_diags/time-series'
+#param.reference_data_path = '/global/cfs/cdirs/e3sm/chengzhu/e3sm_diags_zppy_test_complete_run_output/v2.LR.historical_0101_20240130/post/atm/180x360_aave/ts/daily/15yr'
+param.test_data_path = '/global/cfs/cdirs/e3sm/chengzhu/e3sm_diags_zppy_test_complete_run_output/v2.LR.historical_0101_20240130/post/atm/180x360_aave/ts/daily/15yr'
+param.test_name = 'E3SMv2'
+prefix = '/Users/zhang40/Documents/repos/e3sm_diags/auxiliary_tools/tropical_subseasonal_diags/data1'
+prefix = '/global/cfs/cdirs/e3sm/www/chengzhu/tests/tropical_diags_subsetting'
+
+param.results_dir = os.path.join(prefix, 'tropical_variability_model_obs_refine')
+#param.run_type = "model_vs_model"
+#param.ref_name = 'E3SMv2'
+param.test_start_yr = '2000'
+param.test_end_yr = '2000'
+param.ref_start_yr = '2001'
+param.ref_end_yr = '2001'
+param.save_netcdf = True
+
+runner.sets_to_run = ['tropical_subseasonal']
+runner.run_diags([param])
+
+

auxiliary_tools/tropical_subseasonal_diags/tropical_diags.cfg

@@ -0,0 +1,28 @@
+[#]
+sets = ["tropical_subseasonal"]
+case_id = "wavenumber-frequency"
+variables = ["PRECT"]
+ref_name = "IMERG_Daily"
+reference_name = "IMERG Daily"
+# (2001-2021)"
+regions = ["15S15N"]
+#
+
+[#]
+sets = ["tropical_subseasonal"]
+case_id = "wavenumber-frequency"
+variables = ["FLUT"]
+ref_name = "NOAA-OLR_Daily"
+reference_name = "NOAA OLR Daily"
+# (1979-2021)"
+regions = ["15S15N"]
+#
+#
+#[#]
+#sets = ["tropical_subseasonal"]
+#case_id = "wavenumber-frequency"
+#variables = ["U850"]
+#ref_name = "ERA5_Daily"
+#reference_name = "ERA5 Daily"
+## (1980-2022)"
+#regions = ["15S15N"]

auxiliary_tools/tropical_subseasonal_diags/tropical_subseasonal_driver.py

@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import glob
+import json
+import os
+from typing import TYPE_CHECKING  # , Optional
+
+import numpy as np
+import xarray as xr
+
+import e3sm_diags
+from e3sm_diags.driver import utils
+from e3sm_diags.logger import custom_logger
+from e3sm_diags.plot.cartopy.mp_partition_plot import plot
+import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap, BoundaryNorm
+
+from tropical_subseasonal_plot import plot
+from tropical_subseasonal_viewer import create_viewer
+from zwf import zwf_functions as wf
+
+logger = custom_logger(__name__)
+
+# Script to compute and plot spectral powers of a subseasonal tropical field in 
+#   zonal wavenumber-frequency space.  Both the plot files and files containing the
+#   associated numerical data shown in the plots are created.
+
+# Authors: Jim Benedict and Brian Medeiros
+# Modified by Jill Zhang to integrate into E3SM Diags. 
+
+def find_nearest(array, value):
+    array = np.asarray(array)
+    idx = (np.abs(array - value)).argmin()
+    return idx,array[idx]
+    """Return index of [array] closest in value to [value]
+    Example:
+    array = [ 0.21069679  0.61290182  0.63425412  0.84635244  0.91599191  0.00213826
+              0.17104965  0.56874386  0.57319379  0.28719469]
+    print(find_nearest(array, value=0.5))
+    # 0.568743859261
+
+    """
+
+def wf_analysis(x, **kwargs):
+    """Return zonal wavenumber-frequency power spectra of x.  The returned spectra are:
+    spec_sym:    Raw (non-normalized) power spectrum of the component of x that is symmetric about the equator.
+    spec_asym:   Raw (non-normalized) power spectrum of the component of x that is antisymmetric about the equator.
+    nspec_sym:   Normalized (by a smoothed red-noise background spectrum) power spectrum of the component of x that is symmetric about the equator.
+    nspec_asym:  Normalized (by a smoothed red-noise background spectrum) power spectrum of the component of x that is antisymmetric about the equator.
+
+    The NCL version of 'wkSpaceTime' smooths the symmetric and antisymmetric components
+    along the frequency dimension using a 1-2-1 filter once.  
+
+    """
+    # Get the "raw" spectral power
+    # OPTIONAL kwargs: 
+    # segsize, noverlap, spd, latitude_bounds (tuple: (south, north)), dosymmetries, rmvLowFrq
+
+    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)
+
+    # Following NCL's wkSpaceTime, apply one pass of a 1-2-1 filter along the frequency
+    #   domain to the raw (non-normalized) spectra/um.
+    #   Do not use 0 frequency when smoothing here.
+    #   Use weights that sum to 1 to ensure that smoothing is conservative.
+    z2s = wf.smoothFrq121(z2,1)
+
+    # The background is supposed to be derived from both symmetric & antisymmetric
+    # Inputs to the background spectrum calculation should be z2avg
+    background = wf.smoothBackground_wavefreq(z2avg)
+    # separate components
+    spec_sym = z2s[0,...]
+    spec_asy = z2s[1,...]
+    # normalize:  Following NCL's wkSpaceTime, use lightly smoothed version of spectra/um
+    #             as numerator
+    nspec_sym = spec_sym / background
+    nspec_asy = spec_asy / background
+
+    spec = xr.merge([spec_sym.rename('spec_raw_sym'), spec_asy.rename('spec_raw_asy'), nspec_sym.rename('spec_norm_sym'), nspec_asy.rename('spec_norm_asy'), background.rename('spec_background')], compat='override')
+    spec_all = spec.drop('component')
+    spec_all['spec_raw_sym'].attrs = {"component": "symmetric", "type": "raw"}
+    spec_all['spec_raw_asy'].attrs = {"component": "antisymmetric", "type": "raw"}
+    spec_all['spec_norm_sym'].attrs = {"component": "symmetric", "type": "normalized"}
+    spec_all['spec_norm_asy'].attrs = {"component": "antisymmetric", "type": "normalized"}
+    spec_all['spec_background'].attrs = {"component": "", "type": "background"}
+
+    return spec_all
+
+
+def calculate_spectrum(path, variable):
+
+    var = xr.open_mfdataset(glob.glob(f"{test_data_path}/{variable}_*.nc")).sel(
+        lat=slice(-15, 15))[variable]
+
+    # TODO: subset time
+
+    # Unit conversion
+    if var.name == "PRECT":
+        if var.attrs["units"] == "m/s" or var.attrs["units"] == "m s{-1}":
+            print("\nBEFORE unit conversion: Max/min of data: " + str(var.values.max()) + "   " + str(var.values.min()))
+            var.values = var.values * 1000. * 86400.    # convert m/s to mm/d, do not alter metadata (yet)
+            var.attrs["units"] = "mm/d"                  # adjust metadata to reflect change in units
+            print("\nAFTER unit conversion: Max/min of data: " + str(var.values.max()) + "   " + str(var.values.min()))
+    if var.name == "precipAvg":
+        if var.attrs["units"] == "mm/hr":
+            print("\nBEFORE unit conversion: Max/min of data: " + str(var.values.max()) + "   " + str(var.values.min()))
+            var.values = data.values * 24.               # convert mm/hr to mm/d, do not alter metadata (yet)
+            var.attrs["units"] = "mm/d"                  # adjust metadata to reflect change in units
+            print("\nAFTER unit conversion: Max/min of data: " + str(var.values.max()) + "   " + str(var.values.min()))
+
+    # Wavenumber Frequency Analysis
+    spec_all = wf_analysis(var, **opt)
+    #spec_all.to_netcdf(outDataDir + "/full_spec.nc")
+    return spec_all
+
+#
+# Options ... right now these only go into wk.spacetime_power()
+#
+do_zooming = False    # Set to True to also make plots to zoom into MJO spectral region,
+                     #   in addition to the default (larger) spectral region
+latBound = (-15,15)  # latitude bounds for analysis
+spd      = 1    # SAMPLES PER DAY
+nDayWin  = 96   # Wheeler-Kiladis [WK] temporal window length (days)
+nDaySkip = -60  # time (days) between temporal windows [segments]
+                # negative means there will be overlapping temporal segments
+twoMonthOverlap = -1*nDaySkip
+
+vari = "PRECT"
+srcID = "model"
+outDataDir = "/global/cfs/cdirs/e3sm/www/chengzhu/tests/tropical_diags"
+outDataDir = "/Users/zhang40/Documents/repos/e3sm_diags/auxiliary_tools/tropical_subseasonal_diags/data"
+
+opt = {'segsize': nDayWin,
+       'noverlap': twoMonthOverlap,
+       'spd': spd,
+       'latitude_bounds': latBound,
+       'dosymmetries': True,
+       'rmvLowFrq':True}
+
+
+#datapath = '/global/cfs/cdirs/e3sm/forsyth/E3SMv2/v2.LR.historical_0201/post/atm/180x360_aave/ts/daily/5yr'
+datapath = '/Users/zhang40/Documents/e3sm_diags_data/e3sm_diags_test_data/E3SM_v2_daily'
+
+from e3sm_diags.parameter.core_parameter import CoreParameter
+parameter = CoreParameter()
+
+test_data_path = '/Users/zhang40/Documents/e3sm_diags_data/e3sm_diags_test_data/E3SM_v2_daily'
+parameter.test_data_path = test_data_path
+parameter.test_timeseries_input = True
+parameter.test_start_yr = '2000'
+parameter.test_end_yr = '2014'
+parameter.ref_data_path = test_data_path
+parameter.ref_timeseries_input = True
+parameter.ref_start_yr = '2000'
+parameter.ref_end_yr = '2014'
+parameter.variables = ['PRECT']
+season = "ANN"
+
+test_data = utils.dataset.Dataset(parameter, test=True)
+parameter.test_name_yrs = utils.general.get_name_and_yrs(
+    parameter, test_data, season
+) 
+
+ref_data = utils.dataset.Dataset(parameter, ref=True)
+parameter.ref_name_yrs = utils.general.get_name_and_yrs(
+    parameter, ref_data, season
+) 
+
+for variable in parameter.variables:
+    #test = calculate_spectrum(parameter.test_data_path, variable)
+    #test.to_netcdf("data/full_spec_test.nc")
+    #ref = calculate_spectrum(parameter.ref_data_path, variable)
+    #ref.to_netcdf("data/full_spec_ref.nc")
+    # Below uses intermediate saved files for development
+    test = xr.open_dataset("/Users/zhang40/Documents/repos/e3sm_diags/auxiliary_tools/tropical_subseasonal_diags/data/full_spec_ref.nc").load()
+    ref = xr.open_dataset("/Users/zhang40/Documents/repos/e3sm_diags/auxiliary_tools/tropical_subseasonal_diags/data/full_spec_ref.nc").load()
+    parameter.var_id = variable
+
+    for diff_name in ["raw_sym", "raw_asy", "norm_sym", "norm_asy", "background"]:
+
+         # Compute percentage difference
+         diff = 100 * (test[f"spec_{diff_name}"]-ref[f"spec_{diff_name}"])/ref[f"spec_{diff_name}"]
+         diff.name = f"spec_{diff_name}"
+         diff.attrs.update(test[f"spec_{diff_name}"].attrs)
+         parameter.spec_type = diff_name
+         plot(parameter, test[f"spec_{diff_name}"], ref[f"spec_{diff_name}"], diff)
+         if "norm" in diff_name:
+             parameter.spec_type = f"{diff_name}_zoom"
+             plot(parameter, test[f"spec_{diff_name}"], ref[f"spec_{diff_name}"], diff, do_zoom = True)
+
+
+
+display_name, url = create_viewer('.', parameter)
+print("Viewer Created: ", url)
+