Improve filterby

benchmarks/hypernetx.ipynb

@@ -15,10 +15,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import time\n",
+    "\n",
     "import hypernetx as hnx\n",
-    "import xgi\n",
     "import pandas as pd\n",
-    "import time"
+    "\n",
+    "import xgi"
    ]
   },
   {

benchmarks/networkx.ipynb

@@ -15,9 +15,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import time\n",
+    "\n",
     "import networkx as nx\n",
-    "import xgi\n",
-    "import time"
+    "\n",
+    "import xgi"
    ]
   },
   {

tests/core/test_diviews.py

@@ -221,3 +221,15 @@ def test_isolates():
     assert set(DH.nodes.isolates()) == {0, 1, 2, 3}
     DH.add_edge([{0}, {1, 2}])
     assert set(DH.nodes.isolates()) == {3}
+
+
+def test_diview_custom_filterby(diedgelist2):
+    H = xgi.DiHypergraph(diedgelist2)
+
+    f = lambda val, arg: val % arg == 0
+    assert set(H.edges.filterby("tail_size", 2, mode=f)) == {0, 1}
+
+
+def test_diview_custom_filterby_attr(dihyperwithattrs):
+    f = lambda val, arg: arg in val
+    assert set(dihyperwithattrs.nodes.filterby_attr("color", "l", mode=f)) == {2, 3, 5}

tests/core/test_views.py

@@ -346,3 +346,15 @@ def test_ids_are_immutable(edgelist5):
     H = xgi.Hypergraph(edgelist5)
     H.edges.ids.add(42)
     assert H.edges.ids == {0, 1, 2, 3}
+
+
+def test_view_custom_filterby(edgelist8):
+    H = xgi.Hypergraph(edgelist8)
+
+    f = lambda val, arg: val <= arg**2
+    assert set(H.nodes.filterby("degree", 2, mode=f)) == {2, 3, 4, 5, 6}
+
+
+def test_view_custom_filterby_attr(hyperwithattrs):
+    f = lambda val, arg: arg in val
+    assert set(hyperwithattrs.nodes.filterby_attr("color", "l", mode=f)) == {2, 3, 5}

xgi/core/diviews.py

@@ -190,7 +190,7 @@ def filterby(self, stat, val, mode="eq"):
         val : Any
             Value of the statistic.  Usually a single numeric value.  When mode is
             'between', must be a tuple of exactly two values.
-        mode : str, optional
+        mode : str or function, optional
             How to compare each value to `val`.  Can be one of the following.
 
             * 'eq' (default): Return IDs whose value is exactly equal to `val`.
@@ -201,6 +201,7 @@ def filterby(self, stat, val, mode="eq"):
             * 'geq': Return IDs whose value is greater than or equal to `val`.
             * 'between': In this mode, `val` must be a tuple `(val1, val2)`.  Return IDs
               whose value `v` satisfies `val1 <= v <= val2`.
+            * function, must be able to call `mode(statistic, val)` and have it map to a bool.
 
         See Also
         --------
@@ -256,6 +257,8 @@ def filterby(self, stat, val, mode="eq"):
             bunch = [idx for idx in self if values[idx] >= val]
         elif mode == "between":
             bunch = [node for node in self if val[0] <= values[node] <= val[1]]
+        elif callable(mode):
+            bunch = [idx for idx in self if mode(values[idx], val)]
         else:
             raise ValueError(
                 f"Unrecognized mode {mode}. mode must be one of 'eq', 'neq', 'lt', 'gt', 'leq', 'geq', or 'between'."
@@ -271,9 +274,10 @@ def filterby_attr(self, attr, val, mode="eq", missing=None):
             The name of the attribute
         val : Any
             A single value or, in the case of 'between', a list of length 2
-        mode : str, optional
+        mode : str or function, optional
             Comparison mode. Valid options are 'eq' (default), 'neq', 'lt', 'gt',
-            'leq', 'geq', or 'between'.
+            'leq', 'geq', or 'between'. If a function, must be able to call
+            `mode(attribute, val)` and have it map to a bool.
         missing : Any, optional
             The default value if the attribute is missing. If None (default),
             ignores those IDs.
@@ -323,9 +327,15 @@ def filterby_attr(self, attr, val, mode="eq", missing=None):
                 for idx in self
                 if values[idx] is not None and val[0] <= values[idx] <= val[1]
             ]
+        elif callable(mode):
+            bunch = [
+                idx
+                for idx in self
+                if values[idx] is not None and mode(values[idx], val)
+            ]
         else:
             raise ValueError(
-                f"Unrecognized mode {mode}. mode must be one of 'eq', 'neq', 'lt', 'gt', 'leq', 'geq', or 'between'."
+                f"Unrecognized mode {mode}. mode must be one of 'eq', 'neq', 'lt', 'gt', 'leq', 'geq', 'between', or a callable function."
             )
         return type(self).from_view(self, bunch)
 

xgi/core/views.py

@@ -183,7 +183,7 @@ def filterby(self, stat, val, mode="eq"):
         val : Any
             Value of the statistic.  Usually a single numeric value.  When mode is
             'between', must be a tuple of exactly two values.
-        mode : str, optional
+        mode : str or function, optional
             How to compare each value to `val`.  Can be one of the following.
 
             * 'eq' (default): Return IDs whose value is exactly equal to `val`.
@@ -194,6 +194,7 @@ def filterby(self, stat, val, mode="eq"):
             * 'geq': Return IDs whose value is greater than or equal to `val`.
             * 'between': In this mode, `val` must be a tuple `(val1, val2)`.  Return IDs
               whose value `v` satisfies `val1 <= v <= val2`.
+            * function, must be able to call `mode(statistic, val)` and have it map to a bool.
 
         See Also
         --------
@@ -255,6 +256,8 @@ def filterby(self, stat, val, mode="eq"):
             bunch = [idx for idx in self if values[idx] >= val]
         elif mode == "between":
             bunch = [node for node in self if val[0] <= values[node] <= val[1]]
+        elif callable(mode):
+            bunch = [idx for idx in self if mode(values[idx], val)]
         else:
             raise ValueError(
                 f"Unrecognized mode {mode}. mode must be one of "
@@ -271,9 +274,10 @@ def filterby_attr(self, attr, val, mode="eq", missing=None):
             The name of the attribute
         val : Any
             A single value or, in the case of 'between', a list of length 2
-        mode : str, optional
+        mode : str or function, optional
             Comparison mode. Valid options are 'eq' (default), 'neq', 'lt', 'gt',
-            'leq', 'geq', or 'between'.
+            'leq', 'geq', or 'between'. If a function, must be able to call
+            `mode(attribute, val)` and have it map to a bool.
         missing : Any, optional
             The default value if the attribute is missing. If None (default),
             ignores those IDs.
@@ -323,6 +327,12 @@ def filterby_attr(self, attr, val, mode="eq", missing=None):
                 for idx in self
                 if values[idx] is not None and val[0] <= values[idx] <= val[1]
             ]
+        elif callable(mode):
+            bunch = [
+                idx
+                for idx in self
+                if values[idx] is not None and mode(values[idx], val)
+            ]
         else:
             raise ValueError(
                 f"Unrecognized mode {mode}. mode must be one of "