python类DataArray()的实例源码

def mass_streamfun(self):

        from scipy import integrate

        data = self._obj
#        lonlen = len(data.lon)
        if 'lon' in data.dims:
            data = data.fillna(0).mean('lon')
        levax = data.get_axis_num('lev')
        stream = integrate.cumtrapz(data * np.cos(np.deg2rad(data.lat)), x=data.lev * 1e2, initial=0., axis=levax)
        stream = stream * 2 * np.pi  / cc.g * cc.rearth * 1e-9
        stream = xr.DataArray(stream, coords=data.coords, dims=data.dims)
        stream = stream.rename('ovt')
        stream.attrs['long name'] = 'atmosphere overturning circulation'
        stream.attrs['unit'] = 'Sv (1e9 kg/s)'
        return stream

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)
        # else:
        #     raise ValueError("Data should be dask array.")

    return da

def zeros_like(array, dtype=None, keepmeta=True):
    """Create an array of zeros with the same shape and type as the input array.

    Args:
        array (xarray.DataArray): The shape and data-type of it define
            these same attributes of the output array.
        dtype (data-type, optional): If specified, this function overrides
            the data-type of the output array.
        keepmeta (bool, optional): Whether *coords, attrs, and name of the input
            array are kept in the output one. Default is True.

    Returns:
        array (decode.array): Decode array filled with zeros.
    """
    if keepmeta:
        return xr.zeros_like(array, dtype)
    else:
        return dc.zeros(array.shape, dtype)

def ones_like(array, dtype=None, keepmeta=True):
    """Create an array of ones with the same shape and type as the input array.

    Args:
        array (xarray.DataArray): The shape and data-type of it define
            these same attributes of the output array.
        dtype (data-type, optional): If spacified, this function overrides
            the data-type of the output array.
        keepmeta (bool, optional): Whether *coords, attrs, and name of the input
            array are kept in the output one. Default is True.

    Returns:
        array (decode.array): Decode array filled with ones.
    """
    if keepmeta:
        return xr.ones_like(array, dtype)
    else:
        return dc.ones(array.shape, dtype)

def full_like(array, fill_value, reverse=False, dtype=None, keepmeta=True):
    """Create an array of `fill_value` with the same shape and type as the input array.

    Args:
        array (xarray.DataArray): The shape and data-type of it define
            these same attributes of the output array.
        fill_value (scalar or numpy.ndarray): Fill value or array.
        dtype (data-type, optional): If spacified, this function overrides
            the data-type of the output array.
        keepmeta (bool, optional): Whether *coords, attrs, and name of the input
            array are kept in the output one. Default is True.

    Returns:
        array (decode.array): Decode array filled with `fill_value`.
    """
    if keepmeta:
        return (dc.zeros_like(array) + fill_value).astype(dtype)
    else:
        return dc.full(array.shape, fill_value, dtype)

def empty_like(array, dtype=None, keepmeta=True):
    """Create an array of empty with the same shape and type as the input array.

    Args:
        array (xarray.DataArray): The shape and data-type of it define
            these same attributes of the output array.
        dtype (data-type, optional): If spacified, this function overrides
            the data-type of the output array.
        keepmeta (bool, optional): Whether *coords, attrs, and name of the input
            array are kept in the output one. Default is True.

    Returns:
        array (decode.array): Decode array without initializing entries.
    """
    if keepmeta:
        return dc.empty(array.shape, dtype,
            tcoords=array.dca.tcoords, chcoords=array.dca.chcoords,
            scalarcoords=array.dca.scalarcoords, attrs=array.attrs, name=array.name
        )
    else:
        return dc.empty(array.shape, dtype)

def is_data_dependent(fmto, data):
    """Check whether a formatoption is data dependent

    Parameters
    ----------
    fmto: Formatoption
        The :class:`Formatoption` instance to check
    data: xarray.DataArray
        The data array to use if the :attr:`~Formatoption.data_dependent`
        attribute is a callable

    Returns
    -------
    bool
        True, if the formatoption depends on the data"""
    if callable(fmto.data_dependent):
        return fmto.data_dependent(data)
    return fmto.data_dependent

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 is_triangular(self, var):
        """
        Test if a variable is on a triangular grid

        This method first checks the `grid_type` attribute of the variable (if
        existent) whether it is equal to ``"unstructered"``, then it checks
        whether the bounds are not two-dimensional.

        Parameters
        ----------
        var: xarray.Variable or xarray.DataArray
            The variable to check

        Returns
        -------
        bool
            True, if the grid is triangular, else False"""
        return str(var.attrs.get('grid_type')) == 'unstructured' or \
            self._check_triangular_bounds(var)[0]

def test_auto_update(self):
        """Test the :attr:`psyplot.plotter.Plotter.no_auto_update` attribute"""
        data = xr.DataArray([])
        plotter = TestPlotter(data, auto_update=False)
        self.assertFalse(plotter.no_auto_update)
        data.psy.init_accessor(auto_update=False)
        plotter = TestPlotter(data, auto_update=False)
        self.assertTrue(plotter.no_auto_update)

        plotter.update(fmt1=1)
        self.assertEqual(plotter['fmt1'], '')
        self.assertEqual(plotter._registered_updates['fmt1'], 1)

        plotter.start_update()
        self.assertEqual(plotter['fmt1'], '1')
        self.assertFalse(plotter._registered_updates)

        data.psy.no_auto_update = False
        self.assertFalse(plotter.data.psy.no_auto_update)
        self.assertFalse(plotter.no_auto_update)

def test_data_props_list(self):
        """Test the data properties of Formatoptions with an InteractiveList"""
        data = psyd.InteractiveList([xr.DataArray([]), xr.DataArray([])])
        plot_data = data.copy(True)
        plot_data.extend([xr.DataArray([]), xr.DataArray([])],
                         new_name=True)
        plotter = TestPlotter(data)
        plotter.plot_data = plot_data
        plot_data = plotter.plot_data  # the data might have been copied

        self.assertIs(plotter.fmt1.raw_data, data)
        self.assertIs(plotter.fmt1.data, plot_data)

        # test with index in list
        plotter.fmt1.index_in_list = 1
        self.assertIs(plotter.fmt1.raw_data, data[1])
        self.assertIs(plotter.fmt1.data, plot_data[1])

        # test with index in list of plot_data outside raw_data
        plotter.fmt1.index_in_list = 3
        self.assertIs(plotter.fmt1.data, plot_data[3])

def test_any_decoder(self):
        """Test the decoder property with an InteractiveList"""
        data = psyd.InteractiveList([xr.DataArray([]), xr.DataArray([])])
        plot_data = data.copy(True)
        plot_data.extend([xr.DataArray([]), xr.DataArray([])],
                         new_name=True)
        for arr in data:
            arr.psy.init_accessor(decoder=psyd.CFDecoder(arr.psy.base))
        plotter = TestPlotter(data)
        plotter.plot_data = plot_data
        plot_data = plotter.plot_data  # the data might have been copied

        # test without index in list
        decoder = psyd.CFDecoder(data[0].psy.base)
        plotter.fmt2.decoder = decoder
        for i, d2 in enumerate(plotter.plot_data_decoder):
            self.assertIs(d2, decoder,
                          msg='Decoder %i has been set wrong!' % i)
        self.assertEqual(plotter.fmt2.decoder, plotter.plot_data_decoder)
        self.assertIs(plotter.fmt2.any_decoder, decoder)

def test_has_changed(self):
        """Test the :meth:`psyplot.plotter.Plotter.show_summaries` method"""
        plotter = TestPlotter(xr.DataArray([]), fmt1='something')
        self.assertEqual(plotter['fmt1'], 'something')
        for i in range(1, 4):
            key = 'fmt%i' % i
            fmto = getattr(plotter, key)
            self.assertEqual(plotter.has_changed(key),
                             [fmto.default, plotter[key]],
                             msg="Wrong value for " + key)
        plotter.update()
        self.assertIsNone(plotter.has_changed('fmt1'))
        plotter.update(fmt1='test', fmt3=plotter.fmt3.default, force=True)
        self.assertEqual(plotter.has_changed('fmt1'),
                         ['something', 'test'])
        self.assertIsNone(plotter.has_changed('fmt2'))
        self.assertIsNone(plotter.has_changed('fmt3', include_last=False))
        self.assertEqual(plotter.has_changed('fmt3'),
                         [plotter.fmt3.default, plotter.fmt3.default])

def create_variable(dataset, name, data):
    if isinstance(data, xr.DataArray):
        for i in range(len(data.dims)):
            try:
                if i == 0:  # time
                    ensure_dimension_exists(
                        dataset, data.dims[i], None)
                else:
                    ensure_dimension_exists(
                        dataset, data.dims[i], data.values.shape[i])
            except IOError as err:
                raise IOError(
                    'Error while creating {}: {}'.format(name, err))
        dataset.createVariable(
            name, data.values.dtype, data.dims)
        for key, value in data.attrs.items():
            dataset.variables[name].setncattr(key, value)
    else:
        raise TypeError('data must be of type DataArray')

def to_xarray(self):
        """Convert to xarray.Dataset

        Returns
        -------
        xarray.Dataset
        """
        import xarray as xr
        data_vars = {
            "frequencies": xr.DataArray(self.frequencies, dims="bin"),
            "errors2": xr.DataArray(self.errors2, dims="bin"),
            "bins": xr.DataArray(self.bins, dims=("bin", "x01"))
        }
        coords = {}
        attrs = {
            "underflow": self.underflow,
            "overflow": self.overflow,
            "inner_missed": self.inner_missed,
            "keep_missed": self.keep_missed
        }
        attrs.update(self._meta_data)
        # TODO: Add stats
        return xr.Dataset(data_vars, coords, attrs)

def second_layer_input_matrix(X, models):
    '''Build a second layer model input matrix by taking the
    metadata from X given to the first layer models and forming
    a new matrix from the 1-D predictions of the first layer models
    '''
    preds = predict_many(dict(X=X), to_raster=False,
                         ensemble=models)
    example = preds[0].flat
    input_matrix = np.empty((example.shape[0], len(preds)))
    for j, pred in enumerate(preds):
        input_matrix[:, j] = pred.flat.values[:, 0]
    attrs = X.attrs.copy()
    attrs['old_dims'] = [X[SOIL_MOISTURE].dims] * len(preds)
    attrs['canvas'] = X[SOIL_MOISTURE].canvas
    tags = [tag for tag, _ in models]
    arr = xr.DataArray(input_matrix,
                       coords=[('space', example.space),
                               ('band', tags)],
                       dims=('space', 'band'),
                       attrs=attrs)
    return xr.Dataset(dict(flat=arr), attrs=attrs)

def kmeans_aic(model, X, **kwargs):
    '''AIC (Akaike Information Criterion) for k-means for model selection

    Parameters:
        :model:  An elm.pipeline.Pipeline with KMeans or MiniBatchKMeans as final step in Pipeline
        :X:      The X data that were just given to "fit", or "partial_fit"
        :kwargs: placeholder - ignored

    Returns:
        :AIC: float

    '''

    k, m = model._estimator.cluster_centers_.shape
    if isinstance(X, xr.DataArray):
        n = X.flat.values.shape[0]
    else:
        n = X.shape[0]
    d = model._estimator.inertia_
    aic =  d + 2 * m * k
    delattr(model._estimator, 'labels_')
    return aic

def predict(self, da):
        '''xarray.DataArray version of sklearn.cluster.KMeans.fit.'''

        # compatible with the sklean.cluster.KMeans predict method when the input data is not DataArray
        if not isinstance(da, xr.DataArray):
            return super().predict(da)

        # retrieve parameters
        n_samples = da.shape[0]
        features_shape = da.shape[1:]
        n_features = np.prod(features_shape)

        X = da.data.reshape(n_samples, n_features)# 'data' might be replaced with 'values'.
        # remove NaN values if exists in X
        try:
            X_valid = X[:, self.valid_features_index_]
        except:
            X_valid = X
        samples_dim = da.dims[0]
        samples_coord = {samples_dim: da.coords[samples_dim]}
        labels = xr.DataArray(super().predict(X_valid),
            dims=samples_dim, coords=samples_coord)

        return labels

def plasmaprop(iono:DataArray, f:float,B0:float):
    assert isinstance(iono,DataArray)
    Ne = iono.loc[:,'ne'].astype(float)
    w = 2*np.pi*f

    wp = np.sqrt(Ne*e**2/(eps0*me)) # electron plasma frequency [rad/sec]
    wH = B0*e/me               # electron gyrofrequency [rad/sec]


    if (w <= wp).any(): # else reflection doesn't occur, passes right through (radar freq > MUF)
        reflectionheight = iono.alt_km[abs(w-wp).argmin()]
    else:
        print(f'radar freq {f/1e6:.1f} MHz  >  max. plasma freq {wp.max()/(2*np.pi)/1e6:.1f} MHz: no reflection')
        reflectionheight = None

    return wp,wH,reflectionheight

def test_to_xarray(self):

        tm._skip_if_no_xarray()
        from xarray import DataArray

        p = tm.makePanel()

        result = p.to_xarray()
        self.assertIsInstance(result, DataArray)
        self.assertEqual(len(result.coords), 3)
        assert_almost_equal(list(result.coords.keys()),
                            ['items', 'major_axis', 'minor_axis'])
        self.assertEqual(len(result.dims), 3)

        # idempotency
        assert_panel_equal(result.to_pandas(), p)

def test_to_xarray(self):

        tm._skip_if_no_xarray()
        from xarray import DataArray

        p = tm.makePanel4D()

        result = p.to_xarray()
        self.assertIsInstance(result, DataArray)
        self.assertEqual(len(result.coords), 4)
        assert_almost_equal(list(result.coords.keys()),
                            ['labels', 'items', 'major_axis', 'minor_axis'])
        self.assertEqual(len(result.dims), 4)

        # non-convertible
        self.assertRaises(ValueError, lambda: result.to_pandas())

def test_activation_matrix():
    weights = xr.DataArray(np.array([[0, 1, 0], [1, 0, 0]]),
                           coords={
                               'outcomes': ['o1', 'o2'],
                               'cues': ['c1', 'c2', 'c3']
                           },
                           dims=('outcomes', 'cues'))

    events = [(['c1', 'c2', 'c3'], []),
              (['c1', 'c3'], []),
              (['c2'], []),
              (['c1', 'c1'], [])]
    reference_activations = np.array([[1, 0, 1, 0], [1, 1, 0, 1]])

    with pytest.raises(ValueError):
        activations = activation(events, weights, number_of_threads=1)

    activations = activation(events, weights, number_of_threads=1, remove_duplicates=True)
    activations_mp = activation(events, weights, number_of_threads=3, remove_duplicates=True)

    assert np.allclose(reference_activations, activations)
    assert np.allclose(reference_activations, activations_mp)

def load_subdataset(subdataset, attrs, layer_spec, **reader_kwargs):
    '''Load a single subdataset'''
    import gdal
    data_file = gdal.Open(subdataset)
    np_arr = data_file.ReadAsArray(**reader_kwargs)
    out = _np_arr_to_coords_dims(np_arr,
                 layer_spec,
                 reader_kwargs,
                 geo_transform=None,
                 layer_meta=attrs,
                 handle=data_file)
    np_arr, coords, dims, attrs2 = out
    attrs.update(attrs2)
    return xr.DataArray(data=np_arr,
                        coords=coords,
                        dims=dims,
                        attrs=attrs)

def _append_dataarray_extra_attrs(xarr,**extra_kwargs):
    """Update the dictionnary of attributes a xarray dataarray (xarr.attrs).

    Parameters
    ----------
    xarr : xarray.DataArray
        The function will add extra arguments to xarr.attrs
    **extra_kwargs
        not used

    Returns
    -------
    da : xarray.DataArray
    """
    if not(is_xarray(xarr)):
        raise TypeError('except a xarray.DataArray')
    for kwargs in extra_kwargs:
        xarr.attrs[kwargs] = extra_kwargs[kwargs]
    return xarr

def _grid_location_equals(xarr,grid_location=None):
    """Return True when the xarr grid_location attribute is grid_location

    Parameters
    ----------
    xarr : xarray.DataArray
       xarray dataarray which attributes should be tested.
    grid_location : str
        string describing the grid location : eg 'u','v','t','f'...

    Returns
    -------
    test : bool
        boolean value of the test.

    """
    test = True
    if xarr.attrs.has_key('grid_location'):
        test *= (xarr.attrs['grid_location']==grid_location)
    return test

def _di(scalararray):
    """Return the difference scalararray(i+1) - scalararray(i).

    A priori, for internal use only.

    Parameters
    ----------
    scalararray : xarray.DataArray
        xarray that should be differentiated.

    Returns
    -------
    di : xarray.DataArray
       xarray of difference, defined at point i+1/2
    """
    di = scalararray.shift(x=-1) - scalararray
    return di

def _dj(scalararray):
    """Return the difference scalararray(j+1) - scalararray(j)

    A priori, for internal use only.

    Parameters
    ----------
    scalararray : xarray.DataArray
        xarray that should be differentiated.

    Returns
    -------
    dj : xarray.DataArray
       xarray of difference, defined at point j+1/2
    """
    dj = scalararray.shift(y=-1) - scalararray
    return dj

def _mi(scalararray):
    """Return the average of scalararray(i+1) and scalararray(i)

    A priori, for internal use only.

    Parameters
    ----------
    scalararray : xarray.DataArray
        xarray that should be averaged at i+1/2.

    Returns
    -------
    mi : xarray.DataArray
       averaged xarray, defined at point i+1/2
    """
    mi = ( scalararray.shift(x=-1) + scalararray ) / 2.
    return mi

def _finalize_dataarray_attributes(xarr,**kwargs):
    """Update the dictionary of attibutes of a xarray dataarray.

    Parameters
    ----------
    xarr : xarray.DataArray
       xarray dataarray which attributes will be updated.
    kwargs : dict-like object
       dictionnary of attributes

    Returns
    -------
    xarr : xarray.DataArray
    """
    if isinstance(xarr, xr.DataArray):
        xarr.attrs.update(kwargs)
    if xarr.attrs.has_key('short_name'):
        xarr.name = xarr.attrs['short_name']
    return xarr

def chunk(self,chunks=None):
        """Rechunk all the variables defining the grid.

        Parameters
        ----------
        chunks : dict-like
            dictionnary of sizes of chunk along xarray dimensions.

        Example
        -------
        >>> xr_chunk = {'x':200,'y':200}
        >>> grd = generic_2d_grid(...)
        >>> grd.chunk(xr_chunk)

        """
        for dataname in self._arrays:
            data = self._arrays[dataname]
            if isinstance(data, xr.DataArray):
                self._arrays[dataname] = data.chunk(chunks)