Python xarray.Variable() Examples

def standardize_dims(self, var, dims={}):
        """Replace the coordinate names through x, y, z and t

        Parameters
        ----------
        var: xarray.Variable
            The variable to use the dimensions of
        dims: dict
            The dictionary to use for replacing the original dimensions

        Returns
        -------
        dict
            The dictionary with replaced dimensions"""
        dims = dict(dims)
        name_map = {self.get_xname(var, self.ds.coords): 'x',
                    self.get_yname(var, self.ds.coords): 'y',
                    self.get_zname(var, self.ds.coords): 'z',
                    self.get_tname(var, self.ds.coords): 't'}
        dims = dict(dims)
        for dim in set(dims).intersection(name_map):
            dims[name_map[dim]] = dims.pop(dim)
        return dims 

def _get_variable_point(vname, mask_override):
    # fix for https://github.com/MITgcm/xmitgcm/issues/191
    if vname in mask_override:
        return mask_override[vname]
    dims = _VAR_METADATA[vname]['dims']
    if 'i' in dims and 'j' in dims:
        point = 'c'
    elif 'i_g' in dims and 'j' in dims:
        point = 'w'
    elif 'i' in dims and 'j_g' in dims:
        point = 's'
    elif 'i_g' in dims and 'j_g' in dims:
        raise ValueError("Don't have masks for corner points!")
    else:
        raise ValueError("Variable `%s` is not a horizontal variable." % vname)
    return point 

def _from_dataset_test_variables(self):
        """The variables and coords needed for the from_dataset tests"""
        variables = {
             # 3d-variable
             'v0': xr.Variable(('time', 'ydim', 'xdim'), np.zeros((4, 4, 4))),
             # 2d-variable with time and x
             'v1': xr.Variable(('time', 'xdim', ), np.zeros((4, 4))),
             # 2d-variable with y and x
             'v2': xr.Variable(('ydim', 'xdim', ), np.zeros((4, 4))),
             # 1d-variable
             'v3': xr.Variable(('xdim', ), np.zeros(4))}
        coords = {
            'ydim': xr.Variable(('ydim', ), np.arange(1, 5)),
            'xdim': xr.Variable(('xdim', ), np.arange(4)),
            'time': xr.Variable(
                ('time', ),
                pd.date_range('1999-01-01', '1999-05-01', freq='M').values)}
        return variables, coords 

def test_plot_bounds_2d(self):
        x = np.arange(1, 5)
        y = np.arange(5, 10)
        x2d, y2d = np.meshgrid(x, y)
        x_bnds = np.arange(0.5, 4.51, 1.0)
        y_bnds = np.arange(4.5, 9.51, 1.0)
        # the borders are not modified
        x_bnds[0] = 1.0
        x_bnds[-1] = 4.0
        y_bnds[0] = 5.0
        y_bnds[-1] = 9.0
        x2d_bnds, y2d_bnds = np.meshgrid(x_bnds, y_bnds)
        d = psyd.CFDecoder()
        # test x bounds
        bounds = d.get_plotbounds(xr.Variable(('y', 'x'), x2d))
        self.assertAlmostArrayEqual(bounds, x2d_bnds)

        # test y bounds
        bounds = d.get_plotbounds(xr.Variable(('y', 'x'), y2d))
        self.assertAlmostArrayEqual(bounds, y2d_bnds) 

def _insert_fldmean_bounds(self, da, keepdims=False):
        xcoord = self.get_coord('x')
        ycoord = self.get_coord('y')
        sdims = (self.get_dim('y'), self.get_dim('x'))
        xbounds = np.array([[xcoord.min(), xcoord.max()]])
        ybounds = np.array([[ycoord.min(), ycoord.max()]])
        xdims = (sdims[-1], 'bnds') if keepdims else ('bnds', )
        ydims = (sdims[0], 'bnds') if keepdims else ('bnds', )
        xattrs = xcoord.attrs.copy()
        xattrs.pop('bounds', None)
        yattrs = ycoord.attrs.copy()
        yattrs.pop('bounds', None)
        da.psy.base.coords[xcoord.name + '_bnds'] = xr.Variable(
            xdims, xbounds if keepdims else xbounds[0], attrs=xattrs)
        da.psy.base.coords[ycoord.name + '_bnds'] = xr.Variable(
            ydims, ybounds if keepdims else ybounds[0], attrs=yattrs) 

def base(self):
        """Base dataset this instance gets its data from"""
        if self._base is None:
            if 'variable' in self.arr.dims:
                def to_dataset(i):
                    ret = self.isel(variable=i).to_dataset(
                        name=self.arr.coords['variable'].values[i])
                    try:
                        return ret.drop('variable')
                    except ValueError:  # 'variable' Variable not defined
                        pass
                    return ret
                ds = to_dataset(0)
                if len(self.arr.coords['variable']) > 1:
                    for i in range(1, len(self.arr.coords['variable'])):
                        ds.update(ds.merge(to_dataset(i)))
                self._base = ds
            else:
                self._base = self.arr.to_dataset(
                    name=self.arr.name or self.arr_name)
            self.onbasechange.emit()
        return self._base 

def get_mesh(self, var, coords=None):
        """Get the mesh variable for the given `var`

        Parameters
        ----------
        var: xarray.Variable
            The data source whith the ``'mesh'`` attribute
        coords: dict
            The coordinates to use. If None, the coordinates of the dataset of
            this decoder is used

        Returns
        -------
        xarray.Coordinate
            The mesh coordinate"""
        mesh = var.attrs.get('mesh')
        if mesh is None:
            return None
        if coords is None:
            coords = self.ds.coords
        return coords.get(mesh, self.ds.coords.get(mesh)) 

def read_data(data_handle, domain=None, is_worker=False,
              start=None, stop=None, calendar='standard',
              var_dict=None) -> xr.Dataset:
    """Read data directly from an xarray dataset"""
    varlist = list(data_handle.keys())
    if var_dict is not None:
        data_handle = data_handle.rename(var_dict)
        varlist = list(var_dict.values())
    data_handle = data_handle[varlist]

    if start is not None and stop is not None:
        data_handle = data_handle.sel(time=slice(start, stop))
        dates = data_handle.indexes['time']
        data_handle['day_of_year'] = xr.Variable(('time', ), dates.dayofyear)

    return data_handle 

def can_decode(cls, ds, var):
        """
        Class method to determine whether the object can be decoded by this
        decoder class.

        Parameters
        ----------
        ds: xarray.Dataset
            The dataset that contains the given `var`
        var: xarray.Variable or xarray.DataArray
            The array to decode

        Returns
        -------
        bool
            True if the decoder can decode the given array `var`. Otherwise
            False

        Notes
        -----
        The default implementation returns True for any argument. Subclass this
        method to be specific on what type of data your decoder can decode
        """
        return True 

def _load_GeoTransform(self):
        """Calculate latitude and longitude variable calculated from the
        gdal.Open.GetGeoTransform method"""
        def load_lon():
            return arange(ds.RasterXSize)*b[1]+b[0]

        def load_lat():
            return arange(ds.RasterYSize)*b[5]+b[3]
        ds = self.ds
        b = self.ds.GetGeoTransform()  # bbox, interval
        if with_dask:
            lat = Array(
                {('lat', 0): (load_lat,)}, 'lat', (self.ds.RasterYSize,),
                shape=(self.ds.RasterYSize,), dtype=float)
            lon = Array(
                {('lon', 0): (load_lon,)}, 'lon', (self.ds.RasterXSize,),
                shape=(self.ds.RasterXSize,), dtype=float)
        else:
            lat = load_lat()
            lon = load_lon()
        return Variable(('lat',), lat), Variable(('lon',), lon) 

def sspeed(depth: Union[np.ndarray, xr.Variable],
           latitude: np.ndarray,
           temperature: np.ndarray,
           salinity: np.ndarray) -> np.ndarray:
    """
    Calculates the speed of sound.

    Required Arguments:

    * depth: The depth(s) in meters
    * latitude: The latitude(s) in degrees North
    * temperature: The temperatures(s) in Celsius
    * salinity: The salinity (unitless)
    """

    depth, latitude, temperature, salinity = __validate_depth_lat_temp_sal(
        depth, latitude, temperature, salinity)

    press = __calc_pressure(depth, latitude)

    if salinity.shape != press.shape:
        # pad array shape to match otherwise seawater freaks out
        press = press[..., np.newaxis]

    speed = seawater.svel(salinity, temperature, press)
    return np.squeeze(speed) 

def read_netcdf(data_handle, domain=None, is_worker=False,
                start=None, stop=None, calendar='standard',
                var_dict=None) -> xr.Dataset:
    """Read in a NetCDF file"""
    if '*' in data_handle:
        ds = xr.open_mfdataset(data_handle)
    else:
        ds = xr.open_dataset(data_handle)

    if domain is not None:
        ds = ds.sel({k: domain[k]
                     for k in list(domain.dims.keys())
                     if k in list(ds.dims.keys())})
    else:
        dims_wo_coords = set(ds.dims) - set(ds.coords)
        for d in dims_wo_coords:
            if is_worker:
                logger = logging.getLogger('MetSim')
                logger.warning(
                    'Setting sequential coordinate on dimension {}'.format(d))
            ds[d] = np.arange(0, len(ds[d]))

    if 'time' in ds.coords:
        if isinstance(ds.indexes['time'], xr.CFTimeIndex):
            ds['time'] = ds.indexes['time'].to_datetimeindex()
        ds['time'] = (ds.indexes['time'] -
                      pd.Timedelta('11H59M59S')).round('D')

    if var_dict is not None:
        var_list = list(var_dict.keys())
        ds = ds[var_list]
        ds = ds.rename(var_dict)

    if start is not None or stop is not None:
        ds = ds.sel(time=slice(start, stop))
        dates = ds.indexes['time']
        ds['day_of_year'] = xr.Variable(('time', ), dates.dayofyear)

    return ds 

def test_deepsoundchannel(self):
        self.assertEqual(funcs.deepsoundchannel(xr.Variable(
            data=[0], dims=['depth']), [45], [0.5], [32]), 0) 

def geostrophic_x(h, lat, lon):
    """Calculates the X component of geostrophic currents

    Parameters:
    h -- Sea Surface Height, xarray or netcdf variable, already sliced
    lat -- an array of latitudes, the shape must match that of h
    lon -- an array of longitudes, the shape must match that of h
    """
    if isinstance(lat, xr.Variable):
        lat = lat.values

    if hasattr(h, "dims"):
        dims = h.dims
    else:
        dims = h.dimensions

    dim_order = "".join([d for d in dims if d in 'yx'])

    def f(heights, **kwargs):
        c = metpy.calc.coriolis_parameter(lat * _ureg.degrees)
        if dim_order == "yx":
            dy, dx = kwargs['deltas']
        else:
            dx, dy = kwargs['deltas']

        return metpy.calc.geostrophic_wind(xr.DataArray(heights), c, dx, dy,
                                           dim_order=kwargs['dim_order'])

    return _metpy(f, h, lat, lon, dim_order[0]) 

def geostrophic_y(h, lat, lon):
    """Calculates the Y component of geostrophic currents

    Parameters:
    h -- Sea Surface Height, xarray or netcdf variable, already sliced
    lat -- an array of latitudes, the shape must match that of h
    lon -- an array of longitudes, the shape must match that of h
    """
    if isinstance(lat, xr.Variable):
        lat = lat.values

    if hasattr(h, "dims"):
        dims = h.dims
    else:
        dims = h.dimensions

    dim_order = "".join([d for d in dims if d in 'yx'])

    def f(heights, **kwargs):
        c = metpy.calc.coriolis_parameter(lat * _ureg.degrees)
        if dim_order == "yx":
            dy, dx = kwargs['deltas']
        else:
            dx, dy = kwargs['deltas']

        return metpy.calc.geostrophic_wind(xr.DataArray(heights), c, dx, dy,
                                           dim_order=kwargs['dim_order'])

    return _metpy(f, h, lat, lon, dim_order[1]) 

def qds_month():
    dims = ("quantiles", "month")
    source = xr.Variable(dims=dims, data=np.zeros((5, 12)))
    target = xr.Variable(dims=dims, data=np.ones((5, 12)) * 2)

    return xr.Dataset(
        data_vars={"source": source, "target": target},
        coords={"quantiles": [0, 0.3, 5.0, 7, 1], "month": range(1, 13)},
        attrs={"group": "time.month", "window": 1},
    ) 

def _iternum_to_datetime_variable(iternum, delta_t, ref_date,
                                  calendar, time_dim_name='time'):
    # create time array
    timedata = np.atleast_1d(iternum)*delta_t
    time_attrs = {'standard_name': 'time', 'long_name': 'Time', 'axis': 'T'}
    if ref_date is not None:
        time_attrs['units'] = 'seconds since %s' % ref_date
    else:
        time_attrs['units'] = 'seconds'
    if calendar is not None:
        time_attrs['calendar'] = calendar
    timevar = xr.Variable((time_dim_name,), timedata, time_attrs)
    return timevar 

def test_sspeed(self):
        np.testing.assert_allclose(funcs.sspeed(xr.Variable(
            data=[0], dims=['depth']), [45], [0.5], [32]), 1447.4, rtol=self.relative_tolerance)

        dep = np.array(range(0, 10))
        lat = np.array([[45, 45], [45, 45]])
        temp = 0.5 * np.ones((10, 2, 2))
        sal = 32 * np.ones((10, 2, 2))

        np.testing.assert_allclose(funcs.sspeed(
            dep, lat, temp, sal)[0, 0, 0], 1447.4, rtol=self.relative_tolerance) 

def _faces_coords_to_latlon(ds):
    coords = ds.reset_coords().coords.to_dataset()
    ifac = 4
    jfac = 3
    dim_coords = {}
    for vname in coords.coords:
        if vname[0] == 'i':
            data = np.arange(ifac * coords.dims[vname])
        elif vname[0] == 'j':
            data = np.arange(jfac * coords.dims[vname])
        else:
            data = coords[vname].data
        var = xr.Variable(ds[vname].dims, data, ds[vname].attrs)
        dim_coords[vname] = var
    return xr.Dataset(dim_coords) 

def get_variables(self):
        def load(band):
            band = ds.GetRasterBand(band)
            a = band.ReadAsArray()
            no_data = band.GetNoDataValue()
            if no_data is not None:
                try:
                    a[a == no_data] = a.dtype.type(nan)
                except ValueError:
                    pass
            return a
        ds = self.ds
        dims = ['lat', 'lon']
        chunks = ((ds.RasterYSize,), (ds.RasterXSize,))
        shape = (ds.RasterYSize, ds.RasterXSize)
        variables = OrderedDict()
        for iband in range(1, ds.RasterCount+1):
            band = ds.GetRasterBand(iband)
            dt = dtype(gdal_array.codes[band.DataType])
            if with_dask:
                dsk = {('x', 0, 0): (load, iband)}
                arr = Array(dsk, 'x', chunks, shape=shape, dtype=dt)
            else:
                arr = load(iband)
            attrs = band.GetMetadata_Dict()
            try:
                dt.type(nan)
                attrs['_FillValue'] = nan
            except ValueError:
                no_data = band.GetNoDataValue()
                attrs.update({'_FillValue': no_data} if no_data else {})
            variables['Band%i' % iband] = Variable(dims, arr, attrs)
        variables['lat'], variables['lon'] = self._load_GeoTransform()
        return FrozenOrderedDict(variables) 

def _test_ds(self):
        import xarray as xr
        import pandas as pd
        time = xr.Coordinate('time', pd.to_datetime(
            ['1979-01-01T12:00:00', '1979-01-01T18:00:00',
             '1979-01-01T18:30:00']),
            encoding={'units': 'day as %Y%m%d.%f'})
        var = xr.Variable(('time', 'x'), np.zeros((len(time), 5)))
        return xr.Dataset({'test': var}, {'time': time}) 

def _filter_test_ds(self):
        return xr.Dataset(
            {'v0': xr.Variable(('ydim', 'xdim'), np.zeros((4, 4)),
                               attrs={'test': 1, 'test2': 1}),
             'v1': xr.Variable(('xdim', ), np.zeros(4), attrs={'test': 2,
                                                               'test2': 2}),
             'v2': xr.Variable(('xdim', ), np.zeros(4), attrs={'test': 3,
                                                               'test2': 3})},
            {'ydim': xr.Variable(('ydim', ), np.arange(1, 5)),
             'xdim': xr.Variable(('xdim', ), np.arange(4))}) 

def test_plot_bounds_1d(self):
        """Test to get 2d-interval breaks"""
        x = xr.Variable(('x', ), np.arange(1, 5))
        d = psyd.CFDecoder()
        bounds = d.get_plotbounds(x)
        self.assertAlmostArrayEqual(bounds, np.arange(0.5, 4.51, 1.0)) 

def test_1D_cf_bounds(self):
        """Test whether the CF Conventions for 1D bounaries are correct"""
        final_bounds = np.arange(-180, 181, 30)
        lon = xr.Variable(('lon', ), np.arange(-165, 166, 30),
                          {'bounds': 'lon_bounds'})
        cf_bounds = xr.Variable(('lon', 'bnds'), np.zeros((len(lon), 2)))
        for i in range(len(lon)):
            cf_bounds[i, :] = final_bounds[i:i+2]
        ds = xr.Dataset(coords={'lon': lon, 'lon_bounds': cf_bounds})
        decoder = psyd.CFDecoder(ds)
        self.assertEqual(list(final_bounds),
                         list(decoder.get_plotbounds(lon))) 

def _decode_ds(cls, ds, gridfile=None, decode_coords=True,
                   decode_times=True):
        """
        Static method to decode coordinates and time informations

        This method interpretes absolute time informations (stored with units
        ``'day as %Y%m%d.%f'``) and coordinates

        Parameters
        ----------
        %(CFDecoder.decode_coords.parameters)s
        decode_times : bool, optional
            If True, decode times encoded in the standard NetCDF datetime
            format into datetime objects. Otherwise, leave them encoded as
            numbers.
        decode_coords : bool, optional
            If True, decode the 'coordinates' attribute to identify coordinates
            in the resulting dataset."""
        if decode_coords:
            ds = cls.decode_coords(ds, gridfile=gridfile)
        if decode_times:
            for k, v in six.iteritems(ds.variables):
                # check for absolute time units and make sure the data is not
                # already decoded via dtype check
                if v.attrs.get('units', '') == 'day as %Y%m%d.%f' and (
                        np.issubdtype(v.dtype, np.float64)):
                    decoded = xr.Variable(
                        v.dims, AbsoluteTimeDecoder(v), attrs=v.attrs,
                        encoding=v.encoding)
                    ds.update({k: decoded})
        return ds 

def get_t(self, var, coords=None):
        """
        Get the time coordinate of a variable

        This method searches for the time coordinate in the :attr:`ds`. It
        first checks whether there is one dimension that holds an ``'axis'``
        attribute with 'T', otherwise it looks whether there is an intersection
        between the :attr:`t` attribute and the variables dimensions, otherwise
        it returns the coordinate corresponding to the first dimension of `var`

        Possible types
        --------------
        var: xarray.Variable
            The variable to get the time coordinate for
        coords: dict
            Coordinates to use. If None, the coordinates of the dataset in the
            :attr:`ds` attribute are used.

        Returns
        -------
        xarray.Coordinate or None
            The time coordinate or None if no time coordinate could be found"""
        coords = coords or self.ds.coords
        coord = self.get_variable_by_axis(var, 't', coords)
        if coord is not None:
            return coord
        dimlist = list(self.t.intersection(var.dims).intersection(coords))
        if dimlist:
            if len(dimlist) > 1:
                warn("Found multiple matches for time coordinate in the "
                     "variable: %s. I use %s" % (
                         ', '.join(dimlist), dimlist[0]),
                     PsyPlotRuntimeWarning)
            return coords[dimlist[0]]
        tname = self.get_tname(var)
        if tname is not None:
            return coords.get(tname)
        return None 

def get_z(self, var, coords=None):
        """
        Get the vertical (z-) coordinate of a variable

        This method searches for the z-coordinate in the :attr:`ds`. It first
        checks whether there is one dimension that holds an ``'axis'``
        attribute with 'Z', otherwise it looks whether there is an intersection
        between the :attr:`z` attribute and the variables dimensions, otherwise
        it returns the coordinate corresponding to the third last dimension of
        `var` (or the second last or last if var is two or one-dimensional)

        Possible types
        --------------
        var: xarray.Variable
            The variable to get the z-coordinate for
        coords: dict
            Coordinates to use. If None, the coordinates of the dataset in the
            :attr:`ds` attribute are used.

        Returns
        -------
        xarray.Coordinate or None
            The z-coordinate or None if no z coordinate could be found"""
        coords = coords or self.ds.coords
        coord = self.get_variable_by_axis(var, 'z', coords)
        if coord is not None:
            return coord
        zname = self.get_zname(var)
        if zname is not None:
            return coords.get(zname)
        return None 

def get_y(self, var, coords=None):
        """
        Get the y-coordinate of a variable

        This method searches for the y-coordinate in the :attr:`ds`. It first
        checks whether there is one dimension that holds an ``'axis'``
        attribute with 'Y', otherwise it looks whether there is an intersection
        between the :attr:`y` attribute and the variables dimensions, otherwise
        it returns the coordinate corresponding to the second last dimension of
        `var` (or the last if the dimension of var is one-dimensional)

        Possible types
        --------------
        var: xarray.Variable
            The variable to get the y-coordinate for
        coords: dict
            Coordinates to use. If None, the coordinates of the dataset in the
            :attr:`ds` attribute are used.

        Returns
        -------
        xarray.Coordinate or None
            The y-coordinate or None if it could be found"""
        coords = coords or self.ds.coords
        coord = self.get_variable_by_axis(var, 'y', coords)
        if coord is not None:
            return coord
        return coords.get(self.get_yname(var)) 

def get_x(self, var, coords=None):
        """
        Get the x-coordinate of a variable

        This method searches for the x-coordinate in the :attr:`ds`. It first
        checks whether there is one dimension that holds an ``'axis'``
        attribute with 'X', otherwise it looks whether there is an intersection
        between the :attr:`x` attribute and the variables dimensions, otherwise
        it returns the coordinate corresponding to the last dimension of `var`

        Possible types
        --------------
        var: xarray.Variable
            The variable to get the x-coordinate for
        coords: dict
            Coordinates to use. If None, the coordinates of the dataset in the
            :attr:`ds` attribute are used.

        Returns
        -------
        xarray.Coordinate or None
            The y-coordinate or None if it could be found"""
        coords = coords or self.ds.coords
        coord = self.get_variable_by_axis(var, 'x', coords)
        if coord is not None:
            return coord
        return coords.get(self.get_xname(var)) 

def to_netcdf(ds, *args, **kwargs):
    """
    Store the given dataset as a netCDF file

    This functions works essentially the same as the usual
    :meth:`xarray.Dataset.to_netcdf` method but can also encode absolute time
    units

    Parameters
    ----------
    ds: xarray.Dataset
        The dataset to store
    %(xarray.Dataset.to_netcdf.parameters)s
    """
    to_update = {}
    for v, obj in six.iteritems(ds.variables):
        units = obj.attrs.get('units', obj.encoding.get('units', None))
        if units == 'day as %Y%m%d.%f' and np.issubdtype(
                obj.dtype, np.datetime64):
            to_update[v] = xr.Variable(
                obj.dims, AbsoluteTimeEncoder(obj), attrs=obj.attrs.copy(),
                encoding=obj.encoding)
            to_update[v].attrs['units'] = units
    if to_update:
        ds = ds.copy()
        ds.update(to_update)
    return xarray_api.to_netcdf(ds, *args, **kwargs)