From afb45764185bb280bacaf134ae191c6b39d041a9 Mon Sep 17 00:00:00 2001 From: John Halley Gotway Date: Wed, 7 Feb 2024 09:01:37 -0700 Subject: [PATCH] Feature #2769 tc_diag_driver_v0.11.0 (#2812) * Per #2769, update to tc_diag_driver version 0.11.0. * Per #2769, add tmp_nc_diag_flag config option to retain temporary files. * Per #2769, change default regrdding method from NEAREST to BILIN. * Per #2769, add checks to make sure input files exist. * Fix the handling of diagnostic name and units strings. Both are limited to 7 characters but the units are enclosed in parantheses, increasing the field width to 9. If the strings are too long, truncate them and print warnings like this: WARNING: OutFileInfo::write_cira_diag_vals() -> long diagnostic name "850TANGXXX" truncated to "850TANG"! WARNING: OutFileInfo::write_cira_diag_vals() -> long diagnostic units string "(10^7C/MXXX)" truncated to "(10^7C/M)"! Pad the first 2 columns out to widths of 7 and 9 and set the inter-column spacing between columns 2 and 3 to 0. This could technically introduce parsing problems when the units are 7 characters long and the values have 5 digits, but this is the logic needed to exactly replicate the existing output. * Just whitespace * Per #2769, remove the tmp_nc_diag_flag TC-Diag config file option in favor of the MET_KEEP_TEMP_FILE environment variable option that Howard is adding on his feature_2772_python_embedding_json branch. I did add documentation about that option here even though the functionality is coming from Howard's branch. * Per #2769, remove tmp_nc_diag_flag that has been removed. * Per #2769, remove mention of tmp_nc_diag_flag option that has been removed. * Per #2769, upgrade to METbaseimage version 3.2 to provide SciPy in the Python environment which is needed by TC-Diag. * Per #2769, add SciPy to the list of required Python packages. * Per #2769, tweak AppendixF details. * Per #2769, update release date for MET-12.0.0-beta3 to be 20240207 since we didn't get it out on the 6th. --- .github/jobs/set_job_controls.sh | 2 +- .../build_docker_and_trigger_metplus.yml | 2 +- data/config/TCDiagConfig_default | 4 +- docs/Users_Guide/appendixF.rst | 14 +- docs/Users_Guide/release-notes.rst | 3 +- docs/Users_Guide/tc-diag.rst | 8 +- docs/conf.py | 2 +- internal/scripts/docker/Dockerfile | 2 +- internal/scripts/docker/Dockerfile.copy | 2 +- internal/test_unit/config/TCDiagConfig_ian | 4 +- scripts/python/tc_diag/compute_tc_diag.py | 13 + .../python/tc_diag/config/post_resample.yml | 8 +- .../tc_diag/config/post_resample_nest.yml | 8 +- .../python/tc_diag/tc_diag_driver/__init__.py | 2 +- .../tc_diag_driver/computation_engine.py | 15 + .../tc_diag/tc_diag_driver/diag_vars.py | 418 +++++++++++------- .../tc_diag/tc_diag_driver/met_diag_vars.py | 167 ++++--- .../tc_diag_driver/met_post_process.py | 4 + .../tc_diag_driver/post_resample_driver.py | 8 +- .../python/tc_diag/tc_diag_driver/results.py | 124 ++++-- src/tools/tc_utils/tc_diag/tc_diag.cc | 39 +- src/tools/tc_utils/tc_diag/tc_diag.h | 8 +- 22 files changed, 569 insertions(+), 288 deletions(-) diff --git a/.github/jobs/set_job_controls.sh b/.github/jobs/set_job_controls.sh index 3517fbf65c..9c39a50716 100755 --- a/.github/jobs/set_job_controls.sh +++ b/.github/jobs/set_job_controls.sh @@ -6,7 +6,7 @@ run_unit_tests=false run_diff=false run_update_truth=false met_base_repo=met-base -met_base_tag=v3.1 +met_base_tag=v3.2 input_data_version=develop truth_data_version=develop diff --git a/.github/workflows/build_docker_and_trigger_metplus.yml b/.github/workflows/build_docker_and_trigger_metplus.yml index 23aafdfd76..13606d5ed3 100644 --- a/.github/workflows/build_docker_and_trigger_metplus.yml +++ b/.github/workflows/build_docker_and_trigger_metplus.yml @@ -29,7 +29,7 @@ jobs: env: SOURCE_BRANCH: ${{ steps.get_branch_name.outputs.branch_name }}-lite MET_BASE_REPO: met-base - MET_BASE_TAG: v3.1 + MET_BASE_TAG: v3.2 - name: Push Docker Image run: .github/jobs/push_docker_image.sh diff --git a/data/config/TCDiagConfig_default b/data/config/TCDiagConfig_default index d5061c82c9..a75c7b42b8 100644 --- a/data/config/TCDiagConfig_default +++ b/data/config/TCDiagConfig_default @@ -85,8 +85,8 @@ domain_info = [ // May be set separately in each data "field" array entry // regrid = { - method = NEAREST; - width = 1; + method = BILIN; + width = 2; vld_thresh = 0.5; shape = SQUARE; } diff --git a/docs/Users_Guide/appendixF.rst b/docs/Users_Guide/appendixF.rst index bfdc451e30..dc81f0cd96 100644 --- a/docs/Users_Guide/appendixF.rst +++ b/docs/Users_Guide/appendixF.rst @@ -22,15 +22,19 @@ In order to use Python embedding, a local Python installation must be available 3. **NumPy** Python package -4. **netCDF4** Python package +4. **Pandas** Python package -5. **Pandas** Python package +5. **Xarray** Python package -6. **Xarray** Python package +6. **YAML** Python package -7. **YAML** Python package +7. **SciPy** Python package -Users should be aware that in some cases, the C-language Python header files and libraries may be deleted at the end of the Python installation process, and they may need to confirm their availability prior to compiling MET. Once the user has confirmed the above requirements are satisfied, they can compile the MET software for Python embedding by passing the **\-\-enable-python** option to the **configure** script on the command line. This will link the MET C++ code directly to the Python libraries. The **NumPy** and **netCDF4** Python packages are required by the Python scripts included with the MET software that facilitate the passing of data in memory and the reading and writing of temporary files when Python embedding is used. The **YAML** Python package is required by the tropical cyclone diagnostics Python scripts called by the TC-Diag tool. +8. **netCDF4** Python package + +Users should be aware that in some cases, the C-language Python header files and libraries may be deleted at the end of the Python installation process, and they may need to confirm their availability prior to compiling MET. Once the user has confirmed the above requirements are satisfied, they can compile the MET software for Python embedding by passing the **\-\-enable-python** option to the **configure** script on the command line. This will link the MET C++ code directly to the Python libraries. + +The **NumPy**, **Xarray**, and **Pandas** Python packages are required by the Python scripts included with the MET software that facilitate the passing of data in memory. The *SciPy** and **YAML** Python packages are required by the tropical cyclone diagnostics Python scripts called by the TC-Diag tool. The **netCDF4** package is used for reading and writing temporary files for Python embedding, but only when the **MET_PYTHON_TMP_FORMAT** environment variable is set to `netcdf` at runtime. In addition to using **\-\-enable-python** with **configure** as mentioned above, the following environment variables must also be set prior to executing **configure**: **MET_PYTHON_BIN_EXE**, **MET_PYTHON_CC**, and **MET_PYTHON_LD**. These may either be set as environment variables or as command line options to **configure**. These environment variables are used when building MET to enable the compiler to find the requisite Python executable, header files, and libraries in the user's local filesystem. Fortunately, Python provides a way to set these variables properly. This frees the user from the necessity of having any expert knowledge of the compiling and linking process. Along with the **Python** executable in the users local Python installation, there should be another executable called **python3-config**, whose output can be used to set these environment variables as follows: diff --git a/docs/Users_Guide/release-notes.rst b/docs/Users_Guide/release-notes.rst index 52f3b28558..c762149df0 100644 --- a/docs/Users_Guide/release-notes.rst +++ b/docs/Users_Guide/release-notes.rst @@ -9,11 +9,12 @@ When applicable, release notes are followed by the GitHub issue number which des enhancement, or new feature (`MET GitHub issues `_). Important issues are listed **in bold** for emphasis. -MET Version 12.0.0-beta3 Release Notes (20240206) +MET Version 12.0.0-beta3 Release Notes (20240207) ------------------------------------------------- .. dropdown:: Repository, build, and test + * Enhance METbaseimage to install SciPy Python package needed by the MET TC-Diag tool (`METbaseimage#20 `_). * Remove the SonarQube token from the properties file (`#2757 `_). * Repository cleanup of stale code and configuration consistency (`#2776 `_). * Add new example installation configuration files for Intel compiler users (`#2785 `_). diff --git a/docs/Users_Guide/tc-diag.rst b/docs/Users_Guide/tc-diag.rst index 7bc2448c53..4f1d4630ee 100644 --- a/docs/Users_Guide/tc-diag.rst +++ b/docs/Users_Guide/tc-diag.rst @@ -63,6 +63,8 @@ Optional Arguments for tc_diag 6. The **-v level** option indicates the desired level of verbosity. The contents of "level" will override the default setting of 2. Setting the verbosity to 0 will make the tool run with no log messages, while increasing the verbosity above 1 will increase the amount of logging. +.. note:: Setting the **MET_KEEP_TEMP_FILE** (:numref:`met_keep_temp_file`) environment variable retains the temporary NetCDF cylindrical coordinate files for development, testing, and debugging purposes. + tc_diag Configuration File -------------------------- @@ -155,13 +157,13 @@ Configuring regridding options .. code-block:: none regrid = { - method = NEAREST; - width = 1; + method = BILIN; + width = 2; vld_thresh = 0.5; shape = SQUARE; } -The **regrid** dictionary is common to multiple MET tools and is described in :numref:`config_options`. It specifies how the input data should be regridded to cylindrical coordinates prior to compute diagnostics. It can be specified separately in each **data.field** array entry, described below. The default setting uses nearest neighbor interpolation for all fields. +The **regrid** dictionary is common to multiple MET tools and is described in :numref:`config_options`. It specifies how the input data should be regridded to cylindrical coordinates prior to compute diagnostics. It can be specified separately in each **data.field** array entry, described below. The default setting uses bilinear interpolation for all fields. Configuring Fields, Levels, and Domains ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/conf.py b/docs/conf.py index 6b8fdabfd7..ad28a7dd63 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -24,7 +24,7 @@ verinfo = version release = f'{version}' release_year = '2024' -release_date = f'{release_year}-02-06' +release_date = f'{release_year}-02-07' copyright = f'{release_year}, {author}' # -- General configuration --------------------------------------------------- diff --git a/internal/scripts/docker/Dockerfile b/internal/scripts/docker/Dockerfile index 7ba0bcf0f6..c5daf58492 100644 --- a/internal/scripts/docker/Dockerfile +++ b/internal/scripts/docker/Dockerfile @@ -1,5 +1,5 @@ ARG MET_BASE_REPO=met-base -ARG MET_BASE_TAG=v3.1 +ARG MET_BASE_TAG=v3.2 FROM dtcenter/${MET_BASE_REPO}:${MET_BASE_TAG} MAINTAINER John Halley Gotway diff --git a/internal/scripts/docker/Dockerfile.copy b/internal/scripts/docker/Dockerfile.copy index bcfe4417a6..4d3d9476cd 100644 --- a/internal/scripts/docker/Dockerfile.copy +++ b/internal/scripts/docker/Dockerfile.copy @@ -1,5 +1,5 @@ ARG MET_BASE_REPO=met-base-unit-test -ARG MET_BASE_TAG=v3.1 +ARG MET_BASE_TAG=v3.2 FROM dtcenter/${MET_BASE_REPO}:${MET_BASE_TAG} MAINTAINER John Halley Gotway diff --git a/internal/test_unit/config/TCDiagConfig_ian b/internal/test_unit/config/TCDiagConfig_ian index 267a65d9dd..62a5c1bb5d 100644 --- a/internal/test_unit/config/TCDiagConfig_ian +++ b/internal/test_unit/config/TCDiagConfig_ian @@ -77,8 +77,8 @@ domain_info = [ // May be set separately in each data "field" array entry // regrid = { - method = NEAREST; - width = 1; + method = BILIN; + width = 2; vld_thresh = 0.5; shape = SQUARE; } diff --git a/scripts/python/tc_diag/compute_tc_diag.py b/scripts/python/tc_diag/compute_tc_diag.py index 66b1fef2c6..6a0523784e 100644 --- a/scripts/python/tc_diag/compute_tc_diag.py +++ b/scripts/python/tc_diag/compute_tc_diag.py @@ -85,6 +85,19 @@ def main(): land_filename = os.path.expandvars(args.land_file) data_filename = os.path.expandvars(args.data_file) + # Validate inputs + if not os.path.isfile(config_filename): + print("Error: Input Config File (" + config_filename + ") does not exist!") + sys.exit(1) + + if not os.path.isfile(land_filename): + print("Error: Input Land File (" + land_filename + ")does not exist!") + sys.exit(1) + + if not os.path.isfile(data_filename): + print("Error: Input Data File (" + data_filename + ") does not exist!") + sys.exit(1) + # Print verbose arguments if args.verbose: print("Python Script:\t" + os.path.expandvars(sys.argv[0])) diff --git a/scripts/python/tc_diag/config/post_resample.yml b/scripts/python/tc_diag/config/post_resample.yml index cbac0fa3e0..5ee8519c66 100644 --- a/scripts/python/tc_diag/config/post_resample.yml +++ b/scripts/python/tc_diag/config/post_resample.yml @@ -100,23 +100,25 @@ pressure_independent_computation_specs: callable: *div_vort_func output_vars: [850DVRG, 850VORT] kwargs: {u_name: *u_vert_name, v_name: *v_vert_name, level_hPa: 850, radius_km: *div_vort_max_km} + units: [/S, /S] div_vort_200: callable: *div_vort_func output_vars: [200DVRG, 200VORT] kwargs: {u_name: *u_vert_name, v_name: *v_vert_name, level_hPa: 200, radius_km: *div_vort_max_km} + units: [/S, /S] shear: batch_order: 1 callable: tc_diag_driver.met_diag_vars.shear output_vars: [SHR_MAG, SHR_HDG] - kwargs: {u_name: u, v_name: v, bottom_hPa: 850, top_hPa: 200} + kwargs: {u_name: U, v_name: V, bottom_hPa: 850, top_hPa: 200} unit_converters: [tc_diag_driver.met_post_process.mps_to_kt, pass] units: [KT, DEG] TGRD: batch_order: 1 callable: tc_diag_driver.met_diag_vars.temperature_gradient - kwargs: {u_name: u, v_name: v, bottom_hPa: 850, top_hPa: 700} + kwargs: {u_name: U, v_name: V, bottom_hPa: 850, top_hPa: 700} units: "10^7C/M" sounding_computation_specs: @@ -161,4 +163,4 @@ comment_format: | * 850VORT, 200DVRG averaged from {div_vort_min}-{div_vort_max} km around storm center [x10^7 /s, x10^7 /s] * * 850TANG averaged from {rad_tan_min}-{rad_tan_max} km around storm center [x10 m/s] * * T, R, Z, P averaged from {therm_min}-{therm_max} km around storm center [x10 C, %, dm, mb] * - * TPW averaged from {tpw_min}-{tpw_max} km around storm center [mm] * + * TPW averaged from {tpw_min}-{tpw_max} km around storm center [mm] * \ No newline at end of file diff --git a/scripts/python/tc_diag/config/post_resample_nest.yml b/scripts/python/tc_diag/config/post_resample_nest.yml index d2c397cbd3..fbe3255e55 100644 --- a/scripts/python/tc_diag/config/post_resample_nest.yml +++ b/scripts/python/tc_diag/config/post_resample_nest.yml @@ -100,23 +100,25 @@ pressure_independent_computation_specs: callable: *div_vort_func output_vars: [850DVRG, 850VORT] kwargs: {u_name: *u_vert_name, v_name: *v_vert_name, level_hPa: 850, radius_km: *div_vort_max_km} + units: [/S, /S] div_vort_200: callable: *div_vort_func output_vars: [200DVRG, 200VORT] kwargs: {u_name: *u_vert_name, v_name: *v_vert_name, level_hPa: 200, radius_km: *div_vort_max_km} + units: [/S, /S] shear: batch_order: 1 callable: tc_diag_driver.met_diag_vars.shear output_vars: [SHR_MAG, SHR_HDG] - kwargs: {u_name: u, v_name: v, bottom_hPa: 850, top_hPa: 200} + kwargs: {u_name: U, v_name: V, bottom_hPa: 850, top_hPa: 200} unit_converters: [tc_diag_driver.met_post_process.mps_to_kt, pass] units: [KT, DEG] TGRD: batch_order: 1 callable: tc_diag_driver.met_diag_vars.temperature_gradient - kwargs: {u_name: u, v_name: v, bottom_hPa: 850, top_hPa: 700} + kwargs: {u_name: U, v_name: V, bottom_hPa: 850, top_hPa: 700} units: "10^7C/M" sounding_computation_specs: @@ -161,4 +163,4 @@ comment_format: | * 850VORT, 200DVRG averaged from {div_vort_min}-{div_vort_max} km around storm center [x10^7 /s, x10^7 /s] * * 850TANG averaged from {rad_tan_min}-{rad_tan_max} km around storm center [x10 m/s] * * T, R, Z, P averaged from {therm_min}-{therm_max} km around storm center [x10 C, %, dm, mb] * - * TPW averaged from {tpw_min}-{tpw_max} km around storm center [mm] * + * TPW averaged from {tpw_min}-{tpw_max} km around storm center [mm] * \ No newline at end of file diff --git a/scripts/python/tc_diag/tc_diag_driver/__init__.py b/scripts/python/tc_diag/tc_diag_driver/__init__.py index 3e2f46a3a3..ae6db5f176 100644 --- a/scripts/python/tc_diag/tc_diag_driver/__init__.py +++ b/scripts/python/tc_diag/tc_diag_driver/__init__.py @@ -1 +1 @@ -__version__ = "0.9.0" +__version__ = "0.11.0" diff --git a/scripts/python/tc_diag/tc_diag_driver/computation_engine.py b/scripts/python/tc_diag/tc_diag_driver/computation_engine.py index 2e67698e61..b419ff0800 100644 --- a/scripts/python/tc_diag/tc_diag_driver/computation_engine.py +++ b/scripts/python/tc_diag/tc_diag_driver/computation_engine.py @@ -251,6 +251,21 @@ def get_result_names(computations: List[DiagComputation]) -> List[str]: return names +def get_all_result_units( + pressure_indedpendent: List[DiagComputation], sounding: List[DiagComputation] +) -> Tuple[List[str], List[str]]: + pi_var_units = get_result_units(pressure_indedpendent) + snd_var_units = get_result_units(sounding) + return pi_var_units, snd_var_units + + +def get_result_units(computations: List[DiagComputation]) -> List[str]: + units = [] + for c in computations: + units.extend(c.units) + return units + + def get_computation_batches( pressure_indedpendent: List[DiagComputation], sounding: List[DiagComputation] ) -> List[ComputationBatch]: diff --git a/scripts/python/tc_diag/tc_diag_driver/diag_vars.py b/scripts/python/tc_diag/tc_diag_driver/diag_vars.py index 1de8fe6ef1..37788c309d 100644 --- a/scripts/python/tc_diag/tc_diag_driver/diag_vars.py +++ b/scripts/python/tc_diag/tc_diag_driver/diag_vars.py @@ -10,7 +10,7 @@ import logging import pathlib import sys -from typing import Optional, TextIO, Tuple +from typing import Callable, Optional, TextIO, Tuple import numpy as np import pandas as pd @@ -27,45 +27,51 @@ def debug_cyl_grid_dump( - hour: int, - grib_dataset: xr.Dataset, - cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, - output_filename: str, - grib_var_name: str, - level_hPa: Optional[int] = None, - unit_converter: str = None, - **_kwargs) -> float: + hour: int, + grib_dataset: xr.Dataset, + cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, + output_filename: str, + grib_var_name: str, + level_hPa: Optional[int] = None, + unit_converter: str = None, + **_kwargs, +) -> float: LOGGER.info( "Started debug dump of %s converted to cylindrical grid in file: %s", - grib_var_name, output_filename) + grib_var_name, + output_filename, + ) data = _obtain_data_at_level(grib_dataset, grib_var_name, level_hPa) data = _convert(data, unit_converter) lerp = cylindrical_grid_interpolator data_on_grid = lerp(data) - out_fname = output_filename.format(grib_var_name=grib_var_name, - level_hPa=level_hPa, - hour=hour) + out_fname = output_filename.format( + grib_var_name=grib_var_name, level_hPa=level_hPa, hour=hour + ) data_vars = {grib_var_name: (["theta_radians", "radius_km"], data_on_grid)} theta_radians = lerp.theta_2d_radians[:, 0] radius_km = lerp.rad_2d_km[0, :] - ds = xr.Dataset(data_vars=data_vars, - coords=dict(theta_radians=theta_radians, - radius_km=radius_km)) + ds = xr.Dataset( + data_vars=data_vars, + coords=dict(theta_radians=theta_radians, radius_km=radius_km), + ) ds.to_netcdf(path=out_fname) LOGGER.info( "Finished debug dump of %s converted to cylindrical grid in file: %s", - grib_var_name, output_filename) + grib_var_name, + output_filename, + ) # Have to return a float. return np.nan -def _obtain_data_at_level(grib_dataset: xr.Dataset, - grib_var_name: str, - level_hPa: Optional[int] = None) -> np.ndarray: +def _obtain_data_at_level( + grib_dataset: xr.Dataset, grib_var_name: str, level_hPa: Optional[int] = None +) -> np.ndarray: if level_hPa is None: data = grib_dataset[grib_var_name].values else: @@ -73,8 +79,7 @@ def _obtain_data_at_level(grib_dataset: xr.Dataset, return data -def _convert(data: np.ndarray, - unit_converter: Optional[str] = None) -> np.ndarray: +def _convert(data: np.ndarray, unit_converter: Optional[str] = None) -> np.ndarray: if unit_converter is not None: uc = ce.get_callable_from_import_path(unit_converter) data = uc(data) @@ -83,71 +88,103 @@ def _convert(data: np.ndarray, def mean_in_radius_range( - grib_dataset: xr.Dataset, - cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, - min_radius_km: float, - max_radius_km: float, - grib_var_name: str, - level_hPa: Optional[int] = None, - unit_converter: str = None, - **_kwargs) -> float: - LOGGER.info("Started mean in radius average: %s %s", grib_var_name, - str(level_hPa)) + grib_dataset: xr.Dataset, + cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, + min_radius_km: float, + max_radius_km: float, + grib_var_name: str, + level_hPa: Optional[int] = None, + unit_converter: str = None, + **_kwargs, +) -> float: + LOGGER.info("Started mean in radius average: %s %s", grib_var_name, str(level_hPa)) data = _obtain_data_at_level(grib_dataset, grib_var_name, level_hPa) data = _convert(data, unit_converter) lerp = cylindrical_grid_interpolator data_on_grid = lerp(data) - mean = area_average(data_on_grid, lerp.rad_2d_km, min_radius_km, - max_radius_km) - LOGGER.info("Finished mean in radius average: %s %s", grib_var_name, - str(level_hPa)) + mean = area_average(data_on_grid, lerp.rad_2d_km, min_radius_km, max_radius_km) + LOGGER.info("Finished mean in radius average: %s %s", grib_var_name, str(level_hPa)) return mean -#TODO: Add unit tests, possible move to diag_lib -def area_average(data_on_grid: np.ndarray, radii_2d_km: np.ndarray, - min_radius_km: float, max_radius_km: float) -> float: +# TODO: Add unit tests, possible move to diag_lib +def area_average( + data_on_grid: np.ndarray, + radii_2d_km: np.ndarray, + min_radius_km: float, + max_radius_km: float, +) -> float: mask = (radii_2d_km >= min_radius_km) & (radii_2d_km <= max_radius_km) numerator = np.sum(data_on_grid[mask] * radii_2d_km[mask]) denominator = np.sum(radii_2d_km[mask]) return numerator / denominator -def track_row_lookup(track_row: pd.DataFrame, - column_name: str, - convert_to_0_360: bool = False, - **_kwargs): - LOGGER.info("Started track row lookup: %s convert_to_0_360: %d", - column_name, convert_to_0_360) +def track_row_lookup( + track_row: pd.DataFrame, column_name: str, convert_to_0_360: bool = False, **_kwargs +): + LOGGER.info( + "Started track row lookup: %s convert_to_0_360: %d", + column_name, + convert_to_0_360, + ) val = track_row[column_name][0] if convert_to_0_360: val %= 360 - LOGGER.info("Finished track row lookup: %s convert_to_0_360: %d val: %f", - column_name, convert_to_0_360, val) + LOGGER.info( + "Finished track row lookup: %s convert_to_0_360: %d val: %f", + column_name, + convert_to_0_360, + val, + ) return val -def shear(hour: int, results: hr_results.ForecastHourResults, u_name: str, - v_name: str, bottom_hPa: float, top_hPa: float, - **_kwargs) -> Tuple[float, float]: +def shear( + hour: int, + results: hr_results.ForecastHourResults, + u_name: str, + v_name: str, + bottom_hPa: float, + top_hPa: float, + uv_converter: Optional[Callable[[float], float]] = None, + **_kwargs, +) -> Tuple[float, float]: LOGGER.info( "Started shear calculations hour: %d u_name: %s v_name: %s bottom_hPa: %f top_hPa: %f", - hour, u_name, v_name, bottom_hPa, top_hPa) + hour, + u_name, + v_name, + bottom_hPa, + top_hPa, + ) u = _shear_component(hour, results, u_name, bottom_hPa, top_hPa) v = _shear_component(hour, results, v_name, bottom_hPa, top_hPa) + if uv_converter is not None: + u = uv_converter(u) + v = uv_converter(v) r, theta = _u_v_to_r_theta(u, v) LOGGER.info( "Finished shear calculations hour: %d u_name: %s v_name: %s bottom_hPa: %f top_hPa: %f", - hour, u_name, v_name, bottom_hPa, top_hPa) + hour, + u_name, + v_name, + bottom_hPa, + top_hPa, + ) return r, theta -def _shear_component(hour: int, results: hr_results.ForecastHourResults, - var_name: str, bottom_hPa: float, - top_hPa: float) -> float: +def _shear_component( + hour: int, + results: hr_results.ForecastHourResults, + var_name: str, + bottom_hPa: float, + top_hPa: float, +) -> float: at_all_levels = results.soundings[var_name].sel(forecast_hour=hour) at_bottom = at_all_levels.interp(level_hPa=bottom_hPa) at_top = at_all_levels.interp(level_hPa=top_hPa) @@ -172,13 +209,25 @@ def _u_v_to_r_theta(u: float, v: float) -> Tuple[float, float]: return r, theta -def temperature_gradient(hour: int, results: hr_results.ForecastHourResults, - u_name: str, v_name: str, bottom_hPa: int, - top_hPa: int, tc_lat: float, **_kwargs) -> float: - bottom_u, top_u = _component_at_levels(results, u_name, hour, bottom_hPa, - top_hPa) - bottom_v, top_v = _component_at_levels(results, v_name, hour, bottom_hPa, - top_hPa) +def temperature_gradient( + hour: int, + results: hr_results.ForecastHourResults, + u_name: str, + v_name: str, + bottom_hPa: int, + top_hPa: int, + tc_lat: float, + uv_converter: Optional[Callable[[float], float]] = None, + **_kwargs, +) -> float: + bottom_u, top_u = _component_at_levels(results, u_name, hour, bottom_hPa, top_hPa) + bottom_v, top_v = _component_at_levels(results, v_name, hour, bottom_hPa, top_hPa) + + if uv_converter is not None: + bottom_u = uv_converter(bottom_u) + top_u = uv_converter(top_u) + bottom_v = uv_converter(bottom_v) + top_v = uv_converter(top_v) f = 2.0 * 7.292e-5 * np.sin(tc_lat * 0.017453) r = 287.0 @@ -192,13 +241,15 @@ def temperature_gradient(hour: int, results: hr_results.ForecastHourResults, return tgrd -def _component_at_levels(results: hr_results.ForecastHourResults, - var_name: str, hour: int, bottom_hPa: int, - top_hPa: int) -> Tuple[float, float]: - bottom = results.soundings[var_name].sel(forecast_hour=hour, - level_hPa=bottom_hPa) - top = results.soundings[var_name].sel(forecast_hour=hour, - level_hPa=top_hPa) +def _component_at_levels( + results: hr_results.ForecastHourResults, + var_name: str, + hour: int, + bottom_hPa: int, + top_hPa: int, +) -> Tuple[float, float]: + bottom = results.soundings[var_name].sel(forecast_hour=hour, level_hPa=bottom_hPa) + top = results.soundings[var_name].sel(forecast_hour=hour, level_hPa=top_hPa) return bottom, top @@ -212,8 +263,8 @@ def always_missing(**_kwargs) -> float: @dataclasses.dataclass class LUTExtents: - """Class that stores distance to land LUT meta-data. - """ + """Class that stores distance to land LUT meta-data.""" + ll_lon: float ll_lat: float ur_lon: float @@ -236,6 +287,7 @@ class LandLUT: Bilinear interpolation will be performed on the values from the LUT. The interpolation will be performed across the date-line if necessary. """ + UNITS = "km" def __init__(self, distances: np.ndarray, extents: LUTExtents): @@ -265,11 +317,7 @@ def distance(self, lon: float, lat: float) -> float: w_0_0 = self._inverse_distance(lon, i, lat, j) w_0_1 = self._inverse_distance(lon, i_right, lat, j, wrap_dx=wrap_dx) w_1_0 = self._inverse_distance(lon, i, lat, j_up) - w_1_1 = self._inverse_distance(lon, - i_right, - lat, - j_up, - wrap_dx=wrap_dx) + w_1_1 = self._inverse_distance(lon, i_right, lat, j_up, wrap_dx=wrap_dx) weight_total = w_0_0 + w_0_1 + w_1_0 + w_1_1 w_0_0 /= weight_total @@ -286,12 +334,7 @@ def distance(self, lon: float, lat: float) -> float: return lerp_dist - def _inverse_distance(self, - lon: float, - i: int, - lat: float, - j: int, - wrap_dx=False): + def _inverse_distance(self, lon: float, i: int, lat: float, j: int, wrap_dx=False): if wrap_dx: dx = lon - (self.extents.ur_lon + self.extents.lon_spacing) else: @@ -308,8 +351,7 @@ def _inverse_distance(self, return 1 / np.sqrt(dx**2 + dy**2) -def read_land_lut_file( - lut_filename: pathlib.Path) -> Tuple[np.ndarray, LUTExtents]: +def read_land_lut_file(lut_filename: pathlib.Path) -> Tuple[np.ndarray, LUTExtents]: """Reads a simple ascii LUT of distances from land on a global grid. Args: @@ -333,8 +375,9 @@ def _get_header_info(header_line: str) -> LUTExtents: ll_lat, ur_lat, lat_spacing = map(float, header_tokens[4:7]) nx = int(header_tokens[3]) ny = int(header_tokens[7]) - extents = LUTExtents(ll_lon, ll_lat, ur_lon, ur_lat, nx, ny, lon_spacing, - lat_spacing) + extents = LUTExtents( + ll_lon, ll_lat, ur_lon, ur_lat, nx, ny, lon_spacing, lat_spacing + ) return extents @@ -369,40 +412,57 @@ def get_land_lut(land_lut_filename: pathlib.Path) -> LandLUT: return LandLUT(distances, extents) -def distance_to_land_lookup(land_lut: LandLUT, - hour: int, - track_row: pd.DataFrame, - lon_column_name: str = "lon", - lat_column_name: str = "lat", - **_kwargs) -> float: +def distance_to_land_lookup( + land_lut: LandLUT, + hour: int, + track_row: pd.DataFrame, + lon_column_name: str = "lon", + lat_column_name: str = "lat", + **_kwargs, +) -> float: LOGGER.info( "Started distance to land lookup for hour: %d using column names: %s, %s.", - hour, lon_column_name, lat_column_name) + hour, + lon_column_name, + lat_column_name, + ) lon = track_row_lookup(track_row, lon_column_name, convert_to_0_360=True) lat = track_row_lookup(track_row, lat_column_name) distance = land_lut.distance(lon, lat) LOGGER.info( "Finished distance to land lookup: %f for hour: %d using column names: %s, %s.", - distance, hour, lon_column_name, lat_column_name) + distance, + hour, + lon_column_name, + lat_column_name, + ) return distance -def storm_r_theta(model_track: pd.DataFrame, - hour: int, - model_time: dt.datetime, - time_delta_hours: int = 6, - lon_column_name: str = "lon", - lat_column_name: str = "lat", - tau_column_name: str = "tau", - time_column_name: str = "yyyymmddhh", - **_kwargs) -> Tuple[float, float]: +def storm_r_theta( + model_track: pd.DataFrame, + hour: int, + model_time: dt.datetime, + time_delta_hours: int = 6, + lon_column_name: str = "lon", + lat_column_name: str = "lat", + tau_column_name: str = "tau", + time_column_name: str = "yyyymmddhh", + **_kwargs, +) -> Tuple[float, float]: LOGGER.info( "Started storm r, theta calculation hour:%d, model_time:%s, " "time_delta_hours:%d, lon_column_name:%s, lat_column_name:%s, " - "tau_column_name:%s, time_column_name:%s", hour, model_time, - time_delta_hours, lon_column_name, lat_column_name, tau_column_name, - time_column_name) + "tau_column_name:%s, time_column_name:%s", + hour, + model_time, + time_delta_hours, + lon_column_name, + lat_column_name, + tau_column_name, + time_column_name, + ) mt_rows = model_track.loc[model_track[time_column_name] == model_time] min_tau = min(mt_rows[tau_column_name]) max_tau = max(mt_rows[tau_column_name]) @@ -422,10 +482,8 @@ def storm_r_theta(model_track: pd.DataFrame, upper_offset = 0 delta_t = time_delta_hours - lower_row = mt_rows.loc[mt_rows[tau_column_name] == hour + - lower_offset].iloc[0] - upper_row = mt_rows.loc[mt_rows[tau_column_name] == hour + - upper_offset].iloc[0] + lower_row = mt_rows.loc[mt_rows[tau_column_name] == hour + lower_offset].iloc[0] + upper_row = mt_rows.loc[mt_rows[tau_column_name] == hour + upper_offset].iloc[0] lower_lon = lower_row[lon_column_name] % 360 lower_lat = lower_row[lat_column_name] @@ -444,52 +502,79 @@ def storm_r_theta(model_track: pd.DataFrame, LOGGER.info( "Finished storm r:%f, theta:%f calculation hour:%d, model_time:%s, " "time_delta_hours:%d, lon_column_name:%s, lat_column_name:%s, " - "tau_column_name:%s, time_column_name:%s", r, theta, hour, model_time, - time_delta_hours, lon_column_name, lat_column_name, tau_column_name, - time_column_name) + "tau_column_name:%s, time_column_name:%s", + r, + theta, + hour, + model_time, + time_delta_hours, + lon_column_name, + lat_column_name, + tau_column_name, + time_column_name, + ) return r, theta def radial_and_tangential_area_average( - grib_dataset: xr.Dataset, - cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, - u_name: str, v_name: str, level_hPa: int, min_radius_km: float, - max_radius_km: float, **_kwargs) -> Tuple[float, float]: + grib_dataset: xr.Dataset, + cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, + u_name: str, + v_name: str, + level_hPa: int, + min_radius_km: float, + max_radius_km: float, + **_kwargs, +) -> Tuple[float, float]: """Computes area average of azimuthally averaged radial & tangential winds.""" radial_azavg, tan_azavg, radii = _gen_azimuthal_avg_of_radial_tan_winds( - grib_dataset, cylindrical_grid_interpolator, u_name, v_name, level_hPa) + grib_dataset, cylindrical_grid_interpolator, u_name, v_name, level_hPa + ) # Area average of the azimuthal averages - radial_radius_avg = area_average(radial_azavg, radii, min_radius_km, - max_radius_km) - tan_radius_avg = area_average(tan_azavg, radii, min_radius_km, - max_radius_km) + radial_radius_avg = area_average(radial_azavg, radii, min_radius_km, max_radius_km) + tan_radius_avg = area_average(tan_azavg, radii, min_radius_km, max_radius_km) return radial_radius_avg, tan_radius_avg def divergence_vorticity( - grib_dataset: xr.Dataset, - cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, - u_name: str, v_name: str, level_hPa: int, radius_km: float, - **_kwargs) -> Tuple[float]: + grib_dataset: xr.Dataset, + cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, + u_name: str, + v_name: str, + level_hPa: int, + radius_km: float, + **_kwargs, +) -> Tuple[float]: LOGGER.info( "Started computing divergence and vorticity u_name:%s, v_name:%s, level_hPa:%d, radius_km:%f", - u_name, v_name, level_hPa, radius_km) + u_name, + v_name, + level_hPa, + radius_km, + ) radial_azavg, tan_azavg, radii = _gen_azimuthal_avg_of_radial_tan_winds( - grib_dataset, cylindrical_grid_interpolator, u_name, v_name, level_hPa) + grib_dataset, cylindrical_grid_interpolator, u_name, v_name, level_hPa + ) - divergence, vorticity = _div_vort(radii, radius_km, radial_azavg, - tan_azavg) + divergence, vorticity = _div_vort(radii, radius_km, radial_azavg, tan_azavg) LOGGER.info( "Finished computing divergence:%f and vorticity:%f u_name:%s, v_name:%s, level_hPa:%d, radius_km:%f", - divergence, vorticity, u_name, v_name, level_hPa, radius_km) + divergence, + vorticity, + u_name, + v_name, + level_hPa, + radius_km, + ) return divergence, vorticity -def _div_vort(radii: np.ndarray, radius: float, radial_azavg: np.ndarray, - tan_azavg: np.ndarray) -> Tuple[float, float]: +def _div_vort( + radii: np.ndarray, radius: float, radial_azavg: np.ndarray, tan_azavg: np.ndarray +) -> Tuple[float, float]: closest_radius_km, radius_index = _nearest_radius(radii, radius) radius_m = closest_radius_km * 1000 @@ -499,8 +584,7 @@ def _div_vort(radii: np.ndarray, radius: float, radial_azavg: np.ndarray, return divergence, vorticity -def _nearest_radius(radii: np.ndarray, - desired_radius: float) -> Tuple[float, int]: +def _nearest_radius(radii: np.ndarray, desired_radius: float) -> Tuple[float, int]: if np.nanmax(radii) < desired_radius: raise ValueError( f"Desired radius: {desired_radius} larger than max radius: {np.max(radii)}" @@ -513,10 +597,12 @@ def _nearest_radius(radii: np.ndarray, def _gen_azimuthal_avg_of_radial_tan_winds( - grib_dataset: xr.Dataset, - cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, - u_name: str, v_name: str, - level_hPa: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + grib_dataset: xr.Dataset, + cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, + u_name: str, + v_name: str, + level_hPa: int, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """Produces azimuthally averaged radial & tangential winds from grib u,v""" # Interpolate u, v to the desired pressure level u = grib_dataset[u_name].interp(level=level_hPa).values @@ -527,8 +613,7 @@ def _gen_azimuthal_avg_of_radial_tan_winds( u_cyl = lerp(u) v_cyl = lerp(v) - radial_azavg, tan_azavg = _rad_tan_azavg(u_cyl, v_cyl, - lerp.theta_2d_radians, 0) + radial_azavg, tan_azavg = _rad_tan_azavg(u_cyl, v_cyl, lerp.theta_2d_radians, 0) # Get the 1D array of radii corresponding to the azimuthally averaged data. radii = lerp.rad_2d_km[0, :] @@ -536,14 +621,12 @@ def _gen_azimuthal_avg_of_radial_tan_winds( return radial_azavg, tan_azavg, radii -def _rad_tan_azavg(u_cyl: np.ndarray, v_cyl: np.ndarray, - theta_2d_radians: np.ndarray, - avg_axis: int) -> Tuple[np.ndarray, np.ndarray]: +def _rad_tan_azavg( + u_cyl: np.ndarray, v_cyl: np.ndarray, theta_2d_radians: np.ndarray, avg_axis: int +) -> Tuple[np.ndarray, np.ndarray]: # Convert to radial and tangential wind - radial = u_cyl * np.cos(theta_2d_radians) + v_cyl * np.sin( - theta_2d_radians) - tangential = -u_cyl * np.sin(theta_2d_radians) + v_cyl * np.cos( - theta_2d_radians) + radial = u_cyl * np.cos(theta_2d_radians) + v_cyl * np.sin(theta_2d_radians) + tangential = -u_cyl * np.sin(theta_2d_radians) + v_cyl * np.cos(theta_2d_radians) # Compute the azimuthal averages radial_azavg = cg.azimuthal_average(radial, axis=avg_axis) @@ -552,14 +635,21 @@ def _rad_tan_azavg(u_cyl: np.ndarray, v_cyl: np.ndarray, return radial_azavg, tan_azavg -def average_rmw(grib_dataset: xr.Dataset, - cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, - u_surface_name: str, v_surface_name: str, radius_km: float, - **_kwargs) -> float: +def average_rmw( + grib_dataset: xr.Dataset, + cylindrical_grid_interpolator: cg.CylindricalGridInterpolator, + u_surface_name: str, + v_surface_name: str, + radius_km: float, + **_kwargs, +) -> float: LOGGER.info( "Starting calculation of average rmw u_surface_name:%s " - "v_surface_name:%s radius_km:%f", u_surface_name, v_surface_name, - radius_km) + "v_surface_name:%s radius_km:%f", + u_surface_name, + v_surface_name, + radius_km, + ) u = grib_dataset[u_surface_name].values v = grib_dataset[v_surface_name].values @@ -567,18 +657,22 @@ def average_rmw(grib_dataset: xr.Dataset, u_cyl = lerp(u) v_cyl = lerp(v) - avg_rmw_km = post_cyl_avg_rmw(u_cyl, v_cyl, radius_km, - lerp.rad_2d_km[0, :]) + avg_rmw_km = post_cyl_avg_rmw(u_cyl, v_cyl, radius_km, lerp.rad_2d_km[0, :]) LOGGER.info( "Finished calculation of average rmw:%f u_surface_name:%s " - "v_surface_name:%s radius_km:%f", avg_rmw_km, u_surface_name, - v_surface_name, radius_km) + "v_surface_name:%s radius_km:%f", + avg_rmw_km, + u_surface_name, + v_surface_name, + radius_km, + ) return avg_rmw_km -def post_cyl_avg_rmw(u_cyl: np.ndarray, v_cyl: np.ndarray, radius_km: float, - radii_1d_km: np.ndarray) -> float: +def post_cyl_avg_rmw( + u_cyl: np.ndarray, v_cyl: np.ndarray, radius_km: float, radii_1d_km: np.ndarray +) -> float: speed = np.sqrt(u_cyl**2 + v_cyl**2) # Figure out the index of the largest radius <= radius_km @@ -590,11 +684,11 @@ def post_cyl_avg_rmw(u_cyl: np.ndarray, v_cyl: np.ndarray, radius_km: float, selected_i = i # Slice the speed array so that it no longer includes data for radii > radius_km - speed = speed[:, 0:selected_i + 1] + speed = speed[:, 0 : selected_i + 1] # Get the radius of the max speed at each azimuth max_wind_indices = np.nanargmax(speed, axis=1) rmw_at_each_azimuth = radii_1d_km[max_wind_indices] avg_rmw_km = np.nanmean(rmw_at_each_azimuth) - return avg_rmw_km \ No newline at end of file + return avg_rmw_km diff --git a/scripts/python/tc_diag/tc_diag_driver/met_diag_vars.py b/scripts/python/tc_diag/tc_diag_driver/met_diag_vars.py index 64b3efc0d1..bad59c5844 100644 --- a/scripts/python/tc_diag/tc_diag_driver/met_diag_vars.py +++ b/scripts/python/tc_diag/tc_diag_driver/met_diag_vars.py @@ -5,7 +5,7 @@ import pandas as pd import xarray as xr -from tc_diag_driver import diag_vars +from tc_diag_driver import diag_vars, met_post_process from tc_diag_driver import results as fcresults SHEAR_THETA_UNITS = "degrees" @@ -24,8 +24,9 @@ def kwarg_lookup(lookup_name: str, **kwargs: Dict[str, Any]) -> float: return kwargs[lookup_name] -def dataset_lookup(input_data: xr.Dataset, var_name: str, - **_kwargs: Dict[str, Any]) -> float: +def dataset_lookup( + input_data: xr.Dataset, var_name: str, **_kwargs: Dict[str, Any] +) -> float: var = input_data[var_name][0] return var @@ -34,44 +35,51 @@ def _with_units(data: float, units: str) -> xr.DataArray: return xr.DataArray(data=data, attrs={"units": units}) -def average_rmw(input_data: xr.Dataset, radii_1d: xr.DataArray, - u_surface_name: str, v_surface_name: str, radius_km: float, - **_kwargs: Dict[str, Any]) -> xr.DataArray: +def average_rmw( + input_data: xr.Dataset, + radii_1d: xr.DataArray, + u_surface_name: str, + v_surface_name: str, + radius_km: float, + **_kwargs: Dict[str, Any] +) -> xr.DataArray: usurf = input_data[u_surface_name][0] vsurf = input_data[v_surface_name][0] # The function expects usurf and vsurf to be azimuth x radii, but the NC # files use the opposite convention, so transposing the arrays is necessary. - avg_rmw = diag_vars.post_cyl_avg_rmw(usurf.data.T, vsurf.data.T, radius_km, - radii_1d.data) + avg_rmw = diag_vars.post_cyl_avg_rmw( + usurf.data.T, vsurf.data.T, radius_km, radii_1d.data + ) units = radii_1d.attrs["units"] da_avg_rmw = _with_units(avg_rmw, units) return da_avg_rmw -def area_average(input_data: xr.Dataset, - radii_2d: np.ndarray, - min_radius_km: float, - max_radius_km: float, - var_name: str, - level_hPa: Optional[int] = None, - levels_dimen: Optional[str] = None, - **_kwargs: Dict[str, Any]) -> xr.DataArray: +def area_average( + input_data: xr.Dataset, + radii_2d: np.ndarray, + min_radius_km: float, + max_radius_km: float, + var_name: str, + level_hPa: Optional[int] = None, + levels_dimen: Optional[str] = None, + **_kwargs: Dict[str, Any] +) -> xr.DataArray: if level_hPa is None: var = input_data[var_name][0] else: var = input_data[var_name].sel(**{levels_dimen: level_hPa})[0] - avg = diag_vars.area_average(var.data, radii_2d, min_radius_km, - max_radius_km) + avg = diag_vars.area_average(var.data, radii_2d, min_radius_km, max_radius_km) units = var.attrs["units"] da_avg = _with_units(avg, units) return da_avg -def distance_to_land_lookup(lon: float, lat: float, - land_lut: diag_vars.LandLUT, - **_kwargs: Dict[str, Any]) -> xr.DataArray: +def distance_to_land_lookup( + lon: float, lat: float, land_lut: diag_vars.LandLUT, **_kwargs: Dict[str, Any] +) -> xr.DataArray: distance_km = land_lut.distance(lon, lat) units = land_lut.UNITS da_distance = _with_units(distance_km, units) @@ -79,28 +87,38 @@ def distance_to_land_lookup(lon: float, lat: float, def storm_r_theta( - track: pd.DataFrame, forecast_hour: int, init_time: dt.datetime, - **_kwargs: Dict[str, Any]) -> Tuple[xr.DataArray, xr.DataArray]: - storm_r, storm_theta = diag_vars.storm_r_theta(track, forecast_hour, - init_time) + track: pd.DataFrame, + forecast_hour: int, + init_time: dt.datetime, + **_kwargs: Dict[str, Any] +) -> Tuple[xr.DataArray, xr.DataArray]: + storm_r, storm_theta = diag_vars.storm_r_theta(track, forecast_hour, init_time) da_r = _with_units(storm_r, STORM_R_UNITS) da_theta = _with_units(storm_theta, STORM_THETA_UNITS) return da_r, da_theta def radial_and_tangential_area_average( - input_data: xr.Dataset, radii_1d: xr.DataArray, theta_2d: np.ndarray, - u_name: str, v_name: str, level_hPa: int, min_radius_km: float, - max_radius_km: float, - **_kwargs: Dict[str, Any]) -> Tuple[xr.DataArray, xr.DataArray]: + input_data: xr.Dataset, + radii_1d: xr.DataArray, + theta_2d: np.ndarray, + u_name: str, + v_name: str, + level_hPa: int, + min_radius_km: float, + max_radius_km: float, + **_kwargs: Dict[str, Any] +) -> Tuple[xr.DataArray, xr.DataArray]: u = input_data[u_name] v = input_data[v_name] rad_azavg, tan_azavg = _rad_tan(level_hPa, u, v, theta_2d) - radial_avg = diag_vars.area_average(rad_azavg, radii_1d.values, - min_radius_km, max_radius_km) - tan_avg = diag_vars.area_average(tan_azavg, radii_1d.values, min_radius_km, - max_radius_km) + radial_avg = diag_vars.area_average( + rad_azavg, radii_1d.values, min_radius_km, max_radius_km + ) + tan_avg = diag_vars.area_average( + tan_azavg, radii_1d.values, min_radius_km, max_radius_km + ) units = u.attrs["units"] da_radial = _with_units(radial_avg, units) da_tan = _with_units(tan_avg, units) @@ -108,8 +126,9 @@ def radial_and_tangential_area_average( return da_radial, da_tan -def _rad_tan(level_hPa: int, u: xr.DataArray, v: xr.DataArray, - theta_2d: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: +def _rad_tan( + level_hPa: int, u: xr.DataArray, v: xr.DataArray, theta_2d: np.ndarray +) -> Tuple[np.ndarray, np.ndarray]: u_cyl = u.interp(pressure=level_hPa).values[0, :, :] v_cyl = v.interp(pressure=level_hPa).values[0, :, :] @@ -117,15 +136,22 @@ def _rad_tan(level_hPa: int, u: xr.DataArray, v: xr.DataArray, return rad_azavg, tan_azavg -def div_vort(input_data: xr.Dataset, level_hPa: int, u_name: xr.DataArray, - v_name: xr.DataArray, theta_2d: np.ndarray, - radii_1d: xr.DataArray, radius_km: float, - **_kwargs: Dict[str, Any]) -> Tuple[xr.DataArray, xr.DataArray]: +def div_vort( + input_data: xr.Dataset, + level_hPa: int, + u_name: xr.DataArray, + v_name: xr.DataArray, + theta_2d: np.ndarray, + radii_1d: xr.DataArray, + radius_km: float, + **_kwargs: Dict[str, Any] +) -> Tuple[xr.DataArray, xr.DataArray]: u = input_data[u_name] v = input_data[v_name] rad_azavg, tan_azavg = _rad_tan(level_hPa, u, v, theta_2d) - divergence, vorticity = diag_vars._div_vort(radii_1d, radius_km, rad_azavg, - tan_azavg) + divergence, vorticity = diag_vars._div_vort( + radii_1d, radius_km, rad_azavg, tan_azavg + ) divergence *= DIV_VORT_SCALE vorticity *= DIV_VORT_SCALE @@ -136,11 +162,28 @@ def div_vort(input_data: xr.Dataset, level_hPa: int, u_name: xr.DataArray, return da_div, da_vort -def shear(forecast_hour: int, results: fcresults.ForecastHourResults, - u_name: str, v_name: str, bottom_hPa: int, top_hPa: int, - **_kwargs: Dict[str, Any]) -> Tuple[xr.DataArray, xr.DataArray]: - shear_r, shear_theta = diag_vars.shear(forecast_hour, results, u_name, - v_name, bottom_hPa, top_hPa) +def shear( + forecast_hour: int, + results: fcresults.ForecastHourResults, + u_name: str, + v_name: str, + bottom_hPa: int, + top_hPa: int, + should_convert_uv_from_10kt_to_mps=True, + **_kwargs: Dict[str, Any] +) -> Tuple[xr.DataArray, xr.DataArray]: + converter = None + if should_convert_uv_from_10kt_to_mps: + converter = met_post_process.ten_kt_to_mps + shear_r, shear_theta = diag_vars.shear( + forecast_hour, + results, + u_name, + v_name, + bottom_hPa, + top_hPa, + uv_converter=converter, + ) r_units = results.soundings[u_name].attrs["units"] da_r = _with_units(shear_r, r_units) @@ -149,13 +192,31 @@ def shear(forecast_hour: int, results: fcresults.ForecastHourResults, return da_r, da_theta -def temperature_gradient(forecast_hour: int, - results: fcresults.ForecastHourResults, u_name: str, - v_name: str, bottom_hPa: int, top_hPa: int, - lat: float, **_kwargs: Dict[str, - Any]) -> xr.DataArray: - grad = diag_vars.temperature_gradient(forecast_hour, results, u_name, - v_name, bottom_hPa, top_hPa, lat) +def temperature_gradient( + forecast_hour: int, + results: fcresults.ForecastHourResults, + u_name: str, + v_name: str, + bottom_hPa: int, + top_hPa: int, + lat: float, + should_convert_uv_from_10kt_to_mps=True, + **_kwargs: Dict[str, Any] +) -> xr.DataArray: + converter = None + if should_convert_uv_from_10kt_to_mps: + converter = met_post_process.ten_kt_to_mps + + grad = diag_vars.temperature_gradient( + forecast_hour, + results, + u_name, + v_name, + bottom_hPa, + top_hPa, + lat, + uv_converter=converter, + ) grad *= TGRD_SCALE da_grad = _with_units(grad, TGRD_UNITS) diff --git a/scripts/python/tc_diag/tc_diag_driver/met_post_process.py b/scripts/python/tc_diag/tc_diag_driver/met_post_process.py index 34ebf1be21..c662279e01 100644 --- a/scripts/python/tc_diag/tc_diag_driver/met_post_process.py +++ b/scripts/python/tc_diag/tc_diag_driver/met_post_process.py @@ -12,6 +12,10 @@ def mps_to_10kt(value: float) -> float: return scale10(mps_to_kt(value)) +def ten_kt_to_mps(value: float) -> float: + return (value / 10.0) / MPS_TO_KT + + def k_to_10c(value: float) -> float: return (value + KELVIN_TO_CELCIUS) * 10 diff --git a/scripts/python/tc_diag/tc_diag_driver/post_resample_driver.py b/scripts/python/tc_diag/tc_diag_driver/post_resample_driver.py index 252e54c91e..445470ca52 100644 --- a/scripts/python/tc_diag/tc_diag_driver/post_resample_driver.py +++ b/scripts/python/tc_diag/tc_diag_driver/post_resample_driver.py @@ -203,8 +203,14 @@ def _prep_diag_calculations( snd_comps = ce.diag_computations_from_entry(sounding_computation_specs) pi_result_names, snd_result_names = ce.get_all_result_names(pi_comps, snd_comps) + pi_result_units, snd_result_units = ce.get_all_result_units(pi_comps, snd_comps) results = fcresults.ForecastHourResults( - [forecast_hour], levels_hPa, pi_result_names, snd_result_names + [forecast_hour], + levels_hPa, + pi_result_names, + snd_result_names, + pi_result_units, + snd_result_units, ) batches = ce.get_computation_batches(pi_comps, snd_comps) diff --git a/scripts/python/tc_diag/tc_diag_driver/results.py b/scripts/python/tc_diag/tc_diag_driver/results.py index 32552bd4ce..b81612f87c 100644 --- a/scripts/python/tc_diag/tc_diag_driver/results.py +++ b/scripts/python/tc_diag/tc_diag_driver/results.py @@ -5,58 +5,100 @@ class ForecastHourResults: - """Stores diag var results for a given forecast hour. - """ - def __init__(self, forecast_hours: List[int], levels_hPa: List[int], - pressure_independent_var_names: List[str], - sounding_var_names: List[str]): + """Stores diag var results for a given forecast hour.""" + + def __init__( + self, + forecast_hours: List[int], + levels_hPa: List[int], + pressure_independent_var_names: List[str], + sounding_var_names: List[str], + pressure_independent_units: Optional[List[str]] = None, + sounding_units: Optional[List[str]] = None, + ): self.forecast_hours = forecast_hours self.levels_hPa = levels_hPa soundings_shape = (len(forecast_hours), len(levels_hPa)) - soundings_coords = { - "forecast_hour": forecast_hours, - "level_hPa": levels_hPa - } - self.soundings = self._init_dataset(soundings_shape, soundings_coords, - ("forecast_hour", "level_hPa"), - sounding_var_names) + soundings_coords = {"forecast_hour": forecast_hours, "level_hPa": levels_hPa} + self.soundings = self._init_dataset( + soundings_shape, + soundings_coords, + ("forecast_hour", "level_hPa"), + sounding_var_names, + ) # All the other variables just have forecast hour as the coordinates # For convenience - hour_results_shape = (len(forecast_hours)) + hour_results_shape = len(forecast_hours) hour_results_coords = {"forecast_hour": forecast_hours} self.pressure_independent = self._init_dataset( - hour_results_shape, hour_results_coords, ("forecast_hour"), - pressure_independent_var_names) - - def _init_dataset(self, shape: Tuple[int], coords: Dict["str", List[int]], - dims: List[str], var_names: List[str]) -> xr.Dataset: + hour_results_shape, + hour_results_coords, + ("forecast_hour"), + pressure_independent_var_names, + ) + + if pressure_independent_units is not None: + self._add_predefined_units( + self.pressure_independent, + pressure_independent_var_names, + pressure_independent_units, + ) + + if sounding_units is not None: + self._add_predefined_units( + self.soundings, sounding_var_names, sounding_units + ) + + def _add_predefined_units( + self, ds: xr.Dataset, names: List[str], units: List[str] + ) -> None: + if len(names) != len(units): + raise ValueError( + f"Length of names: {len(names)} does not match units: {len(units)}" + ) + + for name, unit in zip(names, units): + if unit is None: + continue + + ds[name].attrs["units"] = unit + + def _init_dataset( + self, + shape: Tuple[int], + coords: Dict["str", List[int]], + dims: List[str], + var_names: List[str], + ) -> xr.Dataset: data_arrays = {} for var_name in var_names: empty_array = np.full(shape, np.nan) - da = xr.DataArray(empty_array, - name=var_name, - dims=dims, - coords=coords) + da = xr.DataArray(empty_array, name=var_name, dims=dims, coords=coords) data_arrays[var_name] = da return xr.Dataset(data_vars=data_arrays, coords=coords) - def add_pressure_independent_result(self, - var_name: str, - hour: int, - result: Union[float, xr.DataArray], - units: Optional[str] = None) -> None: + def add_pressure_independent_result( + self, + var_name: str, + hour: int, + result: Union[float, xr.DataArray], + units: Optional[str] = None, + ) -> None: da = self.pressure_independent[var_name] self._add_units(da, result, units) da.loc[hour] = result - def _should_add_units(self, array: xr.DataArray, - result: Union[float, xr.DataArray], - provided_units: Optional[str]) -> bool: + def _should_add_units( + self, + array: xr.DataArray, + result: Union[float, xr.DataArray], + provided_units: Optional[str], + ) -> bool: if not hasattr(array, "attrs"): return False @@ -74,20 +116,26 @@ def _should_add_units(self, array: xr.DataArray, return True - def add_sounding_result(self, - var_name: str, - hour: int, - level_hPa: int, - result: float, - units: Optional[str] = None) -> None: + def add_sounding_result( + self, + var_name: str, + hour: int, + level_hPa: int, + result: float, + units: Optional[str] = None, + ) -> None: da = self.soundings[var_name] self._add_units(da, result, units) da.loc[hour, level_hPa] = result - def _add_units(self, array: xr.DataArray, - result: Union[float, xr.DataArray], units: Optional[str]): + def _add_units( + self, + array: xr.DataArray, + result: Union[float, xr.DataArray], + units: Optional[str], + ): if self._should_add_units(array, result, units): if units is not None: array.attrs["units"] = units diff --git a/src/tools/tc_utils/tc_diag/tc_diag.cc b/src/tools/tc_utils/tc_diag/tc_diag.cc index 99575c33c5..332587b9eb 100644 --- a/src/tools/tc_utils/tc_diag/tc_diag.cc +++ b/src/tools/tc_utils/tc_diag/tc_diag.cc @@ -1936,7 +1936,9 @@ void OutFileInfo::write_cira_diag_vals(vector &k, if(!cira_diag_out) return; - // Store lead time information + const char *method_name = "OutFileInfo::write_cira_diag_vals()"; + + // Store lead time information for standard storm and sounding data sections if(write_time) { NumArray times; for(int i=0; in_points(); i++) { @@ -1949,6 +1951,7 @@ void OutFileInfo::write_cira_diag_vals(vector &k, // Variables write AsciiTable output ConcatString cs; + string str; AsciiTable at; int n_row = m.size(); int n_col = bad_data_int; @@ -1963,6 +1966,10 @@ void OutFileInfo::write_cira_diag_vals(vector &k, n_col = 2 + m.at(*it).n(); at.set_size(n_row, n_col); + // Set spacing after the second column to 0 + // since the units column is right-padded + at.set_ics(1, 0); + // Justify columns at.set_column_just(0, LeftJust); at.set_column_just(1, LeftJust); @@ -1976,12 +1983,30 @@ void OutFileInfo::write_cira_diag_vals(vector &k, c = 0; // Diagnostic name - at.set_entry(r, c++, *it); - - // Units - cs << cs_erase << "(" - << get_diag_units(*it) << ")"; - at.set_entry(r, c++, cs); + str = *it; + + // Truncate or pad diagnostic names + if(str.length() > cira_diag_name_width) { + str = (*it).substr(0, cira_diag_name_width); + mlog << Warning << "\n" << method_name << " -> " + << "long diagnostic name \"" << (*it) + << "\" truncated to \"" << str << "\"!\n\n"; + } + str.append(cira_diag_name_width - str.length(), ' '); + at.set_entry(r, c++, str); + + // Units string + str = "(" + get_diag_units(*it) + ")"; + + // Truncate or pad units strings + if(str.length() > cira_diag_units_width) { + str = "(" + get_diag_units(*it).substr(0, cira_diag_units_width - 2) + ")"; + mlog << Warning << "\n" << method_name << " -> " + << "long diagnostic units string \"(" << get_diag_units(*it) + << ")\" truncated to \"" << str << "\"!\n\n"; + } + str.append(cira_diag_units_width - str.length(), ' '); + at.set_entry(r, c++, str); // Diagnostic values for(int i=0; i &); + const DomainInfo &, const std::set &); + void close(); void clear();