Python Dataset.assign_coords方法代码示例

# 需要导入模块: from xarray import Dataset [as 别名]
# 或者: from xarray.Dataset import assign_coords [as 别名]
def normalize_coord_vars(ds: xr.Dataset) -> xr.Dataset:
    """
    Turn potential coordinate variables from data variables into coordinate variables.

    Any data variable is considered a coordinate variable

    * whose name is its only dimension name;
    * whose number of dimensions is two and where the first dimension name is also a variable namd and
      whose last dimension is named "bnds".

    :param ds: The dataset
    :return: The same dataset or a shallow copy with potential coordinate
             variables turned into coordinate variables.
    """

    if 'bnds' not in ds.dims:
        return ds

    coord_var_names = set()
    for data_var_name in ds.data_vars:
        data_var = ds.data_vars[data_var_name]
        if is_coord_var(ds, data_var):
            coord_var_names.add(data_var_name)

    if not coord_var_names:
        return ds

    old_ds = ds
    ds = old_ds.drop(coord_var_names)
    ds = ds.assign_coords(**{bounds_var_name: old_ds[bounds_var_name] for bounds_var_name in coord_var_names})

    return ds

# 需要导入模块: from xarray import Dataset [as 别名]
# 或者: from xarray.Dataset import assign_coords [as 别名]
 def test_concat_coords(self):
     data = Dataset({"foo": ("x", np.random.randn(10))})
     expected = data.assign_coords(c=("x", [0] * 5 + [1] * 5))
     objs = [data.isel(x=slice(5)).assign_coords(c=0), data.isel(x=slice(5, None)).assign_coords(c=1)]
     for coords in ["different", "all", ["c"]]:
         actual = concat(objs, dim="x", coords=coords)
         self.assertDatasetIdentical(expected, actual)
     for coords in ["minimal", []]:
         with self.assertRaisesRegexp(ValueError, "not equal across"):
             concat(objs, dim="x", coords=coords)

# 需要导入模块: from xarray import Dataset [as 别名]
# 或者: from xarray.Dataset import assign_coords [as 别名]
 def test_concat_coords(self):
     data = Dataset({'foo': ('x', np.random.randn(10))})
     expected = data.assign_coords(c=('x', [0] * 5 + [1] * 5))
     objs = [data.isel(x=slice(5)).assign_coords(c=0),
             data.isel(x=slice(5, None)).assign_coords(c=1)]
     for coords in ['different', 'all', ['c']]:
         actual = concat(objs, dim='x', coords=coords)
         self.assertDatasetIdentical(expected, actual)
     for coords in ['minimal', []]:
         with self.assertRaisesRegexp(ValueError, 'not equal across'):
             concat(objs, dim='x', coords=coords)

# 需要导入模块: from xarray import Dataset [as 别名]
# 或者: from xarray.Dataset import assign_coords [as 别名]
def _normalize_lon_360(ds: xr.Dataset) -> xr.Dataset:
    """
    Fix the longitude of the given dataset ``ds`` so that it ranges from -180 to +180 degrees.

    :param ds: The dataset whose longitudes may be given in the range 0 to 360.
    :return: The fixed dataset or the original dataset.
    """

    if 'lon' not in ds.coords:
        return ds

    lon_var = ds.coords['lon']

    if len(lon_var.shape) != 1:
        return ds

    lon_size = lon_var.shape[0]
    if lon_size < 2:
        return ds

    lon_size_05 = lon_size // 2
    lon_values = lon_var.values
    if not np.any(lon_values[lon_size_05:] > 180.):
        return ds

    delta_lon = lon_values[1] - lon_values[0]

    var_names = [var_name for var_name in ds.data_vars]

    ds = ds.assign_coords(lon=xr.DataArray(np.linspace(-180. + 0.5 * delta_lon,
                                                       +180. - 0.5 * delta_lon,
                                                       lon_size),
                                           dims=ds['lon'].dims,
                                           attrs=dict(long_name='longitude',
                                                      standard_name='longitude',
                                                      units='degrees east')))

    ds = adjust_spatial_attrs_impl(ds, True)

    new_vars = dict()
    for var_name in var_names:
        var = ds[var_name]
        if len(var.dims) >= 1 and var.dims[-1] == 'lon':
            values = np.copy(var.values)
            temp = np.copy(values[..., : lon_size_05])
            values[..., : lon_size_05] = values[..., lon_size_05:]
            values[..., lon_size_05:] = temp
            # import matplotlib.pyplot as plt
            # im = values[(len(values.shape) - 2) * [0] + [slice(None), slice(None)]]
            # plt.imshow(im)
            new_vars[var_name] = xr.DataArray(values, dims=var.dims, attrs=var.attrs, encoding=var.encoding)

    return ds.assign(**new_vars)

# 需要导入模块: from xarray import Dataset [as 别名]
# 或者: from xarray.Dataset import assign_coords [as 别名]
def _normalize_lat_lon_2d(ds: xr.Dataset) -> xr.Dataset:
    """
    Detect 2D 'lat', 'lon' variables that span a equi-rectangular grid. Then:
    Drop original 'lat', 'lon' variables
    Rename original dimensions names of 'lat', 'lon' variables, usually ('y', 'x'), to ('lat', 'lon').
    Insert new 1D 'lat', 'lon' coordinate variables with dimensions 'lat' and 'lon', respectively.
    :param ds: some xarray dataset
    :return: a normalized xarray dataset, or the original one
    """
    if not ('lat' in ds and 'lon' in ds):
        return ds

    lat_var = ds['lat']
    lon_var = ds['lon']

    lat_dims = lat_var.dims
    lon_dims = lon_var.dims
    if lat_dims != lon_dims:
        return ds

    spatial_dims = lon_dims
    if len(spatial_dims) != 2:
        return ds

    x_dim_name = spatial_dims[-1]
    y_dim_name = spatial_dims[-2]

    lat_data_1 = lat_var[:, 0]
    lat_data_2 = lat_var[:, -1]
    lon_data_1 = lon_var[0, :]
    lon_data_2 = lon_var[-1, :]

    equal_lat = np.allclose(lat_data_1, lat_data_2, equal_nan=True)
    equal_lon = np.allclose(lon_data_1, lon_data_2, equal_nan=True)

    # Drop lat lon in any case. If note qual_lat and equal_lon subset_spatial_impl will subsequently
    # fail with a ValidationError

    ds = ds.drop(['lon', 'lat'])

    if not (equal_lat and equal_lon):
        return ds

    ds = ds.rename({
        x_dim_name: 'lon',
        y_dim_name: 'lat',
    })

    ds = ds.assign_coords(lon=np.array(lon_data_1), lat=np.array(lat_data_1))

    return ds

# 需要导入模块: from xarray import Dataset [as 别名]
# 或者: from xarray.Dataset import assign_coords [as 别名]
def normalize_missing_time(ds: xr.Dataset) -> xr.Dataset:
    """
    Add a time coordinate variable and their associated bounds coordinate variables
    if temporal CF attributes ``time_coverage_start`` and ``time_coverage_end``
    are given but the time dimension is missing.

    The new time coordinate variable will be named ``time`` with dimension ['time'] and shape [1].
    The time bounds coordinates variable will be named ``time_bnds`` with dimensions ['time', 'bnds'] and shape [1,2].
    Both are of data type ``datetime64``.

    :param ds: Dataset to adjust
    :return: Adjusted dataset
    """
    time_coverage_start = ds.attrs.get('time_coverage_start')
    if time_coverage_start is not None:
        # noinspection PyBroadException
        try:
            time_coverage_start = pd.to_datetime(time_coverage_start)
        except BaseException:
            pass

    time_coverage_end = ds.attrs.get('time_coverage_end')
    if time_coverage_end is not None:
        # noinspection PyBroadException
        try:
            time_coverage_end = pd.to_datetime(time_coverage_end)
        except BaseException:
            pass

    if not time_coverage_start and not time_coverage_end:
        # Can't do anything
        return ds

    if 'time' in ds:
        time = ds.time
        if not time.dims:
            ds = ds.drop('time')
        elif len(time.dims) == 1:
            time_dim_name = time.dims[0]
            is_time_used_as_dim = any([(time_dim_name in ds[var_name].dims) for var_name in ds.data_vars])
            if is_time_used_as_dim:
                # It seems we already have valid time coordinates
                return ds
            time_bnds_var_name = time.attrs.get('bounds')
            if time_bnds_var_name in ds:
                ds = ds.drop(time_bnds_var_name)
            ds = ds.drop('time')
            ds = ds.drop([var_name for var_name in ds.coords if time_dim_name in ds.coords[var_name].dims])

    if time_coverage_start or time_coverage_end:
        # noinspection PyBroadException
        try:
            ds = ds.expand_dims('time')
        except BaseException as e:
            warnings.warn(f'failed to add time dimension: {e}')

        if time_coverage_start and time_coverage_end:
            time_value = time_coverage_start + 0.5 * (time_coverage_end - time_coverage_start)
        else:
            time_value = time_coverage_start or time_coverage_end

        new_coord_vars = dict(time=xr.DataArray([time_value], dims=['time']))

        if time_coverage_start and time_coverage_end:
            has_time_bnds = 'time_bnds' in ds.coords or 'time_bnds' in ds
            if not has_time_bnds:
                new_coord_vars.update(time_bnds=xr.DataArray([[time_coverage_start, time_coverage_end]],
                                                             dims=['time', 'bnds']))

        ds = ds.assign_coords(**new_coord_vars)

        ds.coords['time'].attrs['long_name'] = 'time'
        ds.coords['time'].attrs['standard_name'] = 'time'
        ds.coords['time'].encoding['units'] = 'days since 1970-01-01'
        if 'time_bnds' in ds.coords:
            ds.coords['time'].attrs['bounds'] = 'time_bnds'
            ds.coords['time_bnds'].attrs['long_name'] = 'time'
            ds.coords['time_bnds'].attrs['standard_name'] = 'time'
            ds.coords['time_bnds'].encoding['units'] = 'days since 1970-01-01'

    return ds