Python xarray.where方法代码示例

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def _train(self, ref, hist):
        s_thresh = map_cdf(hist, ref, self.thresh, group=self.group).isel(
            x=0
        )  # Selecting the first threshold.
        # Compute scaling factor on wet-day intensity
        sth = broadcast(s_thresh, hist, group=self.group)
        ws = xr.where(hist >= sth, hist, np.nan)
        wo = xr.where(ref >= self.thresh, ref, np.nan)

        ms = self.group.apply("mean", ws, skipna=True)
        mo = self.group.apply("mean", wo, skipna=True)

        # Adjustment factor
        af = get_correction(ms - s_thresh, mo - self.thresh, MULTIPLICATIVE)
        af.attrs.update(long_name="LOCI adjustment factors")
        s_thresh.attrs.update(long_name="Threshold over modeled data")
        self._make_dataset(hist_thresh=s_thresh, ref_thresh=self.thresh, af=af) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def first_run_1d(arr: Sequence[Union[int, float]], window: int) -> int:
    """Return the index of the first item of a run of at least a given length.

    Parameters
    ----------
    arr : Sequence[Union[int, float]]
      Input array
    window : int
      Minimum duration of consecutive run to accumulate values.

    Returns
    -------
    int
      Index of first item in first valid run. Returns np.nan if there are no valid run.
    """
    v, rl, pos = rle_1d(arr)
    ind = np.where(v * rl >= window, pos, np.inf).min()

    if np.isinf(ind):
        return np.nan
    return ind 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def windowed_run_count_1d(arr: Sequence[bool], window: int) -> int:
    """Return the number of consecutive true values in array for runs at least as long as given duration.

    Parameters
    ----------
    arr : Sequence[bool]
      Input array (bool)
    window : int
      Minimum duration of consecutive run to accumulate values.

    Returns
    -------
    int
      Total number of true values part of a consecutive run at least `window` long.
    """
    v, rl = rle_1d(arr)[:2]
    return np.where(v * rl >= window, rl, 0).sum() 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def get_angles(self, vis):
        """Get the sun and satellite angles from the current dataarray."""
        from pyorbital.astronomy import get_alt_az, sun_zenith_angle
        from pyorbital.orbital import get_observer_look

        lons, lats = vis.attrs['area'].get_lonlats(chunks=vis.data.chunks)
        lons = da.where(lons >= 1e30, np.nan, lons)
        lats = da.where(lats >= 1e30, np.nan, lats)
        sunalt, suna = get_alt_az(vis.attrs['start_time'], lons, lats)
        suna = np.rad2deg(suna)
        sunz = sun_zenith_angle(vis.attrs['start_time'], lons, lats)

        sat_lon, sat_lat, sat_alt = get_satpos(vis)
        sata, satel = get_observer_look(
            sat_lon,
            sat_lat,
            sat_alt / 1000.0,  # km
            vis.attrs['start_time'],
            lons, lats, 0)
        satz = 90 - satel
        return sata, satz, suna, sunz 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def __call__(self, projectables, **kwargs):
        """Generate the composite."""
        data = projectables[0]

        # Default to rough IR thresholds
        # Values below or equal to this are clouds -> opaque white
        tr_min = self.transition_min
        # Values above this are cloud free -> transparent
        tr_max = self.transition_max
        # Gamma correction
        gamma = self.transition_gamma

        slope = 1 / (tr_min - tr_max)
        offset = 1 - slope * tr_min

        alpha = data.where(data > tr_min, 1.)
        alpha = alpha.where(data <= tr_max, 0.)
        alpha = alpha.where((data <= tr_min) | (data > tr_max), slope * data + offset)

        # gamma adjustment
        alpha **= gamma
        res = super(CloudCompositor, self).__call__((data, alpha), **kwargs)
        return res 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def read_dataset(fid, key):
    """Read dataset"""
    dsid = DSET_NAMES[key.name]
    dset = fid["/PWLR/" + dsid]
    if dset.ndim == 3:
        dims = ['y', 'x', 'level']
    else:
        dims = ['y', 'x']
    data = xr.DataArray(da.from_array(dset.value, chunks=CHUNK_SIZE),
                        name=key.name, dims=dims).astype(np.float32)
    data = xr.where(data > 1e30, np.nan, data)

    dset_attrs = dict(dset.attrs)
    data.attrs.update(dset_attrs)

    return data 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def calibrate_counts_to_rad(self, data, key):
        """Calibrate counts to radiances."""
        radiance_units = data.attrs["units"]
        if key.name == 'ir_38':
            data = xr.where(((2 ** 12 - 1 < data) & (data <= 2 ** 13 - 1)),
                            (data * data.attrs.get("warm_scale_factor", 1) +
                             data.attrs.get("warm_add_offset", 0)),
                            (data * data.attrs.get("scale_factor", 1) +
                             data.attrs.get("add_offset", 0))
                            )
        else:
            data = (data * data.attrs.get("scale_factor", 1) +
                    data.attrs.get("add_offset", 0))

        data.attrs["units"] = radiance_units

        return data 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def _categorize_chl(chl):
    return xr.where(chl >= 4., 2,
                    xr.where(chl >= 3.0, 1,
                             xr.where(chl >= 0.0, 0,
                                      np.nan))) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def test_valid_exprs(self):
        namespace = dict(a=np.array([0.1, 0.3, 0.1, 0.7, 0.4, 0.9]),
                         b=np.array([0.2, 0.1, 0.3, 0.2, 0.4, 0.8]),
                         np=np,
                         xr=xr)

        value = compute_array_expr('a + 1', namespace=namespace)
        np.testing.assert_array_almost_equal(value,
                                             np.array([1.1, 1.3, 1.1, 1.7, 1.4, 1.9]))

        value = compute_array_expr('a * b', namespace=namespace)
        np.testing.assert_array_almost_equal(value,
                                             np.array([0.02, 0.03, 0.03, 0.14, 0.16, 0.72]))

        value = compute_array_expr('max(a, b)', namespace=namespace)
        np.testing.assert_array_almost_equal(value,
                                             np.array([0.2, 0.3, 0.3, 0.7, 0.4, 0.9]))

        value = compute_array_expr('a > b', namespace=namespace)
        np.testing.assert_equal(value,
                                np.array([False, True, False, True, False, True]))

        value = compute_array_expr('a == b', namespace=namespace)
        np.testing.assert_equal(value,
                                np.array([False, False, False, False, True, False]))

        # This weirdo expression is a result of translating SNAP conditional expressions to Python.
        value = compute_array_expr('a > 0.35 if a else b', namespace=namespace)
        np.testing.assert_equal(value,
                                np.array([0.2, 0.1, 0.3, 0.7, 0.4, 0.9]))

        # We actually mean
        value = compute_array_expr('where(a > 0.35, a, b)', namespace=namespace)
        np.testing.assert_equal(value,
                                np.array([0.2, 0.1, 0.3, 0.7, 0.4, 0.9])) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def test_conditional(self):
        # The following conditional expr looks wrong but this is how it looks like after translating
        # from SNAP expression 'a >= 0.0 ? a : NaN'
        self.assertEqual(transpile_expr('a >= 0.0 if a else NaN'),
                         'xr.where(a >= 0.0, a, NaN)')
        self.assertEqual(transpile_expr('a >= 0.0 if a else b >= 0.0 if b else NaN'),
                         'xr.where(a >= 0.0, a, xr.where(b >= 0.0, b, NaN))') 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def test_where(self):
        self.assertEqual(transpile_expr('where(a >= 0.0, a, NaN)'),
                         'xr.where(a >= 0.0, a, NaN)')
        self.assertEqual(transpile_expr('xr.where(a >= 0.0, a, NaN)'),
                         'xr.where(a >= 0.0, a, NaN)')
        self.assertEqual(transpile_expr('np.where(a >= 0.0, a, NaN)'),
                         'np.where(a >= 0.0, a, NaN)')
        # xarray.DataArray.where() method:
        self.assertEqual(transpile_expr('a.where(a.x >= 0.0)'),
                         'a.where(a.x >= 0.0)') 

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

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def windowed_run_count(
    da: xr.DataArray,
    window: int,
    dim: str = "time",
    ufunc_1dim: Union[str, bool] = "auto",
) -> xr.DataArray:
    """Return the number of consecutive true values in array for runs at least as long as given duration.

    Parameters
    ----------
    da: xr.DataArray
      Input N-dimensional DataArray (boolean).
    window : int
      Minimum run length.
    dim : str
      Dimension along which to calculate consecutive run (default: 'time').
    ufunc_1dim : Union[str, bool]
      Use the 1d 'ufunc' version of this function : default (auto) will attempt to select optimal
      usage based on number of data points. Using 1D_ufunc=True is typically more efficient
      for dataarray with a small number of gridpoints.

    Returns
    -------
    xr.DataArray
      Total number of true values part of a consecutive runs of at least `window` long.
    """
    if ufunc_1dim == "auto":
        npts = get_npts(da)
        ufunc_1dim = npts <= npts_opt

    if ufunc_1dim:
        out = windowed_run_count_ufunc(da, window)
    else:
        d = rle(da, dim=dim)
        out = d.where(d >= window, 0).sum(dim=dim)
    return out 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def last_run_before_date(
    da: xr.DataArray,
    window: int,
    date: str = "07-01",
    dim: str = "time",
    coord: Optional[Union[bool, str]] = "dayofyear",
):
    """Return the index of the last item of the last run before a given date.

    Parameters
    ----------
    da : xr.DataArray
      Input N-dimensional DataArray (boolean)
    window : int
      Minimum duration of consecutive run to accumulate values.
    date : str
      The date before which to look for the last event.
    dim : str
      Dimension along which to calculate consecutive run (default: 'time').
    coord : Optional[Union[bool, str]]
      If not False, the function returns values along `dim` instead of indexes.
      If `dim` has a datetime dtype, `coord` can also be a str of the name of the
      DateTimeAccessor object to use (ex: 'dayofyear').

    Returns
    -------
    out : xr.DataArray
      Index (or coordinate if `coord` is not False) of last item in last valid run. Returns np.nan if there are no valid run.
    """
    before_date = datetime.strptime(date, "%m-%d").timetuple().tm_yday

    mid_idx = np.where(da.time.dt.dayofyear == before_date)[0]
    if mid_idx.size == 0:  # The date is not within the group. Happens at boundaries.
        return xr.full_like(da.isel(time=0), np.nan, float).drop_vars("time")

    run = da.where(da.time <= da.time[mid_idx][0])
    return last_run(run, window=window, dim=dim, coord=coord) 

# 需要导入模块: import xarray [as 别名]
# 或者: from xarray import where [as 别名]
def rle_1d(
    arr: Union[int, float, bool, Sequence[Union[int, float, bool]]]
) -> Tuple[np.array, np.array, np.array]:
    """Return the length, starting position and value of consecutive identical values.

    Parameters
    ----------
    arr : Sequence[Union[int, float, bool]]
      Array of values to be parsed.

    Returns
    -------
    values : np.array
      The values taken by arr over each run
    run lengths : np.array
      The length of each run
    start position : np.array
      The starting index of each run

    Examples
    --------
    >>> a = [1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3]
    >>> rle_1d(a)
    (array([1, 2, 3]), array([2, 4, 6]), array([0, 2, 6]))
    """
    ia = np.asarray(arr)
    n = len(ia)

    if n == 0:
        e = "run length array empty"
        warn(e)
        # Returning None makes some other 1d func below fail.
        return np.array(np.nan), 0, np.array(np.nan)

    y = np.array(ia[1:] != ia[:-1])  # pairwise unequal (string safe)
    i = np.append(np.where(y), n - 1)  # must include last element position
    rl = np.diff(np.append(-1, i))  # run lengths
    pos = np.cumsum(np.append(0, rl))[:-1]  # positions
    return ia[i], rl, pos