Python computation.apply_ufunc函数代码示例

def test_apply_exclude():

    def concatenate(objects, dim='x'):
        def func(*x):
            return np.concatenate(x, axis=-1)
        result = apply_ufunc(func, *objects,
                             input_core_dims=[[dim]] * len(objects),
                             output_core_dims=[[dim]],
                             exclude_dims={dim})
        if isinstance(result, (xr.Dataset, xr.DataArray)):
            # note: this will fail if dim is not a coordinate on any input
            new_coord = np.concatenate([obj.coords[dim] for obj in objects])
            result.coords[dim] = new_coord
        return result

    arrays = [np.array([1]), np.array([2, 3])]
    variables = [xr.Variable('x', a) for a in arrays]
    data_arrays = [xr.DataArray(v, {'x': c, 'y': ('x', range(len(c)))})
                   for v, c in zip(variables, [['a'], ['b', 'c']])]
    datasets = [xr.Dataset({'data': data_array}) for data_array in data_arrays]

    expected_array = np.array([1, 2, 3])
    expected_variable = xr.Variable('x', expected_array)
    expected_data_array = xr.DataArray(expected_variable, [('x', list('abc'))])
    expected_dataset = xr.Dataset({'data': expected_data_array})

    assert_identical(expected_array, concatenate(arrays))
    assert_identical(expected_variable, concatenate(variables))
    assert_identical(expected_data_array, concatenate(data_arrays))
    assert_identical(expected_dataset, concatenate(datasets))

    # must also be a core dimension
    with pytest.raises(ValueError):
        apply_ufunc(identity, variables[0], exclude_dims={'x'})

def test_apply_exclude():

    def concatenate(objects, dim='x'):
        sig = ([(dim,)] * len(objects), [(dim,)])
        new_coord = np.concatenate(
            [obj.coords[dim] if hasattr(obj, 'coords') else []
             for obj in objects])
        func = lambda *x: np.concatenate(x, axis=-1)
        result = apply_ufunc(func, *objects, signature=sig, exclude_dims={dim})
        if isinstance(result, (xr.Dataset, xr.DataArray)):
            result.coords[dim] = new_coord
        return result

    arrays = [np.array([1]), np.array([2, 3])]
    variables = [xr.Variable('x', a) for a in arrays]
    data_arrays = [xr.DataArray(v, {'x': c, 'y': ('x', range(len(c)))})
                   for v, c in zip(variables, [['a'], ['b', 'c']])]
    datasets = [xr.Dataset({'data': data_array}) for data_array in data_arrays]

    expected_array = np.array([1, 2, 3])
    expected_variable = xr.Variable('x', expected_array)
    expected_data_array = xr.DataArray(expected_variable, [('x', list('abc'))])
    expected_dataset = xr.Dataset({'data': expected_data_array})

    assert_identical(expected_array, concatenate(arrays))
    assert_identical(expected_variable, concatenate(variables))
    assert_identical(expected_data_array, concatenate(data_arrays))
    assert_identical(expected_dataset, concatenate(datasets))

    identity = lambda x: x
    # must also be a core dimension
    with pytest.raises(ValueError):
        apply_ufunc(identity, variables[0], exclude_dims={'x'})

def test_dataset_join():
    import numpy as np
    ds0 = xr.Dataset({'a': ('x', [1, 2]), 'x': [0, 1]})
    ds1 = xr.Dataset({'a': ('x', [99, 3]), 'x': [1, 2]})

    with pytest.raises(TypeError):
        apply_ufunc(operator.add, ds0, ds1, dataset_join='outer')

    def add(a, b, join, dataset_join):
        return apply_ufunc(operator.add, a, b, join=join,
                           dataset_join=dataset_join,
                           dataset_fill_value=np.nan)

    actual = add(ds0, ds1, 'outer', 'inner')
    expected = xr.Dataset({'a': ('x', [np.nan, 101, np.nan]),
                           'x': [0, 1, 2]})
    assert_identical(actual, expected)

    actual = add(ds0, ds1, 'outer', 'outer')
    assert_identical(actual, expected)

    # if variables don't match, join will perform add with np.nan
    ds2 = xr.Dataset({'b': ('x', [99, 3]), 'x': [1, 2]})
    actual = add(ds0, ds2, 'outer', 'inner')
    expected = xr.Dataset({'x': [0, 1, 2]})
    assert_identical(actual, expected)

    actual = add(ds0, ds2, 'outer', 'outer')
    expected = xr.Dataset({'a': ('x', [np.nan, np.nan, np.nan]),
                           'b': ('x', [np.nan, np.nan, np.nan]),
                           'x': [0, 1, 2]})
    assert_identical(actual, expected)

def test_apply_dask_multiple_inputs():
    import dask.array as da

    def covariance(x, y):
        return ((x - x.mean(axis=-1, keepdims=True)) *
                (y - y.mean(axis=-1, keepdims=True))).mean(axis=-1)

    rs = np.random.RandomState(42)
    array1 = da.from_array(rs.randn(4, 4), chunks=(2, 4))
    array2 = da.from_array(rs.randn(4, 4), chunks=(2, 4))
    data_array_1 = xr.DataArray(array1, dims=('x', 'z'))
    data_array_2 = xr.DataArray(array2, dims=('y', 'z'))

    expected = apply_ufunc(
        covariance, data_array_1.compute(), data_array_2.compute(),
        input_core_dims=[['z'], ['z']])
    allowed = apply_ufunc(
        covariance, data_array_1, data_array_2, input_core_dims=[['z'], ['z']],
        dask='allowed')
    assert isinstance(allowed.data, da.Array)
    xr.testing.assert_allclose(expected, allowed.compute())

    parallelized = apply_ufunc(
        covariance, data_array_1, data_array_2, input_core_dims=[['z'], ['z']],
        dask='parallelized', output_dtypes=[float])
    assert isinstance(parallelized.data, da.Array)
    xr.testing.assert_allclose(expected, parallelized.compute())

def xr_polyfit(obj, dim, ix=None, deg=0.5, poly='hermite'):
    """Fit a polynomial of degree ``deg`` using least-squares along ``dim``.

    Parameters
    ----------
    obj : xarray.Dataset or xarray.DataArray
        The object to fit.
    dim : str, optional
        The dimension to fit along.
    ix : {None, int, array_like}, optional
        If ``None``, interpolate the polynomial at the original x points.
        If ``int``, linearly space this many points along the range of the
        original data and interpolate with these.
        If array-like, interpolate at these given points.
    deg : int or float, optional
        The degree of the polynomial to fit. Used directly if integer. If float
        supplied, with ``0.0 < deg < 1.0``, the proportion of the total
        possible degree to use.
    poly : {'chebyshev', 'polynomial', 'legendre',
            'laguerre', hermite}, optional
        The type of polynomial to fit.

    Returns
    -------
    new_xobj : xarray.DataArray or xarray.Dataset

    See Also
    --------
    xr_unispline
    """
    input_core_dims = [(dim,), (dim,)]
    args = (obj[dim], obj)

    if ix is None:
        kwargs = {'ix': ix, 'axis': -1, 'deg': deg, 'poly': poly}
        output_core_dims = [(dim,)]
        return apply_ufunc(_broadcast_polyfit, *args, kwargs=kwargs,
                           input_core_dims=input_core_dims,
                           output_core_dims=output_core_dims)

    if isinstance(ix, int):
        ix = np.linspace(float(obj[dim].min()), float(obj[dim].max()), ix)

    kwargs = {'ix': ix, 'axis': -1, 'deg': deg, 'poly': poly}
    output_core_dims = [('__temp_dim__',)]

    result = apply_ufunc(_broadcast_polyfit, *args, kwargs=kwargs,
                         input_core_dims=input_core_dims,
                         output_core_dims=output_core_dims)
    result['__temp_dim__'] = ix
    return result.rename({'__temp_dim__': dim})

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_apply_input_core_dimension():

    def first_element(obj, dim):
        def func(x):
            return x[..., 0]
        return apply_ufunc(func, obj, input_core_dims=[[dim]])

    array = np.array([[1, 2], [3, 4]])
    variable = xr.Variable(['x', 'y'], array)
    data_array = xr.DataArray(variable, {'x': ['a', 'b'], 'y': [-1, -2]})
    dataset = xr.Dataset({'data': data_array})

    expected_variable_x = xr.Variable(['y'], [1, 2])
    expected_data_array_x = xr.DataArray(expected_variable_x, {'y': [-1, -2]})
    expected_dataset_x = xr.Dataset({'data': expected_data_array_x})

    expected_variable_y = xr.Variable(['x'], [1, 3])
    expected_data_array_y = xr.DataArray(expected_variable_y,
                                         {'x': ['a', 'b']})
    expected_dataset_y = xr.Dataset({'data': expected_data_array_y})

    assert_identical(expected_variable_x, first_element(variable, 'x'))
    assert_identical(expected_variable_y, first_element(variable, 'y'))

    assert_identical(expected_data_array_x, first_element(data_array, 'x'))
    assert_identical(expected_data_array_y, first_element(data_array, 'y'))

    assert_identical(expected_dataset_x, first_element(dataset, 'x'))
    assert_identical(expected_dataset_y, first_element(dataset, 'y'))

    assert_identical(expected_data_array_x,
                     first_element(data_array.groupby('y'), 'x'))
    assert_identical(expected_dataset_x,
                     first_element(dataset.groupby('y'), 'x'))

    def multiply(*args):
        val = args[0]
        for arg in args[1:]:
            val = val * arg
        return val

    # regression test for GH:2341
    with pytest.raises(ValueError):
        apply_ufunc(multiply, data_array, data_array['y'].values,
                    input_core_dims=[['y']], output_core_dims=[['y']])
    expected = xr.DataArray(multiply(data_array, data_array['y']),
                            dims=['x', 'y'], coords=data_array.coords)
    actual = apply_ufunc(multiply, data_array, data_array['y'].values,
                         input_core_dims=[['y'], []], output_core_dims=[['y']])
    assert_identical(expected, actual)

 def original_and_stack_negative(obj):
     func = lambda x: (x, xr.core.npcompat.stack([x, -x], axis=-1))
     sig = ([()], [(), ('sign',)])
     result = apply_ufunc(func, obj, signature=sig)
     if isinstance(result[1], (xr.Dataset, xr.DataArray)):
         result[1].coords['sign'] = [1, -1]
     return result

def xr_diff_u_err(obj, dim):
    """Propagate error through uneven-third-order finite difference derivative.
    If you have calculated a derivative already using ``xr_diff_u``, and you
    have data about the uncertainty on the original data, this function
    propagates that error through to be an error on the derivative.

    Parameters
    ----------
    obj : xarray.Dataset or xarray.DataArray
        The object to differentiate.
    dim : str
        The dimension to differentiate along.

    Returns
    -------
    new_xobj : xarray.DataArray or xarray.Dataset
    """
    kwargs = {'axis': -1}
    input_core_dims = [(dim,), (dim,)]
    output_core_dims = [(dim,)]
    args = (obj[dim], obj)
    return apply_ufunc(_broadcast_diff_u_err, *args,
                       input_core_dims=input_core_dims,
                       output_core_dims=output_core_dims,
                       kwargs=kwargs)

 def stack_negative(obj):
     def func(x):
         return np.stack([x, -x], axis=-1)
     result = apply_ufunc(func, obj, output_core_dims=[['sign']])
     if isinstance(result, (xr.Dataset, xr.DataArray)):
         result.coords['sign'] = [1, -1]
     return result

def xr_diff_u(obj, dim):
    """Uneven-third-order finite difference derivative [1].

    [1] Singh, Ashok K., and B. S. Bhadauria. "Finite difference formulae for
    unequal sub-intervals using lagrange’s interpolation formula."
    International Journal of Mathematics and Analysis 3.17 (2009): 815-827.

    Parameters
    ----------
    obj : xarray.Dataset or xarray.DataArray
        The object to differentiate.
    dim : str
        The dimension to differentiate along.

    Returns
    -------
    new_xobj : xarray.DataArray or xarray.Dataset
    """
    kwargs = {'axis': -1}
    input_core_dims = [(dim,), (dim,)]
    output_core_dims = [(dim,)]
    args = (obj[dim], obj)
    return apply_ufunc(_broadcast_diff_u, *args,
                       input_core_dims=input_core_dims,
                       output_core_dims=output_core_dims,
                       kwargs=kwargs)

 def original_and_stack_negative(obj):
     def func(x):
         return (x, np.stack([x, -x], axis=-1))
     result = apply_ufunc(func, obj, output_core_dims=[[], ['sign']])
     if isinstance(result[1], (xr.Dataset, xr.DataArray)):
         result[1].coords['sign'] = [1, -1]
     return result

def test_vectorize():
    data_array = xr.DataArray([[0, 1, 2], [1, 2, 3]], dims=('x', 'y'))
    expected = xr.DataArray([1, 2], dims=['x'])
    actual = apply_ufunc(pandas_median, data_array,
                         input_core_dims=[['y']],
                         vectorize=True)
    assert_identical(expected, actual)

def xr_interp(obj, dim, ix=100, order=3):
    """Interpolate along axis ``dim`` using :func:`scipy.interpolate.interp1d`.

    Parameters
    ----------
    obj : xarray.Dataset or xarray.DataArray
        The object to interpolate.
    dim : str
        The axis to interpolate along.
    ix : int or array
        If int, interpolate to this many points spaced evenly along the range
        of the original data. If array, interpolate to those points directly.
    order : int
        Supplied to :func:`scipy.interpolate.interp1d` as the order of
        interpolation.

    Returns
    -------
    new_xobj : xarray.DataArray or xarray.Dataset

    See Also
    --------
    xr_pchip
    """

    input_core_dims = [(dim,), (dim,)]
    args = (obj[dim], obj)
    kwargs = {'ix': ix, 'axis': -1, 'order': order}

    if ix is None:
        output_core_dims = [(dim,)]
        return apply_ufunc(_broadcast_interp, *args, kwargs=kwargs,
                           input_core_dims=input_core_dims,
                           output_core_dims=output_core_dims)

    if isinstance(ix, int):
        ix = np.linspace(float(obj[dim].min()), float(obj[dim].max()), ix)

    kwargs['ix'] = ix
    output_core_dims = [('__temp_dim__',)]

    result = apply_ufunc(_broadcast_interp, *args, kwargs=kwargs,
                         input_core_dims=input_core_dims,
                         output_core_dims=output_core_dims)
    result['__temp_dim__'] = ix
    return result.rename({'__temp_dim__': dim})

def xr_filter_wiener(obj, dim, mysize=5, noise=1e-2):
    kwargs = {'mysize': mysize, 'noise': noise, 'axis': -1}
    input_core_dims = [(dim,), (dim,)]
    output_core_dims = [(dim,)]
    args = (obj[dim], obj)
    return apply_ufunc(_broadcast_filter_wiener, *args,
                       input_core_dims=input_core_dims,
                       output_core_dims=output_core_dims,
                       kwargs=kwargs)