Python xarray.broadcast方法代码示例

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def data_for_reg_calcs(values_for_reg_arr):
    lat = [-10., 1., 10., 20.]
    lon = [1., 10.]
    sfc_area = [0.5, 1., 0.5, 0.25]
    land_mask = [1., 1., 0., 1.]

    lat = xr.DataArray(lat, dims=[LAT_STR], coords=[lat])
    lon = xr.DataArray(lon, dims=[LON_STR], coords=[lon])
    sfc_area = xr.DataArray(sfc_area, dims=[LAT_STR], coords=[lat])
    land_mask = xr.DataArray(land_mask, dims=[LAT_STR], coords=[lat])

    sfc_area, _ = xr.broadcast(sfc_area, lon)
    land_mask, _ = xr.broadcast(land_mask, lon)

    da = xr.DataArray(values_for_reg_arr, coords=[lat, lon])
    da.coords[SFC_AREA_STR] = sfc_area
    da.coords[LAND_MASK_STR] = land_mask
    return da 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def test_m2c(PM_da_initialized_1d):
    """Test many-to-control (which can be any other one member) (m2c) comparison basic
    functionality.

    Clean comparison: Remove one control member from ensemble to use as reference.
    Take the remaining members as forecasts."""
    ds = PM_da_initialized_1d
    aforecast, areference = __m2c.function(ds, metric=metric)

    control_member = [0]
    reference = ds.isel(member=control_member).squeeze()
    # drop the member being reference
    ds_dropped = _drop_members(ds, removed_member=ds.member.values[control_member])
    forecast, reference = xr.broadcast(ds_dropped, reference)

    eforecast, ereference = forecast, reference
    # very weak testing on shape
    assert eforecast.size == aforecast.size
    assert ereference.size == areference.size

    assert_equal(eforecast, aforecast)
    assert_equal(ereference, areference) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def _m2c(ds, control_member=None, metric=None):
    """
    Compare all other members forecasts to control member verification.

    Args:
        ds (xarray object): xr.Dataset/xr.DataArray with member and ensemble
                            dimension.
        control_member: list of the one integer member serving as
                        reference. Default 0
        metric (Metric): if deterministic, forecast and reference both have member dim
                      if probabilistic, only forecast has member dim

    Returns:
        xr.object: forecast, reference.
    """
    if control_member is None:
        control_member = [0]
    reference = ds.isel(member=control_member).squeeze()
    # drop the member being reference
    forecast = _drop_members(ds, removed_member=ds.member.values[control_member])
    if not metric.probabilistic:
        forecast, reference = xr.broadcast(forecast, reference)
    elif 'member' in reference.coords:
        del reference['member']
    return forecast, reference 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def _test_ij_bbox_antimeridian(self, conservative: bool):
        def denorm(x):
            return x if x <= 180 else x - 360

        lon = xr.DataArray(np.linspace(175.0, 185.0, 21), dims='columns')
        lat = xr.DataArray(np.linspace(53.0, 58.0, 11), dims='rows')
        lat, lon = xr.broadcast(lat, lon)
        gc = GeoCoding(x=lon, y=lat, x_name='lon', y_name='lat', is_lon_normalized=True)
        ij_bbox = gc.ij_bbox_conservative if conservative else gc.ij_bbox
        self.assertEqual((-1, -1, -1, -1), ij_bbox((0, -50, 30, 0)))
        self.assertEqual((0, 0, 20, 10), ij_bbox((denorm(160), 50, denorm(200), 60)))
        self.assertEqual((0, 0, 20, 6), ij_bbox((denorm(160), 50, denorm(200), 56)))
        self.assertEqual((10, 0, 20, 6), ij_bbox((denorm(180), 50, denorm(200), 56)))
        self.assertEqual((10, 0, 16, 6), ij_bbox((denorm(180), 50, denorm(183), 56)))
        self.assertEqual((10, 1, 16, 6), ij_bbox((denorm(180), 53.5, denorm(183), 56)))
        self.assertEqual((8, 0, 18, 8), ij_bbox((denorm(180), 53.5, denorm(183), 56), ij_border=2))
        self.assertEqual((12, 1, 20, 6), ij_bbox((denorm(181), 53.5, denorm(200), 56)))
        self.assertEqual((12, 1, 18, 6), ij_bbox((denorm(181), 53.5, denorm(184), 56))) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def test_ij_bboxes(self):
        x = xr.DataArray(np.linspace(10.0, 20.0, 21), dims='x')
        y = xr.DataArray(np.linspace(53.0, 58.0, 11), dims='y')
        y, x = xr.broadcast(y, x)
        gc = GeoCoding(x=x, y=y, x_name='x', y_name='y', is_geo_crs=True)

        ij_bboxes = gc.ij_bboxes(np.array([(0.0, -50.0, 30.0, 0.0)]))
        np.testing.assert_almost_equal(ij_bboxes,
                                       np.array([(-1, -1, -1, -1)], dtype=np.int64))

        ij_bboxes = gc.ij_bboxes(np.array([(0.0, 50, 30, 60),
                                           (0.0, 50, 30, 56),
                                           (15, 50, 30, 56),
                                           (15, 50, 18, 56),
                                           (15, 53.5, 18, 56)]))
        np.testing.assert_almost_equal(ij_bboxes,
                                       np.array([(0, 0, 20, 10),
                                                 (0, 0, 20, 6),
                                                 (10, 0, 20, 6),
                                                 (10, 0, 16, 6),
                                                 (10, 1, 16, 6)], dtype=np.int64)) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def weights_3d(a_3d):
    weights = np.cos(np.deg2rad(a_3d.lat))
    _, weights = xr.broadcast(a_3d, weights)
    weights = weights.isel(time=0)
    return weights 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def _preprocess_weights(a, dim, new_dim, weights):
    """Preprocesses weights array to prepare for numpy computation.

    Parameters
    ----------
    a : xarray.Dataset or xarray.DataArray
        One of the arrays over which the function will be applied.
    dim : str, list
        The original dimension(s) to apply the function along.
    new_dim : str
        The newly named dimension after running ``_preprocess_dims``
    weights : xarray.Dataset or xarray.DataArray or None
        Weights to apply to function, matching the dimension size of
        ``new_dim``.
    """
    if weights is None:
        return None
    else:
        # Throw error if there are negative weights.
        if weights.min() < 0:
            raise ValueError(
                'Weights has a minimum below 0. Please submit a weights array '
                'of positive numbers.'
            )
        # Scale weights to vary from 0 to 1.
        weights = weights / weights.max()
        # Check that the weights array has the same size
        # dimension(s) as those being applied over.
        drop_dims = {k: 0 for k in a.dims if k not in new_dim}
        if dict(weights.sizes) != dict(a.isel(drop_dims).sizes):
            raise ValueError(
                f'weights dimension(s) {dim} of size {dict(weights.sizes)} '
                f"does not match DataArray's size "
                f'{dict(a.isel(drop_dims).sizes)}'
            )
        if dict(weights.sizes) != dict(a.sizes):
            # Broadcast weights to full size of main object.
            _, weights = xr.broadcast(a, weights)
        return weights 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [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],
    ) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def test_m2e(PM_da_initialized_1d):
    """Test many-to-ensemble-mean (m2e) comparison basic functionality.

    Clean comparison: Remove one member from ensemble to use as reference.
    Take the remaining members as forecasts."""
    ds = PM_da_initialized_1d
    aforecast, areference = __m2e.function(ds, metric=metric)

    reference_list = []
    forecast_list = []
    for m in ds.member.values:
        forecast = _drop_members(ds, removed_member=[m]).mean('member')
        reference = ds.sel(member=m).squeeze()
        forecast, reference = xr.broadcast(forecast, reference)
        forecast_list.append(forecast)
        reference_list.append(reference)
    reference = xr.concat(reference_list, 'member')
    forecast = xr.concat(forecast_list, 'member')
    forecast['member'] = np.arange(forecast.member.size)
    reference['member'] = np.arange(reference.member.size)

    eforecast, ereference = forecast, reference
    # very weak testing on shape
    assert eforecast.size == aforecast.size
    assert ereference.size == areference.size

    assert_equal(eforecast, aforecast)
    assert_equal(ereference, areference) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def test_m2m(PM_da_initialized_1d):
    """Test many-to-many (m2m) comparison basic functionality.

    Clean comparison: Remove one member from ensemble to use as reference. Take the
    remaining members as forecasts."""
    ds = PM_da_initialized_1d
    aforecast, areference = __m2m.function(ds, metric=metric)

    reference_list = []
    forecast_list = []
    for m in ds.member.values:
        forecast = _drop_members(ds, removed_member=[m])
        forecast['member'] = np.arange(1, 1 + forecast.member.size)
        reference = ds.sel(member=m).squeeze()
        forecast, reference = xr.broadcast(forecast, reference)
        reference_list.append(reference)
        forecast_list.append(forecast)
    supervector_dim = 'forecast_member'
    reference = xr.concat(reference_list, supervector_dim)
    forecast = xr.concat(forecast_list, supervector_dim)
    reference[supervector_dim] = np.arange(reference[supervector_dim].size)
    forecast[supervector_dim] = np.arange(forecast[supervector_dim].size)
    eforecast, ereference = forecast, reference
    # very weak testing here
    assert eforecast.size == aforecast.size
    assert ereference.size == areference.size 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def _m2m(ds, metric=None):
    """Compare all members to all others in turn while leaving out the verification
    ``member``.

    Args:
        ds (xarray object): xr.Dataset/xr.DataArray with ``member`` dimension.
        metric (Metric):
            If deterministic, forecast and reference have ``member`` dim.
            If probabilistic, only forecast has ``member`` dim.

    Returns:
        xr.object: forecast, reference.
    """
    reference_list = []
    forecast_list = []
    for m in ds.member.values:
        forecast = _drop_members(ds, removed_member=[m])
        # set incrementing members to avoid nans from broadcasting
        forecast['member'] = np.arange(1, 1 + forecast.member.size)
        reference = ds.sel(member=m).squeeze()
        # Tiles the singular "reference" member to compare directly to all other members
        if not metric.probabilistic:
            forecast, reference = xr.broadcast(forecast, reference)
        reference_list.append(reference)
        forecast_list.append(forecast)
    reference = xr.concat(reference_list, M2M_MEMBER_DIM)
    forecast = xr.concat(forecast_list, M2M_MEMBER_DIM)
    reference[M2M_MEMBER_DIM] = np.arange(reference[M2M_MEMBER_DIM].size)
    forecast[M2M_MEMBER_DIM] = np.arange(forecast[M2M_MEMBER_DIM].size)
    return forecast, reference 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def test_diff_interp_cubed_sphere(cs, cubed_sphere_connections):
    grid = Grid(cs, face_connections=cubed_sphere_connections)
    face, _ = xr.broadcast(cs.face, cs.data_c)

    face_diff_x = grid.diff(face, "X")
    np.testing.assert_allclose(face_diff_x[:, 0, 0], [-3, 1, 1, 1, 1, 2])
    np.testing.assert_allclose(face_diff_x[:, -1, 0], [-3, 1, 1, 1, 1, 2])

    face_diff_y = grid.diff(face, "Y")
    np.testing.assert_allclose(face_diff_y[:, 0, 0], [-4, -3, -2, -1, 2, 5])
    np.testing.assert_allclose(face_diff_y[:, 0, -1], [-4, -3, -2, -1, 2, 5]) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def test_is_geo_crs_and_is_lon_normalized(self):
        x = xr.DataArray(np.linspace(10.0, 20.0, 21), dims='columns', name='lon')
        y = xr.DataArray(np.linspace(53.0, 58.0, 11), dims='rows', name='lat')
        y, x = xr.broadcast(y, x)
        gc = GeoCoding(x, y)
        self.assertEqual(False, gc.is_geo_crs)
        self.assertEqual(False, gc.is_lon_normalized)
        gc = GeoCoding(x, y, is_geo_crs=True)
        self.assertEqual(True, gc.is_geo_crs)
        self.assertEqual(False, gc.is_lon_normalized)
        gc = GeoCoding(x, y, is_lon_normalized=True)
        self.assertEqual(True, gc.is_geo_crs)
        self.assertEqual(True, gc.is_lon_normalized) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def test_from_dataset_2d(self):
        x = xr.DataArray(np.linspace(10.0, 20.0, 21), dims='columns')
        y = xr.DataArray(np.linspace(53.0, 58.0, 11), dims='rows')
        y, x = xr.broadcast(y, x)
        gc = GeoCoding.from_dataset(xr.Dataset(dict(x=x, y=y)))
        self.assertIsInstance(gc.x, xr.DataArray)
        self.assertIsInstance(gc.y, xr.DataArray)
        self.assertEqual('x', gc.x_name)
        self.assertEqual('y', gc.y_name)
        self.assertEqual(False, gc.is_lon_normalized) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import broadcast [as 别名]
def _test_ij_bbox(self, conservative: bool):
        x = xr.DataArray(np.linspace(10.0, 20.0, 21), dims='x')
        y = xr.DataArray(np.linspace(53.0, 58.0, 11), dims='y')
        y, x = xr.broadcast(y, x)
        gc = GeoCoding(x=x, y=y, x_name='x', y_name='y', is_geo_crs=True)
        ij_bbox = gc.ij_bbox_conservative if conservative else gc.ij_bbox
        self.assertEqual((-1, -1, -1, -1), ij_bbox((0, -50, 30, 0)))
        self.assertEqual((0, 0, 20, 10), ij_bbox((0, 50, 30, 60)))
        self.assertEqual((0, 0, 20, 6), ij_bbox((0, 50, 30, 56)))
        self.assertEqual((10, 0, 20, 6), ij_bbox((15, 50, 30, 56)))
        self.assertEqual((10, 0, 16, 6), ij_bbox((15, 50, 18, 56)))
        self.assertEqual((10, 1, 16, 6), ij_bbox((15, 53.5, 18, 56)))
        self.assertEqual((8, 0, 18, 8), ij_bbox((15, 53.5, 18, 56), ij_border=2))