jejjohnson / jejjohnson/xrtoolz
feat(transforms): Layer-1 Operator wrappers for xrgrad calculus (Partial/Gradient/Divergence/Curl/Laplacian)
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- Avg merge
- 13d 21h
- Merged PRs (30d)
- 3
Description
## Problem / Request
The xrtoolz stack is built on a three-layer model — Layer 0 pure functions, Layer 1 `pipekit.Operator` subclasses, Layer 2 graphs — but the calculus family only exists at Layer 0 (`xrgrad.partial/gradient/divergence/curl/laplacian` are plain functions). Pipeline users can't drop a gradient into a `Sequential` today. The only calculus Operators that exist are domain-specific (`xrtoolz.ocn.Divergence` for spherical U/V, metrics' internal `DivergenceError`), so anything off the ocn happy path means writing a custom wrapper by hand.
## Working context
- **Where the code goes**: `packages/xrtoolz/src/xrtoolz/` (the main package), NOT `packages/xrtoolz-grad`. The workspace deliberately keeps `xrgrad` free of the `pipekit`/`xrcore` dependency (it's the lean, JAX-carrying calculus package); Operators belong in the main package, which already depends on both. Suggested location: a new `transforms/_src/calculus.py`, re-exported through `transforms/__init__.py` (the transforms family is the established home for domain-agnostic Operator wrappers; check its `__init__.py` for the re-export pattern and `__all__` sorting).
- Base class: `Operator` re-exported from `xrcore` (`from xrtoolz import Operator` or see how `transforms` modules import it). Every Operator subclass is a callable implementing `__call__(self, ds: xr.Dataset) -> xr.Dataset`, `get_config() -> dict`, and `__repr__()`. Study one existing operator in `packages/xrtoolz/src/xrtoolz/transforms/` end-to-end (e.g. any encoder/filter class) and copy its structure, docstring shape, and test layout exactly.
- Layer-0 functions to wrap (from `xrgrad`, already a dependency):
- `partial(da, dim, *, geometry, accuracy, method, **geom_kw) -> xr.DataArray`
- `gradient(da, *, dims, geometry, accuracy, method, **geom_kw) -> xr.Dataset` (vars named `d_d` / `_dx`/`_dy`)
- `divergence(ds, components, *, dims, geometry, ...) -> xr.DataArray`
- `curl(ds, (u, v), *, dims, geometry, ...) -> xr.DataArray`
- `laplacian(da, *, dims, geometry, ...) -> xr.DataArray`
- Constructor conventions (enforced house style): `variable=` (not `var=`) for a single input variable name, `variables=` for lists; singular `dim=` accepting `str | Sequence[str]` for dimensions-to-act-over; plural `dims=` only for inherently paired axes (the `(x, y)` pair of `curl` qualifies). Existing precedent: `ocn.Divergence` at `packages/xrtoolz/src/xrtoolz/ocn/operators.py:250` — read it, then keep the generic wrappers *thinner* (no spherical defaults baked in).
- Conventions: Google-style docstrings with Args/Returns/Example; type hints; surgical changes only. Dataset-in/Dataset-out — operators read named variables from the incoming Dataset and merge results back (match how sibling transforms operators handle output naming/overwrites).
- Verification: `cd packages/xrtoolz && uv run pytest tests/ -v`; repo root: `uv run --group lint ruff check .`, `uv run --group lint ruff format --check .`, `make typecheck`.
## Proposed API
```python
class Partial(Operator):
def __init__(self, variable: str, dim: str, *, geometry: str = "cartesian",
accuracy: int = 1, method: str = "central", out_name: str | None = None,
**geom_kw): ...
# __call__: ds -> ds with d_d (or out_name) added
class Gradient(Operator):
def __init__(self, variable: str, *, dim: str | Sequence[str] | None = None,
geometry: str = "cartesian", accuracy: int | tuple[int, ...] = 1,
method: str = "central", **geom_kw): ...
# merges the gradient Dataset's variables into ds
class Divergence(Operator):
def __init__(self, variables: Sequence[str], *, dim: Sequence[str],
geometry: str = "cartesian", out_name: str = "divergence", **kw): ...
class Curl(Operator):
def __init__(self, variables: tuple[str, str], *, dims: tuple[str, str],
geometry: str = "cartesian", out_name: str = "curl", **kw): ...
class Laplacian(Operator):
def __init__(self, variable: str, *, dim: str | Sequence[str] | None = None,
geometry: str = "cartesian", out_name: str | None = None, **kw): ...
```
Notes:
- Pure config state (strings/ints) → these serialise fine; do **not** set `forbid_in_yaml` (that flag is only for live non-serializable state).
- `get_config()` must round-trip every constructor argument including `geom_kw`.
- Name collision: if `ocn.Divergence` and the generic `Divergence` would both be re-exported at the top level, keep the generic ones exported from `xrtoolz.transforms` only (not the root `xrtoolz/__init__.py`) and note the relationship in both docstrings. Do not rename or deprecate `ocn.Divergence` in this issue.
- If sibling audit issues have landed (per-dim `accuracy` tuples, `periodic=`, mixed `geometry` tuples), pass those through transparently — the wrappers forward kwargs, they don't re-validate.
## Acceptance criteria
- [ ] All five Operators importable from `xrtoolz.transforms`; each runs inside a `pipekit.Sequential` chained after an existing transforms operator in a test.
- [ ] `__call__` output naming: defaults documented and tested (`d_d` etc.; `out_name` override respected); input Dataset is not mutated (compare `ds` before/after).
- [ ] `get_config()` round-trips (construct → get_config → reconstruct → identical `repr` and identical output on a fixture).
- [ ] Behaviour parity: each Operator's output equals the corresponding raw `xrgrad` call on the same fixture (one parity test per class).
- [ ] Docstrings with runnable Example blocks; `__repr__` follows the house pattern; lint/format/typecheck clean.
Part of the xrgrad gap-audit **Tier 3 (scope decisions)** epic.
Contributor guide
Research direction
Start with an existing operator under packages/xrtoolz/src/xrtoolz/transforms/, its __init__.py re-export pattern, and ocn/operators.py:250; then inspect the five xrgrad functions being wrapped. Add tests under packages/xrtoolz/tests/ covering Sequential use, naming, immutability, config round-trips, and parity with xrgrad, and run the listed pytest, lint, format, and typecheck commands. Done means all five operators are exported from xrtoolz.transforms with documented examples and clean verification.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- tooling
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100