jejjohnson / jejjohnson/xrtoolz

OB-1.1: xrpatcher integration — PatchDataset + PatchedInference Operators

Open
#151 0 comments 0 reactions 0 assignees View on GitHub
area:code dependencies enhancement
Dominant language
Python
Stars
1
Forks
0
Avg merge
13d 21h
Merged PRs (30d)
3

Description

## Summary

Take `xrpatcher` ([jejjohnson/xrpatcher](https://github.com/jejjohnson/xrpatcher), MIT, ~470 LOC) as a hard top-level dep, re-export `XRDAPatcher`, and add two thin Layer-1 Operators (`PatchDataset`, `PatchedInference`) so tile-wise inference composes naturally with `Sequential`.

## Motivation

ML-based ocean diagnostics (super-resolution, infilling, sub-mesoscale emulation, neural SSH mapping) almost always run on bounded spatial tiles — both because GPU memory caps inference at modest patch sizes and because trained CNNs typically have fixed receptive fields. The standard pattern is patch → infer → reconstruct with overlap blending.

`xr_toolz`'s existing `ModelOp` is per-pixel `(N, F)` flatten-and-predict — the wrong shape for tile-based CNNs. `xrpatcher` already solves the patch-and-reconstruct problem cleanly: `__getitem__(i)` returns one patch, `reconstruct(items, weight=)` stitches with weighted overlap blending, plus runtime caching, domain-limit pre-selection, and full-scan validation.

We hard-dep on `xrpatcher` (the user owns both repos — release coordination is trivial; it's MIT, single dep `tqdm`) and add two Operators, plus a re-export, so tile inference composes with `Sequential`.

## API sketch

```python
# src/xr_toolz/patcher/__init__.py
from xrpatcher import XRDAPatcher
from xr_toolz.patcher._src.operators import PatchDataset, PatchedInference

# Operators
class PatchDataset(Operator):
"""Construct an XRDAPatcher from a Dataset/DataArray."""

def __init__(
self, *,
patches: dict[str, int],
strides: dict[str, int] | None = None,
domain_limits: dict | None = None,
check_full_scan: bool = False,
cache: bool = False,
preload: bool = False,
var: str | None = None,
): ...
def __call__(self, ds: xr.Dataset | xr.DataArray) -> XRDAPatcher: ...

class PatchedInference(Operator):
"""Run a callable per patch; reconstruct via the patcher."""

def __init__(
self, *,
model: Callable[[xr.DataArray], xr.DataArray | np.ndarray],
reconstruct_kwargs: dict | None = None,
progress: bool = True,
): ...
def __call__(self, patcher: XRDAPatcher) -> xr.DataArray: ...
```

Composition example:

```python
Sequential([
PatchDataset(patches={"lat": 64, "lon": 64},
strides={"lat": 32, "lon": 32}, var="ssh"),
PatchedInference(model=cnn,
reconstruct_kwargs={"weight": gaussian_weight}),
])(ds)
```

## Implementation notes

**`PatchDataset` dispatch**:

1. `xr.DataArray` → forward to `XRDAPatcher` directly.
2. `xr.Dataset`:
- `var is not None` → narrow to `ds[var]`.
- `var is None` and `len(ds.data_vars) == 1` → narrow to the single var.
- else → informative `ValueError` listing available vars and the required `var=` kwarg.
3. Construct `XRDAPatcher(da, patches, strides, domain_limits, check_full_scan, cache, preload)`.

**`PatchedInference`**:

```python
def __call__(self, patcher):
iterator = tqdm(patcher) if self.progress else iter(patcher)
outputs = [self.model(patch) for patch in iterator]
return patcher.reconstruct(outputs, **(self.reconstruct_kwargs or {}))
```

**Why two Operators, not `ModelOp.patcher=` retrofit**: per-pixel and tile paradigms shouldn't share an API. `ModelOp` reshape semantics (`(time, lat, lon, feature) → (N, feature)`) directly conflict with tile semantics `(C, H, W) → (C, H, W)`. Sibling abstraction in `xr_toolz.patcher` keeps `ModelOp` untouched. Users who want per-pixel models on tiles compose `PatchedInference(model=ModelOp(rf))`.

**Why hard-dep, not vendor**: the user owns both repos (release coordination trivial); xrpatcher MIT, ~470 LOC, single dep `tqdm`; vendoring would force version skew.

**Pin**: `xrpatcher>=0.x,<0.(x+1)` to insulate from upstream API drift.

**`get_config` round-trip**: when `model` is a Python lambda → emit `{"model": ""}` flag, non-roundtrippable. When `model` is an `Operator` → recurse into its `get_config()` for full round-trip.

## Acceptance criteria

- [ ] `xrpatcher` added as top-level dep in `pyproject.toml` with version pin.
- [ ] `from xr_toolz.patcher import XRDAPatcher` works (re-export).
- [ ] `PatchDataset` on `DataArray` returns `XRDAPatcher` with correct `__len__`; on single-var `Dataset` auto-narrows; on multi-var without `var=` raises informative `ValueError` listing data_vars; with `var="ssh"` narrows correctly.
- [ ] `PatchDataset` with `domain_limits`, `check_full_scan=True`, `cache=True, preload=True` all behave per `xrpatcher` semantics.
- [ ] `PatchedInference` identity model + uniform weight: reconstruct matches input exactly.
- [ ] With overlap (stride < patch): overlap regions averaged correctly.
- [ ] Custom Gaussian weight: weighted reconstruction matches manual weighted-average.
- [ ] Model returning ndarray vs DataArray: both shapes handled.
- [ ] `progress=False` → no tqdm bar.
- [ ] `Sequential([PatchDataset, PatchedInference])` end-to-end identity round-trip.
- [ ] `PatchDataset.get_config()` round-trips identically.

## Out of scope

- `RandomPatchSampler` for ML training data loaders — deferred (xrpatcher iterates deterministically).
- `ModelOp.patcher=` retrofit — declined.
- Vendoring `xrpatcher`.
- Backend-specific tile-inference helpers (PyTorch / JAX hooks) — keep abstraction backend-agnostic; user wraps tensor conversion inside the `model` callable.
- Patcher visualization helpers (`PatcherCoordsPanel`) — not blocking.

## Effort

≈80 LOC + ≈100 LOC tests; single PR.

Contributor guide

Open the contributing guide

Research direction

Start with pyproject.toml and the existing operator/config patterns, then add the package entry point in src/xr_toolz/patcher/__init__.py and implementations under src/xr_toolz/patcher/_src/ as specified. Build tests for dataset dispatch, reconstruction, configuration, progress handling, and Sequential composition; done means every listed acceptance criterion passes without changing ModelOp.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend-api-design, machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.