Python xarray.DataArray() Examples

def ds_time_encoded_cf():
    time_bounds = np.array([[0, 31], [31, 59], [59, 90]])
    bounds = np.array([0, 1])
    time = np.array([15, 46, 74])
    data = np.zeros((3))
    ds = xr.DataArray(data,
                      coords=[time],
                      dims=[TIME_STR],
                      name='a').to_dataset()
    ds[TIME_BOUNDS_STR] = xr.DataArray(time_bounds,
                                       coords=[time, bounds],
                                       dims=[TIME_STR, BOUNDS_STR],
                                       name=TIME_BOUNDS_STR)
    units_str = 'days since 2000-01-01 00:00:00'
    cal_str = 'noleap'
    ds[TIME_STR].attrs['units'] = units_str
    ds[TIME_STR].attrs['calendar'] = cal_str
    return ds 

def test_add_uniform_time_weights():
    time = np.array([15, 46, 74])
    data = np.zeros((3))
    ds = xr.DataArray(data,
                      coords=[time],
                      dims=[TIME_STR],
                      name='a').to_dataset()
    units_str = 'days since 2000-01-01 00:00:00'
    cal_str = 'noleap'
    ds[TIME_STR].attrs['units'] = units_str
    ds[TIME_STR].attrs['calendar'] = cal_str

    with pytest.raises(KeyError):
        ds[TIME_WEIGHTS_STR]

    ds = add_uniform_time_weights(ds)
    time_weights_expected = xr.DataArray(
        [1, 1, 1], coords=ds[TIME_STR].coords, name=TIME_WEIGHTS_STR)
    time_weights_expected.attrs['units'] = 'days'
    assert ds[TIME_WEIGHTS_STR].identical(time_weights_expected) 

def test_dft_real_2d(self):
        """
        Test the real discrete Fourier transform function on one-dimensional
        data. Non-trivial because we need to keep only some of the negative
        frequencies.
        """
        Nx, Ny = 16, 32
        da = xr.DataArray(np.random.rand(Nx, Ny), dims=['x', 'y'],
                          coords={'x': range(Nx), 'y': range(Ny)})
        dx = float(da.x[1] - da.x[0])
        dy = float(da.y[1] - da.y[0])

        daft = xrft.dft(da, real='x')
        npt.assert_almost_equal(daft.values,
                               np.fft.rfftn(da.transpose('y','x')).transpose())
        npt.assert_almost_equal(daft.values,
                               xrft.dft(da, dim=['y'], real='x'))

        actual_freq_x = daft.coords['freq_x'].values
        expected_freq_x = np.fft.rfftfreq(Nx, dx)
        npt.assert_almost_equal(actual_freq_x, expected_freq_x)

        actual_freq_y = daft.coords['freq_y'].values
        expected_freq_y = np.fft.fftfreq(Ny, dy)
        npt.assert_almost_equal(actual_freq_y, expected_freq_y) 

def test_dft_4d(self):
        """Test the discrete Fourier transform on 2D data"""
        N = 16
        da = xr.DataArray(np.random.rand(N,N,N,N),
                         dims=['time','z','y','x'],
                         coords={'time':range(N),'z':range(N),
                                'y':range(N),'x':range(N)}
                         )
        with pytest.raises(ValueError):
            xrft.dft(da.chunk({'time':8}), dim=['y','x'], detrend='linear')
        ft = xrft.dft(da, shift=False)
        npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))

        da_prime = xrft.detrendn(da[:,0].values, [0,1,2]) # cubic detrend over time, y, and x
        npt.assert_almost_equal(xrft.dft(da[:,0].drop('z'),
                                        dim=['time','y','x'],
                                        shift=False, detrend='linear'
                                        ).values,
                                np.fft.fftn(da_prime)) 

def test_cross_phase_2d(self, dask):
        Ny, Nx = (32, 16)
        x = np.linspace(0, 1, num=Nx, endpoint=False)
        y = np.ones(Ny)
        f = 6
        phase_offset = np.pi/2
        signal1 = np.cos(2*np.pi*f*x)  # frequency = 1/(2*pi)
        signal2 = np.cos(2*np.pi*f*x - phase_offset)
        da1 = xr.DataArray(data=signal1*y[:,np.newaxis], name='a',
                          dims=['y','x'], coords={'y':y, 'x':x})
        da2 = xr.DataArray(data=signal2*y[:,np.newaxis], name='b',
                          dims=['y','x'], coords={'y':y, 'x':x})
        with pytest.raises(ValueError):
            xrft.cross_phase(da1, da2, dim=['y','x'])

        if dask:
            da1 = da1.chunk({'x': 16})
            da2 = da2.chunk({'x': 16})
        cp = xrft.cross_phase(da1, da2, dim=['x'])
        actual_phase_offset = cp.sel(freq_x=f).values
        npt.assert_almost_equal(actual_phase_offset, phase_offset) 

def test_ensure_time_as_index_with_change():
    # Time bounds array doesn't index time initially, which gets fixed.
    arr = xr.DataArray([-93], dims=[TIME_STR], coords={TIME_STR: [3]})
    arr[TIME_STR].attrs['units'] = 'days since 2000-01-01 00:00:00'
    arr[TIME_STR].attrs['calendar'] = 'standard'
    ds = arr.to_dataset(name='a')
    ds.coords[TIME_WEIGHTS_STR] = xr.DataArray(
        [1], dims=[TIME_STR], coords={TIME_STR: arr[TIME_STR]}
    )
    ds.coords[TIME_BOUNDS_STR] = xr.DataArray(
        [[3.5, 4.5]], dims=[TIME_STR, BOUNDS_STR],
        coords={TIME_STR: arr[TIME_STR]}
    )
    ds = ds.isel(**{TIME_STR: 0})
    actual = ensure_time_as_index(ds)
    expected = arr.to_dataset(name='a')
    expected.coords[TIME_WEIGHTS_STR] = xr.DataArray(
        [1], dims=[TIME_STR], coords={TIME_STR: arr[TIME_STR]}
    )
    expected.coords[TIME_BOUNDS_STR] = xr.DataArray(
        [[3.5, 4.5]], dims=[TIME_STR, BOUNDS_STR],
        coords={TIME_STR: arr[TIME_STR]}
        )
    xr.testing.assert_identical(actual, expected) 

def test_yearly_average_no_mask():
    times = pd.to_datetime(['2000-06-01', '2000-06-15',
                            '2001-07-04', '2001-10-01', '2001-12-31',
                            '2004-01-01'])
    arr = xr.DataArray(np.random.random((len(times),)),
                       dims=[TIME_STR], coords={TIME_STR: times})
    dt = arr.copy(deep=True)
    dt.values = np.random.random((len(times),))

    actual = yearly_average(arr, dt)

    yr2000 = (arr[0]*dt[0] + arr[1]*dt[1]) / (dt[0] + dt[1])
    yr2001 = ((arr[2]*dt[2] + arr[3]*dt[3] + arr[4]*dt[4]) /
              (dt[2] + dt[3] + dt[4]))
    yr2004 = arr[-1]
    yrs_coord = [2000, 2001, 2004]
    yr_avgs = np.array([yr2000, yr2001, yr2004])
    desired = xr.DataArray(yr_avgs, dims=['year'], coords={'year': yrs_coord})
    xr.testing.assert_allclose(actual, desired) 

def _maybe_cast_to_float64(da):
    """Cast DataArrays to np.float64 if they are of type np.float32.

    Parameters
    ----------
    da : xr.DataArray
        Input DataArray

    Returns
    -------
    DataArray

    """
    if da.dtype == np.float32:
        logging.warning('Datapoints were stored using the np.float32 datatype.'
                        'For accurate reduction operations using bottleneck, '
                        'datapoints are being cast to the np.float64 datatype.'
                        ' For more information see: https://github.com/pydata/'
                        'xarray/issues/1346')
        return da.astype(np.float64)
    else:
        return da 

def test_yearly_average_masked_data():
    times = pd.to_datetime(['2000-06-01', '2000-06-15',
                            '2001-07-04', '2001-10-01', '2001-12-31',
                            '2004-01-01'])
    arr = xr.DataArray(np.random.random((len(times),)),
                       dims=[TIME_STR], coords={TIME_STR: times})
    arr[0] = -999
    arr = arr.where(arr != -999)
    dt = arr.copy(deep=True)
    dt.values = np.random.random((len(times),))

    actual = yearly_average(arr, dt)

    yr2000 = arr[1]
    yr2001 = ((arr[2]*dt[2] + arr[3]*dt[3] + arr[4]*dt[4]) /
              (dt[2] + dt[3] + dt[4]))
    yr2004 = arr[-1]
    yrs_coord = [2000, 2001, 2004]
    yr_avgs = np.array([yr2000, yr2001, yr2004])
    desired = xr.DataArray(yr_avgs, dims=['year'], coords={'year': yrs_coord})
    xr.testing.assert_allclose(actual, desired) 

def test_spacing_tol(test_data_1d):
    da = test_data_1d
    da2 = da.copy().load()

    # Create improperly spaced data
    Nx = 16
    Lx = 1.0
    x  = np.linspace(0, Lx, Nx)
    x[-1] = x[-1] + .001
    da3 = xr.DataArray(np.random.rand(Nx), coords=[x], dims=['x'])

    # This shouldn't raise an error
    xrft.dft(da3, spacing_tol=1e-1)
    # But this should
    with pytest.raises(ValueError):
        xrft.dft(da3, spacing_tol=1e-4) 

def _apply_window(da, dims, window_type='hanning'):
    """Creating windows in dimensions dims."""

    if window_type not in ['hanning']:
        raise NotImplementedError("Only hanning window is supported for now.")

    numpy_win_func = getattr(np, window_type)

    if da.chunks:
        def dask_win_func(n):
            return dsar.from_delayed(
                delayed(numpy_win_func, pure=True)(n),
                (n,), float)
        win_func = dask_win_func
    else:
        win_func = numpy_win_func

    windows = [xr.DataArray(win_func(len(da[d])),
               dims=da[d].dims, coords=da[d].coords) for d in dims]

    return da * reduce(operator.mul, windows[::-1]) 

def _bounds_from_array(arr, dim_name, bounds_name):
    """Get the bounds of an array given its center values.

    E.g. if lat-lon grid center lat/lon values are known, but not the
    bounds of each grid box.  The algorithm assumes that the bounds
    are simply halfway between each pair of center values.
    """
    # TODO: don't assume needed dimension is in axis=0
    # TODO: refactor to get rid of repetitive code
    spacing = arr.diff(dim_name).values
    lower = xr.DataArray(np.empty_like(arr), dims=arr.dims,
                         coords=arr.coords)
    lower.values[:-1] = arr.values[:-1] - 0.5*spacing
    lower.values[-1] = arr.values[-1] - 0.5*spacing[-1]
    upper = xr.DataArray(np.empty_like(arr), dims=arr.dims,
                         coords=arr.coords)
    upper.values[:-1] = arr.values[:-1] + 0.5*spacing
    upper.values[-1] = arr.values[-1] + 0.5*spacing[-1]
    bounds = xr.concat([lower, upper], dim='bounds')
    return bounds.T 

def assert_matching_time_coord(arr1, arr2):
    """Check to see if two DataArrays have the same time coordinate.

    Parameters
    ----------
    arr1 : DataArray or Dataset
        First DataArray or Dataset
    arr2 : DataArray or Dataset
        Second DataArray or Dataset

    Raises
    ------
    ValueError
        If the time coordinates are not identical between the two Datasets
    """
    message = ('Time weights not indexed by the same time coordinate as'
               ' computed data.  This will lead to an improperly computed'
               ' time weighted average.  Exiting.\n'
               'arr1: {}\narr2: {}')
    if not (arr1[TIME_STR].identical(arr2[TIME_STR])):
        raise ValueError(message.format(arr1[TIME_STR], arr2[TIME_STR])) 

def test_dft_2d(self):
        """Test the discrete Fourier transform on 2D data"""
        N = 16
        da = xr.DataArray(np.random.rand(N,N), dims=['x','y'],
                        coords={'x':range(N),'y':range(N)}
                         )
        ft = xrft.dft(da, shift=False)
        npt.assert_almost_equal(ft.values, np.fft.fftn(da.values))

        ft = xrft.dft(da, shift=False, window=True, detrend='constant')
        dim = da.dims
        window = np.hanning(N) * np.hanning(N)[:, np.newaxis]
        da_prime = (da - da.mean(dim=dim)).values
        npt.assert_almost_equal(ft.values, np.fft.fftn(da_prime*window))

        da = xr.DataArray(np.random.rand(N,N), dims=['x','y'],
                         coords={'x':range(N,0,-1),'y':range(N,0,-1)}
                         )
        assert (xrft.power_spectrum(da, shift=False,
                                   density=True) >= 0.).all() 

def test_sel_var():
    time = np.array([0, 31, 59]) + 15
    data = np.zeros((3))
    ds = xr.DataArray(data,
                      coords=[time],
                      dims=[TIME_STR],
                      name=convection_rain.name).to_dataset()
    condensation_rain_alt_name, = condensation_rain.alt_names
    ds[condensation_rain_alt_name] = xr.DataArray(data, coords=[ds.time])
    result = _sel_var(ds, convection_rain)
    assert result.name == convection_rain.name

    result = _sel_var(ds, condensation_rain)
    assert result.name == condensation_rain.name

    with pytest.raises(LookupError):
        _sel_var(ds, precip) 

def extract_months(time, months):
    """Extract times within specified months of the year.

    Parameters
    ----------
    time : xarray.DataArray
         Array of times that can be represented by numpy.datetime64 objects
         (i.e. the year is between 1678 and 2262).
    months : Desired months of the year to include

    Returns
    -------
    xarray.DataArray of the desired times
    """
    inds = _month_conditional(time, months)
    return time.sel(time=inds) 

def ds_with_time_bounds(alt_lat_str, var_name):
    time_bounds = np.array([[0, 31], [31, 59], [59, 90]])
    bounds = np.array([0, 1])
    time = np.array([15, 46, 74])
    data = np.zeros((3, 1, 1))
    lat = [0]
    lon = [0]
    ds = xr.DataArray(data,
                      coords=[time, lat, lon],
                      dims=[TIME_STR, alt_lat_str, LON_STR],
                      name=var_name).to_dataset()
    ds[TIME_BOUNDS_STR] = xr.DataArray(time_bounds,
                                       coords=[time, bounds],
                                       dims=[TIME_STR, BOUNDS_STR],
                                       name=TIME_BOUNDS_STR)
    units_str = 'days since 2000-01-01 00:00:00'
    ds[TIME_STR].attrs['units'] = units_str
    ds[TIME_BOUNDS_STR].attrs['units'] = units_str
    return ds 

def monthly_mean_at_each_ind(monthly_means, sub_monthly_timeseries):
    """Copy monthly mean over each time index in that month.

    Parameters
    ----------
    monthly_means : xarray.DataArray
        array of monthly means
    sub_monthly_timeseries : xarray.DataArray
        array of a timeseries at sub-monthly time resolution

    Returns
    -------
    xarray.DataArray with eath monthly mean value from `monthly_means` repeated
    at each time within that month from `sub_monthly_timeseries`

    See Also
    --------
    monthly_mean_ts : Create timeseries of monthly mean values
    """
    time = monthly_means[TIME_STR]
    start = time.indexes[TIME_STR][0].replace(day=1, hour=0)
    end = time.indexes[TIME_STR][-1]
    new_indices = pd.date_range(start=start, end=end, freq='MS')
    arr_new = monthly_means.reindex(time=new_indices, method='backfill')
    return arr_new.reindex_like(sub_monthly_timeseries, method='pad') 

def monthly_mean_ts(arr):
    """Convert a sub-monthly time-series into one of monthly means.

    Also drops any months with no data in the original DataArray.

    Parameters
    ----------
    arr : xarray.DataArray
        Timeseries of sub-monthly temporal resolution data

    Returns
    -------
    xarray.DataArray
        Array resampled to comprise monthly means

    See Also
    --------
    monthly_mean_at_each_ind : Copy monthly means to each submonthly time

    """
    return arr.resample(
        **{TIME_STR: '1M'}, restore_coord_dims=True
    ).mean(TIME_STR).dropna(TIME_STR) 

def weighted_sum(data, dim=None, weights=None):
    """ Compute weighted sum for xarray objects

    Parameters
    ----------
    data : xarray.Dataset or xarray.DataArray
         xarray object for which to compute `weighted sum`
    dim : str or sequence of str, optional
        Dimension(s) over which to apply sum.
    weights : xarray.DataArray or array-like
        weights to apply. Shape must be broadcastable to shape of data.

    Returns
    -------
    reduced : xarray.Dataset or xarray.DataArray
        New xarray object with weighted sum applied to its data and the indicated
        dimension(s) removed.
    """

    if isinstance(dim, str):
        dim = [dim]

    if isinstance(data, xr.DataArray):
        return weighted_sum_da(data, dim, weights)
    elif isinstance(data, xr.Dataset):
        return data.apply(weighted_sum_da, dim=dim, weights=weights)
    else:
        raise ValueError('Data must be an xarray Dataset or DataArray') 

def weighted_sum_da(da, dim=None, weights=None):

    """ Compute weighted mean for DataArray
    Parameters
    ----------
    da : xarray.DataArray
        DataArray for which to compute `weighted mean`
    dim : str or sequence of str, optional
        Dimension(s) over which to apply mean.
    weights : xarray.DataArray or array-like
        weights to apply. Shape must be broadcastable to shape of da.

    Returns
    -------
    reduced : xarray.DataArray
        New DataArray with mean applied to its data and the indicated
        dimension(s) removed.

    """

    if dim is None:
        dim = list(da.dims)

    if weights is None:
        return da.sum(dim)
    else:
        weights, _ = validate_weights(da, dim, weights)
        return (da * weights).sum(dim) 

def test_dft_1d_time(self):
        """Test the discrete Fourier transform function on timeseries data."""
        time = pd.date_range('2000-01-01', '2001-01-01', closed='left')
        Nt = len(time)
        da = xr.DataArray(np.random.rand(Nt), coords=[time], dims=['time'])

        ft = xrft.dft(da)

        # check that frequencies are correct
        dt = (time[1] - time[0]).total_seconds()
        freq_time_expected = np.fft.fftshift(np.fft.fftfreq(Nt, dt))
        npt.assert_allclose(ft['freq_time'], freq_time_expected) 

def test_window_single_dim():
    # Julius' example
    # https://github.com/rabernat/xrft/issues/16
    data = xr.DataArray(np.random.random([20,30,100]),
                    dims=['time','lat','lon'],
                    coords={'time':range(20),'lat':range(30),'lon':range(100)})
    ps = xrft.power_spectrum(data, dim=['time'], window=True)
    # make sure it works with dask data
    ps = xrft.power_spectrum(data.chunk(), dim=['time'], window=True)
    ps.load() 

def test_data_1d(request):
    """Create one dimensional test DataArray."""
    Nx = 16
    Lx = 1.0
    x = np.linspace(0, Lx, Nx)
    dx = x[1] - x[0]
    coords = None if request.param == 'nocoords' else [x]
    da = xr.DataArray(np.random.rand(Nx), coords=coords, dims=['x'])
    if request.param == 'dask':
        da = da.chunk()
    return da 

def isotropize(ps, fftdim, nfactor=4):
    """
    Isotropize a 2D power spectrum or cross spectrum
    by taking an azimuthal average.

    .. math::
        \text{iso}_{ps} = k_r N^{-1} \sum_{N} |\mathbb{F}(da')|^2

    where :math:`N` is the number of azimuthal bins.

    Parameters
    ----------
    ps : `xarray.DataArray`
        The power spectrum or cross spectrum to be isotropized.
    fftdim : list
        The fft dimensions overwhich the isotropization must be performed.
    nfactor : int, optional
        Ratio of number of bins to take the azimuthal averaging with the
        data size. Default is 4.
    """

    # compute radial wavenumber bins
    k = ps[fftdim[1]]
    l = ps[fftdim[0]]
    N = [k.size, l.size]
    ki, kr = _radial_wvnum(k, l, min(N), nfactor)

    # average azimuthally
    ps = ps.assign_coords(freq_r=np.sqrt(k**2+l**2))
    iso_ps = (ps.groupby_bins('freq_r', bins=ki, labels=kr).mean()
              .rename({'freq_r_bins': 'freq_r'})
             )
    return iso_ps * iso_ps.freq_r 

def _stack_chunks(da, dim, suffix='_segment'):
    """Reshape a DataArray so there is only one chunk along dimension `dim`"""
    data = da.data
    attr = da.attrs
    newdims = []
    newcoords = {}
    newshape = []
    for d in da.dims:
        if d in dim:
            axis_num = da.get_axis_num(d)
            if np.diff(da.chunks[axis_num]).sum() != 0:
                raise ValueError("Chunk lengths need to be the same.")
            n = len(da[d])
            chunklen = da.chunks[axis_num][0]
            coord_rs = da[d].data.reshape((int(n/chunklen),int(chunklen)))
            newdims.append(d + suffix)
            newdims.append(d)
            newshape.append(int(n/chunklen))
            newshape.append(int(chunklen))
            newcoords[d+suffix] = range(int(n/chunklen))
            newcoords[d] = coord_rs[0]
        else:
            newdims.append(d)
            newshape.append(len(da[d]))
            newcoords[d] = da[d].data

    da = xr.DataArray(data.reshape(newshape), dims=newdims, coords=newcoords,
                     attrs=attr)

    return da 

def _apply_detrend(da, axis_num):
    """Wrapper function for applying detrending"""
    if da.chunks:
        func = detrend_wrap(detrendn)
        da = xr.DataArray(func(da.data, axes=axis_num),
                         dims=da.dims, coords=da.coords)
    else:
        if da.ndim == 1:
            da = xr.DataArray(sps.detrend(da),
                             dims=da.dims, coords=da.coords)
        else:
            da = detrendn(da, axes=axis_num)

    return da 

def test_attrs(units, description, dtype_out_vert, expected_units,
               expected_description):
    da = xr.DataArray(None)
    ds = xr.Dataset({'bar': 'foo', 'boo': 'baz'})
    da = _add_metadata_as_attrs(da, units, description, dtype_out_vert)
    ds = _add_metadata_as_attrs(ds, units, description, dtype_out_vert)
    assert expected_units == da.attrs['units']
    assert expected_description == da.attrs['description']
    for name, arr in ds.data_vars.items():
        assert expected_units == arr.attrs['units']
        assert expected_description == arr.attrs['description'] 

def test_rename_grid_attrs_dim_no_coord(ds_with_time_bounds, var_name):
    bounds_dim = 'nv'
    assert bounds_dim not in ds_with_time_bounds
    assert bounds_dim in GRID_ATTRS[BOUNDS_STR]
    # Create DataArray with all dims lacking coords.
    arr = xr.DataArray(ds_with_time_bounds[var_name].values, name='dummy')
    # Insert name to be replaced (its physical meaning doesn't matter here)
    ds = arr.rename({'dim_0': bounds_dim}).to_dataset()
    assert not ds[bounds_dim].coords
    result = grid_attrs_to_aospy_names(ds)
    assert not result[BOUNDS_STR].coords 

def test_ts_non_aospy_names(data_reg_alt_names):
    result = region_land_mask.ts(data_reg_alt_names, **_map_to_alt_names)
    expected = xr.DataArray(data_reg_alt_names.values[3, 0])
    xr.testing.assert_identical(result, expected)