diff --git a/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d.cfg b/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d.cfg new file mode 100644 index 0000000000..5c49d6a73a --- /dev/null +++ b/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d.cfg @@ -0,0 +1,87 @@ +[#] +sets = ["meridional_mean_2d"] +case_id = "MERRA2" +variables = ["T"] +ref_name = "MERRA2" +reference_name = "MERRA2 Reanalysis" +seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +contour_levels = [180,185,190,200,210,220,230,240,250,260,270,280,290,295,300] +diff_levels = [-7,-6,-5,-4,-3,-2,-1,1,2,3,4,5,6,7] + +[#] +sets = ["meridional_mean_2d"] +case_id = "MERRA2" +variables = ["U"] +ref_name = "MERRA2" +reference_name = "MERRA2 Reanalysis" +seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +test_colormap = "PiYG_r" +reference_colormap = "PiYG_r" +contour_levels = [-20, -15, -10, -8, -5, -3, -1, 1, 3, 5, 8, 10, 15, 20] +diff_levels = [ -5, -4, -3, -2, -1,-0.5, -0.2, 0.2, 0.5, 1, 2, 3, 4, 5] + +# [#] +# sets = ["meridional_mean_2d"] +# case_id = "MERRA2" +# variables = ["RELHUM"] +# ref_name = "MERRA2" +# reference_name = "MERRA2 Reanalysis" +# seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +# contour_levels = [5,10,15,20,25,30,40,50,60,70,75,80,85,90,95] +# diff_levels = [-15,-12,-9,-6,-3,-1,1,3,6,9,12,15] + +# [#] +# sets = ["meridional_mean_2d"] +# case_id = "MERRA2" +# variables = ["OMEGA"] +# ref_name = "MERRA2" +# reference_name = "MERRA2" +# seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +# test_colormap = "PiYG_r" +# reference_colormap = "PiYG_r" +# contour_levels = [-35,-30,-25,-20,-15,-10,-5,5,10,15,20,25,30,35] +# diff_levels = [-20,-15,-10,-8,-6,-4,-2,2,4,6,8,10,15,20] + +# [#] +# sets = ["meridional_mean_2d"] +# case_id = "ERA5" +# variables = ["T"] +# ref_name = "ERA5" +# reference_name = "ERA5 Reanalysis" +# seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +# contour_levels = [180,185,190,200,210,220,230,240,250,260,270,280,290,295,300] +# diff_levels = [-7,-6,-5,-4,-3,-2,-1,1,2,3,4,5,6,7] + +# [#] +# sets = ["meridional_mean_2d"] +# case_id = "ERA5" +# variables = ["U"] +# ref_name = "ERA5" +# reference_name = "ERA5 Reanalysis" +# seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +# test_colormap = "PiYG_r" +# reference_colormap = "PiYG_r" +# contour_levels = [-20, -15, -10, -8, -5, -3, -1, 1, 3, 5, 8, 10, 15, 20] +# diff_levels = [ -5, -4, -3, -2, -1,-0.5, -0.2, 0.2, 0.5, 1, 2, 3, 4, 5] + +# [#] +# sets = ["meridional_mean_2d"] +# case_id = "ERA5" +# variables = ["OMEGA"] +# ref_name = "ERA5" +# reference_name = "ERA5 Reanalysis" +# seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +# test_colormap = "PiYG_r" +# reference_colormap = "PiYG_r" +# contour_levels = [-35,-30,-25,-20,-15,-10,-5,5,10,15,20,25,30,35] +# diff_levels = [-20,-15,-10,-8,-6,-4,-2,2,4,6,8,10,15,20] + +# [#] +# sets = ["meridional_mean_2d"] +# case_id = "ERA5" +# variables = ["RELHUM"] +# ref_name = "ERA5" +# reference_name = "ERA5 Reanalysis" +# seasons = ["ANN", "DJF", "MAM", "JJA", "SON"] +# contour_levels = [5,10,15,20,25,30,40,50,60,70,75,80,85,90,95] +# diff_levels = [-15,-12,-9,-6,-3,-1,1,3,6,9,12,15] diff --git a/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d_cdat_regression_test_netcdf.ipynb b/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d_cdat_regression_test_netcdf.ipynb new file mode 100644 index 0000000000..3e1f841a2f --- /dev/null +++ b/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d_cdat_regression_test_netcdf.ipynb @@ -0,0 +1,584 @@ +{ + "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 netcdf4 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": 3, + "metadata": {}, + "outputs": [], + "source": [ + "import glob\n", + "\n", + "import numpy as np\n", + "import xarray as xr\n", + "from e3sm_diags.derivations.derivations import DERIVED_VARIABLES\n", + "\n", + "\n", + "# TODO: Update SET_NAME and SET_DIR\n", + "SET_NAME = \"meridional_mean_2d\"\n", + "SET_DIR = \"657-meridional-mean-2d\"\n", + "\n", + "DEV_PATH = f\"/global/cfs/cdirs/e3sm/www/cdat-migration-fy24/{SET_DIR}/{SET_NAME}/**\"\n", + "DEV_GLOB = sorted(glob.glob(DEV_PATH + \"/*.nc\"))\n", + "DEV_NUM_FILES = len(DEV_GLOB)\n", + "\n", + "MAIN_PATH = f\"/global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/{SET_NAME}/**\"\n", + "MAIN_GLOB = sorted(glob.glob(MAIN_PATH + \"/*.nc\"))\n", + "MAIN_NUM_FILES = len(MAIN_GLOB)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "def _check_if_files_found():\n", + " if DEV_NUM_FILES == 0 or MAIN_NUM_FILES == 0:\n", + " raise IOError(\n", + " \"No files found at DEV_PATH and/or MAIN_PATH. \"\n", + " f\"Please check {DEV_PATH} and {MAIN_PATH}.\"\n", + " )\n", + "\n", + "\n", + "def _check_if_matching_filecount():\n", + " if DEV_NUM_FILES != MAIN_NUM_FILES:\n", + " raise IOError(\n", + " \"Number of files do not match at DEV_PATH and MAIN_PATH \"\n", + " f\"({DEV_NUM_FILES} vs. {MAIN_NUM_FILES}).\"\n", + " )\n", + "\n", + " print(f\"Matching file count ({DEV_NUM_FILES} and {MAIN_NUM_FILES}).\")\n", + "\n", + "\n", + "def _check_if_missing_files():\n", + " missing_count = 0\n", + "\n", + " for fp_main in MAIN_GLOB:\n", + " fp_dev = fp_main.replace(SET_DIR, \"main\")\n", + "\n", + " if fp_dev not in MAIN_GLOB:\n", + " print(f\"No production file found to compare with {fp_dev}!\")\n", + " missing_count += 1\n", + "\n", + " for fp_dev in DEV_GLOB:\n", + " fp_main = fp_main.replace(\"main\", SET_DIR)\n", + "\n", + " if fp_main not in DEV_GLOB:\n", + " print(f\"No development file found to compare with {fp_main}!\")\n", + " missing_count += 1\n", + "\n", + " print(f\"Number of files missing: {missing_count}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "def _get_relative_diffs():\n", + " # We are mainly focusing on relative tolerance here (in percentage terms).\n", + " atol = 0\n", + " rtol = 1e-5\n", + "\n", + " for fp_main in MAIN_GLOB:\n", + " if \"test.nc\" in fp_main or \"ref.nc\" in fp_main:\n", + " fp_dev = fp_main.replace(\"main\", SET_DIR)\n", + "\n", + " print(\"Comparing:\")\n", + " print(f\" * {fp_dev}\")\n", + " print(f\" * {fp_main}\")\n", + "\n", + " ds1 = xr.open_dataset(fp_dev)\n", + " ds2 = xr.open_dataset(fp_main)\n", + "\n", + " var_key = fp_main.split(\"-\")[-3]\n", + " # for 3d vars such as T-200\n", + " var_key.isdigit()\n", + " if var_key.isdigit():\n", + " var_key = fp_main.split(\"-\")[-4]\n", + "\n", + " print(f\" * var_key: {var_key}\")\n", + "\n", + " dev_data = _get_var_data(ds1, var_key)\n", + " main_data = _get_var_data(ds2, var_key)\n", + "\n", + " if dev_data is None or main_data is None:\n", + " print(\" * Could not find variable key in the dataset(s)\")\n", + " continue\n", + "\n", + " try:\n", + " np.testing.assert_allclose(\n", + " dev_data,\n", + " main_data,\n", + " atol=atol,\n", + " rtol=rtol,\n", + " )\n", + " except (KeyError, AssertionError) as e:\n", + " print(f\" {e}\")\n", + " else:\n", + " print(f\" * All close and within relative tolerance ({rtol})\")\n", + "\n", + "\n", + "def _get_var_data(ds: xr.Dataset, var_key: str) -> np.ndarray:\n", + " \"\"\"Get the variable data using a list of matching keys.\n", + "\n", + " The `main` branch saves the dataset using the original variable name,\n", + " while the dev branch saves the variable with the derived variable name.\n", + " The dev branch is performing the expected behavior here.\n", + "\n", + " Parameters\n", + " ----------\n", + " ds : xr.Dataset\n", + " _description_\n", + " var_key : str\n", + " _description_\n", + "\n", + " Returns\n", + " -------\n", + " np.ndarray\n", + " _description_\n", + " \"\"\"\n", + "\n", + " data = None\n", + "\n", + " var_keys = DERIVED_VARIABLES[var_key].keys()\n", + " var_keys = [var_key] + list(sum(var_keys, ()))\n", + "\n", + " for key in var_keys:\n", + " if key in ds.data_vars.keys():\n", + " data = ds[key].values\n", + " break\n", + "\n", + " return data" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 1. Check for matching and equal number of files\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "_check_if_files_found()" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Number of files missing: 0\n" + ] + } + ], + "source": [ + "_check_if_missing_files()" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Matching file count (48 and 48).\n" + ] + } + ], + "source": [ + "_check_if_matching_filecount()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 2 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": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-OMEGA-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-OMEGA-ANN-global_ref.nc\n", + " * var_key: OMEGA\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 6 / 24480 (0.0245%)\n", + "Max absolute difference: 4.56288177e-07\n", + "Max relative difference: 5.34042702e-05\n", + " x: array([[ 0.784679, 0.796702, 0.792179, ..., 0.804814, 0.773629,\n", + " 0.774743],\n", + " [ 1.379984, 1.417249, 1.4173 , ..., 1.432473, 1.391159,...\n", + " y: array([[ 0.784679, 0.796702, 0.792179, ..., 0.804814, 0.773629,\n", + " 0.774743],\n", + " [ 1.379984, 1.417249, 1.4173 , ..., 1.432473, 1.391159,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-OMEGA-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-OMEGA-ANN-global_test.nc\n", + " * var_key: OMEGA\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-OMEGA-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-OMEGA-JJA-global_ref.nc\n", + " * var_key: OMEGA\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 3 / 24480 (0.0123%)\n", + "Max absolute difference: 5.97006945e-07\n", + "Max relative difference: 1.49039703e-05\n", + " x: array([[1.08258 , 1.069207, 1.07484 , ..., 1.090696, 1.087342, 1.09467 ],\n", + " [1.563664, 1.593989, 1.628474, ..., 1.553196, 1.556833, 1.570584],\n", + " [2.220292, 2.307285, 2.364905, ..., 2.215879, 2.220069, 2.228939],...\n", + " y: array([[1.08258 , 1.069207, 1.07484 , ..., 1.090696, 1.087342, 1.09467 ],\n", + " [1.563664, 1.593989, 1.628474, ..., 1.553196, 1.556833, 1.570584],\n", + " [2.220292, 2.307285, 2.364905, ..., 2.215879, 2.220069, 2.228939],...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-OMEGA-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-OMEGA-JJA-global_test.nc\n", + " * var_key: OMEGA\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 2 / 6120 (0.0327%)\n", + "Max absolute difference: 3.39774408e-07\n", + "Max relative difference: 0.00015039\n", + " x: array([[ 1.150119, 1.168407, 1.16202 , ..., 1.153911, 1.138019,\n", + " 1.152851],\n", + " [ 1.567948, 1.599667, 1.601973, ..., 1.42647 , 1.435826,...\n", + " y: array([[ 1.150119, 1.168407, 1.16202 , ..., 1.153911, 1.138019,\n", + " 1.152851],\n", + " [ 1.567948, 1.599667, 1.601973, ..., 1.42647 , 1.435826,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-RELHUM-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-RELHUM-ANN-global_ref.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-RELHUM-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-RELHUM-ANN-global_test.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-RELHUM-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-RELHUM-JJA-global_ref.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-RELHUM-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-RELHUM-JJA-global_test.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-T-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-T-ANN-global_ref.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-T-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-T-ANN-global_test.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-T-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-T-JJA-global_ref.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-T-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-T-JJA-global_test.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-U-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-U-ANN-global_ref.nc\n", + " * var_key: U\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 24480 (0.00408%)\n", + "Max absolute difference: 8.52794937e-08\n", + "Max relative difference: 1.73621834e-05\n", + " x: array([[10.269112, 10.260021, 10.25072 , ..., 10.296745, 10.28736 ,\n", + " 10.278145],\n", + " [12.365973, 12.353776, 12.341225, ..., 12.402582, 12.390261,...\n", + " y: array([[10.269112, 10.260021, 10.25072 , ..., 10.296745, 10.28736 ,\n", + " 10.278145],\n", + " [12.365973, 12.353776, 12.341225, ..., 12.402582, 12.390261,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-U-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-U-ANN-global_test.nc\n", + " * var_key: U\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 6120 (0.0163%)\n", + "Max absolute difference: 1.45889873e-07\n", + "Max relative difference: 1.08735566e-05\n", + " x: array([[ 9.712424, 9.692257, 9.675111, ..., 9.79681 , 9.779955,\n", + " 9.765068],\n", + " [12.482069, 12.457917, 12.434379, ..., 12.581674, 12.560391,...\n", + " y: array([[ 9.712424, 9.692257, 9.675112, ..., 9.79681 , 9.779955,\n", + " 9.765068],\n", + " [12.482069, 12.457917, 12.434379, ..., 12.581674, 12.560391,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-U-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-U-JJA-global_ref.nc\n", + " * var_key: U\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 24480 (0.00408%)\n", + "Max absolute difference: 8.38617442e-08\n", + "Max relative difference: 1.96239722e-05\n", + " x: array([[ 8.311289, 8.302058, 8.292533, ..., 8.338692, 8.329687,\n", + " 8.320503],\n", + " [10.096993, 10.085039, 10.072245, ..., 10.131809, 10.120227,...\n", + " y: array([[ 8.311289, 8.302058, 8.292533, ..., 8.338692, 8.329687,\n", + " 8.320503],\n", + " [10.096994, 10.085039, 10.072245, ..., 10.131809, 10.120227,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/ERA5/ERA5-U-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/ERA5/ERA5-U-JJA-global_test.nc\n", + " * var_key: U\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_ref.nc\n", + " * var_key: OMEGA\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 2 / 9792 (0.0204%)\n", + "Max absolute difference: 3.15701861e-07\n", + "Max relative difference: 2.42041494e-05\n", + " x: array([[0.914233, 0.906507, 0.900358, ..., 0.954717, 0.945675, 0.928962],\n", + " [1.617818, 1.601587, 1.580476, ..., 1.631384, 1.64521 , 1.638424],\n", + " [2.321404, 2.296666, 2.260593, ..., 2.308052, 2.344744, 2.347885],...\n", + " y: array([[0.914233, 0.906507, 0.900358, ..., 0.954717, 0.945675, 0.928962],\n", + " [1.617818, 1.601587, 1.580476, ..., 1.631384, 1.64521 , 1.638424],\n", + " [2.321404, 2.296666, 2.260593, ..., 2.308052, 2.344744, 2.347885],...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-OMEGA-ANN-global_test.nc\n", + " * var_key: OMEGA\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_ref.nc\n", + " * var_key: OMEGA\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 2 / 9792 (0.0204%)\n", + "Max absolute difference: 6.93198768e-07\n", + "Max relative difference: 5.40329309e-05\n", + " x: array([[ 1.283442, 1.27651 , 1.25776 , ..., 1.292833, 1.29537 ,\n", + " 1.287421],\n", + " [ 1.93963 , 1.950937, 1.957762, ..., 1.901146, 1.91872 ,...\n", + " y: array([[ 1.283442, 1.27651 , 1.25776 , ..., 1.292833, 1.29537 ,\n", + " 1.287421],\n", + " [ 1.93963 , 1.950937, 1.957762, ..., 1.901146, 1.91872 ,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-OMEGA-JJA-global_test.nc\n", + " * var_key: OMEGA\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 2 / 6120 (0.0327%)\n", + "Max absolute difference: 3.39774408e-07\n", + "Max relative difference: 0.00015039\n", + " x: array([[ 1.150119, 1.168407, 1.16202 , ..., 1.153911, 1.138019,\n", + " 1.152851],\n", + " [ 1.567948, 1.599667, 1.601973, ..., 1.42647 , 1.435826,...\n", + " y: array([[ 1.150119, 1.168407, 1.16202 , ..., 1.153911, 1.138019,\n", + " 1.152851],\n", + " [ 1.567948, 1.599667, 1.601973, ..., 1.42647 , 1.435826,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_ref.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-RELHUM-ANN-global_test.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_ref.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-RELHUM-JJA-global_test.nc\n", + " * var_key: RELHUM\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-T-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-T-ANN-global_ref.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-T-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-T-ANN-global_test.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-T-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-T-JJA-global_ref.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-T-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-T-JJA-global_test.nc\n", + " * var_key: T\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-U-ANN-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-U-ANN-global_ref.nc\n", + " * var_key: U\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 9792 (0.0102%)\n", + "Max absolute difference: 1.06174015e-07\n", + "Max relative difference: 1.43005542e-05\n", + " x: array([[10.623722, 10.599673, 10.576724, ..., 10.698371, 10.673217,\n", + " 10.64816 ],\n", + " [12.345496, 12.317517, 12.289864, ..., 12.432171, 12.403032,...\n", + " y: array([[10.623722, 10.599673, 10.576724, ..., 10.698371, 10.673217,\n", + " 10.64816 ],\n", + " [12.345496, 12.317517, 12.289864, ..., 12.432171, 12.403032,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-U-ANN-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-U-ANN-global_test.nc\n", + " * var_key: U\n", + " \n", + "Not equal to tolerance rtol=1e-05, atol=0\n", + "\n", + "Mismatched elements: 1 / 6120 (0.0163%)\n", + "Max absolute difference: 1.45889873e-07\n", + "Max relative difference: 1.08735566e-05\n", + " x: array([[ 9.712424, 9.692257, 9.675111, ..., 9.79681 , 9.779955,\n", + " 9.765068],\n", + " [12.482069, 12.457917, 12.434379, ..., 12.581674, 12.560391,...\n", + " y: array([[ 9.712424, 9.692257, 9.675112, ..., 9.79681 , 9.779955,\n", + " 9.765068],\n", + " [12.482069, 12.457917, 12.434379, ..., 12.581674, 12.560391,...\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-U-JJA-global_ref.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-U-JJA-global_ref.nc\n", + " * var_key: U\n", + " * All close and within relative tolerance (1e-05)\n", + "Comparing:\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/657-meridional-mean-2d/meridional_mean_2d/MERRA2/MERRA2-U-JJA-global_test.nc\n", + " * /global/cfs/cdirs/e3sm/www/cdat-migration-fy24/main/meridional_mean_2d/MERRA2/MERRA2-U-JJA-global_test.nc\n", + " * var_key: U\n", + " * All close and within relative tolerance (1e-05)\n" + ] + } + ], + "source": [ + "_get_relative_diffs()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Results\n", + "\n", + "- All files are within rtol 1e-05.\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/657-meridional-mean-2d/657_meridional_mean_2d_run_script.py b/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d_run_script.py new file mode 100644 index 0000000000..1b317b983c --- /dev/null +++ b/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d_run_script.py @@ -0,0 +1,11 @@ +# python -m auxiliary_tools.cdat_regression_testing.657-meridional-mean-2d.657_meridional_mean_2d_run_script +from auxiliary_tools.cdat_regression_testing.base_run_script import run_set + +SET_NAME = "meridional_mean_2d" +SET_DIR = "657-meridional-mean-2d" +CFG_PATH: str | None = None +# CFG_PATH: str | None = "/global/u2/v/vo13/E3SM-Project/e3sm_diags/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d.cfg" +MULTIPROCESSING = True + +# %% +run_set(SET_NAME, SET_DIR, CFG_PATH, MULTIPROCESSING) diff --git a/e3sm_diags/driver/meridional_mean_2d_driver.py b/e3sm_diags/driver/meridional_mean_2d_driver.py old mode 100755 new mode 100644 index a3260b1b3d..e4cbec5b8b --- a/e3sm_diags/driver/meridional_mean_2d_driver.py +++ b/e3sm_diags/driver/meridional_mean_2d_driver.py @@ -1,17 +1,23 @@ from __future__ import annotations +from copy import deepcopy from typing import TYPE_CHECKING -import cdms2 -import cdutil -import MV2 -import numpy - -from e3sm_diags.driver import utils +import xarray as xr +import xcdat as xc + +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 +from e3sm_diags.plot.meridional_mean_2d_plot import plot as plot_func logger = custom_logger(__name__) @@ -20,185 +26,216 @@ MeridionalMean2dParameter, ) +DEFAULT_PLEVS = deepcopy(DEFAULT_PLEVS) -def create_metrics(ref, test, ref_regrid, test_regrid, diff): - """ - Creates the mean, max, min, rmse, corr in a dictionary. + +def run_diag(parameter: MeridionalMean2dParameter) -> MeridionalMean2dParameter: + """Run the meridional_mean_2d diagnostics. + + Parameters + ---------- + parameter : MeridionalMean2dParameter + The parameter for the diagnostic. + + Returns + ------- + MeridionalMean2dParameter + The parameter for the diagnostic with the result (completed or failed) + + Raises + ------ + RuntimeError + If the dimensions of the test and ref variables differ. + RuntimeError + If the test or ref variables do are not 3-D (no Z-axis). """ - orig_bounds = cdms2.getAutoBounds() - cdms2.setAutoBounds(1) - lev = ref.getLevel() - if lev is not None: - lev.setBounds(None) + variables = parameter.variables + seasons = parameter.seasons + ref_name = getattr(parameter, "ref_name", "") - lev = test.getLevel() - if lev is not None: - lev.setBounds(None) + test_ds = Dataset(parameter, data_type="test") + ref_ds = Dataset(parameter, data_type="ref") - lev = test_regrid.getLevel() - if lev is not None: - lev.setBounds(None) + for var_key in variables: + logger.info("Variable: {}".format(var_key)) + parameter.var_id = var_key - lev = ref_regrid.getLevel() - if lev is not None: - lev.setBounds(None) + for season in seasons: + parameter._set_name_yrs_attrs(test_ds, ref_ds, season) - lev = diff.getLevel() - if lev is not None: - lev.setBounds(None) - cdms2.setAutoBounds(orig_bounds) + ds_test = test_ds.get_climo_dataset(var_key, season) + ds_ref = ref_ds.get_ref_climo_dataset(var_key, season, ds_test) - metrics_dict = {} - metrics_dict["ref"] = { - "min": min_cdms(ref), - "max": max_cdms(ref), - "mean": mean(ref, axis="xz"), - } - metrics_dict["test"] = { - "min": min_cdms(test), - "max": max_cdms(test), - "mean": mean(test, axis="xz"), - } + dv_test = ds_test[var_key] + dv_ref = ds_ref[var_key] - metrics_dict["diff"] = { - "min": min_cdms(diff), - "max": max_cdms(diff), - "mean": mean(diff, axis="xz"), - } - metrics_dict["misc"] = { - "rmse": rmse(test_regrid, ref_regrid, axis="xz"), - "corr": corr(test_regrid, ref_regrid, axis="xz"), - } + 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) - return metrics_dict + if is_dims_diff: + raise RuntimeError( + "Dimensions of the test and ref variables are different." + ) + elif not is_vars_3d: + raise RuntimeError( + "The test and/or ref variables are not 3-D (no Z axis)." + ) + elif is_vars_3d: + if not parameter._is_plevs_set(): + parameter.plevs = DEFAULT_PLEVS + _run_diags_3d(parameter, ds_test, ds_ref, season, var_key, ref_name) -def run_diag(parameter: MeridionalMean2dParameter) -> MeridionalMean2dParameter: - variables = parameter.variables - seasons = parameter.seasons - ref_name = getattr(parameter, "ref_name", "") + return parameter - 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 +def _run_diags_3d( + parameter: MeridionalMean2dParameter, + ds_test: xr.Dataset, + ds_ref: xr.Dataset, + season: str, + var_key: str, + ref_name: str, +): + plevs = parameter.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_avg = ds_t_plevs.spatial.average(var_key, axis=["Y"]) + ds_r_plevs_avg = ds_r_plevs.spatial.average(var_key, axis=["Y"]) + + # A placeholder Y axis must be added back to the averaged variables to + # align grids via horizontal regridding. Afterwards, the Y axis is dropped. + ds_t_plevs_avg = _add_placeholder_y_axis(ds_test, ds_t_plevs_avg, var_key) + ds_r_plevs_avg = _add_placeholder_y_axis(ds_ref, ds_r_plevs_avg, var_key) + + ds_t_plevs_rg_avg, ds_r_plevs_rg_avg = align_grids_to_lower_res( + ds_t_plevs_avg, + ds_r_plevs_avg, + var_key, + tool="xesmf", + method="conservative_normed", + axis_to_compare="X", + ) + + # After regridding, squeeze the placeholder Y axis. + ds_t_plevs_avg = ds_t_plevs_avg.squeeze() + ds_r_plevs_avg = ds_r_plevs_avg.squeeze() + ds_t_plevs_rg_avg = ds_t_plevs_rg_avg.squeeze() + ds_r_plevs_rg_avg = ds_r_plevs_rg_avg.squeeze() + + # Get the difference between final regridded variables. + with xr.set_options(keep_attrs=True): + ds_diff_plevs_rg_avg = ds_t_plevs_rg_avg.copy() + ds_diff_plevs_rg_avg[var_key] = ( + ds_t_plevs_rg_avg[var_key] - ds_r_plevs_rg_avg[var_key] ) - 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." - ) - - # For variables with a z-axis. - if mv1.getLevel() and mv2.getLevel(): - # Since the default is now stored in MeridionalMean2dParameter, - # 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 = ZonalMean2dParameter().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 - ) + 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_plevs_rg_avg, + ) - mv1_p = cdutil.averager(mv1_p, axis="y") - mv2_p = cdutil.averager(mv2_p, axis="y") + parameter._set_param_output_attrs( + var_key, season, parameter.regions[0], ref_name, ilev=None + ) + _save_data_metrics_and_plots( + parameter, + plot_func, + var_key, + ds_t_plevs_avg, + ds_r_plevs_avg, + ds_diff_plevs_rg_avg, + metrics_dict, + ) - parameter.output_file = "-".join( - [ref_name, var, season, parameter.regions[0]] - ) - parameter.main_title = str(" ".join([var, season])) - - # Regrid towards the lower resolution of the two - # variables for calculating the difference. - if len(mv1_p.getLongitude()) < len(mv2_p.getLongitude()): - mv1_reg = mv1_p - lev_out = mv1_p.getLevel() - lon_out = mv1_p.getLongitude() - # in order to use regrid tool we need to have at least two latitude bands, so generate new grid first - lat = cdms2.createAxis([0]) - lat.setBounds(numpy.array([-1, 1])) - lat.designateLatitude() - grid = cdms2.createRectGrid(lat, lon_out) - - data_shape = list(mv2_p.shape) - data_shape.append(1) - mv2_reg = MV2.resize(mv2_p, data_shape) - mv2_reg.setAxis(-1, lat) - for i, ax in enumerate(mv2_p.getAxisList()): - mv2_reg.setAxis(i, ax) - - mv2_reg = mv2_reg.regrid(grid, regridTool="regrid2")[..., 0] - # Apply the mask back, since crossSectionRegrid - # doesn't preserve the mask. - mv2_reg = MV2.masked_where(mv2_reg == mv2_reg.fill_value, mv2_reg) - elif len(mv1_p.getLongitude()) > len(mv2_p.getLongitude()): - mv2_reg = mv2_p - lev_out = mv2_p.getLevel() - lon_out = mv2_p.getLongitude() - mv1_reg = mv1_p.crossSectionRegrid(lev_out, lon_out) - # In order to use regrid tool we need to have at least two - # latitude bands, so generate new grid first. - lat = cdms2.createAxis([0]) - lat.setBounds(numpy.array([-1, 1])) - lat.designateLatitude() - grid = cdms2.createRectGrid(lat, lon_out) - - data_shape = list(mv1_p.shape) - data_shape.append(1) - mv1_reg = MV2.resize(mv1_p, data_shape) - mv1_reg.setAxis(-1, lat) - for i, ax in enumerate(mv1_p.getAxisList()): - mv1_reg.setAxis(i, ax) - - mv1_reg = mv1_reg.regrid(grid, regridTool="regrid2")[..., 0] - # Apply the mask back, since crossSectionRegrid - # doesn't preserve the mask. - mv1_reg = MV2.masked_where(mv1_reg == mv1_reg.fill_value, mv1_reg) - else: - mv1_reg = mv1_p - mv2_reg = mv2_p - - diff = mv1_reg - mv2_reg - metrics_dict = create_metrics(mv2_p, mv1_p, mv2_reg, mv1_reg, diff) - - parameter.var_region = "global" - - 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 - ) - # For variables without a z-axis. - elif mv1.getLevel() is None and mv2.getLevel() is None: - raise RuntimeError( - "One of or both data doesn't have z dimention. Aborting." - ) +def _add_placeholder_y_axis(ds_original: xr.Dataset, ds_avg: xr.Dataset, var_key: str): + lat_original = xc.get_dim_coords(ds_original, axis="Y") - return parameter + # Make sure to drop the old Y axis before adding the new Y axis. Otherwise, + # the new Y axis can't be added. + lat_key = lat_original.name + ds_avg_new = ds_avg.drop_dims(lat_key) + + ds_avg_new[var_key] = ds_avg_new[var_key].expand_dims({lat_key: [0]}) + ds_avg_new[lat_key].attrs = lat_original.attrs.copy() + + lat_bnds_key = lat_original.attrs["bounds"] + lat_bnds_original = ds_original[lat_bnds_key].copy() + ds_avg_new[lat_bnds_key] = xr.DataArray( + name=lat_bnds_key, + dims=lat_bnds_original.dims, + data=[[-1, 1]], + attrs=lat_bnds_original.attrs.copy(), + ) + + return ds_avg_new + + +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 = {} + + 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=["X", "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=["X", "Z"]), + } + + 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=["X", "Z"]), + } + + metrics_dict["misc"] = { + "rmse": rmse(ds_test_regrid, ds_ref_regrid, var_key, axis=["X", "Z"]), + "corr": correlation(ds_test_regrid, ds_ref_regrid, var_key, axis=["X", "Z"]), + } + + return metrics_dict diff --git a/e3sm_diags/driver/utils/regrid.py b/e3sm_diags/driver/utils/regrid.py index 13a810e0f0..bed8238462 100644 --- a/e3sm_diags/driver/utils/regrid.py +++ b/e3sm_diags/driver/utils/regrid.py @@ -334,6 +334,7 @@ def align_grids_to_lower_res( var_key: str, tool: REGRID_TOOLS, method: str, + axis_to_compare: str = "Y", ) -> Tuple[xr.Dataset, xr.Dataset]: """Align the grids of two Dataset using the lower resolution of the two. @@ -367,6 +368,9 @@ def align_grids_to_lower_res( regrid2 options: - "conservative" + axis_to_compare : str + The axis to use to compare resolutions to find the lower resolution + of both variables, either "X" or "Y". by default "Y". Returns ------- @@ -392,10 +396,10 @@ def align_grids_to_lower_res( ds_a_new = _drop_unused_ilev_axis(ds_a) ds_b_new = _drop_unused_ilev_axis(ds_b) - lat_a = xc.get_dim_coords(ds_a_new[var_key], axis="Y") - lat_b = xc.get_dim_coords(ds_b_new[var_key], axis="Y") + axis_a = xc.get_dim_coords(ds_a_new[var_key], axis=axis_to_compare) + axis_b = xc.get_dim_coords(ds_b_new[var_key], axis=axis_to_compare) - is_a_lower_res = len(lat_a) <= len(lat_b) + is_a_lower_res = len(axis_a) <= len(axis_b) if is_a_lower_res: output_grid = ds_a_new.regridder.grid diff --git a/e3sm_diags/driver/zonal_mean_2d_driver.py b/e3sm_diags/driver/zonal_mean_2d_driver.py index d45cb2e7bf..90f2a34e71 100755 --- a/e3sm_diags/driver/zonal_mean_2d_driver.py +++ b/e3sm_diags/driver/zonal_mean_2d_driver.py @@ -120,7 +120,6 @@ def _run_diags_3d( parameter._set_param_output_attrs( var_key, season, parameter.regions[0], ref_name, ilev=None ) - _save_data_metrics_and_plots( parameter, plot_func, diff --git a/e3sm_diags/plot/cartopy/meridional_mean_2d_plot.py b/e3sm_diags/plot/cartopy/meridional_mean_2d_plot.py deleted file mode 100644 index 09d3ed2580..0000000000 --- a/e3sm_diags/plot/cartopy/meridional_mean_2d_plot.py +++ /dev/null @@ -1,271 +0,0 @@ -from __future__ import print_function - -import os - -import matplotlib # noqa: E402 -import numpy as np -import numpy.ma as ma - -from e3sm_diags.driver.utils.general import get_output_dir -from e3sm_diags.logger import custom_logger -from e3sm_diags.plot import get_colormap - -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) - lon = 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( - lon, - 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) - ax.set_xticks([0, 60, 120, 180, 240, 300, 359.99]) # , crs=ccrs.PlateCarree()) - ax.set_xticklabels(["0", "60E", "120E", 180, "120W", "60W", "0"]) - # 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) - - ax.invert_yaxis() - ax.set_ylabel("Pressure (mb)") - - # 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 < 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, - ) - fig.text( - panel[n][0] + 0.7635, - panel[n][1] + 0.2107, - "%.2f\n%.2f\n%.2f" % 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 - fig = plt.figure(figsize=parameter.figsize, dpi=parameter.dpi) - # proj = ccrs.PlateCarree(central_longitude=180) - proj = None - - # First two panels - min1 = metrics_dict["test"]["min"] - mean1 = metrics_dict["test"]["mean"] - max1 = metrics_dict["test"]["max"] - - plot_panel( - 0, - fig, - proj, - test, - parameter.contour_levels, - parameter.test_colormap, - (parameter.test_name_yrs, parameter.test_title, test.units), - parameter, - stats=(max1, mean1, min1), - ) - - min2 = metrics_dict["ref"]["min"] - mean2 = metrics_dict["ref"]["mean"] - max2 = metrics_dict["ref"]["max"] - plot_panel( - 1, - fig, - proj, - reference, - parameter.contour_levels, - parameter.reference_colormap, - (parameter.ref_name_yrs, parameter.reference_title, reference.units), - parameter, - stats=(max2, mean2, min2), - ) - - # Third panel - min3 = metrics_dict["diff"]["min"] - mean3 = metrics_dict["diff"]["mean"] - max3 = metrics_dict["diff"]["max"] - - r = metrics_dict["misc"]["rmse"] - c = metrics_dict["misc"]["corr"] - plot_panel( - 2, - fig, - proj, - diff, - parameter.diff_levels, - parameter.diff_colormap, - (None, parameter.diff_title, test.units), - parameter, - stats=(max3, mean3, min3, r, c), - ) - - # 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 - - plt.close() diff --git a/e3sm_diags/plot/meridional_mean_2d_plot.py b/e3sm_diags/plot/meridional_mean_2d_plot.py new file mode 100644 index 0000000000..677bc51858 --- /dev/null +++ b/e3sm_diags/plot/meridional_mean_2d_plot.py @@ -0,0 +1,156 @@ +from typing import List, Tuple, Union + +import matplotlib +import xarray as xr +import xcdat as xc + +from e3sm_diags.driver.utils.type_annotations import MetricsDict +from e3sm_diags.logger import custom_logger +from e3sm_diags.parameter.core_parameter import CoreParameter +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, + _get_c_levels_and_norm, + _make_lon_cyclic, + _save_plot, +) + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # isort:skip # noqa: E402 + +logger = custom_logger(__name__) + + +def plot( + parameter: CoreParameter, + da_test: xr.DataArray, + da_ref: xr.DataArray, + da_diff: xr.DataArray, + metrics_dict: MetricsDict, +): + # Create figure, projection + fig = plt.figure(figsize=parameter.figsize, dpi=parameter.dpi) + + units = metrics_dict["units"] + + min1 = metrics_dict["test"]["min"] # type: ignore + mean1 = metrics_dict["test"]["mean"] # type: ignore + max1 = metrics_dict["test"]["max"] # type: ignore + + _add_colormap( + 0, + da_test, + fig, + parameter, + 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"] # type: ignore + mean2 = metrics_dict["ref"]["mean"] # type: ignore + max2 = metrics_dict["ref"]["max"] # type: ignore + + _add_colormap( + 1, + da_ref, + fig, + parameter, + parameter.reference_colormap, + parameter.contour_levels, + title=(parameter.ref_name_yrs, parameter.reference_title, units), # type: ignore + metrics=(max2, mean2, min2), # type: ignore + ) + + 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 + + _add_colormap( + 2, + da_diff, + fig, + parameter, + parameter.diff_colormap, + parameter.diff_levels, + title=(None, parameter.diff_title, units), # type: ignore + metrics=(max3, mean3, min3, r, c), # type: ignore + ) + + fig.suptitle(parameter.main_title, x=0.5, y=0.96, fontsize=18) + + _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[Union[str, None], str, str], + metrics: Tuple[float, ...], +): + var = _make_lon_cyclic(var) + lon = xc.get_dim_coords(var, axis="X") + plev = xc.get_dim_coords(var, axis="Z") + + var = var.squeeze() + + # Configure contour levels and boundary norm. + # -------------------------------------------------------------------------- + c_levels, norm = _get_c_levels_and_norm(contour_levels) + + # Configure contour plot. + # -------------------------------------------------------------------------- + ax = fig.add_axes(DEFAULT_PANEL_CFG[subplot_num]) + contour_plot = _add_contour_plot( + ax, parameter, var, lon, 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. + # -------------------------------------------------------------------------- + ax.set_xticks([0, 60, 120, 180, 240, 300, 359.99]) + ax.set_xticklabels(["0", "60E", "120E", "180", "120W", "60W", "0"]) + + ax.tick_params(labelsize=8.0, direction="out", width=1) + ax.xaxis.set_ticks_position("bottom") + ax.yaxis.set_ticks_position("left") + + if parameter.plot_log_plevs: + ax.set_yscale("log") + + if parameter.plot_plevs: + plev_ticks = parameter.plevs + plt.yticks(plev_ticks, plev_ticks) + + ax.invert_yaxis() + ax.set_ylabel("Pressure (mb)") + + # Add and configure the color bar. + # -------------------------------------------------------------------------- + _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) diff --git a/e3sm_diags/plot/utils.py b/e3sm_diags/plot/utils.py index 340e7bc27f..ad308dbe5d 100644 --- a/e3sm_diags/plot/utils.py +++ b/e3sm_diags/plot/utils.py @@ -580,7 +580,7 @@ def _add_min_mean_max_text( fontdict = {"fontsize": fontsize} if left_text_pos is None: - left_text_pos = (0.6335, 0.2107) + left_text_pos = (0.6635, 0.2107) if right_text_pos is None: right_text_pos = (0.7635, 0.2107)