diff --git a/config_file.jsonc b/config_file.jsonc
new file mode 100644
index 000000000..4f8f0aa7d
--- /dev/null
+++ b/config_file.jsonc
@@ -0,0 +1,125 @@
+// Configuration for MDTF-diagnostics driver script self-test.
+//
+// Copy this file and customize the settings as needed to run the framework on
+// your own model output without repeating command-line options. Pass it to the
+// framework at the end of the command line (positionally) or with the
+// -f/--input-file flag. Any other explicit command line options will override
+// what's listed here.
+//
+// All text to the right of an unquoted "//" is a comment and ignored, as well
+// as blank lines (JSONC quasi-standard.)
+{
+ "case_list" : [
+ // The cases below correspond to the different sample model data sets. Note
+ // that the MDTF package does not currently support analyzing multiple
+ // models in a single invocation. Comment out or delete the first entry and
+ // uncomment the second to run NOAA-GFDL-AM4 only for the MJO_prop_amp POD,
+ // and likewise for the SM_ET_coupling POD.
+ {
+ "CASENAME" : "CESM2-WACCM_CMIP_historical_r1i1p1f1",
+ "model" : "CESM2-WACCM_CMIP_historical_r1i1p1f1",
+ "convention" : "CMIP",
+ "FIRSTYR" : 1979,
+ "LASTYR" : 2014,
+ "QBOisobar" : 30,
+ "pod_list": ["stc_qbo_enso"]
+ }
+ // {
+ // "CASENAME" : "GFDL.CM4.c96L32.am4g10r8",
+ // "model" : "AM4",
+ // "convention" : "GFDL",
+ // "FIRSTYR" : 1,
+ // "LASTYR" : 10,
+ // "pod_list" : ["MJO_prop_amp"]
+ // }
+ // {
+ // "CASENAME" : "Lmon_GISS-E2-H_historical_r1i1p1",
+ // "model" : "CMIP",
+ // "convention" : "CMIP",
+ // "FIRSTYR" : 1951,
+ // "LASTYR" : 2005,
+ // "pod_list" : ["SM_ET_coupling"]
+ // }
+ // {
+ // "CASENAME" : "NCAR-CAM5.timeslice",
+ // "model" : "CESM",
+ // "convention" : "CMIP",
+ // "FIRSTYR" : 2000,
+ // "LASTYR" : 2004,
+ // "pod_list": ["example"]
+ // }
+ ],
+ // PATHS ---------------------------------------------------------------------
+ // Location of supporting data downloaded when the framework was installed.
+
+ // If a relative path is given, it's resolved relative to the MDTF-diagnostics
+ // code directory. Environment variables (eg, $HOME) can be referenced with a
+ // "$" and will be expended to their current values when the framework runs.
+
+ // Parent directory containing observational data used by individual PODs.
+ "OBS_DATA_ROOT": "../inputdata/obs_data",
+
+ // Parent directory containing results from different models.
+ //"MODEL_DATA_ROOT": "/Volumes/Personal-Folders/CCP-Amy/CMIP6_fromZac/data-for-amy/",
+ //"MODEL_DATA_ROOT": "/Users/delsbury/Desktop/mdtf/scratch/",
+ "MODEL_DATA_ROOT": "/Users/delsbury/Desktop/mdtf/inputdata/model/",
+
+ // Working directory. Defaults to OUTPUT_DIR if blank.
+ "WORKING_DIR": "../wkdir",
+
+ // Directory to write output. The results of each run of the framework will be
+ // put in a subdirectory of this directory.
+ "OUTPUT_DIR": "../wkdir",
+
+ // Location of the Anaconda/miniconda installation to use for managing
+ // dependencies (path returned by running `conda info --base`.) If empty,
+ // framework will attempt to determine location of system's conda installation.
+ "conda_root": "/Users/delsbury/opt/anaconda3",
+
+ // Directory containing the framework-specific conda environments. This should
+ // be equal to the "--env_dir" flag passed to conda_env_setup.sh. If left
+ // blank, the framework will look for its environments in the system default
+ // location.
+ "conda_env_root": "/Users/delsbury/opt/anaconda3/envs",
+
+ // SETTINGS ------------------------------------------------------------------
+ // Any command-line option recognized by the mdtf script (type `mdtf --help`)
+ // can be set here, in the form "flag name": "desired setting".
+
+ // Method used to fetch model data.
+ "data_manager": "Local_File",
+
+ // Type of data that POD(s) will analyze
+ // "single_run" (default) or "multi_run"
+ "data_type": "single_run",
+
+ // Method used to manage dependencies.
+ "environment_manager": "Conda",
+
+ // Settings affecting what output is generated:
+
+ // Set to true to have PODs save postscript figures in addition to bitmaps.
+ "save_ps": false,
+
+ // Set to true to have PODs save netCDF files of processed data.
+ "save_nc": true,
+
+ // Set to true to save HTML and bitmap plots in a .tar file.
+ "make_variab_tar": false,
+
+ // Set to true to overwrite results in OUTPUT_DIR; otherwise results saved
+ // under a unique name.
+ "overwrite": true,
+
+ // Settings used in debugging:
+
+ // Log verbosity level.
+ "verbose": 1,
+
+ // Set to true for framework test. Data is fetched but PODs are not run.
+ "test_mode": false,
+
+ // Set to true for framework test. No external commands are run and no remote
+ // data is copied. Implies test_mode.
+ "dry_run": false
+}
diff --git a/data/fieldlist_CMIP.jsonc b/data/fieldlist_CMIP.jsonc
index 41324bb9f..25144e7a8 100644
--- a/data/fieldlist_CMIP.jsonc
+++ b/data/fieldlist_CMIP.jsonc
@@ -49,6 +49,12 @@
"scalar_coord_templates": {"plev": "ua{value}"},
"ndim": 4
},
+ "u10": {
+ "standard_name": "10 hPa zonal_wind",
+ "units": "m s-1",
+ "scalar_coord_templates": {"plev": "u10{value}"},
+ "ndim": 3
+ },
"va": {
"standard_name": "northward_wind",
"units": "m s-1",
@@ -147,12 +153,22 @@
"units": "Pa",
"ndim": 3
},
+ "tos": {
+ "standard_name": "sea_surface_temperature",
+ "units": "degC",
+ "ndim": 3
+ },
"sfcWind": {
"standard_name": "wind_speed",
"units": "m s-1",
"modifier": "atmos_height",
"ndim": 3
},
+ "tos": {
+ "standard_name": "sea_surface_temperature",
+ "units": "K",
+ "ndim": 3
+ },
// radiative fluxes:
"rsus": {
"standard_name": "surface_upwelling_shortwave_flux_in_air",
diff --git a/diagnostics/stc_qbo_enso/config_file.jsonc b/diagnostics/stc_qbo_enso/config_file.jsonc
new file mode 100644
index 000000000..4f8f0aa7d
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/config_file.jsonc
@@ -0,0 +1,125 @@
+// Configuration for MDTF-diagnostics driver script self-test.
+//
+// Copy this file and customize the settings as needed to run the framework on
+// your own model output without repeating command-line options. Pass it to the
+// framework at the end of the command line (positionally) or with the
+// -f/--input-file flag. Any other explicit command line options will override
+// what's listed here.
+//
+// All text to the right of an unquoted "//" is a comment and ignored, as well
+// as blank lines (JSONC quasi-standard.)
+{
+ "case_list" : [
+ // The cases below correspond to the different sample model data sets. Note
+ // that the MDTF package does not currently support analyzing multiple
+ // models in a single invocation. Comment out or delete the first entry and
+ // uncomment the second to run NOAA-GFDL-AM4 only for the MJO_prop_amp POD,
+ // and likewise for the SM_ET_coupling POD.
+ {
+ "CASENAME" : "CESM2-WACCM_CMIP_historical_r1i1p1f1",
+ "model" : "CESM2-WACCM_CMIP_historical_r1i1p1f1",
+ "convention" : "CMIP",
+ "FIRSTYR" : 1979,
+ "LASTYR" : 2014,
+ "QBOisobar" : 30,
+ "pod_list": ["stc_qbo_enso"]
+ }
+ // {
+ // "CASENAME" : "GFDL.CM4.c96L32.am4g10r8",
+ // "model" : "AM4",
+ // "convention" : "GFDL",
+ // "FIRSTYR" : 1,
+ // "LASTYR" : 10,
+ // "pod_list" : ["MJO_prop_amp"]
+ // }
+ // {
+ // "CASENAME" : "Lmon_GISS-E2-H_historical_r1i1p1",
+ // "model" : "CMIP",
+ // "convention" : "CMIP",
+ // "FIRSTYR" : 1951,
+ // "LASTYR" : 2005,
+ // "pod_list" : ["SM_ET_coupling"]
+ // }
+ // {
+ // "CASENAME" : "NCAR-CAM5.timeslice",
+ // "model" : "CESM",
+ // "convention" : "CMIP",
+ // "FIRSTYR" : 2000,
+ // "LASTYR" : 2004,
+ // "pod_list": ["example"]
+ // }
+ ],
+ // PATHS ---------------------------------------------------------------------
+ // Location of supporting data downloaded when the framework was installed.
+
+ // If a relative path is given, it's resolved relative to the MDTF-diagnostics
+ // code directory. Environment variables (eg, $HOME) can be referenced with a
+ // "$" and will be expended to their current values when the framework runs.
+
+ // Parent directory containing observational data used by individual PODs.
+ "OBS_DATA_ROOT": "../inputdata/obs_data",
+
+ // Parent directory containing results from different models.
+ //"MODEL_DATA_ROOT": "/Volumes/Personal-Folders/CCP-Amy/CMIP6_fromZac/data-for-amy/",
+ //"MODEL_DATA_ROOT": "/Users/delsbury/Desktop/mdtf/scratch/",
+ "MODEL_DATA_ROOT": "/Users/delsbury/Desktop/mdtf/inputdata/model/",
+
+ // Working directory. Defaults to OUTPUT_DIR if blank.
+ "WORKING_DIR": "../wkdir",
+
+ // Directory to write output. The results of each run of the framework will be
+ // put in a subdirectory of this directory.
+ "OUTPUT_DIR": "../wkdir",
+
+ // Location of the Anaconda/miniconda installation to use for managing
+ // dependencies (path returned by running `conda info --base`.) If empty,
+ // framework will attempt to determine location of system's conda installation.
+ "conda_root": "/Users/delsbury/opt/anaconda3",
+
+ // Directory containing the framework-specific conda environments. This should
+ // be equal to the "--env_dir" flag passed to conda_env_setup.sh. If left
+ // blank, the framework will look for its environments in the system default
+ // location.
+ "conda_env_root": "/Users/delsbury/opt/anaconda3/envs",
+
+ // SETTINGS ------------------------------------------------------------------
+ // Any command-line option recognized by the mdtf script (type `mdtf --help`)
+ // can be set here, in the form "flag name": "desired setting".
+
+ // Method used to fetch model data.
+ "data_manager": "Local_File",
+
+ // Type of data that POD(s) will analyze
+ // "single_run" (default) or "multi_run"
+ "data_type": "single_run",
+
+ // Method used to manage dependencies.
+ "environment_manager": "Conda",
+
+ // Settings affecting what output is generated:
+
+ // Set to true to have PODs save postscript figures in addition to bitmaps.
+ "save_ps": false,
+
+ // Set to true to have PODs save netCDF files of processed data.
+ "save_nc": true,
+
+ // Set to true to save HTML and bitmap plots in a .tar file.
+ "make_variab_tar": false,
+
+ // Set to true to overwrite results in OUTPUT_DIR; otherwise results saved
+ // under a unique name.
+ "overwrite": true,
+
+ // Settings used in debugging:
+
+ // Log verbosity level.
+ "verbose": 1,
+
+ // Set to true for framework test. Data is fetched but PODs are not run.
+ "test_mode": false,
+
+ // Set to true for framework test. No external commands are run and no remote
+ // data is copied. Implies test_mode.
+ "dry_run": false
+}
diff --git a/diagnostics/stc_qbo_enso/doc/stc_qbo_enso.rst b/diagnostics/stc_qbo_enso/doc/stc_qbo_enso.rst
new file mode 100644
index 000000000..47fcc9e49
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/doc/stc_qbo_enso.rst
@@ -0,0 +1,209 @@
+.. This is a comment in RestructuredText format (two periods and a space).
+
+.. Note that all "statements" and "paragraphs" need to be separated by a blank
+ line. This means the source code can be hard-wrapped to 80 columns for ease
+ of reading. Multi-line comments or commands like this need to be indented by
+ exactly three spaces.
+
+.. Underline with '='s to set top-level heading:
+ https://docutils.sourceforge.io/docs/user/rst/quickref.html#section-structure
+
+Stratosphere-Troposphere Coupling: QBO and ENSO stratospheric teleconnections
+================================
+
+Last update: 2023-10-03
+
+This script and its helper scripts (“stc_qbo_enso_plottingcodeqbo.py” and
+“stc_qbo_enso_plottingcodeenso.py”) do calculations to assess the representation
+of stratospheric telconnections associated with the Quasi-Biennial Oscillation
+(QBO) and the El Nino Southern Oscillation (ENSO). This POD uses monthly 4D
+(time x plev x lat x lon) zonal wind, 4D meridional wind, 4D temperature, 3D
+(time x lat x lon) sea level pressure, and 3D sea surface temperature data.
+Coupling between the QBO and the boreal polar stratosphere takes place during
+boreal fall and winter whereas coupling between the QBO and the austral polar
+stratosphere takes place mainly during austral spring and summer. By default,
+the POD defines the QBO for NH (SH) analyses using the Oct-Nov (Jul-Aug) 5S-5N
+30 hPa zonal winds. The QBO is thought to influence the polar stratospheres,
+the so-called “polar route,” by modulating the lower stratospheric (~100-50 hPa)
+and middle stratospheric (~20-5 hPa) mid-latitude circulation. The aforementioned
+lower stratospheric teleconnection is also associated with a change in the strength
+and position of the tropospheric jet; the so-called “subtropical route.” In addition,
+evidence continues to show that the QBO directly influences the underlying tropical
+tropospheric circulation, referred to as the “tropical route.” These three
+teleconnections allow the QBO to elicit surface impacts globally. Said teleconnections
+are visualized herein by using a metric of planetary wave propagation (eddy heat flux),
+circulation response (zonal wind), and surface impact (sea level pressure).
+Additionally, metrics of model QBOs (e.g., amplitude, height, width) are produced.
+ENSO’s coupling with the polar stratospheres takes place as the amplitude of ENSO
+maximizes during boreal fall and winter. By default, the POD defines ENSO for NH
+(SH) analyses using the Nov-Mar (Sep-Jan) Nino3.4 SSTs. Though ENSO’s teleconnections
+are global, it interacts with the stratosphere by stimulating tropical-extratropical
+Rossby waves that constructively interfere with the climatological extratropical
+stationary wave mainly over the Pacific, promoting enhanced upward planetary wave
+propagation into the stratosphere. Similar to the QBO code, ENSO’s teleconnections
+are visualized using the eddy heat flux, the zonal wind, and the sea level pressure.
+
+This POD makes six kinds of figures and one text file from provided model data:
+
+- Zonal-mean zonal wind anomalies (deviations from seasonal cycle) composited
+ based on El Nino and La Nina years are shown in red/blue shading. Nina minus
+ Nino differences are shown in shading as well and climatological winds are
+ overlaid on all aforementioned plots in black contours
+- Zonal-mean eddy heat flux anomalies composited based on El Nino and La Nina
+ years are shown in red/blue shading. Nina minus Nino differences are shown
+ in shading as well and climatological heat flux is overlaid on all aforementioned
+ plots in black contours
+- Sea level pressure anomalies composited based on El Nino and La Nina
+ years are shown in red/blue shading. Nina minus Nino differences are shown
+ in shading as well and climatological sea level pressure is overlaid on all aforementioned
+ plots in black contours
+- A text file of QBO metrics (min/mean/max QBO periodicity, easterly/westerly/total
+ amplitude, lowest altitude tropical stratospheric isobar that the QBO reaches,
+ the height or vertical extent of the QBO, and its latitudinal width) is produced.
+- Should the above QBO metrics code detect a QBO in the model data, similar plots as
+ the aforementioned three ENSO plots, but composited based on easterly
+ and westerly QBO years, are made
+
+These plots are made for both hemispheres, with a focus on winter and spring, the seasons
+when upward propagating planetary waves couple the troposphere and stratosphere.
+The metrics are designed to reveal the extratropical circulation response to two forms
+of tropical internal variability, which are generally difficult to represent spontaneously
+in climate models (QBO+ENSO). Though ENSO's representation in climate models as well as the
+representation of its teleconnections has significantly improved over multiple generations
+of CMIP experiments (Bellenger et al. 2014; Planton et al. 2021), it is less clear how
+ENSO's coupling with the polar stratosphere is represented by models. The sea level
+pressure ENSO responses reveal the precursor tropospheric forcing patterns
+(e.g., Aleutian Low response) that should stimulate or reduce upward planetary wave
+propagation into the stratosphere. The zonal wind and eddy heat flux plots reveal if the
+reinforced or suppressed upward planetary wave propagation due to ENSO is actually "felt"
+in the stratosphere and the sea level pressure plots can again be referenced for evidence
+of a downward propagating winter/spring annular mode response to ENSO modulating the polar vortex.
+
+Similar plots to the ones made based on El Nino and La Nina years are made for easterly
+and westerly QBO years if and when a QBO is detected in the model data; e.g., models with
+too-coarse vertical resolution may not simulate a QBO (Zhao et al. 2018). It should be
+interesting to compare the QBO metrics with the representation of QBO teleconnections in
+models. Models struggle to represent several QBO attributes (Richter et al. 2020) and
+since the structure of the QBO (e.g., its amplitude or latitudinal width) is intimately
+tied to the representation of QBO teleconnections (Garfinkel and Hartmann 2011; Hansen
+et al. 2013), models generally have a difficult time representing the extratropical
+impacts of the QBO (Rao et al. 2020).
+
+
+Version & Contact info
+----------------------
+
+- Version/revision information: v1.0 (Oct 2023)
+- Project PIs: Amy H. Butler (NOAA CSL) and Zachary D. Lawrence (CIRES/NOAA PSL)
+- Developer/point of contact: Dillon Elsbury (dillon.elsbury@noaa.gov)
+
+Open source copyright agreement
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The MDTF framework is distributed under the LGPLv3 license (see LICENSE.txt).
+
+
+Functionality
+-------------
+
+This POD is driven by the file ``stc_qbo_enso.py``, with two helper scripts called
+``stc_qbo_enso_plottingcodeenso.py`` and ``stc_qbo_enso_plottingcodeqbo.py``.
+The driver script reads in the model fields, identifies El Nino/La Nina years and
+easterly/westerly QBO years, computes the eddy heat flux, retrieves the QBO metrics,
+and then uses the two helper scripts to make the associated zonal wind, eddy heat
+flux, and sea level pressure plots.
+
+The atmospheric observational data this POD uses is based on ERA5 reanalysis
+(Hersbach, et al., 2020), and includes pre-computed monthly zonal-mean zonal winds, zonal-
+mean eddy heat fluxes, and sea level pressure. The oceanic observational data that
+this POD uses is from HadiSST (Rayner et al. 2003) and includes pre-computed monthly sea
+surface temperature.
+
+
+Required programming language and libraries
+-------------------------------------------
+
+This POD requires Python 3, with the following packages:
+
+- numpy
+- xarray
+- xesmf
+- os
+- matplotlib
+- cartopy
+- scipy
+
+Required model output variables
+-------------------------------
+
+The following monthly mean fields are required:
+
+- Zonal Winds, ``ua`` as ``(time,lev,lat,lon)`` (units: m/s)
+- Meridional Winds, ``va`` as ``(time,lev,lat,lon)`` (units: m/s)
+- Temperature, ``ta`` as ``(time,lev,lat,lon)`` (units: K)
+- Sea level pressure, ``psl`` as ``(time,lat,lon)`` (units: Pa)
+- Sea surface temperature, ``tos`` as ``(time,lat,lon)`` (units: Kelvin)
+
+References
+----------
+
+.. _ref-Bellenger:
+
+ Bellenger, H., Guilyardi, E., Leloup, J., Lengaigne, M., & Vialard, J. (2014).
+ ENSO representation in climate models: From CMIP3 to CMIP5. Climate Dynamics, 42,
+ 1999-2018, https://doi.org/10.1007/s00382-013-1783-z
+
+.. _ref-Planton:
+
+ Planton, Y. Y., Guilyardi, E., Wittenberg, A. T., Lee, J., Gleckler, P. J., Bayr, T.,
+ ... & Voldoire, A. (2021). Evaluating climate models with the CLIVAR 2020 ENSO metrics
+ package. Bulletin of the American Meteorological Society, 102(2), E193-E217,
+ https://doi.org/10.1175/BAMS-D-19-0337.1
+
+.. _ref-Zhao:
+
+ Zhao, M., Golaz, J. C., Held, I. M., Guo, H., Balaji, V., Benson, R., ... & Xiang, B.
+ (2018). The GFDL global atmosphere and land model AM4. 0/LM4. 0: 1. Simulation
+ characteristics with prescribed SSTs. Journal of Advances in Modeling Earth Systems,
+ 10(3), 691-734, https://doi.org/10.1002/2017MS001209
+
+.. _ref-Hersbach:
+
+ Hersbach, H. and coauthors, 2020: The ERA5 global reanalysis. Q J R Meteorol Soc.,
+ 146, 1999-2049, https://doi.org/10.1002/qj.3803
+
+.. _ref-Richter:
+
+ Richter, J. H., Anstey, J. A., Butchart, N., Kawatani, Y., Meehl, G. A., Osprey, S.,
+ & Simpson, I. R. (2020). Progress in simulating the quasi‐biennial oscillation in
+ CMIP models. Journal of Geophysical Research: Atmospheres, 125(8), e2019JD032362,
+ https://doi.org/10.1029/2019JD032362
+
+.. _ref-Garfinkel:
+
+ Garfinkel, C. I., & Hartmann, D. L. (2011). The influence of the quasi-biennial
+ oscillation on the troposphere in winter in a hierarchy of models. Part I: Simplified
+ dry GCMs. Journal of the Atmospheric Sciences, 68(6), 1273-1289,
+ https://doi.org/10.1175/2011JAS3665.1
+
+.. _ref-Hansen:
+ Hansen, F., Matthes, K., & Gray, L. J. (2013). Sensitivity of stratospheric dynamics
+ and chemistry to QBO nudging width in the chemistry‒climate model WACCM. Journal of
+ Geophysical Research: Atmospheres, 118(18), 10-464,
+ https://doi.org/10.1002/jgrd.50812
+
+.. _ref-Rao:
+ Rao, J., Garfinkel, C. I., & White, I. P. (2020). Impact of the quasi-biennial
+ oscillation on the northern winter stratospheric polar vortex in CMIP5/6 models.
+ Journal of Climate, 33(11), 4787-4813, https://doi.org/10.1175/JCLI-D-19-0663.1
+
+More about this POD
+--------------------------
+
+**Statistical testing**
+
+A student's 2-tailed t-test is used to assess how likely it is that the Nina minus
+Nino anomalies and easterly QBO minus westerly QBO anomalies arise by chance for the zonal
+wind, eddy heat flux, and sea level pressure plots. A p-value <= 0.05 is used as the
+threshold for "statistical significance," which is denoted on the aforementioned figures
+in the third row using stippling.
\ No newline at end of file
diff --git a/diagnostics/stc_qbo_enso/make_era5_rean_pod_data.py b/diagnostics/stc_qbo_enso/make_era5_rean_pod_data.py
new file mode 100644
index 000000000..35e30eae2
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/make_era5_rean_pod_data.py
@@ -0,0 +1,190 @@
+# ==============================================================================
+# MDTF Strat-Trop Coupling: Vertical Wave Propagation POD
+# ==============================================================================
+#
+# This file is part of the Strat-Trop Coupling: Vertical Wave Propagation POD
+# of the MDTF code package (see mdtf/MDTF-diagnostics/LICENSE.txt)
+#
+# This script shows a simple recipe for creating the "pre-digested" reanalysis
+# data used for making the reanalysis versions of the POD plots.
+#
+# Data for individual variables obtained from:
+# https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-pressure-levels-monthly-means?tab=form
+#
+# Need the following variables:
+# -----------------------------
+# v-component of wind
+# air temperature
+# geopotential (this is post-processed to geopot height by dividing by 9.80665)
+
+import numpy as np
+import xarray as xr
+import xesmf as xe
+
+### BEGIN: READ INPUT FIELDS ###
+# The following code/paths will have to be adapted for your own system.
+#
+# On my system, the monthly-mean variables are contained in individual
+# files that span all available years in the reanalysis from 1979 to
+# the ~present. I set these as environment variables, but explicitly
+# show the paths in comments as examples
+
+# Some needed functions #
+
+def field_regridding(ds,latname,lonname):
+
+ r""" Regrid input data so that there are 73 latitude points. This grid
+ includes 5S and 5N, which are needed for the QBO analysis. This grid
+ includes 60N, which is needed for the SSW analysis. """
+
+ # Check to see if the latitudes are organized north to south or south to north
+ lat_first = ds[latname].values[0]
+ lat_end = ds[latname].values[-1]
+
+ if lat_first >= lat_end:
+ lats = np.linspace(90,-90,num=73)
+ if lat_end > lat_first:
+ lats = np.linspace(-90,90,num=73)
+
+ # Check to see if the longitudes are organized -180/180 or 0 to 360
+ lon_first = ds[lonname].values[0]
+ print (lon_first, 'lon_first')
+
+ if lon_first < 0:
+ lons = np.linspace(-180,177.5,num=144)
+ if lon_first >= 0:
+ lons = np.linspace(0,357.5,num=144)
+
+ ds_out = xr.Dataset({'lat': (['lat'], lats),'lon': (['lon'], lons),})
+ regridder = xe.Regridder(ds, ds_out, 'bilinear')
+ regridded = regridder(ds)
+ print (regridded, 'regridded')
+
+ return regridded
+
+def compute_total_eddy_heat_flux(varray, tarray, vname, tname):
+
+ r""" Compute the total (all zonal wavenumbers) eddy heat flux
+ using monthly data. Output field has new variable, 'ehf.' """
+
+ # Take the zonal means of v and T
+ dummy = varray.mean('lon')
+
+ eddyv = (varray - varray.mean('lon'))[vname]
+ eddyt = (tarray - tarray.mean('lon'))[tname]
+
+ ehf = np.nanmean(eddyv.values * eddyt.values,axis=-1)
+ dummy[vname].values[:] = ehf
+ dummy = dummy.rename({vname:'ehf'})
+ print (dummy)
+
+ return dummy
+
+# Load the observational data #
+
+sfi = '/Volumes/Personal-Folders/CCP-Dillon/ERA5/stationary/POD/HadISST_sst.nc'
+pfi = '/Volumes/Personal-Folders/CCP-Dillon/ERA5/stationary/POD/cat-era5-prmsl-monmean-91-180.nc'
+ufi = '/Volumes/Personal-Folders/CCP-Dillon/ERA5/stationary/POD/cat-era5-uwnd-monmean-91-180.nc'
+vfi = '/Volumes/Personal-Folders/CCP-Dillon/ERA5/stationary/POD/cat-era5-vwnd-monmean-91-180.nc'
+tfi = '/Volumes/Personal-Folders/CCP-Dillon/ERA5/stationary/POD/cat-era5-air-monmean-91-180.nc'
+
+# Open datasets #
+
+sst_ds = xr.open_dataset(sfi)
+psl_ds = xr.open_dataset(pfi)
+uwnd_ds = xr.open_dataset(ufi)
+vwnd_ds = xr.open_dataset(vfi)
+air_ds = xr.open_dataset(tfi)
+
+# Regrid #
+
+sst_regridded = field_regridding(sst_ds,'latitude','longitude')
+psl_regridded = field_regridding(psl_ds,'lat','lon')
+uwnd_regridded = field_regridding(uwnd_ds,'lat','lon')
+vwnd_regridded = field_regridding(vwnd_ds,'lat','lon')
+air_regridded = field_regridding(air_ds,'lat','lon')
+
+# By the end of this block of code, the vwnd, air, and hgt variables
+# should each contain all available months of meridional wind,
+# air temperature and geopotential height, respectively. They can be
+# lazily loaded with xarray (e.g., after using open_mfdataset)
+### END: READ INPUT FIELDS ###
+
+r""" Compute the total (all zonal wavenumbers) eddy heat flux
+using monthly data. Output field has new variable, 'ehf.' """
+
+# Take the zonal means of v and T
+dummy = vwnd_regridded.mean('lon')
+
+eddyv = (vwnd_regridded - vwnd_regridded.mean('lon'))['vwnd']
+eddyt = (air_regridded - air_regridded.mean('lon'))['air']
+
+ehf = np.nanmean(eddyv.values * eddyt.values,axis=-1)
+dummy['vwnd'].values[:] = ehf
+ehf = dummy.rename({'vwnd':'ehf'})
+ehf.attrs['long_name'] = "Zonal Mean Eddy Heat Flux (v'T')"
+
+r""" Zonally average the zonal wind """
+uzm = uwnd_regridded.mean('lon')
+uzm = uzm.rename({'uwnd':'ua'})
+
+r""" Change name in psl file """
+psl_out = psl_regridded.rename({'prmsl':'psl'})
+
+print ('######### BREAK ##########')
+
+print (sst_regridded)
+print ('sst_regridded')
+print (' ')
+print (ehf)
+print ('ehf')
+print (' ')
+print (uzm)
+print ('uzm')
+print (' ')
+print (psl_ds)
+print ('psl_ds')
+print (' ')
+
+
+# Merge DataArrays into output dataset
+out_ds = xr.merge([ehf, uzm, psl_out])
+print (out_ds)
+out_ds = out_ds.rename({'level':'lev'})
+out_ds.attrs['reanalysis'] = 'ERA5'
+out_ds.attrs['notes'] = 'Fields derived from monthly-mean ERA5 data on pressure levels'
+
+out_ds.psl.attrs['units'] = 'Pa'
+out_ds.ua.lev.attrs['units'] = 'hPa'
+out_ds.ehf.lev.attrs['units'] = 'hPa'
+
+sst_out_ds = sst_regridded
+sst_out_ds.attrs['reanalysis'] = 'HadiSST'
+sst_out_ds.attrs['notes'] = 'Fields derived from monthly-mean HadiSST sea surface temperature'
+sst_out_ds = sst_out_ds.rename({'sst':'tos'})
+
+'''
+# To reduce size of output file without changing results much, will thin the
+# available latitudes from 0.25 to 0.5 degrees. This is optional.
+out_ds = out_ds.isel(lat=slice(0,None,2))
+
+# To reduce size of output file even more, we will only keep latitudes poleward
+# of 30 degrees (since we are primarily interested in the extratropics).
+# This step is also optional.
+out_ds = out_ds.where(np.abs(out_ds.lat) >= 30, drop=True)
+
+# To reduce size of output file even more more, we will use zlib compression
+# on each variable. This step is also optional.
+encoding = {'ta_zm_50': {'dtype':'float32', 'zlib':True, 'complevel':7},
+ 'zg_zm': {'dtype':'float32', 'zlib':True, 'complevel':7},
+ 'ehf_100': {'dtype':'float32', 'zlib':True, 'complevel':7}}
+
+# Save the output file
+filename = 'stc_eddy_heat_fluxes_obs-data.nc'
+out_ds.rename({'level':'lev'}).to_netcdf(filename, encoding=encoding)
+'''
+
+filename = 'stc-qbo-enso-obs-atm.nc'
+out_ds.to_netcdf(filename)
+filename = 'stc-qbo-enso-obs-ocn.nc'
+sst_out_ds.to_netcdf(filename)
diff --git a/diagnostics/stc_qbo_enso/settings.jsonc b/diagnostics/stc_qbo_enso/settings.jsonc
new file mode 100644
index 000000000..f5b1219e1
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/settings.jsonc
@@ -0,0 +1,73 @@
+// Strat-Trop Coupling: Vertical Wave Propagation
+//
+// This POD requires monthly-frequency meridional winds, air temperatures,
+// and geopotential heights with pressure levels in the troposphere
+// and stratosphere.
+//
+{
+ "settings" : {
+ "driver" : "stc_qbo_enso.py",
+ "long_name" : "Metrics of QBO and extratropical circulation impact of QBO and ENSO",
+ "realm" : ["atmos", "ocean"],
+ "description" : "Assess the representation of the QBO",
+ "pod_env_vars" : {
+ // Isobar (hPa) used to define the QBO in the tropical stratosphere
+ // Defaults to 30 hPa
+ "QBOisobar" : "30"
+ },
+ "runtime_requirements": {
+ "python3": ["matplotlib", "numpy", "scipy", "xarray", "xesmf"]
+ }
+ },
+ "data": {
+ "min_frequency": "6hr",
+ "max_frequency": "yr"},
+ "dimensions": {
+ "lat": {"standard_name": "latitude"},
+ "lon": {"standard_name": "longitude"},
+ "lev": {
+ "standard_name": "air_pressure",
+ "units": "Pa",
+ "positive": "down",
+ "axis": "Z"
+ },
+ "time": {"standard_name": "time"}
+ },
+ "varlist": {
+ "tos": {
+ "standard_name" : "sea_surface_temperature",
+ "units" : "degC",
+ "frequency": "mon",
+ "dimensions": ["time", "lat", "lon"],
+ "requirement": "required"
+ },
+ "ua": {
+ "standard_name" : "eastward_wind",
+ "units" : "m s-1",
+ "frequency": "mon",
+ "dimensions": ["time", "lev", "lat", "lon"],
+ "requirement": "required"
+ },
+ "va": {
+ "standard_name" : "northward_wind",
+ "units" : "m s-1",
+ "frequency": "mon",
+ "dimensions": ["time", "lev", "lat", "lon"],
+ "requirement": "required"
+ },
+ "ta": {
+ "standard_name" : "air_temperature",
+ "units" : "K",
+ "frequency": "mon",
+ "dimensions": ["time", "lev", "lat", "lon"],
+ "requirement": "required"
+ },
+ "psl": {
+ "standard_name" : "air_pressure_at_mean_sea_level",
+ "units" : "Pa",
+ "frequency": "mon",
+ "dimensions": ["time", "lat", "lon"],
+ "requirement": "required"
+ }
+ }
+}
\ No newline at end of file
diff --git a/diagnostics/stc_qbo_enso/stc_qbo_enso.html b/diagnostics/stc_qbo_enso/stc_qbo_enso.html
new file mode 100644
index 000000000..c852bab5d
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/stc_qbo_enso.html
@@ -0,0 +1,181 @@
+
+
MDTF example diagnostic
+
+STC QBO and ENSO stratospheric teleconnections
+
+This script and its helper scripts (“stc_qbo_enso_plottingcodeqbo.py” and “stc_qbo_enso_plottingcodeenso.py”) do
+calculations to assess the representation of stratospheric telconnections associated with the Quasi-Biennial
+Oscillation (QBO) and the El Nino Southern Oscillation (ENSO). This POD uses monthly 4D (time x plev x lat x lon)
+zonal wind, 4D meridional wind, 4D temperature, 3D (time x lat x lon) sea level pressure, and 3D sea surface
+temperature data. Coupling between the QBO and the boreal polar stratosphere takes place during boreal fall and
+winter whereas coupling between the QBO and the austral polar stratosphere takes place mainly during austral
+spring and summer. By default, the POD defines the QBO for NH (SH) analyses using the Oct-Nov (Jul-Aug) 5S-5N
+30 hPa zonal winds. The QBO is thought to influence the polar stratospheres, the so-called “polar route,” by
+modulating the lower stratospheric (~100-50 hPa) and middle stratospheric (~20-5 hPa) mid-latitude circulation. The
+aforementioned lower stratospheric teleconnection is also associated with a change in the strength and position of
+the tropospheric jet; the so-called “subtropical route.” In addition, evidence continues to show that the QBO directly
+influences the underlying tropical tropospheric circulation, referred to as the “tropical route.” These three
+teleconnections allow the QBO to elicit surface impacts globally. Said teleconnections are visualized herein by using
+a metric of planetary wave propagation (eddy heat flux), circulation response (zonal wind), and surface impact (sea
+level pressure). Additionally, metrics of model QBOs (e.g., amplitude, height, width) are produced. ENSO’s coupling
+with the polar stratospheres takes place as the amplitude of ENSO maximizes during boreal fall and winter. By
+default, the POD defines ENSO for NH (SH) analyses using the Nov-Mar (Sep-Jan) Nino3.4 SSTs. Though ENSO’s
+teleconnections are truly global, it interacts with the stratosphere by stimulating tropical-extratropical Rossby waves
+that constructively interfere with the climatological extratropical stationary wave, promoting enhanced upward
+planetary wave propagation into the stratosphere. Similar to the QBO code, ENSO’s teleconnections are visualized
+using the eddy heat flux, the zonal wind, and the sea level pressure.
+
+
+{{CASENAME}}
+ENSO teleconnections northern hemisphere
+
+
+
+
+ENSO teleconnections southern hemisphere
+
+
+
+
+QBO teleconnections northern hemisphere
+
+
+
+
+QBO teleconnections southern hemisphere
+
+
+
+
diff --git a/diagnostics/stc_qbo_enso/stc_qbo_enso.py b/diagnostics/stc_qbo_enso/stc_qbo_enso.py
new file mode 100644
index 000000000..0b9ef72ed
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/stc_qbo_enso.py
@@ -0,0 +1,1217 @@
+# ==============================================================================
+# MDTF Strat-Trop Coupling: QBO and ENSO stratospheric teleconnections
+# ==============================================================================
+#
+# This file is part of the Strat-Trop Coupling: QBO and ENSO stratospheric teleconnections
+# POD of the MDTF code package (see mdtf/MDTF-diagnostics/LICENSE.txt)
+#
+# STC QBO and ENSO stratospheric teleconnections
+# Last update: 2023-10-03
+#
+# This script and its helper scripts (“stc_qbo_enso_plottingcodeqbo.py” and “stc_qbo_enso_plottingcodeenso.py”) do
+# calculations to assess the representation of stratospheric telconnections associated with the Quasi-Biennial
+# Oscillation (QBO) and the El Nino Southern Oscillation (ENSO). This POD uses monthly 4D (time x plev x lat x lon)
+# zonal wind, 4D meridional wind, 4D temperature, 3D (time x lat x lon) sea level pressure, and 3D sea surface
+# temperature data. Coupling between the QBO and the boreal polar stratosphere takes place during boreal fall and
+# winter whereas coupling between the QBO and the austral polar stratosphere takes place mainly during austral
+# spring and summer. By default, the POD defines the QBO for NH (SH) analyses using the Oct-Nov (Jul-Aug) 5S-5N
+# 30 hPa zonal winds. The QBO is thought to influence the polar stratospheres, the so-called “polar route,” by
+# modulating the lower stratospheric (~100-50 hPa) and middle stratospheric (~20-5 hPa) mid-latitude circulation. The
+# aforementioned lower stratospheric teleconnection is also associated with a change in the strength and position of
+# the tropospheric jet; the so-called “subtropical route.” In addition, evidence continues to show that the QBO directly
+# influences the underlying tropical tropospheric circulation, referred to as the “tropical route.” These three
+# teleconnections allow the QBO to elicit surface impacts globally. Said teleconnections are visualized herein by using
+# a metric of planetary wave propagation (eddy heat flux), circulation response (zonal wind), and surface impact (sea
+# level pressure). Additionally, metrics of model QBOs (e.g., amplitude, height, width) are produced. ENSO’s coupling
+# with the polar stratospheres takes place as the amplitude of ENSO maximizes during boreal fall and winter. By
+# default, the POD defines ENSO for NH (SH) analyses using the Nov-Mar (Sep-Jan) Nino3.4 SSTs. Though ENSO’s
+# teleconnections are truly global, it interacts with the stratosphere by stimulating tropical-extratropical Rossby waves
+# that constructively interfere with the climatological extratropical stationary wave, promoting enhanced upward
+# planetary wave propagation into the stratosphere. Similar to the QBO code, ENSO’s teleconnections are visualized
+# using the eddy heat flux, the zonal wind, and the sea level pressure.
+#
+# Please see the references for the scientific foundations of this POD.
+#
+# ==============================================================================
+# Version, Contact Info, and License
+# ==============================================================================
+# - Version/revision information: v1.0 (2024-01-23)
+# - PIs: Amy H. Butler, NOAA CSL & Zachary D. Lawrence, CIRES + CU Boulder /NOAA PSL
+# - Developer/point of contact: Dillon Elsbury, dillon.elsbury@noaa.gov
+# - Other contributors: Zachary D. Lawrence, CIRES + CU Boulder / NOAA PSL,
+# zachary.lawrence@noaa.gov; Amy H. Butler, NOAA CSL
+#
+# The MDTF framework is distributed under the LGPLv3 license (see LICENSE.txt).
+#
+# ==============================================================================
+# Functionality
+# ==============================================================================
+# This POD contains three scripts. The primary script is stc_qbo_enso.py. There
+# are two helper scripts with functions used by the primary script that are
+# called stc_qbo_enso_plottingcodeqbo.py and stc_qbo_enso_plottingcodeenso.py.
+# The primary script stc_qbo_enso.py goes through these basic steps:
+# (1) Loads in the reanalysis data and restricts the time period to 1979-2014.
+# The 1979-2014 period is the default period, but this can be altered by
+# modifying the FIRSTYR and LASTYR environment variables, which can be
+# specified in config_file.jsonc and the settings.jsonc.
+# (2) Computes the reanalysis ENSO indices and composities the zonal-mean zonal wind,
+# zonal-mean eddy heat flux, and sea level pressure around the El Nino and
+# La Nina years.
+# (3) Does the same as (2), but for the QBO. The POD then produces reanalysis
+# QBO metrics. By default, the QBO indexing and QBO metrics are calculated by
+# defining the QBO at 30 hPa. This can be altered using the QBOisobar environment
+# variable defined in config_file.jsonc and settings.jsonc.
+# (4) Loads in the model data and restrics the time period to 1979-2014 by default.
+# The vertical ("plev") axis of the 4D fields (zonal wind, meridional wind, and
+# temperature) is then modified, if necessary, to have pressure levels
+# denoted in hPa rather than Pa, which is used in some data sets.
+# (5) Computes the model ENSO indices and composites the zonal-mean zonal wind, zonal-
+# mean eddy heat flux, and sea level pressure around the model El Nino/La Nina years.
+# (6) Runs the QBO metric code using the default definition of the QBO at 30 hPa.
+# If a QBO is detected (some models cannot simulate one), proceeds to compute
+# model QBO indices and composite the zonal-mean zonal wind, zonal-mean eddy heat
+# flux, and sea level pressure around easterly and westerly QBO years. string
+#
+# ==============================================================================
+# Required programming language and libraries
+# ==============================================================================
+# This POD is done fully in python, and primarily makes use of numpy and
+# xarray to read, subset, and transform the data. It also makes use of scipy to
+# calculate the fast fourier transform (FFT) of the tropical stratospheric zonal
+# winds to identify the QBO using its frequency spectrum (20-36 months in
+# observations). matplotlib and cartopy are used for general plotting and
+# creating map plots.
+#
+# ==============================================================================
+# Required model input variables
+# ==============================================================================
+# This POD requires monthly-mean fields of
+# - 4D (time x plev x lat x lon) zonal wind velocity (ua, units: m/s)
+# - 4D (time x plev x lat x lon) meridional wind velocity (va, units: m/s)
+# - 4D (time x plev x lat x lon) temperature (ta, units: Kelvin)
+# - 3D (tmee x lat x lon) sea level pressure (psl, units: Pa)
+# - 3D (time x lat x lon) sea surface temperature (tos, units: Kelvin)
+#
+# ==============================================================================
+# References
+# ==============================================================================
+# QBO metrics related papers:
+#
+# Schenzinger, Verena, et al. "Defining metrics of the Quasi-Biennial Oscillation in
+# global climate models." Geoscientific Model Development 10.6 (2017): 2157-2168
+# doi.org/10.5194/gmd-10-2157-2017
+# Richter, Jadwiga H., et al. "Progress in simulating the quasi‐biennial oscillation
+# in CMIP models." Journal of Geophysical Research: Atmospheres 125.8 (2020):
+# e2019JD032362. doi.org/10.1029/2019JD032362
+# Pascoe, Charlotte L., et al. "The quasi‐biennial oscillation: Analysis using ERA‐40
+# data." Journal of Geophysical Research: Atmospheres 110.D8 (2005).
+# doi.org/10.1029/2004JD004941
+#
+# QBO teleconnections, surface impacts, and indexing methods
+#
+# Gray, Lesley J., et al. "Surface impacts of the quasi biennial oscillation."
+# Atmospheric Chemistry and Physics 18.11 (2018): 8227-8247.
+# doi.org/10.5194/acp-18-8227-2018
+# Rao, Jian, et al. "Southern hemisphere response to the quasi-biennial oscillation
+# in the CMIP5/6 models." Journal of Climate 36.8 (2023): 2603-2623.
+# doi.org/10.1175/JCLI-D-22-0675.1
+#
+# ENSO teleconnections, surface impacts, and indexing methods
+#
+# Domeisen, Daniela IV, Chaim I. Garfinkel, and Amy H. Butler. "The teleconnection
+# of El Niño Southern Oscillation to the stratosphere." Reviews of Geophysics 57.1
+# (2019): 5-47. doi.org/10.1029/2018RG000596
+# Iza, Maddalen, Natalia Calvo, and Elisa Manzini. "The stratospheric pathway of
+# La Niña." Journal of Climate 29.24 (2016): 8899-8914.
+# doi.org/10.1175/JCLI-D-16-0230.1
+
+import os
+import xarray as xr
+import numpy as np
+import xesmf as xe
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+
+from scipy.fft import fft,fftfreq
+from scipy import interpolate
+
+from stc_qbo_enso_plottingcodeqbo import qbo_uzm
+from stc_qbo_enso_plottingcodeqbo import qbo_vt
+from stc_qbo_enso_plottingcodeqbo import qbo_slp
+
+from stc_qbo_enso_plottingcodeenso import enso_uzm
+from stc_qbo_enso_plottingcodeenso import enso_vt
+from stc_qbo_enso_plottingcodeenso import enso_slp
+
+mpl.rcParams['font.family'] = 'sans-serif'
+mpl.rcParams['font.sans-serif'] = 'Roboto'
+mpl.rcParams['font.size'] = 12
+mpl.rcParams['hatch.color']='gray'
+
+def qbo_metrics(ds,QBOisobar):
+
+ r""" Calculates Quasi-Biennial Oscillation (QBO) metrics for the an input
+ zonal wind dataset (dimensions = time x level x latitude x longitude)
+
+ Parameters
+ ----------
+ ds : `xarray.DataArray` or `xarray.Dataset`
+ The input DataArray or Dataset for which to calculate QBO diagnostics for
+
+ Returns
+ -------
+ period_min: scalar
+ The minimum number of months comprising a full QBO cycle (a period of easterlies [EQBO]
+ and westerlies [WQBO])
+
+ period_mean: scalar
+ The average QBO cycle (EQBO + WQBO) duration in months
+
+ period_max: scalar
+ The maximum QBO cycle (EQBO + WQBO) duration in months
+
+ easterly_amp: scalar
+ The average easterly amplitude, arrived at by averaging together the minimum latitudinally
+ averaged 5S-5N 10 hPa monthly zonal wind from each QBO cycle
+
+ westerly_amp: scalar
+ The average westerly amplitude, arrived at by averaging together the minimum latitudinally
+ averaged 5S-5N 10 hPa monthly zonal wind from each QBO cycle
+
+ qbo_amp: scalar
+ The total QBO amplitude, which is estimated by adding half of the mean easterly amplitude
+ (easterly_amp) to half of the mean westerly amplitude (westerly_amp)
+
+ lowest_lev: scalar
+ The lowermost pressure level at which the the latitudinally averaged 5S-5N QBO Fourier
+ amplitude is equal to 10% of its maximum value
+
+ latitudinal_extent: scalar
+ The full width at half maximum of a Gaussian fit to the 10 hPa portion of the QBO
+ Fourier amplitude pressure-latitude cross section
+
+ Notes
+ -----
+ The input xarray variable ds is assumed to have a dimension named "time x lev x lat x lon."
+ E.g., if your data differs in dimension name, e.g., "latitude", use the rename method:
+ ds.rename({'latitude':'lat'})
+ ds.rename({'level':'lev'})
+
+ period_min is required to be >= 14 months. This requirement is used because period_min and
+ period_max are the time bounds that determine which frequencies are most "QBO-like." This
+ step is needed, because say e.g., period_min was allowed to be less than 14 months and ended
+ up being equal to 12 months. Then the annual cycle zonal wind variability would be deemed
+ "QBO-like" and annual cycle variability would be incorporated into the QBO Fourier amplitude
+ calculations, rendering the Fourier amplitude and all of the QBO spatial metrics useless.
+ """
+
+ print ("Running the QBO metrics function 'qbo_metrics'")
+
+ if ds.lev.values[-1] > ds.lev.values[0]:
+ ds = ds.reindex(lev=list(reversed(ds.lev)))
+
+ uwnd = ds.ua.values
+
+ # Subset for 10 hPa #
+ subset = ds.sel(lev=QBOisobar)
+
+ # Select 5S-5N winds #
+ tropical_subset = subset.isel(lat = np.logical_and(subset.lat >= -5, subset.lat <= 5))
+
+ # Latitudinally weight and averaged betwteen 5S-5N
+ qbo = tropical_subset.mean('lat')
+ weights = np.cos(np.deg2rad(tropical_subset.lat.values))
+ interim = np.multiply(tropical_subset.ua.values,weights[np.newaxis,:])
+ interim = np.nansum(interim,axis=1)
+ interim = np.true_divide(interim,np.sum(weights))
+ qbo.ua.values[:] = interim[:]
+
+ # Smooth with five month running mean #
+ qbo = qbo.rolling(time=5, center=True).mean()
+
+ # Identify the indices corresponding to QBO phase changes (e.g., + zonal wind (westerly) -> - zonal wind (easterly))
+ zero_crossings = np.where(np.diff(np.sign(qbo.ua.values)))[0]
+
+ # Using the phase change indices, identify QBO cycles: a set of easterlies and westerlies. After doing this, #
+ # the minimum/maximum/average QBO cycle duration will be retrieved. The first phase change index from #
+ # zero_crossings is excluded in the event that the first QBO cycle is close to making a phase change, which would bias #
+ # the QBO duration statistics low. #
+ store = []
+ cycles = []
+ periods = []
+
+ for i,v in enumerate(zero_crossings):
+
+ # Append pairs of easterly/westerly winds to "store" #
+ if i != 0:
+ tmp = qbo.ua.values[zero_crossings[i-1]+1:zero_crossings[i]+1]
+ store.append(tmp)
+
+ # Retrieve the most recent QBO cycle (easterlies + westerlies) from store. Loop looks for even number indices. #
+ # E.g., if i == 2, one set of QBO easterlies and westerlies has been appended to store already. #
+ if i != 0 and i % 2 == 0:
+ concat = np.concatenate((store[-2],store[-1]))
+
+ # Inserting requirement that each cycle must be at least 14 months. No observed QBO cycle has progressed #
+ # this quickly, but cycles do in the models. This requirement is essential because the minimum and maximum QBO cycle
+ # durations are used to calculate the QBO Fourier amplitude. A minimum QBO cycle duration of, say, 12 months #
+ # would lead to the QBO Fourier amplitude overlapping with the annual cycle Fourier amplitude.
+ if len(concat) >= 14:
+ cycles.append(concat)
+ periods.append(len(concat))
+
+ # Retrieve the minimum/maximum/average QBO cycle duration #
+ period_min = np.round(np.nanmin(periods),1)
+ period_max = np.round(np.nanmax(periods),1)
+ period_mean = np.round(np.nanmean(periods),1)
+
+ print (period_min, "minimum period (months)")
+ print (period_mean, "mean period (months)")
+ print (period_max, "maximum period (months)")
+
+ # Retrieve the minimum/maximum zonal wind from each QBO cycle. Averaging the minima (maxima) gives us the #
+ # easterly (westerly) amplitute #
+
+ easterly_amp = np.round(np.nanmean([np.nanmin(v) for v in cycles]),1)
+ westerly_amp = np.round(np.nanmean([np.nanmax(v) for v in cycles]),1)
+
+ print (easterly_amp, 'easterly amplitude')
+ print (westerly_amp, 'westerly amplitude')
+
+ # Define the 10 hPa amplitude as in Richter et al. (2020)
+ qbo_amp = np.abs(easterly_amp/2) + np.abs(westerly_amp/2)
+ qbo_amp = np.round(qbo_amp,1)
+ print (qbo_amp, '10 hPa amplitude')
+
+ #################################################################################################################
+ # Retrieve the QBO Fourier amplitude, defined as in Pascoe et al. (2005) as the ratio of the QBO power spectrum #
+ # to the power spectrum of the entire zonal wind dataset, multiplied by a metric of the zonal wind variability, #
+ # the standard deviation of the zonal wind dataset. #
+ #################################################################################################################
+
+ # Standard deviation across entire zonal wind dataset #
+ std = np.nanstd(uwnd,axis=0)
+
+ # Define Fourier frequencies comprising data and filter for frequencies between minimum/maximum QBO cycle duration #
+ freq = 1/fftfreq(len(uwnd))
+ arr = np.where((freq > period_min) & (freq < period_max))[0]
+
+ # FFT of entire zonal wind dataset. Square and sum Fourier coefficients. np.abs applies unneeded square root, hence #
+ # np.power to power 2 is used to undo this #
+ y = fft(uwnd,axis=0)
+ amplitudes = np.power(np.abs(y)[:len(y)//2],2)
+
+ # Calculate ratio of QBO power spectrum to full zonal wind power spectrum #
+ quotients = []
+ for i,v in enumerate(ds.lev.values):
+ qbodata = np.nansum(amplitudes[arr,i],axis=0)
+ alldata = np.nansum(amplitudes[1:,i],axis=0)
+ quot = np.true_divide(qbodata,alldata)
+ quotients.append(quot)
+ filtered = np.array(quotients)
+
+ # Ratio of the QBO power spectrum to the power #
+ # spectrum of the entire model dataset (units:%) #
+
+ vmin = 0
+ vmax = 100
+ vlevs = np.linspace(vmin,vmax,num=21)
+
+ from palettable.colorbrewer.diverging import RdBu_11
+ x2,y2 = np.meshgrid(ds.lat.values,ds.lev.values)
+
+ plt.title('Ratio of QBO power spectrum\n to zonal wind spectrum (%)')
+ plt.contourf(x2,y2,filtered*100,vmin=vmin,vmax=vmax,levels=vlevs,cmap='coolwarm')
+ plt.semilogy()
+ plt.gca().invert_yaxis()
+ if np.nanmax(ds.lev.values) > 2000:
+ plt.ylabel('Pressure (Pa)')
+ if np.nanmax(ds.lev.values) < 2000:
+ plt.ylabel('Pressure (hPa)')
+ plt.xlabel('Latitude')
+ plt.colorbar()
+
+ # Retrieve the Fourier amplitude by multiplying aforementioned ratio by standard deviation of zonal wind #
+ fa = np.multiply(filtered,std)
+
+ #################################################################################################################
+ # Do the Fourier amplitude calculations between 5S and 5N to retrieve spatial metrics (e.g., latitudinal width) #
+ # of the QBO #
+ #################################################################################################################
+
+ # hmax fixed at 10 hPa per Richter et al. (2020, JGR) #
+ hmax = np.where(ds.lev.values == qbo.lev.values)[0]
+
+ # Retreive the indices of lats between 5S and 5N #
+ lat_hits = [i for i,v in enumerate(ds.lat.values) if v >= -5 and v <=5]
+
+ # Retrieve the Fourier amplitude profile averaged latitudinally (w/weighting) between 5S and 5N #
+ weights = np.cos(np.deg2rad(ds.lat.values[lat_hits]))
+ interim = np.multiply(fa[:,lat_hits],weights[np.newaxis,:])
+ interim2 = np.nansum(interim,axis=1)
+ height_profile = np.true_divide(interim2,np.sum(weights))
+
+ # Retrieve half the max Fourier amplitude and 10% of the max Fourier amplitude #
+ half_max = np.max(height_profile)/2
+ qbo_base = np.max(height_profile)*0.1
+
+ # Interpolate the equatorial Fourier amplitude profile to have 10000 vertical grid points, enough #
+ # points so that each isobar can be selected to a one tenth of a hPa accuracy #
+ f = interpolate.interp1d(ds.lev.values,height_profile)
+ xnew = np.linspace(ds.lev.values[0],ds.lev.values[-1],num=10000)
+ ynew = f(xnew)
+
+ # Of the 20,000 vertical grid points, find the one corresponding to hmax. #
+ hmax_idx = (np.abs(xnew - ds.lev.values[hmax])).argmin()
+
+ # The lower and upper portions of the fourier amplitude tropical wind height profile, #
+ # which has been interpolated to 10000 grid points. #
+ lower_portion = ynew[:hmax_idx]
+ upper_portion = ynew[hmax_idx:]
+
+ # Same as above, but these are lists of the isobars corresponding to the above fourier amplitudes. #
+ lower_portion_isobar = xnew[:hmax_idx]
+ upper_portion_isobar = xnew[hmax_idx:]
+
+ # Retrieve the indices in the upper/lower portions corresponding to half the fourier max. #
+ lower_vertical_extent = (np.abs(lower_portion - half_max)).argmin()
+ upper_vertical_extent = (np.abs(upper_portion - half_max)).argmin()
+
+ # Find the upper/lower portion isboars corresponding to the half fourier max values identified above. #
+ bottom = lower_portion_isobar[lower_vertical_extent]
+ top = upper_portion_isobar[upper_vertical_extent]
+
+ # Convert the isobars into altitudes in meters. #
+ sfc = 1000 # hPa
+ bottomz = np.log(bottom/sfc)*-7000
+ topz = np.log(top/sfc)*-7000
+
+ # Obtain the vertical extent by differencing the bottomz and topz. #
+ vertical_extent = (topz - bottomz)/1000
+ vertical_extent = np.round(vertical_extent,1)
+ print (vertical_extent, "vertical_extent")
+
+ # Retrieve the lowest lev the QBO extends to using 10% of the maximum Fourier amplitude #
+ # "lower" for CMIP6 datasets and "upper" for ERA5 reanalysis
+ lowest_lev = (lower_portion_isobar[(np.abs(lower_portion - qbo_base)).argmin()])
+ lowest_lev = np.round(lowest_lev,1)
+ print (lowest_lev, "lowest_lev")
+
+ ##############################################################################################
+ ##############################################################################################
+ ##############################################################################################
+
+ # Retrieve the latitudinal extent #
+
+ # https://www.geeksforgeeks.org/python-gaussian-fit/
+ xdata = ds.lat.values
+ ydata = fa[hmax][0]
+ ydata[0] = 0
+ ydata[-1] = 0
+
+ # Recast xdata and ydata into numpy arrays so we can use their handy features
+ xdata = np.array(xdata)
+ ydata = np.array(ydata)
+
+ from scipy.optimize import curve_fit
+
+ def gauss(x, H, A, x0, sigma):
+ return H + A * np.exp(-(x - x0) ** 2 / (2 * sigma ** 2))
+
+ def gauss_fit(x, y):
+ mean = sum(x * y) / sum(y)
+ sigma = np.sqrt(sum(y * (x - mean) ** 2) / sum(y))
+ popt, pcov = curve_fit(gauss, x, y, p0=[min(y), max(y), mean, sigma])
+ return popt
+
+ out = gauss(xdata, *gauss_fit(xdata, ydata))
+
+ f = interpolate.interp1d(ds.lat.values,out)
+ xnew = np.linspace(ds.lat.values[0],ds.lat.values[-1],num=10000)
+ ynew = f(xnew)
+
+ lower_portion = ynew[:5000]
+ upper_portion = ynew[5000:]
+
+ lower_portion_lat = xnew[:5000]
+ upper_portion_lat = xnew[5000:]
+
+ lat1 = lower_portion_lat[(np.abs(lower_portion - (np.max(out)/2))).argmin()]
+ lat2 = upper_portion_lat[(np.abs(upper_portion - (np.max(out)/2))).argmin()]
+
+ latitudinal_extent = np.abs(lat1) + np.abs(lat2)
+ latitudinal_extent = np.round(latitudinal_extent,1)
+
+ print (latitudinal_extent, "latitudinal_extent")
+
+
+ if period_min != period_max:
+ print ('Based on period statistics, dataset is likely to have a QBO')
+ qbo_switch = 1
+ else:
+ print ('Persistent stratospheric easterlies detected - dataset likely does not have QBO')
+ qbo_switch = 0
+
+ metrics = ['minimum period: %s (months)' % period_min,
+ 'mean period: %s (months)' % period_mean,
+ 'maximum period: %s (months)' % period_max,
+ 'easterly amplitude: %s (m/s)' % easterly_amp,
+ 'westerly amplitude: %s (m/s)' % westerly_amp,
+ 'QBO amplitude: %s (m/s)' % qbo_amp,
+ 'lowest QBO level: %s (hPa)' % lowest_lev,
+ 'vertical extent: %s (kilometers)' % vertical_extent,
+ 'latitudinal extent of QBO: %s (degrees)' % latitudinal_extent]
+
+ return metrics, qbo_switch
+
+
+def compute_total_eddy_heat_flux(v, T):
+ r""" Compute the total zonal mean eddy heat flux from meridonal winds
+ and temperatures. The eddy heat flux is calculated as:
+ ehf = zonal_mean( (v - zonal_mean(v)) * (T - zonal_mean(T)))
+
+ Parameters
+ ----------
+ v : `xarray.DataArray`
+ The meridional wind component. Assumed to have the same dimensions as T
+
+ T : `xarray.DataArray`
+ The air temperature. Assumed to have the same dimensions as v
+
+ Returns
+ -------
+ ehf : `xarray.DataArray`
+ The zonal mean eddy heat flux
+
+ Notes
+ -----
+ The input fields v and T are assumed to have dimensions named "lat"
+ and "lon". E.g., if your data has dimensions "latitude" and/or "longitude",
+ use the rename method:
+ ds.rename({'latitude':'lat','longitude':'lon'})
+
+ Ideally v and T would be provided on the same latitude/longitude grid.
+ In practice this is not necessarily the case as some models provide
+ different variables at cell-centers and cell-edges. If this is found
+ to be the case, this function will use xesmf to do bilinear regridding
+ of the meridional wind component to match the grid of the temperatures.
+
+ """
+
+ # Take the zonal means of v and T
+ v_zm = v.mean('lon')
+ T_zm = T.mean('lon')
+
+ # If v and T are on same grid, can multiply the eddies and take zonal mean
+ if (np.array_equal(v.lat,T.lat)) and (np.array_equal(v.lon, T.lon)):
+ ehf = ((v - v_zm) * (T - T_zm)).mean('lon')
+
+ # If v and T are on different grids, interpolate v to T's grid beforehand
+ else:
+ # Set up xESMF regridder with necessary grid-defining datasets
+ print('*** Interpolating v to same grid as T')
+ in_grid = xr.Dataset(
+ {
+ "lat": (["lat"], v.lat.values),
+ "lon": (["lon"], v.lon.values),
+ }
+ )
+
+ out_grid = xr.Dataset(
+ {
+ "lat": (["lat"], T.lat.values),
+ "lon": (["lon"], T.lon.values),
+ }
+ )
+ regridder = xe.Regridder(in_grid, out_grid, "bilinear")
+
+ ehf = (regridder(v - v_zm)*(T - T_zm)).mean('lon')
+
+ return ehf
+
+##################################################################################################
+##################################################################################################
+##################################################################################################
+
+########################
+# --- BEGIN SCRIPT --- #
+########################
+print('\n=======================================')
+print('BEGIN stc_qbo_enso.py ')
+print('=======================================\n')
+
+# Parse MDTF-set environment variables
+print('*** Parse MDTF-set environment variables ...')
+CASENAME = os.environ['CASENAME']
+FIRSTYR = os.environ['FIRSTYR']
+LASTYR = os.environ['LASTYR']
+WK_DIR = os.environ['WK_DIR']
+OBS_DATA = os.environ['OBS_DATA']
+QBOisobar = os.environ['QBOisobar']
+
+###########################################################################
+################################ observations #############################
+###########################################################################
+
+print(f'*** Now working on obs data\n------------------------------')
+obs_file_atm = OBS_DATA+'/stc-qbo-enso-obs-atm.nc'
+obs_file_ocn = OBS_DATA+'/stc-qbo-enso-obs-ocn.nc'
+
+print(f'*** Reading obs data from {obs_file_atm}')
+obs_atm = xr.open_dataset(obs_file_atm)
+
+print (obs_atm, 'obs_atm')
+
+print(f'*** Reading obs data from {obs_file_ocn}')
+obs_sst = xr.open_dataset(obs_file_ocn)
+
+print (obs_sst, 'obs_sst')
+
+# Subset the data for the user defined first and last years #
+
+obs_atm = obs_atm.sel(time=slice(str(FIRSTYR),str(LASTYR)))
+obs_sst = obs_sst.sel(time=slice(str(FIRSTYR),str(LASTYR)))
+
+# Create the POD figures directory
+
+plot_dir = f'{WK_DIR}/obs/'
+
+################################################
+print ('*** Running the observed ENSO indexing')
+################################################
+
+# Extract the tropical domain #
+ENSO = obs_sst.isel(lat = np.logical_and(obs_sst.lat >= -5, obs_sst.lat <= 5))
+
+# Extract date and longitude info from ENSO dataset #
+date_first = obs_sst.time[0]
+date_last = obs_sst.time[-1]
+lon_first = obs_sst.lon.values[0]
+
+# Identify the correct ENSO longitudinal grid #
+if lon_first < 0:
+ ENSO = ENSO.sel(lon=slice(-170,-120))
+else:
+ ENSO = ENSO.sel(lon=slice(190,240))
+
+# Latitudinally average the ENSO data #
+weighted_mean = ENSO.mean('lat')
+weights = np.cos(np.deg2rad(ENSO.lat.values))
+interim = np.multiply(ENSO.tos.values,weights[np.newaxis,:,np.newaxis])
+interim = np.nansum(interim,axis=1)
+interim = np.true_divide(interim,np.sum(weights))
+weighted_mean.tos.values[:] = interim[:]
+
+def enso_index(seasonal):
+
+ # Create 5-month seasonally averaged standardized ENSO anomalies. Weight each month by number of days comprising month #
+ day_in_month_weights = seasonal.time.dt.days_in_month.values[:5]/np.sum(seasonal.time.dt.days_in_month.values[:5])
+ sstindex = np.reshape(seasonal.tos.values,(int(len(seasonal.tos.values)/5),5))
+ sstindex = np.nanmean(np.multiply(sstindex,day_in_month_weights[np.newaxis,:]),axis=1)
+ anom = np.subtract(sstindex,np.nanmean(sstindex))
+ anom = np.true_divide(anom,np.nanstd(sstindex))
+
+ # Get the unique years from "seasonal" and then remove the last one, which is not needed
+ years = [v for v in set(np.sort(seasonal.time.dt.year.values))]
+ nina_years = [years[i] for i,v in enumerate(anom) if v <= -1]
+ nino_years = [years[i] for i,v in enumerate(anom) if v >= 1]
+
+ return nina_years, nino_years
+
+# Subsample ENSO data for NH #
+seasonal = weighted_mean.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+seasonal = seasonal.sel(time=seasonal.time.dt.month.isin([11,12,1,2,3])).mean('lon')
+nh_nina, nh_nino = enso_index(seasonal)
+seasonal.close()
+
+# Subsample ENSO data for SH #
+seasonal = weighted_mean.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+seasonal = seasonal.sel(time=seasonal.time.dt.month.isin([9,10,11,12,1])).mean('lon')
+sh_nina, sh_nino = enso_index(seasonal)
+seasonal.close()
+
+# Store the Nina/Nino years in a dictionary to call later #
+enso_dict = {}
+enso_dict['NH'] = nh_nina, nh_nino
+enso_dict['SH'] = sh_nina, sh_nino
+
+##########################################################################
+# Define ENSO plotting parameters to be passed to the plotting functions #
+##########################################################################
+
+nh_enso_uzm = obs_atm.ua.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+nh_enso_vtzm = obs_atm.ehf.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+nh_enso_psl = obs_atm.psl.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+nh_enso_titles = ['November','December','January','February','March']
+nh_enso_plot_months = [11,12,1,2,3]
+nh_enso_axes = [0,90,1000,1]
+nh_enso_psl_axes = [-180, 180, 20, 90]
+
+sh_enso_uzm = obs_atm.ua.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+sh_enso_vtzm = obs_atm.ehf.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+sh_enso_psl = obs_atm.psl.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+sh_enso_titles = ['September','October','November','December','January']
+sh_enso_plot_months = [9,10,11,12,1]
+sh_enso_axes = [-90,0,1000,1]
+sh_enso_psl_axes = [-180, 180, -90, -20]
+
+uzm_dict = {}
+uzm_dict['NH'] = nh_enso_uzm, nh_enso_titles, nh_enso_plot_months, nh_enso_axes
+uzm_dict['SH'] = sh_enso_uzm, sh_enso_titles, sh_enso_plot_months, sh_enso_axes
+
+vtzm_dict = {}
+vtzm_dict['NH'] = nh_enso_vtzm, nh_enso_titles, nh_enso_plot_months, nh_enso_axes
+vtzm_dict['SH'] = sh_enso_vtzm, sh_enso_titles, sh_enso_plot_months, sh_enso_axes
+
+psl_dict = {}
+psl_dict['NH'] = nh_enso_psl, nh_enso_titles, nh_enso_plot_months, ccrs.NorthPolarStereo(), nh_enso_psl_axes
+psl_dict['SH'] = sh_enso_psl, sh_enso_titles, sh_enso_plot_months, ccrs.SouthPolarStereo(), sh_enso_psl_axes
+
+###############################################
+print ('*** Running the observed QBO indexing')
+###############################################
+
+print (QBOisobar, "QBOisobar")
+
+# Subset atmospheric data for user defined isobar #
+subset = obs_atm.sel(lev=QBOisobar)
+
+# Select 5S-5N winds #
+tropical_subset = subset.interp(lat=[-5,-2.5,0,2.5,5])
+
+# Latitudinally weight and average #
+qbo = tropical_subset.mean('lat')
+weights = np.cos(np.deg2rad(tropical_subset.lat.values))
+interim = np.multiply(tropical_subset.ua.values,weights[np.newaxis,:])
+interim = np.nansum(interim,axis=1)
+interim = np.true_divide(interim,np.sum(weights))
+qbo.ua.values[:] = interim[:]
+
+def qbo_index(seasonal):
+
+ # Create 2-month seasonally averaged standardized QBO anomalies. Weight each month by number of days comprising month #
+ day_in_month_weights = seasonal.time.dt.days_in_month.values[:2]/np.sum(seasonal.time.dt.days_in_month.values[:2])
+ qboindex = np.reshape(seasonal.ua.values,(int(len(seasonal.ua.values)/2),2))
+ qboindex = np.nanmean(np.multiply(qboindex,day_in_month_weights[np.newaxis,:]),axis=1)
+ anom = np.subtract(qboindex,np.nanmean(qboindex))
+ anom = np.true_divide(anom,np.nanstd(qboindex))
+
+ # Get the unique years from "seasonal" and then remove the last one, which is not needed
+ years = [v for v in set(np.sort(seasonal.time.dt.year.values))]
+ eqbo_years = [years[i] for i,v in enumerate(anom) if v <= -1]
+ wqbo_years = [years[i] for i,v in enumerate(anom) if v >= 1]
+
+ return eqbo_years, wqbo_years
+
+# Subsample QBO data for NH #
+seasonal = qbo.sel(time=qbo.time.dt.month.isin([10,11])).mean('lon')
+nh_eqbo, nh_wqbo = qbo_index(seasonal)
+seasonal.close()
+
+# Subsample QBO data for SH #
+seasonal = qbo.sel(time=qbo.time.dt.month.isin([7,8])).mean('lon')
+sh_eqbo, sh_wqbo = qbo_index(seasonal)
+seasonal.close()
+
+# Store the Nina/Nino years in a dictionary to call later #
+qbo_dict = {}
+qbo_dict['NH'] = nh_eqbo, nh_wqbo
+qbo_dict['SH'] = sh_eqbo, sh_wqbo
+
+# Extract date and longitude info from QBO dataset #
+date_first = obs_atm.time[0]
+date_last = obs_atm.time[-1]
+lon_first = obs_atm.lon.values[0]
+
+#########################################################################
+# Define QBO plotting parameters to be passed to the plotting functions #
+#########################################################################
+
+nh_qbo_uzm = obs_atm.ua.sel(time=slice('%s-10-01' % date_first.dt.year.values, '%s-02-28' % date_last.dt.year.values))
+nh_qbo_vtzm = obs_atm.ehf.sel(time=slice('%s-10-01' % date_first.dt.year.values, '%s-02-28' % date_last.dt.year.values))
+nh_qbo_psl = obs_atm.psl.sel(time=slice('%s-10-01' % date_first.dt.year.values, '%s-02-28' % date_last.dt.year.values))
+nh_qbo_titles = ['October','November','December','January','February']
+nh_qbo_plot_months = [10,11,12,1,2]
+nh_qbo_axes = [0,90,1000,1]
+nh_qbo_psl_axes = [-180, 180, 20, 90]
+
+sh_qbo_uzm = obs_atm.ua.sel(time=slice('%s-07-01' % date_first.dt.year.values, '%s-11-30' % date_last.dt.year.values))
+sh_qbo_vtzm = obs_atm.ehf.sel(time=slice('%s-07-01' % date_first.dt.year.values, '%s-11-30' % date_last.dt.year.values))
+sh_qbo_psl = obs_atm.psl.sel(time=slice('%s-07-01' % date_first.dt.year.values, '%s-11-30' % date_last.dt.year.values))
+sh_qbo_titles = ['July','August','September','October','November']
+sh_qbo_plot_months = [7,8,9,10,11]
+sh_qbo_axes = [-90,0,1000,1]
+sh_qbo_psl_axes = [-180, 180, -90, -20]
+
+uzm_qbo_dict = {}
+uzm_qbo_dict['NH'] = nh_qbo_uzm, nh_qbo_titles, nh_qbo_plot_months, nh_qbo_axes
+uzm_qbo_dict['SH'] = sh_qbo_uzm, sh_qbo_titles, sh_qbo_plot_months, sh_qbo_axes
+
+vtzm_qbo_dict = {}
+vtzm_qbo_dict['NH'] = nh_qbo_vtzm, nh_qbo_titles, nh_qbo_plot_months, nh_qbo_axes
+vtzm_qbo_dict['SH'] = sh_qbo_vtzm, sh_qbo_titles, sh_qbo_plot_months, sh_qbo_axes
+
+psl_qbo_dict = {}
+psl_qbo_dict['NH'] = nh_qbo_psl, nh_qbo_titles, nh_qbo_plot_months, ccrs.NorthPolarStereo(), nh_qbo_psl_axes
+psl_qbo_dict['SH'] = sh_qbo_psl, sh_qbo_titles, sh_qbo_plot_months, ccrs.SouthPolarStereo(), sh_qbo_psl_axes
+
+hemispheres = ['NH','SH']
+
+for hemi in hemispheres:
+
+ ###############################################
+ print ('*** Calling the observed ENSO indices')
+ ###############################################
+ obs_nina, obs_nino = enso_dict[hemi]
+
+ print (obs_nina,'obs_nina')
+ print (obs_nino, 'obs_nino')
+
+ ###################################################################
+ print ('*** Running the observed ENSO zonal mean zonal wind calcs')
+ ###################################################################
+
+ obstos_plot = f'{plot_dir}/obs-enso34-uzm-{FIRSTYR}-{LASTYR}-%s.png' % hemi
+ out_fig, out_ax = enso_uzm(uzm_dict[hemi][0],obs_nina,obs_nino,uzm_dict[hemi][1],uzm_dict[hemi][2],uzm_dict[hemi][3])
+ out_fig.savefig(obstos_plot,dpi=700)
+
+ ############################################################
+ print ('*** Running the observed ENSO eddy heat flux calcs')
+ ############################################################
+ obsvt_plot = f'{plot_dir}/obs-enso34-vt-{FIRSTYR}-{LASTYR}-%s.png' % hemi
+ out_fig, out_ax = enso_vt(vtzm_dict[hemi][0],obs_nina,obs_nino,vtzm_dict[hemi][1],vtzm_dict[hemi][2],vtzm_dict[hemi][3])
+ out_fig.savefig(obsvt_plot,dpi=700)
+
+ ##########################################################
+ print ('*** Running the observed ENSO sea level pressure')
+ ##########################################################
+ obsps_plot = f'{plot_dir}/obs-enso34-psl-{FIRSTYR}-{LASTYR}-%s.png' % hemi
+ out_fig, out_ax = enso_slp(psl_dict[hemi][0],obs_nina,obs_nino,psl_dict[hemi][1],psl_dict[hemi][2],psl_dict[hemi][3],psl_dict[hemi][4])
+ out_fig.savefig(obsps_plot,dpi=700)
+
+ ##############################################
+ print ('*** Calling the observed QBO indices')
+ ##############################################
+ obs_eqbo, obs_wqbo = qbo_dict[hemi]
+
+ print (obs_eqbo,'obs_eqbo')
+ print (obs_wqbo, 'obs_wqbo')
+
+ #####################################################################
+ print ('*** Running the observed QBO zonal mean zonal wind plotting')
+ #####################################################################
+ uzm_plot = f'{plot_dir}/obs-qbo{QBOisobar}hPa-uzm-{FIRSTYR}-{LASTYR}-%s.png' % hemi
+ out_fig, out_ax = qbo_uzm(uzm_qbo_dict[hemi][0],obs_eqbo,obs_wqbo,QBOisobar,uzm_qbo_dict[hemi][1],uzm_qbo_dict[hemi][2],uzm_qbo_dict[hemi][3])
+ out_fig.savefig(uzm_plot,dpi=700)
+
+ #########################################################################
+ print ('*** Running the observed QBO zonal mean eddy heat flux plotting')
+ #########################################################################
+ vtzm_plot = f'{plot_dir}/obs-qbo{QBOisobar}hPa-vt-{FIRSTYR}-{LASTYR}-%s.png' % hemi
+ out_fig, out_ax = qbo_vt(vtzm_qbo_dict[hemi][0],obs_eqbo,obs_wqbo,QBOisobar,vtzm_qbo_dict[hemi][1],vtzm_qbo_dict[hemi][2],vtzm_qbo_dict[hemi][3])
+ out_fig.savefig(vtzm_plot,dpi=700)
+
+ ##################################################################
+ print ('*** Running the observed QBO sea level pressure plotting')
+ ##################################################################
+ psl_plot = f'{plot_dir}/obs-qbo{QBOisobar}hPa-psl-{FIRSTYR}-{LASTYR}-%s.png' % hemi
+ out_fig, out_ax = qbo_slp(psl_qbo_dict[hemi][0],obs_eqbo,obs_wqbo,QBOisobar,psl_qbo_dict[hemi][1],psl_qbo_dict[hemi][2],psl_qbo_dict[hemi][3],psl_qbo_dict[hemi][4])
+ out_fig.savefig(psl_plot,dpi=700)
+
+
+print('*** Running the observed QBO metrics')
+metricsout, switch = qbo_metrics(obs_atm,QBOisobar)
+
+filepath = f'{plot_dir}/obs-qbo{QBOisobar}hPa-metrics-{FIRSTYR}-{LASTYR}.txt'
+with open(filepath, 'w') as file_handler:
+ file_handler.write(f"{'QBO metrics: periodicity and spatial characteristics'}\n")
+ file_handler.write(f"{' '}\n")
+ for item in metricsout:
+ file_handler.write(f"{item}\n")
+
+###############################################
+# Tidy up by closing the open xarray datasets #
+###############################################
+
+obs_atm.close()
+obs_sst.close()
+ENSO.close()
+weighted_mean.close()
+nh_enso_uzm.close()
+nh_enso_vtzm.close()
+nh_enso_psl.close()
+sh_enso_uzm.close()
+sh_enso_vtzm.close()
+sh_enso_psl.close()
+subset.close()
+tropical_subset.close()
+qbo.close()
+nh_qbo_uzm.close()
+nh_qbo_vtzm.close()
+nh_qbo_psl.close()
+sh_qbo_uzm.close()
+sh_qbo_vtzm.close()
+sh_qbo_psl.close()
+
+###########################################################################
+################################## model ##################################
+###########################################################################
+
+plot_dir = f'{WK_DIR}/model/'
+
+# Read the input model data
+print(f'*** Now starting work on {CASENAME}\n------------------------------')
+print('*** Reading variables ...')
+
+sfi = os.environ['TOS_FILE']
+pfi = os.environ['PSL_FILE']
+ufi = os.environ['UA_FILE']
+vfi = os.environ['VA_FILE']
+tfi = os.environ['TA_FILE']
+
+tos = xr.open_dataset(sfi)
+psl = xr.open_dataset(pfi)
+ua = xr.open_dataset(ufi)
+va = xr.open_dataset(vfi)
+ta = xr.open_dataset(tfi)
+
+# Compute the diagnostics (note, here we assume that all model variables are the same length in time)
+mod_firstyr = ua.time.dt.year.values[0]
+mod_lastyr = ua.time.dt.year.values[-1]
+print(mod_firstyr,mod_lastyr)
+print (FIRSTYR,LASTYR)
+
+ps = psl.sel(time=slice(str(FIRSTYR),str(LASTYR)))
+uas = ua.sel(time=slice(str(FIRSTYR),str(LASTYR))).mean('lon')
+vas = va.sel(time=slice(str(FIRSTYR),str(LASTYR)))
+tas = ta.sel(time=slice(str(FIRSTYR),str(LASTYR)))
+toss = tos.sel(time=slice(str(FIRSTYR),str(LASTYR)))
+
+print(f'***Determine whether model pressure levels are in Pa or hPa, convert to hPa')
+if getattr(uas.lev,'units') == 'Pa':
+ print(f'**Converting pressure levels to hPa')
+ uas = uas.assign_coords({"lev": (uas.lev/100.)})
+ uas.lev.attrs['units'] = 'hPa'
+ vas = vas.assign_coords({"lev": (vas.lev/100.)})
+ vas.lev.attrs['units'] = 'hPa'
+ tas = tas.assign_coords({"lev": (tas.lev/100.)})
+ tas.lev.attrs['units'] = 'hPa'
+
+if getattr(ps.psl,'units') == 'Pa':
+ print(f'**Converting pressure levels to hPa')
+ ps.psl.attrs['units'] = 'hPa'
+ ps.psl.values[:] = ps.psl.values/100.
+
+# Create the POD figures directory
+plot_dir = f'{WK_DIR}/model/'
+
+#############################################
+print ('*** Running the model ENSO indexing')
+#############################################
+
+# Extract the tropical domain #
+ENSO = toss.isel(lat = np.logical_and(toss.lat >= -5, toss.lat <= 5))
+
+# Extract date and longitude info from ENSO dataset #
+date_first = toss.time[0]
+date_last = toss.time[-1]
+lon_first = toss.lon.values[0]
+
+# Identify the correct ENSO longitudinal grid #
+if lon_first < 0:
+ ENSO = ENSO.sel(lon=slice(-170,-120))
+else:
+ ENSO = ENSO.sel(lon=slice(190,240))
+
+# Latitudinally average the ENSO data #
+weighted_mean = ENSO.mean('lat')
+weights = np.cos(np.deg2rad(ENSO.lat.values))
+interim = np.multiply(ENSO.tos.values,weights[np.newaxis,:,np.newaxis])
+interim = np.nansum(interim,axis=1)
+interim = np.true_divide(interim,np.sum(weights))
+weighted_mean.tos.values[:] = interim[:]
+
+# Subsample ENSO data for NH #
+seasonal = weighted_mean.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+seasonal = seasonal.sel(time=seasonal.time.dt.month.isin([11,12,1,2,3])).mean('lon')
+nh_nina, nh_nino = enso_index(seasonal)
+seasonal.close()
+
+# Subsample ENSO data for SH #
+seasonal = weighted_mean.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+seasonal = seasonal.sel(time=seasonal.time.dt.month.isin([9,10,11,12,1])).mean('lon')
+sh_nina, sh_nino = enso_index(seasonal)
+seasonal.close()
+
+# Store the Nina/Nino years in a dictionary to call later #
+model_enso_dict = {}
+model_enso_dict['NH'] = nh_nina, nh_nino
+model_enso_dict['SH'] = sh_nina, sh_nino
+
+##########################################################################
+# Define ENSO plotting parameters to be passed to the plotting functions #
+##########################################################################
+
+#########################################################
+print ('*** Doing the model eddy heat flux calculations')
+#########################################################
+vt = compute_total_eddy_heat_flux(vas.va,tas.ta)
+
+model_nh_enso_uzm = uas.ua.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+model_nh_enso_vtzm = vt.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+model_nh_enso_psl = ps.psl.sel(time=slice('%s-11-01' % date_first.dt.year.values, '%s-03-31' % date_last.dt.year.values))
+model_nh_enso_titles = ['November','December','January','February','March']
+model_nh_enso_plot_months = [11,12,1,2,3]
+model_nh_enso_axes = [0,90,1000,1]
+model_nh_enso_psl_axes = [-180, 180, 20, 90]
+
+model_sh_enso_uzm = uas.ua.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+model_sh_enso_vtzm = vt.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+model_sh_enso_psl = ps.psl.sel(time=slice('%s-09-01' % date_first.dt.year.values, '%s-01-31' % date_last.dt.year.values))
+model_sh_enso_titles = ['September','October','November','December','January']
+model_sh_enso_plot_months = [9,10,11,12,1]
+model_sh_enso_axes = [-90,0,1000,1]
+model_sh_enso_psl_axes = [-180, 180, -90, -20]
+
+model_uzm_dict = {}
+model_uzm_dict['NH'] = model_nh_enso_uzm, model_nh_enso_titles, model_nh_enso_plot_months, model_nh_enso_axes
+model_uzm_dict['SH'] = model_sh_enso_uzm, model_sh_enso_titles, model_sh_enso_plot_months, model_sh_enso_axes
+
+model_vtzm_dict = {}
+model_vtzm_dict['NH'] = model_nh_enso_vtzm, model_nh_enso_titles, model_nh_enso_plot_months, model_nh_enso_axes
+model_vtzm_dict['SH'] = model_sh_enso_vtzm, model_sh_enso_titles, model_sh_enso_plot_months, model_sh_enso_axes
+
+model_psl_dict = {}
+model_psl_dict['NH'] = model_nh_enso_psl, model_nh_enso_titles, model_nh_enso_plot_months, ccrs.NorthPolarStereo(), model_nh_enso_psl_axes
+model_psl_dict['SH'] = model_sh_enso_psl, model_sh_enso_titles, model_sh_enso_plot_months, ccrs.SouthPolarStereo(), model_sh_enso_psl_axes
+
+#########################################################################
+# Define QBO plotting parameters to be passed to the plotting functions #
+#########################################################################
+
+model_nh_qbo_uzm = uas.ua.sel(time=slice('%s-10-01' % date_first.dt.year.values, '%s-02-28' % date_last.dt.year.values))
+model_nh_qbo_vtzm = vt.sel(time=slice('%s-10-01' % date_first.dt.year.values, '%s-02-28' % date_last.dt.year.values))
+model_nh_qbo_psl = ps.psl.sel(time=slice('%s-10-01' % date_first.dt.year.values, '%s-02-28' % date_last.dt.year.values))
+model_nh_qbo_titles = ['October','November','December','January','February']
+model_nh_qbo_plot_months = [10,11,12,1,2]
+model_nh_qbo_axes = [0,90,1000,1]
+model_nh_qbo_psl_axes = [-180, 180, 20, 90]
+
+model_sh_qbo_uzm = uas.ua.sel(time=slice('%s-07-01' % date_first.dt.year.values, '%s-11-30' % date_last.dt.year.values))
+model_sh_qbo_vtzm = vt.sel(time=slice('%s-07-01' % date_first.dt.year.values, '%s-11-30' % date_last.dt.year.values))
+model_sh_qbo_psl = ps.psl.sel(time=slice('%s-07-01' % date_first.dt.year.values, '%s-11-30' % date_last.dt.year.values))
+model_sh_qbo_titles = ['July','August','September','October','November']
+model_sh_qbo_plot_months = [7,8,9,10,11]
+model_sh_qbo_axes = [-90,0,1000,1]
+model_sh_qbo_psl_axes = [-180, 180, -90, -20]
+
+model_uzm_qbo_dict = {}
+model_uzm_qbo_dict['NH'] = model_nh_qbo_uzm, model_nh_qbo_titles, model_nh_qbo_plot_months, model_nh_qbo_axes
+model_uzm_qbo_dict['SH'] = model_sh_qbo_uzm, model_sh_qbo_titles, model_sh_qbo_plot_months, model_sh_qbo_axes
+print (model_uzm_qbo_dict)
+
+model_vtzm_qbo_dict = {}
+model_vtzm_qbo_dict['NH'] = model_nh_qbo_vtzm, model_nh_qbo_titles, model_nh_qbo_plot_months, model_nh_qbo_axes
+model_vtzm_qbo_dict['SH'] = model_sh_qbo_vtzm, model_sh_qbo_titles, model_sh_qbo_plot_months, model_sh_qbo_axes
+print (model_vtzm_qbo_dict)
+
+model_psl_qbo_dict = {}
+model_psl_qbo_dict['NH'] = model_nh_qbo_psl, model_nh_qbo_titles, model_nh_qbo_plot_months, ccrs.NorthPolarStereo(), model_nh_qbo_psl_axes
+model_psl_qbo_dict['SH'] = model_sh_qbo_psl, model_sh_qbo_titles, model_sh_qbo_plot_months, ccrs.SouthPolarStereo(), model_sh_qbo_psl_axes
+print (model_psl_qbo_dict)
+
+
+hemispheres = ['NH','SH']
+
+for hemi in hemispheres:
+
+ ############################################
+ print ('*** Calling the model ENSO indices')
+ ############################################
+ model_nina, model_nino = model_enso_dict[hemi]
+
+ print (model_nina,'model_nina')
+ print (model_nino, 'model_nino')
+
+ ################################################################
+ print ('*** Running the model ENSO zonal mean zonal wind calcs')
+ ################################################################
+ out_plot = f'{plot_dir}/{CASENAME}-{FIRSTYR}-{LASTYR}-enso34-uzm-%s.png' % hemi
+ out_fig, out_ax = enso_uzm(model_uzm_dict[hemi][0],model_nina,model_nino,model_uzm_dict[hemi][1],model_uzm_dict[hemi][2],model_uzm_dict[hemi][3])
+ out_fig.savefig(out_plot,dpi=700)
+
+ #########################################################
+ print ('*** Running the model ENSO eddy heat flux calcs')
+ #########################################################
+ out_plot = f'{plot_dir}/{CASENAME}-{FIRSTYR}-{LASTYR}-enso34-vt-%s.png' % hemi
+ out_fig, out_ax = enso_vt(model_vtzm_dict[hemi][0],model_nina,model_nino,model_vtzm_dict[hemi][1],model_vtzm_dict[hemi][2],model_vtzm_dict[hemi][3])
+ out_fig.savefig(out_plot,dpi=700)
+
+ #######################################################
+ print ('*** Running the model ENSO sea level pressure')
+ #######################################################
+ out_plot = f'{plot_dir}/{CASENAME}-{FIRSTYR}-{LASTYR}-enso34-psl-%s.png' % hemi
+ out_fig, out_ax = enso_slp(model_psl_dict[hemi][0],model_nina,model_nino,model_psl_dict[hemi][1],model_psl_dict[hemi][2],model_psl_dict[hemi][3],model_psl_dict[hemi][4])
+ out_fig.savefig(out_plot,dpi=700)
+
+##########################################
+print('*** Running the model QBO metrics')
+##########################################
+metricsout, switch = qbo_metrics(uas,QBOisobar)
+
+filepath = f'{plot_dir}/{CASENAME}-{FIRSTYR}-{LASTYR}-qbo{QBOisobar}hPa-metrics.txt'
+with open(filepath, 'w') as file_handler:
+ file_handler.write(f"{'QBO metrics: periodicity and spatial characteristics'}\n")
+ file_handler.write(f"{' '}\n")
+ for item in metricsout:
+ file_handler.write(f"{item}\n")
+
+if switch == 1:
+
+ ###################################################################################
+ print ('*** A model QBO was detected so POD is now running the model QBO indexing')
+ ###################################################################################
+
+ print (QBOisobar, "QBOisobar")
+
+ # Subset atmospheric data for user defined isobar #
+ subset = uas.sel(lev=QBOisobar)
+
+ # Select 5S-5N winds #
+ tropical_subset = subset.interp(lat=[-5,-2.5,0,2.5,5])
+
+ # Latitudinally weight and average #
+ qbo = tropical_subset.mean('lat')
+ weights = np.cos(np.deg2rad(tropical_subset.lat.values))
+ interim = np.multiply(tropical_subset.ua.values,weights[np.newaxis,:])
+ interim = np.nansum(interim,axis=1)
+ interim = np.true_divide(interim,np.sum(weights))
+ qbo.ua.values[:] = interim[:]
+
+ def qbo_index(seasonal):
+
+ # Create 2-month seasonally averaged standardized QBO anomalies. Weight each month by number of days comprising month #
+ day_in_month_weights = seasonal.time.dt.days_in_month.values[:2]/np.sum(seasonal.time.dt.days_in_month.values[:2])
+ qboindex = np.reshape(seasonal.ua.values,(int(len(seasonal.ua.values)/2),2))
+ qboindex = np.nanmean(np.multiply(qboindex,day_in_month_weights[np.newaxis,:]),axis=1)
+ anom = np.subtract(qboindex,np.nanmean(qboindex))
+ anom = np.true_divide(anom,np.nanstd(qboindex))
+
+ # Get the unique years from "seasonal" and then remove the last one, which is not needed
+ years = [v for v in set(np.sort(seasonal.time.dt.year.values))]
+ eqbo_years = [years[i] for i,v in enumerate(anom) if v <= -1]
+ wqbo_years = [years[i] for i,v in enumerate(anom) if v >= 1]
+
+ return eqbo_years, wqbo_years
+
+ # Subsample QBO data for NH #
+ seasonal = qbo.sel(time=qbo.time.dt.month.isin([10,11]))
+ model_nh_eqbo, model_nh_wqbo = qbo_index(seasonal)
+ seasonal.close()
+
+ # Subsample QBO data for SH #
+ seasonal = qbo.sel(time=qbo.time.dt.month.isin([7,8]))
+ model_sh_eqbo, model_sh_wqbo = qbo_index(seasonal)
+ seasonal.close()
+
+ for hemi in hemispheres:
+
+ # Store the Nina/Nino years in a dictionary to call later #
+ model_qbo_dict = {}
+ model_qbo_dict['NH'] = model_nh_eqbo, model_nh_wqbo
+ model_qbo_dict['SH'] = model_sh_eqbo, model_sh_wqbo
+
+ ############################################
+ print ('*** Running the model QBO indexing')
+ ############################################
+ model_eqbo, model_wqbo = model_qbo_dict[hemi]
+
+ print (model_eqbo, 'model_eqbo')
+ print (model_wqbo, 'model_wqbo')
+
+ ##################################################################
+ print ('*** Running the model QBO zonal mean zonal wind plotting')
+ ##################################################################
+ out_plot = f'{plot_dir}/{CASENAME}-{FIRSTYR}-{LASTYR}-qbo{QBOisobar}hPa-uzm-%s.png' % hemi
+ out_fig, out_ax = qbo_uzm(model_uzm_qbo_dict[hemi][0],model_eqbo,model_wqbo,QBOisobar,model_uzm_qbo_dict[hemi][1],model_uzm_qbo_dict[hemi][2],model_uzm_qbo_dict[hemi][3])
+ out_fig.savefig(out_plot,dpi=700)
+
+ ######################################################################
+ print ('*** Running the model QBO zonal mean eddy heat flux plotting')
+ ######################################################################
+ out_plot = f'{plot_dir}/{CASENAME}-{FIRSTYR}-{LASTYR}-qbo{QBOisobar}hPa-vt-%s.png' % hemi
+ out_fig, out_ax = qbo_vt(model_vtzm_qbo_dict[hemi][0],model_eqbo,model_wqbo,QBOisobar,model_vtzm_qbo_dict[hemi][1],model_vtzm_qbo_dict[hemi][2],model_vtzm_qbo_dict[hemi][3])
+ out_fig.savefig(out_plot,dpi=700)
+
+ print ('*** Running the model QBO sea level pressure plotting')
+ out_plot = f'{plot_dir}/{CASENAME}-{FIRSTYR}-{LASTYR}-qbo{QBOisobar}hPa-psl-%s.png' % hemi
+ out_fig, out_ax = qbo_slp(model_psl_qbo_dict[hemi][0],model_eqbo,model_wqbo,QBOisobar,model_psl_qbo_dict[hemi][1],model_psl_qbo_dict[hemi][2],model_psl_qbo_dict[hemi][3],model_psl_qbo_dict[hemi][4])
+ out_fig.savefig(out_plot,dpi=700)
+
+if switch == 0:
+
+ print ("No QBO detected in the model data. As a result, QBO Ubar, v'T', ans SLP plots were not made.")
+
+###################################
+# Prepare the output dictionaries #
+###################################
+
+vt_data = {}
+uzm_data = {}
+slp_data = {}
+
+###########################
+# Saving some of the data #
+###########################
+
+vt_data['NH'] = vt.sel(lat = np.logical_and(vt.lat >= 0, vt.lat <= 90))
+vt_data['SH'] = vt.sel(lat = np.logical_and(vt.lat >= -90, vt.lat <= 0))
+vt_out = xr.concat([vt_data['SH'], vt_data['NH']], dim='hemi')
+vt_out.name = 'vt_out'
+vt_out.attrs['units'] = 'Km s**-1'
+vt_out.attrs['long_name'] = 'Pole to equator zonal-mean eddy heat flux'
+
+uzm_data['NH'] = uas.ua.sel(lat = np.logical_and(uas.lat >= 0, uas.lat <= 90))
+uzm_data['SH'] = uas.ua.sel(lat = np.logical_and(uas.lat >= -90, uas.lat <= 0))
+uzm_out = xr.concat([uzm_data['SH'], uzm_data['NH']], dim='hemi')
+uzm_out.name = 'uzm_out'
+uzm_out.attrs['units'] = 'm s**-1'
+uzm_out.attrs['long_name'] = 'Pole to equator zonal-mean zonal wind'
+
+slp_data['NH'] = ps.psl.sel(lat = np.logical_and(ps.lat >= 20, ps.lat <= 90))
+slp_data['SH'] = ps.psl.sel(lat = np.logical_and(ps.lat >= -90, ps.lat <= -20))
+slp_out = xr.concat([slp_data['SH'], slp_data['NH']], dim='hemi')
+slp_out.name = 'slp_out'
+slp_out.attrs['units'] = 'hPa'
+slp_out.attrs['long_name'] = 'Pole to 20N/S sea level pressure'
+
+qbo_out = uas.ua.interp(lat=[-5,-2.5,0,2.5,5]).sel(lev=QBOisobar)
+qbo_out.name = 'qbo_out'
+qbo_out.attrs['units'] = 'm s**-1'
+qbo_out.attrs['long_name'] = f'5S to 5N {QBOisobar} hPa zonal-mean zonal wind'
+print (qbo_out, 'qbo_out')
+
+out_ds = xr.merge([vt_out,uzm_out,slp_out,qbo_out])
+print ('OUT_DS')
+print (out_ds)
+print (' ')
+print (' ')
+
+print('*** Preparing to save derived data')
+data_dir = f'{WK_DIR}/model/netCDF'
+outfile = data_dir+f'/{CASENAME}_qbo-enso_diagnostics.nc'
+
+encoding = {'vt_out': {'dtype':'float32'},
+ 'uzm_out': {'dtype':'float32'},
+ 'slp_out': {'dtype':'float32'},
+ 'qbo_out': {'dtype':'float32'}}
+
+print(f'*** Saving qbo-enso diagnostic data to {outfile}')
+out_ds.to_netcdf(outfile, encoding=encoding)
+
+
+print('\n=====================================')
+print('END stc_qbo_enso.py ')
+print('=====================================\n')
\ No newline at end of file
diff --git a/diagnostics/stc_qbo_enso/stc_qbo_enso_plottingcodeenso.py b/diagnostics/stc_qbo_enso/stc_qbo_enso_plottingcodeenso.py
new file mode 100644
index 000000000..35f797de7
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/stc_qbo_enso_plottingcodeenso.py
@@ -0,0 +1,767 @@
+'''
+This module contains functions used in the Stratospheric QBO and ENSO POD.
+
+Contains:
+ enso_slp: plots sea level pressure response to ENSO as a function of month and ENSO phase
+ enso_uzm: plots the zonal-mean zonal wind response to ENSO as a function of month and ENSO phase
+ enso_vt: plots the zonally averaged eddy heat flux responses to the ENSO as a function of month and ENSO phase
+'''
+
+
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import matplotlib.path as mpath
+from cartopy.util import add_cyclic_point
+from scipy import stats
+
+##################################################################################################
+##################################################################################################
+##################################################################################################
+
+def enso_uzm(uzm,negative_indices,positive_indices,titles,plot_months,axes):
+
+ r""" Compute the zonal mean zonal wind response to ENSO. Anomalies are defined as
+ deviations from the seasonal cycle and composites of anomalies are made during La Nina years,
+ El Nino years, and then their difference is shown. Stippling is used to denote statistical
+ significance on the La Nina minus El Nino composites.
+
+ Parameters
+ ----------
+ uzm : xarray.DataArray
+ The zonal mean zonal wind.
+
+ negative_indices : list
+ A list of La Nina years.
+
+ positive_indices : list
+ A list of El Nino years.
+
+ titles : list of strings
+ A list of month names that will be used as the titles for each column of subplots
+
+ plot_months : list
+ A list of numbers corresponding to each month (e.g., 10 = october)
+
+ axes : A list of numbers used to set the pressure and latitude limits of the subplots.
+ [0,90,1000,1] are used for the N. hemisphere and [-90,0,1000,1] for S. hemisphere
+
+ Returns
+ -------
+ 3 row by 6 column plot of La Nina uzm anomalies (top row), El Nino uzm anomalies (middle),
+ and their difference (bottom row) with stippling highlighting differences between El Nino
+ and La Nina winds statistically significant at the 95% level using a two-sided t-test.
+
+ Notes
+ -----
+ The input field uzm is assumed to have dimensions named "lat" and "lev".
+ E.g., if your data has dimensions "latitude" and/or "level",
+ use the rename method:
+ ds.rename({'latitude':'lat','level':'lev'})
+ """
+
+ nina_out = []
+ nino_out = []
+ diff_out = []
+ sigs_out = []
+ clim_out = []
+
+ for mon in plot_months:
+
+ clim = uzm.sel(time=uzm.time.dt.month.isin([mon]))
+ clim_out.append(clim.mean('time').values)
+
+ anom = clim.groupby("time.month") - clim.groupby("time.month").mean("time")
+
+ if mon == 7 or mon == 8 or mon == 9 or mon == 10 or mon == 11 or mon == 12:
+ tmp_negative_indices = np.add(negative_indices,0)
+ tmp_positive_indices = np.add(positive_indices,0)
+ if mon == 1 or mon == 2 or mon == 3:
+ tmp_negative_indices = np.add(negative_indices,1)
+ tmp_positive_indices = np.add(positive_indices,1)
+
+ nina_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_negative_indices]))
+ nino_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_positive_indices]))
+
+ t,p = stats.ttest_ind(nina_tmp.values,nino_tmp.values,axis=0,nan_policy='omit')
+
+ sigs_out.append(np.subtract(1,p))
+ diff_out.append(np.subtract(nina_tmp.mean('time').values,nino_tmp.mean('time').values))
+ nina_out.append(nina_tmp.mean('time').values)
+ nino_out.append(nino_tmp.mean('time').values)
+
+ clim.close()
+ anom.close()
+ nina_tmp.close()
+ nino_tmp.close()
+ uzm.close()
+
+ nina_out = np.array(nina_out)
+ nino_out = np.array(nino_out)
+ diff_out = np.array(diff_out)
+ sigs_out = np.array(sigs_out)
+ clim_out = np.array(clim_out)
+
+ ############# Begin the plotting ############
+
+ fig, ax = plt.subplots()
+
+ mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
+
+ vmin = -10
+ vmax = 10
+ vlevs = np.linspace(vmin,vmax,num=21)
+ vlevs = [v for v in vlevs if v != 0]
+ ticks = [vmin,vmin/2,0,vmax/2,vmax]
+
+ cmin = -200
+ cmax = 200
+ clevs = np.linspace(cmin,cmax,num=41)
+ clevs = [v for v in clevs if v != 0]
+
+ plt.suptitle('ENSO zonal-mean zonal wind (m/s)',fontsize=12,fontweight='normal')
+
+ # Add colormap #
+ from palettable.colorbrewer.diverging import RdBu_11
+ cmap1=RdBu_11.mpl_colormap.reversed()
+
+ x, y = np.meshgrid(uzm.lat.values,uzm.lev.values)
+ uzm.close()
+
+ cols = [0,1,2,3,4]
+
+ for i in cols:
+
+ print (i)
+
+ # Nina #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(0,cols[i]))
+ plt.title('%s' % titles[i],fontsize=10,y=0.93,fontweight='normal')
+
+ cs = plt.contourf(x,y,nina_out[i],cmap=cmap1,levels=vlevs,extend="both",vmin=vmin,vmax=vmax,zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nina (%s seasons)' % int(len(negative_indices)),fontsize=10)
+
+ # Nino #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(1,cols[i]))
+
+ cs = plt.contourf(x,y,nino_out[i],cmap=cmap1,levels=vlevs,extend="both",vmin=vmin,vmax=vmax,zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nino (%s seasons)' % int(len(positive_indices)),fontsize=10)
+
+ # Diff: Nina minus Nino #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(2,cols[i]))
+
+ cs = plt.contourf(x,y,diff_out[i],cmap=cmap1,levels=vlevs,extend="both",vmin=vmin,vmax=vmax,zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ plt.xlabel('Latitude',fontsize=8,fontweight='normal')
+
+ sig_levs = [0.95,1]
+ mpl.rcParams['hatch.linewidth'] = 0.2
+ plt.contourf(x,y,sigs_out[i],colors='black',vmin=0.95,vmax=1,levels=sig_levs,hatches=['......','......'],alpha=0.0)
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nina - Nino',fontsize=10)
+
+ # Add colorbar #
+
+ cb_ax = fig.add_axes([0.365, 0.04, 0.30, 0.015])
+ cbar = fig.colorbar(cs, cax=cb_ax, ticks=ticks,orientation='horizontal')
+ cbar.ax.tick_params(labelsize=8, width=1)
+ cbar.ax.set_xticklabels(ticks,weight='normal')
+
+ plt.subplots_adjust(top=0.86,bottom=0.16,hspace=0.5,wspace=0.55,left=0.08,right=0.95)
+ return fig, ax
+
+##################################################################################################
+##################################################################################################
+##################################################################################################
+
+def enso_vt(vt,negative_indices,positive_indices,titles,plot_months,axes):
+
+ r""" Compute the zonal mean eddy heat flux response to ENSO. Anomalies are defined as
+ deviations from the seasonal cycle and composites of anomalies are made during La Nina years,
+ El Nino years, and then their difference is shown. Stippling is used to denote statistical
+ significance on the La Nina minus El Nino composites.
+
+ Parameters
+ ----------
+ vt : xarray.DataArray
+ The zonal mean eddy heat flux. This quantity is calculated using the
+ compute_total_eddy_heat_flux function given in the driver script stc_qbo_enso.py
+
+ negative_indices : list
+ A list of La Nina years.
+
+ positive_indices : list
+ A list of El Nino years.
+
+ titles : list of strings
+ A list of month names that will be used as the titles for each column of subplots
+
+ plot_months : list
+ A list of numbers corresponding to each month (e.g., 10 = october)
+
+ axes : A list of numbers used to set the pressure and latitude limits of the subplots.
+ [0,90,1000,1] are used for the N. hemisphere and [-90,0,1000,1] for S. hemisphere
+
+ Returns
+ -------
+ 3 row by 6 column plot of La Nina vt anomalies (top row), El Nino vt anomalies (middle),
+ and their difference (bottom row) with stippling highlighting differences between El Nino
+ and La Nina winds statistically significant at the 95% level using a two-sided t-test.
+
+ Notes
+ -----
+ The input field vt is assumed to have dimensions named "lat" and "lev".
+ E.g., if your data has dimensions "latitude" and/or "level",
+ use the rename method:
+ ds.rename({'latitude':'lat','level':'lev'})
+ """
+
+ nina_out = []
+ nino_out = []
+ diff_out = []
+ sigs_out = []
+ clim_out = []
+
+ for mon in plot_months:
+
+ clim = vt.sel(time=vt.time.dt.month.isin([mon]))
+ clim_out.append(clim.mean('time').values)
+
+ anom = clim.groupby("time.month") - clim.groupby("time.month").mean("time")
+
+ if mon == 7 or mon == 8 or mon == 9 or mon == 10 or mon == 11 or mon == 12:
+ tmp_negative_indices = np.add(negative_indices,0)
+ tmp_positive_indices = np.add(positive_indices,0)
+ if mon == 1 or mon == 2 or mon == 3:
+ tmp_negative_indices = np.add(negative_indices,1)
+ tmp_positive_indices = np.add(positive_indices,1)
+
+ nina_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_negative_indices]))
+ nino_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_positive_indices]))
+
+ t,p = stats.ttest_ind(nina_tmp.values,nino_tmp.values,axis=0,nan_policy='omit')
+
+ sigs_out.append(np.subtract(1,p))
+ diff_out.append(np.subtract(nina_tmp.mean('time').values,nino_tmp.mean('time').values))
+ nina_out.append(nina_tmp.mean('time').values)
+ nino_out.append(nino_tmp.mean('time').values)
+
+ clim.close()
+ anom.close()
+ nina_tmp.close()
+ nino_tmp.close()
+ vt.close()
+
+ nina_out = np.array(nina_out)
+ nino_out = np.array(nino_out)
+ diff_out = np.array(diff_out)
+ sigs_out = np.array(sigs_out)
+ clim_out = np.array(clim_out)
+
+ ############# Begin the plotting ############
+
+ fig, ax = plt.subplots()
+
+ mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
+
+ blevs = []
+
+ blevs.append(-2)
+ blevs.append(-6)
+ blevs.append(-10)
+ blevs.append(-25)
+ blevs.append(-50)
+ blevs.append(-100)
+ #blevs.append(-200)
+
+ blevs.append(2)
+ blevs.append(6)
+ blevs.append(10)
+ blevs.append(25)
+ blevs.append(50)
+ blevs.append(100)
+ #blevs.append(200)
+
+
+ blevs = np.sort(blevs)
+ print (blevs)
+
+ cmin = -200
+ cmax = 200
+ clevs = np.linspace(cmin,cmax,num=41)
+ clevs = [v for v in clevs if v != 0]
+
+ plt.suptitle('ENSO zonal-mean eddy heat flux (Km/s)',fontsize=12,fontweight='normal')
+
+ # Add colormap #
+ from palettable.colorbrewer.diverging import RdBu_11
+ cmap1=RdBu_11.mpl_colormap.reversed()
+
+ x, y = np.meshgrid(vt.lat.values,vt.lev.values)
+
+ cols = [0,1,2,3,4]
+
+ for i in cols:
+
+ print (i)
+
+ # Nina #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(0,cols[i]))
+ plt.title('%s' % titles[i],fontsize=10,y=0.93,fontweight='normal')
+
+ cs = plt.contourf(x,y,nina_out[i],blevs,norm = mpl.colors.SymLogNorm(linthresh=2,linscale=1,vmin=-100,vmax=100),cmap=cmap1,extend="both",zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nina (%s seasons)' % int(len(negative_indices)),fontsize=10)
+
+ # Nino #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(1,cols[i]))
+
+ cs = plt.contourf(x,y,nino_out[i],blevs,norm = mpl.colors.SymLogNorm(linthresh=2,linscale=1,vmin=-100,vmax=100),cmap=cmap1,extend="both",zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nino (%s seasons)' % int(len(positive_indices)),fontsize=10)
+
+ # Diff: Nina minus Nino #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(2,cols[i]))
+
+ cs = plt.contourf(x,y,diff_out[i],blevs,norm = mpl.colors.SymLogNorm(linthresh=2,linscale=1,vmin=-100,vmax=100),cmap=cmap1,extend="both",zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ plt.xlabel('Latitude',fontsize=8,fontweight='normal')
+
+ sig_levs = [0.95,1]
+ mpl.rcParams['hatch.linewidth'] = 0.2
+ plt.contourf(x,y,sigs_out[i],colors='black',vmin=0.95,vmax=1,levels=sig_levs,hatches=['......','......'],alpha=0.0)
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nina - Nino',fontsize=10)
+
+ # Add colorbar #
+
+ oticks = [-100,-50,-25,-10,-6,-2,2,6,10,25,50,100]
+ cb_ax = fig.add_axes([0.365, 0.04, 0.30, 0.015])
+ cbar = fig.colorbar(cs, cax=cb_ax, ticks=oticks,orientation='horizontal')
+ cbar.ax.tick_params(labelsize=6, width=1)
+ cbar.ax.set_xticklabels(oticks,weight='normal')
+
+ plt.subplots_adjust(top=0.86,bottom=0.16,hspace=0.5,wspace=0.55,left=0.08,right=0.95)
+ return fig, ax
+
+##################################################################################################
+##################################################################################################
+##################################################################################################
+
+
+def enso_slp(ps,negative_indices,positive_indices,titles,plot_months,projection,axes):
+
+ r""" Compute the sea level pressure response to ENSO. Anomalies are defined as
+ deviations from the seasonal cycle and composites of anomalies are made during La Nina years,
+ El Nino years, and then their difference is shown. Stippling is used to denote statistical
+ significance on the La Nina minus El Nino composites.
+
+ Parameters
+ ----------
+ ps : xarray.DataArray
+ The sea level pressure.
+
+ negative_indices : list
+ A list of La Nina years.
+
+ positive_indices : list
+ A list of El Nino years.
+
+ titles : list of strings
+ A list of month names that will be used as the titles for each column of subplots
+
+ plot_months : list
+ A list of numbers corresponding to each month (e.g., 10 = october)
+
+ projection : ccrs.NorthPolarStereo() or ccrs.SouthPolarStereo()
+
+ axes : A list of numbers used to set the longitude and latitude bounds of the subplots.
+ [-180, 180, 20, 90] are used for the N. hemisphere and [-180, 180, -90, -20] for S. hemisphere
+
+ Returns
+ -------
+ 3 row by 6 column plot of La Nina ps anomalies (top row), El Nino ps anomalies (middle),
+ and their difference (bottom row) with stippling highlighting differences between El Nino
+ and La Nina winds statistically significant at the 95% level using a two-sided t-test.
+
+ Notes
+ -----
+ The input field ps is assumed to have dimensions named "lat" and "lon".
+ E.g., if your data has dimensions "latitude" and/or "level",
+ use the rename method:
+ ds.rename({'latitude':'lat','longitude':'lon'})
+
+ The input field ps is expected to have units of hPa. Directly below, the code
+ will check to see if the units are Pa instead, and if they are, convert them to hPa.
+ """
+
+ if getattr(ps,'units') == 'Pa':
+ print(f'**Converting pressure levels to hPa')
+ ps.attrs['units'] = 'hPa'
+ ps.values[:] = ps.values/100.
+
+ print (np.nanmin(ps.values))
+ print (np.nanmedian(ps.values))
+ print (np.nanmean(ps.values))
+ print (np.nanmax(ps.values))
+
+ nina_out = []
+ nino_out = []
+ diff_out = []
+ sigs_out = []
+ clim_out = []
+
+ for mon in plot_months:
+
+ clim = ps.sel(time=ps.time.dt.month.isin([mon]))
+ clim_out.append(clim.mean('time').values)
+
+ anom = clim.groupby("time.month") - clim.groupby("time.month").mean("time")
+
+ if mon == 7 or mon == 8 or mon == 9 or mon == 10 or mon == 11 or mon == 12:
+ tmp_negative_indices = np.add(negative_indices,0)
+ tmp_positive_indices = np.add(positive_indices,0)
+ if mon == 1 or mon == 2 or mon == 3:
+ tmp_negative_indices = np.add(negative_indices,1)
+ tmp_positive_indices = np.add(positive_indices,1)
+
+ nina_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_negative_indices]))
+ nino_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_positive_indices]))
+
+ t,p = stats.ttest_ind(nina_tmp.values,nino_tmp.values,axis=0,nan_policy='omit')
+
+ sigs_out.append(np.subtract(1,p))
+ diff_out.append(np.subtract(nina_tmp.mean('time').values,nino_tmp.mean('time').values))
+ nina_out.append(nina_tmp.mean('time').values)
+ nino_out.append(nino_tmp.mean('time').values)
+
+ clim.close()
+ anom.close()
+ nina_tmp.close()
+ nino_tmp.close()
+ ps.close()
+
+ nina_out = np.array(nina_out)
+ nino_out = np.array(nino_out)
+ diff_out = np.array(diff_out)
+ sigs_out = np.array(sigs_out)
+ clim_out = np.array(clim_out)
+
+ ############# Begin the plotting ############
+
+ fig, ax = plt.subplots()
+
+ mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
+
+ vmin = -10
+ vmax = 10
+ vlevs = np.linspace(vmin,vmax,num=21)
+ vlevs = [v for v in vlevs if v != 0]
+ ticks = [vmin,vmin/2,0,vmax/2,vmax]
+
+ cmin = 900
+ cmax = 1100
+ clevs = np.linspace(cmin,cmax,num=21)
+
+ plt.suptitle('Nina - Nino sea level pressure (hPa)',fontsize=12,fontweight='normal')
+
+ # Add colormap #
+
+ from palettable.colorbrewer.diverging import RdBu_11
+ cmap1=RdBu_11.mpl_colormap.reversed()
+
+ lons = ps.lon.values
+ lats = ps.lat.values
+
+ cols = [0,1,2,3,4]
+
+ for i in cols:
+
+ print (i)
+
+ ########
+ # Nina #
+ ########
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(0,cols[i]), projection=projection)
+ ax1.set_extent(axes, ccrs.PlateCarree())
+
+ plt.title('%s' % titles[i],fontsize=10,y=0.93,fontweight='normal')
+
+ # Plot style features #
+
+ ax1.coastlines(linewidth=0.25)
+ theta = np.linspace(0, 2*np.pi, 100)
+ center, radius = [0.5, 0.5], 0.5
+ verts = np.vstack([np.sin(theta), np.cos(theta)]).T
+ circle = mpath.Path(verts * radius + center)
+ ax1.set_boundary(circle, transform=ax1.transAxes)
+ pos1 = ax1.get_position()
+ plt.title("%s" % titles[i], fontsize=10,fontweight='normal',y=0.98)
+ cyclic_z, cyclic_lon = add_cyclic_point(nina_out[i], coord=lons)
+
+ # Plot anomalies #
+
+ contourf = ax1.contourf(cyclic_lon, lats, cyclic_z,transform=ccrs.PlateCarree(),cmap=cmap1,vmin=vmin,vmax=vmax,levels=vlevs,extend='both',zorder=1)
+
+ # Overlay the climatology #
+
+ cyclic_clim, cyclic_lon = add_cyclic_point(clim_out[i], coord=lons)
+ cs = ax1.contour(cyclic_lon, lats, cyclic_clim,transform=ccrs.PlateCarree(),colors='k',linewidths=0.5,vmin=cmin,vmax=cmax,levels=clevs,extend='both',zorder=3)
+
+ plt.rc('font',weight='normal')
+ plt.clabel(cs,cs.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+ plt.rc('font',weight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nina (%s seasons)' % int(len(negative_indices)),fontsize=10)
+ ax2.spines['top'].set_visible(False)
+ ax2.spines['right'].set_visible(False)
+ ax2.spines['bottom'].set_visible(False)
+ ax2.spines['left'].set_visible(False)
+ ax2.get_yaxis().set_ticks([])
+
+ ########
+ # Nino #
+ ########
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(1,cols[i]), projection=projection)
+ ax1.set_extent(axes, ccrs.PlateCarree())
+
+ # Plot style features #
+
+ ax1.coastlines(linewidth=0.25)
+ theta = np.linspace(0, 2*np.pi, 100)
+ center, radius = [0.5, 0.5], 0.5
+ verts = np.vstack([np.sin(theta), np.cos(theta)]).T
+ circle = mpath.Path(verts * radius + center)
+ ax1.set_boundary(circle, transform=ax1.transAxes)
+ pos1 = ax1.get_position()
+ plt.title("%s" % titles[i], fontsize=10,fontweight='normal',y=0.98)
+ cyclic_z, cyclic_lon = add_cyclic_point(nino_out[i], coord=lons)
+
+ # Plot anomalies #
+
+ contourf = ax1.contourf(cyclic_lon, lats, cyclic_z,transform=ccrs.PlateCarree(),cmap=cmap1,vmin=vmin,vmax=vmax,levels=vlevs,extend='both',zorder=1)
+
+ # Overlay the climatology #
+
+ cyclic_clim, cyclic_lon = add_cyclic_point(clim_out[i], coord=lons)
+ cs = ax1.contour(cyclic_lon, lats, cyclic_clim,transform=ccrs.PlateCarree(),colors='k',linewidths=0.5,vmin=cmin,vmax=cmax,levels=clevs,extend='both',zorder=3)
+
+ plt.rc('font',weight='normal')
+ plt.clabel(cs,cs.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+ plt.rc('font',weight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nino (%s seasons)' % int(len(positive_indices)),fontsize=10)
+ ax2.spines['top'].set_visible(False)
+ ax2.spines['right'].set_visible(False)
+ ax2.spines['bottom'].set_visible(False)
+ ax2.spines['left'].set_visible(False)
+ ax2.get_yaxis().set_ticks([])
+
+ ##############
+ # Difference #
+ ##############
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(2,cols[i]), projection=projection)
+ ax1.set_extent(axes, ccrs.PlateCarree())
+
+ # Plot style features #
+
+ ax1.coastlines(linewidth=0.25)
+ theta = np.linspace(0, 2*np.pi, 100)
+ center, radius = [0.5, 0.5], 0.5
+ verts = np.vstack([np.sin(theta), np.cos(theta)]).T
+ circle = mpath.Path(verts * radius + center)
+ ax1.set_boundary(circle, transform=ax1.transAxes)
+ pos1 = ax1.get_position()
+ plt.title("%s" % titles[i], fontsize=10,fontweight='normal',y=0.98)
+ cyclic_z, cyclic_lon = add_cyclic_point(diff_out[i], coord=lons)
+
+ # Plot anomalies #
+
+ contourf = ax1.contourf(cyclic_lon, lats, cyclic_z,transform=ccrs.PlateCarree(),cmap=cmap1,vmin=vmin,vmax=vmax,levels=vlevs,extend='both',zorder=1)
+
+ # Statistical significance #
+
+ sig_levs = [0.95,1]
+ mpl.rcParams['hatch.linewidth'] = 0.2
+ cyclic_sig, cyclic_lontmp = add_cyclic_point(sigs_out[i], coord=lons)
+ ax1.contourf(cyclic_lon, lats, cyclic_sig,transform=ccrs.PlateCarree(),colors='black',vmin=0.95,vmax=1,levels=sig_levs,hatches=['......','......'],alpha=0.0,zorder=2)
+
+ # Overlay the climatology #
+
+ cyclic_clim, cyclic_lon = add_cyclic_point(clim_out[i], coord=lons)
+ cs = ax1.contour(cyclic_lon, lats, cyclic_clim,transform=ccrs.PlateCarree(),colors='k',linewidths=0.5,vmin=cmin,vmax=cmax,levels=clevs,extend='both',zorder=3)
+
+ plt.rc('font',weight='normal')
+ plt.clabel(cs,cs.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+ plt.rc('font',weight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('Nina - Nino',fontsize=10)
+ ax2.spines['top'].set_visible(False)
+ ax2.spines['right'].set_visible(False)
+ ax2.spines['bottom'].set_visible(False)
+ ax2.spines['left'].set_visible(False)
+ ax2.get_yaxis().set_ticks([])
+
+ # Add colorbar #
+
+ cb_ax = fig.add_axes([0.35, 0.05, 0.30, 0.015])
+ cbar = fig.colorbar(contourf, cax=cb_ax, ticks=ticks,orientation='horizontal')
+ cbar.ax.tick_params(labelsize=8, width=1)
+ cbar.ax.set_xticklabels(ticks,weight='normal')
+
+
+ plt.subplots_adjust(top=0.86,bottom=0.09,hspace=0.3,wspace=0.0,left=0.02,right=0.94)
+
+ return fig, ax
\ No newline at end of file
diff --git a/diagnostics/stc_qbo_enso/stc_qbo_enso_plottingcodeqbo.py b/diagnostics/stc_qbo_enso/stc_qbo_enso_plottingcodeqbo.py
new file mode 100644
index 000000000..7b311821a
--- /dev/null
+++ b/diagnostics/stc_qbo_enso/stc_qbo_enso_plottingcodeqbo.py
@@ -0,0 +1,774 @@
+'''
+This module contains functions used in the Stratospheric QBO and ENSO POD.
+
+Contains:
+ qbo_slp: plots sea level pressure response to QBO as a function of month and QBO phase
+ qbo_uzm: plots the zonal-mean zonal wind response to QBO as a function of month and QBO phase
+ qbo_vt: plots the zonally averaged eddy heat flux response to the QBO as a function of month and QBO phase
+'''
+
+import numpy as np
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import cartopy.crs as ccrs
+import matplotlib.path as mpath
+from cartopy.util import add_cyclic_point
+
+from scipy import stats
+
+##################################################################################################
+##################################################################################################
+##################################################################################################
+
+def qbo_uzm(uzm,negative_indices,positive_indices,QBOisobar,titles,plot_months,axes):
+
+ r""" Compute the zonal mean zonal wind response to the QBO. Anomalies are defined as
+ deviations from the seasonal cycle and composites of anomalies are made during easterly QBO,
+ westerly QBO, and then their difference is shown. Stippling is used to denote statistical
+ significance on the EQBO minus WQBO composites.
+
+ Parameters
+ ----------
+ uzm : xarray.DataArray
+ The zonal mean zonal wind.
+
+ negative_indices : list
+ A list of easterly QBO years.
+
+ positive_indices : list
+ A list of westerly QBO years.
+
+ QBOisobar : int
+ An integer defined by the user in the config.jsonc file specifying what isobar
+ is used to index the QBO
+
+ titles : list of strings
+ A list of month names that will be used as the titles for each column of subplots
+
+ plot_months : list
+ A list of numbers corresponding to each month (e.g., 10 = october)
+
+ axes : A list of numbers used to set the pressure and latitude limits of the subplots.
+ [0,90,1000,1] are used for the N. hemisphere and [-90,0,1000,1] for S. hemisphere
+
+ Returns
+ -------
+ 3 row by 6 column plot of easterly QBO uzm anomalies (top row), westerly QBO uzm
+ anomalies (middle), and their difference (bottom row) with stippling highlighting differences
+ between EQBO and WQBO winds statistically significant at the 95% level using a two-sided t-test.
+
+ Notes
+ -----
+ The input field uzm is assumed to have dimensions named "lat" and "lev".
+ E.g., if your data has dimensions "latitude" and/or "level",
+ use the rename method:
+ ds.rename({'latitude':'lat','level':'lev'})
+ """
+
+ eqbo_out = []
+ wqbo_out = []
+ diff_out = []
+ sigs_out = []
+ clim_out = []
+
+ for mon in plot_months:
+
+ clim = uzm.sel(time=uzm.time.dt.month.isin([mon]))
+ clim_out.append(clim.mean('time').values)
+
+ anom = clim.groupby("time.month") - clim.groupby("time.month").mean("time")
+
+ if mon == 7 or mon == 8 or mon == 9 or mon == 10 or mon == 11 or mon == 12:
+ tmp_negative_indices = np.add(negative_indices,0)
+ tmp_positive_indices = np.add(positive_indices,0)
+ if mon == 1 or mon == 2 or mon == 3:
+ tmp_negative_indices = np.add(negative_indices,1)
+ tmp_positive_indices = np.add(positive_indices,1)
+
+ eqbo_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_negative_indices]))
+ wqbo_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_positive_indices]))
+
+ t,p = stats.ttest_ind(eqbo_tmp.values,wqbo_tmp.values,axis=0,nan_policy='omit')
+
+ sigs_out.append(np.subtract(1,p))
+ diff_out.append(np.subtract(eqbo_tmp.mean('time').values,wqbo_tmp.mean('time').values))
+ eqbo_out.append(eqbo_tmp.mean('time').values)
+ wqbo_out.append(wqbo_tmp.mean('time').values)
+
+ clim.close()
+ anom.close()
+ eqbo_tmp.close()
+ wqbo_tmp.close()
+ uzm.close()
+
+ eqbo_out = np.array(eqbo_out)
+ wqbo_out = np.array(wqbo_out)
+ diff_out = np.array(diff_out)
+ sigs_out = np.array(sigs_out)
+ clim_out = np.array(clim_out)
+
+ ############# Begin the plotting ############
+
+ fig, ax = plt.subplots()
+
+ mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
+
+ vmin = -10
+ vmax = 10
+ vlevs = np.linspace(vmin,vmax,num=21)
+ vlevs = [v for v in vlevs if v != 0]
+ ticks = [vmin,vmin/2,0,vmax/2,vmax]
+
+ cmin = -200
+ cmax = 200
+ clevs = np.linspace(cmin,cmax,num=41)
+ clevs = [v for v in clevs if v != 0]
+
+ plt.suptitle('QBO (5S-5N index @ %s hPa) zonal-mean zonal wind (m/s)' % int(QBOisobar),fontsize=12,fontweight='normal')
+
+ # Add colormap #
+ from palettable.colorbrewer.diverging import RdBu_11
+ cmap1=RdBu_11.mpl_colormap.reversed()
+
+ x, y = np.meshgrid(uzm.lat.values,uzm.lev.values)
+
+ cols = [0,1,2,3,4]
+
+ for i in cols:
+
+ # eqbo #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(0,cols[i]))
+ plt.title('%s' % titles[i],fontsize=10,y=0.93,fontweight='normal')
+
+ cs = plt.contourf(x,y,eqbo_out[i],cmap=cmap1,levels=vlevs,extend="both",vmin=vmin,vmax=vmax,zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,5,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('eqbo (%s seasons)' % int(len(negative_indices)),fontsize=10)
+
+ # wqbo #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(1,cols[i]))
+
+ cs = plt.contourf(x,y,wqbo_out[i],cmap=cmap1,levels=vlevs,extend="both",vmin=vmin,vmax=vmax,zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,5,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('wqbo (%s seasons)' % int(len(positive_indices)),fontsize=10)
+
+ # Diff: eqbo minus wqbo #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(2,cols[i]))
+
+ cs = plt.contourf(x,y,diff_out[i],cmap=cmap1,levels=vlevs,extend="both",vmin=vmin,vmax=vmax,zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,5,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ plt.xlabel('Latitude',fontsize=8,fontweight='normal')
+
+ sig_levs = [0.95,1]
+ mpl.rcParams['hatch.linewidth'] = 0.2
+ plt.contourf(x,y,sigs_out[i],colors='black',vmin=0.95,vmax=1,levels=sig_levs,hatches=['......','......'],alpha=0.0)
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('eqbo - wqbo',fontsize=10)
+
+ # Add colorbar #
+
+ cb_ax = fig.add_axes([0.365, 0.04, 0.30, 0.015])
+ cbar = fig.colorbar(cs, cax=cb_ax, ticks=ticks,orientation='horizontal')
+ cbar.ax.tick_params(labelsize=8, width=1)
+ cbar.ax.set_xticklabels(ticks,weight='normal')
+
+ plt.subplots_adjust(top=0.86,bottom=0.16,hspace=0.5,wspace=0.55,left=0.08,right=0.95)
+ return fig, ax
+
+##################################################################################################
+##################################################################################################
+##################################################################################################
+
+def qbo_vt(vt,negative_indices,positive_indices,QBOisobar,titles,plot_months,axes):
+
+ r""" Compute the zonal mean eddy heat flux response to the QBO. Anomalies are defined as
+ deviations from the seasonal cycle and composites of anomalies are made during easterly QBO,
+ westerly QBO, and then their difference is shown. Stippling is used to denote statistical
+ significance on the EQBO minus WQBO composites.
+
+ Parameters
+ ----------
+ vt : xarray.DataArray
+ The zonal mean eddy heat flux. This quantity is calculated using the
+ compute_total_eddy_heat_flux function given in the driver script stc_qbo_enso.py
+
+ negative_indices : list
+ A list of easterly QBO years.
+
+ positive_indices : list
+ A list of westerly QBO years.
+
+ QBOisobar : int
+ An integer defined by the user in the config.jsonc file specifying what isobar
+ is used to index the QBO
+
+ titles : list of strings
+ A list of month names that will be used as the titles for each column of subplots
+
+ plot_months : list
+ A list of numbers corresponding to each month (e.g., 10 = october)
+
+ axes : A list of numbers used to set the pressure and latitude limits of the subplots.
+ [0,90,1000,1] are used for the N. hemisphere and [-90,0,1000,1] for S. hemisphere
+
+ Returns
+ -------
+ 3 row by 6 column plot of easterly QBO vt anomalies (top row), westerly QBO vt
+ anomalies (middle), and their difference (bottom row) with stippling highlighting differences
+ between EQBO and WQBO winds statistically significant at the 95% level using a two-sided t-test.
+
+ Notes
+ -----
+ The input field vt is assumed to have dimensions named "lat" and "lev".
+ E.g., if your data has dimensions "latitude" and/or "level",
+ use the rename method:
+ ds.rename({'latitude':'lat','level':'lev'})
+ """
+
+ eqbo_out = []
+ wqbo_out = []
+ diff_out = []
+ sigs_out = []
+ clim_out = []
+
+ for mon in plot_months:
+
+ clim = vt.sel(time=vt.time.dt.month.isin([mon]))
+ clim_out.append(clim.mean('time').values)
+
+ anom = clim.groupby("time.month") - clim.groupby("time.month").mean("time")
+
+ # For NH, QBO index is based on Oct-Nov year. The plot will show Oct-Feb. Note that Jan-Feb are selected using QBO index year + 1
+ # For SH, QBO index is based on Jul-Aug year. The plot will show Jul-Nov.
+
+ if mon == 7 or mon == 8 or mon == 9 or mon == 10 or mon == 11 or mon == 12:
+ tmp_negative_indices = np.add(negative_indices,0)
+ tmp_positive_indices = np.add(positive_indices,0)
+ if mon == 1 or mon == 2 or mon == 3:
+ tmp_negative_indices = np.add(negative_indices,1)
+ tmp_positive_indices = np.add(positive_indices,1)
+
+ eqbo_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_negative_indices]))
+ wqbo_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_positive_indices]))
+
+ t,p = stats.ttest_ind(eqbo_tmp.values,wqbo_tmp.values,axis=0,nan_policy='omit')
+
+ sigs_out.append(np.subtract(1,p))
+ diff_out.append(np.subtract(eqbo_tmp.mean('time').values,wqbo_tmp.mean('time').values))
+ eqbo_out.append(eqbo_tmp.mean('time').values)
+ wqbo_out.append(wqbo_tmp.mean('time').values)
+
+ clim.close()
+ anom.close()
+ eqbo_tmp.close()
+ wqbo_tmp.close()
+ vt.close()
+
+ eqbo_out = np.array(eqbo_out)
+ wqbo_out = np.array(wqbo_out)
+ diff_out = np.array(diff_out)
+ sigs_out = np.array(sigs_out)
+ clim_out = np.array(clim_out)
+
+ ############# Begin the plotting ############
+
+ fig, ax = plt.subplots()
+
+ mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
+
+ blevs = []
+
+ blevs.append(-2)
+ blevs.append(-6)
+ blevs.append(-10)
+ blevs.append(-25)
+ blevs.append(-50)
+ blevs.append(-100)
+
+ blevs.append(2)
+ blevs.append(6)
+ blevs.append(10)
+ blevs.append(25)
+ blevs.append(50)
+ blevs.append(100)
+
+ blevs = np.sort(blevs)
+ print (blevs)
+
+ cmin = -200
+ cmax = 200
+ clevs = np.linspace(cmin,cmax,num=41)
+ clevs = [v for v in clevs if v != 0]
+
+ plt.suptitle('QBO (5S-5N index @ %s hPa) zonal-mean eddy heat flux (Km/s)' % int(QBOisobar),fontsize=12,fontweight='normal')
+
+ # Add colormap #
+ from palettable.colorbrewer.diverging import RdBu_11
+ cmap1=RdBu_11.mpl_colormap.reversed()
+
+ x, y = np.meshgrid(vt.lat.values,vt.lev.values)
+
+ cols = [0,1,2,3,4]
+
+ for i in cols:
+
+ print (i)
+
+ # eqbo #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(0,cols[i]))
+ plt.title('%s' % titles[i],fontsize=10,y=0.93,fontweight='normal')
+
+ cs = plt.contourf(x,y,eqbo_out[i],blevs,norm = mpl.colors.SymLogNorm(linthresh=2,linscale=1,vmin=-100,vmax=100),cmap=cmap1,extend="both",zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,5,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('eqbo (%s seasons)' % int(len(negative_indices)),fontsize=10)
+
+ # wqbo #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(1,cols[i]))
+
+ cs = plt.contourf(x,y,wqbo_out[i],blevs,norm = mpl.colors.SymLogNorm(linthresh=2,linscale=1,vmin=-100,vmax=100),cmap=cmap1,extend="both",zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,5,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('wqbo (%s seasons)' % int(len(positive_indices)),fontsize=10)
+
+ # Diff: eqbo minus wqbo #
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(2,cols[i]))
+
+ cs = plt.contourf(x,y,diff_out[i],blevs,norm = mpl.colors.SymLogNorm(linthresh=2,linscale=1,vmin=-100,vmax=100),cmap=cmap1,extend="both",zorder=1)
+
+ mpl.rcParams["lines.linewidth"] = 0.2
+ mpl.rcParams["lines.dashed_pattern"] = 10, 3
+ black = plt.contour(x,y,clim_out[i],colors='k',levels=clevs,extend="both",vmin=cmin,vmax=cmax,zorder=3)
+ plt.clabel(black,black.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+
+ plt.semilogy()
+ yticks = [1,5,10,100,1000]
+ plt.yticks(yticks,yticks,fontsize=6,fontweight='normal')
+ xticks = [-90,-60,-30,0,30,60,90]
+ plt.xticks(xticks,xticks,fontsize=6,fontweight='normal')
+ plt.gca().invert_yaxis()
+ plt.axis(axes)
+ if i == 0:
+ plt.ylabel('Pressure (hPa)',fontsize=8,fontweight='normal')
+ plt.xlabel('Latitude',fontsize=8,fontweight='normal')
+
+ sig_levs = [0.95,1]
+ mpl.rcParams['hatch.linewidth'] = 0.2
+ plt.contourf(x,y,sigs_out[i],colors='black',vmin=0.95,vmax=1,levels=sig_levs,hatches=['......','......'],alpha=0.0)
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('eqbo - wqbo',fontsize=10)
+
+ # Add colorbar #
+
+ oticks = [-100,-50,-25,-10,-6,-2,2,6,10,25,50,100]
+ cb_ax = fig.add_axes([0.365, 0.04, 0.30, 0.015])
+ cbar = fig.colorbar(cs, cax=cb_ax, ticks=oticks,orientation='horizontal')
+ cbar.ax.tick_params(labelsize=8, width=1)
+ cbar.ax.set_xticklabels(oticks,weight='normal')
+
+ plt.subplots_adjust(top=0.86,bottom=0.16,hspace=0.5,wspace=0.55,left=0.08,right=0.95)
+
+ return fig, ax
+##################################################################################################
+##################################################################################################
+##################################################################################################
+
+def qbo_slp(ps,negative_indices,positive_indices,QBOisobar,titles,plot_months,projection,axes):
+
+ r""" Compute the sea level pressure response to QBO. Anomalies are defined as
+ deviations from the seasonal cycle and composites of anomalies are made during easterly QBO years,
+ westerly QBO years, and then their difference is shown. Stippling is used to denote statistical
+ significance on the EQBO minus WQBO composites.
+
+ Parameters
+ ----------
+ ps : xarray.DataArray
+ The sea level pressure.
+
+ negative_indices : list
+ A list of easterly QBO years.
+
+ positive_indices : list
+ A list of westerly QBO years.
+
+ QBOisobar : int
+ An integer defined by the user in the config.jsonc file specifying what isobar
+ is used to index the QBO
+
+ titles : list of strings
+ A list of month names that will be used as the titles for each column of subplots
+
+ plot_months : list
+ A list of numbers corresponding to each month (e.g., 10 = october)
+
+ projection : ccrs.NorthPolarStereo() or ccrs.SouthPolarStereo()
+
+ axes : A list of numbers used to set the longitude and latitude bounds of the subplots.
+ [-180, 180, 20, 90] are used for the N. hemisphere and [-180, 180, -90, -20] for S. hemisphere
+
+ Returns
+ -------
+ 3 row by 6 column plot of easterly QBO ps anomalies (top row), westerly QBO ps anomalies (middle),
+ and their difference (bottom row) with stippling highlighting differences between EQBO
+ and WQBO winds statistically significant at the 95% level using a two-sided t-test.
+
+ Notes
+ -----
+ The input field ps is assumed to have dimensions named "lat" and "lon".
+ E.g., if your data has dimensions "latitude" and/or "level",
+ use the rename method:
+ ds.rename({'latitude':'lat','longitude':'lon'})
+
+ The input field ps is expected to have units of hPa. Directly below, the code
+ will check to see if the units are Pa instead, and if they are, convert them to hPa.
+ """
+
+ if getattr(ps,'units') == 'Pa':
+ print(f'**Converting pressure levels to hPa')
+ ps.attrs['units'] = 'hPa'
+ ps.values[:] = ps.values/100.
+
+ print (np.nanmin(ps.values))
+ print (np.nanmedian(ps.values))
+ print (np.nanmean(ps.values))
+ print (np.nanmax(ps.values))
+
+ eqbo_out = []
+ wqbo_out = []
+ diff_out = []
+ sigs_out = []
+ clim_out = []
+
+ for mon in plot_months:
+
+ clim = ps.sel(time=ps.time.dt.month.isin([mon]))
+ clim_out.append(clim.mean('time').values)
+
+ anom = clim.groupby("time.month") - clim.groupby("time.month").mean("time")
+
+ if mon == 7 or mon == 8 or mon == 9 or mon == 10 or mon == 11 or mon == 12:
+ tmp_negative_indices = np.add(negative_indices,0)
+ tmp_positive_indices = np.add(positive_indices,0)
+ if mon == 1 or mon == 2 or mon == 3:
+ tmp_negative_indices = np.add(negative_indices,1)
+ tmp_positive_indices = np.add(positive_indices,1)
+
+ eqbo_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_negative_indices]))
+ wqbo_tmp = anom.sel(time=anom.time.dt.year.isin([tmp_positive_indices]))
+
+ t,p = stats.ttest_ind(eqbo_tmp.values,wqbo_tmp.values,axis=0,nan_policy='omit')
+
+ sigs_out.append(np.subtract(1,p))
+ diff_out.append(np.subtract(eqbo_tmp.mean('time').values,wqbo_tmp.mean('time').values))
+ eqbo_out.append(eqbo_tmp.mean('time').values)
+ wqbo_out.append(wqbo_tmp.mean('time').values)
+
+ clim.close()
+ anom.close()
+ eqbo_tmp.close()
+ wqbo_tmp.close()
+ ps.close()
+
+ eqbo_out = np.array(eqbo_out)
+ wqbo_out = np.array(wqbo_out)
+ diff_out = np.array(diff_out)
+ sigs_out = np.array(sigs_out)
+ clim_out = np.array(clim_out)
+
+ ############# Begin the plotting ############
+
+ fig, ax = plt.subplots()
+
+ mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
+
+ vmin = -10
+ vmax = 10
+ vlevs = np.linspace(vmin,vmax,num=21)
+ vlevs = [v for v in vlevs if v != 0]
+ ticks = [vmin,vmin/2,0,vmax/2,vmax]
+
+ cmin = 900
+ cmax = 1100
+ clevs = np.linspace(cmin,cmax,num=21)
+
+ plt.suptitle('QBO (5S-5N index @ %s hPa) sea level pressure (hPa)' % QBOisobar,fontsize=12,fontweight='normal')
+
+ # Add colormap #
+
+ from palettable.colorbrewer.diverging import RdBu_11
+ cmap1=RdBu_11.mpl_colormap.reversed()
+
+ lons = ps.lon.values
+ lats = ps.lat.values
+
+ cols = [0,1,2,3,4]
+
+ for i in cols:
+
+ print (i)
+
+ ########
+ # eqbo #
+ ########
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(0,cols[i]), projection=projection)
+ ax1.set_extent(axes, ccrs.PlateCarree())
+ plt.title('%s' % titles[i],fontsize=10,y=0.93,fontweight='normal')
+
+ # Plot style features #
+
+ ax1.coastlines(linewidth=0.25)
+ theta = np.linspace(0, 2*np.pi, 100)
+ center, radius = [0.5, 0.5], 0.5
+ verts = np.vstack([np.sin(theta), np.cos(theta)]).T
+ circle = mpath.Path(verts * radius + center)
+ ax1.set_boundary(circle, transform=ax1.transAxes)
+ pos1 = ax1.get_position()
+ plt.title("%s" % titles[i], fontsize=10,fontweight='normal',y=0.98)
+ cyclic_z, cyclic_lon = add_cyclic_point(eqbo_out[i], coord=lons)
+
+ # Plot anomalies #
+
+ contourf = ax1.contourf(cyclic_lon, lats, cyclic_z,transform=ccrs.PlateCarree(),cmap=cmap1,vmin=vmin,vmax=vmax,levels=vlevs,extend='both',zorder=1)
+
+ # Overlay the climatology #
+
+ cyclic_clim, cyclic_lon = add_cyclic_point(clim_out[i], coord=lons)
+ cs = ax1.contour(cyclic_lon, lats, cyclic_clim,transform=ccrs.PlateCarree(),colors='k',linewidths=0.5,vmin=cmin,vmax=cmax,levels=clevs,extend='both',zorder=3)
+
+ plt.rc('font',weight='normal')
+ plt.clabel(cs,cs.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+ plt.rc('font',weight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('eqbo (%s seasons)' % int(len(negative_indices)),fontsize=10)
+ ax2.spines['top'].set_visible(False)
+ ax2.spines['right'].set_visible(False)
+ ax2.spines['bottom'].set_visible(False)
+ ax2.spines['left'].set_visible(False)
+ ax2.get_yaxis().set_ticks([])
+
+ ########
+ # wqbo #
+ ########
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(1,cols[i]), projection=projection)
+ ax1.set_extent(axes, ccrs.PlateCarree())
+
+ # Plot style features #
+
+ ax1.coastlines(linewidth=0.25)
+ theta = np.linspace(0, 2*np.pi, 100)
+ center, radius = [0.5, 0.5], 0.5
+ verts = np.vstack([np.sin(theta), np.cos(theta)]).T
+ circle = mpath.Path(verts * radius + center)
+ ax1.set_boundary(circle, transform=ax1.transAxes)
+ pos1 = ax1.get_position()
+ plt.title("%s" % titles[i], fontsize=10,fontweight='normal',y=0.98)
+ cyclic_z, cyclic_lon = add_cyclic_point(wqbo_out[i], coord=lons)
+
+ # Plot anomalies #
+
+ contourf = ax1.contourf(cyclic_lon, lats, cyclic_z,transform=ccrs.PlateCarree(),cmap=cmap1,vmin=vmin,vmax=vmax,levels=vlevs,extend='both',zorder=1)
+
+ # Overlay the climatology #
+
+ cyclic_clim, cyclic_lon = add_cyclic_point(clim_out[i], coord=lons)
+ cs = ax1.contour(cyclic_lon, lats, cyclic_clim,transform=ccrs.PlateCarree(),colors='k',linewidths=0.5,vmin=cmin,vmax=cmax,levels=clevs,extend='both',zorder=3)
+
+ plt.rc('font',weight='normal')
+ plt.clabel(cs,cs.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+ plt.rc('font',weight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('wqbo (%s seasons)' % int(len(positive_indices)),fontsize=10)
+ ax2.spines['top'].set_visible(False)
+ ax2.spines['right'].set_visible(False)
+ ax2.spines['bottom'].set_visible(False)
+ ax2.spines['left'].set_visible(False)
+ ax2.get_yaxis().set_ticks([])
+
+ ##############
+ # Difference #
+ ##############
+
+ ax1 = plt.subplot2grid(shape=(3,5), loc=(2,cols[i]), projection=projection)
+ ax1.set_extent(axes, ccrs.PlateCarree())
+
+ # Plot style features #
+
+ ax1.coastlines(linewidth=0.25)
+ theta = np.linspace(0, 2*np.pi, 100)
+ center, radius = [0.5, 0.5], 0.5
+ verts = np.vstack([np.sin(theta), np.cos(theta)]).T
+ circle = mpath.Path(verts * radius + center)
+ ax1.set_boundary(circle, transform=ax1.transAxes)
+ pos1 = ax1.get_position()
+ plt.title("%s" % titles[i], fontsize=10,fontweight='normal',y=0.98)
+ cyclic_z, cyclic_lon = add_cyclic_point(diff_out[i], coord=lons)
+
+ # Plot anomalies #
+
+ contourf = ax1.contourf(cyclic_lon, lats, cyclic_z,transform=ccrs.PlateCarree(),cmap=cmap1,vmin=vmin,vmax=vmax,levels=vlevs,extend='both',zorder=1)
+
+ # Statistical significance #
+
+ sig_levs = [0.95,1]
+ mpl.rcParams['hatch.linewidth'] = 0.2
+ cyclic_sig, cyclic_lontmp = add_cyclic_point(sigs_out[i], coord=lons)
+ ax1.contourf(cyclic_lon, lats, cyclic_sig,transform=ccrs.PlateCarree(),colors='black',vmin=0.95,vmax=1,levels=sig_levs,hatches=['......','......'],alpha=0.0,zorder=2)
+
+ # Overlay the climatology #
+
+ cyclic_clim, cyclic_lon = add_cyclic_point(clim_out[i], coord=lons)
+ cs = ax1.contour(cyclic_lon, lats, cyclic_clim,transform=ccrs.PlateCarree(),colors='k',linewidths=0.5,vmin=cmin,vmax=cmax,levels=clevs,extend='both',zorder=3)
+
+ plt.rc('font',weight='normal')
+ plt.clabel(cs,cs.levels[:],inline=1,fmt='%1.0f',fontsize=4,colors='k',inline_spacing=1)
+ plt.rc('font',weight='normal')
+
+ if i == 4:
+ ax2 = ax1.twinx()
+ yticks = [0,0.5,1.0]
+ ylabels = ['','','']
+ ax2.set_yticks(yticks)
+ ax2.set_yticklabels(ylabels, fontsize=8, fontweight = 'normal')
+ ax2.set_ylabel('eqbo - wqbo',fontsize=10)
+ ax2.spines['top'].set_visible(False)
+ ax2.spines['right'].set_visible(False)
+ ax2.spines['bottom'].set_visible(False)
+ ax2.spines['left'].set_visible(False)
+ ax2.get_yaxis().set_ticks([])
+
+ # Add colorbar #
+
+ cb_ax = fig.add_axes([0.35, 0.05, 0.30, 0.015])
+ cbar = fig.colorbar(contourf, cax=cb_ax, ticks=ticks,orientation='horizontal')
+ cbar.ax.tick_params(labelsize=8, width=1)
+ cbar.ax.set_xticklabels(ticks,weight='normal')
+
+
+ plt.subplots_adjust(top=0.86,bottom=0.09,hspace=0.3,wspace=0.0,left=0.02,right=0.94)
+
+ return fig, ax
\ No newline at end of file