Align arguments, auto align in broadcast and concat (#963)

* align(): added exclude and index optional parameters
broadcast(): added exclude and copy optional parameters. Auto-align misaligned attrs

This is unfinished work - see TODO

* Recommitted changes after rebaselining

* Added tests for broadcast on DataArray with copy=True|False and exclude= parameters

* Removed copy flag from broadcast (now it's always False).
Added/fixed unit tests

* Fixed bug in dataset broadcast. Completed unit tests on align() and broadcast().

* Cleaned up code in broadcast().
Use source_ndarray() in no-copy tests

* concat() to auto-align inputs.
Fixed failing unit test on no-copy.
Updated What's New.

* tweaked What's New

Author: crusaderky

doc/whats-new.rst

@@ -28,6 +28,12 @@ Enhancements
   xarray objects (:issue:`432`). See :ref:`dictionary IO <dictionary io>`
   for more details. By `Julia Signell <https://github.com/jsignell>`_.
 
+- Added ``exclude`` and ``indexes`` optional parameters to :py:func:`~xarray.align`,
+  and ``exclude`` optional parameter to :py:func:`~xarray.broadcast`.
+  By `Guido Imperiale <https://github.com/crusaderky>`_.
+- :py:func:`~xarray.broadcast` and :py:func:`~xarray.concat` will now auto-align inputs,
+  using ``join=outer``. By `Guido Imperiale <https://github.com/crusaderky>`_.
+
 Bug fixes
 ~~~~~~~~~
 

xarray/core/alignment.py

@@ -9,7 +9,7 @@
 from .common import _maybe_promote
 from .pycompat import iteritems, OrderedDict
 from .utils import is_full_slice, is_dict_like
-from .variable import Variable, Coordinate, broadcast_variables
+from .variable import Variable, Coordinate
 
 
 def _get_joiner(join):
@@ -71,25 +71,17 @@ def align(*objects, **kwargs):
         If ``copy=True``, the returned objects contain all new variables. If
         ``copy=False`` and no reindexing is required then the aligned objects
         will include original variables.
+    exclude : sequence of str, optional
+        Dimensions that must be excluded from alignment
+    indexes : dict-like, optional
+        Any indexes explicitly provided with the `indexes` argument should be
+        used in preference to the aligned indexes.
 
     Returns
     -------
     aligned : same as *objects
         Tuple of objects with aligned coordinates.
     """
-    return partial_align(*objects, exclude=None, **kwargs)
-
-
-def partial_align(*objects, **kwargs):
-    """partial_align(*objects, join='inner', copy=True, indexes=None,
-                     exclude=set())
-
-    Like align, but don't align along dimensions in exclude. Any indexes
-    explicitly provided with the `indexes` argument should be used in preference
-    to the aligned indexes.
-
-    Not public API.
-    """
     join = kwargs.pop('join', 'inner')
     copy = kwargs.pop('copy', True)
     indexes = kwargs.pop('indexes', None)
@@ -247,21 +239,26 @@ def var_indexers(var, indexers):
     return reindexed
 
 
-def broadcast(*args):
+def broadcast(*args, **kwargs):
     """Explicitly broadcast any number of DataArray or Dataset objects against
     one another.
 
     xarray objects automatically broadcast against each other in arithmetic
     operations, so this function should not be necessary for normal use.
 
+    If no change is needed, the input data is returned to the output without
+    being copied.
+
     Parameters
     ----------
-    *args: DataArray or Dataset objects
+    *args : DataArray or Dataset objects
         Arrays to broadcast against each other.
+    exclude : sequence of str, optional
+        Dimensions that must not be broadcasted
 
     Returns
     -------
-    broadcast: tuple of xarray objects
+    broadcast : tuple of xarray objects
         The same data as the input arrays, but with additional dimensions
         inserted so that all data arrays have the same dimensions and shape.
 
@@ -322,36 +319,48 @@ def broadcast(*args):
     from .dataarray import DataArray
     from .dataset import Dataset
 
-    all_indexes = _get_all_indexes(args)
-    for k, v in all_indexes.items():
-        if not all(v[0].equals(vi) for vi in v[1:]):
-            raise ValueError('cannot broadcast arrays: the %s index is not '
-                             'aligned (use xarray.align first)' % k)
+    exclude = kwargs.pop('exclude', None)
+    if exclude is None:
+        exclude = set()
+    if kwargs:
+        raise TypeError('broadcast() got unexpected keyword arguments: %s'
+                        % list(kwargs))
+
+    args = align(*args, join='outer', copy=False, exclude=exclude)
 
     common_coords = OrderedDict()
     dims_map = OrderedDict()
     for arg in args:
         for dim in arg.dims:
-            if dim not in common_coords:
+            if dim not in common_coords and dim not in exclude:
                 common_coords[dim] = arg.coords[dim].variable
                 dims_map[dim] = common_coords[dim].size
 
+    def _expand_dims(var):
+        # Add excluded dims to a copy of dims_map    
+        var_dims_map = dims_map.copy()
+        for dim in exclude:
+            try:
+                var_dims_map[dim] = var.shape[var.dims.index(dim)]
+            except ValueError:
+                # dim not in var.dims
+                pass
+
+        return var.expand_dims(var_dims_map)
+
     def _broadcast_array(array):
-        data = array.variable.expand_dims(dims_map)
+        data = _expand_dims(array.variable)
         coords = OrderedDict(array.coords)
         coords.update(common_coords)
-        dims = tuple(common_coords)
-        return DataArray(data, coords, dims, name=array.name,
+        return DataArray(data, coords, data.dims, name=array.name,
                          attrs=array.attrs, encoding=array.encoding)
 
     def _broadcast_dataset(ds):
-        data_vars = OrderedDict()
-        for k in ds.data_vars:
-            data_vars[k] = ds.variables[k].expand_dims(dims_map)
-
+        data_vars = OrderedDict(
+            (k, _expand_dims(ds.variables[k]))
+            for k in ds.data_vars)
         coords = OrderedDict(ds.coords)
         coords.update(common_coords)
-
         return Dataset(data_vars, coords, ds.attrs)
 
     result = []
@@ -367,6 +376,7 @@ def _broadcast_dataset(ds):
 
 
 def broadcast_arrays(*args):
+    import warnings
     warnings.warn('xarray.broadcast_arrays is deprecated: use '
                   'xarray.broadcast instead', DeprecationWarning, stacklevel=2)
     return broadcast(*args)

xarray/core/combine.py

@@ -3,6 +3,7 @@
 import pandas as pd
 
 from . import utils
+from .alignment import align
 from .merge import merge
 from .pycompat import iteritems, OrderedDict, basestring
 from .variable import Variable, as_variable, Coordinate, concat as concat_vars
@@ -202,10 +203,9 @@ def _dataset_concat(datasets, dim, data_vars, coords, compat, positions):
         raise ValueError("compat=%r invalid: must be 'equals' "
                          "or 'identical'" % compat)
 
-    # don't bother trying to work with datasets as a generator instead of a
-    # list; the gains would be minimal
-    datasets = [as_dataset(ds) for ds in datasets]
     dim, coord = _calc_concat_dim_coord(dim)
+    datasets = [as_dataset(ds) for ds in datasets]
+    datasets = align(*datasets, join='outer', copy=False, exclude=[dim])
 
     concat_over = _calc_concat_over(datasets, dim, data_vars, coords)
 

xarray/core/merge.py

@@ -1,6 +1,6 @@
 import pandas as pd
 
-from .alignment import partial_align
+from .alignment import align
 from .utils import Frozen, is_dict_like
 from .variable import as_variable, default_index_coordinate
 from .pycompat import (basestring, OrderedDict)
@@ -322,7 +322,7 @@ def is_alignable(obj):
         # https://github.com/pydata/xarray/issues/943
         return input_objects
 
-    aligned = partial_align(*targets, join=join, copy=copy, indexes=indexes)
+    aligned = align(*targets, join=join, copy=copy, indexes=indexes)
 
     for position, key, aligned_obj in zip(positions, keys, aligned):
         if key is no_key:

xarray/test/test_combine.py

@@ -106,6 +106,14 @@ def test_concat_size0(self):
         actual = concat(split_data[::-1], 'dim1')
         self.assertDatasetIdentical(data, actual)
 
+    def test_concat_autoalign(self):
+        ds1 = Dataset({'foo': DataArray([1, 2], coords={'x': [1, 2]})})
+        ds2 = Dataset({'foo': DataArray([1, 2], coords={'x': [1, 3]})})        
+        actual = concat([ds1, ds2], 'y')
+        expected = Dataset({'foo': DataArray([[1, 2, np.nan], [1, np.nan, 2]],
+                                             dims=['y', 'x'], coords={'y': [0, 1], 'x': [1, 2, 3]})})
+        self.assertDatasetIdentical(expected, actual)
+
     def test_concat_errors(self):
         data = create_test_data()
         split_data = [data.isel(dim1=slice(3)),
@@ -129,14 +137,6 @@ def test_concat_errors(self):
             data1['foo'] = ('bar', np.random.randn(10))
             concat([data0, data1], 'dim1')
 
-        with self.assertRaisesRegexp(ValueError, 'not equal across datasets'):
-            data0, data1 = deepcopy(split_data)
-            data1['dim2'] = 2 * data1['dim2']
-            concat([data0, data1], 'dim1', coords='minimal')
-
-        with self.assertRaisesRegexp(ValueError, 'it is not 1-dimensional'):
-            concat([data0, data1], 'dim1')
-
         with self.assertRaisesRegexp(ValueError, 'compat.* invalid'):
             concat(split_data, 'dim1', compat='foobar')
 

xarray/test/test_dataarray.py

@@ -1618,6 +1618,60 @@ def test_align_dtype(self):
         c, d = align(a, b, join='outer')
         self.assertEqual(c.dtype, np.float32)
 
+    def test_align_copy(self):
+        x = DataArray([1, 2, 3], coords=[('a', [1, 2, 3])])
+        y = DataArray([1, 2], coords=[('a', [3, 1])])
+        
+        expected_x2 = x
+        expected_y2 = DataArray([2, np.nan, 1], coords=[('a', [1, 2, 3])])
+
+        x2, y2 = align(x, y, join='outer', copy=False)
+        self.assertDataArrayIdentical(expected_x2, x2)
+        self.assertDataArrayIdentical(expected_y2, y2)
+        assert source_ndarray(x2.data) is source_ndarray(x.data)
+    
+        x2, y2 = align(x, y, join='outer', copy=True)
+        self.assertDataArrayIdentical(expected_x2, x2)
+        self.assertDataArrayIdentical(expected_y2, y2)
+        assert source_ndarray(x2.data) is not source_ndarray(x.data)
+
+        # Trivial align - 1 element
+        x = DataArray([1, 2, 3], coords=[('a', [1, 2, 3])])
+        x2, = align(x, copy=False)
+        self.assertDataArrayIdentical(x, x2)
+        assert source_ndarray(x2.data) is source_ndarray(x.data)
+    
+        x2, = align(x, copy=True)
+        self.assertDataArrayIdentical(x, x2)
+        assert source_ndarray(x2.data) is not source_ndarray(x.data)
+
+    def test_align_exclude(self):
+        x = DataArray([[1, 2], [3, 4]], coords=[('a', [-1, -2]), ('b', [3, 4])])
+        y = DataArray([[1, 2], [3, 4]], coords=[('a', [-1, 20]), ('b', [5, 6])])
+        z = DataArray([1], dims=['a'], coords={'a': [20], 'b': 7})
+        
+        x2, y2, z2 = align(x, y, z, join='outer', exclude=['b'])
+        expected_x2 = DataArray([[3, 4], [1, 2], [np.nan, np.nan]], coords=[('a', [-2, -1, 20]), ('b', [3, 4])])
+        expected_y2 = DataArray([[np.nan, np.nan], [1, 2], [3, 4]], coords=[('a', [-2, -1, 20]), ('b', [5, 6])])
+        expected_z2 = DataArray([np.nan, np.nan, 1], dims=['a'], coords={'a': [-2, -1, 20], 'b': 7})
+        self.assertDataArrayIdentical(expected_x2, x2)
+        self.assertDataArrayIdentical(expected_y2, y2)
+        self.assertDataArrayIdentical(expected_z2, z2)        
+
+    def test_align_indexes(self):
+        x = DataArray([1, 2, 3], coords=[('a', [-1, 10, -2])])
+        y = DataArray([1, 2], coords=[('a', [-2, -1])])
+
+        x2, y2 = align(x, y, join='outer', indexes={'a': [10, -1, -2]})
+        expected_x2 = DataArray([2, 1, 3], coords=[('a', [10, -1, -2])])
+        expected_y2 = DataArray([np.nan, 2, 1], coords=[('a', [10, -1, -2])])
+        self.assertDataArrayIdentical(expected_x2, x2)
+        self.assertDataArrayIdentical(expected_y2, y2)
+
+        x2, = align(x, join='outer', indexes={'a': [-2, 7, 10, -1]})
+        expected_x2 = DataArray([3, np.nan, 2, 1], coords=[('a', [-2, 7, 10, -1])])
+        self.assertDataArrayIdentical(expected_x2, x2)
+
     def test_broadcast_arrays(self):
         x = DataArray([1, 2], coords=[('a', [-1, -2])], name='x')
         y = DataArray([1, 2], coords=[('b', [3, 4])], name='y')
@@ -1636,9 +1690,45 @@ def test_broadcast_arrays(self):
         self.assertDataArrayIdentical(expected_x2, x2)
         self.assertDataArrayIdentical(expected_y2, y2)
 
-        z = DataArray([1, 2], coords=[('a', [-10, 20])])
-        with self.assertRaisesRegexp(ValueError, 'cannot broadcast'):
-            broadcast(x, z)
+    def test_broadcast_arrays_misaligned(self):
+        # broadcast on misaligned coords must auto-align
+        x = DataArray([[1, 2], [3, 4]], coords=[('a', [-1, -2]), ('b', [3, 4])])
+        y = DataArray([1, 2], coords=[('a', [-1, 20])])
+        expected_x2 = DataArray([[3, 4], [1, 2], [np.nan, np.nan]], coords=[('a', [-2, -1, 20]), ('b', [3, 4])])
+        expected_y2 = DataArray([[np.nan, np.nan], [1, 1], [2, 2]], coords=[('a', [-2, -1, 20]), ('b', [3, 4])])
+        x2, y2 = broadcast(x, y)
+        self.assertDataArrayIdentical(expected_x2, x2)
+        self.assertDataArrayIdentical(expected_y2, y2)
+
+    def test_broadcast_arrays_nocopy(self):
+        # Test that input data is not copied over in case no alteration is needed
+        x = DataArray([1, 2], coords=[('a', [-1, -2])], name='x')
+        y = DataArray(3, name='y')
+        expected_x2 = DataArray([1, 2], coords=[('a', [-1, -2])], name='x')
+        expected_y2 = DataArray([3, 3], coords=[('a', [-1, -2])], name='y')
+
+        x2, y2 = broadcast(x, y)
+        self.assertDataArrayIdentical(expected_x2, x2)
+        self.assertDataArrayIdentical(expected_y2, y2)
+        assert source_ndarray(x2.data) is source_ndarray(x.data)
+        
+        # single-element broadcast (trivial case)
+        x2, = broadcast(x)
+        self.assertDataArrayIdentical(x, x2)
+        assert source_ndarray(x2.data) is source_ndarray(x.data)
+
+    def test_broadcast_arrays_exclude(self):
+        x = DataArray([[1, 2], [3, 4]], coords=[('a', [-1, -2]), ('b', [3, 4])])
+        y = DataArray([1, 2], coords=[('a', [-1, 20])])
+        z = DataArray(5, coords={'b': 5})
+        
+        x2, y2, z2 = broadcast(x, y, z, exclude=['b'])
+        expected_x2 = DataArray([[3, 4], [1, 2], [np.nan, np.nan]], coords=[('a', [-2, -1, 20]), ('b', [3, 4])])
+        expected_y2 = DataArray([np.nan, 1, 2], coords=[('a', [-2, -1, 20])])
+        expected_z2 = DataArray([5, 5, 5], dims=['a'], coords={'a': [-2, -1, 20], 'b': 5})
+        self.assertDataArrayIdentical(expected_x2, x2)
+        self.assertDataArrayIdentical(expected_y2, y2)
+        self.assertDataArrayIdentical(expected_z2, z2)
 
     def test_broadcast_coordinates(self):
         # regression test for GH649