FIX: boolean subset with incompatible axes raises an error (closes #1085)

gdementen · gdementen · commit 4c0bab360b68 · 2024-07-16T15:28:05.000+02:00
diff --git a/doc/source/changes/version_0_34_3.rst.inc b/doc/source/changes/version_0_34_3.rst.inc
@@ -55,7 +55,8 @@ Miscellaneous improvements
 Fixes
 ^^^^^
 
-* fixed something (closes :issue:`1`).
+* using a boolean array as a filter to take a subset of another array now raise an error when the
+  two arrays have incompatible axes instead of producing wrong result (closes :issue:`1085`).
 
 * fixed converting a scalar Array (an Array with 0 dimensions) to string with numpy 1.22+.
 
diff --git a/larray/core/axis.py b/larray/core/axis.py
@@ -2878,6 +2878,22 @@ def _key_to_axis_indices_dict(self, key):
                 extra_key_axes = axis_key.axes - self
                 if extra_key_axes:
                     raise ValueError(f"boolean subset key contains more axes ({axis_key.axes}) than array ({self})")
+
+                # TODO: factorize with check_compatible
+                for i, subset_axis in enumerate(axis_key.axes):
+                    array_axis = self.get_by_pos(subset_axis, i)
+                    if not array_axis.iscompatible(subset_axis):
+                        msg = f"""boolean subset array has incompatible axes with array:
+array axes: {self}
+subset array axes: {axis_key.axes}
+incompatible axes:
+    array axis:
+        {array_axis!r}
+    subset array axis:
+        {subset_axis!r}
+"""
+                        raise ValueError(msg)
+
                 # nonzero (currently) returns a tuple of IGroups containing 1D Arrays (one IGroup per axis)
                 filtered_key.extend(axis_key.nonzero())
             # drop slice(None) and Ellipsis since they are meaningless because of guess_axis.
diff --git a/larray/tests/test_array.py b/larray/tests/test_array.py
@@ -771,7 +771,8 @@ def test_getitem_bool_larray_key_arr_whout_bool_axis():
     raw = arr.data
 
     # all dimensions
-    res = arr[arr < 5]
+    filter_ = arr < 5
+    res = arr[filter_]
     assert isinstance(res, Array)
     assert res.ndim == 1
     assert_nparray_equal(res.data, raw[raw < 5])
@@ -786,6 +787,20 @@ def test_getitem_bool_larray_key_arr_whout_bool_axis():
     raw_d1, raw_d3 = raw_key.nonzero()
     assert_nparray_equal(res.data, raw[raw_d1, :, raw_d3])
 
+    # filter with smaller axis than array
+    filter_ = arr < 10
+    filter2 = filter_['c0,c2,c3']
+    with must_raise(ValueError, """boolean subset array has incompatible axes with array:
+array axes: {a, b, c}
+subset array axes: {a, b, c}
+incompatible axes:
+    array axis:
+        Axis(['c0', 'c1', 'c2', 'c3'], 'c')
+    subset array axis:
+        Axis(['c0', 'c2', 'c3'], 'c')
+"""):
+        _ = arr[filter2]
+
     # using an Axis object
     arr = ndtest('a=a0,a1;b=0..3')
     raw = arr.data
