Python xarray.apply_ufunc方法代码示例

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def _determine_input_core_dims(dim, weights):
    """
    Determine input_core_dims based on type of dim and weights.

    Parameters
    ----------
    dim : str, list
        The dimension(s) to apply the correlation along.
    weights : xarray.Dataset or xarray.DataArray or None
        Weights matching dimensions of ``dim`` to apply during the function.

    Returns
    -------
    list of lists
        input_core_dims used for xr.apply_ufunc.
    """
    if not isinstance(dim, list):
        dim = [dim]
    # build input_core_dims depending on weights
    if weights is None:
        input_core_dims = [dim, dim, [None]]
    else:
        input_core_dims = [dim, dim, dim]
    return input_core_dims 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def windowed_run_events_ufunc(x: Sequence[bool], window: int) -> xr.apply_ufunc:
    """Dask-parallel version of windowed_run_events_1d, ie the number of runs at least as long as given duration.

    Parameters
    ----------
    x : Sequence[bool]
      Input array (bool)
    window : int
      Minimum run length

    Returns
    -------
    out : func
      A function operating along the time dimension of a dask-array.
    """
    return xr.apply_ufunc(
        windowed_run_events_1d,
        x,
        input_core_dims=[["time"]],
        vectorize=True,
        dask="parallelized",
        output_dtypes=[np.int],
        keep_attrs=True,
        kwargs={"window": window},
    ) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def longest_run_ufunc(x: Sequence[bool]) -> xr.apply_ufunc:
    """Dask-parallel version of longest_run_1d, ie the maximum number of consecutive true values in
    array.

    Parameters
    ----------
    x : Sequence[bool]
      Input array (bool)

    Returns
    -------
    out : func
      A function operating along the time dimension of a dask-array.
    """
    return xr.apply_ufunc(
        longest_run_1d,
        x,
        input_core_dims=[["time"]],
        vectorize=True,
        dask="parallelized",
        output_dtypes=[np.int],
        keep_attrs=True,
    ) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def psd_fft(varr):
    _T = len(varr.coords['frame'])
    ns = _T // 2 + 1
    if _T % 2 == 0:
        freq_crd = np.linspace(0, 0.5, ns)
    else:
        freq_crd = np.linspace(0, 0.5 * (_T - 1) / _T, ns)
    print("computing psd of input")
    varr_fft = xr.apply_ufunc(
        fftw.rfft,
        varr.chunk(dict(frame=-1)),
        input_core_dims=[['frame']],
        output_core_dims=[['freq']],
        dask='allowed',
        output_sizes=dict(freq=ns),
        output_dtypes=[np.complex_])
    varr_fft = varr_fft.assign_coords(freq=freq_crd)
    varr_psd = 1 / _T * np.abs(varr_fft)**2
    return varr_psd 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def psd_welch(varr):
    _T = len(varr.coords['frame'])
    ns = _T // 2 + 1
    if _T % 2 == 0:
        freq_crd = np.linspace(0, 0.5, ns)
    else:
        freq_crd = np.linspace(0, 0.5 * (_T - 1) / _T, ns)
    varr_psd = xr.apply_ufunc(
            _welch,
            varr.chunk(dict(frame=-1)),
            input_core_dims=[['frame']],
            output_core_dims=[['freq']],
            dask='parallelized',
            vectorize=True,
            kwargs=dict(nperseg=_T),
            output_sizes=dict(freq=ns),
            output_dtypes=[varr.dtype])
    varr_psd = varr_psd.assign_coords(freq=freq_crd)
    return varr_psd 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def get_noise_welch(varr,
                    noise_range=(0.25, 0.5),
                    noise_method='logmexp',
                    compute=True):
    print("estimating noise")
    sn = xr.apply_ufunc(
        noise_welch,
        varr.chunk(dict(frame=-1)),
        input_core_dims=[['frame']],
        dask='parallelized',
        vectorize=True,
        kwargs=dict(noise_range=noise_range, noise_method=noise_method),
        output_dtypes=[varr.dtype])
    if compute:
        sn = sn.compute()
    return sn 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def ks_refine(varr, seeds, sig=0.05):
    print("selecting seeds")
    varr_sub = varr.sel(
        spatial=[tuple(hw) for hw in seeds[['height', 'width']].values])
    print("performing KS test")
    ks = xr.apply_ufunc(
        lambda x: kstest(zscore(x), 'norm')[1],
        varr_sub.chunk(dict(frame=-1, spatial='auto')),
        input_core_dims=[['frame']],
        vectorize=True,
        dask='parallelized',
        output_dtypes=[float])
    mask = ks < sig
    mask_df = mask.to_pandas().rename('mask_ks').reset_index()
    seeds = pd.merge(seeds, mask_df, on=['height', 'width'], how='left')
    return seeds 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def __eq__(self, other):
        if isinstance(other, Longitude):
            return (self.hemisphere == other.hemisphere and
                    self.longitude == other.longitude)
        else:
            return xr.apply_ufunc(np.equal, other, self) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def __lt__(self, other):
        if isinstance(other, Longitude):
            if self.hemisphere == 'W':
                if other.hemisphere == 'E':
                    return True
                else:
                    return self.longitude > other.longitude
            else:
                if other.hemisphere == 'W':
                    return False
                else:
                    return self.longitude < other.longitude
        else:
            return xr.apply_ufunc(np.greater, other, self) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def __gt__(self, other):
        if isinstance(other, Longitude):
            if self.hemisphere == 'W':
                if other.hemisphere == 'E':
                    return False
                else:
                    return self.longitude < other.longitude
            else:
                if other.hemisphere == 'W':
                    return True
                else:
                    return self.longitude > other.longitude
        else:
            return xr.apply_ufunc(np.less, other, self) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def __le__(self, other):
        if isinstance(other, Longitude):
            return self < other or self == other
        else:
            return xr.apply_ufunc(np.greater_equal, other, self) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def __ge__(self, other):
        if isinstance(other, Longitude):
            return self > other or self == other
        else:
            return xr.apply_ufunc(np.less_equal, other, self) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def xr_pearson_r(x, y, dim):
    """pearson_r implementation using xarray and minimal numpy only."""
    return xr.apply_ufunc(
        pearson_correlation_gufunc,
        x,
        y,
        input_core_dims=[[dim], [dim]],
        dask='parallelized',
        output_dtypes=[float],
    ) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def median_absolute_error(a, b, dim, skipna=False):
    """
    Median Absolute Error.

    Parameters
    ----------
    a : xarray.Dataset or xarray.DataArray
        Labeled array(s) over which to apply the function.
    b : xarray.Dataset or xarray.DataArray
        Labeled array(s) over which to apply the function.
    dim : str, list
        The dimension(s) to apply the median absolute error along.
    skipna : bool
        If True, skip NaNs when computing function.

    Returns
    -------
    xarray.Dataset or xarray.DataArray
        Median Absolute Error.

    See Also
    --------
    sklearn.metrics.median_absolute_error
    xarray.apply_ufunc
    xskillscore.core.np_deterministic._median_absolute_error

    """
    dim, axis = _preprocess_dims(dim)

    return xr.apply_ufunc(
        _median_absolute_error,
        a,
        b,
        input_core_dims=[dim, dim],
        kwargs={'axis': axis, 'skipna': skipna},
        dask='parallelized',
        output_dtypes=[float],
    ) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import apply_ufunc [as 别名]
def xr_crps_gaussian(observations, mu, sig):
    """
    xarray version of properscoring.crps_gaussian: Continuous Ranked
     Probability Score with a Gaussian distribution.
    Parameters
    ----------
    observations : xarray.Dataset or xarray.DataArray
        The observations or set of observations.
    mu : xarray.Dataset or xarray.DataArray
        The mean of the forecast normal distribution.
    sig : xarray.Dataset or xarray.DataArray
        The standard deviation of the forecast distribution.
    Returns
    -------
    xarray.Dataset or xarray.DataArray
    See Also
    --------
    properscoring.crps_gaussian
    xarray.apply_ufunc
    """
    # check if same dimensions
    if isinstance(mu, (int, float)):
        mu = xr.DataArray(mu)
    if isinstance(sig, (int, float)):
        sig = xr.DataArray(sig)
    if mu.dims != observations.dims:
        observations, mu = xr.broadcast(observations, mu)
    if sig.dims != observations.dims:
        observations, sig = xr.broadcast(observations, sig)
    return xr.apply_ufunc(
        crps_gaussian,
        observations,
        mu,
        sig,
        input_core_dims=[[], [], []],
        dask='parallelized',
        output_dtypes=[float],
    )