Make gen_vals and apex_vals call style consistent with other huracanpy functions. Pass arrays when called as a function, pass variable names when called as a method

Author: leosaffin

huracanpy/_accessor.py

@@ -170,7 +170,7 @@ def get_pace(
             wind_units=wind_units,
             **kwargs,
         )
-        return pace_values
+        return pace_values, model
 
     def add_pace(
         self,
@@ -556,11 +556,16 @@ def get_track_duration(self, time_name="time", track_id_name="track_id"):
         )
 
     def get_gen_vals(self, time_name="time", track_id_name="track_id"):
-        return diags.get_gen_vals(self._dataset, time_name, track_id_name)
+        return diags.get_gen_vals(
+            self._dataset, self._dataset[time_name], self._dataset[track_id_name]
+        )
 
     def get_apex_vals(self, varname, track_id_name="track_id", stat="max"):
         return diags.get_apex_vals(
-            self._dataset, varname, track_id_name=track_id_name, stat=stat
+            self._dataset,
+            variable=self._dataset[varname],
+            track_id=self._dataset[track_id_name],
+            stat=stat,
         )
 
     # ---- climato

huracanpy/diags/_track_stats.py

@@ -2,6 +2,9 @@
 Module containing functions to compute track statistics
 """
 
+import numpy as np
+import pandas as pd
+
 
 def get_track_duration(time, track_ids):
     """
@@ -25,39 +28,50 @@ def get_track_duration(time, track_ids):
     return duration
 
 
-def get_gen_vals(tracks, time_name="time", track_id_name="track_id"):
+def get_gen_vals(tracks, time, track_id):
     """
     Shows the attributes for the genesis point of each track
 
     Parameters
     ----------
     tracks : xarray.DataSet
+    time : array_like
+    track_id : xarray.Dataset
 
     Returns
     -------
     xarray.Dataset
         Dataset containing only genesis points, with track_id as index.
 
     """
+    # It is 470 times much faster to switch to a dataframe...
+    # Use the sortby/groupby with pandas to find the relevant indices in the original
+    # Dataset by passing an index (named idx to not clash with "index")
+    df = pd.DataFrame(
+        data=dict(
+            idx=np.arange(len(track_id)),
+            time=np.array(time),
+            track_id=np.array(track_id),
+        )
+    )
+    idx = np.array(df.sort_values("time").groupby("track_id").first().idx)
 
-    return (
-        tracks.to_dataframe()
-        .sort_values(time_name)
-        .groupby(track_id_name)
-        .first()
-        .to_xarray()
-    )  # It is 470 times much faster to switch to a dataframe...
+    # Could check that track_id is 1d, but the function would already have failed by now
+    # if not
+    dim = track_id.dims[0]
+    return tracks.isel(**{dim: idx})
 
 
-def get_apex_vals(tracks, varname, stat="max", track_id_name="track_id"):
+def get_apex_vals(tracks, variable, track_id, stat="max"):
     """
     Shows the attribute for the extremum point of each track
 
     Parameters
     ----------
     tracks : xarray.DataSet
-    var : str
+    variable : array_like
         The extremum variable
+    track_id : xarray.DataArray
     stat : str, optional
         Type of extremum. Can be "min" or "max". The default is "max".
 
@@ -82,10 +96,16 @@ def get_apex_vals(tracks, varname, stat="max", track_id_name="track_id"):
     else:
         raise NotImplementedError("stat not recognized. Please use one of {min, max}")
 
-    return (
-        tracks.to_dataframe()
-        .sort_values(varname, ascending=asc)
-        .groupby(track_id_name)
-        .first()
-        .to_xarray()
-    )  # It is 350 times much faster to switch to a dataframe..
+    # It is 350 times much faster to switch to a dataframe.
+    # Use the same trick as with gen_vals
+    df = pd.DataFrame(
+        data=dict(
+            idx=np.arange(len(variable)),
+            var=np.array(variable),
+            track_id=np.array(track_id),
+        )
+    )
+    idx = np.array(df.sort_values("var", ascending=asc).groupby("track_id").first().idx)
+
+    dim = track_id.dims[0]
+    return tracks.isel(**{dim: idx})

tests/test_accessor.py

@@ -1,8 +1,9 @@
 import pytest
 
-import huracanpy
-
 import numpy as np
+import xarray as xr
+
+import huracanpy
 
 
 # %% DataArrayAccessor
@@ -77,6 +78,13 @@ def test_nunique():
         ),
         (huracanpy.diags.get_freq, ["track_id"], "freq", {}),
         (huracanpy.diags.get_tc_days, ["time", "track_id"], "tc_days", {}),
+        # (huracanpy.diags.get_gen_vals, ["all", "time", "track_id"], "gen_vals", {}),
+        # (
+        #     huracanpy.diags.get_apex_vals,
+        #     ["all", "wind10", "track_id"],
+        #     "apex_vals",
+        #     {"varname": "wind10"},
+        # ),
     ],
 )
 def test_accessor_methods_match_functions(
@@ -93,13 +101,20 @@ def test_accessor_methods_match_functions(
             "track_duration",
             "freq",
             "tc_days",
+            "gen_vals",
+            "apex_vals",
         ]:
             pytest.skip(f"Accessor function add_{accessor_name} does not exist")
         elif accessor_name in ["ace"] and "sum_by" in accessor_function_kwargs:
             pytest.skip(f"sum_by not a valid argument for add_{accessor_name}")
 
     # Call the huracanpy function
-    result = function(*[tracks_csv[var] for var in function_args])
+    # Get the function arguments as arrays. Use "all" as a wildcard for the full dataset
+    function_args = [
+        tracks_csv[var] if not var == "all" else tracks_csv for var in function_args
+    ]
+    result = function(*function_args)
+
     # Call the accessor method
     result_accessor = getattr(tracks_csv.hrcn, f"{call_type}_{accessor_name}")(
         **accessor_function_kwargs
@@ -127,14 +142,18 @@ def test_accessor_methods_match_functions(
 
 
 # %% DatasetAccessor
+# Currently keeping tests here that return more than just a DataArray as the testing is
+# less generic
 def test_get_methods(tracks_csv):
     """Test get_ accessors output is same as function"""
     data = tracks_csv
 
     ## - pace
-    pace_acc = data.hrcn.get_pace(pressure_name="slp", wind_name="wind10")
+    pace_acc, _ = data.hrcn.get_pace(pressure_name="slp", wind_name="wind10")
     pace_fct, model_fct = huracanpy.tc.pace(data.slp, data.wind10)
-    assert not any(pace_acc != pace_fct), "accessor output differs from function output"
+    np.testing.assert_array_equal(
+        pace_acc, pace_fct, err_msg="accessor output differs from function output"
+    )
 
     ## - time components
     year_acc, month_acc, day_acc, hour_acc = data.hrcn.get_time_components(
@@ -143,10 +162,18 @@ def test_get_methods(tracks_csv):
     year_fct, month_fct, day_fct, hour_fct = huracanpy.utils.get_time_components(
         data.time
     )
-    assert all(year_acc == year_fct), "Year component does not match"
-    assert all(month_acc == month_fct), "Month component does not match"
-    assert all(day_acc == day_fct), "Day component does not match"
-    assert all(hour_acc == hour_fct), "Hour component does not match"
+    np.testing.assert_array_equal(
+        year_acc, year_fct, err_msg="Year component does not match"
+    )
+    np.testing.assert_array_equal(
+        month_acc, month_fct, err_msg="Month component does not match"
+    )
+    np.testing.assert_array_equal(
+        day_acc, day_fct, err_msg="Day component does not match"
+    )
+    np.testing.assert_array_equal(
+        hour_acc, hour_fct, err_msg="Hour component does not match"
+    )
 
     ## - track pace
     pace_acc, _ = data.hrcn.get_pace(wind_name="wind10", sum_by="track_id")
@@ -162,21 +189,17 @@ def test_get_methods(tracks_csv):
         time_name="time",
         track_id_name="track_id",
     )
-    gen_vals_fct = huracanpy.diags.get_gen_vals(
-        data,
-    )
-    assert gen_vals_acc.equals(
-        gen_vals_fct
-    ), "Genesis Values accessor output differs from function output"
+    gen_vals_fct = huracanpy.diags.get_gen_vals(data, data.time, data.track_id)
+    xr.testing.assert_equal(gen_vals_acc, gen_vals_fct)
 
     ## - Apex Values
     apex_vals_acc = data.hrcn.get_apex_vals(
         track_id_name="track_id", varname="wind10", stat="max"
     )
-    apex_vals_fct = huracanpy.diags.get_apex_vals(data, varname="wind10", stat="max")
-    assert apex_vals_acc.equals(
-        apex_vals_fct
-    ), "Genesis Values accessor output differs from function output"
+    apex_vals_fct = huracanpy.diags.get_apex_vals(
+        data, data.wind10, data.track_id, stat="max"
+    )
+    xr.testing.assert_equal(apex_vals_acc, apex_vals_fct)
 
 
 def test_interp_methods():

tests/test_diags/test_track_stats.py

@@ -11,13 +11,13 @@ def test_duration():
 
 def test_gen_vals():
     data = huracanpy.load(huracanpy.example_csv_file, source="csv")
-    G = huracanpy.diags.get_gen_vals(data)
+    G = huracanpy.diags.get_gen_vals(data, data.time, data.track_id)
     assert G.time.dt.day.mean() == 10
 
 
-def test_extremum_vals():
-    data = huracanpy.load(huracanpy.example_csv_file, source="csv")
-    M = huracanpy.diags.get_apex_vals(data, "wind10", "max")
-    m = huracanpy.diags.get_apex_vals(data, "slp", "min")
+def test_apex_vals():
+    data = huracanpy.load(huracanpy.example_csv_file)
+    M = huracanpy.diags.get_apex_vals(data, data.wind10, data.track_id, "max")
+    m = huracanpy.diags.get_apex_vals(data, data.slp, data.track_id, "min")
     assert M.time.dt.day.mean() == 15
     assert m.lat.mean() == -27