xscalers 0.1.0__tar.gz

xscalers-0.1.0/PKG-INFO

@@ -0,0 +1,544 @@
+Metadata-Version: 2.4
+Name: xscalers
+Version: 0.1.0
+Summary: sklearn-shaped, dim-aware scaling transformers for xarray and pytensor.xtensor
+Keywords: xarray,scaling,preprocessing,pymc,pytensor,panel-data,time-series
+Author: Will Dean
+Author-email: Will Dean <wd60622@gmail.com>
+License-Expression: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Typing :: Typed
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: POSIX
+Classifier: Operating System :: Unix
+Classifier: Operating System :: MacOS
+Requires-Dist: numpy>=2.0
+Requires-Dist: xarray>=2024.10.0
+Requires-Dist: pytensor>=3.0 ; extra == 'pytensor'
+Requires-Python: >=3.12
+Project-URL: Bug Tracker, https://github.com/williambdean/xscalers/issues
+Project-URL: Homepage, https://github.com/williambdean/xscalers
+Project-URL: Repository, https://github.com/williambdean/xscalers
+Provides-Extra: pytensor
+Description-Content-Type: text/markdown
+
+# xscalers
+
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![Tests](https://github.com/williambdean/xscalers/actions/workflows/ci.yaml/badge.svg)](https://github.com/williambdean/xscalers/actions/workflows/ci.yaml)
+[![PyPI version](https://badge.fury.io/py/xscalers.svg)](https://badge.fury.io/py/xscalers)
+
+sklearn-shaped, dim-aware scaling transformers for `xarray` and `pytensor.xtensor`.
+
+**Fit is eager xarray** — the fitted state is an `xr.Dataset` you can inspect
+and edit — and **transform is pure arithmetic by dimension name** that runs
+identically on `DataArray`s and symbolic `pytensor.xtensor` variables. All v1
+scalers are affine, so model code binds only to a composed `(center, scale)`
+pair. Nothing custom: every container and serialization format is a native
+ecosystem object. Every example below is executed against this README by the
+test suite, so it stays copy-paste-runnable.
+
+## Why
+
+Hand-writing `(x - mean) / std` is fine until the regression's inputs are an
+`xr.Dataset` with mixed grains and the model is PyMC. Three familiar pains:
+
+- sklearn's transformers assume stacked 2-D arrays and hide fitted state in
+  `mean_`-style attributes, which neither preserves dims nor hands you the
+  fitted constants;
+- the scaling constants fed into a `pytensor.xtensor` graph are usually
+  hand-copied, so converting a model contribution back to real terms goes
+  silently wrong;
+- xarray itself ships no scaling transformer.
+
+xscalers fills exactly that gap: affine scaling over dimension names on
+native xarray objects, fitted state as an inspectable `xr.Dataset` artifact,
+the same arithmetic lifting into a symbolic `pytensor.xtensor` graph, and
+one `unscale` verb that converts model contributions back to real units,
+refusing where the algebra is not affine instead of guessing.
+
+The frame is a regression loop you can toggle: point the same pipeline at a
+different transformation (`MeanScaler` vs `ZScore` vs `FixedScaler`, per
+variable) and regenerate the model-ready `xr.Dataset` without touching model
+code, then invert predictions and contributions back to real terms with the
+same fitted state.
+
+## The contract
+
+```text
+class BaseTransformer:
+    dims: tuple[str, ...] | None                       # dims to reduce over; None = all
+    def fit(self, X) -> xr.Dataset                     # fitted state, inspectable + editable
+    def transform(self, X, state) -> X                 # state: Mapping[str, DataArray]
+    def inverse_transform(self, X, state) -> X
+```
+
+Dataset and DataTree inputs are handled per variable — each variable keeps
+its own remaining dims, which is what makes mixed-grain inputs work with no
+special casing.
+
+## The basics
+
+```python
+import xarray as xr
+
+from xscalers.datasets import load_mixed_grain, load_national
+from xscalers.preprocessing import MeanScaler
+
+national: xr.Dataset = load_national()
+
+scaler = MeanScaler(("date",))
+state = scaler.fit(national["media"])
+scaled = scaler.transform(national["media"], state)
+
+state
+```
+
+```text
+<xarray.Dataset> Size: 96B
+Dimensions:  (channel: 3)
+Coordinates:
+  * channel  (channel) <U6 72B 'search' 'social' 'fb'
+Data variables:
+    mean     (channel) float64 24B 49.68 48.81 53.79
+Attributes:
+    kind:         mean
+    dims:         ['date']
+    scaled_from:  media
+    spec_sha:     af4a187be0e9
+    recipe:       MeanScaler(dims=['date'])
+```
+
+State is a value on the record — inspect it, save it, pass it to another
+process. Hand-editing is a manual override (the stamp records the recipe,
+not your edits).
+
+Mixed grains need no special casing: `fit` a Dataset and each variable gets
+its own state:
+
+```python
+import xarray as xr
+
+from xscalers.compose import VariableTransformer
+from xscalers.datasets import load_mixed_grain
+from xscalers.preprocessing import FixedScaler, MeanScaler, ZScore
+
+mixed: xr.Dataset = load_mixed_grain()
+transformer = VariableTransformer(
+    {
+        "media_national": MeanScaler(dims=("date",)),
+        "media_digital": ZScore(dims=("date",)),
+        "media_tv": FixedScaler(scale=1_000_000),
+    },
+    remainder="passthrough",
+)
+state = transformer.fit(mixed)              # one Dataset state per variable
+scaled = transformer.transform(mixed, state)  # each var scaled by its own recipe
+```
+
+## Derived metrics — new variables, even across grains
+
+`Derive` turns selected pieces of existing variables into a *new* variable
+that later steps treat like anything else. The two pieces don't have to come
+from the same variable or the same grain — a national series over a
+product-grain control broadcasts up to product grain, so `tv_per_control`
+below is a product-grain variable built from a national input.
+
+Ops are a closed, spec-serializable vocabulary — `op="ratio | diff | product"`
+— and selector pieces take `reduce="sum | mean | first | last"` for
+share-of-total style derivations:
+
+```python
+import xarray as xr
+
+from xscalers.compose import VariableTransformer
+from xscalers.datasets import load_mixed_grain
+from xscalers.pipeline import Pipeline
+from xscalers.preprocessing import Derive, FunctionTransformer, Splitter, ZScore
+
+mixed: xr.Dataset = load_mixed_grain()
+mixed["media_national"].dims                  # (date, media) — national grain
+mixed["controls_prod"].dims                    # (date, product) — product grain
+
+pipe = Pipeline(
+    [
+        (
+            "derive",
+            Derive(
+                name="tv_per_control",
+                left=("media_national", {"media": "tv"}),
+                right=("controls_prod", {}),   # {} = the whole variable
+                op="ratio",
+            ),
+        ),
+        ("log", FunctionTransformer({"tv_per_control": "log"})),  # named funcs
+        ("scale", VariableTransformer({"tv_per_control": ZScore(dims=("date",))})),
+    ]
+)
+state = pipe.fit(mixed)
+scaled = pipe.transform(mixed, state)
+scaled["tv_per_control"].dims                  # (date, product) — broadcast, then scaled
+
+back = pipe.inverse_transform(scaled, state)
+import numpy as np
+np.testing.assert_allclose(back["controls_prod"], mixed["controls_prod"])
+```
+
+Because the z-score reduces over `date` only, each product keeps its own
+mean and standard deviation — the per-variable, per-group treatment that
+mixed-grain data needs, with no reshaping.
+
+`inverse_transform` drops the derived variable and restores the inputs — so
+the whole chain (derive → log → z-score — unzscore → exp → un-derive)
+round-trips. Note that a log makes the target path non-affine: `pipe.affine()`
+will refuse such pipelines, and real-terms contributions must go through the
+batched `inverse_transform` instead of the `(β, scale)` shortcut.
+
+The same step derives within one variable too, when the variable carries an
+extra coordinate dimension — `controls (date, product, control)` with
+`control = [level, reference, seasonality]`:
+
+```python
+# continues above — `mixed` carries over
+from xscalers.preprocessing import Derive
+
+share = Derive(
+    name="level_diff",                           # op="diff" this time:
+    left=("controls", {"control": "level"}), # level minus reference,
+    right=("controls", {"control": "reference"}),
+    op="diff",                                   # not a ratio
+)
+derived = share.transform(mixed, share.fit(mixed))
+derived["level_diff"].dims                     # (date, product)
+```
+
+## Dropping variables and coordinate values
+
+`VariableDropper` is the drop verb — the xarray answer to sklego's `ColumnDropper`.
+`vars=` removes whole variables; `labels=` removes values from a coordinate
+dataset-wide (the complement of keeping, which is `Selector`'s job):
+
+```python
+# continues above — `mixed` carries over
+from xscalers.preprocessing import VariableDropper
+
+prune = VariableDropper(vars=["media_national", "controls_prod"], labels={"control": ["reference"]})
+pruned = prune.transform(mixed, prune.fit(mixed))
+"media_national" in pruned.data_vars          # False — raw inputs are gone
+"reference" in pruned.coords["control"].values  # False — every variable shrank with it
+```
+
+The motivating shape is pruning raw inputs once a derived metric supersedes
+them, so later steps never have to enumerate what they don't care about:
+
+```python
+# continues above — `mixed` carries over
+from xscalers.compose import VariableTransformer
+from xscalers.pipeline import Pipeline
+from xscalers.preprocessing import Derive, VariableDropper, ZScore
+
+pipe = Pipeline(
+    [
+        (
+            "derive",
+            Derive(
+                name="tv_per_control",
+                left=("media_national", {"media": "tv"}),
+                right=("controls_prod", {}),
+                op="ratio",
+            ),
+        ),
+        ("prune", VariableDropper(vars=["media_national", "controls_prod"])),
+        ("scale", VariableTransformer({"tv_per_control": ZScore(dims=("date",))})),
+    ]
+)
+model_ready = pipe.fit_transform(mixed)
+sorted(model_ready.data_vars)                 # raw inputs pruned, ratio kept
+```
+
+Dropping destroys data, so — like `Selector` — `VariableDropper.inverse_transform`
+raises instead of guessing, and `pipe.inverse_transform` refuses for the whole
+chain. Real-terms conversion is unaffected: `pipe.unscale(contrib,
+var="tv_per_control")` still walks that one variable backwards. `fit` validates
+rather than failing late (sklego's contract): unknown names raise `KeyError`,
+dropping *every* variable raises `ValueError`, and a non-`Dataset` input raises
+`TypeError` (DataTree node drops are deferred).
+
+If you need the reverse, a pipeline that scales what it keeps, `Splitter` is
+the invertible way to split one variable into differently-treated pieces, and
+`VariableTransformer(..., remainder="drop")` is the leftover-dropping sibling.
+
+
+## Splitting one variable into differently-treated pieces
+
+Different treatment *within* one media variable (digital vs TV) works by
+splitting into pieces, scaling each, then recombining. Real inputs are a
+plain Dataset of variables, so pieces become *sibling variables*
+(`media_digital`, `media_linear`), each keeping its slice of the coordinate
+under a piece-specific dim name; `inverse_transform` recombines and restores
+the original coordinate order:
+
+```python
+from xscalers.compose import VariableTransformer
+from xscalers.datasets import load_media_channels
+from xscalers.pipeline import Pipeline
+from xscalers.preprocessing import FixedScaler, MeanScaler, Splitter
+
+ds = load_media_channels()
+
+pipe = Pipeline(
+    [
+        (
+            "split",
+            Splitter(
+                "media",
+                {
+                    "digital": {"channel": ["search", "social", "fb"]},
+                    "linear": {"channel": ["tv", "stream"]},
+                },
+            ),
+        ),
+        (
+            "scale",
+            VariableTransformer(
+                {
+                    "media_digital": MeanScaler(dims=("date",)),
+                    "media_linear": FixedScaler(scale=1_000),
+                },
+            ),
+        ),
+    ]
+)
+state = pipe.fit(ds)
+scaled = pipe.transform(ds, state)
+
+back = pipe.inverse_transform(scaled, state)   # pieces recombine; coords restored
+import numpy as np
+np.testing.assert_allclose(back["media"], ds["media"])
+```
+
+## Saving the fitted state
+
+State is **data** — saving it is native xarray IO, in any of the three shapes
+you'll meet in practice:
+
+**1. A standalone one-file artifact** — `to_netcdf`, next to your recipe:
+
+```python
+import os
+import xarray as xr
+import numpy as np
+from xscalers.datasets import load_national
+from xscalers.preprocessing import MeanScaler
+
+national: xr.Dataset = load_national()
+scaler = MeanScaler(dims=("date",))
+state = scaler.fit(national["media"])
+
+state.to_netcdf("media_state.nc")                 # state is a Dataset — native IO
+reloaded = xr.open_dataset("media_state.nc")      # ...read back natively
+scaled = scaler.transform(national["media"], reloaded)
+back = scaler.inverse_transform(scaled, reloaded)
+np.testing.assert_allclose(back, national["media"])
+os.remove("media_state.nc")                       # tidy up after this example
+```
+
+**2. Alongside your (large) zarr inputs** — append the state tree as groups
+of the existing store; `append_state` exists only because
+`DataTree.to_zarr(group=...)` is not implemented upstream — reading back is
+native:
+
+```python
+# continues above — `xr`, `mixed` and the fitted `transformer` carry over
+from pathlib import Path
+from xscalers.compose import VariableTransformer
+from xscalers.container import state_tree
+from xscalers.datasets import load_mixed_grain
+from xscalers.persistence import append_state
+from xscalers.preprocessing import MeanScaler
+
+store: Path = Path("inputs.zarr")             # your existing input store
+mixed: xr.Dataset = load_mixed_grain()       # pretend it already lives there
+mixed.to_zarr(store)                         # ... here we create it first
+
+transformer: VariableTransformer = VariableTransformer({"media_national": MeanScaler(dims=("date",))}, remainder="drop")
+state = transformer.fit(mixed)
+append_state(state_tree({"scale": state}), store, group="scalers")  # state under "scale"
+
+import xarray as xr
+reloaded_tree = xr.open_datatree(store, group="scalers")
+reloaded_state = {
+    name: node.dataset for name, node in reloaded_tree.children["scale"].children.items()
+}
+scaled_again = transformer.transform(mixed, reloaded_state)
+import shutil; shutil.rmtree(store)          # tidy up after this example
+```
+
+**3. Next to a pymc model** — the state and the `InferenceData` are separate
+artifacts sharing a directory; recipe (JSON via `to_spec`) is versioned
+config, state is data, never the two mixed:
+
+```python
+# requires-pymc
+import pymc as pm
+import arviz as az
+from xscalers.datasets import load_national
+from xscalers.preprocessing import MeanScaler
+
+national = load_national()
+scaler = MeanScaler(dims=("date",))
+state = scaler.fit(national["target"])
+y = scaler.transform(national["target"], state)   # scaled input into the model
+
+import xarray as xr
+with pm.Model() as model:
+    beta = pm.Normal("beta", 0, 1)
+    sigma = pm.HalfNormal("sigma", 1)
+    # numpy views into the scaled dataset operate the same as scaled arrays
+    pm.Normal("obs", beta * float(y.mean()), sigma, observed=np.asarray(y))
+    idata = pm.sample_prior_predictive(draws=10, random_seed=0)
+
+import numpy as np
+idata.to_netcdf("posterior.nc")                   # model results...
+state.to_netcdf("scaler_state.nc")                # ...and its scaling, side by side
+idata_back = xr.open_datatree("posterior.nc")     # DataTree in, DataTree out
+state_back = xr.open_dataset("scaler_state.nc")
+
+# real-terms conversion: the affine pair (center, scale) — MeanScaler is
+# divisive, so its composed view is (center=0, scale=mean):
+posterior_beta: float = float(idata_back["prior"]["beta"].mean())
+contribution_scaled: float = posterior_beta * float(np.asarray(y).mean())
+contributions_original_scale = contribution_scaled * float(state_back["mean"])
+import os
+for name in ("posterior.nc", "scaler_state.nc"):
+    os.remove(name)
+```
+
+The fitted pipeline serializes as *recipe + artifact*:
+
+```python
+# continues above — `state` carries over
+import os
+from xscalers.container import state_tree
+
+state_tree({"scale": state}).to_netcdf("fitted_state.nc")   # whole pipeline: one artifact
+```
+
+## Converting back to real units — one verb
+
+Model code never hand-writes `× scale + center`: the affine pair is an
+*internal* detail. `Pipeline.unscale` applies the composed pair per variable
+(eager xarray contributions or symbolic `pytensor.xtensor` in the graph);
+when a pipeline is non-affine (e.g. a log-derived metric), `unscale` walks
+only *that variable's own* steps backwards, stopping at the step that created
+it — so the one call stays honest where hand-written `β·scale` algebra goes
+silently wrong:
+
+```python
+import numpy as np
+import xarray as xr
+from xscalers.compose import VariableTransformer
+from xscalers.datasets import load_mixed_grain
+from xscalers.pipeline import Pipeline
+from xscalers.preprocessing import Derive, FixedScaler, FunctionTransformer, ZScore
+import os
+
+pipe = Pipeline(
+    [
+        (
+            "derive",
+            Derive(
+                name="tv_per_control",
+                left=("media_national", {"media": "tv"}),
+                right=("controls_prod", {}),
+                op="ratio",
+            ),
+        ),
+        ("log", FunctionTransformer({"tv_per_control": "log"})),
+        ("scale", VariableTransformer({"tv_per_control": FixedScaler(scale=1.0)})),
+    ]
+)
+state = pipe.fit(load_mixed_grain())
+contrib = xr.DataArray(np.full((104, 4), -0.75), dims=("date", "product"))
+
+y_real = pipe.unscale(contrib, var="tv_per_control")   # one verb; located = exp(−0.75)
+np.testing.assert_allclose(y_real, np.exp(contrib))
+```
+
+## The flagship flow: fit on xarray, unscale a pymc graph
+
+Requires the optional `pytensor` extra (`pip install "xscalers[pytensor]"`).
+
+```python
+import pytensor.xtensor as ptx
+from xscalers.compose import VariableTransformer
+from xscalers.datasets import load_national
+from xscalers.pipeline import Pipeline
+from xscalers.preprocessing import MeanScaler
+
+scaler = MeanScaler(dims=("date",))
+data = load_national()
+pipe = Pipeline([("scale", VariableTransformer({"media": scaler}))])
+state = pipe.fit(data)                                     # fitted pipeline state
+
+contrib = ptx.xtensor("channel_contrib", dims=("date", "channel"))
+spend_real = pipe.unscale(contrib, var="media", state=state)   # symbolic: contrib * mean
+```
+
+After sampling, contributions arrive as xarray with `(chain, draw, ...)` —
+the same state broadcasts over those extra dims by name.
+
+## Writing it out: post-fit step access
+
+```python
+from xscalers.datasets import load_national
+from xscalers.pipeline import make_pipeline
+from xscalers.preprocessing import MeanScaler
+
+pipe = make_pipeline(MeanScaler(dims=("date",)), MeanScaler(dims=("date",)))
+assert [name for name, _ in pipe.steps] == ["meanscaler", "meanscaler-1"]
+pipe.fit(load_national())                      # sklearn-style fitted pipeline
+
+bound = pipe["meanscaler"]                     # transformer + its fitted state slice
+scaled_again = bound.transform(load_national())  # one fitted step on fresh data
+back = bound.inverse_transform(scaled_again)     # ...and its inverse
+```
+
+## For sklearn users
+
+The vocabulary is deliberately sklearn-shaped:
+
+| sklearn | xscalers | notes |
+| --- | --- | --- |
+| `fit` / `transform` / `inverse_transform` | same | states are `xr.Dataset`s, not attributes |
+| `ColumnTransformer` | `VariableTransformer` | per-variable treatment over named data variables |
+| `ColumnDropper` (sklego) | `VariableDropper` | over `drop_vars` / `drop_sel`; `Selector` covers the keeping side |
+| `Pipeline` | `Pipeline` | named steps, chained states, reversed inverse |
+| `make_pipeline` | `make_pipeline` | the composition verb (no `>>` chaining: it hides the recipe/state split) |
+| `FeatureUnion` | `make_union` | many pipelines over one dataset, collision policy, covariates-only by design |
+| `"passthrough"` | `"passthrough"` | identical semantics |
+
+The deliberate departure: state is **explicit and passed around**, not hidden
+as `scaler.mean_`, which is what lets you edit fitted values and reuse them
+across backends. The counterpart to `transform` for model code is
+**`unscale`**: one verb, correct for affine pipelines and honest for
+non-affine ones.
+
+## Install
+
+```text
+pip install xscalers            # core (xarray, numpy)
+```
+
+## License
+
+MIT