Merge pull request #746 from martinfleis/categorical-lag-fix

BUG: fix categorical lag for custom index

libpysal/graph/_spatial_lag.py

@@ -91,9 +91,12 @@ def _lag_spatial(graph, y, categorical=False, ties="raise"):
     ):
         categorical = True
     if categorical:
+        if isinstance(y, np.ndarray):
+            y = pd.Series(y, index=graph.unique_ids)
+
         df = pd.DataFrame(data=graph.adjacency)
-        df["neighbor_label"] = y[graph.adjacency.index.get_level_values(1)]
-        df["own_label"] = y[graph.adjacency.index.get_level_values(0)]
+        df["neighbor_label"] = y.loc[graph.adjacency.index.get_level_values(1)].values
+        df["own_label"] = y.loc[graph.adjacency.index.get_level_values(0)].values
         df["neighbor_idx"] = df.index.get_level_values(1)
         df["focal_idx"] = df.index.get_level_values(0)
         gb = df.groupby(["focal", "neighbor_label"]).count().groupby(level="focal")

libpysal/graph/tests/test_spatial_lag.py

@@ -65,3 +65,9 @@ def test_ties_raise(self):
         with pytest.raises(ValueError, match="There are 2 ties that must be broken"):
             self.yc[3] = "a"  # create ties
             _lag_spatial(self.gc, self.yc, categorical=True)
+
+    def test_categorical_custom_index(self):
+        expected = np.array(["bar", "foo", "bar", "foo"])
+        np.testing.assert_array_equal(
+            expected, self.g.lag(["foo", "bar", "foo", "foo"])
+        )