Feat/add piecewise linear #558
Closed
Conversation
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
- 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
Sign up for free
to 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.