Extend xarray method coverage on Variable / LinearExpression

[FBumann]

Summary

Variable, LinearExpression/QuadraticExpression and Constraint each expose a hand-picked subset of xarray.Dataset methods (via varwrap / exprwrap / conwrap). The three lists have drifted apart, several common arithmetic methods are missing entirely, and a few existing ones diverge from xarray — especially around datetime indexes.

This issue collects the analysis and a design sketch for closing the gaps. All of the proposed methods preserve linearity (none multiplies two variables), so they are mostly bookkeeping.

---

A. Cross-class inconsistencies (existing methods)

The same wrapper is exposed on some classes but not others, with no clear reason:

method	Variable	Expression	Constraint
compute	yes	no	no
chunk (dask)	no	yes	yes
reindex / reindex_like	no	yes	yes
astype	no	yes	no
drop / drop_vars	no	yes	no
reset_index	no	yes	no
rename_dims	no	yes	yes
shift fill_value default	yes (variables.py:1268)	no (expressions.py:1474)	no

The two that actually bite users:

• Variable.reindex / reindex_like missing — you can reindex an expression or constraint onto a master index but not a variable. Natural operation for datetime work (aligning a variable to a full snapshot index).
• Expression.shift / Constraint.shift have no fill_value while Variable.shift defaults to the linopy fill value. Shifted slots get coeffs=NaN, vars=-1; it mostly survives downstream fillna(0), but it is the root of the diff issue below.

B. Existing methods lacking features vs. xarray (datetime focus)

1. diff (expressions.py:1238) — implemented as self - self.shift({dim: n}). Two divergences from xarray.diff:
   • Length: xarray trims the dimension to N-n; linopy keeps N, so the first n rows are spurious (self[0] - empty, a leading garbage term still carrying the original coordinate label).
   • No label argument — xarray has label="upper"|"lower"; linopy is hard-wired to "upper".
2. groupby (expressions.py:1258) — typed only for DataArray/Series/DataFrame. The string-accessor form (groupby("time.month"), "time.season") only reaches the slow fallback; the fast reindex summation (expressions.py:229) requires Series/DataFrame/DataArray. groupby_bins not exposed. LinearExpressionGroupby only implements .sum() / .map() / .roll() — no .mean() / .first() / .last() / .count().
3. rolling (expressions.py:1289) — LinearExpressionRolling only implements .sum() (expressions.py:307). No .mean(), no .construct(). Integer-window only.
4. No frequency-aware shift — neither linopy nor xarray support shift(time="1D"); pandas does. See the datetime-shift section below.

sel/isel delegate cleanly, so datetime partial-string indexing (sel(time="2030"), method="nearest") already works.

C. Missing methods, ranked by usefulness

Only linearity-preserving methods are candidates. Excluded as non-linear/meaningless: max, min, median, std, var, prod, cumprod, quantile, rank, idxmax/idxmin, argmax/argmin, clip.

Tier 1 — high value

• mean — linear (sum/n), ubiquitous, missing on both Variable and LinearExpression.
• resample — datetime aggregation (hourly to daily). Core energy-modeling need. Absent.
• reindex / reindex_like on Variable — fixes the asymmetry in A.
• weighted — expr.weighted(w).sum(); maps onto PyPSA snapshot_weightings.

Tier 2 — useful

• coarsen — positional block aggregation (time=24).
• rolling().mean() and groupby().mean() — extend the existing reducers.
• transpose — missing on both; needs care to keep _term/_factor last.
• compute on Expression, chunk on Variable — close the dask asymmetry from A.
• dropna — drop missing coordinate slices.

Tier 3 — nice-to-have / niche

• astype / reset_index / rename_dims on Variable (consistency).
• sortby, squeeze, head/tail/thin, pad.
• interp / interp_like (genuinely linear but niche).
• groupby_bins.

---

Implementation design notes (Tier 1 + Tier 2)

Why none of this is hard

A linopy linear expression is just a list of terms plus a constant — 3*x[a] + 2*x[b] + 5 — stored as three aligned arrays: coeffs, vars (integer labels), const.

Every method below does one of three harmless things to that list:

1. Regroup terms — collect terms from several cells into one. (sum, resample, coarsen, groupby)
2. Copy terms — the same term appears in several output cells. (rolling)
3. Move or drop whole cells — pick, fill, or discard cells without touching the terms inside. (reindex, shift, dropna, transpose)

and some also rescale — multiply every coefficient and the constant by a number/array. (mean, weighted, the .mean() variants)

None of this multiplies two variables together, so the result is always still a valid linear (or quadratic) expression. linopy already owns every building block: regrouping is the term-stacking trick behind .sum(); copying is the window trick behind .rolling().sum(); moving cells is ordinary .sel()/.reindex(); rescaling is ordinary *//.

Tier 1

• mean = sum divided by how many things were summed: gen.mean("time") == gen.sum("time") / 3. Divide by the count of non-missing entries (matches xarray skipna=True and linopy's own sum, which already skips -1 variables). All-missing slice -> 0/0 = NaN, as xarray.
• resample aggregates a datetime axis by calendar period. Each period's expression = sum of that period's terms = a groupby keyed by "which period". Ask pandas for the period label of each timestamp, reuse the existing groupby(...).sum() fast path; .mean() divides by the count. Decide: keep empty periods as 0 (xarray parity) or drop them. Forward closed/label/origin.
• reindex / reindex_like (Variable) — new slots get the missing-sentinel (labels=-1, bounds NaN), same fill Expression/Constraint already use. One line each: reindex = varwrap(Dataset.reindex, fill_value=FILL_VALUE).
• weighted — gen.weighted(w).sum("time") == (gen * w).sum("time"). Pure sugar over * and sum; .mean() also divides by w.sum(). Just a small wrapper object.

Tier 2

• coarsen — resample's positional cousin (group every N rows). Reshape into blocks via xarray's coarsen().construct(), then sum each block; .mean() divides by block size. Handle boundary="trim"|"pad".
• rolling().mean() / groupby().mean() — rolling/groupby sum already exist; .mean() divides by the window size / group size (or the valid count at window edges).
• transpose — pure axis reorder, no arithmetic. Must reorder only real dims and keep the internal _term/_factor last. Plain wrap for Variable.
• compute / chunk — no math, dask plumbing; add only to make the three classes consistent.
• dropna — drops "missing" coordinate slices. "Missing" in linopy is not raw NaN: for Variable it is labels == -1; for an expression it is isnull() (all terms empty AND const NaN). Build the drop-mask from isnull().

Datetime-aware shift

Today shift is integer/positional only — shift(time=1) moves everything one row. Fine for a regular hourly index, used for storage balance soc[t] == soc[t-1] + charge[t]. But on irregular snapshots (variable time resolution, clustered representative periods) "one row" is not "one hour". There are three distinct operators:

1. Integer shift — what exists today.
2. Time shift — move by an actual duration, shift(time="1h"): for each timestamp find the row exactly one hour earlier; no such row -> missing. The correct operator for irregular grids and what storage/ramping constraints want. Basically a reindex onto "time minus offset" labels, so nearly free once Variable has reindex. Collapses to integer shift on a regular index.
3. Index shift — keep the data, relabel the time axis by an offset (pandas shift(freq=...)).

pandas handles month-end/DST for offsets like "1ME". Adding datetime shift is also a good moment to fix the missing fill_value on Expression.shift / Constraint.shift (A above).

---

Open questions

• mean: divide by the non-missing count, or by the raw length?
• resample / coarsen: keep empty periods, or drop them?
• rolling().mean() at the edges: divide by the window size, or by the valid count?
• shift(time="1h"): time shift or index shift (or expose both)?
• Add the new methods to QuadraticExpression and Constraint too, or just Variable / LinearExpression?
• Should this ship as one PR per tier, or per method?

[FBumann]

Implementation checklist — xarray method coverage

Tracking checklist for PyPSA/linopy#703.

Decisions (resolved)

• mean divisor — raw length (sum / sizes[dim]). Contingent on refac: Harmonize linopy operations and introduce a new predictable and strict convention #591 keeping the explicit-zero-fill convention (alignment gaps become real 0 expressions, no lingering nulls); revisit if refac: Harmonize linopy operations and introduce a new predictable and strict convention #591 reverses course. Note: masked-variable slots (vars == -1) are not flagged by isnull() — that structural-absence inconsistency is orthogonal and tracked separately.
• resample / coarsen — keep empty periods. Regular gap-free axis; an empty bin is an empty expression. resample().mean() of an empty bin (0/0) is defined as an empty expression. coarsen divisibility handled via xarray's boundary= arg.
• rolling().mean() at edges — valid count. Divide each window by its non-padded entries (window padding is genuine nan). Keeps mean * count == sum, consistent with the current rolling().sum() partial-sum behaviour. Not affected by refac: Harmonize linopy operations and introduce a new predictable and strict convention #591.
• shift — one shift, internal dispatch. shift(time=1) and shift(time="1h") both work; shift dispatches on argument type to a dedicated, separately-tested _shift_freq helper. xarray has no native freq shift — parse the freq with pd.tseries.frequencies.to_offset (pandas is already a dep), shift the coordinate, reindex back with FILL_VALUE.

Scope (resolved)

Aggregation methods (mean, weighted, resample, coarsen,
rolling().mean(), groupby().mean()) land on BaseExpression — so
LinearExpression and QuadraticExpression stay synced by inheritance — and on
Variable as to_linexpr() delegations. They do not go on Constraint: a
constraint is a finished relation, not something you aggregate, and it already
has no sum / groupby / rolling.

method	Variable	LinearExpr	QuadExpr	Constraint
mean, weighted, resample, coarsen, *.mean()	yes	yes	yes	no
transpose	yes	yes	yes	yes
dropna	yes	yes	yes	optional (low value)
reindex / reindex_like	add	has	has	has
datetime shift	yes	yes	yes	defer (has integer shift)

Implement each method once on BaseExpression; QuadraticExpression inherits
it. Per-method check: _factor-safety — holds automatically since all touch
only coordinate dims; transpose must keep _factor / _term last.

Phase 1 — consistency fixes (cheap one-liners)

• Variable.reindex / Variable.reindex_like (with FILL_VALUE)
• LinearExpression.shift / Constraint.shift — add fill_value default

Phase 2 — Tier 1 methods

• mean
• weighted wrapper with .sum() / .mean()
• resample with .sum() / .mean() (reuse groupby fast path)

Phase 3 — Tier 2 methods

• transpose (keep _term / _factor last)
• dropna (mask from isnull())
• coarsen with .sum() / .mean()
• LinearExpressionRolling.mean()
• LinearExpressionGroupby.mean()

Phase 4 — datetime-aware shift

• frequency-aware (time) shift

Phase 5 — existing-method fixes (Part B)

• groupby — accept string accessor ("time.month"), route through fast path
• diff — xarray-parity: trim length + label arg (breaking; needs decision)
• expose groupby_bins (optional)

Cross-cutting

• docstrings borrowed/adapted from xarray for every new method
• documentation / example notebook coverage
• release_notes.rst entry

[FBumann]

Those methods should probably be buildt on top of #717, only be unavaillable for "legacy". THis will implementing them much easier.

[FBumann]

Cross-referencing from #717 (v1 arithmetic convention) — §13 of arithmetics-design/convention.md already settles the first of your "Open questions": mean divides by the count of present slots, not raw length, matching xarray skipna=True and v1's "absent ≠ zero" rule (§1).

The downstream edge case — mean over an all-absent slice (0/0) — isn't covered by the §13 text. For the eventual implementation:

Propagate NaN at the slot, do not absorb to 0.

Rationale:

• sum(none) = 0 works because 0 is the additive identity — the unique "right" answer when nothing is summed.
• mean(none) has no equivalent: an arithmetic mean over zero values is genuinely undefined. Returning 0 would silently fold "absent" back into "present zero" — exactly the §1 / An absent variable is indistinguishable from a zero variable #712 violation v1 is built to close.
• xarray's default already returns NaN here, so matching it is free.

So mean should be: an all-absent slice produces one NaN entry, and the caller resolves it explicitly with .fillna(...) if they want a different behaviour.

Same logic likely applies to rolling().mean() at the edges and coarsen().mean() on empty blocks (NaN, not 0), but those are separate design calls in this issue's list and worth confirming individually with the implementation PR.

[FBumann]

Note

Generated with Claude Code.

Design note carried over from the v1 convention (§13, #717), so it isn't lost when these reductions land here.

When the absent-aware reductions are added, they follow §13's skip-absent rule — collapse the dimension by skipping absent slots, rather than NaN-propagating (§6) or treating absent as zero (§1):

• mean divides by the count of present slots, not all of them — dividing by all would treat an absent slot as a zero term, which §1 forbids.
• resample / coarsen aggregations: the same skip-absent principle, per bucket.

Add these under v1 only. They are new operations with no legacy behaviour, so legacy never gains them — it is a frozen compatibility layer removed at 1.0, and there is no reason to give it new operations or to design fill-0 semantics for them. So: no legacy/v1 fork, no deprecation, nothing to change for them at the 1.0 removal. They're fully decoupled from the v1 rollout — v1 doesn't gate them and they don't gate v1.

§13 in #717 now keeps only a forward-pointer to this issue.