Make accelerated functions equivalent to naive slow approaches and add tests to verify

Author: leosaffin

huracanpy/_subset.py

@@ -1,28 +1,36 @@
+import numpy as np
 import xarray as xr
 
 __all__ = ["trackswhere", "sel_id"]
 
 
-def sel_id(data, track_id, track_ids):
+def sel_id(tracks, track_ids, track_id):
     """Select an individual track from a set of tracks by ID
 
     Parameters
     ----------
-    data : xarray.Dataset
-    track_id : scalar
+    tracks : xarray.Dataset
     track_ids : xarray.DataArray
+        The track_ids corresponding to the tracks Dataset
+    track_id : Any
+        The track ID to match in track_ids. Must be the same type as the track_ids.
+        Usually `int` or `str`
 
     Returns
     -------
     xarray.Dataset
 
     """
-    df = data.to_dataframe()
-    track = df[track_ids == track_id]
-    return track.to_xarray()
+    if track_ids.ndim != 1:
+        raise ValueError("track_ids must be 1d")
+
+    dim = track_ids.dims[0]
+    idx = np.where(track_ids == track_id)[0]
+
+    return tracks.isel(**{dim: idx})
 
 
-def trackswhere(tracks, condition):
+def trackswhere(tracks, track_ids, condition):
     """Subset tracks from the input
 
     e.g select all tracks that are solely in the Northern hemisphere
@@ -31,7 +39,8 @@ def trackswhere(tracks, condition):
     Parameters
     ----------
     tracks : xarray.Dataset
-    condition : function
+    track_ids : xarray.DataArray
+    condition : callable
         A function that takes an `xarray.Dataset` of an individual track and returns
         True or False
 
@@ -41,6 +50,9 @@ def trackswhere(tracks, condition):
         A dataset with the subset of tracks from the input that match the given criteria
 
     """
+    if track_ids.ndim != 1:
+        raise ValueError("track_ids must be 1d")
+
     track_groups = tracks.groupby("track_id")
 
     if callable(condition):
@@ -50,7 +62,4 @@ def trackswhere(tracks, condition):
         track for n, (track_id, track) in enumerate(track_groups) if is_match[n]
     ]
 
-    if len(tracks.time.dims) == 1:
-        raise ValueError("trackswhere input must have exactly 1 time dimension")
-
-    return xr.concat(track_groups, dim=tracks.time.dims[0])
+    return xr.concat(track_groups, dim=track_ids.dims[0])

huracanpy/assess/_match.py

@@ -12,7 +12,7 @@ def match(tracksets, names=["1", "2"], max_dist=300, min_overlap=0):
 
     Parameters
     ----------
-    tracksets : list
+    tracksets : list[xarray.Dataset]
         list of track datasets to match together. Must be of length two or more.
     names : list, optional
         list of track datasets names. Must be the same size as tracksets. The default is ['1','2'].
@@ -23,7 +23,7 @@ def match(tracksets, names=["1", "2"], max_dist=300, min_overlap=0):
 
     Returns
     -------
-    pd.DataFrame
+    pandas.DataFrame
         Dataframe containing the matching tracks with
             the id from corresponding datasets
             the number of matching time steps (if only two datasets provided)

huracanpy/assess/_overlap.py

@@ -9,9 +9,9 @@ def overlap(tracks1, tracks2, matches=None):
 
     Parameters
     ----------
-    tracks1 (pd.Dataframe)
-    tracks2 (pd.Dataframe)
-    matches (pd.Dataframe): The output from match_tracks on tracks1 and tracks2.
+    tracks1 (xarray.Dataset)
+    tracks2 (xarray.Dataset)
+    matches (pandas.Dataframe): The output from match_tracks on tracks1 and tracks2.
         If None, match_tracks is run on tracks1 and tracks2.
 
     Returns
@@ -20,7 +20,7 @@ def overlap(tracks1, tracks2, matches=None):
         Match dataset with added deltas in days
     """
     if matches is None:
-        matches = match(tracks1, tracks2)
+        matches = match([tracks1, tracks2])
     c1, c2 = matches.columns[:2].str.slice(3)
     tracks1, tracks2 = tracks1.to_dataframe(), tracks2.to_dataframe()
     matches = (

huracanpy/calc/_lifecycle.py

@@ -29,7 +29,10 @@ def get_time_from_genesis(time, track_ids):
     )
     time_from_start = data_df.time_actual - data_df.time_gen
     return (
-        time_from_start.to_xarray().rename({"index": "obs"}).rename("time_from_genesis")
+        time_from_start.to_xarray()
+        .rename({"index": track_ids.dims[0]})
+        .drop(track_ids.dims[0])
+        .rename("time_from_genesis")
     )
 
 
@@ -40,7 +43,7 @@ def get_time_from_apex(time, track_ids, intensity_var, stat="max"):
     Parameters
     ----------
     time : array_like
-    track_ids : array_like
+    track_ids : xarray.DataArray
     intensity_var : array_like
     stat : str, optional
         Take either the maxima ("max") or minima ("min") of `intensity_var`. Default is
@@ -65,6 +68,7 @@ def get_time_from_apex(time, track_ids, intensity_var, stat="max"):
     time_from_extr = data_df.time_actual - data_df.time_extr
     return (
         time_from_extr.to_xarray()
-        .rename({"index": time.dims[0]})
+        .rename({"index": track_ids.dims[0]})
+        .drop(track_ids.dims[0])
         .rename("time_from_extremum")
     )

huracanpy/calc/_track_stats.py

@@ -60,7 +60,12 @@ def get_gen_vals(tracks, time, track_id):
     # 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})
+    tracks = tracks.isel(**{dim: idx})
+
+    # Promote track_id to a coordinate and remove record
+    return tracks.assign_coords(**{track_id.name: tracks[track_id.name]}).swap_dims(
+        **{dim: track_id.name}
+    )
 
 
 def get_apex_vals(tracks, variable, track_id, stat="max"):
@@ -109,4 +114,9 @@ def get_apex_vals(tracks, variable, track_id, stat="max"):
     idx = np.array(df.sort_values("var", ascending=asc).groupby("track_id").first().idx)
 
     dim = track_id.dims[0]
-    return tracks.isel(**{dim: idx})
+    tracks = tracks.isel(**{dim: idx})
+
+    # Promote track_id to a coordinate and remove record
+    return tracks.assign_coords(**{track_id.name: tracks[track_id.name]}).swap_dims(
+        **{dim: track_id.name}
+    )

tests/test_accel.py

@@ -0,0 +1,140 @@
+"""
+Test functions that use tricks to speed up their code produce the same result as the
+slower method
+"""
+
+from haversine import haversine_vector, Unit
+import numpy as np
+import xarray as xr
+
+import huracanpy
+
+
+def test_accel_sel_id(tracks_csv):
+    result = huracanpy.sel_id(tracks_csv, tracks_csv.track_id, 0)
+
+    expected = tracks_csv.groupby("track_id")[0]
+
+    xr.testing.assert_identical(result, expected)
+
+
+def test_accel_trackswhere():
+    # TODO accelerate trackswhere
+    pass
+
+
+def test_accel_get_gen_vals(tracks_csv):
+    result = huracanpy.calc.get_gen_vals(
+        tracks_csv, tracks_csv.time, tracks_csv.track_id
+    )
+
+    expected = tracks_csv.groupby("track_id").first()
+
+    xr.testing.assert_identical(result, expected)
+
+
+def test_accel_get_apex_vals(tracks_csv):
+    result = huracanpy.calc.get_apex_vals(
+        tracks_csv, tracks_csv.wind10, tracks_csv.track_id
+    )
+
+    expected = tracks_csv.sortby("wind10", ascending=False).groupby("track_id").first()
+
+    xr.testing.assert_identical(result, expected)
+
+
+def test_accel_get_time_from_genesis(tracks_csv):
+    result = huracanpy.calc.get_time_from_genesis(tracks_csv.time, tracks_csv.track_id)
+
+    track_groups = tracks_csv.groupby("track_id")
+    expected = []
+    for track_id, track in track_groups:
+        expected.append(track.time - track.time[0])
+
+    expected = xr.concat(expected, dim="record")
+    expected = expected.rename("time_from_genesis")
+
+    xr.testing.assert_identical(result, expected)
+
+
+def test_accel_get_time_from_apex(tracks_csv):
+    result = huracanpy.calc.get_time_from_apex(
+        tracks_csv.time, tracks_csv.track_id, tracks_csv.wind10
+    )
+
+    track_groups = tracks_csv.groupby("track_id")
+    expected = []
+    for track_id, track in track_groups:
+        idx = track.wind10.argmax()
+        expected.append(track.time - track.time[idx])
+
+    expected = xr.concat(expected, dim="record")
+    expected = expected.rename("time_from_extremum")
+
+    xr.testing.assert_identical(result, expected)
+
+
+def test_accel_match():
+    ref = huracanpy.load(huracanpy.example_csv_file)
+    tracks = ref.where(ref.track_id < 2, drop=True)
+    tracks = tracks.where(tracks.time.dt.hour == 0, drop=True)
+    tracks["lon"] = tracks.lon + 0.5
+    tracks["lat"] = tracks.lat + 0.5
+
+    result = huracanpy.assess.match([tracks, ref])
+
+    max_dist = 300
+    track_id1 = []
+    track_id2 = []
+    npoints = []
+    dist = []
+
+    for track_id, track in tracks.groupby("track_id"):
+        for track_id_ref, track_ref in ref.groupby("track_id"):
+            # Match times
+            track_ = track.where(track.time.isin(track_ref.time), drop=True)
+
+            if len(track_.time) > 0:
+                track_ref_ = track_ref.where(track_ref.time.isin(track.time), drop=True)
+
+                yx_track = np.array([track_.lat, track_.lon]).T
+                yx_ref = np.array([track_ref_.lat, track_ref_.lon]).T
+
+                dists = haversine_vector(yx_track, yx_ref, Unit.KILOMETERS)
+
+                matches = dists < max_dist
+                if matches.any():
+                    track_id1.append(track_id)
+                    track_id2.append(track_id_ref)
+
+                    dists_track = dists[matches]
+                    npoints.append(len(dists_track))
+                    dist.append(np.mean(dists_track))
+
+    np.testing.assert_equal(result.id_1, np.array(track_id1))
+    np.testing.assert_equal(result.id_2, np.array(track_id2))
+    np.testing.assert_equal(result.temp, np.array(npoints))
+    np.testing.assert_allclose(result.dist, np.array(dist), rtol=1e-12)
+
+
+def test_accel_overlap():
+    ref = huracanpy.load(huracanpy.example_csv_file)
+    tracks = ref.where(ref.track_id < 2, drop=True)
+    tracks = tracks.where(tracks.time.dt.hour == 0, drop=True)
+    tracks["lon"] = tracks.lon + 0.5
+    tracks["lat"] = tracks.lat + 0.5
+
+    result = huracanpy.assess.overlap(tracks, ref)
+
+    delta_start = []
+    delta_end = []
+
+    for n, row in result.iterrows():
+        track = tracks.where(tracks.track_id == row.id_1, drop=True)
+        track_ref = ref.where(ref.track_id == row.id_2, drop=True)
+
+        delta_start.append((track_ref.time[0] - track.time[0]) / np.timedelta64(1, "D"))
+        delta_end.append((track_ref.time[-1] - track.time[-1]) / np.timedelta64(1, "D"))
+
+    np.testing.assert_equal(result.delta_start, np.array(delta_start))
+    np.testing.assert_equal(result.delta_end, np.array(delta_end))

tests/test_subset.py

@@ -7,7 +7,7 @@ def test_trackswhere():
     tracks["category"] = huracanpy.tc.get_pressure_cat(tracks.slp, slp_units="Pa")
 
     tracks_subset = huracanpy.trackswhere(
-        tracks, lambda track: track.category.max() >= 2
+        tracks, tracks.track_id, lambda track: track.category.max() >= 2
     )
 
     assert set(tracks_subset.track_id.data) == {0, 2}