scikit-learn on xarray (xrops.learn)

The stack has one scikit-learn bridge: pipekit-array's Estimator, with to_rows / from_rows, which chooses samples by dim name on any carrier (geostack OP-O3). On a single DataArray in a pipeline, use it directly.

xrops.learn adds what is xarray's own. XarrayEstimator puts a Dataset's variables side by side as one feature table, each on its own grid, and fits a DataTree per node or pooled over its nodes. Importing xrops.learn registers the .sklearn accessor on DataArray, Dataset and DataTree. scikit-learn is imported when an estimator is fitted, so install the xrtoolz-ops[sklearn] extra.

XarrayEstimator

XarrayEstimator(estimator: Any, *, sample_dims: Hashable | Sequence[Hashable], feature_dim: str | None = None, nan_policy: NanPolicy = 'propagate', method: str = 'transform', tree_mode: TreeMode = 'per_node')

Bases: Operator

A scikit-learn-style estimator fitted and applied on xarray objects.

The sample dims are chosen by name; every other dim of a variable is a feature. A Dataset's variables are put side by side as the columns of one table, so each may have its own feature grid. Outputs come back as:

• one value per sample (predict): a DataArray over the sample dims;
• the fitted features (a scaler, inverse_transform): the input's variables on their own grids, as a DataArray or a Dataset;
• a new output dim (PCA scores, predict_proba): one DataArray with feature_dim where the first variable's feature dims were.

Parameters:

Name	Type	Description	Default
estimator	Any	Any object with scikit-learn's fit and the methods you call. It is cloned on fit; the original is never fitted.	required
sample_dims	Hashable | Sequence[Hashable]	The dim (or dims) whose entries are samples.	required
feature_dim	str | None	The name of a new output dim, for estimators whose output columns are not the input's features ("component" for PCA). predict_proba defaults to "class".	None
nan_policy	NanPolicy	"propagate", "raise", "mask_samples" or "mask_features", as in pipekit_array.Estimator.	'propagate'
method	str	What calling the operator runs once fitted: "transform", "predict" or "predict_proba".	'transform'
tree_mode	TreeMode	What fitting on a DataTree means: "per_node" fits one estimator per node (nodes may have different grids); "pool_samples" fits one on every node's samples (nodes share their variables and feature grids).	'per_node'

Attributes:

Name	Type	Description
estimator_	Any	The fitted scikit-learn estimator (None after a "per_node" DataTree fit; see estimators_).
estimators_	dict[str, XarrayEstimator] | None	{path: fitted XarrayEstimator} after a "per_node" DataTree fit.

Raises:

Type	Description
ValueError	An unknown tree_mode, nan_policy or method.

Example

Decompose a (time, lat, lon) cube with PCA and recover the grid:

>>> import numpy as np
>>> import xarray as xr
>>> from sklearn.decomposition import PCA
>>> rng = np.random.default_rng(0)
>>> da = xr.DataArray(rng.normal(size=(8, 3, 4)), dims=("time", "lat", "lon"))
>>> est = XarrayEstimator(PCA(3), sample_dims="time", feature_dim="component")
>>> scores = est.fit(da).transform(da)
>>> scores.dims, scores.shape
(('time', 'component'), (8, 3))
>>> est.inverse_transform(scores).dims
('time', 'lat', 'lon')
>>> est.estimator_.components_.shape
(3, 12)

Scale every variable of a Dataset on its own grid:

>>> from sklearn.preprocessing import StandardScaler
>>> ds = xr.Dataset({"u": da, "v": da.isel(lat=0)})
>>> scaler = XarrayEstimator(StandardScaler(), sample_dims="time")
>>> out = scaler.fit(ds).transform(ds)
>>> {name: out[name].dims for name in out.data_vars}
{'u': ('time', 'lat', 'lon'), 'v': ('time', 'lon')}

Source code in packages/xrtoolz-ops/src/xrops/learn/_src/estimator.py

def __init__(
    self,
    estimator: Any,
    *,
    sample_dims: Hashable | Sequence[Hashable],
    feature_dim: str | None = None,
    nan_policy: NanPolicy = "propagate",
    method: str = "transform",
    tree_mode: TreeMode = "per_node",
) -> None:
    if tree_mode not in get_args(TreeMode):
        raise ValueError(
            f"tree_mode must be one of {get_args(TreeMode)}; got {tree_mode!r}"
        )
    self.estimator = estimator
    self.sample_dims = (
        (sample_dims,)
        if isinstance(sample_dims, str) or not isinstance(sample_dims, Sequence)
        else tuple(sample_dims)
    )
    self.feature_dim = feature_dim
    self.nan_policy = nan_policy
    self.method = method
    self.tree_mode = tree_mode
    self._table_estimator()  # validates nan_policy and method now
    self.fitted_: Estimator | None = None
    self.layout_: Layout | None = None
    self.estimators_: dict[str, XarrayEstimator] | None = None

estimator_ property

estimator_: Any

The fitted scikit-learn estimator.

fit

fit(x: DataArray | Dataset | DataTree, y: Any = None) -> Self

Learn the estimator from x (and y, for a supervised one).

Parameters:

Name	Type	Description	Default
x	DataArray | Dataset | DataTree	A DataArray, Dataset or DataTree with the sample dims.	required
y	Any	Targets: a DataArray over the sample dims (plus at most one target dim); for a DataTree, a DataTree with one variable per node or {path: DataArray}.	None

Returns:

Type	Description
Self	This operator, fitted.

Raises:

Type	Description
ValueError	A variable lacks a sample dim, or ("pool_samples") the nodes do not share their feature grids.

Source code in packages/xrtoolz-ops/src/xrops/learn/_src/estimator.py

def fit(self, x: xr.DataArray | xr.Dataset | xr.DataTree, y: Any = None) -> Self:
    """Learn the estimator from ``x`` (and ``y``, for a supervised one).

    Args:
        x: A DataArray, Dataset or DataTree with the sample dims.
        y: Targets: a DataArray over the sample dims (plus at most one
            target dim); for a DataTree, a DataTree with one variable
            per node or ``{path: DataArray}``.

    Returns:
        This operator, fitted.

    Raises:
        ValueError: A variable lacks a sample dim, or (``"pool_samples"``)
            the nodes do not share their feature grids.
    """
    self.fitted_ = self.layout_ = self.estimators_ = None
    if isinstance(x, xr.DataTree):
        nodes = _nodes(x)
        if self.tree_mode == "per_node":
            self.estimators_ = {
                path: self._node_estimator().fit(ds, _node_target(y, path))
                for path, ds in nodes.items()
            }
            return self
        tables, targets, layout = [], [], None
        for path, ds in nodes.items():
            table, layout, samples = pack(ds, self.sample_dims, layout)
            tables.append(table)
            target = _node_target(y, path)
            if target is not None:
                targets.append(pack_target(target, samples))
        assert layout is not None
        self._fit_table(
            concat_tables(tables),
            concat_tables(targets) if targets else None,
            layout,
        )
        return self
    table, layout, samples = pack(x, self.sample_dims)
    self._fit_table(table, None if y is None else pack_target(y, samples), layout)
    return self

transform

transform(x: DataArray | Dataset | DataTree) -> Any

The estimator's transform of x, rebuilt as xarray.

Source code in packages/xrtoolz-ops/src/xrops/learn/_src/estimator.py

def transform(self, x: xr.DataArray | xr.Dataset | xr.DataTree) -> Any:
    """The estimator's ``transform`` of ``x``, rebuilt as xarray."""
    return self._run("transform", x)

predict

predict(x: DataArray | Dataset | DataTree) -> Any

One prediction per sample: a DataArray over the sample dims.

Source code in packages/xrtoolz-ops/src/xrops/learn/_src/estimator.py

def predict(self, x: xr.DataArray | xr.Dataset | xr.DataTree) -> Any:
    """One prediction per sample: a DataArray over the sample dims."""
    return self._run("predict", x)

predict_proba

predict_proba(x: DataArray | Dataset | DataTree) -> Any

Class probabilities per sample, along feature_dim ("class").

Source code in packages/xrtoolz-ops/src/xrops/learn/_src/estimator.py

def predict_proba(self, x: xr.DataArray | xr.Dataset | xr.DataTree) -> Any:
    """Class probabilities per sample, along ``feature_dim`` (``"class"``)."""
    return self._run("predict_proba", x)

inverse_transform

inverse_transform(z: DataArray | Dataset | DataTree) -> Any

Map transform's output back to the fitted variables and grids.

Parameters:

Name	Type	Description	Default
z	DataArray | Dataset | DataTree	transform's output: a DataArray with the sample dims and one output dim (scores); or, for a one-to-one transform (a scaler), the fitted variables on their grids, as a DataArray or Dataset; or a DataTree of either.	required

Returns:

Type	Description
Any	The fitted variables on their grids; features dropped by
Any	"mask_features" are NaN.

Source code in packages/xrtoolz-ops/src/xrops/learn/_src/estimator.py

def inverse_transform(self, z: xr.DataArray | xr.Dataset | xr.DataTree) -> Any:
    """Map ``transform``'s output back to the fitted variables and grids.

    Args:
        z: ``transform``'s output: a DataArray with the sample dims and
            one output dim (scores); or, for a one-to-one transform (a
            scaler), the fitted variables on their grids, as a DataArray
            or Dataset; or a DataTree of either.

    Returns:
        The fitted variables on their grids; features dropped by
        ``"mask_features"`` are NaN.
    """
    return self._run("inverse_transform", z)

fit_transform

fit_transform(x: DataArray | Dataset | DataTree, y: Any = None) -> Any

fit(x, y) then transform(x).

This is transform of the fitted estimator, not scikit-learn's fit_transform; the two differ for estimators that refit inside transform (NMF solves for W again, so its scores move by about 1e-3).

Source code in packages/xrtoolz-ops/src/xrops/learn/_src/estimator.py

def fit_transform(
    self, x: xr.DataArray | xr.Dataset | xr.DataTree, y: Any = None
) -> Any:
    """``fit(x, y)`` then ``transform(x)``.

    This is ``transform`` of the fitted estimator, not scikit-learn's
    ``fit_transform``; the two differ for estimators that refit inside
    ``transform`` (NMF solves for W again, so its scores move by about
    1e-3).
    """
    return self.fit(x, y).transform(x)

TreeMode module-attribute

TreeMode = Literal['per_node', 'pool_samples']