improved handling of boolean variables and dedicated unit test for it

quaquel · quaquel · commit 87e6134f0d09 · 2025-09-22T16:14:13.000+02:00
diff --git a/ema_workbench/analysis/scenario_discovery_util.py b/ema_workbench/analysis/scenario_discovery_util.py
@@ -221,26 +221,24 @@ def _in_box(x, boxlim):
     category dtype
 
     """
-    x_numbered = x.select_dtypes(np.number)
-    boxlim_numbered = boxlim.select_dtypes(np.number)
+    number_like = [np.number, np.bool]
+    x_numbered = x.select_dtypes(number_like)
+    boxlim_numbered = boxlim.select_dtypes(number_like)
     logical = (boxlim_numbered.loc[0, :].values <= x_numbered.values) & (
         x_numbered.values <= boxlim_numbered.loc[1, :].values
     )
     logical = logical.all(axis=1)
 
     # TODO:: how to speed this up
-    for column, values in x.select_dtypes(exclude=np.number).items():
+
+    for column, values in x.select_dtypes(exclude=number_like).items():
         entries = boxlim.loc[0, column]
-        if values.dtype == np.dtype(np.bool):
-            l = x[column] == entries
-            logical = logical & l
-        else:
-            not_present  = set(values.cat.categories.values) - entries
+        not_present  = set(values.cat.categories.values) - entries
 
-            if not_present:
-                # what other options do we have here....
-                l = pd.isnull(x[column].cat.remove_categories(list(entries)))  # noqa: E741
-                logical = l & logical
+        if not_present:
+            # what other options do we have here....
+            l = pd.isnull(x[column].cat.remove_categories(list(entries)))  # noqa: E741
+            logical = l & logical
     return logical
 
 
diff --git a/ema_workbench/examples/sd_cart_wcm.py b/ema_workbench/examples/sd_cart_wcm.py
@@ -19,6 +19,8 @@
 results = load_results(fn)
 x, outcomes = results
 
+x = x.drop(["scenario", "model", "policy"], axis=1)
+
 ooi = "throughput_Rotterdam"
 outcome = outcomes[ooi] / default_flow
 y = outcome < 1
diff --git a/test/test_analysis/test_scenario_discovery_util.py b/test/test_analysis/test_scenario_discovery_util.py
@@ -89,6 +89,32 @@ def test_in_box(self):
         result = x.loc[logical]
         self.assertTrue(np.all(correct_result == result))
 
+        x = pd.DataFrame(
+            [
+                (0.1, 0, "a", True),
+                (1.1, 1, "a", True),
+                (2.1, 2, "b", True),
+                (3.1, 3, "b", True),
+                (4.1, 4, "c", False),
+                (5.1, 5, "c", False),
+                (6.1, 6, "d", False),
+                (7.1, 7, "d", False),
+                (8.1, 8, "e", False),
+                (9.1, 9, "e", False),
+            ],
+            columns=["a", "b", "c", "d"],
+        )
+        boxlim = pd.DataFrame(
+            [(1.2, 0, {"a", "b", "c"}, True), (8.0, 7, {"a", "b", "c"}, True)], columns=["a", "b", "c", "d"]
+        )
+        x["c"] = x["c"].astype("category")
+
+        correct_result = x.loc[[2, 3], :]
+        logical = sdutil._in_box(x, boxlim)
+        result = x.loc[logical]
+        self.assertTrue(np.all(correct_result == result))
+
+
     def test_make_box(self):
         x = pd.DataFrame([(0, 1, 2), (2, 5, 6), (3, 2, 1)], columns=["a", "b", "c"])
 
