Adding Comparison and Indexing methods to DataFrameMixin

mesa_frames/abstract/mixin.py

@@ -152,6 +152,15 @@ class DataFrameMixin(ABC):
     def _df_remove(self, df: DataFrame, mask: Mask, index_cols: str) -> DataFrame:
         return self._df_get_masked_df(df, index_cols, mask, negate=True)
 
+    @abstractmethod
+    def _df_and(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
     @abstractmethod
     def _df_add(
         self,
@@ -258,6 +267,15 @@ def _df_drop_duplicates(
         keep: Literal["first", "last", False] = "first",
     ) -> DataFrame: ...
 
+    @abstractmethod
+    def _df_ge(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
     @abstractmethod
     def _df_get_bool_mask(
         self,
@@ -282,6 +300,9 @@ def _df_groupby_cumcount(
         self, df: DataFrame, by: str | list[str], name: str = "cum_count"
     ) -> Series: ...
 
+    @abstractmethod
+    def _df_index(self, df: DataFrame, index_name: str) -> Index: ...
+
     @abstractmethod
     def _df_iterator(self, df: DataFrame) -> Iterator[dict[str, Any]]: ...
 
@@ -302,6 +323,24 @@ def _df_join(
         suffix="_right",
     ) -> DataFrame: ...
 
+    @abstractmethod
+    def _df_lt(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_mod(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
     @abstractmethod
     def _df_mul(
         self,
@@ -337,6 +376,23 @@ def _df_norm(
         include_cols: bool = False,
     ) -> Series | DataFrame: ...
 
+    @abstractmethod
+    def _df_or(
+        self,
+        df: DataFrame,
+        other: DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> DataFrame: ...
+
+    @abstractmethod
+    def _df_reindex(
+        self,
+        df: DataFrame,
+        other: Sequence[Hashable] | DataFrame,
+        index_cols: str | list[str],
+    ) -> DataFrame: ...
+
     @abstractmethod
     def _df_rename_columns(
         self,

mesa_frames/abstract/space.py

@@ -1341,13 +1341,15 @@ def get_neighborhood(
             neighbors_df = self._df_drop_duplicates(neighbors_df, self._pos_col_names)
 
         # Filter out-of-bound neighbors
-        neighbors_df = self._df_get_masked_df(
-            neighbors_df,
-            mask=self._df_all(
-                (neighbors_df[self._pos_col_names] < self._dimensions)
-                & (neighbors_df >= 0)
-            ),
+        mask = self._df_all(
+            self._df_and(
+                self._df_lt(
+                    neighbors_df[self._pos_col_names], self._dimensions, axis="columns"
+                ),
+                neighbors_df >= 0,
+            )
         )
+        neighbors_df = self._df_get_masked_df(neighbors_df, mask=mask)
 
         if include_center:
             center_df = self._df_rename_columns(
@@ -1409,7 +1411,15 @@ def out_of_bounds(self, pos: GridCoordinate | GridCoordinates) -> DataFrame:
             raise ValueError("This method is only valid for non-torus grids")
         pos_df = self._get_df_coords(pos, check_bounds=False)
         out_of_bounds = self._df_all(
-            (pos_df < 0) | (pos_df >= self._dimensions),
+            self._df_or(
+                pos_df < 0,
+                self._df_ge(
+                    pos_df,
+                    self._dimensions,
+                    axis="columns",
+                    index_cols=self._pos_col_names,
+                ),
+            ),
             name="out_of_bounds",
         )
         return self._df_concat(objs=[pos_df, out_of_bounds], how="horizontal")
@@ -1452,7 +1462,7 @@ def torus_adj(self, pos: GridCoordinate | GridCoordinates) -> DataFrame:
             The adjusted coordinates
         """
         df_coords = self._get_df_coords(pos)
-        df_coords = df_coords % self._dimensions
+        df_coords = self._df_mod(df_coords, self._dimensions, axis="columns")
         return df_coords
 
     def _calculate_differences(
@@ -1603,13 +1613,11 @@ def _get_df_coords(
                 if agents.n_unique() != len(agents):
                     raise ValueError("Some agents are present multiple times")
         if agents is not None:
-            return self._df_reset_index(
-                self._df_get_masked_df(
-                    self._agents, index_cols="agent_id", mask=agents
-                ),
-                index_cols="agent_id",
-                drop=True,
+            df = self._df_get_masked_df(
+                self._agents, index_cols="agent_id", mask=agents
             )
+            df = self._df_reindex(df, agents, "agent_id")
+            return self._df_reset_index(df, index_cols="agent_id", drop=True)
         if isinstance(pos, DataFrame):
             return pos[self._pos_col_names]
         elif (

mesa_frames/concrete/pandas/mixin.py

@@ -1,4 +1,4 @@
-from collections.abc import Collection, Hashable, Iterator, Sequence
+from collections.abc import Callable, Collection, Hashable, Iterator, Sequence
 from typing import Literal
 
 import numpy as np
@@ -19,6 +19,21 @@
     ) -> pd.DataFrame:
         return df.add(other=other, axis=axis)
 
+    def _df_and(
+        self,
+        df: pd.DataFrame,
+        other: pd.DataFrame | Sequence[float | int],
+        axis: Literal["index"] | Literal["columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> pd.DataFrame:
+        return self._df_logical_operation(
+            df=df,
+            other=other,
+            operation=lambda x, y: x & y,
+            axis=axis,
+            index_cols=index_cols,
+        )
+
     def _df_all(
         self,
         df: pd.DataFrame,
@@ -137,6 +152,15 @@
     ) -> pd.DataFrame:
         return df.drop_duplicates(subset=subset, keep=keep)
 
+    def _df_ge(
+        self,
+        df: pd.DataFrame,
+        other: pd.DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> pd.DataFrame:
+        return df.ge(other, axis=axis)
+
     def _df_get_bool_mask(
         self,
         df: pd.DataFrame,
@@ -192,6 +216,16 @@
     ) -> pd.Series:
         return df.groupby(by).cumcount().rename(name) + 1
 
+    def _df_index(self, df: pd.DataFrame, index_col: str | list[str]) -> pd.Index:
+        if (
+            index_col is None
+            or df.index.name == index_col
+            or df.index.names == index_col
+        ):
+            return df.index
+        else:
+            return df.set_index(index_col).index
+
     def _df_iterator(self, df: pd.DataFrame) -> Iterator[dict[str, Any]]:
         for index, row in df.iterrows():
             row_dict = row.to_dict()
@@ -246,6 +280,50 @@
         else:
             return df
 
+    def _df_lt(
+        self,
+        df: pd.DataFrame,
+        other: pd.DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> pd.DataFrame:
+        return df.lt(other, axis=axis)
+
+    def _df_logical_operation(
+        self,
+        df: pd.DataFrame,
+        other: pd.DataFrame | Sequence[bool],
+        operation: Callable[
+            [pd.DataFrame, Sequence[bool] | pd.DataFrame], pd.DataFrame
+        ],
+        axis: Literal["index"] | Literal["columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> pd.DataFrame:
+        if isinstance(other, pd.DataFrame):
+            if index_cols is not None:
+                if df.index.name != index_cols:
+                    df = df.set_index(index_cols)
+                if other.index.name != index_cols:
+                    other = other.set_index(index_cols)
+            other = other.reindex(df.index, fill_value=np.nan)
+            return operation(df, other)
+        else:  # Sequence[bool]
+            other = pd.Series(other)
+            if axis == "index":
+                other.index = df.index
+                return operation(df, other.values[:, None]).astype(bool)
+            else:
+                return operation(df, other.values[None, :]).astype(bool)
+
+    def _df_mod(
+        self,
+        df: pd.DataFrame,
+        other: pd.DataFrame | Sequence[float | int],
+        axis: Literal["index", "columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> pd.DataFrame:
+        return df.mod(other, axis=axis)
+
     def _df_mul(
         self,
         df: pd.DataFrame,
@@ -285,6 +363,31 @@
         else:
             return srs
 
+    def _df_or(
+        self,
+        df: pd.DataFrame,
+        other: pd.DataFrame | Sequence[bool],
+        axis: Literal["index"] | Literal["columns"] = "index",
+        index_cols: str | list[str] | None = None,
+    ) -> pd.DataFrame:
+        return self._df_logical_operation(
+            df=df,
+            other=other,
+            operation=lambda x, y: x | y,
+            axis=axis,
+            index_cols=index_cols,
+        )
+
+    def _df_reindex(
+        self,
+        df: pd.DataFrame,
+        other: Sequence[Hashable] | pd.DataFrame,
+        index_cols: str | list[str],
+    ) -> pd.DataFrame:
+        df = df.reindex(other)
+        df.index.name = index_cols
+        return df
+
     def _df_rename_columns(
         self,
         df: pd.DataFrame,