Python xarray.broadcast() Examples

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 

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) 

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)) 

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))) 

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 

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) 

def varr_to_float32(varr):
    varr = varr.astype(np.float32, copy=False)
    varr_max = varr.max()
    varr_min = varr.min()
    varr, varr_min_bd = xr.broadcast(varr, varr_min)
    varr_norm = varr - varr_min_bd
    del varr_min_bd
    gc.collect()
    varr_norm, varr_denom = xr.broadcast(varr_norm, (varr_max - varr_min))
    varr_norm = varr_norm / varr_denom
    del varr_denom
    return varr_norm 

def rle(da: xr.DataArray, dim: str = "time", max_chunk: int = 1_000_000):
    n = len(da[dim])
    i = xr.DataArray(np.arange(da[dim].size), dims=dim).chunk({"time": 1})
    ind = xr.broadcast(i, da)[0].chunk(da.chunks)
    b = ind.where(~da)  # find indexes where false
    end1 = (
        da.where(b[dim] == b[dim][-1], drop=True) * 0 + n
    )  # add additional end value index (deal with end cases)
    start1 = (
        da.where(b[dim] == b[dim][0], drop=True) * 0 - 1
    )  # add additional start index (deal with end cases)
    b = xr.concat([start1, b, end1], dim)

    # Ensure bfill operates on entire (unchunked) time dimension
    # Determine appropraite chunk size for other dims - do not exceed 'max_chunk' total size per chunk (default 1000000)
    ndims = len(b.shape)
    chunk_dim = b[dim].size
    # divide extra dims into equal size
    # Note : even if calculated chunksize > dim.size result will have chunk==dim.size
    chunksize_ex_dims = None
    if ndims > 1:
        chunksize_ex_dims = np.round(np.power(max_chunk / chunk_dim, 1 / (ndims - 1)))
    chunks = dict()
    chunks[dim] = -1
    for dd in b.dims:
        if dd != dim:
            chunks[dd] = chunksize_ex_dims
    b = b.chunk(chunks)

    # back fill nans with first position after
    z = b.bfill(dim=dim)
    # calculate lengths
    d = z.diff(dim=dim) - 1
    d = d.where(d >= 0)
    return d 

def from_xy(cls,
                xy: Tuple[xr.DataArray, xr.DataArray],
                xy_names: Tuple[str, str] = None) -> 'GeoCoding':
        """
        Return new geo-coding for given *dataset*.

        :param xy: Tuple of x and y coordinate variables.
        :param xy_names: Optional tuple of the x- and y-coordinate variables in *dataset*.
        :return: The source dataset's geo-coding.
        """
        x, y = xy

        if xy_names is None:
            xy_names = x.name, y.name
        x_name, y_name = xy_names
        if x_name is None or y_name is None:
            raise ValueError(f'unable to determine x and y coordinate variable names')

        if x.ndim == 1 and y.ndim == 1:
            y, x = xr.broadcast(y, x)
        if x.ndim != 2 or y.ndim != 2:
            raise ValueError(
                f'coordinate variables {x_name!r} and {y_name!r} must both have either one or two dimensions')

        if x.shape != y.shape or x.dims != y.dims:
            raise ValueError(f"coordinate variables {x_name!r} and {y_name!r} must have same shape and dimensions")

        height, width = x.shape
        if width < 2 or height < 2:
            raise ValueError(f"size in each dimension of {x_name!r} and {y_name!r} must be greater two")

        is_geo_crs = _is_geo_crs(x_name, y_name)
        is_lon_normalized = False
        if is_geo_crs:
            x, is_lon_normalized = _maybe_normalise_2d_lon(x)

        return GeoCoding(x=x, y=y,
                         x_name=x_name,
                         y_name=y_name,
                         is_geo_crs=is_geo_crs,
                         is_lon_normalized=is_lon_normalized) 

def setUp(self) -> None:
        lon = xr.DataArray(np.linspace(10., 20., 11), dims='x')
        lat = xr.DataArray(np.linspace(50., 60., 11), dims='y')
        lat, lon = xr.broadcast(lat, lon)
        self.lon_values = lon.values
        self.lat_values = lat.values 

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)) 

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) 

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]) 

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 

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 

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) 

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],
    ) 

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 

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 

def distance(
    da: Union[xarray.DataArray, xarray.Dataset],
    *,
    lon: Union[float, Sequence[float], xarray.DataArray],
    lat: Union[float, Sequence[float], xarray.DataArray],
):
    """Return distance to a point in meters.

    Parameters
    ----------
    da : Union[xarray.DataArray, xarray.Dataset]
      Input data.
    lon : Union[float, Sequence[float], xarray.DataArray]
      Longitude coordinate.
    lat : Union[float, Sequence[float], xarray.DataArray]
      Latitude coordinate.

    Returns
    -------
    xarray.DataArray
      Distance in meters to point.

    Note
    ----
    To get the indices from closest point, use:
    >>> da = xr.open_dataset(path_to_pr_file).pr
    >>> d = distance(da, lon=-75, lat=45)
    >>> k = d.argmin()
    >>> i, j, _ = np.unravel_index(k, d.shape)
    """
    ptdim = lat.dims[0]

    g = Geod(ellps="WGS84")  # WGS84 ellipsoid - decent globally

    def func(lons, lats, lon, lat):
        return g.inv(lons, lats, lon, lat)[2]

    out = xarray.apply_ufunc(
        func,
        *xarray.broadcast(da.lon.load(), da.lat.load(), lon, lat),
        input_core_dims=[[ptdim]] * 4,
        output_core_dims=[[ptdim]],
    )
    out.attrs["units"] = "m"
    return out