diff --git a/huracanpy/diags/track_stats.py b/huracanpy/diags/track_stats.py index c970ab9..90598f4 100644 --- a/huracanpy/diags/track_stats.py +++ b/huracanpy/diags/track_stats.py @@ -2,12 +2,9 @@ Module containing functions to compute track statistics """ -# import xarray as xr -# import pint -# from metpy.xarray import preprocess_and_wrap from metpy.units import units -from huracanpy.utils.ace import ace_by_point +from huracanpy.utils.ace import ace_by_point, pace_by_point def ace_by_track( @@ -16,6 +13,7 @@ def ace_by_track( threshold=34 * units("knots"), wind_units="m s-1", keep_ace_by_point=False, + ace_varname="ace", ): r"""Calculate accumulate cyclone energy (ACE) for each track @@ -38,7 +36,11 @@ def ace_by_track( assume it is in these units before converting to knots keep_ace_by_point : bool, default=False If True the ACE calculated from each point of the input wind is saved in the - input tracks dataset as "ace" + input tracks dataset as `ace_varname` + ace_varname : str, default="ace" + The name to give the variable for ACE at each point added to the `tracks` + dataset. Change this if you want to have a different variable name or want to + avoid overwriting an existing variable in the dataset named `ace` Returns ------- @@ -46,16 +48,91 @@ def ace_by_track( The ACE for each track in wind """ - tracks["ace"] = ace_by_point(wind, threshold, wind_units) + tracks[ace_varname] = ace_by_point(wind, threshold, wind_units) - ace_by_storm = tracks.groupby("track_id").map(lambda x: x.ace.sum()) + ace_by_storm = tracks.groupby("track_id").map(lambda x: x[ace_varname].sum()) if not keep_ace_by_point: - del tracks["ace"] + del tracks[ace_varname] return ace_by_storm +def pace_by_track( + tracks, + pressure, + wind=None, + model=None, + threshold_wind=None, + threshold_pressure=None, + wind_units="m s-1", + keep_pace_by_point=False, + pace_varname="pace", + **kwargs, +): + """Calculate a pressure-based accumulated cyclone energy (PACE) for each individual + point + + PACE is calculated the same way as ACE, but the wind is derived from fitting a + pressure-wind relationship and calculating wind values from pressure using this fit + + Example + ------- + This function can be called in two ways + + 1. Pass the pressure and wind to fit a pressure-wind relationship to the data and + then calculate pace from the winds derived from this fit + + >>> pace, pw_model = pace_by_point(pressure, wind) + + The default model to fit is a quadratic polynomial + (:py:class:`numpy.polynomial.polynomial.Polynomial` with `deg=2`) + + 2. Pass just the pressure and an already fit model to calculate the wind speeds from + this model + + >>> pace, _ = pace_by_point(pressure, model=pw_model) + + Parameters + ---------- + tracks : xarray.Dataset + pressure : array_like + wind : array_like, optional + model : str, class, or object, optional + threshold_wind : scalar, optional + threshold_pressure : scalar, optional + wind_units : str, default="m s-1" + keep_pace_by_point : bool, default=False + If True the PACE calculated from each point of the input wind is saved in the + input tracks dataset as `pace_varname` + pace_varname : str, default="pace" + **kwargs + + Returns + ------- + pace_values : array_like + + model : object + + """ + tracks[pace_varname], model = pace_by_point( + pressure, + wind=wind, + model=model, + threshold_wind=threshold_wind, + threshold_pressure=threshold_pressure, + wind_units=wind_units, + **kwargs, + ) + + pace_by_storm = tracks.groupby("track_id").map(lambda x: x[pace_varname].sum()) + + if not keep_pace_by_point: + del tracks[pace_varname] + + return pace_by_storm, model + + def duration(time, track_ids): """ Compute the duration of each track diff --git a/huracanpy/utils/ace.py b/huracanpy/utils/ace.py index a62eb81..ff02f4e 100644 --- a/huracanpy/utils/ace.py +++ b/huracanpy/utils/ace.py @@ -2,6 +2,7 @@ Module containing functions to compute ACE """ +from numpy.polynomial.polynomial import Polynomial import xarray as xr import pint from metpy.xarray import preprocess_and_wrap @@ -51,11 +52,113 @@ def _ace_by_point(wind, threshold=34 * units("knots"), wind_units="m s-1"): wind = wind * units(wind_units) wind = wind.to(units("knots")) - if not isinstance(threshold, pint.Quantity) or threshold.unitless: - threshold = threshold * units(wind_units) + if threshold is not None: + if not isinstance(threshold, pint.Quantity) or threshold.unitless: + threshold = threshold * units(wind_units) - wind[wind < threshold] = 0 * units("knots") + wind[wind < threshold] = 0 * units("knots") ace_values = (wind**2.0) * 1e-4 return ace_values + + +def pace_by_point( + pressure, + wind=None, + model=None, + threshold_wind=None, + threshold_pressure=None, + wind_units="m s-1", + **kwargs, +): + """Calculate a pressure-based accumulated cyclone energy (PACE) for each individual + point + + PACE is calculated the same way as ACE, but the wind is derived from fitting a + pressure-wind relationship and calculating wind values from pressure using this fit + + Example + ------- + This function can be called in two ways + + 1. Pass the pressure and wind to fit a pressure-wind relationship to the data and + then calculate pace from the winds derived from this fit + + >>> pace, pw_model = pace_by_point(pressure, wind) + + The default model to fit is a quadratic polynomial + (:py:class:`numpy.polynomial.polynomial.Polynomial` with `deg=2`) + + 2. Pass just the pressure and an already fit model to calculate the wind speeds from + this model + + >>> pace, _ = pace_by_point(pressure, model=pw_model) + + Parameters + ---------- + pressure : array_like + wind : array_like, optional + model : str or class, optional + threshold_wind : scalar, optional + threshold_pressure : scalar, optional + wind_units : str, default="m s-1" + **kwargs + + Returns + ------- + pace_values : array_like + + model : object + + """ + model_wind, model = pressure_wind_relationship( + pressure, wind=wind, model=model, **kwargs + ) + pace_values = ace_by_point( + model_wind, threshold=threshold_wind, wind_units=wind_units + ) + + if threshold_pressure is not None: + pace_values[pressure > threshold_pressure] = 0.0 + + return pace_values, model + + +def pressure_wind_relationship(pressure, wind=None, model=None, **kwargs): + if isinstance(model, str): + if model.lower() == "z2021": + model = pw_z2021 + elif model.lower() == "holland": + model = pw_holland + + elif wind is not None: + if model is None: + # Here, we calculate a quadratic P/W fit based off of the "control" + if "deg" not in kwargs: + kwargs["deg"] = 2 + model = Polynomial.fit(pressure, wind, **kwargs) + + else: + model = model.fit(pressure, wind, **kwargs) + + elif model is None: + raise ValueError( + "Need to specify either wind or model to calculate pressure-wind relation" + ) + + wind_from_fit = model(pressure) + + return wind_from_fit, model + + +# Pre-determined pressure-wind relationships +_z2021 = Polynomial([1.43290190e01, 5.68356519e-01, -1.05371378e-03]) + + +def pw_z2021(pressure): + return _z2021(1010.0 - pressure) + + +def pw_holland(pressure): + return 2.3 * (1010.0 - pressure) ** 0.76 diff --git a/tests/test_diags/test_track_stats.py b/tests/test_diags/test_track_stats.py index 6686518..2fdb9f4 100644 --- a/tests/test_diags/test_track_stats.py +++ b/tests/test_diags/test_track_stats.py @@ -11,6 +11,24 @@ def test_ace(tracks_csv): assert isinstance(ace.data, np.ndarray) +def test_pace(tracks_csv): + # Pass wind values to fit a (quadratic) model to the pressure-wind relationship + pace, model = huracanpy.diags.track_stats.pace_by_track( + tracks_csv, tracks_csv.slp, wind=tracks_csv.wind10 + ) + + np.testing.assert_allclose(pace, np.array([4.34978137, 2.65410482, 6.09892875])) + + # Call with the already fit model instead of wind values + pace, _ = huracanpy.diags.track_stats.pace_by_track( + tracks_csv, + tracks_csv.slp, + model=model, + ) + + np.testing.assert_allclose(pace, np.array([4.34978137, 2.65410482, 6.09892875])) + + def test_duration(): data = huracanpy.load(huracanpy.example_csv_file, tracker="csv") d = huracanpy.diags.track_stats.duration(data.time, data.track_id)