Skip to content

Update computation.py to use Python 3 function signatures #2756

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 12, 2019
Merged

Update computation.py to use Python 3 function signatures #2756

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
51bdbb9
Update computation.py to use Python 3 function signatures
shoyer Feb 8, 2019
a985217
Lint / py35 fixup
shoyer Feb 8, 2019
77b7744
Merge branch 'master' into computation-py3
shoyer Feb 11, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions setup.cfg
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@ ignore_missing_imports = True
ignore_missing_imports = True
[mypy-Nio.*]
ignore_missing_imports = True
[mypy-nc_time_axis.*]
ignore_missing_imports = True
[mypy-numpy.*]
ignore_missing_imports = True
[mypy-netCDF4.*]
Expand Down
213 changes: 92 additions & 121 deletions xarray/core/computation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@
from collections import Counter, OrderedDict
from distutils.version import LooseVersion
from typing import (
AbstractSet, Any, Dict, Iterable, List, Mapping, Union, Tuple,
TYPE_CHECKING, TypeVar
AbstractSet, Any, Callable, Iterable, List, Mapping, Optional, Sequence,
Tuple, TYPE_CHECKING, Union,
)

import numpy as np
Expand Down Expand Up @@ -190,20 +190,19 @@ def build_output_coords(
return output_coords


def apply_dataarray_ufunc(func, *args, **kwargs):
"""apply_dataarray_ufunc(func, *args, signature, join='inner',
exclude_dims=frozenset())
def apply_dataarray_vfunc(
func,
*args,
signature,
join='inner',
exclude_dims=frozenset(),
keep_attrs=False
):
"""Apply a variable level function over DataArray, Variable and/or ndarray
objects.
"""
from .dataarray import DataArray

signature = kwargs.pop('signature')
join = kwargs.pop('join', 'inner')
exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET)
keep_attrs = kwargs.pop('keep_attrs', True)
if kwargs:
raise TypeError('apply_dataarray_ufunc() got unexpected keyword '
'arguments: %s' % list(kwargs))

if len(args) > 1:
args = deep_align(args, join=join, copy=False, exclude=exclude_dims,
raise_on_invalid=False)
Expand Down Expand Up @@ -261,15 +260,19 @@ def assert_and_return_exact_match(all_keys):
}


def join_dict_keys(objects, how='inner'):
# type: (Iterable[Union[Mapping, Any]], str) -> Iterable
def join_dict_keys(
objects: Iterable[Union[Mapping, Any]], how: str = 'inner',
) -> Iterable:
joiner = _JOINERS[how]
all_keys = [obj.keys() for obj in objects if hasattr(obj, 'keys')]
return joiner(all_keys)


def collect_dict_values(objects, keys, fill_value=None):
# type: (Iterable[Union[Mapping, Any]], Iterable, Any) -> List[list]
def collect_dict_values(
objects: Iterable[Union[Mapping, Any]],
keys: Iterable,
fill_value: object = None,
) -> List[list]:
return [[obj.get(key, fill_value)
if is_dict_like(obj)
else obj
Expand Down Expand Up @@ -298,17 +301,12 @@ def _unpack_dict_tuples(
return out


def apply_dict_of_variables_ufunc(func, *args, **kwargs):
"""apply_dict_of_variables_ufunc(func, *args, signature, join='inner',
fill_value=None):
def apply_dict_of_variables_vfunc(
func, *args, signature, join='inner', fill_value=None
):
"""Apply a variable level function over dicts of DataArray, DataArray,
Variable and ndarray objects.
"""
signature = kwargs.pop('signature')
join = kwargs.pop('join', 'inner')
fill_value = kwargs.pop('fill_value', None)
if kwargs:
raise TypeError('apply_dict_of_variables_ufunc() got unexpected '
'keyword arguments: %s' % list(kwargs))

args = [_as_variables_or_variable(arg) for arg in args]
names = join_dict_keys(args, how=join)
grouped_by_name = collect_dict_values(args, names, fill_value)
Expand All @@ -323,8 +321,10 @@ def apply_dict_of_variables_ufunc(func, *args, **kwargs):
return result_vars


def _fast_dataset(variables, coord_variables):
# type: (OrderedDict[Any, Variable], Mapping[Any, Variable]) -> Dataset
def _fast_dataset(
variables: 'OrderedDict[Any, Variable]',
coord_variables: Mapping[Any, Variable],
) -> 'Dataset':
"""Create a dataset as quickly as possible.

Beware: the `variables` OrderedDict is modified INPLACE.
Expand All @@ -335,21 +335,20 @@ def _fast_dataset(variables, coord_variables):
return Dataset._from_vars_and_coord_names(variables, coord_names)


def apply_dataset_ufunc(func, *args, **kwargs):
"""apply_dataset_ufunc(func, *args, signature, join='inner',
dataset_join='inner', fill_value=None,
exclude_dims=frozenset(), keep_attrs=False):

If dataset_join != 'inner', a non-default fill_value must be supplied
by the user. Otherwise a TypeError is raised.
def apply_dataset_vfunc(
func,
*args,
signature,
join='inner',
dataset_join='exact',
fill_value=_NO_FILL_VALUE,
exclude_dims=frozenset(),
keep_attrs=False
):
"""Apply a variable level function over Dataset, dict of DataArray,
DataArray, Variable and/or ndarray objects.
"""
from .dataset import Dataset
signature = kwargs.pop('signature')
join = kwargs.pop('join', 'inner')
dataset_join = kwargs.pop('dataset_join', 'inner')
fill_value = kwargs.pop('fill_value', None)
exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET)
keep_attrs = kwargs.pop('keep_attrs', False)
first_obj = args[0] # we'll copy attrs from this in case keep_attrs=True

if (dataset_join not in _JOINS_WITHOUT_FILL_VALUES and
Expand All @@ -358,17 +357,14 @@ def apply_dataset_ufunc(func, *args, **kwargs):
'data variables with apply_ufunc, you must supply the '
'dataset_fill_value argument.')

if kwargs:
raise TypeError('apply_dataset_ufunc() got unexpected keyword '
'arguments: %s' % list(kwargs))
if len(args) > 1:
args = deep_align(args, join=join, copy=False, exclude=exclude_dims,
raise_on_invalid=False)

list_of_coords = build_output_coords(args, signature, exclude_dims)
args = [getattr(arg, 'data_vars', arg) for arg in args]

result_vars = apply_dict_of_variables_ufunc(
result_vars = apply_dict_of_variables_vfunc(
func, *args, signature=signature, join=dataset_join,
fill_value=fill_value)

Expand Down Expand Up @@ -402,7 +398,10 @@ def _iter_over_selections(obj, dim, values):
yield obj_sel


def apply_groupby_ufunc(func, *args):
def apply_groupby_func(func, *args):
"""Apply a dataset or datarray level function over GroupBy, Dataset,
DataArray, Variable and/or ndarray objects.
"""
from .groupby import GroupBy, peek_at
from .variable import Variable

Expand Down Expand Up @@ -444,7 +443,7 @@ def apply_groupby_ufunc(func, *args):

def unified_dim_sizes(
variables: Iterable[Variable],
exclude_dims: AbstractSet = frozenset(),
exclude_dims: AbstractSet = frozenset()
) -> 'OrderedDict[Any, int]':

dim_sizes = OrderedDict() # type: OrderedDict[Any, int]
Expand Down Expand Up @@ -516,21 +515,20 @@ def broadcast_compat_data(variable, broadcast_dims, core_dims):
return data


def apply_variable_ufunc(func, *args, **kwargs):
"""apply_variable_ufunc(func, *args, signature, exclude_dims=frozenset())
def apply_variable_ufunc(
func,
*args,
signature,
exclude_dims=frozenset(),
dask='forbidden',
output_dtypes=None,
output_sizes=None,
keep_attrs=False
):
"""Apply a ndarray level function over Variable and/or ndarray objects.
"""
from .variable import Variable, as_compatible_data

signature = kwargs.pop('signature')
exclude_dims = kwargs.pop('exclude_dims', _DEFAULT_FROZEN_SET)
dask = kwargs.pop('dask', 'forbidden')
output_dtypes = kwargs.pop('output_dtypes', None)
output_sizes = kwargs.pop('output_sizes', None)
keep_attrs = kwargs.pop('keep_attrs', False)
if kwargs:
raise TypeError('apply_variable_ufunc() got unexpected keyword '
'arguments: %s' % list(kwargs))

dim_sizes = unified_dim_sizes((a for a in args if hasattr(a, 'dims')),
exclude_dims=exclude_dims)
broadcast_dims = tuple(dim for dim in dim_sizes
Expand Down Expand Up @@ -661,14 +659,8 @@ def _apply_with_dask_atop(func, args, input_dims, output_dims, signature,
new_axes=output_sizes)


def apply_array_ufunc(func, *args, **kwargs):
"""apply_array_ufunc(func, *args, dask='forbidden')
"""
dask = kwargs.pop('dask', 'forbidden')
if kwargs:
raise TypeError('apply_array_ufunc() got unexpected keyword '
'arguments: %s' % list(kwargs))

def apply_array_ufunc(func, *args, dask='forbidden'):
"""Apply a ndarray level function over ndarray objects."""
if any(isinstance(arg, dask_array_type) for arg in args):
if dask == 'forbidden':
raise ValueError('apply_ufunc encountered a dask array on an '
Expand All @@ -687,23 +679,23 @@ def apply_array_ufunc(func, *args, **kwargs):
return func(*args)


def apply_ufunc(func, *args, **kwargs):
"""apply_ufunc(func : Callable,
*args : Any,
input_core_dims : Optional[Sequence[Sequence]] = None,
output_core_dims : Optional[Sequence[Sequence]] = ((),),
exclude_dims : Collection = frozenset(),
vectorize : bool = False,
join : str = 'exact',
dataset_join : str = 'exact',
dataset_fill_value : Any = _NO_FILL_VALUE,
keep_attrs : bool = False,
kwargs : Mapping = None,
dask : str = 'forbidden',
output_dtypes : Optional[Sequence] = None,
output_sizes : Optional[Mapping[Any, int]] = None)

Apply a vectorized function for unlabeled arrays on xarray objects.
def apply_ufunc(
func: Callable,
*args: Any,
input_core_dims: Optional[Sequence[Sequence]] = None,
output_core_dims: Optional[Sequence[Sequence]] = ((),),
exclude_dims: AbstractSet = frozenset(),
vectorize: bool = False,
join: str = 'exact',
dataset_join: str = 'exact',
dataset_fill_value: object = _NO_FILL_VALUE,
keep_attrs: bool = False,
kwargs: Mapping = None,
dask: str = 'forbidden',
output_dtypes: Optional[Sequence] = None,
output_sizes: Optional[Mapping[Any, int]] = None
) -> Any:
"""Apply a vectorized function for unlabeled arrays on xarray objects.

The function will be mapped over the data variable(s) of the input
arguments using xarray's standard rules for labeled computation, including
Expand Down Expand Up @@ -907,22 +899,6 @@ def earth_mover_distance(first_samples,
from .dataarray import DataArray
from .variable import Variable

input_core_dims = kwargs.pop('input_core_dims', None)
output_core_dims = kwargs.pop('output_core_dims', ((),))
vectorize = kwargs.pop('vectorize', False)
join = kwargs.pop('join', 'exact')
dataset_join = kwargs.pop('dataset_join', 'exact')
keep_attrs = kwargs.pop('keep_attrs', False)
exclude_dims = kwargs.pop('exclude_dims', frozenset())
dataset_fill_value = kwargs.pop('dataset_fill_value', _NO_FILL_VALUE)
kwargs_ = kwargs.pop('kwargs', None)
dask = kwargs.pop('dask', 'forbidden')
output_dtypes = kwargs.pop('output_dtypes', None)
output_sizes = kwargs.pop('output_sizes', None)
if kwargs:
raise TypeError('apply_ufunc() got unexpected keyword arguments: %s'
% list(kwargs))

if input_core_dims is None:
input_core_dims = ((),) * (len(args))
elif len(input_core_dims) != len(args):
Expand All @@ -931,14 +907,17 @@ def earth_mover_distance(first_samples,
'the number of arguments. Given input_core_dims: {}, '
'number of args: {}.'.format(input_core_dims, len(args)))

if kwargs is None:
kwargs = {}

signature = _UFuncSignature(input_core_dims, output_core_dims)

if exclude_dims and not exclude_dims <= signature.all_core_dims:
raise ValueError('each dimension in `exclude_dims` must also be a '
'core dimension in the function signature')

if kwargs_:
func = functools.partial(func, **kwargs_)
if kwargs:
func = functools.partial(func, **kwargs)

if vectorize:
if signature.all_core_dims:
Expand All @@ -950,14 +929,11 @@ def earth_mover_distance(first_samples,
'dimensions.')
func = np.vectorize(func,
otypes=output_dtypes,
signature=signature.to_gufunc_string(),
excluded=set(kwargs))
signature=signature.to_gufunc_string())
else:
func = np.vectorize(func,
otypes=output_dtypes,
excluded=set(kwargs))
func = np.vectorize(func, otypes=output_dtypes)

variables_ufunc = functools.partial(apply_variable_ufunc, func,
variables_vfunc = functools.partial(apply_variable_ufunc, func,
signature=signature,
exclude_dims=exclude_dims,
keep_attrs=keep_attrs,
Expand All @@ -966,7 +942,6 @@ def earth_mover_distance(first_samples,
output_sizes=output_sizes)

if any(isinstance(a, GroupBy) for a in args):
# kwargs has already been added into func
this_apply = functools.partial(apply_ufunc, func,
input_core_dims=input_core_dims,
output_core_dims=output_core_dims,
Expand All @@ -976,31 +951,29 @@ def earth_mover_distance(first_samples,
dataset_fill_value=dataset_fill_value,
keep_attrs=keep_attrs,
dask=dask)
return apply_groupby_ufunc(this_apply, *args)
return apply_groupby_func(this_apply, *args)
elif any(is_dict_like(a) for a in args):
return apply_dataset_ufunc(variables_ufunc, *args,
return apply_dataset_vfunc(variables_vfunc, *args,
signature=signature,
join=join,
exclude_dims=exclude_dims,
fill_value=dataset_fill_value,
dataset_join=dataset_join,
fill_value=dataset_fill_value,
keep_attrs=keep_attrs)
elif any(isinstance(a, DataArray) for a in args):
return apply_dataarray_ufunc(variables_ufunc, *args,
return apply_dataarray_vfunc(variables_vfunc, *args,
signature=signature,
join=join,
exclude_dims=exclude_dims,
keep_attrs=keep_attrs)
elif any(isinstance(a, Variable) for a in args):
return variables_ufunc(*args)
return variables_vfunc(*args)
else:
return apply_array_ufunc(func, *args, dask=dask)


def dot(*arrays, **kwargs):
""" dot(*arrays, dims=None)

Generalized dot product for xarray objects. Like np.einsum, but
def dot(*arrays, dims=None, **kwargs):
"""Generalized dot product for xarray objects. Like np.einsum, but
provides a simpler interface based on array dimensions.

Parameters
Expand Down Expand Up @@ -1036,8 +1009,6 @@ def dot(*arrays, **kwargs):
from .dataarray import DataArray
from .variable import Variable

dims = kwargs.pop('dims', None)

if any(not isinstance(arr, (Variable, DataArray)) for arr in arrays):
raise TypeError('Only xr.DataArray and xr.Variable are supported.'
'Given {}.'.format([type(arr) for arr in arrays]))
Expand Down