refac: Harmonize linopy operations and introduce a new predictable and strict convention

[FBumann]

Harmonize linopy arithmetic with legacy/v1 convention transition

This PR harmonizes coordinate alignment and NaN handling in linopy arithmetic, with a non-breaking transition layer.

linopy.options["arithmetic_convention"]

Two modes:

• "legacy" (default) — reproduces current master behavior exactly. Emits LinopyDeprecationWarning on every legacy codepath.
• "v1" — strict exact-join semantics: mismatched coordinates raise ValueError with helpful messages. NaN in user data raises immediately.

linopy.options["arithmetic_convention"] = "v1"

v1 behavior

Exact coordinate matching — Arithmetic operators (+, -, *, /) require matching coordinates on shared dimensions. Mismatched coordinates raise ValueError with suggestions:

x[i=0,1,2] + y[i=1,2,3]  # ValueError: use .add(y, join="inner")

Named methods with explicit join — .add(), .sub(), .mul(), .div(), .le(), .ge(), .eq() accept join= parameter:

x.add(y, join="inner")      # intersection
x.add(y, join="outer")      # union with fill
x.add(y, join="left")       # keep x's coordinates
x.add(y, join="override")   # positional alignment

Free broadcasting — Constants can introduce new dimensions. All algebraic laws hold.

---

v1 NaN convention

NaN means "absent term" — never a numeric value.

How NaN enters

Only two sources:

1. mask= argument at construction (add_variables, add_constraints)
2. Structural operations: .shift(), .where(), .reindex(), .reindex_like(), .unstack() (with missing combinations)

Operations that do not produce NaN: .roll() (circular), .sel() / .isel() (subset), .drop_sel() (drops), .expand_dims() / .broadcast_like() (broadcast existing data).

How NaN propagates

NaN marks an individual term as absent, not the entire coordinate. When expressions are combined (e.g., x*2 + y.shift(time=1)), each term is independent — an absent term from y.shift does not mask the valid x term at the same coordinate.

A coordinate is only fully absent when all terms have vars=-1 and const is NaN. This is what isnull() checks.

What raises

Any user-supplied NaN at an API boundary — in constants, factors, or constraint RHS — raises ValueError immediately:

x + data_with_nans        # ValueError
x * factor_with_nans      # ValueError
expr >= rhs_with_nans     # ValueError

Users handle NaN explicitly:

# Fill before arithmetic — you choose the fill value
x + data.fillna(0)        # NaN = "no offset"
x * factor.fillna(0)      # NaN = "exclude this term"
x * factor.fillna(1)      # NaN = "no scaling"

# Mask constraints — .sel() (preferred) or mask=
valid = rhs.notnull()
m.add_constraints(expr.sel(i=valid) <= rhs.sel(i=valid), name="c")
# or
m.add_constraints(expr <= rhs.fillna(0), mask=rhs.notnull(), name="c")

Why this is consistent

• lhs >= rhs ≡ lhs - rhs >= 0 — RHS follows the same rules as any constant
• No dual role for NaN: internal NaN (from shift, mask=) is structural; user NaN is always an error
• Absent terms, not absent coordinates: combining a valid expression with a partially-absent one preserves the valid part

Implementation details

• FILL_VALUE = {"vars": -1, "coeffs": NaN, "const": NaN} — all float fields use NaN for absence, integer fields use -1
• NaN checks in _add_constant, _apply_constant_op, to_constraint
• Piecewise internals use .fillna(0) on breakpoint data

---

Legacy behavior (default, backward-compatible)

• Size-aware alignment: override when sizes match, left-join otherwise
• NaN as neutral element: filled with 0 (add/sub/mul) or 1 (div)
• Constraint RHS: NaN means "no constraint", preserved through subtraction

merge() behavior

merge() enforces exact matching on shared user-dimension coordinates in v1 mode. Helper dims (_term, _factor) and the concat dim are excluded from this check. The actual xr.concat uses join="outer".

In legacy mode, merge uses override when shared user dims have matching sizes, outer otherwise.

Source changes

File	Change
config.py	LinopyDeprecationWarning, arithmetic_convention setting
expressions.py	All arithmetic paths branch on convention; merge() pre-validates user-dim coords under v1; to_constraint has separate legacy/v1 paths; NaN validation at API boundaries
common.py	align() reads convention (legacy→inner, v1→exact)
variables.py	Scalar fast path in __mul__, explicit TypeError in __div__, .reindex() methods
piecewise.py	.fillna(0) on breakpoint data for v1 compatibility
monkey_patch_xarray.py	DataArray/Dataset arithmetic with linopy types
model.py	Convention-aware model methods

Documentation

Notebook	Content
arithmetic-convention.ipynb	Coordinate alignment rules, join parameter, migration guide
missing-data.ipynb	NaN convention principles, fillna patterns, masking with .sel() and mask=, legacy comparison
_nan-edge-cases.ipynb	Dev notebook: investigation of shift, roll, where, reindex, isnull, arithmetic on shifted expressions, FILL_VALUE internals

Test structure

• Marker-based separation: @pytest.mark.v1_only and @pytest.mark.legacy_only for convention-specific tests
• Shared fixtures in conftest.py with auto-convention switching
• Tests validate: strict matching, explicit join, NaN raises, NaN propagation in shifted expressions, piecewise/SOS under both conventions

Rollout plan

1. This PR: Default "legacy" — nothing breaks
2. Downstream: Users opt in with linopy.options["arithmetic_convention"] = "v1"
3. linopy v1: Flip default to "v1" and drop legacy mode

Open questions

• from_tuples / linexpr() — Currently follows the global convention. In practice always called with same-coord variables, so convention doesn't matter. Low-priority.
• nan_as_mask option — A future global setting could allow NaN in user data to be treated as masking (equivalent to .fillna(0) + mask=data.notnull()). Deferred for team discussion.
• Pipe operator — Only linopy objects, or also constants? (follow-up PR)

Test plan

• All tests pass under both conventions (3450 passed)
• Legacy tests validate backward compatibility
• v1 tests validate strict coordinate matching and NaN raises
• Piecewise and SOS tests run under both conventions
• Documentation notebooks execute cleanly

🤖 Generated with Claude Code

[FBumann]

@FabianHofmann Im quite happy with the notebook now. It showcases the convention and its consequences.
Tests need some work though. And migration as well.
Looking forward to your opinion on the convention

[FBumann]

The convention should be "exact" for all of +, -, *, /, with an additional check that neither side may introduce dimensions the other doesn't have — also for all operations.

Why "exact" instead of "inner" for * and /

"exact" still broadcasts freely over dimensions that only exist on one side — it only enforces strict matching on shared dimensions. So the common scaling pattern works fine:

cost = xr.DataArray([10, 20], coords=[("tech", ["wind", "solar"])])
capacity  # dims: (tech=["wind", "solar"], region=["A", "B"])

cost * capacity  # ✓ tech matches exactly, region broadcasts freely

"inner" is dangerous: if coords on a shared dimension don't match due to a typo or upstream change, it silently drops values. The explicit and safe way to subset before multiplying is:

capacity.sel(tech=["wind", "solar"]) * renewable_cost

No operation should introduce new dimensions

Neither side of any arithmetic operation should be allowed to introduce dimensions the other doesn't have. The same problem applies to + and - as to * and / — new dimensions silently expand the optimization problem in unintended ways:

cost_expr      # dims: (tech, time)
regional_expr  # dims: (tech, time, region)

cost_expr + regional_expr  # ✗ silently expands to (tech, time, region)

capacity  # dims: (tech, region, time)
risk      # dims: (tech, scenario)
risk * capacity  # ✗ silently expands to (tech, region, time, scenario)

An explicit pre-check on all operations:

asymmetric_dims = set(other.dims).symmetric_difference(set(self.dims))
if asymmetric_dims:
    raise ValueError(f"Operation introduces new dimensions: {asymmetric_dims}")

Summary

Operation	Convention
+, -, *, /	"exact" on shared dims; neither side may introduce dims the other doesn't have

[coroa]

The convention should be "exact" for all of +, -, *, /, with an additional check that neither side may introduce dimensions the other doesn't have — also for all operations.

Let's clearly differentiate between dimensions and labels.

labels

I agree with "exact" for labels by default, but we need an easy way to have inner or outer joining characteristics. I found the pyoframe conventions
strange at the beginning, but they grew on me:

x + y.keep_extras() to say that an outer join is in order and mismatches should fill with 0.

x + y.drop_extras() to say that you want an outer inner join.
x.drop_extras() + y does the same, though.

I have in a different project used | 0 to indicate keep_extras ie (x + y | 0).

dimensions

i am actually fond of the ability to auto broadcast over different dimensions. and would want to keep that (actually my main problem with pyoframe).

your first example actually implicitly assumes broadcasting.

[FBumann]

Dimensions and broadcasting

I agree that auto broadcasting is helpful in some cases.
I'm happy with allowing broadcasting of constants. We could allow this always...?
But I would enforce that the constant never has more dims than the variable/expression.
Or is there a use case for this?

So the full convention requires two separate things:
1. "exact" join — shared dims must have matching coords (xarray handles this)
2. Subset dim check — the constant side’s dims must be a subset of the variable/expression (custom pre-check needed)

labels

I'm not sure if I like this approach, as it's needs careful state management of the flags on expressions. The flag (keep or drop extras) needs to be handled.
I would rather enforce to reindex or fill data to the correct index.
I think aligning is the correct approach:

import linopy

# outer join — fill gaps with 0 before adding
x_aligned, y_aligned = linopy.align(x, y, join="outer", fill_value=0)
x_aligned + y_aligned

# inner join — drop non-matching coords before adding
x_aligned, y_aligned = linopy.align(x, y, join="inner")
x_aligned + y_aligned

Combining disjoint expressions would then still need the explicit methods though.
I'm interested about your take on this

[FBumann]

The proposed convention for all arithmetic operations in linopy:
1. "exact" join by default — shared coords must match exactly, raises on mismatch
2. Subset dim check — constants may introduce dimensions the variable/expression doesn’t have
3. No implicit inner join — use .sel() explicitly instead
4. Outer join with fill — use x + (y | 0) or .add(join="outer", fill_value=0)
The escape hatches in order of preference: .sel() for subsetting, | 0 for inline fill, named method .add(join=...) for everything else. No context manager needed.​​​​​​​​​​​​​​​​

I'm not sure how to implement the | operator yet. Might need some sort of flag/state for defered indexing

[FBumann]

I thought about the pipe operator:
I think it should only work with linopy internal types (Variables/expression), not constants (scalar, numpy, pandas, dataarray), as this would need monkey patching a lot and hard to get stable.

Would this be an issue for you?

[brynpickering]

This is definitely a worthwhile update. How we've handled it in calliope is to allow for dims to have mismatched labels but to require a default value for each parameter, expr, variable which will be used when there is a mismatch. This aligns with keep_extras in pyoframe except that the default value can change (e.g., for technology efficiency, you would rather have a default value of 1 than of 0). It is the same as the fill_value proposed on aligning mentioned here, but extended to include vars and exprs.

Raising exceptions on misalignment can be a pain to deal with. We often have patchwork data that we want to combine with a clear fill value in mind that is equivalent to "no effect". I can see alignment being used a reasonable amount to deal with this. The issue is that you end up with lots of intermediate arrays in memory for you to achieve alignment. Having a pre-defined default or the ability to define a default on applying arithmetic operations would keep in-memory arrays to a minimum. This would be easy to then integrate with #561.

On fill values: this PR only makes them available on parameters ("data") but I'd say they're also useful for decision variables and expressions, too. E.g., when creating a total cost expression you might want to combine several instances of var1 * cost_param1 + var2 * cost_param2 + ... where dimension labels don't align between var1 and param1 nor between var1 * cost_param1 and var2 * cost_param2 but that you just always want the default to be zero whenever there is misalignment since you'll later be summing all this to define your objective function.

[FBumann]

@brynpickering As i understand your comment, you are mostly ok with the convention, but have thoughts on nan data. Am I right? THis is what im currently not that opinionated about, so im happy to hear your suggestions.

[FBumann]

@brynpickering Thanks for the detailed feedback.

Fill values depend on context (efficiency=1 vs cost=0) — this is id like to make it a user choice via .fillna(value) rather than picking a default.

Misalignment doesn't have to raise. The escape hatches are lightweight:

total = cost_a.add(cost_b, join="outer")  # absent terms → zero at solve time
scaled = var * efficiency.fillna(1)        # you choose the fill

As far as I know, Memory overhead is the same either way — fillna() on a parameter is not much larger than the parameter itself, and join="outer" does alignment inside xarray's concat anyway. But maybe im wrong here.

For expressions with mismatched labels, join="outer" already works: absent terms carry NaN coefficients that are filtered out when flattening for the solver. No explicit fill value needed.

Maybe adding a fill_value= parameter on .mul() / .add() resolves your issue if the fillna + join= pattern proves too verbose in practice?

[FBumann]

@coroa @FabianHofmann Im quite happy with this. But i think it desperately needs discussion.

[brynpickering]

@FBumann your fillna suggestion is fine for parameters, but doesn't address nan filling later down the line of variables or expressions. As far as I understand linopy objects, it isn't possible to nan fill in the same way. This risks NaNs making their way into arithmetic when one would rather have some numeric value that has no effect to be there.