WIP: indexing with broadcasting

xarray/core/variable.py

@@ -422,62 +422,40 @@ def _broadcast_indexes(self, key):
         indexers: list of integer, array-like, or slice. This is aligned
             along self.dims.
         """
-        if not utils.is_dict_like(key):
-            key = {self.dims[0]: key}
-        example_v = None
-        indexes = OrderedDict()
-        for k, v in key.items():
-            if not isinstance(v, (integer_types, slice, Variable)):
-                if not hasattr(key, 'ndim'):  # convert list or tuple
-                    v = np.array(v)
-            if example_v is None and isinstance(v, Variable):
-                example_v = v
-            indexes[k] = v
-
-        # When all the keys are array or integer, slice
-        if example_v is None:
-            # more than one (unlabelled) arrays
-            if len([v for k, v in indexes.items()
-                    if not isinstance(v, (integer_types, slice))]) > 1:
-                raise IndexError("Indexing with multiple unlabelled arrays "
-                                 "is not allowed.")
-            # multi-dimensional unlabelled array
-            if any([v.ndim > 1 for k, v in indexes.items()
-                    if not isinstance(v, integer_types)]):
-                raise IndexError("Indexing with a multi-dimensional unlabelled"
-                                 "array is not allowed.")
-            # convert the array into Variable
-            for k, v in indexes.items():
-                if not hasattr(v, 'dims'):
-                    indexes[k] = type(self)([k], v)
-                    example_v = v
-
-        for k, v in indexes.items():
-            # Found unlabelled array
-            if not isinstance(v, (Variable, integer_types, slice)):
-                if (v.ndim > example_v.ndim or
-                        any([example_v.ndim != v.ndim for k, v
-                             in indexes.items() if isinstance(v, Variable)])):
-                    raise IndexError("Broadcasting failed because dimensions "
-                                     "does not match.")
-                else:
-                    _, indexes[k], _ = _broadcast_compat_data(example_v, v)
-
-        index_tuple = tuple(indexes.get(d, slice(None)) for d in self.dims)
-        index_tuple = indexing.expanded_indexer(index_tuple, self.ndim)
-
-        # comput dims
-        dims = []
-        for i, d in enumerate(self.dims):
-            if d in indexes.keys():
-                if isinstance(v, Variable):
-                    for d in v.dims:
-                        if d not in dims:
-                            dims.append(d)
-            else:
-                dims.append(d)
+        key = self._item_key_to_tuple(key)  # key is a tuple
+        # key is a tuple of full size
+        key = indexing.expanded_indexer(key, self.ndim)
+        basic_indexing_types = integer_types + (slice,)
+        if all([isinstance(k, basic_indexing_types) for k in key]):
+            return self._broadcast_indexes_basic(key)
+        else:
+            return self._broadcast_indexes_advanced(key)
+
+    def _broadcast_indexes_basic(self, key):
+        dims = tuple(dim for k, dim in zip(key, self.dims)
+                     if not isinstance(k, integer_types))
+        return dims, key
+
+    def _broadcast_indexes_advanced(self, key):
+        variables = []
 
-        return dims, index_tuple
+        for dim, value in zip(self.dims, key):
+            if isinstance(value, slice):
+                value = np.arange(self.sizes[dim])[value]
+
+            try:  # TODO we need our own Exception.
+                variable = as_variable(value, name=dim)
+            except ValueError as e:
+                if "cannot set variable" in str(e):
+                    raise IndexError("Unlabelled multi-dimensional array "
+                                     "cannot be used for indexing.")
+                else:
+                    raise e
+            variables.append(variable)
+        variables = _broadcast_compat_variables(*variables)
+        dims = variables[0].dims  # all variables have the same dims
+        key = tuple(variable.data for variable in variables)
+        return dims, key
 
     def getitem2(self, key):
         """Return a new Array object whose contents are consistent with

xarray/tests/test_variable.py

@@ -709,75 +709,96 @@ def test_items(self):
         v[range(10), range(11)] = 1
         self.assertArrayEqual(v.values, np.ones((10, 11)))
 
-    def test_getitem2(self):
+    def test_getitem2_basic(self):
         v = self.cls(['x', 'y'], [[0, 1, 2], [3, 4, 5]])
 
-        with self.assertRaisesRegexp(IndexError, "Indexing with multiple"):
-            v.getitem2(dict(x=[0, 1], y=[0, 1]))
+        v_new = v.getitem2(dict(x=0))
+        self.assertTrue(v_new.dims == ('y', ))
+        self.assertArrayEqual(v_new, v._data[0])
 
-        with self.assertRaisesRegexp(IndexError, "Indexing with a multi-"):
-            v.getitem2([[0, 1], [1, 2]])
+        v_new = v.getitem2(dict(x=0, y=slice(None)))
+        self.assertTrue(v_new.dims == ('y', ))
+        self.assertArrayEqual(v_new, v._data[0])
+
+        v_new = v.getitem2(dict(x=0, y=1))
+        self.assertTrue(v_new.dims == ())
+        self.assertArrayEqual(v_new, v._data[0, 1])
+
+        v_new = v.getitem2(dict(y=1))
+        self.assertTrue(v_new.dims == ('x', ))
+        self.assertArrayEqual(v_new, v._data[:, 1])
+
+        # tuple argument
+        v_new = v.getitem2((slice(None), 1))
+        self.assertTrue(v_new.dims == ('x', ))
+        self.assertArrayEqual(v_new, v._data[:, 1])
+
+    def test_getitem2_advanced(self):
+        v = self.cls(['x', 'y'], [[0, 1, 2], [3, 4, 5]])
+
+        # orthogonal indexing
+        v_new = v.getitem2(([0, 1], [1, 0]))
+        self.assertTrue(v_new.dims == ('x', 'y'))
+        self.assertArrayEqual(v_new, v._data[[0, 1]][:, [1, 0]])
 
-        dims, index_tuple = v._broadcast_indexes([0, 1])
-        self.assertTrue(dims == ['x', 'y'])
-        self.assertTrue(np.allclose(index_tuple[0], [0, 1]))
-        self.assertTrue(index_tuple[1] == slice(None, None, None))
         v_new = v.getitem2([0, 1])
         self.assertTrue(v_new.dims == ('x', 'y'))
         self.assertArrayEqual(v_new, v._data[[0, 1]])
 
         ind = Variable(['a', 'b'], [[0, 1, 1], [1, 1, 0]])
-        dims, index_tuple = v._broadcast_indexes(ind)
-        self.assertTrue(dims == ['a', 'b', 'y'])
-        self.assertTrue(np.allclose(index_tuple[0], [[0, 1, 1], [1, 1, 0]]))
-        self.assertTrue(index_tuple[1] == slice(None, None, None))
         v_new = v.getitem2(ind)
         self.assertTrue(v_new.dims == ('a', 'b', 'y'))
         self.assertArrayEqual(v_new, v._data[([0, 1, 1], [1, 1, 0]), :])
 
         ind = Variable(['a', 'b'], [[0, 1, 2], [2, 1, 0]])
-        dims, index_tuple = v._broadcast_indexes(dict(y=ind))
-        self.assertTrue(dims == ['x', 'a', 'b'])
-        self.assertTrue(len(index_tuple) == 2)
-        self.assertTrue(index_tuple[0] == slice(None, None, None))
-        self.assertTrue(np.allclose(index_tuple[1], [[0, 1, 2], [2, 1, 0]]))
         v_new = v.getitem2(dict(y=ind))
         self.assertTrue(v_new.dims == ('x', 'a', 'b'))
         self.assertArrayEqual(v_new, v._data[:, ([0, 1, 2], [2, 1, 0])])
 
-        # with broadcast
+        # with mixed arguments
         ind = Variable(['a'], [0, 1])
-        dims, index_tuple = v._broadcast_indexes(dict(x=[0, 1], y=ind))
-        self.assertTrue(dims == ['a'])
-        self.assertTrue(np.allclose(index_tuple[0], [0, 1]))
-        self.assertTrue(np.allclose(index_tuple[1], [0, 1]))
         v_new = v.getitem2(dict(x=[0, 1], y=ind))
-        self.assertArrayEqual(v_new, v._data[[0, 1], [0, 1]])
+        self.assertTrue(v_new.dims == ('x', 'a'))
+        self.assertArrayEqual(v_new, v._data[[0, 1]][:, [0, 1]])
 
         ind = Variable(['a', 'b'], [[0, 0], [1, 1]])
-        dims, index_tuple = v._broadcast_indexes(dict(x=[[1, 0], [1, 0]],
-                                                      y=ind))
-        self.assertTrue(dims == ['a', 'b'])
-        self.assertTrue(np.allclose(index_tuple[0], [[1, 0], [1, 0]]))
-        self.assertTrue(np.allclose(index_tuple[1], [[0, 0], [1, 1]]))
-        v_new = v.getitem2(dict(x=[[1, 0], [1, 0]], y=ind))
-        self.assertArrayEqual(v_new,
-                              v._data[([1, 0], [1, 0]), ([0, 0], [1, 1])])
-
-        # broadcast impossible case
-        with self.assertRaisesRegexp(IndexError, "Broadcasting failed "):
-            ind = Variable(['a'], [0, 1])
-            v.getitem2(dict(x=[[1, 0], [1, 0]], y=ind))
+        v_new = v.getitem2(dict(x=[1, 0], y=ind))
+        self.assertTrue(v_new.dims == ('x', 'a', 'b'))
+        self.assertArrayEqual(v_new, v._data[[1, 0]][:, ind])
 
         # with integer
         ind = Variable(['a', 'b'], [[0, 0], [1, 1]])
-        dims, index_tuple = v._broadcast_indexes(dict(x=0, y=ind))
-        self.assertTrue(dims == ['a', 'b'])
-        self.assertTrue(np.allclose(index_tuple[0], 0))
-        self.assertTrue(np.allclose(index_tuple[1], [[0, 0], [1, 1]]))
         v_new = v.getitem2(dict(x=0, y=ind))
-        self.assertArrayEqual(v_new,
-                              v._data[0, ([0, 0], [1, 1])])
+        self.assertTrue(v_new.dims == ('a', 'b'))
+        self.assertArrayEqual(v_new[0], v._data[0][[0, 0]])
+        self.assertArrayEqual(v_new[1], v._data[0][[1, 1]])
+
+        # with slice
+        ind = Variable(['a', 'b'], [[0, 0], [1, 1]])
+        v_new = v.getitem2(dict(x=slice(None), y=ind))
+        self.assertTrue(v_new.dims == ('x', 'a', 'b'))
+        self.assertArrayEqual(v_new, v._data[:, [[0, 0], [1, 1]]])
+
+        ind = Variable(['a', 'b'], [[0, 0], [1, 1]])
+        v_new = v.getitem2(dict(x=ind, y=slice(None)))
+        self.assertTrue(v_new.dims == ('a', 'b', 'y'))
+        self.assertArrayEqual(v_new, v._data[[[0, 0], [1, 1]], :])
+
+        ind = Variable(['a', 'b'], [[0, 0], [1, 1]])
+        v_new = v.getitem2(dict(x=ind, y=slice(None, 1)))
+        self.assertTrue(v_new.dims == ('a', 'b', 'y'))
+        self.assertArrayEqual(v_new, v._data[[[0, 0], [1, 1]], slice(None, 1)])
+
+    def test_getitem2_error(self):
+        v = self.cls(['x', 'y'], [[0, 1, 2], [3, 4, 5]])
+
+        with self.assertRaisesRegexp(IndexError, "Unlabelled multi-"):
+            v.getitem2([[0, 1], [1, 2]])
+
+        with self.assertRaisesRegexp(ValueError, "operands cannot be "):
+            ind_x = Variable(['a', 'b'], [[0, 0], [1, 1]])
+            ind_y = Variable(['a'], [0])
+            v.getitem2((ind_x, ind_y))
 
     def test_isel(self):
         v = Variable(['time', 'x'], self.d)