Sklearn patch-wise Recipe¶
This recipe shows how XarrayEstimator composes with patch-wise xarray workflows: fit local sklearn estimators independently on each spatial patch, or fit one global estimator and run memory-bounded inference on patches that preserve the fitted feature axis.
What you'll learn:
- How per-patch fitting works when each spatial tile has its own local statistics.
- Why global-fit / patch-infer workflows must patch along the sample axis so every patch keeps the same feature layout seen during fitting.
- Where
XRDAPatcherslots into the loop without any additional xrtoolz operator.
In [1]:
Copied!
import subprocess
import sys
try:
import google.colab # noqa: F401
IN_COLAB = True
except ImportError:
IN_COLAB = False
if IN_COLAB:
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
"-q",
"xrtoolz @ git+https://github.com/jejjohnson/xrtoolz@main",
],
check=True,
)
import subprocess
import sys
try:
import google.colab # noqa: F401
IN_COLAB = True
except ImportError:
IN_COLAB = False
if IN_COLAB:
subprocess.run(
[
sys.executable,
"-m",
"pip",
"install",
"-q",
"xrtoolz @ git+https://github.com/jejjohnson/xrtoolz@main",
],
check=True,
)
In [2]:
Copied!
import importlib.util
import warnings
import numpy as np
import xarray as xr
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from xrops.learn import XarrayEstimator
# Render xarray reprs as HTML and keep every section collapsed by
# default — readers expand only what they need; rendered notebooks stay
# tidy on the docs site.
_ = xr.set_options(
display_style="html",
display_expand_data=False,
display_expand_attrs=False,
display_expand_coords=False,
display_expand_data_vars=False,
display_expand_groups=False,
display_expand_indexes=False,
)
warnings.filterwarnings("ignore", message=r".*IProgress.*")
try:
from IPython import get_ipython
ipython = get_ipython()
except ImportError:
ipython = None
if ipython is not None and importlib.util.find_spec("watermark") is not None:
ipython.run_line_magic("load_ext", "watermark")
ipython.run_line_magic("watermark", "-v -m -p numpy,xarray,sklearn,xrops")
else:
print("watermark extension not installed; skipping reproducibility readout.")
import importlib.util
import warnings
import numpy as np
import xarray as xr
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from xrops.learn import XarrayEstimator
# Render xarray reprs as HTML and keep every section collapsed by
# default — readers expand only what they need; rendered notebooks stay
# tidy on the docs site.
_ = xr.set_options(
display_style="html",
display_expand_data=False,
display_expand_attrs=False,
display_expand_coords=False,
display_expand_data_vars=False,
display_expand_groups=False,
display_expand_indexes=False,
)
warnings.filterwarnings("ignore", message=r".*IProgress.*")
try:
from IPython import get_ipython
ipython = get_ipython()
except ImportError:
ipython = None
if ipython is not None and importlib.util.find_spec("watermark") is not None:
ipython.run_line_magic("load_ext", "watermark")
ipython.run_line_magic("watermark", "-v -m -p numpy,xarray,sklearn,xrops")
else:
print("watermark extension not installed; skipping reproducibility readout.")
Python implementation: CPython Python version : 3.13.12 IPython version : 9.16.1 numpy : 2.5.2 xarray : 2026.7.0 sklearn: 1.9.0 xrops : 0.4.0 Compiler : Clang 21.1.4 OS : Linux Release : 6.8.0-1044-azure Machine : x86_64 Processor : x86_64 CPU cores : 16 Architecture: 64bit
Synthetic sea-surface-height cube¶
The examples use a tiny time × lat × lon DataArray so the patch shapes are visible. Real workflows replace the hand-written isel slices below with XRDAPatcher(da=ssh, patches=..., strides=...).
In [3]:
Copied!
time = np.arange(6)
lat = np.linspace(-30.0, 30.0, 4)
lon = np.linspace(140.0, 143.0, 4)
values = (
np.sin(time[:, None, None] / 2.0)
+ lat[None, :, None] / 100.0
+ lon[None, None, :] / 1000.0
)
ssh = xr.DataArray(
values,
dims=("time", "lat", "lon"),
coords={"time": time, "lat": lat, "lon": lon},
name="ssh",
attrs={"units": "m"},
)
ssh
time = np.arange(6)
lat = np.linspace(-30.0, 30.0, 4)
lon = np.linspace(140.0, 143.0, 4)
values = (
np.sin(time[:, None, None] / 2.0)
+ lat[None, :, None] / 100.0
+ lon[None, None, :] / 1000.0
)
ssh = xr.DataArray(
values,
dims=("time", "lat", "lon"),
coords={"time": time, "lat": lat, "lon": lon},
name="ssh",
attrs={"units": "m"},
)
ssh
Out[3]:
<xarray.DataArray 'ssh' (time: 6, lat: 4, lon: 4)> Size: 768B -0.16 -0.159 -0.158 -0.157 0.04 0.041 ... 0.8405 0.8415 1.038 1.039 1.04 1.041 Coordinates: (3) Attributes: (1)
Pattern 1 — per-patch local fit¶
Use this pattern when each patch should learn its own local statistics, such as local scaling before a regional EOF/PCA. With XRDAPatcher, the loop shape is:
from xrpatcher import XRDAPatcher
patcher = XRDAPatcher(da=ssh, patches={"lat": 64, "lon": 64}, strides={"lat": 64, "lon": 64})
scaled = [
XarrayEstimator(StandardScaler(), sample_dims="time").fit_transform(patcher[i])
for i in range(len(patcher))
]
out = patcher.reconstruct(scaled)
In [4]:
Copied!
spatial_slices = [
{"lat": slice(0, 2), "lon": slice(0, 2)},
{"lat": slice(0, 2), "lon": slice(2, 4)},
{"lat": slice(2, 4), "lon": slice(0, 2)},
{"lat": slice(2, 4), "lon": slice(2, 4)},
]
local_scaled_patches = [
XarrayEstimator(StandardScaler(), sample_dims="time").fit_transform(ssh.isel(**slc))
for slc in spatial_slices
]
local_scaled = xr.combine_by_coords(local_scaled_patches)["ssh"]
local_scaled
spatial_slices = [
{"lat": slice(0, 2), "lon": slice(0, 2)},
{"lat": slice(0, 2), "lon": slice(2, 4)},
{"lat": slice(2, 4), "lon": slice(0, 2)},
{"lat": slice(2, 4), "lon": slice(2, 4)},
]
local_scaled_patches = [
XarrayEstimator(StandardScaler(), sample_dims="time").fit_transform(ssh.isel(**slc))
for slc in spatial_slices
]
local_scaled = xr.combine_by_coords(local_scaled_patches)["ssh"]
local_scaled
Out[4]:
<xarray.DataArray 'ssh' (time: 6, lat: 4, lon: 4)> Size: 768B -1.897 -1.897 -1.897 -1.897 -1.897 ... -0.1167 -0.1167 -0.1167 -0.1167 -0.1167 Coordinates: (3) Attributes: (1)
Pattern 2 — global fit, patch-wise inference¶
A fitted sklearn estimator expects the same feature count during inference that it saw during fitting. For global-fit / patch-infer workflows, patch along the sample axis and keep the feature axis intact. Here the grid is stacked to point × time, PCA is fit globally over all points, and inference runs over point patches that all keep the same time features.
In [5]:
Copied!
point_series = ssh.stack(point=("lat", "lon")).transpose("point", "time")
global_pca = XarrayEstimator(
PCA(n_components=2), sample_dims="point", feature_dim="component"
).fit(point_series)
point_slices = [slice(0, 8), slice(8, 16)]
score_patches = [
global_pca.transform(point_series.isel(point=slc)) for slc in point_slices
]
global_scores = xr.concat(score_patches, dim="point").sortby("point")
global_scores
point_series = ssh.stack(point=("lat", "lon")).transpose("point", "time")
global_pca = XarrayEstimator(
PCA(n_components=2), sample_dims="point", feature_dim="component"
).fit(point_series)
point_slices = [slice(0, 8), slice(8, 16)]
score_patches = [
global_pca.transform(point_series.isel(point=slc)) for slc in point_slices
]
global_scores = xr.concat(score_patches, dim="point").sortby("point")
global_scores
Out[5]:
<xarray.DataArray (point: 16, component: 2)> Size: 256B -0.7385 1.943e-16 -0.7361 2.082e-16 ... 0.7361 1.943e-16 0.7385 2.498e-16 Coordinates: (3) Dimensions without coordinates: component
The equivalent XRDAPatcher recipe patches the stacked sample dimension and reconstructs the component scores:
from xrpatcher import XRDAPatcher
point_series = ssh.stack(point=("lat", "lon")).transpose("point", "time")
pca = XarrayEstimator(PCA(n_components=10), sample_dims="point", feature_dim="component").fit(point_series)
patcher = XRDAPatcher(da=point_series, patches={"point": 4096}, strides={"point": 4096})
score_patches = [pca.transform(patcher[i]) for i in range(len(patcher))]
scores = patcher.reconstruct(score_patches)
The important invariant is that each patch still has the fitted feature axis (time here). No new xrtoolz operator is needed because XRDAPatcher yields DataArray patches and XarrayEstimator consumes DataArray inputs directly.
Where this lands once OB-1.1 ships¶
The patch loops above use hand-written isel slices to keep this notebook executable without extra dependencies. The actual XRDAPatcher integration (operator wrappers, overlap-blended reconstruct, runtime caching) is tracked under issue OB-1.1 (#151). After that lands, the loops below shorten to:
from xrtoolz.patcher import PatchDataset, PatchedInference
Sequential([
PatchDataset(patches={"lat": 64, "lon": 64}, var="ssh"),
PatchedInference(model=lambda da: scaler.fit_transform(da)),
])(ds)
Until then, the manual isel patterns shown here capture the same compositional shape on xr.DataArray only — no patcher dependency.