fixed arr[bool_axis[bool_key]] (closes #735)

Author: gdementen

doc/source/changes/version_0_31.rst.inc

@@ -33,4 +33,6 @@ Miscellaneous improvements
 Fixes
 ^^^^^
 
-* fixed something (closes :issue:`1`).
+* fixed taking a subset of an array with boolean labels for an axis if the user explicitly specify the axis
+  (closes :issue:`735`). When the user does not specify the axis, it currently fails but it is unclear what to do in
+  that case (see :issue:`794`).

larray/core/axis.py

@@ -887,9 +887,6 @@ def index(self, key):
             # stop is inclusive in the input key and exclusive in the output !
             stop = mapping[key.stop] + 1 if key.stop is not None else None
             return slice(start, stop, key.step)
-        # XXX: bool LArray do not pass through???
-        elif isinstance(key, np.ndarray) and key.dtype.kind is 'b':
-            return key
         elif isinstance(key, (tuple, list, OrderedSet)):
             # TODO: the result should be cached
             # Note that this is faster than array_lookup(np.array(key), mapping)
@@ -2730,7 +2727,7 @@ def _key_to_igroups(self, key):
             if isinstance(axis_key, LArray) and np.issubdtype(axis_key.dtype, np.bool_):
                 extra_key_axes = axis_key.axes - self
                 if extra_key_axes:
-                    raise ValueError("subset key contains more axes ({}) than array ({})"
+                    raise ValueError("boolean subset key contains more axes ({}) than array ({})"
                                      .format(axis_key.axes, self))
                 # nonzero (currently) returns a tuple of IGroups containing 1D LArrays (one IGroup per axis)
                 nonboolkey.extend(axis_key.nonzero())

larray/tests/test_array.py

@@ -2,6 +2,7 @@
 from __future__ import absolute_import, division, print_function
 
 import os
+import re
 import sys
 
 import pytest
@@ -30,6 +31,7 @@
 # Test Value Strings #
 # ================== #
 
+
 def test_value_string_split():
     assert_array_equal(_to_ticks('M,F'), np.asarray(['M', 'F']))
     assert_array_equal(_to_ticks('M, F'), np.asarray(['M', 'F']))
@@ -721,7 +723,7 @@ def test_getitem_abstract_positional(array):
         array[X.age.i[2, 3], X.age.i[1, 5]]
 
 
-def test_getitem_bool_larray_key():
+def test_getitem_bool_larray_key_arr_whout_bool_axis():
     arr = ndtest((3, 2, 4))
     raw = arr.data
 
@@ -732,8 +734,8 @@ def test_getitem_bool_larray_key():
     assert_array_equal(res, raw[raw < 5])
 
     # missing dimension
-    filt = arr['b1'] % 5 == 0
-    res = arr[filt]
+    filter_ = arr['b1'] % 5 == 0
+    res = arr[filter_]
     assert isinstance(res, LArray)
     assert res.ndim == 2
     assert res.shape == (3, 2)
@@ -752,6 +754,26 @@ def test_getitem_bool_larray_key():
     assert_array_equal(res, raw[:, :2])
 
 
+def test_getitem_bool_larray_key_arr_wh_bool_axis():
+    gender = Axis([False, True], 'gender')
+    arr = LArray([0.1, 0.2], gender)
+    id_axis = Axis('id=0..3')
+    key = LArray([True, False, True, True], id_axis)
+    expected = LArray([0.2, 0.1, 0.2, 0.2], id_axis)
+
+    # LGroup using the real axis
+    assert_larray_equal(arr[gender[key]], expected)
+
+    # LGroup using an AxisReference
+    assert_larray_equal(arr[X.gender[key]], expected)
+
+    # this test checks that the current behavior does not change unintentionally...
+    # ... but I am unsure the current behavior is what we actually want
+    msg = re.escape("boolean subset key contains more axes ({id}) than array ({gender})")
+    with pytest.raises(ValueError, match=msg):
+        arr[key]
+
+
 def test_getitem_bool_larray_and_group_key():
     arr = ndtest((3, 6, 4)).set_labels('b', '0..5')
 
@@ -769,14 +791,34 @@ def test_getitem_bool_larray_and_group_key():
     assert_array_equal(res, expected)
 
 
-def test_getitem_bool_ndarray_key(array):
+def test_getitem_bool_ndarray_key_arr_whout_bool_axis(array):
     raw = array.data
     res = array[raw < 5]
     assert isinstance(res, LArray)
     assert res.ndim == 1
     assert_array_equal(res, raw[raw < 5])
 
 
+def test_getitem_bool_ndarray_key_arr_wh_bool_axis():
+    gender = Axis([False, True], 'gender')
+    arr = LArray([0.1, 0.2], gender)
+    key = np.array([True, False, True, True])
+    expected = arr.i[[1, 0, 1, 1]]
+
+    # LGroup using the real axis
+    assert_larray_equal(arr[gender[key]], expected)
+
+    # LGroup using an AxisReference
+    assert_larray_equal(arr[X.gender[key]], expected)
+
+    # raw key => ???
+    # this test checks that the current behavior does not change unintentionally...
+    # ... but I am unsure the current behavior is what we actually want
+    msg = re.escape("boolean key with a different shape ((4,)) than array ((2,))")
+    with pytest.raises(ValueError, match=msg):
+        arr[key]
+
+
 def test_getitem_bool_anonymous_axes():
     a = ndtest([Axis(2), Axis(3), Axis(4), Axis(5)])
     mask = ones(a.axes[1, 3], dtype=bool)