add navdata.find_wildcard_indexes

betaBison · betaBison · commit 171f8db306da · 2022-08-31T22:05:25.000-07:00
diff --git a/gnss_lib_py/parsers/navdata.py b/gnss_lib_py/parsers/navdata.py
@@ -575,6 +575,63 @@ def in_rows(self, rows):
             raise KeyError(", ".join(missing_rows) + " row(s) are" \
                            + " missing from NavData object.")
 
+    def find_wildcard_indexes(self, wildcards, max_allow = None):
+        """Searches for indexes matching wildcard search input.
+
+        For example, a search for ``x_*_m`` would find ``x_rx_m`` or
+        ``x_sv_m`` or ``x_alpha_beta_gamma_m`` depending on the rows
+        existing in the NavData instance.
+
+        Will return an error no index is found matching the wildcard or
+        if more than ``max_allow`` indexes are found.
+
+        Currently only allows for a single wildcard per index.
+
+        Parameters
+        ----------
+        wildcards : array-like or str
+            List/tuple/np.ndarray/set of indexes for which to search.
+        max_allow : int or None
+            Maximum number of valid indexes to allow before throwing an
+            error. If None, then no limit is placed.
+
+        Returns
+        -------
+        wildcard_indexes : dict
+            Dictionary of the form {"search_term", [indexes,...]},
+
+        """
+
+        if isinstance(wildcards,str):
+            wildcards = [wildcards]
+        if not isinstance(wildcards, (list,tuple,np.ndarray,set)):
+            raise TypeError("wildcards input in find_wildcard_indexes" \
+                         +  " must be array-like or single string")
+        if not (isinstance(max_allow,int) or max_allow is None):
+            raise TypeError("max_allow input in find_wildcard_indexes" \
+                          + " must be an integer or None.")
+
+        wildcard_indexes = {}
+
+        for wildcard in wildcards:
+            if not isinstance(wildcard,str):
+                raise TypeError("wildcards must be strings")
+            if wildcard.count("*") != 1:
+                raise RuntimeError("One wildcard '*' and only one "\
+                          + "wildcard must be present in search string")
+            indexes = [row for row in self.rows
+                   if row.startswith(wildcard.split("*",maxsplit=1)[0])
+                    and row.endswith(wildcard.split("*",maxsplit=1)[1])]
+            if max_allow is not None and len(indexes) > max_allow:
+                raise KeyError("More than " + str(max_allow) \
+                             + " possible row indexes for "  + wildcard)
+            if len(indexes) == 0:
+                raise KeyError("Missing " + wildcard + " row.")
+
+            wildcard_indexes[wildcard] = indexes
+
+        return wildcard_indexes
+
     def pandas_df(self):
         """Return pandas DataFrame equivalent to class
 
diff --git a/gnss_lib_py/utils/visualizations.py b/gnss_lib_py/utils/visualizations.py
@@ -284,23 +284,8 @@ def plot_skyplot(navdata, receiver_state, save=False, prefix="",
     receiver_state.in_rows(["gps_millis"])
 
     # check for receiver_state indexes
-    rx_idxs = {"x_*_m" : [],
-               "y_*_m" : [],
-               "z_*_m" : [],
-               }
-    for name, indexes in rx_idxs.items():
-        indexes = [row for row in receiver_state.rows
-                      if row.startswith(name.split("*",maxsplit=1)[0])
-                       and row.endswith(name.split("*",maxsplit=1)[1])]
-        if len(indexes) > 1:
-            raise KeyError("Multiple possible row indexes for " \
-                         + name \
-                         + ". Unable to resolve for plot_skyplot().")
-        if len(indexes) == 0:
-            raise KeyError("Missing required " + name + " row for " \
-                        + "plot_skyplot().")
-        # must call dictionary to avoid pass by value
-        rx_idxs[name] = indexes[0]
+    rx_idxs = receiver_state.find_wildcard_indexes(["x_*_m","y_*_m",
+                                                    "z_*_m"],max_allow=1)
 
     if "el_sv_deg" not in navdata.rows or "az_sv_deg" not in navdata.rows:
         sv_el_az = None
@@ -312,9 +297,9 @@ def plot_skyplot(navdata, receiver_state, save=False, prefix="",
             # find time index for receiver_state NavData instance
             rx_t_idx = np.argmin(np.abs(receiver_state["gps_millis"] - timestamp))
 
-            pos_rx_m = receiver_state[[rx_idxs["x_*_m"],
-                                       rx_idxs["y_*_m"],
-                                       rx_idxs["z_*_m"]],
+            pos_rx_m = receiver_state[[rx_idxs["x_*_m"][0],
+                                       rx_idxs["y_*_m"][0],
+                                       rx_idxs["z_*_m"][0]],
                                        rx_t_idx].reshape(1,-1)
 
             timestep_el_az = ecef_to_el_az(pos_rx_m, pos_sv_m)
@@ -457,26 +442,14 @@ def plot_map(*args, sections=0, save=False, prefix="",
                           + "NavData.")
 
         # check for lat/lon indexes
-        traj_idxs = {}
-        traj_idxs_possibilities = ["lat_*_deg", "lon_*_deg"]
-        for name in traj_idxs_possibilities:
-            indexes = [row for row in traj_data.rows
-                          if row.startswith(name.split("*",maxsplit=1)[0])
-                           and row.endswith(name.split("*",maxsplit=1)[1])]
-            if len(indexes) > 1:
-                raise KeyError("Multiple possible row indexes for " \
-                             + name \
-                             + ". Unable to resolve for plot_map().")
-            if len(indexes) == 0:
-                raise KeyError("Missing required " + name + " row for " \
-                            + "plot_map().")
-            # must call dictionary to avoid pass by value
-            traj_idxs[name] = indexes[0]
-
-        label_name = _get_label({"":"_".join((traj_idxs["lat_*_deg"].split("_"))[1:-1])})
-
-        data = {"latitude" : traj_data[traj_idxs["lat_*_deg"]],
-                "longitude" : traj_data[traj_idxs["lon_*_deg"]],
+        traj_idxs = traj_data.find_wildcard_indexes(["lat_*_deg",
+                                                     "lon_*_deg"],
+                                                     max_allow=1)
+
+        label_name = _get_label({"":"_".join((traj_idxs["lat_*_deg"][0].split("_"))[1:-1])})
+
+        data = {"latitude" : traj_data[traj_idxs["lat_*_deg"][0]],
+                "longitude" : traj_data[traj_idxs["lon_*_deg"][0]],
                 "Trajectory" : [label_name] * len(traj_data),
                 }
         if sections >= 2:
diff --git a/tests/parsers/test_navdata.py b/tests/parsers/test_navdata.py
@@ -1877,3 +1877,91 @@ def test_large_int():
     navdata["numbers"] = test_list
 
     np.testing.assert_array_equal(navdata["numbers"], test_list)
+
+def test_find_wildcard_indexes(data):
+    """Tests find_wildcard_indexes
+
+    """
+
+    all_matching = data.rename({"names" : "x_alpha_m",
+                                "integers" : "x_beta_m",
+                                "floats" : "x_gamma_m",
+                                "strings" : "x_zeta_m"})
+    expected = ["x_alpha_m","x_beta_m","x_gamma_m","x_zeta_m"]
+
+    indexes = all_matching.find_wildcard_indexes("x_*_m")
+    assert indexes["x_*_m"] == expected
+    expect_pass_allows = [None,12,4]
+    for max_allow in expect_pass_allows:
+        indexes = all_matching.find_wildcard_indexes("x_*_m",max_allow)
+        assert indexes["x_*_m"] == expected
+
+    expect_fail_allows = [0,-1,3,2,1]
+    for max_allow in expect_fail_allows:
+        with pytest.raises(KeyError) as excinfo:
+            all_matching.find_wildcard_indexes("x_*_m",max_allow)
+        assert "More than " + str(max_allow) in str(excinfo.value)
+        assert "x_*_m" in str(excinfo.value)
+
+    multi = data.rename({"names" : "x_alpha_m",
+                         "integers" : "x_beta_m",
+                         "floats" : "y_alpha_deg",
+                         "strings" : "x_zeta_deg"})
+    expected = {"x_*_m" : ["x_alpha_m","x_beta_m"],
+                "y_*_deg" : ["y_alpha_deg"]}
+
+    expect_pass_allows = [None,2,4]
+    for max_allow in expect_pass_allows:
+        indexes = multi.find_wildcard_indexes(["x_*_m","y_*_deg"],
+                                               max_allow)
+        assert indexes == expected
+
+    expect_pass_allows = [None,2,4]
+    for max_allow in expect_pass_allows:
+        indexes = multi.find_wildcard_indexes(tuple(["x_*_m","y_*_deg"]),
+                                                     max_allow)
+        assert indexes == expected
+
+    expect_pass_allows = [None,2,4]
+    for max_allow in expect_pass_allows:
+        indexes = multi.find_wildcard_indexes(set(["x_*_m","y_*_deg"]),
+                                                     max_allow)
+        assert indexes == expected
+
+    expect_pass_allows = [None,2,4]
+    for max_allow in expect_pass_allows:
+        indexes = multi.find_wildcard_indexes(np.array(["x_*_m",
+                                                        "y_*_deg"]),
+                                                        max_allow)
+        assert indexes == expected
+
+    expect_fail_allows = [0,-1,1]
+    for max_allow in expect_fail_allows:
+        with pytest.raises(KeyError) as excinfo:
+            multi.find_wildcard_indexes(["x_*_m","y_*_deg"],max_allow)
+        assert "More than " + str(max_allow) in str(excinfo.value)
+        assert "x_*_m" in str(excinfo.value)
+
+    with pytest.raises(KeyError) as excinfo:
+        multi.find_wildcard_indexes(["z_*_m"])
+    assert "Missing " in str(excinfo.value)
+    assert "z_*_m" in str(excinfo.value)
+
+    with pytest.raises(TypeError) as excinfo:
+        multi.find_wildcard_indexes(1.0)
+    assert "find_wildcard_indexes " in str(excinfo.value)
+    assert "array-like" in str(excinfo.value)
+
+    with pytest.raises(TypeError) as excinfo:
+        multi.find_wildcard_indexes([1.0])
+    assert "wildcards must be strings" in str(excinfo.value)
+
+    with pytest.raises(RuntimeError) as excinfo:
+        multi.find_wildcard_indexes("x_*_*")
+    assert "One wildcard" in str(excinfo.value)
+
+    incorrect_max_allow = [3.,"hi",[]]
+    for max_allow in incorrect_max_allow:
+        with pytest.raises(TypeError) as excinfo:
+            multi.find_wildcard_indexes("x_*_m",max_allow)
+        assert "max_allow" in str(excinfo.value)
diff --git a/tests/utils/test_visualizations.py b/tests/utils/test_visualizations.py
@@ -395,7 +395,7 @@ def test_plot_skyplot(navdata, state_estimate):
         with pytest.raises(KeyError) as excinfo:
             viz.plot_skyplot(navdata, state_double, save=False)
         assert row.replace("rx","*") in str(excinfo.value)
-        assert "Multiple" in str(excinfo.value)
+        assert "More than 1" in str(excinfo.value)
 
 def test_get_label():
     """Test for getting nice labels.
@@ -493,4 +493,4 @@ def test_plot_map(gtruth, state_estimate):
         with pytest.raises(KeyError) as excinfo:
             viz.plot_map(gtruth, state_double, save=False)
         assert row.replace("rx","*") in str(excinfo.value)
-        assert "Multiple" in str(excinfo.value)
+        assert "More than 1" in str(excinfo.value)
