diff --git a/wrf_analysis_toolkit/pyproject.toml b/wrf_analysis_toolkit/pyproject.toml index cd8f93f..4a4763e 100644 --- a/wrf_analysis_toolkit/pyproject.toml +++ b/wrf_analysis_toolkit/pyproject.toml @@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta" [project] name = "wrf_analysis_toolkit" -version = "2.4.2" +version = "2.5.0" requires-python = ">=3.10" dependencies = [ "netCDF4>=1.7", diff --git a/wrf_analysis_toolkit/wrf_analysis_toolkit/GetSensVar.py b/wrf_analysis_toolkit/wrf_analysis_toolkit/GetSensVar.py index dff3d98..c98bf1c 100644 --- a/wrf_analysis_toolkit/wrf_analysis_toolkit/GetSensVar.py +++ b/wrf_analysis_toolkit/wrf_analysis_toolkit/GetSensVar.py @@ -4,10 +4,23 @@ from wrf import to_np, getvar, g_geoht, interplevel import numpy as np +from copy import deepcopy import wrf_analysis_toolkit.SensibleVariables as sv import wrf_analysis_toolkit.Frontogenesis as Frontogenesis +from wrf_analysis_toolkit.utils import destagger_var, project_vector + +MOMENTUM_TEND_DICT = { + "tendhadv": {"var_u": "RU_TEND_HADV", "var_v": "RV_TEND_HADV"}, + "tendvadv": {"var_u": "RU_TEND_VADV", "var_v": "RU_TEND_VADV"}, + "tendpgf": {"var_u": "RU_TEND_PGF", "var_v": "RU_TEND_PGF"}, + "tendcor": {"var_u": "RU_TEND_COR", "var_v": "RU_TEND_COR"}, + "tendcurv": {"var_u": "RU_TEND_CURV", "var_v": "RU_TEND_CURV"}, + "tendfpbl": {"var_u": "RU_TENDF_PBL", "var_v": "RU_TENDF_PBL"}, + "tendfcu": {"var_u": "RU_TENDF_CU", "var_v": "RU_TENDF_CU"}, + "tendfdiff": {"var_u": "RU_TENDF_DIFF", "var_v": "RU_TENDF_DIFF"} +} def GetSensVar(ncfile, svariable, windbarbs=0, time=0, varprevv=None): u = v = varv = None @@ -32,9 +45,10 @@ def GetSensVar(ncfile, svariable, windbarbs=0, time=0, varprevv=None): # Converts accumulated rain to "hourly" rain (given hourly time indices) if varprevv is not None: var.values = var.values - varprevv - + # For 3D +value variables, interpolated at interpvalue of interpvar elif svariable.dim == 4: + interpvar = getvar(ncfile, svariable.interpvar, timeidx=time) if svariable.wrfname is not None: d4var = getvar(ncfile, svariable.wrfname, timeidx=time) @@ -45,7 +59,42 @@ def GetSensVar(ncfile, svariable, windbarbs=0, time=0, varprevv=None): F3D = Frontogenesis.frontogenesis3D(ncfile, time) d4var = getvar(ncfile, svariable.interpvar, timeidx=time) d4var.values = F3D + elif any([k in svariable.outfile.lower() for k in MOMENTUM_TEND_DICT]): + for k in MOMENTUM_TEND_DICT: + if k in svariable.outfile.lower(): + tend_name = k + break + print(f"Extracting variables to calculate {tend_name}") + var_u = getvar(ncfile, MOMENTUM_TEND_DICT[tend_name]["var_u"], timeidx=time) + attrs = var_u.attrs + var_u = destagger_var(var_u, meta=False) + var_v = getvar(ncfile, MOMENTUM_TEND_DICT[tend_name]["var_v"], timeidx=time) + var_v = destagger_var(var_v, meta=False) + ua = getvar(ncfile, "ua", timeidx=time) + va = getvar(ncfile, "va", timeidx=time) + wspd = getvar(ncfile, "wspd", timeidx=time) + + # Calculate the variable projected onto the unit vector of + # horizontal winds, to calculate the along-flow values + # This is quickest if calculated using arrays without metadata, + # and the metadata are copied from the interpvar afterwards + d4var = deepcopy(interpvar) + print(f"Projecting {tend_name} onto unit wind vector") + d4var.values = project_vector(var_u, var_v, ua, va, wspd) + d4var.attrs.update(attrs) + + else: + if svariable.wrfname is not None: + raise ValueError(f"Failed to extract variable {svariable.wrfname}") + else: + raise ValueError(f"Failed to extract variable for {svariable.outfile}") + + # Destagger the variable if it is staggered + d4var = destagger_var(d4var, meta_var=interpvar, meta=True) + + # interpolate variable var = interplevel(d4var, interpvar, svariable.interpvalue) + # Special variable computation if "AirTempDif6h" in svariable.outfile: # Temperature difference in 6h @@ -59,6 +108,7 @@ def GetSensVar(ncfile, svariable, windbarbs=0, time=0, varprevv=None): else: varv = np.append(varprevv[1:], [var.values], axis=0) var.values = var.values - varprevv[0] + elif "AirTempDif12h" in svariable.outfile: # Temperature difference in 12h if varprevv is None: @@ -82,6 +132,7 @@ def GetSensVar(ncfile, svariable, windbarbs=0, time=0, varprevv=None): elif svariable.outfile in ["InstRain"]: # InstRain (R) from SimRadarReflectivity1km (dBZ) using Marshall-Palmer: Z = 10^(dBZ/10) = 200*R^1.6 var.values = (0.005 * 10 ** (0.1 * var.values)) ** (0.625) + if windbarbs: # Get wind speed components at interpvalue ua = getvar(ncfile, "ua", timeidx=time) diff --git a/wrf_analysis_toolkit/wrf_analysis_toolkit/SensibleVariables.py b/wrf_analysis_toolkit/wrf_analysis_toolkit/SensibleVariables.py index 1bdea88..c572a7f 100644 --- a/wrf_analysis_toolkit/wrf_analysis_toolkit/SensibleVariables.py +++ b/wrf_analysis_toolkit/wrf_analysis_toolkit/SensibleVariables.py @@ -981,7 +981,6 @@ def create_GeoPotHeight_at( range_max=128, ) -# Frontgenesis Frontogenesis925 = svariable( dim=4, ptitle="Petterssen Frontogenesis at 925 hPa [K/(100km 3h)]", @@ -1257,6 +1256,29 @@ def create_WindSpeed_at(interpvalue, range_min=0, range_max=60, nticks=12, nlevs WindSpeed500 = create_WindSpeed_at(500) WindSpeed300 = create_WindSpeed_at(300, range_min=0, range_max=80) +def create_U_at(interpvalue, range_min=0, range_max=60, nticks=12, nlevs=12): + return svariable( + dim=4, + wrfname="U", + ptitle=f"Wind Speed in West-East direction at {interpvalue} hPa [m/s]", + outfile=f"U{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + colormap=get_cmap("YlGnBu"), + nticks=nticks, + nlevs=nlevs, + range_min=range_min, + range_max=range_max, + windbarbs=True, + ) + +U925 = create_U_at(925) +U850 = create_U_at(850) +U700 = create_U_at(700) +U500 = create_U_at(500) +U300 = create_U_at(300, range_min=0, range_max=80) + + # Potential Vorticity pv_range_min = -10 pv_range_max = 10 @@ -1507,3 +1529,333 @@ def create_PotentialVorticity_at( range_min=-60, range_max=40, ) + +# QVapor +def create_QVapor_at(interpvalue, range_min=0, range_max=0.01, nticks=11, nlevs=11): + return svariable( + dim=4, + wrfname="QVAPOR", + ptitle=f"Water Vapour Mixing Ratio at {interpvalue} hPa [kg kg-1]", + outfile=f"QVapor{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + colormap=get_cmap("YlGnBu"), + nticks=nticks, + nlevs=nlevs, + range_min=range_min, + range_max=range_max, + windbarbs=True, + ) +QVapor925 = create_QVapor_at(925) +QVapor850 = create_QVapor_at(850) +QVapor700 = create_QVapor_at(700) +QVapor500 = create_QVapor_at(500) +QVapor300 = create_QVapor_at(300) + +# Sensible variables for analysing momentum tendency terms +# Raw variables, in x- or y-direction, not projected onto unit vector of wind +# Mass Tendency 2D, U direction +mu_range_min = -100 +mu_range_max = 100 +mu_max_frac = 0.55 + min(0.45, (pv_range_max / (pv_range_max - pv_range_min))) +mu_min_frac = 0.55 + min(0, (pv_range_min / (pv_range_max - pv_range_min))) +mu_nticks = 21 +mu_nlevs = 21 +UHorizAdvMomentum = svariable( + dim=4, + wrfname="RU_TEND_HADV", + ptitle=f"Horizontal advection of zonal coupled momentum term in X [Pa m s-2]", + outfile=f"UHorizAdvMomentum", + colormap=get_cmap("YlGnBu"), +) + +def create_UHorizAdvMomentum_at( + interpvalue, + range_min=mu_range_min, + range_max=mu_range_max, + nticks=mu_nticks, + nlevs=mu_nlevs +): + min_frac = 0.55 + min(0, (range_min / (range_max - range_min))) + max_frac = 0.55 + min(0.45, (range_max / (range_max - range_min))) + return svariable( + dim=4, + wrfname="RU_TEND_HADV", + ptitle=f"Horizontal advection of zonal coupled momentum term in X at {interpvalue} hPa [Pa m s-2]", + outfile=f"UHorizAdvMomentum{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + colormap=cmr.get_sub_cmap("PuOr", min_frac, max_frac, N=nlevs), + nticks=nticks, + nlevs=nlevs, + range_min=range_min, + range_max=range_max, + ) +UMassTendency925 = create_UHorizAdvMomentum_at(925) #, range_min=-5, range_max=5) +UMassTendency850 = create_UHorizAdvMomentum_at(850) #, range_min=-5, range_max=5) +UMassTendency700 = create_UHorizAdvMomentum_at(700) #, range_min=-5, range_max=5) +UMassTendency500 = create_UHorizAdvMomentum_at(500) #, range_min=-5, range_max=5) +UMassTendency300 = create_UHorizAdvMomentum_at(300) #, range_min=-10, range_max=10) + +# Combined terms projected onto unit vector of wind +# MOMENTUM_TEND_DICT = { +# "tend_hadv": {"var_u": "ru_tend_hadv", "var_v": "rv_tend_hadv"}, +# "tend_vadv": {"var_u": "ru_tend_vadv", "var_v": "ru_tend_vadv"}, +# "tend_pgf": {"var_u": "ru_tend_pgf", "var_v": "rv_tend_pgf"}, +# "tend_cor": {"var_u": "ru_tend_cor", "var_v": "rv_tend_cor"}, +# "tend_curv": {"var_u": "ru_tend_curv", "var_v": "rv_tend_curv"}, +# "tendf_pbl": {"var_u": "ru_tendf_pbl", "var_v": "rv_tendf_pbl"}, +# "tendf_cu": {"var_u": "ru_tendf_cu", "var_v": "rv_tendf_cu"}, +# "tendf_diff": {"var_u": "ru_tendf_diff", "var_v": "rv_tendf_diff"} +# } +tend_scale = "bounds" +tend_bounds = [-400, -200, -100, -50, -25, 25, 50, 100, 200, 400] +tend_colormap = ListedColormap( + [ + "darkred", + "red", + "darkorange", + "gold", + "white", + "greenyellow", + "limegreen", + "forestgreen", + "darkgreen", + ] +) +tend_range_min=-400 +tend_range_max=400 + +def create_TendHADV_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + """ + Calculates: + (ru_tend_hadv * U + rv_tend_hadv * V) / |wind_spd| + + GetSensVar works out what to calculate based on outname rather than wrfname + """ + return svariable( + dim=4, + wrfname=None, + ptitle=f"Horizontal advection of zonal coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendHADV{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendHADV925 = create_TendHADV_at(925) #, range_min=-5, range_max=5) +TendHADV850 = create_TendHADV_at(850) #, range_min=-5, range_max=5) +TendHADV700 = create_TendHADV_at(700) #, range_min=-5, range_max=5) +TendHADV500 = create_TendHADV_at(500) #, range_min=-5, range_max=5) +TendHADV300 = create_TendHADV_at(300) #, range_min=-10, range_max=10) + + + +def create_TendVADV_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + return svariable( + dim=4, + wrfname=None, + ptitle=f"Vertical advection of zonal coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendVADV{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendVADV925 = create_TendVADV_at(925) #, range_min=-5, range_max=5) +TendVADV850 = create_TendVADV_at(850) #, range_min=-5, range_max=5) +TendVADV700 = create_TendVADV_at(700) #, range_min=-5, range_max=5) +TendVADV500 = create_TendVADV_at(500) #, range_min=-5, range_max=5) +TendVADV300 = create_TendVADV_at(300) #, range_min=-10, range_max=10) + +def create_TendPGF_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + return svariable( + dim=4, + wrfname=None, + ptitle=f"Pressure Gradient Force zonal coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendPGF{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendPGF925 = create_TendPGF_at(925) #, range_min=-5, range_max=5) +TendPGF850 = create_TendPGF_at(850) #, range_min=-5, range_max=5) +TendPGF700 = create_TendPGF_at(700) #, range_min=-5, range_max=5) +TendPGF500 = create_TendPGF_at(500) #, range_min=-5, range_max=5) +TendPGF300 = create_TendPGF_at(300) #, range_min=-10, range_max=10) + +def create_TendCOR_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + return svariable( + dim=4, + wrfname=None, + ptitle=f"Coriolis zonal coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendCOR{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendCOR925 = create_TendCOR_at(925) #, range_min=-5, range_max=5) +TendCOR850 = create_TendCOR_at(850) #, range_min=-5, range_max=5) +TendCOR700 = create_TendCOR_at(700) #, range_min=-5, range_max=5) +TendCOR500 = create_TendCOR_at(500) #, range_min=-5, range_max=5) +TendCOR300 = create_TendCOR_at(300) #, range_min=-10, range_max=10) + +def create_TendCURV_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + return svariable( + dim=4, + wrfname=None, + ptitle=f"Curvature zonal coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendCURV{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendCURV925 = create_TendCURV_at(925) #, range_min=-5, range_max=5) +TendCURV850 = create_TendCURV_at(850) #, range_min=-5, range_max=5) +TendCURV700 = create_TendCURV_at(700) #, range_min=-5, range_max=5) +TendCURV500 = create_TendCURV_at(500) #, range_min=-5, range_max=5) +TendCURV300 = create_TendCURV_at(300) #, range_min=-10, range_max=10) + + +def create_TendfPBL_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + return svariable( + dim=4, + wrfname=None, + ptitle=f"PBL coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendfPBL{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendfPBL925 = create_TendfPBL_at(925) #, range_min=-5, range_max=5) +TendfPBL850 = create_TendfPBL_at(850) #, range_min=-5, range_max=5) +TendfPBL700 = create_TendfPBL_at(700) #, range_min=-5, range_max=5) +TendfPBL500 = create_TendfPBL_at(500) #, range_min=-5, range_max=5) +TendfPBL300 = create_TendfPBL_at(300) #, range_min=-10, range_max=10) + +def create_TendfCU_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + return svariable( + dim=4, + wrfname=None, + ptitle=f"Microphysics coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendfCU{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendfCU925 = create_TendfCU_at(925) #, range_min=-5, range_max=5) +TendfCU850 = create_TendfCU_at(850) #, range_min=-5, range_max=5) +TendfCU700 = create_TendfCU_at(700) #, range_min=-5, range_max=5) +TendfCU500 = create_TendfCU_at(500) #, range_min=-5, range_max=5) +TendfCU300 = create_TendfCU_at(300) #, range_min=-10, range_max=10) + +def create_TendfDIFF_at( + interpvalue, + scale=tend_scale, + bounds=tend_bounds, + colormap=tend_colormap, + range_min=tend_range_min, + range_max=tend_range_max +): + return svariable( + dim=4, + wrfname=None, + ptitle=f"Diffusion coupled momentum term projected onto unit vector at {interpvalue} hPa [Pa m s-2]", + outfile=f"TendfDIFF{interpvalue}", + interpvar="pressure", + interpvalue=interpvalue, + scale=scale, + bounds=bounds, + colormap=colormap, + range_min=range_min, + range_max=range_max, + ) + +TendfDIFF925 = create_TendfDIFF_at(925) #, range_min=-5, range_max=5) +TendfDIFF850 = create_TendfDIFF_at(850) #, range_min=-5, range_max=5) +TendfDIFF700 = create_TendfDIFF_at(700) #, range_min=-5, range_max=5) +TendfDIFF500 = create_TendfDIFF_at(500) #, range_min=-5, range_max=5) +TendfDIFF300 = create_TendfDIFF_at(300) #, range_min=-10, range_max=10) + diff --git a/wrf_analysis_toolkit/wrf_analysis_toolkit/utils.py b/wrf_analysis_toolkit/wrf_analysis_toolkit/utils.py index bcb7d14..75cc9bd 100644 --- a/wrf_analysis_toolkit/wrf_analysis_toolkit/utils.py +++ b/wrf_analysis_toolkit/wrf_analysis_toolkit/utils.py @@ -6,7 +6,9 @@ import re from copy import deepcopy from datetime import datetime, timedelta -from wrf import CoordPair, ll_to_xy +from wrf import CoordPair, ll_to_xy, destagger +import numpy as np +from xarray import Dataset, DataArray import wrf_analysis_toolkit.SensibleVariables as sv @@ -200,3 +202,68 @@ def latlon_check(ncfile: Dataset, latlon: tuple): raise ValueError( f"Point ({lat}, {lon}) is outside the WRF domain" ) + +def destagger_var( + var: DataArray, + meta_var: DataArray | None=None, + meta: bool=False +): + stagger_dim = None + for i, dim in enumerate(var.dims): + if dim.endswith("_stag"): + stagger_dim = i + break + + # Destagger the variable if it is staggered + if stagger_dim: + try: + var_out = destagger(var, stagger_dim, meta=meta) + if meta: + if meta_var is None: + raise ValueError("Need a sample meta_var to attatch coordinate fields to destaggared array") + # Manually assign coordinates from meta_var because these aren't done automatically + var_out = var_out.assign_coords(coords=meta_var.coords) + var_out.attrs.update(stagger=meta_var.attrs['stagger'], coordinates=meta_var.attrs['coordinates']) + var_out["Time"] = meta_var.Time + + return var_out + + except: + raise ValueError("Unable to destagger variable") + + # Otherwise return the original variable unchanged + else: + return var + +def project_vector( + var_u: np.ndarray, + var_v: np.ndarray, + wind_u: np.ndarray, + wind_v: np.ndarray, + wind_spd = None, +): + """ + Function to project a given vector (with terms in the X and Y directions) + onto the unit vectors of wind in the X and Y directions (U and V), + in order to calculate the along-flow values of the vector + + Inputs: + - var_u: values of variable of in interest in the x direction, must be destaggered + - var_v: values of variable of in interest in the y direction, must be destaggered + - wind_u: values of wind speed in the x direction, must be destaggered + - wind_v: values of wind speed in the y direction, must be destaggered + - wind_spd: (optional) magnitude of wind spped, will be calculated from vectors if not given + """ + if any([ + var_u.shape != var_v.shape, + var_u.shape != wind_u.shape, + var_u.shape != wind_v.shape + ]): + raise ValueError("project_vector: Shape of vectors do not match, do they need destaggering?") + + if wind_spd is None: + wind_mag = np.sqrt(np.square(wind_u) + np.square(wind_v)) + else: + wind_mag = wind_spd + + return (var_u*wind_u + var_v*wind_v) / wind_mag