Merge pull request #41 from leosaffin/pace

Add PACE calculation mimicking ACE calculation

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,24 +36,103 @@ 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
     -------
     array_like
         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

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

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)