diff --git a/.vscode/e3sm_diags.code-workspace b/.vscode/e3sm_diags.code-workspace index c7d968567..7cf614852 100644 --- a/.vscode/e3sm_diags.code-workspace +++ b/.vscode/e3sm_diags.code-workspace @@ -58,7 +58,7 @@ "configurations": [ { "name": "Python: Current File", - "type": "python", + "type": "debugpy", "request": "launch", "program": "${file}", "console": "integratedTerminal", diff --git a/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_regression_test_netcdf.ipynb b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_regression_test_netcdf.ipynb new file mode 100644 index 000000000..4157e1d8a --- /dev/null +++ b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_regression_test_netcdf.ipynb @@ -0,0 +1,520 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# CDAT Migration Regression Testing Notebook (`.nc` files)\n", + "\n", + "This notebook is used to perform regression testing between the development and\n", + "production versions of a diagnostic set.\n", + "\n", + "## How it works\n", + "\n", + "It compares the relative differences (%) between ref and test variables between\n", + "the dev and `main` branches.\n", + "\n", + "## How to use\n", + "\n", + "PREREQUISITE: The diagnostic set's netCDF stored in `.json` files in two directories\n", + "(dev and `main` branches).\n", + "\n", + "1. Make a copy of this notebook under `auxiliary_tools/cdat_regression_testing/`.\n", + "2. Run `mamba create -n cdat_regression_test -y -c conda-forge \"python<3.12\" xarray dask pandas matplotlib-base ipykernel`\n", + "3. Run `mamba activate cdat_regression_test`\n", + "4. Update `SET_DIR` and `SET_NAME` in the copy of your notebook.\n", + "5. Run all cells IN ORDER.\n", + "6. Review results for any outstanding differences (>=1e-5 relative tolerance).\n", + " - Debug these differences (e.g., bug in metrics functions, incorrect variable references, etc.)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup Code\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "import glob\n", + "\n", + "import numpy as np\n", + "import xarray as xr\n", + "\n", + "# TODO: Update SET_NAME and SET_DIR\n", + "SET_NAME = \"zonal_mean_2d\"\n", + "SET_DIR = \"655-zonal-mean-2d\"\n", + "\n", + "DEV_PATH = f\"/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/{SET_DIR}/{SET_NAME}/**\"\n", + "MAIN_PATH = f\"/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/{SET_NAME}/**\"\n", + "\n", + "DEV_GLOB = sorted(glob.glob(DEV_PATH + \"/*.nc\"))\n", + "MAIN_GLOB = sorted(glob.glob(MAIN_PATH + \"/*.nc\"))\n", + "\n", + "if len(DEV_GLOB) == 0 or len(MAIN_GLOB) == 0:\n", + " raise IOError(\"No files found at DEV_PATH and/or MAIN_PATH.\")\n", + "\n", + "dev_num_files = len(DEV_GLOB)\n", + "main_num_files = len(MAIN_GLOB)\n", + "if dev_num_files != main_num_files:\n", + " raise IOError(\n", + " f\"Number of files do not match at DEV_PATH ({dev_num_files}) and MAIN_PATH ({main_num_files}).\"\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "def _get_var_to_filepath_map():\n", + " var_to_file = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))\n", + "\n", + " for dev_file, main_file in zip(DEV_GLOB, MAIN_GLOB):\n", + " if \"relative difference\" in dev_file:\n", + " continue\n", + "\n", + " # Example:\n", + " # \"/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/660-cosp-histogram/cosp_histogram/ISCCP-COSP/ISCCPCOSP-COSP_HISTOGRAM_ISCCP-ANN-global_test.nc\"\n", + " file_arr = dev_file.split(\"/\")\n", + "\n", + " # Example: \"test\"\n", + " data_type = dev_file.split(\"_\")[-1].split(\".nc\")[0]\n", + "\n", + " # Skip comparing `.nc` \"diff\" files because comparing relative diffs of\n", + " # does not make sense.\n", + " if data_type == \"test\" or data_type == \"ref\":\n", + " # Example: [ERA5, OMEGA, ANN, global_ref.nc]\n", + " filename = file_arr[-1].split(\"-\")\n", + " # Example: ERA5\n", + " model = filename[0]\n", + " # Example: OMEGA\n", + " var_key = filename[1]\n", + "\n", + " season = \"JJA\" if \"JJA\" in dev_file else \"ANN\"\n", + "\n", + " var_to_file[model][var_key][data_type][season] = (dev_file, main_file)\n", + "\n", + " return var_to_file\n", + "\n", + "\n", + "def _get_relative_diffs(var_to_filepath):\n", + " # Absolute tolerance of 0 and relative tolerance of 1e-5.\n", + " # We are mainly focusing on relative tolerance here (in percentage terms).\n", + " atol = 0\n", + " rtol = 1e-5\n", + "\n", + " for _, var_keys in var_to_filepath.items():\n", + " for var_key, data_types in var_keys.items():\n", + " for _, seasons in data_types.items():\n", + " for _, filepaths in seasons.items():\n", + " print(\"Comparing:\")\n", + " print(filepaths[0], \"\\n\", filepaths[1])\n", + " ds1 = xr.open_dataset(filepaths[0])\n", + " ds2 = xr.open_dataset(filepaths[1])\n", + "\n", + " try:\n", + " np.testing.assert_allclose(\n", + " ds1[var_key].values,\n", + " ds2[var_key].values,\n", + " atol=atol,\n", + " rtol=rtol,\n", + " )\n", + " except AssertionError as e:\n", + " print(e)\n", + " else:\n", + " print(f\" * All close and within relative tolerance ({rtol})\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Compare the netCDF files between branches\n", + "\n", + "- Compare \"ref\" and \"test\" files\n", + "- \"diff\" files are ignored because getting relative diffs for these does not make sense (relative diff will be above tolerance)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "var_to_filepaths = _get_var_to_filepath_map()" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "defaultdict(.()>,\n", + " {'ERA5': defaultdict(...()>,\n", + " {'OMEGA': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_test.nc')}}),\n", + " 'Q': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_test.nc')}}),\n", + " 'RELHUM': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_test.nc')}}),\n", + " 'T': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-JJA-global_test.nc')}}),\n", + " 'U': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-JJA-global_test.nc')}})}),\n", + " 'MERRA2': defaultdict(...()>,\n", + " {'OMEGA': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_test.nc')}}),\n", + " 'Q': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_test.nc')}}),\n", + " 'RELHUM': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_test.nc')}}),\n", + " 'T': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_test.nc')}}),\n", + " 'U': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_test.nc')}})})})" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "var_to_filepaths" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-ANN-global_test.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 3 / 3600 (0.0833%)\n", + "Max absolute difference: 6.00982503e-06\n", + "Max relative difference: 2.63548592e-05\n", + " x: array([[-0.224216, -0.225982, -0.230203, ..., -0.535621, -0.538314,\n", + " -0.539618],\n", + " [-0.291696, -0.292976, -0.298023, ..., -0.461387, -0.462255,...\n", + " y: array([[-0.224216, -0.225982, -0.230203, ..., -0.535621, -0.538314,\n", + " -0.539618],\n", + " [-0.291696, -0.292976, -0.298023, ..., -0.461387, -0.462255,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-OMEGA-JJA-global_test.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 3600 (0.0278%)\n", + "Max absolute difference: 4.69767184e-06\n", + "Max relative difference: 1.44081722e-05\n", + " x: array([[-3.842963e-02, -3.952561e-02, -4.294728e-02, ..., -6.250532e-02,\n", + " -6.345372e-02, -6.394634e-02],\n", + " [-2.693705e-01, -2.642957e-01, -2.585470e-01, ..., -7.585498e-04,...\n", + " y: array([[-3.842963e-02, -3.952561e-02, -4.294728e-02, ..., -6.250532e-02,\n", + " -6.345372e-02, -6.394633e-02],\n", + " [-2.693705e-01, -2.642957e-01, -2.585470e-01, ..., -7.585498e-04,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-Q-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-RELHUM-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-T-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-T-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/ERA5/ERA5-U-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/ERA5/ERA5-U-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_test.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 3 / 3600 (0.0833%)\n", + "Max absolute difference: 6.00982503e-06\n", + "Max relative difference: 2.63548592e-05\n", + " x: array([[-0.224216, -0.225982, -0.230203, ..., -0.535621, -0.538314,\n", + " -0.539618],\n", + " [-0.291696, -0.292976, -0.298023, ..., -0.461387, -0.462255,...\n", + " y: array([[-0.224216, -0.225982, -0.230203, ..., -0.535621, -0.538314,\n", + " -0.539618],\n", + " [-0.291696, -0.292976, -0.298023, ..., -0.461387, -0.462255,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_test.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 3600 (0.0278%)\n", + "Max absolute difference: 4.69767184e-06\n", + "Max relative difference: 1.44081722e-05\n", + " x: array([[-3.842963e-02, -3.952561e-02, -4.294728e-02, ..., -6.250532e-02,\n", + " -6.345372e-02, -6.394634e-02],\n", + " [-2.693705e-01, -2.642957e-01, -2.585470e-01, ..., -7.585498e-04,...\n", + " y: array([[-3.842963e-02, -3.952561e-02, -4.294728e-02, ..., -6.250532e-02,\n", + " -6.345372e-02, -6.394633e-02],\n", + " [-2.693705e-01, -2.642957e-01, -2.585470e-01, ..., -7.585498e-04,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-Q-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-T-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-ANN-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d/MERRA2/MERRA2-U-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n" + ] + } + ], + "source": [ + "_get_relative_diffs(var_to_filepaths)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Results\n", + "\n", + "All datasets are <= relative tolerance of 1e-5, with only 1 to 3 elements mismatching.\n", + "\n", + "This gives me great confidence that the changes are good to go.\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "cdat_regression_test", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_run_script.py b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_run_script.py new file mode 100644 index 000000000..704720997 --- /dev/null +++ b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_run_script.py @@ -0,0 +1,6 @@ +from auxiliary_tools.cdat_regression_testing.base_run_script import run_set + +SET_NAME = "zonal_mean_2d" +SET_DIR = "655-zonal-mean-2d" + +run_set(SET_NAME, SET_DIR) diff --git a/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_stratosphere_regression_test_netcdf.ipynb b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_stratosphere_regression_test_netcdf.ipynb new file mode 100644 index 000000000..522f24a56 --- /dev/null +++ b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_stratosphere_regression_test_netcdf.ipynb @@ -0,0 +1,556 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# CDAT Migration Regression Testing Notebook (`.nc` files)\n", + "\n", + "This notebook is used to perform regression testing between the development and\n", + "production versions of a diagnostic set.\n", + "\n", + "## How it works\n", + "\n", + "It compares the relative differences (%) between ref and test variables between\n", + "the dev and `main` branches.\n", + "\n", + "## How to use\n", + "\n", + "PREREQUISITE: The diagnostic set's netCDF stored in `.json` files in two directories\n", + "(dev and `main` branches).\n", + "\n", + "1. Make a copy of this notebook under `auxiliary_tools/cdat_regression_testing/`.\n", + "2. Run `mamba create -n cdat_regression_test -y -c conda-forge \"python<3.12\" xarray dask pandas matplotlib-base ipykernel`\n", + "3. Run `mamba activate cdat_regression_test`\n", + "4. Update `SET_DIR` and `SET_NAME` in the copy of your notebook.\n", + "5. Run all cells IN ORDER.\n", + "6. Review results for any outstanding differences (>=1e-5 relative tolerance).\n", + " - Debug these differences (e.g., bug in metrics functions, incorrect variable references, etc.)\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup Code\n" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "import glob\n", + "\n", + "import numpy as np\n", + "import xarray as xr\n", + "\n", + "# TODO: Update SET_NAME and SET_DIR\n", + "SET_NAME = \"zonal_mean_2d_stratosphere\"\n", + "SET_DIR = \"655-zonal-mean-2d-stratosphere\"\n", + "\n", + "DEV_PATH = f\"/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/{SET_DIR}/{SET_NAME}/**\"\n", + "MAIN_PATH = f\"/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/{SET_NAME}/**\"\n", + "\n", + "DEV_GLOB = sorted(glob.glob(DEV_PATH + \"/*.nc\"))\n", + "MAIN_GLOB = sorted(glob.glob(MAIN_PATH + \"/*.nc\"))\n", + "\n", + "if len(DEV_GLOB) == 0 or len(MAIN_GLOB) == 0:\n", + " raise IOError(\"No files found at DEV_PATH and/or MAIN_PATH.\")\n", + "\n", + "dev_num_files = len(DEV_GLOB)\n", + "main_num_files = len(MAIN_GLOB)\n", + "if dev_num_files != main_num_files:\n", + " raise IOError(\n", + " f\"Number of files do not match at DEV_PATH ({dev_num_files}) and MAIN_PATH ({main_num_files}).\"\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "def _get_var_to_filepath_map():\n", + " var_to_file = defaultdict(lambda: defaultdict(lambda: defaultdict(dict)))\n", + "\n", + " for dev_file, main_file in zip(DEV_GLOB, MAIN_GLOB):\n", + " if \"relative difference\" in dev_file:\n", + " continue\n", + "\n", + " # Example:\n", + " # \"/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/660-cosp-histogram/cosp_histogram/ISCCP-COSP/ISCCPCOSP-COSP_HISTOGRAM_ISCCP-ANN-global_test.nc\"\n", + " file_arr = dev_file.split(\"/\")\n", + "\n", + " # Example: \"test\"\n", + " data_type = dev_file.split(\"_\")[-1].split(\".nc\")[0]\n", + "\n", + " # Skip comparing `.nc` \"diff\" files because comparing relative diffs of\n", + " # does not make sense.\n", + " if data_type == \"test\" or data_type == \"ref\":\n", + " # Example: [ERA5, OMEGA, ANN, global_ref.nc]\n", + " filename = file_arr[-1].split(\"-\")\n", + " # Example: ERA5\n", + " model = filename[0]\n", + " # Example: OMEGA\n", + " var_key = filename[1]\n", + "\n", + " season = \"JJA\" if \"JJA\" in dev_file else \"ANN\"\n", + "\n", + " var_to_file[model][var_key][data_type][season] = (dev_file, main_file)\n", + "\n", + " return var_to_file\n", + "\n", + "\n", + "def _get_relative_diffs(var_to_filepath):\n", + " # Absolute tolerance of 0 and relative tolerance of 1e-5.\n", + " # We are mainly focusing on relative tolerance here (in percentage terms).\n", + " atol = 0\n", + " rtol = 1e-5\n", + "\n", + " for _, var_keys in var_to_filepath.items():\n", + " for var_key, data_types in var_keys.items():\n", + " for _, seasons in data_types.items():\n", + " for _, filepaths in seasons.items():\n", + " print(\"Comparing:\")\n", + " print(filepaths[0], \"\\n\", filepaths[1])\n", + " ds1 = xr.open_dataset(filepaths[0])\n", + " ds2 = xr.open_dataset(filepaths[1])\n", + "\n", + " try:\n", + " np.testing.assert_allclose(\n", + " ds1[var_key].values,\n", + " ds2[var_key].values,\n", + " atol=atol,\n", + " rtol=rtol,\n", + " )\n", + " except AssertionError as e:\n", + " print(e)\n", + " else:\n", + " print(f\" * All close and within relative tolerance ({rtol})\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Compare the netCDF files between branches\n", + "\n", + "- Compare \"ref\" and \"test\" files\n", + "- \"diff\" files are ignored because getting relative diffs for these does not make sense (relative diff will be above tolerance)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "var_to_filepaths = _get_var_to_filepath_map()" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "defaultdict(.()>,\n", + " {'ERA5': defaultdict(...()>,\n", + " {'OMEGA': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_test.nc')}}),\n", + " 'Q': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_test.nc')}}),\n", + " 'RELHUM': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_test.nc')}}),\n", + " 'T': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_test.nc')}}),\n", + " 'U': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_test.nc')}})}),\n", + " 'MERRA2': defaultdict(...()>,\n", + " {'OMEGA': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_test.nc')}}),\n", + " 'Q': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_test.nc')}}),\n", + " 'RELHUM': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_test.nc')}}),\n", + " 'T': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_test.nc')}}),\n", + " 'U': defaultdict(dict,\n", + " {'ref': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_ref.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_ref.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_ref.nc')},\n", + " 'test': {'ANN': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_test.nc'),\n", + " 'JJA': ('/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_test.nc',\n", + " '/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_test.nc')}})})})" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "var_to_filepaths" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-OMEGA-JJA-global_test.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 1800 (0.0556%)\n", + "Max absolute difference: 1.69393882e-08\n", + "Max relative difference: 8.69285391e-05\n", + " x: array([[ 0.070015, 0.069959, 0.069908, ..., -0.031941, -0.030957,\n", + " -0.030484],\n", + " [ 0.083583, 0.083959, 0.08449 , ..., -0.045252, -0.044022,...\n", + " y: array([[ 0.070015, 0.069959, 0.069908, ..., -0.031941, -0.030957,\n", + " -0.030484],\n", + " [ 0.083583, 0.083959, 0.08449 , ..., -0.045252, -0.044022,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-Q-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-RELHUM-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-T-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/ERA5/ERA5-U-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-OMEGA-JJA-global_test.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 1800 (0.0556%)\n", + "Max absolute difference: 1.69393882e-08\n", + "Max relative difference: 8.69285391e-05\n", + " x: array([[ 0.070015, 0.069959, 0.069908, ..., -0.031941, -0.030957,\n", + " -0.030484],\n", + " [ 0.083583, 0.083959, 0.08449 , ..., -0.045252, -0.044022,...\n", + " y: array([[ 0.070015, 0.069959, 0.069908, ..., -0.031941, -0.030957,\n", + " -0.030484],\n", + " [ 0.083583, 0.083959, 0.08449 , ..., -0.045252, -0.044022,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-Q-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-RELHUM-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_ref.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-T-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_ref.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 16 / 3610 (0.443%)\n", + "Max absolute difference: 1.12679693e-07\n", + "Max relative difference: 0.06101062\n", + " x: array([[ 6.962348e-07, 8.762654e-01, 1.605988e+00, ..., 2.470426e+00,\n", + " 1.311890e+00, 1.454603e-06],\n", + " [ 1.235800e-06, 1.004880e+00, 1.845646e+00, ..., 2.164340e+00,...\n", + " y: array([[ 6.962348e-07, 8.762654e-01, 1.605988e+00, ..., 2.470426e+00,\n", + " 1.311890e+00, 1.454603e-06],\n", + " [ 1.227703e-06, 1.004880e+00, 1.845646e+00, ..., 2.164340e+00,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_ref.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_ref.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 15 / 3610 (0.416%)\n", + "Max absolute difference: 2.09578261e-07\n", + "Max relative difference: 0.00898989\n", + " x: array([[-9.994902e-07, 1.855954e+00, 3.415055e+00, ..., -1.025078e+00,\n", + " -5.548744e-01, -9.440600e-07],\n", + " [-9.595097e-07, 1.685548e+00, 3.103145e+00, ..., -8.308557e-01,...\n", + " y: array([[-9.994902e-07, 1.855954e+00, 3.415055e+00, ..., -1.025078e+00,\n", + " -5.548744e-01, -9.440600e-07],\n", + " [-9.595177e-07, 1.685548e+00, 3.103145e+00, ..., -8.308557e-01,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_test.nc \n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_test.nc\n", + " * All close and within relative tolerance (1e-05)\n" + ] + } + ], + "source": [ + "_get_relative_diffs(var_to_filepaths)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Results\n", + "\n", + "Tehre are two datasets not within the relative tolernace:\n", + "\n", + "```python\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_ref.nc\n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-SON-global_ref.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 16 / 3610 (0.443%)\n", + "Max absolute difference: 1.12679693e-07\n", + "Max relative difference: 0.06101062\n", + " x: array([[ 6.962348e-07, 8.762654e-01, 1.605988e+00, ..., 2.470426e+00,\n", + " 1.311890e+00, 1.454603e-06],\n", + " [ 1.235800e-06, 1.004880e+00, 1.845646e+00, ..., 2.164340e+00,...\n", + " y: array([[ 6.962348e-07, 8.762654e-01, 1.605988e+00, ..., 2.470426e+00,\n", + " 1.311890e+00, 1.454603e-06],\n", + " [ 1.227703e-06, 1.004880e+00, 1.845646e+00, ..., 2.164340e+00,...\n", + "Comparing:\n", + "/global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/655-zonal-mean-2d-stratosphere/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_ref.nc\n", + " /global/cfs/projectdirs/e3sm/e3sm_diags_cdat_test/main/zonal_mean_2d_stratosphere/MERRA2/MERRA2-U-JJA-global_ref.nc\n", + "\n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 15 / 3610 (0.416%)\n", + "Max absolute difference: 2.09578261e-07\n", + "Max relative difference: 0.00898989\n", + " x: array([[-9.994902e-07, 1.855954e+00, 3.415055e+00, ..., -1.025078e+00,\n", + " -5.548744e-01, -9.440600e-07],\n", + " [-9.595097e-07, 1.685548e+00, 3.103145e+00, ..., -8.308557e-01,...\n", + " y: array([[-9.994902e-07, 1.855954e+00, 3.415055e+00, ..., -1.025078e+00,\n", + " -5.548744e-01, -9.440600e-07],\n", + " [-9.595177e-07, 1.685548e+00, 3.103145e+00, ..., -8.308557e-01,...\n", + "```\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "cdat_regression_test", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_stratosphere_run_script.py b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_stratosphere_run_script.py new file mode 100644 index 000000000..c8ff04512 --- /dev/null +++ b/auxiliary_tools/cdat_regression_testing/655-zonal-mean-2d/655_zonal_mean_2d_stratosphere_run_script.py @@ -0,0 +1,6 @@ +from auxiliary_tools.cdat_regression_testing.base_run_script import run_set + +SET_NAME = "zonal_mean_2d_stratosphere" +SET_DIR = "655-zonal-mean-2d-stratosphere" + +run_set(SET_NAME, SET_DIR) diff --git a/e3sm_diags/driver/utils/dataset_xr.py b/e3sm_diags/driver/utils/dataset_xr.py index f3706701d..3702e89be 100644 --- a/e3sm_diags/driver/utils/dataset_xr.py +++ b/e3sm_diags/driver/utils/dataset_xr.py @@ -385,10 +385,10 @@ def _get_climo_dataset(self, season: str) -> xr.Dataset: filepath = self._get_climo_filepath(season) ds = self._open_climo_dataset(filepath) - if self.var in ds.variables: - pass - elif self.var in self.derived_vars_map: + if self.var in self.derived_vars_map: ds = self._get_dataset_with_derived_climo_var(ds) + elif self.var in ds.data_vars.keys(): + pass else: raise IOError( f"Variable '{self.var}' was not in the file '{filepath}', nor was " diff --git a/e3sm_diags/driver/utils/regrid.py b/e3sm_diags/driver/utils/regrid.py index 0d8f39372..efe83f601 100644 --- a/e3sm_diags/driver/utils/regrid.py +++ b/e3sm_diags/driver/utils/regrid.py @@ -122,8 +122,8 @@ def has_z_axis(data_var: xr.DataArray) -> bool: return False -def get_z_axis(data_var: xr.DataArray) -> xr.DataArray: - """Gets the Z axis coordinates. +def get_z_axis(obj: xr.Dataset | xr.DataArray) -> xr.DataArray: + """Gets the Z axis coordinates from an xarray object. Returns True if: - Data variable has a "Z" axis in the cf-xarray mapping dict @@ -133,8 +133,8 @@ def get_z_axis(data_var: xr.DataArray) -> xr.DataArray: Parameters ---------- - data_var : xr.DataArray - The data variable. + obj : xr.Dataset | xr.DataArray + The xarray Dataset or DataArray. Returns ------- @@ -148,18 +148,19 @@ def get_z_axis(data_var: xr.DataArray) -> xr.DataArray: - https://cdms.readthedocs.io/en/latest/_modules/cdms2/axis.html#AbstractAxis.isLevel """ try: - z_coords = xc.get_dim_coords(data_var, axis="Z") + z_coords = xc.get_dim_coords(obj, axis="Z") + return z_coords except KeyError: pass - for coord in data_var.coords.values(): + for coord in obj.coords.values(): if coord.name in ["lev", "plev", "depth"]: return coord raise KeyError( - f"No Z axis coordinate were found in the '{data_var.name}' " - "Make sure the variable has Z axis coordinates" + f"No Z axis coordinates were found in the {type(obj)}. Make sure the " + f"{type(obj)} has Z axis coordinates." ) @@ -517,6 +518,7 @@ def _hybrid_to_plevs( pressure_grid = xc.create_grid(z=z_axis) pressure_coords = _hybrid_to_pressure(ds, var_key) + # Keep the "axis" and "coordinate" attributes for CF mapping. with xr.set_options(keep_attrs=True): result = ds.regridder.vertical( @@ -527,6 +529,10 @@ def _hybrid_to_plevs( target_data=pressure_coords, ) + # Vertical regriding sets the units to "mb", but the original units + # should be preserved. + result[var_key].attrs["units"] = ds[var_key].attrs["units"] + return result diff --git a/e3sm_diags/driver/zonal_mean_2d_driver.py b/e3sm_diags/driver/zonal_mean_2d_driver.py index 0ddb14519..da2aa26c0 100755 --- a/e3sm_diags/driver/zonal_mean_2d_driver.py +++ b/e3sm_diags/driver/zonal_mean_2d_driver.py @@ -1,283 +1,287 @@ -import os - -import cdms2 -import cdutil -import MV2 -import numpy - -import e3sm_diags -from e3sm_diags.driver import utils +import copy +from typing import Tuple + +import xarray as xr +import xcdat as xc # noqa: F401 + +from e3sm_diags.driver.utils.dataset_xr import Dataset +from e3sm_diags.driver.utils.io import _save_data_metrics_and_plots +from e3sm_diags.driver.utils.regrid import ( + align_grids_to_lower_res, + has_z_axis, + regrid_z_axis_to_plevs, +) +from e3sm_diags.driver.utils.type_annotations import MetricsDict from e3sm_diags.logger import custom_logger -from e3sm_diags.metrics import corr, max_cdms, mean, min_cdms, rmse -from e3sm_diags.parameter.zonal_mean_2d_parameter import ZonalMean2dParameter -from e3sm_diags.plot import plot +from e3sm_diags.metrics.metrics import correlation, rmse, spatial_avg +from e3sm_diags.parameter.zonal_mean_2d_parameter import ( + DEFAULT_PLEVS, + ZonalMean2dParameter, +) +from e3sm_diags.plot.cartopy.zonal_mean_2d_plot import plot as plot_func logger = custom_logger(__name__) - -def create_metrics(ref, test, ref_regrid, test_regrid, diff): - """Creates the mean, max, min, rmse, corr in a dictionary""" - orig_bounds = cdms2.getAutoBounds() - cdms2.setAutoBounds(1) - lev = ref.getLevel() - if lev is not None: - lev.setBounds(None) - - lev = test.getLevel() - if lev is not None: - lev.setBounds(None) - - lev = test_regrid.getLevel() - if lev is not None: - lev.setBounds(None) - - lev = ref_regrid.getLevel() - if lev is not None: - lev.setBounds(None) - - lev = diff.getLevel() - if lev is not None: - lev.setBounds(None) - cdms2.setAutoBounds(orig_bounds) - - metrics_dict = {} - metrics_dict["ref"] = { - "min": min_cdms(ref), - "max": max_cdms(ref), - "mean": mean(ref, axis="yz"), - } - metrics_dict["test"] = { - "min": min_cdms(test), - "max": max_cdms(test), - "mean": mean(test, axis="yz"), - } - - metrics_dict["diff"] = { - "min": min_cdms(diff), - "max": max_cdms(diff), - "mean": mean(diff, axis="yz"), - } - metrics_dict["misc"] = { - "rmse": rmse(test_regrid, ref_regrid, axis="yz"), - "corr": corr(test_regrid, ref_regrid, axis="yz"), - } - - return metrics_dict +DEFAULT_PLEVS = copy.deepcopy(DEFAULT_PLEVS) def run_diag( - parameter: ZonalMean2dParameter, default_plevs=ZonalMean2dParameter().plevs + parameter: ZonalMean2dParameter, default_plevs=DEFAULT_PLEVS ) -> ZonalMean2dParameter: variables = parameter.variables seasons = parameter.seasons ref_name = getattr(parameter, "ref_name", "") - regions = parameter.regions - - test_data = utils.dataset.Dataset(parameter, test=True) - ref_data = utils.dataset.Dataset(parameter, ref=True) - - for season in seasons: - # Get the name of the data, appended with the years averaged. - parameter.test_name_yrs = utils.general.get_name_and_yrs( - parameter, test_data, season - ) - parameter.ref_name_yrs = utils.general.get_name_and_yrs( - parameter, ref_data, season - ) - - # Get land/ocean fraction for masking. - try: - land_frac = test_data.get_climo_variable("LANDFRAC", season) - ocean_frac = test_data.get_climo_variable("OCNFRAC", season) - except Exception: - mask_path = os.path.join( - e3sm_diags.INSTALL_PATH, "acme_ne30_ocean_land_mask.nc" - ) - with cdms2.open(mask_path) as f: - land_frac = f("LANDFRAC") - ocean_frac = f("OCNFRAC") - - for var in variables: - logger.info("Variable: {}".format(var)) - parameter.var_id = var - - mv1 = test_data.get_climo_variable(var, season) - mv2 = ref_data.get_climo_variable(var, season) - - parameter.viewer_descr[var] = ( - mv1.long_name - if hasattr(mv1, "long_name") - else "No long_name attr in test data." - ) - - # Special case, cdms didn't properly convert mask with fill value - # -999.0, filed issue with Denis. - if ref_name == "WARREN": - # This is cdms2 fix for bad mask, Denis' fix should fix this. - mv2 = MV2.masked_where(mv2 == -0.9, mv2) - # The following should be moved to a derived variable. - if ref_name == "AIRS": - # This is cdms2 fix for bad mask, Denis' fix should fix this. - mv2 = MV2.masked_where(mv2 > 1e20, mv2) - if ref_name == "WILLMOTT" or ref_name == "CLOUDSAT": - # This is cdms2 fix for bad mask, Denis' fix should fix this. - mv2 = MV2.masked_where(mv2 == -999.0, mv2) - - # The following should be moved to a derived variable. - if var == "PRECT_LAND": - days_season = { - "ANN": 365, - "DJF": 90, - "MAM": 92, - "JJA": 92, - "SON": 91, - } - # mv1 = mv1 * days_season[season] * 0.1 # following AMWG - # Approximate way to convert to seasonal cumulative - # precipitation, need to have solution in derived variable, - # unit convert from mm/day to cm. - mv2 = ( - mv2 / days_season[season] / 0.1 - ) # Convert cm to mm/day instead. - mv2.units = "mm/day" - - # For variables with a z-axis. - if mv1.getLevel() and mv2.getLevel(): - # Since the default is now stored in `default_plevs`, - # we must get it from there if the plevs param is blank. - plevs = parameter.plevs - if (isinstance(plevs, numpy.ndarray) and not plevs.all()) or ( - not isinstance(plevs, numpy.ndarray) and not plevs - ): - plevs = default_plevs - logger.info(f"Selected pressure level: {plevs}") - - mv1_p = utils.general.convert_to_pressure_levels( - mv1, plevs, test_data, var, season - ) - mv2_p = utils.general.convert_to_pressure_levels( - mv2, plevs, ref_data, var, season - ) - # Note this is a special case to handle small values of stratosphere specific humidity. - # The general derived variable process converts specific humidity to units [g/kg] - # Following converts from g/kg to ppm - - if ( - parameter.current_set == "zonal_mean_2d_stratosphere" - and parameter.var_id == "Q" - ): - mv1_p = mv1_p * 1000.0 - mv1_p.units = "ppm" - mv2_p = mv2_p * 1000.0 - mv2_p.units = "ppm" - # Regrid towards the lower resolution of the two - # variables for calculating the difference. - mv1_p_reg, mv2_p_reg = utils.general.regrid_to_lower_res( - mv1_p, - mv2_p, - parameter.regrid_tool, - parameter.regrid_method, - ) + if not parameter._is_plevs_set(): + parameter.plevs = default_plevs + + test_ds = Dataset(parameter, data_type="test") + ref_ds = Dataset(parameter, data_type="ref") - diff_p = mv1_p_reg - mv2_p_reg - diff = cdutil.averager(diff_p, axis="x") + for var_key in variables: + logger.info("Variable: {}".format(var_key)) + parameter.var_id = var_key - mv1_p = cdutil.averager(mv1_p, axis="x") - mv2_p = cdutil.averager(mv2_p, axis="x") + for season in seasons: + parameter._set_name_yrs_attrs(test_ds, ref_ds, season) - # Make sure mv1_p_reg and mv2_p_reg have same mask - mv1_p_reg = mv2_p_reg + diff_p - mv2_p_reg = mv1_p_reg - diff_p + ds_test = test_ds.get_climo_dataset(var_key, season) + ds_ref = ref_ds.get_ref_climo_dataset(var_key, season, ds_test) - mv1_reg = cdutil.averager(mv1_p_reg, axis="x") - mv2_reg = cdutil.averager(mv2_p_reg, axis="x") + # Store the variable's DataArray objects for reuse. + dv_test = ds_test[var_key] + dv_ref = ds_ref[var_key] - parameter.output_file = "-".join( - [ref_name, var, season, parameter.regions[0]] + is_vars_3d = has_z_axis(dv_test) and has_z_axis(dv_ref) + is_dims_diff = has_z_axis(dv_test) != has_z_axis(dv_ref) + + if is_dims_diff: + raise RuntimeError( + "The dimensions of the test and reference variables are different, " + f"({dv_test.dims} vs. {dv_ref.dims})." + ) + elif is_vars_3d: + _run_diags_3d( + parameter, + ds_test, + ds_ref, + season, + var_key, + ref_name, ) - parameter.main_title = str(" ".join([var, season])) + else: + raise RuntimeError("Dimensions of the two variables are different.") - if parameter.diff_type == "relative": - diff = diff / mv2_reg * 100.0 + return parameter - # Use mv2_p and mv1_p on the original horizonal grids for visualization and their own metrics - # Use mv2_reg and mv1_reg for rmse and correlation coefficient calculation - metrics_dict = create_metrics(mv2_p, mv1_p, mv2_reg, mv1_reg, diff) - parameter.var_region = "global" +def _run_diags_3d( + parameter: ZonalMean2dParameter, + ds_test: xr.Dataset, + ds_ref: xr.Dataset, + season: str, + var_key: str, + ref_name: str, +): + plevs = parameter.plevs + logger.info(f"Selected pressure level: {plevs}") + + ds_t_plevs = regrid_z_axis_to_plevs(ds_test, var_key, plevs) + ds_r_plevs = regrid_z_axis_to_plevs(ds_ref, var_key, plevs) + + ds_t_plevs = _convert_g_kg_to_ppm_units(parameter, ds_t_plevs, var_key) + ds_r_plevs = _convert_g_kg_to_ppm_units(parameter, ds_r_plevs, var_key) + + # Calculate the spatial average of the variables on their original pressure + # level grids. These variables are used for the visualization on original + # horizonal grids for metrics output. + ds_t_plevs_avg = ds_t_plevs.spatial.average(var_key, axis="X") + ds_r_plevs_avg = ds_r_plevs.spatial.average(var_key, axis="X") + + # Align the grids for the variables and calculate their spatial averages + # and the difference between their spatial averages. These variables are + # used for calculating rmse and correlation. + ( + ds_t_plevs_rg_avg, + ds_r_plevs_rg_avg, + ds_diff_rg_avg, + ) = _get_avg_for_regridded_datasets(parameter, ds_t_plevs, ds_r_plevs, var_key) + + metrics_dict = _create_metrics_dict( + var_key, + ds_t_plevs_avg, + ds_t_plevs_rg_avg, + ds_r_plevs_avg, + ds_r_plevs_rg_avg, + ds_diff_rg_avg, + ) + + # Set parameter attributes for output files. + parameter.var_region = "global" + parameter.output_file = "-".join([ref_name, var_key, season, parameter.regions[0]]) + parameter.main_title = str(" ".join([var_key, season])) + + _save_data_metrics_and_plots( + parameter, + plot_func, + var_key, + ds_t_plevs_avg, + ds_r_plevs_avg, + ds_diff_rg_avg, + metrics_dict, + ) + + +def _convert_g_kg_to_ppm_units( + parameter: ZonalMean2dParameter, ds: xr.Dataset, var_key: str +) -> xr.Dataset: + """Adjust the units for a variable to handle special cases. + + This is a special case to handle small values of stratosphere specific + humidity. The general derived variable process converts specific humidity to + units [g/kg]. This function converts from "g/kg" to "ppm". + + Parameters + ---------- + parameter : ZonalMean2dParameter + The parameter object. + ds : xr.Dataset + The dataset. + var_key : str + They key of the variable. + + Returns + ------- + xr.Dataset + The dataset with units converted. + """ + ds_new = ds.copy() + + with xr.set_options(keep_attrs=True): + if ( + parameter.current_set == "zonal_mean_2d_stratosphere" + and parameter.var_id == "Q" + ): + ds_new[var_key] = ds_new[var_key] * 1000.0 + ds_new[var_key].attrs["units"] = "ppm" + + return ds_new + + +def _get_avg_for_regridded_datasets( + parameter: ZonalMean2dParameter, + ds_test_plevs: xr.Dataset, + ds_ref_plevs: xr.Dataset, + var_key: str, +) -> Tuple[xr.Dataset, xr.Dataset, xr.Dataset]: + """Get the average and difference between averages for the plevs datasets. + + Parameters + ---------- + parameter : ZonalMean2dParameter + The parameter object. + ds_test_plevs : xr.Dataset + The test dataset on pressure level coordinates. + ds_ref_plevs : xr.Dataset + The reference dataset on pressure level coordinates. + var_key : str + The key of the variable. + + Returns + ------- + Tuple[xr.Dataset, xr.Dataset, xr.Dataset] + A tuple consisting of the average of the test dataset, the + average of the ref dataset, and the difference between averages. + """ + ds_test_rg, ds_ref_rg = align_grids_to_lower_res( + ds_test_plevs, + ds_ref_plevs, + var_key, + parameter.regrid_tool, + parameter.regrid_method, + ) + + # Get the difference between the regridded variables and use it to + # make sure the regridded variables have the same mask. + with xr.set_options(keep_attrs=True): + ds_diff_rg = ds_test_rg.copy() + ds_diff_rg[var_key] = ds_test_rg[var_key] - ds_ref_rg[var_key] + + ds_test_rg[var_key] = ds_ref_rg[var_key] + ds_diff_rg[var_key] + ds_ref_rg[var_key] = ds_test_rg[var_key] - ds_diff_rg[var_key] + + # Calculate the spatial averages for the masked variables. + ds_test_rg_avg = ds_test_rg.spatial.average(var_key, axis=["X"]) + ds_ref_rg_avg = ds_ref_rg.spatial.average(var_key, axis=["X"]) + + # Calculate the spatial average for the differences + ds_diff_rg_avg = ds_diff_rg.spatial.average(var_key, axis=["X"]) + + if parameter.diff_type == "relative": + ds_diff_rg_avg = ds_diff_rg_avg / ds_ref_rg_avg * 100.0 + ds_diff_rg_avg[var_key].attrs["units"] = "%" + + return ds_test_rg_avg, ds_test_rg_avg, ds_diff_rg_avg + + +def _create_metrics_dict( + var_key: str, + ds_test: xr.Dataset, + ds_test_regrid: xr.Dataset, + ds_ref: xr.Dataset, + ds_ref_regrid: xr.Dataset, + ds_diff: xr.Dataset, +) -> MetricsDict: + """Calculate metrics using the variable in the datasets. + + Metrics include min value, max value, spatial average (mean), standard + deviation, correlation (pearson_r), and RMSE. + + Parameters + ---------- + var_key : str + The variable key. + ds_test : xr.Dataset + The test dataset. + ds_test_regrid : xr.Dataset + The regridded test Dataset. + ds_ref : xr.Dataset + The reference dataset. + ds_ref_regrid : xr.Dataset + The regridded reference dataset. + ds_diff : xr. Dataset + The difference between ``ds_test_regrid`` and ``ds_ref_regrid``. + + Returns + ------- + Metrics + A dictionary with the key being a string and the value being either + a sub-dictionary (key is metric and value is float) or a string + ("unit"). + """ + metrics_dict = {} - plot( - parameter.current_set, - mv2_p, - mv1_p, - diff, - metrics_dict, - parameter, - ) - utils.general.save_ncfiles( - parameter.current_set, mv1_p, mv2_p, diff, parameter - ) + metrics_dict["units"] = ds_test[var_key].attrs["units"] + metrics_dict["ref"] = { + "min": ds_ref[var_key].min().item(), + "max": ds_test[var_key].max().item(), + "mean": spatial_avg(ds_ref, var_key, axis=["Y", "Z"]), + } + metrics_dict["test"] = { + "min": ds_test[var_key].min().item(), + "max": ds_test[var_key].max().item(), + "mean": spatial_avg(ds_test, var_key, axis=["Y", "Z"]), + } - # For variables without a z-axis. - elif mv1.getLevel() is None and mv2.getLevel() is None: - for region in regions: - logger.info(f"Selected region: {region}") - - mv1_domain = utils.general.select_region( - region, mv1, land_frac, ocean_frac, parameter - ) - mv2_domain = utils.general.select_region( - region, mv2, land_frac, ocean_frac, parameter - ) - - parameter.output_file = "-".join([ref_name, var, season, region]) - parameter.main_title = str(" ".join([var, season, region])) - - # Regrid towards the lower resolution of the two - # variables for calculating the difference. - mv1_reg, mv2_reg = utils.general.regrid_to_lower_res( - mv1_domain, - mv2_domain, - parameter.regrid_tool, - parameter.regrid_method, - ) - - # Special case. - if var == "TREFHT_LAND" or var == "SST": - if ref_name == "WILLMOTT": - mv2_reg = MV2.masked_where( - mv2_reg == mv2_reg.fill_value, mv2_reg - ) - land_mask = MV2.logical_or(mv1_reg.mask, mv2_reg.mask) - mv1_reg = MV2.masked_where(land_mask, mv1_reg) - mv2_reg = MV2.masked_where(land_mask, mv2_reg) - - diff = mv1_reg - mv2_reg - metrics_dict = create_metrics( - mv2_domain, mv1_domain, mv2_reg, mv1_reg, diff - ) - parameter.var_region = region - - plot( - parameter.current_set, - mv2_domain, - mv1_domain, - diff, - metrics_dict, - parameter, - ) - utils.general.save_ncfiles( - parameter.current_set, - mv1_domain, - mv2_domain, - diff, - parameter, - ) + metrics_dict["diff"] = { + "min": ds_diff[var_key].min().item(), + "max": ds_diff[var_key].max().item(), + "mean": spatial_avg(ds_diff, var_key, axis=["Y", "Z"]), + } - else: - raise RuntimeError( - "Dimensions of the two variables are different. Aborting." - ) + metrics_dict["misc"] = { + "rmse": rmse(ds_test_regrid, ds_ref_regrid, var_key, axis=["Y", "Z"]), + "corr": correlation(ds_test_regrid, ds_ref_regrid, var_key, axis=["Y", "Z"]), + } - return parameter + return metrics_dict diff --git a/e3sm_diags/driver/zonal_mean_2d_stratosphere_driver.py b/e3sm_diags/driver/zonal_mean_2d_stratosphere_driver.py index d3eecac86..4fc38b712 100755 --- a/e3sm_diags/driver/zonal_mean_2d_stratosphere_driver.py +++ b/e3sm_diags/driver/zonal_mean_2d_stratosphere_driver.py @@ -1,19 +1,16 @@ -from e3sm_diags.driver.zonal_mean_2d_driver import create_metrics as base_create_metrics +import copy + from e3sm_diags.driver.zonal_mean_2d_driver import run_diag as base_run_diag from e3sm_diags.parameter.zonal_mean_2d_parameter import ZonalMean2dParameter from e3sm_diags.parameter.zonal_mean_2d_stratosphere_parameter import ( + DEFAULT_PLEVS, ZonalMean2dStratosphereParameter, ) - -def create_metrics(ref, test, ref_regrid, test_regrid, diff): - """Creates the mean, max, min, rmse, corr in a dictionary""" - return base_create_metrics(ref, test, ref_regrid, test_regrid, diff) +DEFAULT_PLEVS = copy.deepcopy(DEFAULT_PLEVS) def run_diag( parameter: ZonalMean2dStratosphereParameter, ) -> ZonalMean2dParameter: - return base_run_diag( - parameter, default_plevs=ZonalMean2dStratosphereParameter().plevs - ) + return base_run_diag(parameter, default_plevs=DEFAULT_PLEVS) diff --git a/e3sm_diags/metrics/metrics.py b/e3sm_diags/metrics/metrics.py index 68e5a4fc1..68410f08a 100644 --- a/e3sm_diags/metrics/metrics.py +++ b/e3sm_diags/metrics/metrics.py @@ -1,8 +1,9 @@ """This module stores functions to calculate metrics using Xarray objects.""" from __future__ import annotations -from typing import List +from typing import List, Literal +import numpy as np import xarray as xr import xcdat as xc import xskillscore as xs @@ -11,27 +12,12 @@ logger = custom_logger(__name__) -AXES = ["X", "Y"] - - -def get_weights(ds: xr.Dataset): - """Get weights for the X and Y spatial axes. - - Parameters - ---------- - ds : xr.Dataset - The dataset. - - Returns - ------- - xr.DataArray - Weights for the specified axis. - """ - return ds.spatial.get_weights(axis=["X", "Y"]) +Axis = Literal["X", "Y", "Z"] +DEFAULT_AXIS: List[Axis] = ["X", "Y"] def spatial_avg( - ds: xr.Dataset, var_key: str, as_list: bool = True + ds: xr.Dataset, var_key: str, axis: List[Axis] = DEFAULT_AXIS, as_list: bool = True ) -> List[float] | xr.DataArray: """Compute a variable's weighted spatial average. @@ -40,10 +26,14 @@ def spatial_avg( ds : xr.Dataset The dataset containing the variable. var_key : str - The key of the varible. + The key of the variable in the dataset. + axis : List[Axis] + The list of axes to use for the computation, by default ["X", "Y"]. + Valid axes including "X", "Y", and "Z". as_list : bool Return the spatial average as a list of floats, by default True. - If False, return an xr.DataArray. + If False, return an xr.DataArray. Must be True to be serializable for + writing out to a `.json` metrics file. Returns ------- @@ -59,8 +49,11 @@ def spatial_avg( ----- Replaces `e3sm_diags.metrics.mean`. """ - ds_avg = ds.spatial.average(var_key, axis=AXES, weights="generate") - results = ds_avg[var_key] + dv = ds[var_key].copy() + weights = _get_weights(ds, var_key, axis) + dims = _get_dims(dv, axis) + + results = dv.weighted(weights).mean(dims, keep_attrs=True) if as_list: return results.data.tolist() @@ -68,7 +61,7 @@ def spatial_avg( return results -def std(ds: xr.Dataset, var_key: str) -> List[float]: +def std(ds: xr.Dataset, var_key: str, axis: List[Axis] = DEFAULT_AXIS) -> List[float]: """Compute the weighted standard deviation for a variable. Parameters @@ -77,6 +70,9 @@ def std(ds: xr.Dataset, var_key: str) -> List[float]: The dataset containing the variable. var_key : str The key of the variable. + axis : List[Axis] + The list of axes to use for the computation, by default ["X", "Y"]. + Valid strings include "X", "Y", and "Z". Returns ------- @@ -94,15 +90,17 @@ def std(ds: xr.Dataset, var_key: str) -> List[float]: """ dv = ds[var_key].copy() - weights = ds.spatial.get_weights(axis=AXES, data_var=var_key) - dims = _get_dims(dv, axis=AXES) + weights = _get_weights(ds, var_key, axis) + dims = _get_dims(dv, axis) result = dv.weighted(weights).std(dim=dims, keep_attrs=True) return result.data.tolist() -def correlation(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float]: +def correlation( + ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str, axis: List[Axis] = DEFAULT_AXIS +) -> List[float]: """Compute the correlation coefficient between two variables. This function uses the Pearson correlation coefficient. Refer to [1]_ for @@ -116,6 +114,9 @@ def correlation(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float] The second dataset. var_key: str The key of the variable. + axis : List[Axis] + The list of axes to use for the computation, by default ["X", "Y"]. + Valid axes including "X", "Y", and "Z". Returns ------- @@ -135,9 +136,9 @@ def correlation(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float] var_b = ds_b[var_key] # Dimensions, bounds, and coordinates should be identical between datasets, - # so use the first dataset and variable to get dimensions and weights. - dims = _get_dims(var_a, axis=AXES) - weights = get_weights(ds_a) + # so use the first dataset and first variable to get dimensions and weights. + dims = _get_dims(var_a, axis) + weights = _get_weights(ds_a, var_key, axis) result = xs.pearson_r(var_a, var_b, dim=dims, weights=weights, skipna=True) results_list = result.data.tolist() @@ -145,7 +146,9 @@ def correlation(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float] return results_list -def rmse(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float]: +def rmse( + ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str, axis: List[Axis] = DEFAULT_AXIS +) -> List[float]: """Calculates the root mean square error (RMSE) between two variables. Parameters @@ -156,6 +159,9 @@ def rmse(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float]: The second dataset. var_key: str The key of the variable. + axis : List[Axis] + The list of axes to use for the computation, by default ["X", "Y"]. + Valid axes including "X", "Y", and "Z". Returns ------- @@ -170,9 +176,9 @@ def rmse(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float]: var_b = ds_b[var_key] # Dimensions, bounds, and coordinates should be identical between datasets, - # so use the first dataset and variable to get dimensions and weights. - dims = _get_dims(var_a, axis=AXES) - weights = get_weights(ds_a) + # so use the first dataset and first variable to get dimensions and weights. + dims = _get_dims(var_a, axis) + weights = _get_weights(ds_a, var_key, axis) result = xs.rmse(var_a, var_b, dim=dims, weights=weights, skipna=True) results_list = result.data.tolist() @@ -180,7 +186,94 @@ def rmse(ds_a: xr.Dataset, ds_b: xr.Dataset, var_key: str) -> List[float]: return results_list -def _get_dims(da: xr.DataArray, axis: List[str]): +def _get_weights(ds: xr.Dataset, var_key: str, axis: List[Axis] = DEFAULT_AXIS): + """Get weights for the specified axes of the variable. + + Parameters + ---------- + ds : xr.Dataset + The dataset containing the variable. + var_key : str + The key of the variable. + axis : List[Axis] + The list of axes to use for the computation, by default ["X", "Y"]. + Valid axes including "X", "Y", and "Z". + + Returns + ------- + xr.DataArray + Weights for the specified axis. + + Notes + ----- + xCDAT's ``ds.spatial.get_weights()`` method only supports rectilinear grids + # ("X", "Y") as of v0.6.1. This function computes weights for "X" and Y" + using xCDAT, then calculates "Z" weights separately before combining all + weights into a single matrix. + """ + spatial_wts = None + vertical_wts = None + + spatial_axis = [key for key in axis if key != "Z"] + if len(spatial_axis) > 0: + spatial_wts = ds.spatial.get_weights(spatial_axis, data_var=var_key) + spatial_wts = spatial_wts.fillna(0) + + if "Z" in axis: + vertical_wts = _get_z_weights(ds, var_key) + + if spatial_wts is not None and vertical_wts is not None: + return spatial_wts * vertical_wts + elif spatial_wts is not None: + return spatial_wts + elif vertical_wts is not None: + return vertical_wts + + +def _get_z_weights(ds: xr.Dataset, var_key: str) -> xr.DataArray: + """Get the Z axis weights using Z axis bounds. + + Parameters + ---------- + ds : xr.Dataset + The dataset containing a Z axis. + var_key : str + The key of the variable to get weights for. + + Returns + ------- + xr.DataArray + Weights for the Z axis. + + Raises + ------ + RuntimeError + If the dataset has no Z axis bounds. + + Notes + ----- + xCDAT spatial average get_weights() method does not support the Z axis as of + v0.6.1. This function is temporarily used until get_weights() supports the + Z axis. The logic is the same as xCDAT. + + * Related issue: https://github.com/xCDAT/xcdat/issues/596 + * Source: https://github.com/xCDAT/xcdat/blob/main/xcdat/spatial.py#L479-L495C16 + """ + try: + z_bnds = ds.bounds.get_bounds("Z", var_key) + except KeyError: + raise RuntimeError( + f"The dataset for {var_key} has no Z axis bounds to get weights for the Z " + "axis." + ) + + weights = np.abs(z_bnds[:, 1] - z_bnds[:, 0]) + weights = weights.fillna(0) + + return weights + + +def _get_dims(da: xr.DataArray, axis: List[Axis]) -> List[str]: """Get the dimensions for an axis in an xarray.DataArray. The dimensions are passed to the ``dim`` argument in xarray or xarray-based @@ -190,8 +283,9 @@ def _get_dims(da: xr.DataArray, axis: List[str]): ---------- da : xr.DataArray The array. - axis : List[str] - A list of axis strings. + axis : List[Axis] + The list of axes to get dimensions for. Valid strings + include "X", "Y", and "Z". Returns ------- diff --git a/e3sm_diags/parameter/core_parameter.py b/e3sm_diags/parameter/core_parameter.py index 78711f857..128fe8123 100644 --- a/e3sm_diags/parameter/core_parameter.py +++ b/e3sm_diags/parameter/core_parameter.py @@ -5,6 +5,8 @@ import sys from typing import TYPE_CHECKING, Any, Dict, List, Tuple +import numpy as np + from e3sm_diags.derivations.derivations import DerivedVariablesMap from e3sm_diags.driver.utils.climo_xr import CLIMO_FREQ from e3sm_diags.driver.utils.regrid import REGRID_TOOLS @@ -290,6 +292,14 @@ def _set_name_yrs_attrs( self.test_name_yrs = ds_test.get_name_yrs_attr(season) self.ref_name_yrs = ds_ref.get_name_yrs_attr(season) + def _is_plevs_set(self): + if (isinstance(self.plevs, np.ndarray) and not self.plevs.all()) or ( + not isinstance(self.plevs, np.ndarray) and not self.plevs + ): + return False + + return True + def _run_diag(self) -> List[Any]: """Run the diagnostics for each set in the parameter. diff --git a/e3sm_diags/parameter/zonal_mean_2d_parameter.py b/e3sm_diags/parameter/zonal_mean_2d_parameter.py index 0a4260c3c..554c35ce6 100644 --- a/e3sm_diags/parameter/zonal_mean_2d_parameter.py +++ b/e3sm_diags/parameter/zonal_mean_2d_parameter.py @@ -1,16 +1,20 @@ -import numpy +import copy + +import numpy as np from e3sm_diags.driver.utils.general import monotonic from .core_parameter import CoreParameter +DEFAULT_PLEVS = np.linspace(50, 1000, 20).tolist() + class ZonalMean2dParameter(CoreParameter): def __init__(self): super(ZonalMean2dParameter, self).__init__() # Override existing attributes # ============================= - self.plevs = numpy.linspace(50, 1000, 20).tolist() + self.plevs = copy.deepcopy(DEFAULT_PLEVS) self.plot_log_plevs = False self.plot_plevs = False # Granulating plevs causes duplicate plots in this case. diff --git a/e3sm_diags/parameter/zonal_mean_2d_stratosphere_parameter.py b/e3sm_diags/parameter/zonal_mean_2d_stratosphere_parameter.py index 937ae6436..bb8adb350 100644 --- a/e3sm_diags/parameter/zonal_mean_2d_stratosphere_parameter.py +++ b/e3sm_diags/parameter/zonal_mean_2d_stratosphere_parameter.py @@ -1,12 +1,16 @@ -import numpy +import copy + +import numpy as np from e3sm_diags.parameter.zonal_mean_2d_parameter import ZonalMean2dParameter +DEFAULT_PLEVS = np.logspace(0, 2.0, num=10).tolist() + class ZonalMean2dStratosphereParameter(ZonalMean2dParameter): def __init__(self): super(ZonalMean2dStratosphereParameter, self).__init__() # Override existing attributes # ============================= - self.plevs = numpy.logspace(0, 2.0, num=10).tolist() + self.plevs = copy.deepcopy(DEFAULT_PLEVS) self.plot_log_plevs = True diff --git a/e3sm_diags/plot/cartopy/zonal_mean_2d_plot.py b/e3sm_diags/plot/cartopy/zonal_mean_2d_plot.py index 83a9b48a8..a72bf5dce 100644 --- a/e3sm_diags/plot/cartopy/zonal_mean_2d_plot.py +++ b/e3sm_diags/plot/cartopy/zonal_mean_2d_plot.py @@ -1,293 +1,187 @@ -from __future__ import print_function - -import os +from typing import List, Optional, Tuple import matplotlib import numpy as np -import numpy.ma as ma -from cartopy.mpl.ticker import LatitudeFormatter +import xarray as xr +import xcdat as xc -from e3sm_diags.driver.utils.general import get_output_dir +from e3sm_diags.driver.utils.type_annotations import MetricsDict from e3sm_diags.logger import custom_logger -from e3sm_diags.plot import get_colormap +from e3sm_diags.parameter.core_parameter import CoreParameter +from e3sm_diags.parameter.zonal_mean_2d_parameter import DEFAULT_PLEVS +from e3sm_diags.plot.utils import ( + DEFAULT_PANEL_CFG, + _add_colorbar, + _add_contour_plot, + _add_min_mean_max_text, + _add_rmse_corr_text, + _configure_titles, + _configure_x_and_y_axes, + _get_c_levels_and_norm, + _save_plot, +) matplotlib.use("Agg") -import matplotlib.colors as colors # isort:skip # noqa: E402 import matplotlib.pyplot as plt # isort:skip # noqa: E402 logger = custom_logger(__name__) -plotTitle = {"fontsize": 11.5} -plotSideTitle = {"fontsize": 9.5} - -# Position and sizes of subplot axes in page coordinates (0 to 1) -panel = [ - (0.1691, 0.6810, 0.6465, 0.2258), - (0.1691, 0.3961, 0.6465, 0.2258), - (0.1691, 0.1112, 0.6465, 0.2258), -] - -# Border padding relative to subplot axes for saving individual panels -# (left, bottom, right, top) in page coordinates -border = (-0.06, -0.03, 0.13, 0.03) - - -def add_cyclic(var): - lon = var.getLongitude() - return var(longitude=(lon[0], lon[0] + 360.0, "coe")) - - -def get_ax_size(fig, ax): - bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) - width, height = bbox.width, bbox.height - width *= fig.dpi - height *= fig.dpi - return width, height - - -def plot_panel(n, fig, proj, var, clevels, cmap, title, parameters, stats=None): - # var_min = float(var.min()) - # var_max = float(var.max()) - # var_mean = cdutil.averager(var, axis='xy', weights='generate') - # var = add_cyclic(var) - var.getLongitude() - lat = var.getLatitude() - plev = var.getLevel() - var = ma.squeeze(var.asma()) - - # Contour levels - levels = None - norm = None - if len(clevels) > 0: - levels = [-1.0e8] + clevels + [1.0e8] - norm = colors.BoundaryNorm(boundaries=levels, ncolors=256) - - # Contour plot - ax = fig.add_axes(panel[n], projection=proj) - cmap = get_colormap(cmap, parameters) - p1 = ax.contourf( - lat, - plev, - var, - # transform=ccrs.PlateCarree(), - norm=norm, - levels=levels, - cmap=cmap, - extend="both", - ) - ax.set_aspect("auto") - # ax.coastlines(lw=0.3) - if title[0] is not None: - ax.set_title(title[0], loc="left", fontdict=plotSideTitle) - if title[1] is not None: - ax.set_title(title[1], fontdict=plotTitle) - if title[2] is not None: - ax.set_title( - title[2], loc="right", fontdict=plotSideTitle - ) # loc="right" doesn't work for polar projection - # ax.set_xticks([0, 60, 120, 180, 240, 300, 359.99], crs=ccrs.PlateCarree()) - # ax.set_xticks([-180, -120, -60, 0, 60, 120, 180], crs=ccrs.PlateCarree()) - ax.set_xticks([-90, -60, -30, 0, 30, 60, 90]) # , crs=ccrs.PlateCarree()) - ax.set_xlim(-90, 90) - # lon_formatter = LongitudeFormatter( - # zero_direction_label=True, number_format='.0f') - LatitudeFormatter() - # ax.xaxis.set_major_formatter(lon_formatter) - # ax.xaxis.set_major_formatter(lat_formatter) - ax.tick_params(labelsize=8.0, direction="out", width=1) - ax.xaxis.set_ticks_position("bottom") - ax.yaxis.set_ticks_position("left") - if parameters.plot_log_plevs: - ax.set_yscale("log") - if parameters.plot_plevs: - plev_ticks = parameters.plevs - # plev_ticks = plev_ticks[::-1] - plt.yticks(plev_ticks, plev_ticks) - if not parameters.plot_log_plevs and not parameters.plot_plevs: - # Below 4 lines are to specify the pressure axis and show the 50 mb tick at the top, given default plevs np.linspace(50, 1000, 20) - if parameters.plevs == np.linspace(50, 1000, 20).tolist(): - plev_ticks = parameters.plevs - new_ticks = [plev_ticks[0]] + plev_ticks[1::2] - new_ticks = [int(x) for x in new_ticks] - plt.yticks(new_ticks, new_ticks) - plt.ylabel("pressure (mb)") - # ax.set_yscale('log') - ax.invert_yaxis() - - # Color bar - cbax = fig.add_axes((panel[n][0] + 0.6635, panel[n][1] + 0.0215, 0.0326, 0.1792)) - cbar = fig.colorbar(p1, cax=cbax) - w, h = get_ax_size(fig, cbax) - - if levels is None: - cbar.ax.tick_params(labelsize=9.0, length=0) - - else: - maxval = np.amax(np.absolute(levels[1:-1])) - if maxval < 0.01: - fmt = "%.1e" - pad = 35 - elif maxval < 10.0: - fmt = "%5.2f" - pad = 25 - elif maxval < 100.0: - fmt = "%5.1f" - pad = 25 - else: - fmt = "%6.1f" - pad = 30 - cbar.set_ticks(levels[1:-1]) - labels = [fmt % level for level in levels[1:-1]] - cbar.ax.set_yticklabels(labels, ha="right") - cbar.ax.tick_params(labelsize=9.0, pad=pad, length=0) - - # Min, Mean, Max - fig.text( - panel[n][0] + 0.6635, - panel[n][1] + 0.2107, - "Max\nMean\nMin", - ha="left", - fontdict=plotSideTitle, - ) - - # if positive Max is smaller than 0.01, use scientific notation - if stats[0] < 0.01 and stats[0] > 0: - stats_fmt = "%.e\n%.e\n%.e" - else: - stats_fmt = "%.2f\n%.2f\n%.2f" - - fig.text( - panel[n][0] + 0.7635, - panel[n][1] + 0.2107, - stats_fmt % stats[0:3], - ha="right", - fontdict=plotSideTitle, - ) - - # RMSE, CORR - if len(stats) == 5: - fig.text( - panel[n][0] + 0.6635, - panel[n][1] - 0.0105, - "RMSE\nCORR", - ha="left", - fontdict=plotSideTitle, - ) - fig.text( - panel[n][0] + 0.7635, - panel[n][1] - 0.0105, - "%.2f\n%.2f" % stats[3:5], - ha="right", - fontdict=plotSideTitle, - ) - # plt.yscale('log') - # plt.gca().invert_yaxis() - - -def plot(reference, test, diff, metrics_dict, parameter): - # Create figure, projection +# Configs for x axis ticks and x axis limits. +X_TICKS = np.array([-90, -60, -30, 0, 30, 60, 90]) +X_LIM = -90, 90 + + +def plot( + parameter: CoreParameter, + da_test: xr.DataArray, + da_ref: xr.DataArray, + da_diff: xr.DataArray, + metrics_dict: MetricsDict, +): + """Plot the variable's metrics generated by the zonal_mean_2d set. + + Parameters + ---------- + parameter : CoreParameter + The CoreParameter object containing plot configurations. + da_test : xr.DataArray + The test data. + da_ref : xr.DataArray + The reference data. + da_diff : xr.DataArray + The difference between `da_test` and `da_ref` (both are regridded to + the lower resolution of the two beforehand). + metrics_dict : Metrics + The metrics. + """ fig = plt.figure(figsize=parameter.figsize, dpi=parameter.dpi) - # proj = ccrs.PlateCarree(central_longitude=180) - proj = None + fig.suptitle(parameter.main_title, x=0.5, y=0.96, fontsize=18) + + # The variable units. + units = metrics_dict["units"] - # First two panels - min1 = metrics_dict["test"]["min"] - mean1 = metrics_dict["test"]["mean"] - max1 = metrics_dict["test"]["max"] + # Add the first subplot for test data. + min1 = metrics_dict["test"]["min"] # type: ignore + mean1 = metrics_dict["test"]["mean"] # type: ignore + max1 = metrics_dict["test"]["max"] # type: ignore - plot_panel( + _add_colormap( 0, + da_test, fig, - proj, - test, - parameter.contour_levels, - parameter.test_colormap, - (parameter.test_name_yrs, parameter.test_title, test.units), parameter, - stats=(max1, mean1, min1), + parameter.test_colormap, + parameter.contour_levels, + title=(parameter.test_name_yrs, parameter.test_title, units), # type: ignore + metrics=(max1, mean1, min1), # type: ignore ) - min2 = metrics_dict["ref"]["min"] - mean2 = metrics_dict["ref"]["mean"] - max2 = metrics_dict["ref"]["max"] - plot_panel( + # Add the second and third subplots for ref data and the differences, + # respectively. + min2 = metrics_dict["ref"]["min"] # type: ignore + mean2 = metrics_dict["ref"]["mean"] # type: ignore + max2 = metrics_dict["ref"]["max"] # type: ignore + + _add_colormap( 1, + da_ref, fig, - proj, - reference, - parameter.contour_levels, - parameter.reference_colormap, - (parameter.ref_name_yrs, parameter.reference_title, reference.units), parameter, - stats=(max2, mean2, min2), + parameter.reference_colormap, + parameter.contour_levels, + title=(parameter.ref_name_yrs, parameter.reference_title, units), # type: ignore + metrics=(max2, mean2, min2), # type: ignore ) - # Third panel - min3 = metrics_dict["diff"]["min"] - mean3 = metrics_dict["diff"]["mean"] - max3 = metrics_dict["diff"]["max"] - if parameter.diff_type == "relative": - test.units = "%" + min3 = metrics_dict["diff"]["min"] # type: ignore + mean3 = metrics_dict["diff"]["mean"] # type: ignore + max3 = metrics_dict["diff"]["max"] # type: ignore + r = metrics_dict["misc"]["rmse"] # type: ignore + c = metrics_dict["misc"]["corr"] # type: ignore - r = metrics_dict["misc"]["rmse"] - c = metrics_dict["misc"]["corr"] - plot_panel( + _add_colormap( 2, + da_diff, fig, - proj, - diff, - parameter.diff_levels, - parameter.diff_colormap, - (None, parameter.diff_title, test.units), parameter, - stats=(max3, mean3, min3, r, c), + parameter.diff_colormap, + parameter.diff_levels, + title=(None, parameter.diff_title, da_diff.attrs["units"]), # + metrics=(max3, mean3, min3, r, c), # type: ignore ) - # Figure title - fig.suptitle(parameter.main_title, x=0.5, y=0.96, fontsize=18) - - # Save figure - for f in parameter.output_format: - f = f.lower().split(".")[-1] - fnm = os.path.join( - get_output_dir(parameter.current_set, parameter), - parameter.output_file + "." + f, - ) - plt.savefig(fnm) - # Get the filename that the user has passed in and display that. - fnm = os.path.join( - get_output_dir(parameter.current_set, parameter), - parameter.output_file + "." + f, - ) - logger.info(f"Plot saved in: {fnm}") - - # Save individual subplots - for f in parameter.output_format_subplot: - fnm = os.path.join( - get_output_dir(parameter.current_set, parameter), - parameter.output_file, - ) - page = fig.get_size_inches() - i = 0 - for p in panel: - # Extent of subplot - subpage = np.array(p).reshape(2, 2) - subpage[1, :] = subpage[0, :] + subpage[1, :] - subpage = subpage + np.array(border).reshape(2, 2) - subpage = list(((subpage) * page).flatten()) # type: ignore - extent = matplotlib.transforms.Bbox.from_extents(*subpage) - # Save subplot - fname = fnm + ".%i." % (i) + f - plt.savefig(fname, bbox_inches=extent) - - orig_fnm = os.path.join( - get_output_dir(parameter.current_set, parameter), - parameter.output_file, - ) - fname = orig_fnm + ".%i." % (i) + f - logger.info(f"Sub-plot saved in: {fname}") - - i += 1 + _save_plot(fig, parameter) plt.close() + + +def _add_colormap( + subplot_num: int, + var: xr.DataArray, + fig: plt.Figure, + parameter: CoreParameter, + color_map: str, + contour_levels: List[float], + title: Tuple[Optional[str], str, str], + metrics: Tuple[float, ...], +): + lat = xc.get_dim_coords(var, axis="Y") + plev = xc.get_dim_coords(var, axis="Z") + var = var.squeeze() + + # Configure contour levels + # -------------------------------------------------------------------------- + c_levels, norm = _get_c_levels_and_norm(contour_levels) + + # Add the contour plot + # -------------------------------------------------------------------------- + ax = fig.add_axes(DEFAULT_PANEL_CFG[subplot_num], projection=None) + + contour_plot = _add_contour_plot( + ax, parameter, var, lat, plev, color_map, None, norm, c_levels + ) + + # Configure the aspect ratio and plot titles. + # -------------------------------------------------------------------------- + ax.set_aspect("auto") + _configure_titles(ax, title) + + # Configure x and y axis. + # -------------------------------------------------------------------------- + _configure_x_and_y_axes(ax, X_TICKS, None, None, parameter.current_set) + ax.set_xlim(X_LIM) + + if parameter.plot_log_plevs: + ax.set_yscale("log") + + if parameter.plot_plevs: + plev_ticks = parameter.plevs + plt.yticks(plev_ticks, plev_ticks) + + # For default plevs, specify the pressure axis and show the 50 mb tick + # at the top. + if ( + not parameter.plot_log_plevs + and not parameter.plot_plevs + and parameter.plevs == DEFAULT_PLEVS + ): + plev_ticks = parameter.plevs + new_ticks = [plev_ticks[0]] + plev_ticks[1::2] + new_ticks = [int(x) for x in new_ticks] + plt.yticks(new_ticks, new_ticks) + + plt.ylabel("pressure (mb)") + ax.invert_yaxis() + + # Add and configure the color bar. + # -------------------------------------------------------------------------- + _add_colorbar(fig, subplot_num, DEFAULT_PANEL_CFG, contour_plot, c_levels) + + # Add metrics text. + # -------------------------------------------------------------------------- + # Min, Mean, Max + _add_min_mean_max_text(fig, subplot_num, DEFAULT_PANEL_CFG, metrics) + + if len(metrics) == 5: + _add_rmse_corr_text(fig, subplot_num, DEFAULT_PANEL_CFG, metrics) diff --git a/e3sm_diags/plot/cartopy/zonal_mean_2d_stratosphere_plot.py b/e3sm_diags/plot/cartopy/zonal_mean_2d_stratosphere_plot.py index e9bab3713..004f3c93d 100644 --- a/e3sm_diags/plot/cartopy/zonal_mean_2d_stratosphere_plot.py +++ b/e3sm_diags/plot/cartopy/zonal_mean_2d_stratosphere_plot.py @@ -1,7 +1,15 @@ -from __future__ import print_function +import xarray as xr +from e3sm_diags.driver.utils.type_annotations import MetricsDict +from e3sm_diags.parameter.core_parameter import CoreParameter from e3sm_diags.plot.cartopy.zonal_mean_2d_plot import plot as base_plot -def plot(reference, test, diff, metrics_dict, parameter): - return base_plot(reference, test, diff, metrics_dict, parameter) +def plot( + parameter: CoreParameter, + da_test: xr.DataArray, + da_ref: xr.DataArray, + da_diff: xr.DataArray, + metrics_dict: MetricsDict, +): + return base_plot(parameter, da_test, da_ref, da_diff, metrics_dict) diff --git a/e3sm_diags/plot/lat_lon_plot.py b/e3sm_diags/plot/lat_lon_plot.py index 05dda58a3..d4d648f82 100644 --- a/e3sm_diags/plot/lat_lon_plot.py +++ b/e3sm_diags/plot/lat_lon_plot.py @@ -1,13 +1,31 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, List, Tuple +import cartopy.crs as ccrs +import cartopy.feature as cfeature import matplotlib import xarray as xr +import xcdat as xc +from e3sm_diags.derivations.default_regions_xr import REGION_SPECS from e3sm_diags.logger import custom_logger from e3sm_diags.parameter.core_parameter import CoreParameter -from e3sm_diags.plot.utils import _add_colormap, _save_plot +from e3sm_diags.plot.utils import ( + DEFAULT_PANEL_CFG, + _add_colorbar, + _add_contour_plot, + _add_grid_res_info, + _add_min_mean_max_text, + _add_rmse_corr_text, + _configure_titles, + _configure_x_and_y_axes, + _get_c_levels_and_norm, + _get_x_ticks, + _get_y_ticks, + _make_lon_cyclic, + _save_plot, +) if TYPE_CHECKING: from e3sm_diags.driver.lat_lon_driver import MetricsDict @@ -102,3 +120,119 @@ def plot( _save_plot(fig, parameter) plt.close() + + +def _add_colormap( + subplot_num: int, + var: xr.DataArray, + fig: plt.Figure, + parameter: CoreParameter, + color_map: str, + contour_levels: List[float], + title: Tuple[str | None, str, str], + metrics: Tuple[float, ...], +): + """Adds a colormap containing the variable data and metrics to the figure. + + This function is used by: + - `lat_lon_plot.py` + - `aerosol_aeronet_plot.py` (when refactored). + + Parameters + ---------- + subplot_num : int + The subplot number. + var : xr.DataArray + The variable to plot. + fig : plt.Figure + The figure object to add the subplot to. + parameter : CoreParameter + The CoreParameter object containing plot configurations. + color_map : str + The colormap styling to use (e.g., "cet_rainbow.rgb"). + contour_levels : List[float] + The map contour levels. + title : Tuple[str | None, str, str] + A tuple of strings to form the title of the colormap, in the format + ( years, title, units). + metrics : Tuple[float, ...] + A tuple of metrics for this subplot. + """ + var = _make_lon_cyclic(var) + lat = xc.get_dim_coords(var, axis="Y") + lon = xc.get_dim_coords(var, axis="X") + + var = var.squeeze() + + # Configure contour levels and boundary norm. + # -------------------------------------------------------------------------- + c_levels, norm = _get_c_levels_and_norm(contour_levels) + + # Get region info and X and Y plot ticks. + # -------------------------------------------------------------------------- + region_key = parameter.regions[0] + region_specs = REGION_SPECS[region_key] + + # Get the region's domain slices for latitude and longitude if set, or + # use the default value. If both are not set, then the region type is + # considered "global". + lat_slice = region_specs.get("lat", (-90, 90)) # type: ignore + lon_slice = region_specs.get("lon", (0, 360)) # type: ignore + + # Boolean flags for configuring plots. + is_global_domain = lat_slice == (-90, 90) and lon_slice == (0, 360) + is_lon_full = lon_slice == (0, 360) + + # Determine X and Y ticks using longitude and latitude domains respectively. + lon_west, lon_east = lon_slice + x_ticks = _get_x_ticks(lon_west, lon_east, is_global_domain, is_lon_full) + + lat_south, lat_north = lat_slice + y_ticks = _get_y_ticks(lat_south, lat_north) + + # Get the cartopy projection based on region info. + # -------------------------------------------------------------------------- + projection = ccrs.PlateCarree() + if is_global_domain or is_lon_full: + projection = ccrs.PlateCarree(central_longitude=180) + + # Get the figure Axes object using the projection above. + # -------------------------------------------------------------------------- + ax = fig.add_axes(DEFAULT_PANEL_CFG[subplot_num], projection=projection) + ax.set_extent([lon_west, lon_east, lat_south, lat_north], crs=projection) + contour_plot = _add_contour_plot( + ax, parameter, var, lon, lat, color_map, ccrs.PlateCarree(), norm, c_levels + ) + + # Configure the aspect ratio and coast lines. + # -------------------------------------------------------------------------- + # Full world would be aspect 360/(2*180) = 1 + ax.set_aspect((lon_east - lon_west) / (2 * (lat_north - lat_south))) + ax.coastlines(lw=0.3) + + if not is_global_domain and "RRM" in region_key: + ax.coastlines(resolution="50m", color="black", linewidth=1) + state_borders = cfeature.NaturalEarthFeature( + category="cultural", + name="admin_1_states_provinces_lakes", + scale="50m", + facecolor="none", + ) + ax.add_feature(state_borders, edgecolor="black") + + # Configure the titles, x and y axes, and colorbar. + # -------------------------------------------------------------------------- + _configure_titles(ax, title) + _configure_x_and_y_axes( + ax, x_ticks, y_ticks, ccrs.PlateCarree(), parameter.current_set + ) + _add_colorbar(fig, subplot_num, DEFAULT_PANEL_CFG, contour_plot, c_levels) + + # Add metrics text to the figure. + # -------------------------------------------------------------------------- + _add_min_mean_max_text(fig, subplot_num, DEFAULT_PANEL_CFG, metrics) + + if len(metrics) == 5: + _add_rmse_corr_text(fig, subplot_num, DEFAULT_PANEL_CFG, metrics) + + _add_grid_res_info(fig, subplot_num, region_key, lat, lon, DEFAULT_PANEL_CFG) diff --git a/e3sm_diags/plot/utils.py b/e3sm_diags/plot/utils.py index e2d9de5cb..e365c2995 100644 --- a/e3sm_diags/plot/utils.py +++ b/e3sm_diags/plot/utils.py @@ -4,15 +4,14 @@ from typing import List, Tuple import cartopy.crs as ccrs -import cartopy.feature as cfeature import matplotlib +import matplotlib.contour as mcontour import numpy as np import xarray as xr import xcdat as xc from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter from matplotlib.transforms import Bbox -from e3sm_diags.derivations.default_regions_xr import REGION_SPECS from e3sm_diags.driver.utils.general import get_output_dir from e3sm_diags.logger import custom_logger from e3sm_diags.parameter.core_parameter import CoreParameter @@ -25,11 +24,12 @@ logger = custom_logger(__name__) # Plot title and side title configurations. -PLOT_TITLE = {"fontsize": 11.5} -PLOT_SIDE_TITLE = {"fontsize": 9.5} +MAIN_TITLE_FONTSIZE = 11.5 +SECONDARY_TITLE_FONTSIZE = 9.5 # Position and sizes of subplot axes in page coordinates (0 to 1) -PANEL = [ +PanelConfig = List[Tuple[float, float, float, float]] +DEFAULT_PANEL_CFG: PanelConfig = [ (0.1691, 0.6810, 0.6465, 0.2258), (0.1691, 0.3961, 0.6465, 0.2258), (0.1691, 0.1112, 0.6465, 0.2258), @@ -37,10 +37,27 @@ # Border padding relative to subplot axes for saving individual panels # (left, bottom, right, top) in page coordinates -BORDER_PADDING = (-0.06, -0.03, 0.13, 0.03) +DEFAULT_BORDER_PADDING = (-0.06, -0.03, 0.13, 0.03) + +# Sets that use the lat_lon formatter to configure the X and Y axes of the plot. +SETS_USING_LAT_LON_FORMATTER = [ + "lat_lon", + "lat_lon_land", + "lat_lon_river", + "diurnal_cycle", + "enso_diags", + "meridional_mean_2d", + "streamflow", + "tc_analysis", +] -def _save_plot(fig: plt.Figure, parameter: CoreParameter): +def _save_plot( + fig: plt.Figure, + parameter: CoreParameter, + panel_configs: PanelConfig = DEFAULT_PANEL_CFG, + border_padding: Tuple[float, float, float, float] = DEFAULT_BORDER_PADDING, +): """Save the plot using the figure object and parameter configs. This function creates the output filename to save the plot. It also @@ -52,6 +69,13 @@ def _save_plot(fig: plt.Figure, parameter: CoreParameter): The plot figure. parameter : CoreParameter The CoreParameter with file configurations. + panel_configs : PanelConfig + A list of panel configs consisting of positions and sizes, with each + element representing a panel. By default, set to ``DEFAULT_PANEL_CFG``. + border_padding : Tuple[float, float, float, float] + A tuple of border padding configs (left, bottom, right, top) for each + panel relative to the subplot axes. By default, set to + ``DEFAULT_BORDER_PADDING``. """ for f in parameter.output_format: f = f.lower().split(".")[-1] @@ -64,9 +88,9 @@ def _save_plot(fig: plt.Figure, parameter: CoreParameter): # Save individual subplots if parameter.ref_name == "": - panels = [PANEL[0]] + panels = [panel_configs[0]] else: - panels = PANEL + panels = panel_configs for f in parameter.output_format_subplot: fnm = os.path.join( @@ -79,7 +103,7 @@ def _save_plot(fig: plt.Figure, parameter: CoreParameter): # Extent of subplot subpage = np.array(panel).reshape(2, 2) subpage[1, :] = subpage[0, :] + subpage[1, :] - subpage = subpage + np.array(BORDER_PADDING).reshape(2, 2) + subpage = subpage + np.array(border_padding).reshape(2, 2) subpage = list(((subpage) * page).flatten()) # type: ignore extent = Bbox.from_extents(*subpage) @@ -95,221 +119,25 @@ def _save_plot(fig: plt.Figure, parameter: CoreParameter): logger.info(f"Sub-plot saved in: {fname}") -def _add_colormap( - subplot_num: int, - var: xr.DataArray, - fig: plt.Figure, - parameter: CoreParameter, - color_map: str, - contour_levels: List[float], - title: Tuple[str | None, str, str], - metrics: Tuple[float, ...], -): - """Adds a colormap containing the variable data and metrics to the figure. - - This function is used by: - - `lat_lon_plot.py` - - `aerosol_aeronet_plot.py` (TODO) - - Parameters - ---------- - subplot_num : int - The subplot number. - var : xr.DataArray - The variable to plot. - fig : plt.Figure - The figure object to add the subplot to. - parameter : CoreParameter - The CoreParameter object containing plot configurations. - color_map : str - The colormap styling to use (e.g., "cet_rainbow.rgb"). - contour_levels : List[float] - The map contour levels. - title : Tuple[str | None, str, str] - A tuple of strings to form the title of the colormap, in the format - ( years, title, units). - metrics : Tuple[float, ...] - A tuple of metrics for this subplot. - """ - var = _make_lon_cyclic(var) - lat = xc.get_dim_coords(var, axis="Y") - lon = xc.get_dim_coords(var, axis="X") - - var = var.squeeze() - - # Configure contour levels - # -------------------------------------------------------------------------- - c_levels = None - norm = None - - if len(contour_levels) > 0: - c_levels = [-1.0e8] + contour_levels + [1.0e8] - norm = colors.BoundaryNorm(boundaries=c_levels, ncolors=256) - - # Configure plot tickets based on longitude and latitude. - # -------------------------------------------------------------------------- - region_key = parameter.regions[0] - region_specs = REGION_SPECS[region_key] - - # Get the region's domain slices for latitude and longitude if set, or - # use the default value. If both are not set, then the region type is - # considered "global". - lat_slice = region_specs.get("lat", (-90, 90)) # type: ignore - lon_slice = region_specs.get("lon", (0, 360)) # type: ignore - - # Boolean flags for configuring plots. - is_global_domain = lat_slice == (-90, 90) and lon_slice == (0, 360) - is_lon_full = lon_slice == (0, 360) - - # Determine X and Y ticks using longitude and latitude domains respectively. - lon_west, lon_east = lon_slice - x_ticks = _get_x_ticks(lon_west, lon_east, is_global_domain, is_lon_full) - - lat_south, lat_north = lat_slice - y_ticks = _get_y_ticks(lat_south, lat_north) - - # Add the contour plot. - # -------------------------------------------------------------------------- - projection = ccrs.PlateCarree() - if is_global_domain or is_lon_full: - projection = ccrs.PlateCarree(central_longitude=180) - - ax = fig.add_axes(PANEL[subplot_num], projection=projection) - ax.set_extent([lon_west, lon_east, lat_south, lat_north], crs=projection) - color_map = get_colormap(color_map, parameter) - p1 = ax.contourf( - lon, - lat, - var, - transform=ccrs.PlateCarree(), - norm=norm, - levels=c_levels, - cmap=color_map, - extend="both", - ) - - # Configure the aspect ratio and coast lines. - # -------------------------------------------------------------------------- - # Full world would be aspect 360/(2*180) = 1 - ax.set_aspect((lon_east - lon_west) / (2 * (lat_north - lat_south))) - ax.coastlines(lw=0.3) - - if not is_global_domain and "RRM" in region_key: - ax.coastlines(resolution="50m", color="black", linewidth=1) - state_borders = cfeature.NaturalEarthFeature( - category="cultural", - name="admin_1_states_provinces_lakes", - scale="50m", - facecolor="none", - ) - ax.add_feature(state_borders, edgecolor="black") - - # Configure the titles. - # -------------------------------------------------------------------------- - if title[0] is not None: - ax.set_title(title[0], loc="left", fontdict=PLOT_SIDE_TITLE) - if title[1] is not None: - ax.set_title(title[1], fontdict=PLOT_TITLE) - if title[2] is not None: - ax.set_title(title[2], loc="right", fontdict=PLOT_SIDE_TITLE) - - # Configure x and y axis. - # -------------------------------------------------------------------------- - ax.set_xticks(x_ticks, crs=ccrs.PlateCarree()) - ax.set_yticks(y_ticks, crs=ccrs.PlateCarree()) - - lon_formatter = LongitudeFormatter(zero_direction_label=True, number_format=".0f") - lat_formatter = LatitudeFormatter() - ax.xaxis.set_major_formatter(lon_formatter) - ax.yaxis.set_major_formatter(lat_formatter) - - ax.tick_params(labelsize=8.0, direction="out", width=1) - - ax.xaxis.set_ticks_position("bottom") - ax.yaxis.set_ticks_position("left") - - # Add and configure the color bar. - # -------------------------------------------------------------------------- - cbax = fig.add_axes( - (PANEL[subplot_num][0] + 0.6635, PANEL[subplot_num][1] + 0.0215, 0.0326, 0.1792) - ) - cbar = fig.colorbar(p1, cax=cbax) - - if c_levels is None: - cbar.ax.tick_params(labelsize=9.0, length=0) - else: - cbar.set_ticks(c_levels[1:-1]) - - label_format, pad = _get_contour_label_format_and_pad(c_levels) - labels = [label_format % level for level in c_levels[1:-1]] - cbar.ax.set_yticklabels(labels, ha="right") - cbar.ax.tick_params(labelsize=9.0, pad=pad, length=0) - - # Add metrics text. - # -------------------------------------------------------------------------- - # Min, Mean, Max - fig.text( - PANEL[subplot_num][0] + 0.6635, - PANEL[subplot_num][1] + 0.2107, - "Max\nMean\nMin", - ha="left", - fontdict=PLOT_SIDE_TITLE, - ) - - fmt_m = [] - - # Print in scientific notation if value is greater than 10^5 - for i in range(len(metrics[0:3])): - fs = "1e" if metrics[i] > 100000.0 else "2f" - fmt_m.append(fs) - - fmt_metrics = f"%.{fmt_m[0]}\n%.{fmt_m[1]}\n%.{fmt_m[2]}" - - fig.text( - PANEL[subplot_num][0] + 0.7635, - PANEL[subplot_num][1] + 0.2107, - # "%.2f\n%.2f\n%.2f" % stats[0:3], - fmt_metrics % metrics[0:3], - ha="right", - fontdict=PLOT_SIDE_TITLE, - ) - - # RMSE, CORR - if len(metrics) == 5: - fig.text( - PANEL[subplot_num][0] + 0.6635, - PANEL[subplot_num][1] - 0.0105, - "RMSE\nCORR", - ha="left", - fontdict=PLOT_SIDE_TITLE, - ) - fig.text( - PANEL[subplot_num][0] + 0.7635, - PANEL[subplot_num][1] - 0.0105, - "%.2f\n%.2f" % metrics[3:5], - ha="right", - fontdict=PLOT_SIDE_TITLE, - ) - - # Add grid resolution info. - # -------------------------------------------------------------------------- +def _add_grid_res_info(fig, subplot_num, region_key, lat, lon, panel_configs): if subplot_num == 2 and "RRM" in region_key: dlat = lat[2] - lat[1] dlon = lon[2] - lon[1] fig.text( - PANEL[subplot_num][0] + 0.4635, - PANEL[subplot_num][1] - 0.04, + panel_configs[subplot_num][0] + 0.4635, + panel_configs[subplot_num][1] - 0.04, "Resolution: {:.2f}x{:.2f}".format(dlat, dlon), ha="left", - fontdict=PLOT_SIDE_TITLE, + fontdict={"fontsize": SECONDARY_TITLE_FONTSIZE}, ) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ def _make_lon_cyclic(var: xr.DataArray): """Make the longitude axis cyclic by adding a new coordinate point with 360. This function appends a new longitude coordinate point by taking the last - coordinate point and adding 360 to it. + coordinate point and adding 360 to it. It is used for sets such as lat_lon. Parameters ---------- @@ -332,6 +160,89 @@ def _make_lon_cyclic(var: xr.DataArray): return new_var +def _get_c_levels_and_norm( + contour_levels: List[float], +) -> Tuple[List[float] | None, colors.BoundaryNorm | None]: + """Get the contour levels and boundary norm. + + If custom contour_levels are used (> 0 elements), then adjust contour_levels with + endpoints and add a boundary norm. + + Parameters + ---------- + contour_levels : List[float] + The contour levels. + + Returns + ------- + Tuple[List[float] | None, colors.BoundaryNorm | None] + A tuple of optional contour levels and boundary norm. + """ + c_levels = None + norm = None + + if len(contour_levels) > 0: + c_levels = [-1.0e8] + contour_levels + [1.0e8] + norm = colors.BoundaryNorm(boundaries=c_levels, ncolors=256) + + return c_levels, norm + + +def _add_contour_plot( + ax: matplotlib.axes.Axes, + parameter: CoreParameter, + var: xr.DataArray, + x: xr.DataArray, + y: xr.DataArray, + color_map: str, + projection: ccrs.PlateCarree | None, + norm: colors.BoundaryNorm | None, + c_levels: List[float] | None, +) -> mcontour.QuadContourSet: + """Add the contour plot to the figure axes object. + + Parameters + ---------- + ax : matplotlib.axes.Axes + The figure axes object. + parameter : CoreParameter + The CoreParameter object containing plot configurations. + var : xr.DataArray + The variable to plot. + x : xr.DataArray + The coordinates of the X axis for the plot. + y : xr.DataArray + The coordinates of the Y axis for the plot. + color_map : str + The color map file path. + projection : ccrs.PlateCarree | None + The optional cartopy projection. + norm : colors.BoundaryNorm | None + The optional norm boundaries. + c_levels : List[float] | None + The optional contour levels. + + Returns + ------- + mcontour.QuadContourSet + The contour plot object. + """ + cmap = get_colormap(color_map, parameter) + + c_plot = ax.contourf( + x, + y, + var, + cmap=cmap, + transform=projection, + norm=norm, + levels=c_levels, + extend="both", + ) + + return c_plot + + def _get_x_ticks( lon_west: float, lon_east: float, is_global_domain: bool, is_lon_full: bool ) -> np.ndarray: @@ -427,6 +338,139 @@ def _determine_tick_step(degrees_covered: float) -> int: return 1 +def _configure_titles( + ax: matplotlib.axes.Axes, + title: Tuple[str | None, str, str], + main_fontsize: float = MAIN_TITLE_FONTSIZE, + secondary_fontsize: float = SECONDARY_TITLE_FONTSIZE, +): + """Configure the axes titles. + + Parameters + ---------- + ax : matplotlib.axes.Axes + The figure axes object. + title : Tuple[str | None, str, str] + A tuple of strings to form the title of the colormap, in the format + ( years, title, units). + main_fontsize : float + The main title font size, by default 11.5. + secondary_fontsize : float + The secondary title font sizes, by default 9.5. + + Returns + ------- + matplotlib.axes.Axes + The axes objects. + """ + if title[0] is not None: + ax.set_title(title[0], loc="left", fontdict={"fontsize": secondary_fontsize}) + if title[1] is not None: + ax.set_title(title[1], fontdict={"fontsize": main_fontsize}) + if title[2] is not None: + # NOTE: loc="right" doesn't work for polar projection + ax.set_title(title[2], loc="right", fontdict={"fontsize": secondary_fontsize}) + + +def _configure_x_and_y_axes( + ax: matplotlib.axes.Axes, + x_ticks: np.ndarray, + y_ticks: np.ndarray | None, + projection: ccrs.PlateCarree | None, + set_name: str, +): + """Configure the X and Y axes. + + Parameters + ---------- + ax : matplotlib.axes.Axes + The figure axes object. + x_ticks : np.ndarray + The array of X ticks. + y_ticks : np.ndarray | None + The optional array of Y ticks. Some set plotters pass None to configure + the Y axis ticks using other ticks such as Z axis plevs instead. + projection : ccrs.PlateCarree | None + The optional cartopy projection to use for X and Y ticks. + set_name : set_name + The name of the current set which determines whether the latitude and + longitude major formatters are used. + """ + # For `ax.set_xticks` and `ax.set_yticks`, `crs` cannot be `None` and we + # must split up arguments passed to these methods using a conditional + # statement. Otherwise, this error is raised: `ValueError: Incorrect use of + # keyword argument 'crs'. Keyword arguments other than 'minor' modify the + # text labels and can only be used if 'labels' are passed as well.` + if projection is not None: + ax.set_xticks(x_ticks, crs=projection) + + if y_ticks is not None: + ax.set_yticks(y_ticks, crs=projection) + else: + ax.set_xticks(x_ticks) + + if y_ticks is not None: + ax.set_yticks(y_ticks) + + ax.tick_params(labelsize=8.0, direction="out", width=1) + + ax.xaxis.set_ticks_position("bottom") + ax.yaxis.set_ticks_position("left") + + if set_name in SETS_USING_LAT_LON_FORMATTER: + lon_formatter = LongitudeFormatter( + zero_direction_label=True, number_format=".0f" + ) + lat_formatter = LatitudeFormatter() + ax.xaxis.set_major_formatter(lon_formatter) + ax.yaxis.set_major_formatter(lat_formatter) + + +def _add_colorbar( + fig: plt.Figure, + subplot_num: int, + panel_configs: PanelConfig, + contour_plot: mcontour.QuadContourSet, + c_levels: List[float] | None, +): + """Configure the colorbar on a colormap. + + Parameters + ---------- + fig : plt.Figure + The figure object. + subplot_num : int + The subplot number. + panel_configs : PanelConfig + A list of panel configs consisting of positions and sizes, with each + element representing a panel. + contour_plot : mcontour.QuadContourSet + The contour plot object. + c_levels : List[float] | None + The optional contour levels used to configure the colorbar. + """ + cbax = fig.add_axes( + ( + panel_configs[subplot_num][0] + 0.6635, + panel_configs[subplot_num][1] + 0.0215, + 0.0326, + 0.1792, + ) + ) + + cbar = fig.colorbar(contour_plot, cax=cbax) + + if c_levels is None: + cbar.ax.tick_params(labelsize=9.0, length=0) + else: + cbar.set_ticks(c_levels[1:-1]) + + label_format, pad = _get_contour_label_format_and_pad(c_levels) + labels = [label_format % level for level in c_levels[1:-1]] + cbar.ax.set_yticklabels(labels, ha="right") + cbar.ax.tick_params(labelsize=9.0, pad=pad, length=0) + + def _get_contour_label_format_and_pad(c_levels: List[float]) -> Tuple[str, int]: """Get the label format and padding for each contour level. @@ -442,7 +486,10 @@ def _get_contour_label_format_and_pad(c_levels: List[float]) -> Tuple[str, int]: """ maxval = np.amax(np.absolute(c_levels[1:-1])) - if maxval < 0.2: + if maxval < 0.01: + fmt = "%.1e" + pad = 35 + elif maxval < 0.2: fmt = "%5.3f" pad = 28 elif maxval < 10.0: @@ -459,3 +506,129 @@ def _get_contour_label_format_and_pad(c_levels: List[float]) -> Tuple[str, int]: pad = 30 return fmt, pad + + +def _add_min_mean_max_text( + fig: plt.Figure, + subplot_num: int, + panel_configs: PanelConfig, + metrics: Tuple[float, ...], + set_name: str | None = None, + fontsize: float = SECONDARY_TITLE_FONTSIZE, +): + """Add min, mean, and max text to the figure. + + Parameters + ---------- + fig : plt.Figure + The figure object. + subplot_num : int + The subplot number. + panel_configs : PanelConfig + A list of panel configs consisting of positions and sizes, with each + element representing a panel. + metrics : Tuple[float, ...] + The tuple of metrics, with the first three elements being max, mean, + and min. + set_name : str | None + The optional set name used to determine float format, by default None. + fontsize : float + The text font size, by default 9.5. + """ + fontdict = {"fontsize": fontsize} + + fig.text( + panel_configs[subplot_num][0] + 0.6635, + panel_configs[subplot_num][1] + 0.2107, + "Max\nMean\nMin", + ha="left", + fontdict=fontdict, + ) + + fmt_metrics = _get_float_format(metrics, set_name) + + fig.text( + panel_configs[subplot_num][0] + 0.7635, + panel_configs[subplot_num][1] + 0.2107, + fmt_metrics % metrics[0:3], + ha="right", + fontdict=fontdict, + ) + + +def _get_float_format(metrics: Tuple[float, ...], set_name: str | None) -> str: + """Get the float format for string text based on decimal places of metrics. + + Parameters + ---------- + metrics : Tuple[float, ...] + The tuple of metrics, with the first three elements being max, mean, and + min. + set_name : str | None + The optional name of the set. + + Returns + ------- + str + The float format. + """ + # FIXME: This conditional code was ported over from two plot functions and + # can be implemented better. + if set_name in ["zonal_mean_2d", "zonal_mean_2d_stratosphere"]: + # if positive Max is smaller than 0.01, use scientific notation + if metrics[0] < 0.01 and metrics[0] > 0: + float_format = "%.e\n%.e\n%.e" + else: + float_format = "%.2f\n%.2f\n%.2f" + else: + fmt_m = [] + + # Print in scientific notation if value is greater than 10^5 + for i in range(len(metrics[0:3])): + fs = "1e" if metrics[i] > 100000.0 else "2f" + fmt_m.append(fs) + + float_format = f"%.{fmt_m[0]}\n%.{fmt_m[1]}\n%.{fmt_m[2]}" + + return float_format + + +def _add_rmse_corr_text( + fig: plt.Figure, + subplot_num: int, + panel_configs: PanelConfig, + metrics: Tuple[float, ...], + fontsize: float = SECONDARY_TITLE_FONTSIZE, +): + """Add RMSE and CORR metrics text to the figure. + + Parameters + ---------- + fig : plt.Figure + The figure object. + subplot_num : int + The subplot number. + panel_configs : PanelConfig + A list of panel configs consisting of positions and sizes, with each + element representing a panel. + metrics : Tuple[float, ...] + The tuple of metrics, with the last two elements being RMSE and CORR. + fontsize : float + The text font size, by default 9.5. + """ + fontdict = {"fontsize": fontsize} + + fig.text( + panel_configs[subplot_num][0] + 0.6635, + panel_configs[subplot_num][1] - 0.0105, + "RMSE\nCORR", + ha="left", + fontdict=fontdict, + ) + fig.text( + panel_configs[subplot_num][0] + 0.7635, + panel_configs[subplot_num][1] - 0.0105, + "%.2f\n%.2f" % metrics[3:5], + ha="right", + fontdict=fontdict, + ) diff --git a/tests/e3sm_diags/driver/utils/test_dataset_xr.py b/tests/e3sm_diags/driver/utils/test_dataset_xr.py index 4b112161e..664b8df76 100644 --- a/tests/e3sm_diags/driver/utils/test_dataset_xr.py +++ b/tests/e3sm_diags/driver/utils/test_dataset_xr.py @@ -778,7 +778,7 @@ def test_returns_climo_dataset_using_derived_var_directly_from_dataset_and_repla # time scalar variable using Xarray because it just throws the error # below. We might need to use another library like netCDF4 to create # a dummy dataset. - ds_precst = xr.Dataset( + ds_src = xr.Dataset( coords={ **spatial_coords, }, @@ -788,7 +788,7 @@ def test_returns_climo_dataset_using_derived_var_directly_from_dataset_and_repla dims="time", data=0, ), - "PRECST": xr.DataArray( + "SOURCE_VAR": xr.DataArray( xr.DataArray( data=np.array( [ @@ -806,17 +806,17 @@ def test_returns_climo_dataset_using_derived_var_directly_from_dataset_and_repla "ref", "climo", self.data_path, "2000", "2001" ) parameter.ref_file = "pr_200001_200112.nc" - ds_precst.to_netcdf(f"{self.data_path}/{parameter.ref_file}") + ds_src.to_netcdf(f"{self.data_path}/{parameter.ref_file}") ds = Dataset(parameter, data_type="ref") - result = ds.get_climo_dataset("PRECST", season="ANN") - expected = ds_precst.squeeze(dim="time").drop_vars("time") + result = ds.get_climo_dataset("SO", season="ANN") + expected = ds_src.squeeze(dim="time").drop_vars("time") xr.testing.assert_identical(result, expected) def test_returns_climo_dataset_using_derived_var_directly_from_dataset(self): - ds_precst = xr.Dataset( + ds_src = xr.Dataset( coords={ **spatial_coords, "time": xr.DataArray( @@ -839,7 +839,7 @@ def test_returns_climo_dataset_using_derived_var_directly_from_dataset(self): }, data_vars={ **spatial_bounds, - "PRECST": xr.DataArray( + "SOURCE_VAR": xr.DataArray( xr.DataArray( data=np.array( [ @@ -857,12 +857,12 @@ def test_returns_climo_dataset_using_derived_var_directly_from_dataset(self): "ref", "climo", self.data_path, "2000", "2001" ) parameter.ref_file = "pr_200001_200112.nc" - ds_precst.to_netcdf(f"{self.data_path}/{parameter.ref_file}") + ds_src.to_netcdf(f"{self.data_path}/{parameter.ref_file}") ds = Dataset(parameter, data_type="ref") - result = ds.get_climo_dataset("PRECST", season="ANN") - expected = ds_precst.squeeze(dim="time").drop_vars("time") + result = ds.get_climo_dataset("SOURCE_VAR", season="ANN") + expected = ds_src.squeeze(dim="time").drop_vars("time") xr.testing.assert_identical(result, expected) diff --git a/tests/e3sm_diags/driver/utils/test_regrid.py b/tests/e3sm_diags/driver/utils/test_regrid.py index 870de6a6a..88ce8bb35 100644 --- a/tests/e3sm_diags/driver/utils/test_regrid.py +++ b/tests/e3sm_diags/driver/utils/test_regrid.py @@ -334,7 +334,7 @@ def test_regrids_hybrid_levels_to_pressure_levels_with_existing_z_bounds(self): # updating the arrays and attributes of data variables and coordinates. expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"]) expected["so"].data[:] = np.nan - expected["so"].attrs["units"] = "mb" + expected["so"].attrs["units"] = "ppt" expected["lev"].attrs = { "axis": "Z", "coordinate": "vertical", @@ -364,7 +364,7 @@ def test_regrids_hybrid_levels_to_pressure_levels_with_generated_z_bounds(self): # updating the arrays and attributes of data variables and coordinates. expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"]) expected["so"].data[:] = np.nan - expected["so"].attrs["units"] = "mb" + expected["so"].attrs["units"] = "ppt" expected["lev"].attrs = { "axis": "Z", "coordinate": "vertical", @@ -393,7 +393,7 @@ def test_regrids_hybrid_levels_to_pressure_levels_with_Pa_units(self): # updating the arrays and attributes of data variables and coordinates. expected = ds.sel(lev=[800, 200]).drop_vars(["ps", "hyam", "hybm"]) expected["so"].data[:] = np.nan - expected["so"].attrs["units"] = "mb" + expected["so"].attrs["units"] = "ppt" expected["lev"].attrs = { "axis": "Z", "coordinate": "vertical", diff --git a/tests/e3sm_diags/fixtures.py b/tests/e3sm_diags/fixtures.py index bf0b7249e..4f6da58da 100644 --- a/tests/e3sm_diags/fixtures.py +++ b/tests/e3sm_diags/fixtures.py @@ -63,6 +63,7 @@ def generate_lev_dataset( name="so", data=np.ones((5, 4, 4, 4)), coords={"time": time_decoded, "lev": lev, "lat": lat, "lon": lon}, + attrs={"units": "ppt"}, ), }, coords={ diff --git a/tests/e3sm_diags/metrics/test_metrics.py b/tests/e3sm_diags/metrics/test_metrics.py index 141b5e0d6..fd4e8a339 100644 --- a/tests/e3sm_diags/metrics/test_metrics.py +++ b/tests/e3sm_diags/metrics/test_metrics.py @@ -3,7 +3,7 @@ import xarray as xr from xarray.testing import assert_allclose -from e3sm_diags.metrics.metrics import correlation, get_weights, rmse, spatial_avg, std +from e3sm_diags.metrics.metrics import _get_weights, correlation, rmse, spatial_avg, std class TestGetWeights: @@ -11,23 +11,30 @@ class TestGetWeights: def setup(self): self.ds = xr.Dataset( coords={ + "lev": xr.DataArray( + data=[0, 1, 2], + dims="lev", + attrs={"bounds": "lev_bnds", "axis": "Z"}, + ), "lat": xr.DataArray( data=[0, 1], dims="lat", attrs={"bounds": "lat_bnds", "axis": "Y"} ), "lon": xr.DataArray( data=[0, 1], dims="lon", attrs={"bounds": "lon_bnds", "axis": "X"} ), - "time": xr.DataArray(data=[1, 2, 3], dims="time"), }, ) self.ds["ts"] = xr.DataArray( data=np.array([[[1, 2], [1, 2]], [[np.nan, 1], [1, 2]], [[2, 1], [1, 2]]]), - coords={"lat": self.ds.lat, "lon": self.ds.lon, "time": self.ds.time}, - dims=["time", "lat", "lon"], + coords={"lat": self.ds.lat, "lon": self.ds.lon, "lev": self.ds.lev}, + dims=["lev", "lat", "lon"], ) # Bounds are used to generate weights. + self.ds["lev_bnds"] = xr.DataArray( + [[0, 1], [1, 2], [2, 3]], dims=["lev", "bnds"] + ) self.ds["lat_bnds"] = xr.DataArray([[0, 1], [1, 2]], dims=["lat", "bnds"]) self.ds["lon_bnds"] = xr.DataArray([[0, 1], [1, 2]], dims=["lon", "bnds"]) @@ -39,10 +46,48 @@ def test_returns_weights_for_x_y_axes(self): ), coords={"lon": self.ds.lon, "lat": self.ds.lat}, ) - result = get_weights(self.ds) + result = _get_weights(self.ds, "ts", axis=["X", "Y"]) assert_allclose(expected, result) + def test_returns_weights_for_x_y_z_axes(self): + expected = xr.DataArray( + dims=["lon", "lat", "lev"], + data=np.array( + [ + [ + [0.01745241, 0.01745241, 0.01745241], + [0.01744709, 0.01744709, 0.01744709], + ], + [ + [0.01745241, 0.01745241, 0.01745241], + [0.01744709, 0.01744709, 0.01744709], + ], + ] + ), + coords={"lon": self.ds.lon, "lat": self.ds.lat, "lev": self.ds.lev}, + ) + result = _get_weights(self.ds, "ts", axis=["X", "Y", "Z"]) + + assert_allclose(expected, result) + + def test_returns_weights_for_z_axis(self): + expected = xr.DataArray( + dims="lev", + data=np.array([1, 1, 1], dtype="float64"), + coords={"lev": self.ds.lev}, + ) + result = _get_weights(self.ds, "ts", axis=["Z"]) + + assert_allclose(expected, result) + + def test_raises_error_if_z_bounds_not_found(self): + ds = self.ds.copy() + ds = ds.drop_vars("lev_bnds") + + with pytest.raises(RuntimeError): + _get_weights(ds, "ts", axis=["Z"]) + class TestSpatialAvg: @pytest.fixture(autouse=True)