a[bool_expr] now uses LArray broadcasting

gdementen · gdementen · commit de387e55ec8e · 2015-01-28T16:09:32.000+01:00
diff --git a/larray/core.py b/larray/core.py
@@ -989,6 +989,8 @@ def __setitem__(self, key, value, collapse_slices=True):
         data = np.asarray(self)
 
         if isinstance(key, np.ndarray) and np.issubdtype(key.dtype, bool):
+            if isinstance(key, LArray):
+                key = key.broadcast_with(self.axes)
             data[key] = value
             return
 
@@ -1036,7 +1038,7 @@ def reshape_like(self, target):
 
     def broadcast_with(self, target):
         """
-        returns an LArray that is broadcastable with target
+        returns an LArray that is (numpy) broadcastable with target
         target can be either an LArray or any collection of Axis
 
         * all common axes must be either 1 or the same length
diff --git a/larray/tests/test_la.py b/larray/tests/test_la.py
@@ -712,13 +712,33 @@ def test_setitem_ndarray(self):
         self._assert_equal_raw(la, raw)
 
     def test_setitem_bool_array_key(self):
+        age, geo, sex, lipro = self.larray.axes
+
         # LArray key
+        # a) same shape
         la = self.larray.copy()
         raw = self.array.copy()
         la[la < 5] = 0
         raw[raw < 5] = 0
         self._assert_equal_raw(la, raw)
 
+        # b) numpy-broadcastable shape
+        la = self.larray.copy()
+        raw = self.array.copy()
+        key = la[sex['F,']] < 5
+        self.assertEqual(key.ndim, 4)
+        la[key] = 0
+        raw[raw[:, :, [1]] < 5] = 0
+        self._assert_equal_raw(la, raw)
+
+        # c) LArray-broadcastable shape (missing axis)
+        la = self.larray.copy()
+        key = la[sex['F']] < 5
+        self.assertEqual(key.ndim, 3)
+        la[key] = 0
+        # same raw as last test
+        self._assert_equal_raw(la, raw)
+
         # ndarray key
         la = self.larray.copy()
         raw = self.array.copy()
