Source code for climpred.horizon

import numpy as np
import xarray as xr


def _last_item_cond_true(cond, dim):
    """Return the final item from cond that evaluates to True.

    Args:
        cond (xr.DataArray, xr.Dataset): User-defined boolean array where True means
            the system is predictable at the given lead. E.g., this could be based on
            the dynamical forecast beating a reference forecast, p values, confidence
            intervals, etc. cond should contain the dimension dim at the minimum.
        dim (str): Dimension to check for condition == True over.

    Returns:
        xr.DataArray, xr.Dataset: ``dim`` value until condition is True.

    """
    # force DataArray because isel (when transforming to dim space) requires DataArray
    if isinstance(cond, xr.Dataset):
        was_dataset = True
        cond = cond.to_array()
    else:
        was_dataset = False
    # index last True
    reached = cond.argmin(dim)
    # fix below one
    reached = reached.where(reached >= 1, np.nan)
    # reset where always true to len(lead)
    reached = reached.where(~cond.all("lead"), other=cond[dim].size)
    # fix locations where always nan to nan
    mask = cond.notnull().all("lead")
    reached = reached.where(mask, other=np.nan)
    # shift back into coordinate space
    # problem: cannot convert nan to idx in isel
    # therefore set to dim:0 and mask again afterwards
    reached_notnull = reached.notnull()  # remember where not masked
    reached = reached.where(
        reached.notnull(), other=cond.isel({dim: 0})
    )  # set nan to dim:0
    # take one index before calculated by argmin
    reached_dim_space = cond[dim].isel(
        {dim: reached.astype(int) - 1}
    )  # to not break conversion to dim space
    reached_dim_space = reached_dim_space.where(
        reached_notnull, other=np.nan
    )  # cleanup replace dim:0 with nan again
    if was_dataset:
        reached_dim_space = reached_dim_space.to_dataset(dim="variable").squeeze(
            drop=True
        )
        if "lead" in reached_dim_space.coords:
            reached_dim_space = reached_dim_space.drop_vars("lead")
    return reached_dim_space


[docs]def horizon(cond): """Calculate the predictability horizon based on a condition ```cond``. Args: cond (xr.DataArray, xr.Dataset): User-defined boolean array where True means the system is predictable at the given lead. E.g., this could be based on the dynamical forecast beating a reference forecast, p values, confidence intervals, etc. cond should contain the dimension lead at the minimum. Returns: xr.DataArray, xr.Dataset: predictability horizon reduced by ``lead`` dimension. Example: >>> skill = PerfectModelEnsemble.verify( ... metric="acc", ... comparison="m2e", ... dim=["init", "member"], ... reference=["persistence"], ... ) >>> horizon(skill.sel(skill="initialized") > skill.sel(skill="persistence")) <xarray.Dataset> Dimensions: () Data variables: tos float64 15.0 Attributes: units: years standard_name: forecast_period long_name: Lead description: Forecast period is the time interval between the forecast... >>> bskill = PerfectModelEnsemble.bootstrap( ... metric="acc", ... comparison="m2e", ... dim=["init", "member"], ... reference="uninitialized", ... iterations=201, ... resample_dim="init", ... ) >>> horizon(bskill.sel(skill="uninitialized", results="p") <= 0.05) <xarray.Dataset> Dimensions: () Coordinates: skill <U13 'uninitialized' results <U12 'p' Data variables: tos float64 2.0 Attributes: units: years standard_name: forecast_period long_name: Lead description: Forecast period is the time interval between the forecast... """ ph = _last_item_cond_true(cond, "lead") if isinstance(ph, xr.DataArray): ph.attrs["units"] = cond["lead"].attrs["units"] elif isinstance(ph, xr.Dataset): for v in ph.data_vars: ph[v].attrs["units"] = cond["lead"].attrs["units"] return ph