Uh oh!
There was an error while loading. Please reload this page.
Feat/add piecewise linear - #558
Closed
FBumann wants to merge 9 commits into
Closed
Conversation
- Add return type annotations to add_sos_constraints and add_sos - Fix iterate_slices call with list instead of tuple - Fix groupby for multi-dimensional SOS using stack/unstack - Add defensive check in remove_sos_constraints - Clarify direct API support (Gurobi only) in docs
Files Modified 1. linopy/model.py - Added add_piecewise_constraint method (lines 593-827) 2. test/test_piecewise_constraints.py - Created comprehensive test suite (21 tests) Method Signature def add_piecewise_constraint( self, vars: Variable | dict[str, Variable], breakpoints: DataArray, link_dim: str | None = None, dim: str = "breakpoint", mask: DataArray | None = None, name: str | None = None, ) -> Constraint Features Implemented - Single Variable Support: Pass a single Variable directly - Multiple Variables: Pass a dict of Variables with link_dim to link them - Auto-detect link_dim: When vars is a dict, automatically detects which breakpoints dimension matches the dict keys - NaN Masking: Auto-detects masked values from NaN in breakpoints - Explicit Masking: User-provided mask support - Multi-dimensional: Works with variables that have additional coordinates (generators, timesteps, etc.) - Auto-naming: Generates names like pwl0, pwl1 automatically - Custom naming: User can specify custom names SOS2 Formulation The method creates: 1. Lambda (λ) variables with bounds [0, 1] for each breakpoint 2. SOS2 constraint ensuring at most two adjacent λ values are non-zero 3. Convexity constraint: Σλ = 1 4. Linking constraints: var = Σ(λ × breakpoint) for each variable Test Results All 21 tests pass including: - Basic single/multiple variable cases - Auto-detection of link_dim - Masking (NaN and explicit) - Multi-dimensional cases - Input validation errors - LP file output - Solver integration tests with Gurobi
Performance Improvements 1. Single Variable Case Handled Directly: No longer normalizes to dict, handles the single variable case with a direct path. 2. Single Linking Constraint for Dict Case: Instead of creating N separate linking constraints (one per variable), now creates a single constraint that covers all variables using merge() to stack variable expressions along link_dim. Code Quality Improvements 1. Added Constants in linopy/constants.py: - PWL_LAMBDA_SUFFIX = "_lambda" - PWL_CONVEX_SUFFIX = "_convex" - PWL_LINK_SUFFIX = "_link" - DEFAULT_BREAKPOINT_DIM = "breakpoint" 2. Replaced magic strings with constants throughout the implementation and tests. Result The constraint structure is now: - Before: pwl0_link_power, pwl0_link_efficiency (separate constraints) - After: pwl0_link (single constraint covering all variables) Constraints: ['pwl0_convex', 'pwl0_link'] Constraint `pwl0_link` [var: 2, generator: 2]: [power, gen1]: +1 power[gen1] - 0 λ[0] - 50 λ[1] - 100 λ[2] - 150 λ[3] = 0 [power, gen2]: +1 power[gen2] - 0 λ[0] - 50 λ[1] - 100 λ[2] - 150 λ[3] = 0 [efficiency, gen1]: +1 efficiency[gen1] - 0 λ[0] - 0.85 λ[1] - 0.92 λ[2] - 0.88 λ[3] = 0 [efficiency, gen2]: +1 efficiency[gen2] - 0 λ[0] - 0.85 λ[1] - 0.92 λ[2] - 0.88 λ[3] = 0
1. Single Variable Case: Use breakpoints.coords directly
# Before: Creating pd.Index objects in a loop
lambda_coords = [
pd.Index(breakpoints.coords[d].values, name=d) for d in lambda_dims
]
# After: Pass coords directly
lambda_var = self.add_variables(
lower=0, upper=1, coords=breakpoints.coords, ...
)
2. Dict Case: Build stacked expression directly
# Before: Multiple intermediate objects
for k in link_coords:
expr = var.to_linexpr() # Creates LinearExpression
expr_data = expr.data.expand_dims({link_dim: [k]}) # Creates Dataset
var_exprs.append(LinearExpression(expr_data, self)) # Creates another LinearExpression
stacked_vars_expr = merge(var_exprs, dim=link_dim) # Merges all
# After: Direct Dataset construction from labels
labels_list = []
for k in link_coords:
labels_list.append(var.labels.expand_dims({link_dim: [k]}))
stacked_labels = xr.concat(labels_list, dim=link_dim) # Single concat
# Build Dataset directly
stacked_expr_data = Dataset({
"coeffs": xr.ones_like(stacked_labels, dtype=float).expand_dims(TERM_DIM),
"vars": stacked_labels.expand_dims(TERM_DIM),
"const": xr.zeros_like(...),
})
stacked_vars_expr = LinearExpression(stacked_expr_data, self) # Single object
3. Combined validation with expression building
Variable existence validation now happens in the same loop that collects labels, avoiding a separate validation pass.
Key Benefits:
- Fewer intermediate objects: Avoids creating N LinearExpression objects + merge
- Direct Dataset construction: Builds the final structure in one shot
- Single xr.concat call: Instead of multiple expand_dims + merge operations
- Removed merge import: No longer needed 1. Skip NaN Check Parameter (skip_nan_check: bool = False)
# Before: Always computed O(n) scan
if mask is None:
mask = ~breakpoints.isnull()
# After: Skip when user knows data is clean
m.add_piecewise_constraint(x, breakpoints, dim='bp', skip_nan_check=True)
2. Counter-Based Name Generation
# Before: Loop until finding unused name - O(n) worst case
i = 0
while f"pwl{i}{PWL_LAMBDA_SUFFIX}" in self.variables:
i += 1
# After: O(1) counter increment
name = f"pwl{self._pwlCounter}"
self._pwlCounter += 1
3. Expression Support
# Now supports LinearExpression, not just Variable
m.add_piecewise_constraint(x + y, breakpoints, dim='bp')
# Also works in dict form
m.add_piecewise_constraint({
'total': x + y,
'cost': cost_expr
}, breakpoints, link_dim='var', dim='bp')
4. Refactored with Helper Methods
Extracted common logic into reusable private methods:
- _to_linexpr() - Convert Variable/LinearExpression to LinearExpression
- _compute_pwl_mask() - Handle mask computation with skip_nan_check
- _resolve_pwl_link_dim() - Auto-detect or validate link_dim
- _build_stacked_expr() - Build stacked expression from dict
Code Structure (Before vs After)
Before: ~250 lines, duplicated logic between single/dict cases
After: ~180 lines main method + 4 helper methods, DRY code
Test Coverage
- 22 tests total (added tests for expression support and skip_nan_check)
- All tests pass
- Linting passes
This was referenced Mar 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes # (if applicable).
Changes proposed in this Pull Request
Checklist
doc.doc/release_notes.rstof the upcoming release is included.