From 1a9dac93375687a3ec9fbac4337ee15e9cfe1cd2 Mon Sep 17 00:00:00 2001 From: John Halley Gotway Date: Fri, 15 Dec 2023 12:02:17 -0700 Subject: [PATCH] Per #2769, update to tc_diag_driver version 0.11.0. --- .../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 ++++-- 9 files changed, 493 insertions(+), 261 deletions(-) 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