Repository: pydata/xarray

Issue:
apply_ufunc produces illegal coordinate sizes
If `func` changes the size of the core dimension, `apply_ufunc(func, ..)` will only change the data variable dimension but keep the coordinate dimension and value unchanged. The resulting DataArray cannot be saved to a NetCDF file due to dimension inconsistency. Please see this [GitHub gist](https://gist.github.com/JiaweiZhuang/89a6bce7d5adc473134c7014bec065a0) for a simple illustration. 

Would it be more natural to drop the core coordinate by default? This is safer (will not produce illegal NetCDF file) and also makes more physical sense (the core coordinate is likely to change).

I'd like use `apply_ufunc` to track metadata in my xESMF package, but the output DataArray will have incorrect `lon` and `lat` coordinate dimension and values. It is easy to manually correct this afterwards but I want to bring up this issue.

#### Output of ``xr.show_versions()``
<details>

commit: None
python: 3.6.2.final.0
python-bits: 64
OS: Darwin
OS-release: 16.7.0
machine: x86_64
processor: i386
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: en_US.UTF-8

xarray: 0.10.0
pandas: 0.22.0
numpy: 1.13.3
scipy: 1.0.0
netCDF4: 1.3.1
h5netcdf: 0.5.0
Nio: None
bottleneck: 1.2.1
cyordereddict: None
dask: 0.17.0
matplotlib: 2.0.2
cartopy: 0.15.1
seaborn: 0.8.0
setuptools: 38.2.5
pip: 9.0.1
conda: None
pytest: 3.2.5
IPython: 6.1.0
sphinx: 1.6.5
</details>


Discussion:
shoyer: I plan to merge this shortly, unless anyone else has feedback here

Patch:
diff --git a/.gitignore b/.gitignore
index 92e488ed..2a016bb9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -38,6 +38,7 @@ nosetests.xml
 .ropeproject/
 .tags*
 .testmon*
+.tmontmp/
 .pytest_cache
 dask-worker-space/
 
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index aef80a2b..f3db96f9 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -46,6 +46,11 @@ Enhancements
   to manage its version strings. (:issue:`1300`).
   By `Joe Hamman <https://github.com/jhamman>`_.
 
+- `:py:class:`Dataset`s align `:py:class:`DataArray`s to coords that are explicitly
+  passed into the constructor, where previously an error would be raised.
+  (:issue:`674`)
+  By `Maximilian Roos <https://github.com/maxim-lian`_.
+
 - :py:meth:`~DataArray.sel`, :py:meth:`~DataArray.isel` & :py:meth:`~DataArray.reindex`,
   (and their :py:class:`Dataset` counterparts) now support supplying a ``dict``
   as a first argument, as an alternative to the existing approach 
@@ -62,6 +67,10 @@ Enhancements
 Bug fixes
 ~~~~~~~~~
 
+- :py:func:`apply_ufunc` now directly validates output variables
+  (:issue:`1931`).
+  By `Stephan Hoyer <https://github.com/shoyer>`_.
+
 - Fixed a bug where ``to_netcdf(..., unlimited_dims='bar')`` yielded NetCDF
   files with spurious 0-length dimensions (i.e. ``b``, ``a``, and ``r``)
   (:issue:`2134`).
diff --git a/xarray/core/computation.py b/xarray/core/computation.py
index f6e22dfe..ebbce114 100644
--- a/xarray/core/computation.py
+++ b/xarray/core/computation.py
@@ -513,7 +513,7 @@ def broadcast_compat_data(variable, broadcast_dims, core_dims):
 def apply_variable_ufunc(func, *args, **kwargs):
     """apply_variable_ufunc(func, *args, signature, exclude_dims=frozenset())
     """
-    from .variable import Variable
+    from .variable import Variable, as_compatible_data
 
     signature = kwargs.pop('signature')
     exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET)
@@ -559,20 +559,42 @@ def apply_variable_ufunc(func, *args, **kwargs):
                              'apply_ufunc: {}'.format(dask))
     result_data = func(*input_data)
 
-    if signature.num_outputs > 1:
-        output = []
-        for dims, data in zip(output_dims, result_data):
-            var = Variable(dims, data)
-            if keep_attrs and isinstance(args[0], Variable):
-                var.attrs.update(args[0].attrs)
-            output.append(var)
-        return tuple(output)
-    else:
-        dims, = output_dims
-        var = Variable(dims, result_data)
+    if signature.num_outputs == 1:
+        result_data = (result_data,)
+    elif (not isinstance(result_data, tuple) or
+            len(result_data) != signature.num_outputs):
+        raise ValueError('applied function does not have the number of '
+                         'outputs specified in the ufunc signature. '
+                         'Result is not a tuple of {} elements: {!r}'
+                         .format(signature.num_outputs, result_data))
+
+    output = []
+    for dims, data in zip(output_dims, result_data):
+        data = as_compatible_data(data)
+        if data.ndim != len(dims):
+            raise ValueError(
+                'applied function returned data with unexpected '
+                'number of dimensions: {} vs {}, for dimensions {}'
+                .format(data.ndim, len(dims), dims))
+
+        var = Variable(dims, data, fastpath=True)
+        for dim, new_size in var.sizes.items():
+            if dim in dim_sizes and new_size != dim_sizes[dim]:
+                raise ValueError(
+                    'size of dimension {!r} on inputs was unexpectedly '
+                    'changed by applied function from {} to {}. Only '
+                    'dimensions specified in ``exclude_dims`` with '
+                    'xarray.apply_ufunc are allowed to change size.'
+                    .format(dim, dim_sizes[dim], new_size))
+
         if keep_attrs and isinstance(args[0], Variable):
             var.attrs.update(args[0].attrs)
-        return var
+        output.append(var)
+
+    if signature.num_outputs == 1:
+        return output[0]
+    else:
+        return tuple(output)
 
 
 def _apply_with_dask_atop(func, args, input_dims, output_dims, signature,
@@ -719,7 +741,8 @@ def apply_ufunc(func, *args, **kwargs):
         Core dimensions on the inputs to exclude from alignment and
         broadcasting entirely. Any input coordinates along these dimensions
         will be dropped. Each excluded dimension must also appear in
-        ``input_core_dims`` for at least one argument.
+        ``input_core_dims`` for at least one argument. Only dimensions listed
+        here are allowed to change size between input and output objects.
     vectorize : bool, optional
         If True, then assume ``func`` only takes arrays defined over core
         dimensions as input and vectorize it automatically with
@@ -777,15 +800,38 @@ def apply_ufunc(func, *args, **kwargs):
 
     Examples
     --------
-    For illustrative purposes only, here are examples of how you could use
-    ``apply_ufunc`` to write functions to (very nearly) replicate existing
-    xarray functionality:
 
-    Calculate the vector magnitude of two arguments::
+    Calculate the vector magnitude of two arguments:
+
+    >>> def magnitude(a, b):
+    ...     func = lambda x, y: np.sqrt(x ** 2 + y ** 2)
+    ...     return xr.apply_ufunc(func, a, b)
+
+    You can now apply ``magnitude()`` to ``xr.DataArray`` and ``xr.Dataset``
+    objects, with automatically preserved dimensions and coordinates, e.g.,
+
+    >>> array = xr.DataArray([1, 2, 3], coords=[('x', [0.1, 0.2, 0.3])])
+    >>> magnitude(array, -array)
+    <xarray.DataArray (x: 3)>
+    array([1.414214, 2.828427, 4.242641])
+    Coordinates:
+      * x        (x) float64 0.1 0.2 0.3
+
+    Plain scalars, numpy arrays and a mix of these with xarray objects is also
+    supported:
+
+    >>> magnitude(4, 5)
+    5.0
+    >>> magnitude(3, np.array([0, 4]))
+    array([3., 5.])
+    >>> magnitude(array, 0)
+    <xarray.DataArray (x: 3)>
+    array([1., 2., 3.])
+    Coordinates:
+      * x        (x) float64 0.1 0.2 0.3
 
-        def magnitude(a, b):
-            func = lambda x, y: np.sqrt(x ** 2 + y ** 2)
-            return xr.apply_func(func, a, b)
+    Other examples of how you could use ``apply_ufunc`` to write functions to
+    (very nearly) replicate existing xarray functionality:
 
     Compute the mean (``.mean``) over one dimension::
 
@@ -795,7 +841,7 @@ def apply_ufunc(func, *args, **kwargs):
                                input_core_dims=[[dim]],
                                kwargs={'axis': -1})
 
-    Inner product over a specific dimension::
+    Inner product over a specific dimension (like ``xr.dot``)::
 
         def _inner(x, y):
             result = np.matmul(x[..., np.newaxis, :], y[..., :, np.newaxis])
diff --git a/xarray/core/merge.py b/xarray/core/merge.py
index 40a58d6e..f823717a 100644
--- a/xarray/core/merge.py
+++ b/xarray/core/merge.py
@@ -360,7 +360,17 @@ def merge_data_and_coords(data, coords, compat='broadcast_equals',
     """Used in Dataset.__init__."""
     objs = [data, coords]
     explicit_coords = coords.keys()
-    return merge_core(objs, compat, join, explicit_coords=explicit_coords)
+    indexes = dict(extract_indexes(coords))
+    return merge_core(objs, compat, join, explicit_coords=explicit_coords,
+                      indexes=indexes)
+
+
+def extract_indexes(coords):
+    """Yields the name & index of valid indexes from a mapping of coords"""
+    for name, variable in coords.items():
+        variable = as_variable(variable, name=name)
+        if variable.dims == (name,):
+            yield name, variable.to_index()
 
 
 def assert_valid_explicit_coords(variables, dims, explicit_coords):


Test patch:
diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
index bbcc02ba..37f97a81 100644
--- a/xarray/tests/test_computation.py
+++ b/xarray/tests/test_computation.py
@@ -752,6 +752,94 @@ def test_vectorize_dask():
     assert_identical(expected, actual)
 
 
+def test_output_wrong_number():
+    variable = xr.Variable('x', np.arange(10))
+
+    def identity(x):
+        return x
+
+    def tuple3x(x):
+        return (x, x, x)
+
+    with raises_regex(ValueError, 'number of outputs'):
+        apply_ufunc(identity, variable, output_core_dims=[(), ()])
+
+    with raises_regex(ValueError, 'number of outputs'):
+        apply_ufunc(tuple3x, variable, output_core_dims=[(), ()])
+
+
+def test_output_wrong_dims():
+    variable = xr.Variable('x', np.arange(10))
+
+    def add_dim(x):
+        return x[..., np.newaxis]
+
+    def remove_dim(x):
+        return x[..., 0]
+
+    with raises_regex(ValueError, 'unexpected number of dimensions'):
+        apply_ufunc(add_dim, variable, output_core_dims=[('y', 'z')])
+
+    with raises_regex(ValueError, 'unexpected number of dimensions'):
+        apply_ufunc(add_dim, variable)
+
+    with raises_regex(ValueError, 'unexpected number of dimensions'):
+        apply_ufunc(remove_dim, variable)
+
+
+def test_output_wrong_dim_size():
+    array = np.arange(10)
+    variable = xr.Variable('x', array)
+    data_array = xr.DataArray(variable, [('x', -array)])
+    dataset = xr.Dataset({'y': variable}, {'x': -array})
+
+    def truncate(array):
+        return array[:5]
+
+    def apply_truncate_broadcast_invalid(obj):
+        return apply_ufunc(truncate, obj)
+
+    with raises_regex(ValueError, 'size of dimension'):
+        apply_truncate_broadcast_invalid(variable)
+    with raises_regex(ValueError, 'size of dimension'):
+        apply_truncate_broadcast_invalid(data_array)
+    with raises_regex(ValueError, 'size of dimension'):
+        apply_truncate_broadcast_invalid(dataset)
+
+    def apply_truncate_x_x_invalid(obj):
+        return apply_ufunc(truncate, obj, input_core_dims=[['x']],
+                           output_core_dims=[['x']])
+
+    with raises_regex(ValueError, 'size of dimension'):
+        apply_truncate_x_x_invalid(variable)
+    with raises_regex(ValueError, 'size of dimension'):
+        apply_truncate_x_x_invalid(data_array)
+    with raises_regex(ValueError, 'size of dimension'):
+        apply_truncate_x_x_invalid(dataset)
+
+    def apply_truncate_x_z(obj):
+        return apply_ufunc(truncate, obj, input_core_dims=[['x']],
+                           output_core_dims=[['z']])
+
+    assert_identical(xr.Variable('z', array[:5]),
+                     apply_truncate_x_z(variable))
+    assert_identical(xr.DataArray(array[:5], dims=['z']),
+                     apply_truncate_x_z(data_array))
+    assert_identical(xr.Dataset({'y': ('z', array[:5])}),
+                     apply_truncate_x_z(dataset))
+
+    def apply_truncate_x_x_valid(obj):
+        return apply_ufunc(truncate, obj, input_core_dims=[['x']],
+                           output_core_dims=[['x']], exclude_dims={'x'})
+
+    assert_identical(xr.Variable('x', array[:5]),
+                     apply_truncate_x_x_valid(variable))
+    assert_identical(xr.DataArray(array[:5], dims=['x']),
+                     apply_truncate_x_x_valid(data_array))
+    assert_identical(xr.Dataset({'y': ('x', array[:5])}),
+                     apply_truncate_x_x_valid(dataset))
+
+
 @pytest.mark.parametrize('use_dask', [True, False])
 def test_dot(use_dask):
     if use_dask:
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index 2dad40ae..a0d316d7 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -4193,6 +4193,45 @@ def test_isin_dataset():
         ds.isin(ds)
 
 
+@pytest.mark.parametrize('unaligned_coords', (
+    {'x': [2, 1, 0]},
+    {'x': (['x'], np.asarray([2, 1, 0]))},
+    {'x': (['x'], np.asarray([1, 2, 0]))},
+    {'x': pd.Index([2, 1, 0])},
+    {'x': Variable(dims='x', data=[0, 2, 1])},
+    {'x': IndexVariable(dims='x', data=[0, 1, 2])},
+    {'y': 42},
+    {'y': ('x', [2, 1, 0])},
+    {'y': ('x', np.asarray([2, 1, 0]))},
+    {'y': (['x'], np.asarray([2, 1, 0]))},
+))
+@pytest.mark.parametrize('coords', (
+    {'x': ('x', [0, 1, 2])},
+    {'x': [0, 1, 2]},
+))
+def test_dataset_constructor_aligns_to_explicit_coords(
+        unaligned_coords, coords):
+
+    a = xr.DataArray([1, 2, 3], dims=['x'], coords=unaligned_coords)
+
+    expected = xr.Dataset(coords=coords)
+    expected['a'] = a
+
+    result = xr.Dataset({'a': a}, coords=coords)
+
+    assert_equal(expected, result)
+
+
+@pytest.mark.parametrize('unaligned_coords', (
+    {'y': ('b', np.asarray([2, 1, 0]))},
+))
+def test_constructor_raises_with_invalid_coords(unaligned_coords):
+
+    with pytest.raises(ValueError,
+                       message='not a subset of the DataArray dimensions'):
+        xr.DataArray([1, 2, 3], dims=['x'], coords=unaligned_coords)
+
+
 def test_dir_expected_attrs(data_set):
 
     some_expected_attrs = {'pipe', 'mean', 'isnull', 'var1',
@@ -4205,7 +4244,7 @@ def test_dir_non_string(data_set):
     # add a numbered key to ensure this doesn't break dir
     data_set[5] = 'foo'
     result = dir(data_set)
-    assert not (5 in result)
+    assert 5 not in result
 
     # GH2172
     sample_data = np.random.uniform(size=[2, 2000, 10000])