3720

pending completion

Build # 3720

Build Type

push

travis-ci

Committed by

jhamman

Commit Message

<a href="https://github.com/pydata/xarray/commit/<a class=hub.com/pydata/xarray/commit/14b5f1ca84b41696b7e8e3ec5e07d8064acc9a55">14b5f1ca8<a href="https://github.com/pydata/xarray/commit/14b5f1ca84b41696b7e8e3ec5e07d8064acc9a55">&quot;&gt;Load nonindex coords ahead of concat() (#1551)

* Load non-index coords to memory ahead of concat

* Update unit test after #1522

* Minimise loads on concat. Extend new concat logic to data_vars.

* Trivial tweaks

* Added unit tests

Fix loads when vars are found different halfway through

* Add xfail for #1586

* Revert &amp;quot;Add xfail for #1586&quot;

This reverts commit </a><a class="double-link" href="https://github.com/pydata/xarray/commit/<a class="double-link" href="https://github.com/pydata/xarray/commit/f99313cc7b266964cb540c4dbd858f701461eb2b">f99313cc7</a>">f99313cc7</a><a href="https://github.com/pydata/xarray/commit/14b5f1ca84b41696b7e8e3ec5e07d8064acc9a55">.

Run Details

90 of 90 new or added lines in 2 files covered. (100.0%)

17291 of 17741 relevant lines covered (97.46%)

7.24 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

97.35

/xarray/core/combine.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import warnings

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, IndexVariable, concat as concat_vars


def concat(objs, dim=None, data_vars='all', coords='different',
           compat='equals', positions=None, indexers=None, mode=None,
           concat_over=None):
    """Concatenate xarray objects along a new or existing dimension.

    Parameters
    ----------
    objs : sequence of Dataset and DataArray objects
        xarray objects to concatenate together. Each object is expected to
        consist of variables and coordinates with matching shapes except for
        along the concatenated dimension.
    dim : str or DataArray or pandas.Index
        Name of the dimension to concatenate along. This can either be a new
        dimension name, in which case it is added along axis=0, or an existing
        dimension name, in which case the location of the dimension is
        unchanged. If dimension is provided as a DataArray or Index, its name
        is used as the dimension to concatenate along and the values are added
        as a coordinate.
    data_vars : {'minimal', 'different', 'all' or list of str}, optional
        These data variables will be concatenated together:
          * 'minimal': Only data variables in which the dimension already
            appears are included.
          * 'different': Data variables which are not equal (ignoring
            attributes) across all datasets are also concatenated (as well as
            all for which dimension already appears). Beware: this option may
            load the data payload of data variables into memory if they are not
            already loaded.
          * 'all': All data variables will be concatenated.
          * list of str: The listed data variables will be concatenated, in
            addition to the 'minimal' data variables.
        If objects are DataArrays, data_vars must be 'all'.
    coords : {'minimal', 'different', 'all' o list of str}, optional
        These coordinate variables will be concatenated together:
          * 'minimal': Only coordinates in which the dimension already appears
            are included.
          * 'different': Coordinates which are not equal (ignoring attributes)
            across all datasets are also concatenated (as well as all for which
            dimension already appears). Beware: this option may load the data
            payload of coordinate variables into memory if they are not already
            loaded.
          * 'all': All coordinate variables will be concatenated, except
            those corresponding to other dimensions.
          * list of str: The listed coordinate variables will be concatenated,
            in addition the 'minimal' coordinates.
    compat : {'equals', 'identical'}, optional
        String indicating how to compare non-concatenated variables and
        dataset global attributes for potential conflicts. 'equals' means
        that all variable values and dimensions must be the same;
        'identical' means that variable attributes and global attributes
        must also be equal.
    positions : None or list of integer arrays, optional
        List of integer arrays which specifies the integer positions to which
        to assign each dataset along the concatenated dimension. If not
        supplied, objects are concatenated in the provided order.
    indexers, mode, concat_over : deprecated

    Returns
    -------
    concatenated : type of objs

    See also
    --------
    merge
    auto_combine
    """
    # TODO: add join and ignore_index arguments copied from pandas.concat
    # TODO: support concatenating scalar coordinates even if the concatenated
    # dimension already exists
    from .dataset import Dataset
    from .dataarray import DataArray

    try:
        first_obj, objs = utils.peek_at(objs)
    except StopIteration:
        raise ValueError('must supply at least one object to concatenate')

    if dim is None:
        warnings.warn('the `dim` argument to `concat` will be required '
                      'in a future version of xarray; for now, setting it to '
                      "the old default of 'concat_dim'",
                      FutureWarning, stacklevel=2)
        dim = 'concat_dims'

    if indexers is not None:  # pragma: nocover
        warnings.warn('indexers has been renamed to positions; the alias '
                      'will be removed in a future version of xarray',
                      FutureWarning, stacklevel=2)
        positions = indexers

    if mode is not None:
        raise ValueError('`mode` is no longer a valid argument to '
                         'xarray.concat; it has been split into the `data_vars` '
                         'and `coords` arguments')
    if concat_over is not None:
        raise ValueError('`concat_over` is no longer a valid argument to '
                         'xarray.concat; it has been split into the `data_vars` '
                         'and `coords` arguments')

    if isinstance(first_obj, DataArray):
        f = _dataarray_concat
    elif isinstance(first_obj, Dataset):
        f = _dataset_concat
    else:
        raise TypeError('can only concatenate xarray Dataset and DataArray '
                        'objects, got %s' % type(first_obj))
    return f(objs, dim, data_vars, coords, compat, positions)


def _calc_concat_dim_coord(dim):
    """
    Infer the dimension name and 1d coordinate variable (if appropriate)
    for concatenating along the new dimension.
    """
    if isinstance(dim, basestring):
        coord = None
    elif not hasattr(dim, 'dims'):
        # dim is not a DataArray or IndexVariable
        dim_name = getattr(dim, 'name', None)
        if dim_name is None:
            dim_name = 'concat_dim'
        coord = IndexVariable(dim_name, dim)
        dim = dim_name
    elif not hasattr(dim, 'name'):
        coord = as_variable(dim).to_index_variable()
        dim, = coord.dims
    else:
        coord = dim
        dim, = coord.dims
    return dim, coord


def _calc_concat_over(datasets, dim, data_vars, coords):
    """
    Determine which dataset variables need to be concatenated in the result,
    and which can simply be taken from the first dataset.
    """
    # Return values
    concat_over = set()
    equals = {}

    if dim in datasets[0]:
        concat_over.add(dim)
    for ds in datasets:
        concat_over.update(k for k, v in ds.variables.items()
                           if dim in v.dims)

    def process_subset_opt(opt, subset):
        if isinstance(opt, basestring):
            if opt == 'different':
                # all nonindexes that are not the same in each dataset
                for k in getattr(datasets[0], subset):
                    if k not in concat_over:
                        # Compare the variable of all datasets vs. the one
                        # of the first dataset. Perform the minimum amount of
                        # loads in order to avoid multiple loads from disk while
                        # keeping the RAM footprint low.
                        v_lhs = datasets[0].variables[k].load()
                        # We'll need to know later on if variables are equal.
                        computed = []
                        for ds_rhs in datasets[1:]:
                            v_rhs = ds_rhs.variables[k].compute()
                            computed.append(v_rhs)
                            if not v_lhs.equals(v_rhs):
                                concat_over.add(k)
                                equals[k] = False
                                # computed variables are not to be re-computed
                                # again in the future
                                for ds, v in zip(datasets[1:], computed):
                                    ds.variables[k].data = v.data
                                break
                        else:
                            equals[k] = True

            elif opt == 'all':
                concat_over.update(set(getattr(datasets[0], subset)) -
                                   set(datasets[0].dims))
            elif opt == 'minimal':
                pass
            else:
                raise ValueError("unexpected value for %s: %s" % (subset, opt))
        else:
            invalid_vars = [k for k in opt
                            if k not in getattr(datasets[0], subset)]
            if invalid_vars:
                if subset == 'coords':
                    raise ValueError(
                        'some variables in coords are not coordinates on '
                        'the first dataset: %s' % invalid_vars)
                else:
                    raise ValueError(
                        'some variables in data_vars are not data variables on '
                        'the first dataset: %s' % invalid_vars)
            concat_over.update(opt)

    process_subset_opt(data_vars, 'data_vars')
    process_subset_opt(coords, 'coords')
    return concat_over, equals


def _dataset_concat(datasets, dim, data_vars, coords, compat, positions):
    """
    Concatenate a sequence of datasets along a new or existing dimension
    """
    from .dataset import Dataset

    if compat not in ['equals', 'identical']:
        raise ValueError("compat=%r invalid: must be 'equals' "
                         "or 'identical'" % compat)

    dim, coord = _calc_concat_dim_coord(dim)
    # Make sure we're working on a copy (we'll be loading variables)
    datasets = [ds.copy() for ds in datasets]
    datasets = align(*datasets, join='outer', copy=False, exclude=[dim])

    concat_over, equals = _calc_concat_over(datasets, dim, data_vars, coords)

    def insert_result_variable(k, v):
        assert isinstance(v, Variable)
        if k in datasets[0].coords:
            result_coord_names.add(k)
        result_vars[k] = v

    # create the new dataset and add constant variables
    result_vars = OrderedDict()
    result_coord_names = set(datasets[0].coords)
    result_attrs = datasets[0].attrs
    result_encoding = datasets[0].encoding

    for k, v in datasets[0].variables.items():
        if k not in concat_over:
            insert_result_variable(k, v)

    # check that global attributes and non-concatenated variables are fixed
    # across all datasets
    for ds in datasets[1:]:
        if (compat == 'identical' and
                not utils.dict_equiv(ds.attrs, result_attrs)):
            raise ValueError('dataset global attributes not equal')
        for k, v in iteritems(ds.variables):
            if k not in result_vars and k not in concat_over:
                raise ValueError('encountered unexpected variable %r' % k)
            elif (k in result_coord_names) != (k in ds.coords):
                raise ValueError('%r is a coordinate in some datasets but not '
                                 'others' % k)
            elif k in result_vars and k != dim:
                # Don't use Variable.identical as it internally invokes
                # Variable.equals, and we may already know the answer
                if compat == 'identical' and not utils.dict_equiv(
                        v.attrs, result_vars[k].attrs):
                    raise ValueError(
                        'variable %s not identical across datasets' % k)

                # Proceed with equals()
                try:
                    # May be populated when using the "different" method
                    is_equal = equals[k]
                except KeyError:
                    result_vars[k].load()
                    is_equal = v.equals(result_vars[k])
                if not is_equal:
                    raise ValueError(
                        'variable %s not equal across datasets' % k)


    # we've already verified everything is consistent; now, calculate
    # shared dimension sizes so we can expand the necessary variables
    dim_lengths = [ds.dims.get(dim, 1) for ds in datasets]
    non_concat_dims = {}
    for ds in datasets:
        non_concat_dims.update(ds.dims)
    non_concat_dims.pop(dim, None)

    def ensure_common_dims(vars):
        # ensure each variable with the given name shares the same
        # dimensions and the same shape for all of them except along the
        # concat dimension
        common_dims = tuple(pd.unique([d for v in vars for d in v.dims]))
        if dim not in common_dims:
            common_dims = (dim,) + common_dims
        for var, dim_len in zip(vars, dim_lengths):
            if var.dims != common_dims:
                common_shape = tuple(non_concat_dims.get(d, dim_len)
                                     for d in common_dims)
                var = var.set_dims(common_dims, common_shape)
            yield var

    # stack up each variable to fill-out the dataset (in order)
    for k in datasets[0].variables:
        if k in concat_over:
            vars = ensure_common_dims([ds.variables[k] for ds in datasets])
            combined = concat_vars(vars, dim, positions)
            insert_result_variable(k, combined)

    result = Dataset(result_vars, attrs=result_attrs)
    result = result.set_coords(result_coord_names)
    result.encoding = result_encoding

    if coord is not None:
        # add concat dimension last to ensure that its in the final Dataset
        result[coord.name] = coord

    return result


def _dataarray_concat(arrays, dim, data_vars, coords, compat,
                      positions):
    arrays = list(arrays)

    if data_vars != 'all':
        raise ValueError('data_vars is not a valid argument when '
                         'concatenating DataArray objects')

    datasets = []
    for n, arr in enumerate(arrays):
        if n == 0:
            name = arr.name
        elif name != arr.name:
            if compat == 'identical':
                raise ValueError('array names not identical')
            else:
                arr = arr.rename(name)
        datasets.append(arr._to_temp_dataset())

    ds = _dataset_concat(datasets, dim, data_vars, coords, compat,
                         positions)
    return arrays[0]._from_temp_dataset(ds, name)


def _auto_concat(datasets, dim=None):
    if len(datasets) == 1:
        return datasets[0]
    else:
        if dim is None:
            ds0 = datasets[0]
            ds1 = datasets[1]
            concat_dims = set(ds0.dims)
            if ds0.dims != ds1.dims:
                dim_tuples = set(ds0.dims.items()) - set(ds1.dims.items())
                concat_dims = set(i for i, _ in dim_tuples)
            if len(concat_dims) > 1:
                concat_dims = set(d for d in concat_dims
                                  if not ds0[d].equals(ds1[d]))
            if len(concat_dims) > 1:
                raise ValueError('too many different dimensions to '
                                 'concatenate: %s' % concat_dims)
            elif len(concat_dims) == 0:
                raise ValueError('cannot infer dimension to concatenate: '
                                 'supply the ``concat_dim`` argument '
                                 'explicitly')
            dim, = concat_dims
        return concat(datasets, dim=dim)


_CONCAT_DIM_DEFAULT = '__infer_concat_dim__'


def auto_combine(datasets,
                 concat_dim=_CONCAT_DIM_DEFAULT,
                 compat='no_conflicts'):
    """Attempt to auto-magically combine the given datasets into one.

    This method attempts to combine a list of datasets into a single entity by
    inspecting metadata and using a combination of concat and merge.

    It does not concatenate along more than one dimension or sort data under any
    circumstances. It does align coordinates, but different variables on
    datasets can cause it to fail under some scenarios. In complex cases, you
    may need to clean up your data and use ``concat``/``merge`` explicitly.

    ``auto_combine`` works well if you have N years of data and M data
    variables, and each combination of a distinct time period and set of data
    variables is saved its own dataset.

    Parameters
    ----------
    datasets : sequence of xarray.Dataset
        Dataset objects to merge.
    concat_dim : str or DataArray or Index, optional
        Dimension along which to concatenate variables, as used by
        :py:func:`xarray.concat`. You only need to provide this argument if the
        dimension along which you want to concatenate is not a dimension in
        the original datasets, e.g., if you want to stack a collection of
        2D arrays along a third dimension.
        By default, xarray attempts to infer this argument by examining
        component files. Set ``concat_dim=None`` explicitly to disable
        concatenation.
    compat : {'identical', 'equals', 'broadcast_equals',
              'no_conflicts'}, optional
        String indicating how to compare variables of the same name for
        potential conflicts:

        - 'broadcast_equals': all values must be equal when variables are
          broadcast against each other to ensure common dimensions.
        - 'equals': all values and dimensions must be the same.
        - 'identical': all values, dimensions and attributes must be the
          same.
        - 'no_conflicts': only values which are not null in both datasets
          must be equal. The returned dataset then contains the combination
          of all non-null values.

    Returns
    -------
    combined : xarray.Dataset

    See also
    --------
    concat
    Dataset.merge
    """
    from toolz import itertoolz
    if concat_dim is not None:
        dim = None if concat_dim is _CONCAT_DIM_DEFAULT else concat_dim
        grouped = itertoolz.groupby(lambda ds: tuple(sorted(ds.data_vars)),
                                    datasets).values()
        concatenated = [_auto_concat(ds, dim=dim) for ds in grouped]
    else:
        concatenated = datasets
    merged = merge(concatenated, compat=compat)
    return merged

1	from __future__ import absolute_import	8✔
2	from __future__ import division	8✔
3	from __future__ import print_function	8✔
4	import warnings	8✔
5
6	import pandas as pd	8✔
7
8	from . import utils	8✔
9	from .alignment import align	8✔
10	from .merge import merge	8✔
11	from .pycompat import iteritems, OrderedDict, basestring	8✔
12	from .variable import Variable, as_variable, IndexVariable, concat as concat_vars	8✔
13
14
15	def concat(objs, dim=None, data_vars='all', coords='different',	8✔
16	compat='equals', positions=None, indexers=None, mode=None,
17	concat_over=None):
18	"""Concatenate xarray objects along a new or existing dimension.
19
20	Parameters
21	----------
22	objs : sequence of Dataset and DataArray objects
23	xarray objects to concatenate together. Each object is expected to
24	consist of variables and coordinates with matching shapes except for
25	along the concatenated dimension.
26	dim : str or DataArray or pandas.Index
27	Name of the dimension to concatenate along. This can either be a new
28	dimension name, in which case it is added along axis=0, or an existing
29	dimension name, in which case the location of the dimension is
30	unchanged. If dimension is provided as a DataArray or Index, its name
31	is used as the dimension to concatenate along and the values are added
32	as a coordinate.
33	data_vars : {'minimal', 'different', 'all' or list of str}, optional
34	These data variables will be concatenated together:
35	* 'minimal': Only data variables in which the dimension already
36	appears are included.
37	* 'different': Data variables which are not equal (ignoring
38	attributes) across all datasets are also concatenated (as well as
39	all for which dimension already appears). Beware: this option may
40	load the data payload of data variables into memory if they are not
41	already loaded.
42	* 'all': All data variables will be concatenated.
43	* list of str: The listed data variables will be concatenated, in
44	addition to the 'minimal' data variables.
45	If objects are DataArrays, data_vars must be 'all'.
46	coords : {'minimal', 'different', 'all' o list of str}, optional
47	These coordinate variables will be concatenated together:
48	* 'minimal': Only coordinates in which the dimension already appears
49	are included.
50	* 'different': Coordinates which are not equal (ignoring attributes)
51	across all datasets are also concatenated (as well as all for which
52	dimension already appears). Beware: this option may load the data
53	payload of coordinate variables into memory if they are not already
54	loaded.
55	* 'all': All coordinate variables will be concatenated, except
56	those corresponding to other dimensions.
57	* list of str: The listed coordinate variables will be concatenated,
58	in addition the 'minimal' coordinates.
59	compat : {'equals', 'identical'}, optional
60	String indicating how to compare non-concatenated variables and
61	dataset global attributes for potential conflicts. 'equals' means
62	that all variable values and dimensions must be the same;
63	'identical' means that variable attributes and global attributes
64	must also be equal.
65	positions : None or list of integer arrays, optional
66	List of integer arrays which specifies the integer positions to which
67	to assign each dataset along the concatenated dimension. If not
68	supplied, objects are concatenated in the provided order.
69	indexers, mode, concat_over : deprecated
70
71	Returns
72	-------
73	concatenated : type of objs
74
75	See also
76	--------
77	merge
78	auto_combine
79	"""
80	# TODO: add join and ignore_index arguments copied from pandas.concat
81	# TODO: support concatenating scalar coordinates even if the concatenated
82	# dimension already exists
83	from .dataset import Dataset	8✔
84	from .dataarray import DataArray	8✔
85
86	try:	8✔
87	first_obj, objs = utils.peek_at(objs)	8✔
88	except StopIteration:	8✔
89	raise ValueError('must supply at least one object to concatenate')	8✔
90
91	if dim is None:	8✔
92	warnings.warn('the `dim` argument to `concat` will be required '	×
93	'in a future version of xarray; for now, setting it to '
94	"the old default of 'concat_dim'",
95	FutureWarning, stacklevel=2)
96	dim = 'concat_dims'	×
97
98	if indexers is not None: # pragma: nocover
99	warnings.warn('indexers has been renamed to positions; the alias '
100	'will be removed in a future version of xarray',
101	FutureWarning, stacklevel=2)
102	positions = indexers
103
104	if mode is not None:	8✔
105	raise ValueError('`mode` is no longer a valid argument to '	8✔
106	'xarray.concat; it has been split into the `data_vars` '
107	'and `coords` arguments')
108	if concat_over is not None:	8✔
109	raise ValueError('`concat_over` is no longer a valid argument to '	8✔
110	'xarray.concat; it has been split into the `data_vars` '
111	'and `coords` arguments')
112
113	if isinstance(first_obj, DataArray):	8✔
114	f = _dataarray_concat	8✔
115	elif isinstance(first_obj, Dataset):	8✔
116	f = _dataset_concat	8✔
117	else:
118	raise TypeError('can only concatenate xarray Dataset and DataArray '	×
119	'objects, got %s' % type(first_obj))
120	return f(objs, dim, data_vars, coords, compat, positions)	8✔
121
122
123	def _calc_concat_dim_coord(dim):	8✔
124	"""
125	Infer the dimension name and 1d coordinate variable (if appropriate)
126	for concatenating along the new dimension.
127	"""
128	if isinstance(dim, basestring):	8✔
129	coord = None	8✔
130	elif not hasattr(dim, 'dims'):	8✔
131	# dim is not a DataArray or IndexVariable
132	dim_name = getattr(dim, 'name', None)	8✔
133	if dim_name is None:	8✔
134	dim_name = 'concat_dim'	8✔
135	coord = IndexVariable(dim_name, dim)	8✔
136	dim = dim_name	8✔
137	elif not hasattr(dim, 'name'):	8✔
138	coord = as_variable(dim).to_index_variable()	8✔
139	dim, = coord.dims	8✔
140	else:
141	coord = dim	8✔
142	dim, = coord.dims	8✔
143	return dim, coord	8✔
144
145
146	def _calc_concat_over(datasets, dim, data_vars, coords):	8✔
147	"""
148	Determine which dataset variables need to be concatenated in the result,
149	and which can simply be taken from the first dataset.
150	"""
151	# Return values
152	concat_over = set()	8✔
153	equals = {}	8✔
154
155	if dim in datasets[0]:	8✔
156	concat_over.add(dim)	8✔
157	for ds in datasets:	8✔
158	concat_over.update(k for k, v in ds.variables.items()	8✔
159	if dim in v.dims)
160
161	def process_subset_opt(opt, subset):	8✔
162	if isinstance(opt, basestring):	8✔
163	if opt == 'different':	8✔
164	# all nonindexes that are not the same in each dataset
165	for k in getattr(datasets[0], subset):	8✔
166	if k not in concat_over:	8✔
167	# Compare the variable of all datasets vs. the one
168	# of the first dataset. Perform the minimum amount of
169	# loads in order to avoid multiple loads from disk while
170	# keeping the RAM footprint low.
171	v_lhs = datasets[0].variables[k].load()	8✔
172	# We'll need to know later on if variables are equal.
173	computed = []	8✔
174	for ds_rhs in datasets[1:]:	8✔
175	v_rhs = ds_rhs.variables[k].compute()	8✔
176	computed.append(v_rhs)	8✔
177	if not v_lhs.equals(v_rhs):	8✔
178	concat_over.add(k)	8✔
179	equals[k] = False	8✔
180	# computed variables are not to be re-computed
181	# again in the future
182	for ds, v in zip(datasets[1:], computed):	8✔
183	ds.variables[k].data = v.data	8✔
184	break	8✔
185	else:
186	equals[k] = True	8✔
187
188	elif opt == 'all':	8✔
189	concat_over.update(set(getattr(datasets[0], subset)) -	8✔
190	set(datasets[0].dims))
191	elif opt == 'minimal':	8✔
192	pass	8✔
193	else:
194	raise ValueError("unexpected value for %s: %s" % (subset, opt))	8✔
195	else:
196	invalid_vars = [k for k in opt	8✔
197	if k not in getattr(datasets[0], subset)]
198	if invalid_vars:	8✔
199	if subset == 'coords':	8✔
200	raise ValueError(	8✔
201	'some variables in coords are not coordinates on '
202	'the first dataset: %s' % invalid_vars)
203	else:
204	raise ValueError(	×
205	'some variables in data_vars are not data variables on '
206	'the first dataset: %s' % invalid_vars)
207	concat_over.update(opt)	8✔
208
209	process_subset_opt(data_vars, 'data_vars')	8✔
210	process_subset_opt(coords, 'coords')	8✔
211	return concat_over, equals	8✔
212
213
214	def _dataset_concat(datasets, dim, data_vars, coords, compat, positions):	8✔
215	"""
216	Concatenate a sequence of datasets along a new or existing dimension
217	"""
218	from .dataset import Dataset	8✔
219
220	if compat not in ['equals', 'identical']:	8✔
221	raise ValueError("compat=%r invalid: must be 'equals' "	8✔
222	"or 'identical'" % compat)
223
224	dim, coord = _calc_concat_dim_coord(dim)	8✔
225	# Make sure we're working on a copy (we'll be loading variables)
226	datasets = [ds.copy() for ds in datasets]	8✔
227	datasets = align(*datasets, join='outer', copy=False, exclude=[dim])	8✔
228
229	concat_over, equals = _calc_concat_over(datasets, dim, data_vars, coords)	8✔
230
231	def insert_result_variable(k, v):	8✔
232	assert isinstance(v, Variable)	8✔
233	if k in datasets[0].coords:	8✔
234	result_coord_names.add(k)	8✔
235	result_vars[k] = v	8✔
236
237	# create the new dataset and add constant variables
238	result_vars = OrderedDict()	8✔
239	result_coord_names = set(datasets[0].coords)	8✔
240	result_attrs = datasets[0].attrs	8✔
241	result_encoding = datasets[0].encoding	8✔
242
243	for k, v in datasets[0].variables.items():	8✔
244	if k not in concat_over:	8✔
245	insert_result_variable(k, v)	8✔
246
247	# check that global attributes and non-concatenated variables are fixed
248	# across all datasets
249	for ds in datasets[1:]:	8✔
250	if (compat == 'identical' and	8✔
251	not utils.dict_equiv(ds.attrs, result_attrs)):
252	raise ValueError('dataset global attributes not equal')	8✔
253	for k, v in iteritems(ds.variables):	8✔
254	if k not in result_vars and k not in concat_over:	8✔
255	raise ValueError('encountered unexpected variable %r' % k)	8✔
256	elif (k in result_coord_names) != (k in ds.coords):	8✔
257	raise ValueError('%r is a coordinate in some datasets but not '	8✔
258	'others' % k)
259	elif k in result_vars and k != dim:	8✔
260	# Don't use Variable.identical as it internally invokes
261	# Variable.equals, and we may already know the answer
262	if compat == 'identical' and not utils.dict_equiv(	8✔
263	v.attrs, result_vars[k].attrs):
264	raise ValueError(	×
265	'variable %s not identical across datasets' % k)
266
267	# Proceed with equals()
268	try:	8✔
269	# May be populated when using the "different" method
270	is_equal = equals[k]	8✔
271	except KeyError:	8✔
272	result_vars[k].load()	8✔
273	is_equal = v.equals(result_vars[k])	8✔
274	if not is_equal:	8✔
275	raise ValueError(	8✔
276	'variable %s not equal across datasets' % k)
277
278
279	# we've already verified everything is consistent; now, calculate
280	# shared dimension sizes so we can expand the necessary variables
281	dim_lengths = [ds.dims.get(dim, 1) for ds in datasets]	8✔
282	non_concat_dims = {}	8✔
283	for ds in datasets:	8✔
284	non_concat_dims.update(ds.dims)	8✔
285	non_concat_dims.pop(dim, None)	8✔
286
287	def ensure_common_dims(vars):	8✔
288	# ensure each variable with the given name shares the same
289	# dimensions and the same shape for all of them except along the
290	# concat dimension
291	common_dims = tuple(pd.unique([d for v in vars for d in v.dims]))	8✔
292	if dim not in common_dims:	8✔
293	common_dims = (dim,) + common_dims	8✔
294	for var, dim_len in zip(vars, dim_lengths):	8✔
295	if var.dims != common_dims:	8✔
296	common_shape = tuple(non_concat_dims.get(d, dim_len)	8✔
297	for d in common_dims)
298	var = var.set_dims(common_dims, common_shape)	8✔
299	yield var	8✔
300
301	# stack up each variable to fill-out the dataset (in order)
302	for k in datasets[0].variables:	8✔
303	if k in concat_over:	8✔
304	vars = ensure_common_dims([ds.variables[k] for ds in datasets])	8✔
305	combined = concat_vars(vars, dim, positions)	8✔
306	insert_result_variable(k, combined)	8✔
307
308	result = Dataset(result_vars, attrs=result_attrs)	8✔
309	result = result.set_coords(result_coord_names)	8✔
310	result.encoding = result_encoding	8✔
311
312	if coord is not None:	8✔
313	# add concat dimension last to ensure that its in the final Dataset
314	result[coord.name] = coord	8✔
315
316	return result	8✔
317
318
319	def _dataarray_concat(arrays, dim, data_vars, coords, compat,	8✔
320	positions):
321	arrays = list(arrays)	8✔
322
323	if data_vars != 'all':	8✔
324	raise ValueError('data_vars is not a valid argument when '	8✔
325	'concatenating DataArray objects')
326
327	datasets = []	8✔
328	for n, arr in enumerate(arrays):	8✔
329	if n == 0:	8✔
330	name = arr.name	8✔
331	elif name != arr.name:	8✔
332	if compat == 'identical':	8✔
333	raise ValueError('array names not identical')	8✔
334	else:
335	arr = arr.rename(name)	8✔
336	datasets.append(arr._to_temp_dataset())	8✔
337
338	ds = _dataset_concat(datasets, dim, data_vars, coords, compat,	8✔
339	positions)
340	return arrays[0]._from_temp_dataset(ds, name)	8✔
341
342
343	def _auto_concat(datasets, dim=None):	8✔
344	if len(datasets) == 1:	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
345	return datasets[0]	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
346	else:
347	if dim is None:	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
348	ds0 = datasets[0]	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
349	ds1 = datasets[1]	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
350	concat_dims = set(ds0.dims)	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
351	if ds0.dims != ds1.dims:	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
352	dim_tuples = set(ds0.dims.items()) - set(ds1.dims.items())	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
353	concat_dims = set(i for i, _ in dim_tuples)	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
354	if len(concat_dims) > 1:	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
355	concat_dims = set(d for d in concat_dims	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
356	if not ds0[d].equals(ds1[d]))
357	if len(concat_dims) > 1:	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
358	raise ValueError('too many different dimensions to '	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
359	'concatenate: %s' % concat_dims)
360	elif len(concat_dims) == 0:	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
361	raise ValueError('cannot infer dimension to concatenate: '	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
362	'supply the ``concat_dim`` argument '
363	'explicitly')
364	dim, = concat_dims	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
365	return concat(datasets, dim=dim)	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
366
367
368	_CONCAT_DIM_DEFAULT = '__infer_concat_dim__'	8✔
369
370
371	def auto_combine(datasets,	8✔
372	concat_dim=_CONCAT_DIM_DEFAULT,
373	compat='no_conflicts'):
374	"""Attempt to auto-magically combine the given datasets into one.
375
376	This method attempts to combine a list of datasets into a single entity by
377	inspecting metadata and using a combination of concat and merge.
378
379	It does not concatenate along more than one dimension or sort data under any
380	circumstances. It does align coordinates, but different variables on
381	datasets can cause it to fail under some scenarios. In complex cases, you
382	may need to clean up your data and use ``concat``/``merge`` explicitly.
383
384	``auto_combine`` works well if you have N years of data and M data
385	variables, and each combination of a distinct time period and set of data
386	variables is saved its own dataset.
387
388	Parameters
389	----------
390	datasets : sequence of xarray.Dataset
391	Dataset objects to merge.
392	concat_dim : str or DataArray or Index, optional
393	Dimension along which to concatenate variables, as used by
394	:py:func:`xarray.concat`. You only need to provide this argument if the
395	dimension along which you want to concatenate is not a dimension in
396	the original datasets, e.g., if you want to stack a collection of
397	2D arrays along a third dimension.
398	By default, xarray attempts to infer this argument by examining
399	component files. Set ``concat_dim=None`` explicitly to disable
400	concatenation.
401	compat : {'identical', 'equals', 'broadcast_equals',
402	'no_conflicts'}, optional
403	String indicating how to compare variables of the same name for
404	potential conflicts:
405
406	- 'broadcast_equals': all values must be equal when variables are
407	broadcast against each other to ensure common dimensions.
408	- 'equals': all values and dimensions must be the same.
409	- 'identical': all values, dimensions and attributes must be the
410	same.
411	- 'no_conflicts': only values which are not null in both datasets
412	must be equal. The returned dataset then contains the combination
413	of all non-null values.
414
415	Returns
416	-------
417	combined : xarray.Dataset
418
419	See also
420	--------
421	concat
422	Dataset.merge
423	"""
424	from toolz import itertoolz	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
425	if concat_dim is not None:	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
426	dim = None if concat_dim is _CONCAT_DIM_DEFAULT else concat_dim	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
427	grouped = itertoolz.groupby(lambda ds: tuple(sorted(ds.data_vars)),	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
428	datasets).values()
429	concatenated = [_auto_concat(ds, dim=dim) for ds in grouped]	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
430	else:
431	concatenated = datasets	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
432	merged = merge(concatenated, compat=compat)	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔
433	return merged	6 all except CONDA_ENV=py34 and CONDA_ENV=py27-min ✔

pydata / xarray / 3720

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous