Python xarray.set_options() Examples

def test_repr_HC(
    hind_ds_initialized_1d, hist_ds_uninitialized_1d, observations_ds_1d, display_style,
):
    """Test html repr."""
    with xr.set_options(display_style=display_style):
        he = HindcastEnsemble(hind_ds_initialized_1d)
        display(he)
        he = he.add_uninitialized(hist_ds_uninitialized_1d)
        display(he)
        he = he.add_observations(observations_ds_1d, 'OBS1')
        display(he)
        he = he.add_observations(observations_ds_1d, 'OBS2')
        display(he)
        # no uninit
        he = HindcastEnsemble(hind_ds_initialized_1d)
        he = he.add_observations(observations_ds_1d, 'OBS1')
        display(he) 

def test_repr_html(self):
        """Test if the function _repr_html is generating html."""
        idata = load_arviz_data("centered_eight")
        display_style = OPTIONS["display_style"]
        xr.set_options(display_style="html")
        html = idata._repr_html_()  # pylint: disable=protected-access

        assert html is not None
        assert "<div" in html
        for group in idata._groups:  # pylint: disable=protected-access
            assert group in html
            xr_data = getattr(idata, group)
            for item, _ in xr_data.items():
                assert item in html
        specific_style = ".xr-wrap{width:700px!important;}"
        assert specific_style in html

        xr.set_options(display_style="text")
        html = idata._repr_html_()  # pylint: disable=protected-access
        assert escape(idata.__repr__()) in html
        xr.set_options(display_style=display_style) 

def invert(x: xr.DataArray, kind: Optional[str] = None):
    """Invert a DataArray either additively (-x) or multiplicatively (1/x).

    If kind is not given, default to the one stored in the "kind" attribute of x.
    """
    kind = kind or x.get("kind", None)
    with xr.set_options(keep_attrs=True):
        if kind == ADDITIVE:
            return -x
        if kind == MULTIPLICATIVE:
            return 1 / x
        raise ValueError 

def cmip3_day_tas():
    # xr.set_options(enable_cftimeindex=False)
    ds = xr.open_dataset(
        os.path.join(TESTS_DATA, "cmip3", "tas.sresb1.giss_model_e_r.run1.atm.da.nc")
    )
    yield ds.tas
    ds.close() 

def mon_series(series, mon_triangular):
    def _mon_series(values, name):
        """Random time series whose mean varies over a monthly cycle."""
        x = series(values, name)
        m = mon_triangular
        factor = series(m[x.time.dt.month - 1], name)

        with xr.set_options(keep_attrs=True):
            return apply_correction(x, factor, x.kind)

    return _mon_series 

def isothermality(tasmin: xarray.DataArray, tasmax: xarray.DataArray, freq: str = "YS"):
    r"""Isothermality

    The mean diurnal range divided by the annual temperature range.

    Parameters
    ----------
    tasmin : xarray.DataArray
      Average daily minimum temperature [℃] or [K] at daily, weekly, or monthly frequency.
    tasmax : xarray.DataArray
      Average daily maximum temperature [℃] or [K] at daily, weekly, or monthly frequency.
    freq : str
      Resampling frequency; Defaults to "YS".

    Returns
    -------
    xarray.DataArray
      The isothermality value expressed as a percent.

    Notes
    -----
    According to the ANUCLIM user-guide https://fennerschool.anu.edu.au/files/anuclim61.pdf (ch. 6), input
    values should be at a weekly (or monthly) frequency.  However, the xclim.indices implementation here will calculate
    the output with input data with daily frequency as well. As such weekly or monthly input values, if desired, should
    be calculated prior to calling the function.
    """

    dtr = daily_temperature_range(tasmin=tasmin, tasmax=tasmax, freq=freq)
    etr = extreme_temperature_range(tasmin=tasmin, tasmax=tasmax, freq=freq)
    with xarray.set_options(keep_attrs=True):
        iso = dtr / etr * 100
        iso.attrs["units"] = "%"
    return iso 

def _adjust(self, sim, interp="linear"):
        sth = broadcast(self.ds.hist_thresh, sim, group=self.group, interp=interp)
        factor = broadcast(self.ds.af, sim, group=self.group, interp=interp)
        with xr.set_options(keep_attrs=True):
            scen = (factor * (sim - sth) + self.ds.ref_thresh).clip(min=0)
        return scen 

def __init__(self, **kwargs):
        self.old = xr.set_options(**kwargs).old 

def apply_correction(x: xr.DataArray, factor: xr.DataArray, kind: Optional[str] = None):
    """Apply the additive or multiplicative correction/adjustment factors.

    If kind is not given, default to the one stored in the "kind" attribute of factor.
    """
    kind = kind or factor.get("kind", None)
    with xr.set_options(keep_attrs=True):
        if kind == ADDITIVE:
            out = x + factor
        elif kind == MULTIPLICATIVE:
            out = x * factor
        else:
            raise ValueError
    return out 

def get_correction(x: xr.DataArray, y: xr.DataArray, kind: str):
    """Return the additive or multiplicative correction/adjustment factors."""
    with xr.set_options(keep_attrs=True):
        if kind == ADDITIVE:
            out = y - x
        elif kind == MULTIPLICATIVE:
            out = y / x
        else:
            raise ValueError("kind must be + or *.")

    if isinstance(out, xr.DataArray):
        out.attrs["kind"] = kind
    return out 

def test_repr_PM(PM_da_initialized_1d, PM_da_control_1d, display_style):
    """Test html and text repr."""
    with xr.set_options(display_style=display_style):
        pm = PerfectModelEnsemble(PM_da_initialized_1d)
        display(pm)
        pm = pm.add_control(PM_da_control_1d)
        display(pm)
        pm = pm.generate_uninitialized()
        display(pm) 

def __init__(self, data):
        super().__init__()
        xr.set_options(display_style='html')
        pn.extension(raw_css=[CSS])
        self.data = data
        self.panel = pn.pane.HTML(min_width=500, max_height=250, css_classes=['xrviz-scroll'])
        self.panel.object = "Description Section"
        self.setup() 

def __exit__(self, *exc):
        xr.set_options(**self.old)
        return 

def precip_seasonality(pr: xarray.DataArray,):
    r""" ANUCLIM Precipitation Seasonality (C of V)

    The annual precipitation Coefficient of Variation (C of V) expressed in percent. Calculated as the standard deviation
    of precipitation values for a given year expressed as a percentage of the mean of those values.


    Parameters
    ----------
    pr : xarray.DataArray
      Total precipitation rate at daily, weekly, or monthly frequency.
      Units need to be defined as a rate (e.g. mm d-1, mm week-1).

    Returns
    -------
    xarray.DataArray
      The coefficient of variation of precipitation values.

    Examples
    --------
    The following would compute for each grid cell of file `pr.day.nc` the annual precipitation seasonality:

    >>> import xclim.indices as xci
    >>> p = xr.open_dataset('pr.day.nc')
    >>> pday_seasonality = xci.precip_seasonality(p)

    >>> p_weekly = xci.precip_accumulation(p, freq='7D')
    >>> p_weekly.attrs['units'] = "mm/week" # input units need to be a rate
    >>> pweek_seasonality = xci.precip_seasonality(p_weekly)

    Notes
    -----
    According to the ANUCLIM user-guide https://fennerschool.anu.edu.au/files/anuclim61.pdf (ch. 6), input
    values should be at a weekly (or monthly) frequency.  However, the xclim.indices implementation here will calculate
    the result with input data with daily frequency as well. As such weekly or monthly input values, if desired,
    should be calculated prior to calling the function.

    If input units are in mm s-1 (or equivalent) values are converted to mm/day to avoid potentially small denominator
    values.

    """

    # If units in mm/sec convert to mm/days to avoid potentially small denominator
    if units2pint(pr) == units("mm / s"):
        pr = convert_units_to(pr, "mm d-1")

    with xarray.set_options(keep_attrs=True):
        seas = 100 * _anuclim_coeff_var(pr)

    seas.attrs["units"] = "percent"
    return seas 

def tg_mean_wetdry_quarter(
    tas: xarray.DataArray,
    pr: xarray.DataArray,
    op: str = None,
    input_freq: str = None,
    freq: str = "YS",
):
    r""" ANUCLIM Mean temperature of wettest/driest quarter

    The wettest (or driest) quarter of the year is determined, and the mean temperature of this period is calculated.
    If the input data frequency is "daily" or "weekly" quarters are defined as 13 week periods, otherwise are 3 months.

    Parameters
    ----------
    tas : xarray.DataArray
      Mean temperature [℃] or [K] at daily, weekly, or monthly frequency.
    pr : xarray.DataArray
      Total precipitation rate at daily, weekly, or monthly frequency.
    op : str {'wettest', 'driest'}
      Operation to perform: 'wettest' calculate for the wettest quarter; 'driest' calculate for the driest quarter.
    input_freq : str
      Input data time frequency - One of 'daily', 'weekly' or 'monthly'.
    freq : str
      Resampling frequency; Defaults to "YS".

    Returns
    -------
    xarray.DataArray
       Mean temperature values of the wettest/driest quarter of each year.

    Notes
    -----
    According to the ANUCLIM user-guide https://fennerschool.anu.edu.au/files/anuclim61.pdf (ch. 6), input
    values should be at a weekly (or monthly) frequency.  However, the xclim.indices implementation here will calculate
    the result with input data with daily frequency as well. As such weekly or monthly input values, if desired,
    should be calculated prior to calling the function.

    """
    tas_qrt = _to_quarter(input_freq, tas=tas)
    pr_qrt = _to_quarter(input_freq, pr=pr)

    xr_op = _xr_argops[op]
    with xarray.set_options(keep_attrs=True):
        out = _from_other_arg(criteria=pr_qrt, output=tas_qrt, op=xr_op, freq=freq)
        out.attrs = tas.attrs
        return out 

def prcptot_warmcold_quarter(
    pr: xarray.DataArray,
    tas: xarray.DataArray,
    op: str = None,
    input_freq: str = None,
    freq="YS",
):
    r""" ANUCLIM Total precipitation of warmest/coldest quarter

    The warmest (or coldest) quarter of the year is determined, and the total
    precipitation of this period is calculated.  If the input data frequency is "daily" or "weekly" quarters
    are defined as 13 week periods, otherwise are 3 months.

    Parameters
    ----------
    pr : xarray.DataArray
      Total precipitation rate at daily, weekly, or monthly frequency.
    tas : xarray.DataArray
      Mean temperature [℃] or [K] at daily, weekly, or monthly frequency.
    op : str
      Operation to perform: 'warmest' calculate for the warmest quarter ; 'coldest' calculate for the coldest quarter.
    input_freq : str
      Input data time frequency - One of 'daily', 'weekly' or 'monthly'.
    freq : str
      Resampling frequency; Defaults to "YS".

    Returns
    -------
    xarray.DataArray
       Total precipitation values of the warmest/coldest quarter of each year.

    Notes
    -----
    According to the ANUCLIM user-guide https://fennerschool.anu.edu.au/files/anuclim61.pdf (ch. 6), input
    values should be at a weekly (or monthly) frequency.  However, the xclim.indices implementation here will calculate
    the result with input data with daily frequency as well. As such weekly or monthly input values, if desired,
    should be calculated prior to calling the function.

    """
    # determine input data frequency
    tas_qrt = _to_quarter(input_freq, tas=tas)
    pr_qrt = _to_quarter(input_freq, pr=pr)

    xr_op = _xr_argops[op]
    with xarray.set_options(keep_attrs=True):
        out = _from_other_arg(criteria=tas_qrt, output=pr_qrt, op=xr_op, freq=freq)
        out.attrs = pr_qrt.attrs
        return out 

def _to_quarter(freq, pr=None, tas=None):
    """Convert daily, weekly or monthly time series to quarterly time series according to ANUCLIM specifications."""

    if freq.upper().startswith("D"):
        if tas is not None:
            tas = tg_mean(tas, freq="7D")

        if pr is not None:
            pr = precip_accumulation(pr, freq="7D")
            pr.attrs["units"] = "mm/week"

        freq = "W"

    if freq.upper().startswith("W"):
        window = 13
        u = units.week

    elif freq.upper().startswith("M"):
        window = 3
        u = units.month

    else:
        raise NotImplementedError(
            f'Unknown input time frequency "{freq}": must be one of "daily", "weekly" or "monthly".'
        )

    if tas is not None:
        tas = ensure_chunk_size(tas, time=np.ceil(window / 2))
    if pr is not None:
        pr = ensure_chunk_size(pr, time=np.ceil(window / 2))

    with xarray.set_options(keep_attrs=True):
        if pr is not None:
            pr = pint_multiply(pr, 1 * u, "mm")
            out = pr.rolling(time=window, center=False).sum()
            out.attrs = pr.attrs
            out.attrs["units"] = "mm"

        if tas is not None:
            out = tas.rolling(time=window, center=False).mean(
                allow_lazy=True, skipna=False
            )
            out.attrs = tas.attrs

    out = ensure_chunk_size(out, time=-1)
    return out