diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a3eeae73..cbac968b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -25,6 +25,7 @@ repos: hooks: - id: blackdoc additional_dependencies: ['black==24.8.0'] + exclude: dev-scripts/quadratic_constraints_plan\.md - repo: https://github.com/codespell-project/codespell rev: v2.4.1 hooks: diff --git a/dev-scripts/quadratic_constraint_example.py b/dev-scripts/quadratic_constraint_example.py new file mode 100644 index 00000000..6c739c7e --- /dev/null +++ b/dev-scripts/quadratic_constraint_example.py @@ -0,0 +1,29 @@ +""" +Quadratic Constraints Example +""" + +import linopy + +m = linopy.Model() + +# Variables +x = m.add_variables(lower=0, upper=10, name="x") +y = m.add_variables(lower=0, upper=10, name="y") + +# Linear constraint +m.add_constraints(x + y <= 8, name="linear_budget") + +# Quadratic constraints +m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") +m.add_quadratic_constraints(x * y, "<=", 10, name="mixed_term") + +# Objective: maximize x + 2y +m.add_objective(x + 2 * y, sense="max") + +# Solve +m.solve(solver_name="gurobi") + +# Results +print(f"x = {x.solution.values.item():.4f}") +print(f"y = {y.solution.values.item():.4f}") +print(f"Objective = {m.objective.value:.4f}") diff --git a/dev-scripts/quadratic_constraints_plan.md b/dev-scripts/quadratic_constraints_plan.md new file mode 100644 index 00000000..14857480 --- /dev/null +++ b/dev-scripts/quadratic_constraints_plan.md @@ -0,0 +1,603 @@ + + + +# Quadratic Constraints and Expressions in Linopy + +## Design Document and Implementation Plan + +### Executive Summary + +This document outlines a plan to extend linopy with support for **quadratic constraints** (QCP/QCQP). Currently, linopy supports: +- Linear constraints (`Constraint` class) +- Linear expressions (`LinearExpression` class) +- Quadratic expressions (`QuadraticExpression` class) — **only for objectives** + +The goal is to enable quadratic constraints of the form: + +``` +x'Qx + a'x ≤ b (or ≥, =) +``` + +This feature would make linopy suitable for a broader class of optimization problems including convex QCPs, non-convex QCQPs (where supported by solvers), and Second-Order Cone Programs (SOCPs). + +--- + +## 1. Current Architecture Analysis + +### 1.1 Linear Expression (`expressions.py`) + +The `LinearExpression` class stores: +```python +Dataset { + 'coeffs': DataArray[float] # shape: (..., _term) + 'vars': DataArray[int] # shape: (..., _term), variable labels + 'const': DataArray[float] # shape: (...) +} +``` + +Key method for constraint creation (`expressions.py:843-866`): +```python +def to_constraint(self, sign: SignLike, rhs: ConstantLike) -> Constraint: + all_to_lhs = (self - rhs).data + data = assign_multiindex_safe( + all_to_lhs[["coeffs", "vars"]], + sign=sign, + rhs=-all_to_lhs.const + ) + return constraints.Constraint(data, model=self.model) +``` + +### 1.2 Quadratic Expression (`expressions.py`) + +The `QuadraticExpression` class adds a `_factor` dimension (size=2) for storing two variables per quadratic term: +```python +Dataset { + 'coeffs': DataArray[float] # shape: (..., _factor, _term) + 'vars': DataArray[int] # shape: (..., _factor, _term) + 'const': DataArray[float] # shape: (...) +} +``` + +Currently, `QuadraticExpression.to_constraint()` raises `NotImplementedError` (`expressions.py:1805-1808`). + +### 1.3 Linear Constraint (`constraints.py`) + +The `Constraint` class stores: +```python +Dataset { + 'coeffs': DataArray[float] # LHS coefficients + 'vars': DataArray[int] # Variable labels + 'sign': DataArray[str] # '=', '<=', '>=' + 'rhs': DataArray[float] # Right-hand side + 'labels': DataArray[int] # Constraint labels (-1 if masked) + 'dual': DataArray[float] # [OPTIONAL] Dual values +} +``` + +### 1.4 Solver Support for Quadratic Constraints + +| Solver | QCP Support | QCQP Support | Non-Convex Support | +|--------|-------------|--------------|-------------------| +| Gurobi | ✅ Yes | ✅ Yes | ✅ Yes (v9.0+) | +| CPLEX | ✅ Yes | ✅ Yes | ⚠️ Limited | +| MOSEK | ✅ Yes | ✅ Yes | ❌ Convex only | +| Xpress | ✅ Yes | ✅ Yes | ⚠️ Limited | +| COPT | ✅ Yes | ✅ Yes | ⚠️ Limited | +| SCIP | ✅ Yes | ✅ Yes | ✅ Yes | +| HiGHS | ❌ No | ❌ No | ❌ No | +| GLPK | ❌ No | ❌ No | ❌ No | +| CBC | ❌ No | ❌ No | ❌ No | + +**Key Insight**: HiGHS (a common default solver) does NOT support quadratic constraints. This has implications for default behavior and error handling. + +--- + +## 2. Proposed Design + +### 2.1 New Class: `QuadraticConstraint` + +Create a new `QuadraticConstraint` class parallel to `Constraint`: + +```python +class QuadraticConstraint: + """ + A quadratic constraint of the form: x'Qx + a'x ≤ b (or ≥, =) + + Dataset structure: + { + 'quad_coeffs': DataArray[float] # shape: (..., _factor, _qterm) + 'quad_vars': DataArray[int] # shape: (..., _factor, _qterm) + 'lin_coeffs': DataArray[float] # shape: (..., _term) + 'lin_vars': DataArray[int] # shape: (..., _term) + 'sign': DataArray[str] # '=', '<=', '>=' + 'rhs': DataArray[float] # Right-hand side constant + 'labels': DataArray[int] # Constraint labels + 'dual': DataArray[float] # [OPTIONAL] Dual values (only for convex) + } + """ +``` + +**Design Rationale**: +- Separate `quad_*` and `lin_*` arrays to allow efficient handling of purely linear terms +- Use `_qterm` dimension (distinct from `_term`) for quadratic terms +- Maintain API consistency with `Constraint` class + +### 2.2 Container Class: `QuadraticConstraints` + +Add a container class analogous to `Constraints`: + +```python +class QuadraticConstraints: + """ + Container for multiple QuadraticConstraint objects. + Provides dict-like access and aggregation properties. + """ +``` + +### 2.3 Model Integration + +Extend the `Model` class: + +```python +class Model: + def __init__(self, ...): + self.constraints = Constraints() # Linear constraints + self.quadratic_constraints = QuadraticConstraints() # NEW + + def add_quadratic_constraints( + self, + lhs: QuadraticExpression | Callable, + sign: SignLike, + rhs: ConstantLike, + name: str | None = None, + coords: CoordsLike | None = None, + mask: MaskLike | None = None, + ) -> QuadraticConstraint: + """Add quadratic constraint(s) to the model.""" + + @property + def has_quadratic_constraints(self) -> bool: + """Return True if model has any quadratic constraints.""" + + @property + def type(self) -> str: + """Return problem type: 'LP', 'QP', 'MILP', 'MIQP', 'QCP', 'QCQP', etc.""" +``` + +### 2.4 Expression API Changes + +Implement `QuadraticExpression.to_constraint()`: + +```python +def to_constraint(self, sign: SignLike, rhs: ConstantLike) -> QuadraticConstraint: + """ + Convert quadratic expression to a quadratic constraint. + + Parameters + ---------- + sign : str + Constraint sense: '<=', '>=', or '=' + rhs : float or array-like + Right-hand side constant + + Returns + ------- + QuadraticConstraint + """ +``` + +Enable comparison operators on `QuadraticExpression`: +```python +# These would create QuadraticConstraint objects +quad_expr <= 10 # Works (returns QuadraticConstraint) +quad_expr >= 5 # Works +quad_expr == 0 # Works +``` + +--- + +## 3. Implementation Details + +### 3.1 Data Storage for Quadratic Constraints + +**Option A: Unified Storage** (simpler, less efficient) +```python +# Store everything with _factor dimension, linear terms have vars[_factor=1] = -1 +Dataset { + 'coeffs': DataArray[float] # shape: (..., _factor, _term) + 'vars': DataArray[int] # shape: (..., _factor, _term) + 'sign': DataArray[str] + 'rhs': DataArray[float] + 'labels': DataArray[int] +} +``` + +**Option B: Split Storage** (recommended, more efficient) +```python +# Separate linear and quadratic terms +Dataset { + 'quad_coeffs': DataArray[float] # shape: (..., _factor, _qterm) + 'quad_vars': DataArray[int] # shape: (..., _factor, _qterm) + 'lin_coeffs': DataArray[float] # shape: (..., _term) + 'lin_vars': DataArray[int] # shape: (..., _term) + 'sign': DataArray[str] + 'rhs': DataArray[float] + 'labels': DataArray[int] +} +``` + +**Recommendation**: Option B provides clearer separation, easier debugging, and more efficient matrix construction for solvers that handle linear and quadratic parts separately. + +### 3.2 Matrix Representation + +Add to `MatrixAccessor`: + +```python +@property +def Qc(self) -> list[tuple[csc_matrix, ndarray, float, str]]: + """ + List of quadratic constraint matrices. + + Returns list of tuples: (Q_i, a_i, b_i, sense_i) + where constraint i is: x'Q_i x + a_i'x {sense_i} b_i + """ + +@property +def qc_labels(self) -> ndarray: + """Labels of quadratic constraints.""" +``` + +### 3.3 Solver Export Functions + +#### LP File Format + +The LP file format supports quadratic constraints in the `QCROWS` section: + +``` +Subject To + c1: x + y <= 10 + +QCROWS + qc1: [ x^2 + 2 x * y + y^2 ] + x + y <= 5 +End +``` + +Add function: +```python +def quadratic_constraints_to_file( + m: Model, + f: BufferedWriter, + progress: bool = False, + explicit_coordinate_names: bool = False, +) -> None: + """Write quadratic constraints to LP file.""" +``` + +#### Direct API Export + +**Gurobi** (`addQConstr` or matrix interface): +```python +def to_gurobipy(m: Model, env=None, ...): + # ... existing code ... + + # Add quadratic constraints + for qc in m.quadratic_constraints: + model.addQConstr(Q, sense, rhs, name) +``` + +**MOSEK** (`putqconk`): +```python +def to_mosek(m: Model, task=None, ...): + # ... existing code ... + + # Add quadratic constraints + for k, (Q, a, b, sense) in enumerate(M.Qc): + task.putqconk(k, Q.row, Q.col, Q.data) + task.putarow(k, a.nonzero()[0], a[a.nonzero()]) +``` + +### 3.4 Solution Handling + +Quadratic constraints may have dual values (for convex problems): + +```python +class QuadraticConstraint: + @property + def dual(self) -> DataArray | None: + """ + Dual values for the quadratic constraint. + + Note: Only available for convex quadratic constraints + and when the solver provides them. + """ +``` + +--- + +## 4. API Design Considerations + +### 4.1 Consistency with Existing API + +The API should feel natural to existing linopy users: + +```python +import linopy as lp + +m = lp.Model() +x = m.add_variables(coords=[range(3)], name='x') +y = m.add_variables(name='y') + +# Linear constraint (existing) +m.add_constraints(x.sum() <= 10, name='linear_budget') + +# Quadratic constraint (new - Option A: via add_constraints) +m.add_constraints(x @ x + y <= 5, name='quad_con') + +# Quadratic constraint (new - Option B: via add_quadratic_constraints) +m.add_quadratic_constraints(x @ x + y <= 5, name='quad_con') +``` + +**Question for discussion**: Should quadratic constraints be added via: +- **Option A**: Same `add_constraints()` method (auto-detect based on expression type) +- **Option B**: Separate `add_quadratic_constraints()` method + +**Recommendation**: Start with **Option B** for clarity, with Option A as a future enhancement. This makes the API explicit about what type of constraint is being added. + +### 4.2 Operator Overloading + +Enable natural syntax on `QuadraticExpression`: + +```python +# All should return QuadraticConstraint +x * x <= 10 +(x @ x) + y >= 5 +2 * x * y == 0 +``` + +### 4.3 Error Handling + +```python +# Clear error for unsupported solvers +m.solve(solver='highs') +# Raises: "Solver 'highs' does not support quadratic constraints. +# Use one of: ['gurobi', 'cplex', 'mosek', 'xpress', 'copt', 'scip']" + +# Warning for non-convex constraints with convex-only solvers +m.solve(solver='mosek') +# Warning: "MOSEK requires convex quadratic constraints. +# Non-convex constraints may cause solver failure." +``` + +--- + +## 5. File Structure Changes + +### 5.1 New Files + +None required - extend existing modules. + +### 5.2 Modified Files + +| File | Changes | +|------|---------| +| `expressions.py` | Implement `QuadraticExpression.to_constraint()` | +| `constraints.py` | Add `QuadraticConstraint` and `QuadraticConstraints` classes | +| `model.py` | Add `add_quadratic_constraints()`, `quadratic_constraints` property | +| `io.py` | Add LP file export for quadratic constraints | +| `solvers.py` | Add `QUADRATIC_CONSTRAINT_SOLVERS` list | +| `matrices.py` | Add `Qc` property for quadratic constraint matrices | +| `constants.py` | Add any new constants (e.g., `QTERM_DIM = "_qterm"`) | + +### 5.3 New Test Files + +- `test/test_quadratic_constraint.py` — Unit tests for QuadraticConstraint class +- `test/test_quadratic_optimization.py` — Integration tests with solvers + +--- + +## 6. Implementation Phases + +### Phase 1: Core Data Structures (Week 1-2) + +1. Add `QTERM_DIM` constant +2. Implement `QuadraticConstraint` class with basic functionality: + - `__init__`, `__repr__`, `__getitem__` + - Properties: `labels`, `sign`, `rhs`, `lhs`, `mask`, etc. + - Methods: `to_polars()`, `flat` +3. Implement `QuadraticConstraints` container +4. Add `QuadraticExpression.to_constraint()` method + +### Phase 2: Model Integration (Week 2-3) + +1. Add `Model.quadratic_constraints` property +2. Implement `Model.add_quadratic_constraints()` method +3. Update `Model.type` property for QCP/QCQP detection +4. Add `Model.has_quadratic_constraints` property +5. Update constraint label management + +### Phase 3: Solver Export (Week 3-4) + +1. Extend LP file writer with `QCROWS` section +2. Update `to_gurobipy()` for quadratic constraints +3. Update `to_mosek()` for quadratic constraints +4. Update other direct-API solvers (CPLEX, Xpress, COPT) +5. Add solver compatibility checks + +### Phase 4: Solution Handling (Week 4-5) + +1. Parse quadratic constraint duals from solver results +2. Map duals back to constraint coordinates +3. Add `QuadraticConstraint.dual` property + +### Phase 5: Testing & Documentation (Week 5-6) + +1. Comprehensive unit tests +2. Integration tests with each supported solver +3. Update documentation and examples +4. Add tutorial notebook + +--- + +## 7. Code Examples + +### 7.1 Basic Usage + +```python +import linopy as lp + +# Create model +m = lp.Model() + +# Add variables +x = m.add_variables(lower=0, upper=10, name='x') +y = m.add_variables(lower=0, upper=10, name='y') + +# Quadratic objective (already supported) +m.add_objective(x**2 + y**2 + x + y) + +# Linear constraint (already supported) +m.add_constraints(x + y >= 1, name='sum_bound') + +# NEW: Quadratic constraint +m.add_quadratic_constraints(x**2 + y**2 <= 25, name='circle') + +# Solve with quadratic-capable solver +m.solve(solver='gurobi') + +print(f"x = {x.solution.values}") +print(f"y = {y.solution.values}") +``` + +### 7.2 Multi-dimensional Quadratic Constraints + +```python +import linopy as lp +import pandas as pd + +m = lp.Model() + +# Index set +times = pd.Index(range(24), name='time') + +# Variables +power = m.add_variables(coords=[times], name='power') +reserve = m.add_variables(coords=[times], name='reserve') + +# Quadratic constraint at each time step +# power²[t] + reserve²[t] <= capacity[t] +capacity = [100] * 24 +qc = m.add_quadratic_constraints( + power**2 + reserve**2 <= capacity, + name='capacity_limit' +) + +print(qc) +# QuadraticConstraint `capacity_limit` [time: 24]: +# ------------------------------------------- +# 0: power[0]² + reserve[0]² <= 100 +# 1: power[1]² + reserve[1]² <= 100 +# ... +``` + +### 7.3 Rule-based Quadratic Constraints + +```python +def capacity_rule(m, t): + """Quadratic capacity constraint at time t.""" + return m['power'][t]**2 + m['reserve'][t]**2 <= capacity[t] + +m.add_quadratic_constraints(capacity_rule, coords=[times], name='capacity') +``` + +--- + +## 8. Design Decisions (Resolved) + +### Q1: Unified vs Separate `add_constraints` method? + +**Decision**: Use separate `add_quadratic_constraints()` method. + +```python +m.add_quadratic_constraints(x**2 <= 10) # Quadratic +m.add_constraints(x <= 10) # Linear +``` + +**Rationale**: Explicit API is clearer. Auto-detection can be added later by routing `add_constraints()` to `add_quadratic_constraints()` when a `QuadraticExpression` is detected. + +### Q2: Storage location for quadratic constraints? + +**Decision**: Separate containers. + +```python +m.constraints # Linear only +m.quadratic_constraints # Quadratic only +``` + +**Rationale**: Simpler implementation, matches current pattern for variables/constraints, avoids complexity of mixed container. + +### Q3: How to handle mixed linear+quadratic in same named constraint group? + +**Decision**: Each named constraint should be uniformly linear or quadratic. Mixed cases require two separate constraints. + +### Q4: Convexity checking? + +**Decision**: Defer to solver. + +**Rationale**: Avoids computational overhead and complex eigenvalue analysis. Solvers like Gurobi and MOSEK provide clear error messages for non-convex constraints. We can add clear documentation about convexity requirements. + +--- + +## 9. References + +### Solver Documentation + +- [Gurobi Quadratic Constraints](https://docs.gurobi.com/projects/optimizer/en/current/concepts/modeling/constraints.html) +- [MOSEK Conic Constraints](https://docs.mosek.com/latest/pythonapi/tutorial-cqo-shared.html) +- [CPLEX Quadratic Constraints](https://www.ibm.com/docs/en/icos/latest?topic=programming-adding-quadratic-constraints) + +### Related Issues/PRs + +- Current `NotImplementedError` in `expressions.py:1805-1808` +- Test showing limitation: `test/test_quadratic_expression.py:331` + +### Mathematical Background + +A quadratically constrained quadratic program (QCQP) has the form: + +``` +minimize (1/2)x'Q₀x + c'x +subject to (1/2)x'Qᵢx + aᵢ'x ≤ bᵢ for i = 1,...,m + Ax = b + l ≤ x ≤ u +``` + +Where: +- Q₀ is the objective quadratic matrix (may be zero for QCP) +- Qᵢ are constraint quadratic matrices +- aᵢ are constraint linear coefficient vectors +- bᵢ are constraint right-hand sides + +For **convex** QCPs, all Qᵢ must be positive semi-definite (PSD). + +--- + +## 10. Summary + +Adding quadratic constraint support to linopy is a significant but feasible enhancement. The key design decisions are: + +1. **New class**: `QuadraticConstraint` parallel to `Constraint` +2. **Split storage**: Separate `quad_*` and `lin_*` arrays for efficiency +3. **Explicit API**: `add_quadratic_constraints()` method +4. **Solver filtering**: Clear error messages for unsupported solvers +5. **Phased implementation**: Core → Model → Export → Tests + +This enhancement would expand linopy's capabilities to cover: +- Convex QCPs (portfolio optimization, geometric programming) +- QCQPs (facility location, engineering design) +- SOCPs via quadratic constraint reformulation + +--- + +*Document Version: 1.0* +*Date: 2025-01-25* +*Status: Draft for Discussion* diff --git a/dev-scripts/quadratic_constraints_remaining_tasks.md b/dev-scripts/quadratic_constraints_remaining_tasks.md new file mode 100644 index 00000000..389b77bf --- /dev/null +++ b/dev-scripts/quadratic_constraints_remaining_tasks.md @@ -0,0 +1,271 @@ +# Quadratic Constraints - Remaining Tasks + +## Implementation Status + +### ✅ Completed (Phase 1 - Core Implementation) + +1. **Core Data Structures** + - `QTERM_DIM` constant in `constants.py` + - `QuadraticConstraint` class in `constraints.py` + - `QuadraticConstraints` container class + - `QuadraticExpression.to_constraint()` method + +2. **Model Integration** + - `Model.quadratic_constraints` property + - `Model.add_quadratic_constraints()` method + - `Model.has_quadratic_constraints` property + - `Model.type` property updated for QCP/QCQP detection + +3. **Solver Support** + - `QUADRATIC_CONSTRAINT_SOLVERS` list in `solvers.py` + - Solver validation in `Model.solve()` - rejects unsupported solvers + +4. **Export Functionality** + - LP file export via `quadratic_constraints_to_file()` in `io.py` + - Gurobi direct API export via updated `to_gurobipy()` + +5. **Tests** + - `test/test_quadratic_constraint.py` - 23 unit tests + - Updated `test_quadratic_to_constraint` in `test_quadratic_expression.py` + +--- + +## 🔲 Remaining Tasks (Phase 2) + +### High Priority + +#### 1. Matrix Accessor for Quadratic Constraints +**File:** `linopy/matrices.py` + +Add `Qc` property to `MatrixAccessor` class for quadratic constraint matrices. + +```python +@property +def Qc(self) -> list[scipy.sparse.csc_matrix]: + """Return list of Q matrices for quadratic constraints.""" + # Each quadratic constraint has its own Q matrix + pass + + +@property +def qc_linear(self) -> scipy.sparse.csc_matrix: + """Return linear coefficients for quadratic constraints.""" + pass + + +@property +def qc_sense(self) -> np.ndarray: + """Return sense array for quadratic constraints.""" + pass + + +@property +def qc_rhs(self) -> np.ndarray: + """Return RHS values for quadratic constraints.""" + pass +``` + +#### 2. MOSEK Direct API Support +**File:** `linopy/io.py` - `to_mosek()` function + +Add quadratic constraint support to MOSEK export: +```python +# After linear constraints section +if len(m.quadratic_constraints): + for name in m.quadratic_constraints: + qcon = m.quadratic_constraints[name] + # Use task.appendcone() or task.putqconk() for quadratic constraints +``` + +Reference: [MOSEK Python API - Quadratic Constraints](https://docs.mosek.com/latest/pythonapi/tutorial-qcqo.html) + +#### 3. HiGHSpy Validation +**File:** `linopy/io.py` - `to_highspy()` function + +Add explicit error if model has quadratic constraints: +```python +if len(m.quadratic_constraints): + raise ValueError( + "HiGHS does not support quadratic constraints. " + "Use a solver that supports QCP: gurobi, cplex, mosek, xpress, copt, scip" + ) +``` + +#### 4. Solution Retrieval for Quadratic Constraints +**Files:** `linopy/solvers.py`, `linopy/constraints.py` + +- Add dual value retrieval for quadratic constraints (where supported) +- Store duals in `QuadraticConstraint.dual` property +- Update solver result parsing + +### Medium Priority + +#### 5. Multi-dimensional Quadratic Constraints +**File:** `linopy/constraints.py`, `linopy/model.py` + +Currently, quadratic constraints are primarily scalar. Add support for: +- Broadcasting over coordinates (like linear constraints) +- `iterate_slices()` support for memory-efficient processing +- Coordinate-based indexing + +Example API: +```python +# Should work with coordinates +m.add_quadratic_constraints( + x * x + y * y, "<=", 100, name="qc" # where x, y have dims=['time', 'node'] +) +``` + +#### 6. Constraint Modification Methods +**File:** `linopy/constraints.py` + +Add methods to `QuadraticConstraint`: +```python +def modify_rhs(self, new_rhs: ConstantLike) -> None: + """Modify the right-hand side of the constraint.""" + + +def modify_coeffs(self, new_coeffs: xr.DataArray) -> None: + """Modify coefficients of the constraint.""" +``` + +#### 7. netCDF Serialization +**File:** `linopy/io.py` - `to_netcdf()` and `read_netcdf()` + +Add quadratic constraints to model serialization: +```python +# In to_netcdf() +qcons = [ + with_prefix(qcon.data, f"quadratic_constraints-{name}") + for name, qcon in m.quadratic_constraints.items() +] + +# In read_netcdf() +# Parse quadratic_constraints-* prefixed datasets +``` + +### Low Priority + +#### 8. Convexity Checking (Optional) +**File:** `linopy/constraints.py` or new `linopy/analysis.py` + +Add optional convexity verification: +```python +def check_convexity(self) -> bool: + """ + Check if quadratic constraint is convex. + + A quadratic constraint x'Qx + a'x <= b is convex if Q is + positive semidefinite. + """ + # Extract Q matrix + # Check eigenvalues or use Cholesky decomposition + pass +``` + +#### 9. Constraint Printing Improvements +**File:** `linopy/constraints.py` + +Enhance `_format_single_constraint()` for better display: +- Handle large constraints with truncation +- Add option for matrix form display +- Support LaTeX output + +#### 10. Documentation +**Files:** `doc/` directory + +- Add quadratic constraints section to user guide +- Document supported solvers and their limitations +- Add examples for common QCP formulations (portfolio optimization, etc.) + +--- + +## Testing Tasks + +### Unit Tests to Add + +1. **Multi-dimensional constraints** - Test with coordinates +2. **Edge cases** - Empty constraints, single term, all linear terms +3. **Numerical precision** - Very small/large coefficients +4. **Memory efficiency** - Large constraint sets with `iterate_slices` + +### Integration Tests + +1. **Solver round-trip** - Create model, solve, verify solution +2. **File format round-trip** - Write LP, read back, compare +3. **Cross-solver consistency** - Same problem, multiple solvers + +### Solver-Specific Tests + +```python +@pytest.mark.parametrize("solver", ["gurobi", "mosek", "cplex"]) +def test_qcp_solve(solver): + """Test solving QCP with different solvers.""" + if solver not in available_solvers: + pytest.skip(f"{solver} not available") + # ... test code +``` + +--- + +## Code Quality Tasks + +1. **Type hints** - Ensure all new functions have complete type annotations +2. **Docstrings** - Add NumPy-style docstrings to all public methods +3. **Linting** - Run `ruff check --fix` on all modified files +4. **MyPy** - Fix any type errors in new code + +--- + +## Architecture Notes + +### Data Structure + +``` +QuadraticConstraint.data (xarray.Dataset): +├── quad_coeffs: (_qterm, _factor) float64 # Quadratic term coefficients +├── quad_vars: (_qterm, _factor) int64 # Variable indices for quad terms +├── lin_coeffs: (_term) float64 # Linear term coefficients +├── lin_vars: (_term) int64 # Variable indices for linear terms +├── sign: str # "<=", ">=", or "=" +├── rhs: float64 # Right-hand side value +└── labels: int64 # Constraint label/index +``` + +### LP File Format + +``` +qc0: ++3.0 x0 ++4.0 x1 ++ [ ++1.0 x0 ^ 2 ++2.0 x0 * x1 ++1.0 x1 ^ 2 +] +<= +100.0 +``` + +### Solver Compatibility Matrix + +| Solver | QCP Support | API Method | +|--------|-------------|------------| +| Gurobi | ✅ | `addQConstr()` | +| CPLEX | ✅ | `add_quadratic_constraint()` | +| MOSEK | ✅ | `putqconk()` | +| Xpress | ✅ | `addConstraint()` with quadratic | +| COPT | ✅ | `addQConstr()` | +| SCIP | ✅ | LP file import | +| HiGHS | ❌ | Not supported | + +--- + +## Suggested Task Order for Next Agent + +1. **Matrix Accessor** (`matrices.py`) - Enables programmatic access to constraint data +2. **MOSEK support** (`io.py`) - Important solver for QCP +3. **Multi-dimensional constraints** - Core functionality improvement +4. **netCDF serialization** - Model persistence +5. **Documentation** - User-facing docs + +Each task is relatively independent and can be completed in a single session. diff --git a/dev-scripts/quadratic_constraints_status.md b/dev-scripts/quadratic_constraints_status.md new file mode 100644 index 00000000..19c93c68 --- /dev/null +++ b/dev-scripts/quadratic_constraints_status.md @@ -0,0 +1,169 @@ +# Quadratic Constraints Implementation Status + +**Date:** 2024-11-25 +**Status:** Phase 1 Complete ✅ + +## Related Documents + +- **Original Plan:** `dev-scripts/quadratic_constraints_plan.md` - Detailed design document +- **Remaining Tasks:** `dev-scripts/quadratic_constraints_remaining_tasks.md` - What's left to do + +--- + +## Summary + +Quadratic constraints (QCP/QCQP) have been successfully added to linopy. Users can now create and solve optimization problems with constraints of the form: + +``` +x'Qx + a'x ≤ b (or ≥, =) +``` + +## Usage Example + +```python +import linopy + +m = linopy.Model() + +# Variables +x = m.add_variables(lower=0, name="x") +y = m.add_variables(lower=0, name="y") + +# Linear constraint (existing) +m.add_constraints(x + y <= 10, name="budget") + +# Quadratic constraint (NEW!) +m.add_quadratic_constraints( + x * x + 2 * x * y + y * y + 3 * x + 4 * y, "<=", 100, name="quadratic_budget" +) + +# Objective +m.add_objective(x + 2 * y) + +# Solve (with Gurobi, MOSEK, CPLEX, etc.) +m.solve(solver_name="gurobi") +``` + +## What Was Implemented + +### Core Components + +| Component | File | Status | +|-----------|------|--------| +| `QTERM_DIM` constant | `constants.py` | ✅ | +| `QuadraticConstraint` class | `constraints.py` | ✅ | +| `QuadraticConstraints` container | `constraints.py` | ✅ | +| `QuadraticExpression.to_constraint()` | `expressions.py` | ✅ | +| `Model.add_quadratic_constraints()` | `model.py` | ✅ | +| `Model.quadratic_constraints` property | `model.py` | ✅ | +| `Model.has_quadratic_constraints` | `model.py` | ✅ | +| Model type detection (QCLP/QCQP) | `model.py` | ✅ | +| Solver validation | `model.py` | ✅ | +| `QUADRATIC_CONSTRAINT_SOLVERS` list | `solvers.py` | ✅ | +| LP file export | `io.py` | ✅ | +| Gurobi direct export | `io.py` | ✅ | +| Unit tests | `test/test_quadratic_constraint.py` | ✅ | + +### Supported Solvers + +| Solver | Support | Notes | +|--------|---------|-------| +| Gurobi | ✅ | Full support via `addQConstr()` | +| CPLEX | ✅ | Via LP file | +| MOSEK | ✅ | Via LP file (direct API pending) | +| Xpress | ✅ | Via LP file | +| COPT | ✅ | Via LP file | +| SCIP | ✅ | Via LP file | +| HiGHS | ❌ | Does not support QC - validation error raised | + +### Model Type Strings + +The `Model.type` property now returns: + +| Type | Meaning | +|------|---------| +| `LP` | Linear constraints, linear objective | +| `QP` | Linear constraints, quadratic objective | +| `QCLP` | Quadratic constraints, linear objective | +| `QCQP` | Quadratic constraints, quadratic objective | +| `MILP` | Mixed-integer linear | +| `MIQP` | Mixed-integer quadratic objective | +| `MIQCLP` | Mixed-integer with quadratic constraints | +| `MIQCQP` | Mixed-integer QC with quadratic objective | + +--- + +## What's NOT Yet Implemented + +See `dev-scripts/quadratic_constraints_remaining_tasks.md` for full details. + +### High Priority +1. Matrix accessor (`Qc` property in `matrices.py`) +2. MOSEK direct API support +3. HiGHS explicit validation in `to_highspy()` +4. Solution/dual retrieval for quadratic constraints + +### Medium Priority +5. Multi-dimensional quadratic constraints (with coordinates) +6. netCDF serialization +7. Constraint modification methods + +### Low Priority +8. Convexity checking (optional) +9. Documentation + +--- + +## Testing + +All 23 unit tests pass: + +```bash +pytest test/test_quadratic_constraint.py -v +# 23 passed +``` + +Existing quadratic expression tests updated: +```bash +pytest test/test_quadratic_expression.py -v +# 32 passed (test_quadratic_to_constraint updated) +``` + +--- + +## Design Decisions Made + +1. **Separate method**: Using `add_quadratic_constraints()` instead of overloading `add_constraints()` - clearer API, can add auto-detection later + +2. **Separate container**: `Model.quadratic_constraints` is separate from `Model.constraints` - cleaner code, explicit handling + +3. **Defer convexity**: Convexity checking deferred to solver - avoids false positives, solvers handle it better + +4. **LP format**: Using Gurobi-style LP format with `[ ]` brackets for quadratic terms + +--- + +## Files Modified + +``` +linopy/ +├── constants.py # +QTERM_DIM +├── constraints.py # +QuadraticConstraint, +QuadraticConstraints (~500 lines) +├── expressions.py # +to_constraint() for QuadraticExpression +├── model.py # +add_quadratic_constraints(), properties, validation +├── solvers.py # +QUADRATIC_CONSTRAINT_SOLVERS +└── io.py # +quadratic_constraints_to_file(), updated to_gurobipy() + +test/ +├── test_quadratic_constraint.py # NEW - 23 tests +└── test_quadratic_expression.py # Updated 1 test +``` + +--- + +## Next Steps for Code Agent + +1. Pick a task from `quadratic_constraints_remaining_tasks.md` +2. Tasks are independent - any order works +3. Recommended first: Matrix Accessor (enables other features) +4. Run `pytest test/test_quadratic_constraint.py` to verify nothing breaks diff --git a/doc/api.rst b/doc/api.rst index 6011aa81..c6fc543d 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -17,6 +17,7 @@ Creating a model model.Model model.Model.add_variables model.Model.add_constraints + model.Model.add_quadratic_constraints model.Model.add_objective model.Model.linexpr model.Model.remove_constraints @@ -109,6 +110,50 @@ Constraints constraints.Constraints.to_matrix +QuadraticConstraint +------------------- + +``QuadraticConstraint`` stores quadratic constraints of the form x'Qx + a'x <= b (or >=, =). + +.. autosummary:: + :toctree: generated/ + + constraints.QuadraticConstraint + constraints.QuadraticConstraint.quad_coeffs + constraints.QuadraticConstraint.quad_vars + constraints.QuadraticConstraint.lin_coeffs + constraints.QuadraticConstraint.lin_vars + constraints.QuadraticConstraint.lhs + constraints.QuadraticConstraint.sign + constraints.QuadraticConstraint.rhs + constraints.QuadraticConstraint.dual + constraints.QuadraticConstraint.flat + constraints.QuadraticConstraint.to_polars + + +QuadraticConstraints +-------------------- + +``QuadraticConstraints`` is a container for storing multiple quadratic constraint arrays. + +.. autosummary:: + :toctree: generated/ + + constraints.QuadraticConstraints + constraints.QuadraticConstraints.add + constraints.QuadraticConstraints.remove + constraints.QuadraticConstraints.inequalities + constraints.QuadraticConstraints.equalities + constraints.QuadraticConstraints.sanitize_zeros + constraints.QuadraticConstraints.sanitize_missings + constraints.QuadraticConstraints.flat + constraints.QuadraticConstraints.ncons + constraints.QuadraticConstraints.labels + constraints.QuadraticConstraints.sign + constraints.QuadraticConstraints.rhs + constraints.QuadraticConstraints.dual + + IO functions ============ diff --git a/doc/index.rst b/doc/index.rst index f05f89f3..fa9e9eb6 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -109,6 +109,7 @@ This package is published under MIT license. creating-variables creating-expressions creating-constraints + quadratic-constraints manipulating-models testing-framework transport-tutorial diff --git a/doc/quadratic-constraints.nblink b/doc/quadratic-constraints.nblink new file mode 100644 index 00000000..5b90a26b --- /dev/null +++ b/doc/quadratic-constraints.nblink @@ -0,0 +1,3 @@ +{ + "path": "../examples/quadratic-constraints.ipynb" +} diff --git a/doc/release_notes.rst b/doc/release_notes.rst index d127cffe..5b8b613b 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -3,6 +3,14 @@ Release Notes .. Upcoming Version +* Add support for quadratic constraints (QCP/QCQP). Use `model.add_quadratic_constraints()` to add constraints of the form x'Qx + a'x <= b. Supported solvers: Gurobi, MOSEK, COPT. Gurobi also supports nonconvex quadratic constraints. +* Add MPS export support for models with quadratic constraints (requires Gurobi). +* Add netCDF serialization support for models with quadratic constraints. +* Add MOSEK direct API support for quadratic constraints. +* Add dual values for quadratic constraints (use `QCPDual=1` solver option with Gurobi). +* Add `model.quadratic_constraints` container for inspecting quadratic constraints. +* Add `model.matrices` accessor properties for quadratic constraints: `qclabels`, `qc_sense`, `qc_rhs`, `Qc`, `qc_linear`. +* Fix `expression.shape`to always exclude the _term dimension * Fix compatibility for xpress versions below 9.6 (regression) Version 0.5.8 diff --git a/examples/quadratic-constraints.ipynb b/examples/quadratic-constraints.ipynb new file mode 100644 index 00000000..ce67fbef --- /dev/null +++ b/examples/quadratic-constraints.ipynb @@ -0,0 +1,382 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro", + "metadata": {}, + "source": [ + "# Quadratic Constraints\n", + "\n", + "This example demonstrates how to create and solve models with quadratic constraints (QCP/QCQP). A quadratic constraint has the form:\n", + "\n", + "$$ x^T Q x + a^T x \\leq b $$\n", + "\n", + "Quadratic constraints work seamlessly with coordinates, just like linear constraints." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:49.055855Z", + "start_time": "2025-12-03T14:04:45.460505Z" + } + }, + "outputs": [], + "source": [ + "import pandas as pd\n", + "\n", + "import linopy" + ] + }, + { + "cell_type": "markdown", + "id": "setup-header", + "metadata": {}, + "source": [ + "## Problem Setup\n", + "\n", + "We model a resource allocation problem over multiple time periods. For each time $t$, we want to maximize $x_t + 2y_t$ subject to:\n", + "- A budget constraint: $x_t + y_t \\leq b_t$\n", + "- A risk constraint (quadratic): $x_t^2 + y_t^2 \\leq r_t$\n", + "\n", + "where the budget $b_t$ and risk limit $r_t$ vary over time." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "model-setup", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:49.840053Z", + "start_time": "2025-12-03T14:04:49.669343Z" + } + }, + "outputs": [], + "source": [ + "m = linopy.Model()\n", + "\n", + "time = pd.Index(range(5), name=\"time\")\n", + "\n", + "x = m.add_variables(lower=0, coords=[time], name=\"x\")\n", + "y = m.add_variables(lower=0, coords=[time], name=\"y\")" + ] + }, + { + "cell_type": "markdown", + "id": "constraints-header", + "metadata": {}, + "source": [ + "## Adding Constraints\n", + "\n", + "Linear constraints work as usual. We use a `pd.Series` for the time-varying budget:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "linear-constraint", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:50.203648Z", + "start_time": "2025-12-03T14:04:50.079589Z" + } + }, + "outputs": [], + "source": [ + "budget = pd.Series([8, 9, 10, 11, 12], index=time)\n", + "\n", + "m.add_constraints(x + y <= budget, name=\"budget\")" + ] + }, + { + "cell_type": "markdown", + "id": "qc-header", + "metadata": {}, + "source": [ + "Quadratic constraints use `add_quadratic_constraints`. The risk limits also vary over time:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "quad-constraint", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:50.639153Z", + "start_time": "2025-12-03T14:04:50.544579Z" + } + }, + "outputs": [], + "source": [ + "risk_limit = pd.Series([20, 25, 30, 35, 40], index=time)\n", + "\n", + "m.add_quadratic_constraints(x**2 + y**2, \"<=\", risk_limit, name=\"risk\")" + ] + }, + { + "cell_type": "markdown", + "id": "objective-header", + "metadata": {}, + "source": [ + "## Objective and Solve\n", + "\n", + "We maximize the sum over all time periods:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "objective", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:51.366506Z", + "start_time": "2025-12-03T14:04:51.276869Z" + } + }, + "outputs": [], + "source": [ + "m.add_objective((x + 2 * y).sum(), sense=\"max\")\n", + "m" + ] + }, + { + "cell_type": "markdown", + "id": "duals-header", + "metadata": {}, + "source": [ + "To retrieve dual values for quadratic constraints with Gurobi, set `QCPDual=1`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "solve", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:52.311740Z", + "start_time": "2025-12-03T14:04:51.994541Z" + } + }, + "outputs": [], + "source": [ + "m.solve(solver_name=\"gurobi\", QCPDual=1)" + ] + }, + { + "cell_type": "markdown", + "id": "results-header", + "metadata": {}, + "source": [ + "## Results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "results", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:57.706433Z", + "start_time": "2025-12-03T14:04:56.543039Z" + } + }, + "outputs": [], + "source": [ + "m.solution.to_dataframe().plot(kind=\"bar\", ylabel=\"Optimal Value\", rot=0);" + ] + }, + { + "cell_type": "markdown", + "id": "inspect-header", + "metadata": {}, + "source": [ + "## Inspecting Quadratic Constraints\n", + "\n", + "Quadratic constraints are stored in `m.quadratic_constraints`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "qc-container", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:58.261183Z", + "start_time": "2025-12-03T14:04:58.251877Z" + } + }, + "outputs": [], + "source": [ + "m.quadratic_constraints" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "qc-single", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:58.556810Z", + "start_time": "2025-12-03T14:04:58.521816Z" + } + }, + "outputs": [], + "source": [ + "m.quadratic_constraints[\"risk\"]" + ] + }, + { + "cell_type": "markdown", + "id": "qc-duals-header", + "metadata": {}, + "source": [ + "Dual values are available via the `.dual` property:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "qc-duals", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:58.728502Z", + "start_time": "2025-12-03T14:04:58.714225Z" + } + }, + "outputs": [], + "source": [ + "m.quadratic_constraints[\"risk\"].dual" + ] + }, + { + "cell_type": "markdown", + "id": "bilinear-header", + "metadata": {}, + "source": [ + "## Bilinear Terms\n", + "\n", + "Quadratic constraints can also include cross-product terms like $xy$. Note that bilinear constraints are nonconvex and require a solver that supports them (e.g., Gurobi):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bilinear", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:04:59.651299Z", + "start_time": "2025-12-03T14:04:59.297236Z" + } + }, + "outputs": [], + "source": [ + "m2 = linopy.Model()\n", + "\n", + "x = m2.add_variables(lower=0, name=\"x\")\n", + "y = m2.add_variables(lower=0, name=\"y\")\n", + "\n", + "# Bilinear constraint: xy <= 4\n", + "m2.add_quadratic_constraints(x * y, \"<=\", 4, name=\"bilinear\")\n", + "\n", + "m2.add_objective(x + y, sense=\"max\")\n", + "m2.add_constraints(x <= 5)\n", + "m2.add_constraints(y <= 5)\n", + "\n", + "m2.solve(solver_name=\"gurobi\")\n", + "\n", + "print(f\"x = {float(x.solution):.2f}, y = {float(y.solution):.2f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "w5c0l9z5efi", + "metadata": {}, + "source": "## Mixed Linear and Quadratic Terms\n\nQuadratic constraints can combine both quadratic and linear terms in the same constraint. For example, $x^2 + 2x + y \\leq 10$:" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "gtgvi4i0mzl", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:05:02.498083Z", + "start_time": "2025-12-03T14:05:02.254357Z" + } + }, + "outputs": [], + "source": [ + "m3 = linopy.Model()\n", + "\n", + "x = m3.add_variables(lower=0, name=\"x\")\n", + "y = m3.add_variables(lower=0, name=\"y\")\n", + "\n", + "# Mixed constraint: x² + 2x + y <= 10\n", + "m3.add_quadratic_constraints(x**2 + 2 * x + y, \"<=\", 10, name=\"mixed\")\n", + "\n", + "m3.add_objective(x + y, sense=\"max\")\n", + "\n", + "m3.solve(solver_name=\"gurobi\")\n", + "\n", + "print(f\"x = {float(x.solution):.2f}, y = {float(y.solution):.2f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "5gisroyuys4", + "metadata": {}, + "source": "## Equality Constraints\n\nQuadratic equality constraints use `\"==\"`. This example constrains a point to lie exactly on a circle:" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "x7tzlfma7fq", + "metadata": { + "ExecuteTime": { + "end_time": "2025-12-03T14:05:05.330521Z", + "start_time": "2025-12-03T14:05:05.202101Z" + } + }, + "outputs": [], + "source": [ + "m4 = linopy.Model()\n", + "\n", + "x = m4.add_variables(lower=-5, upper=5, name=\"x\")\n", + "y = m4.add_variables(lower=-5, upper=5, name=\"y\")\n", + "\n", + "# Point must lie on circle of radius 2\n", + "m4.add_quadratic_constraints(x**2 + y**2, \"==\", 4, name=\"circle\")\n", + "\n", + "# Maximize x + y (find point on circle furthest in direction (1,1))\n", + "m4.add_objective(x + y, sense=\"max\")\n", + "\n", + "m4.solve(solver_name=\"gurobi\")\n", + "\n", + "print(f\"x = {float(x.solution):.4f}, y = {float(y.solution):.4f}\")\n", + "print(f\"x² + y² = {float(x.solution) ** 2 + float(y.solution) ** 2:.4f}\")" + ] + }, + { + "cell_type": "markdown", + "id": "ru3y88un3an", + "metadata": {}, + "source": "## Convexity Considerations\n\nQuadratic constraints have important convexity properties that affect which solvers can handle them:\n\n- **Convex** (most solvers): $x^T Q x + a^T x \\leq b$ where $Q$ is positive semidefinite (e.g., sum of squares like $x^2 + y^2$)\n- **Nonconvex** (requires specialized solvers like Gurobi):\n - $x^T Q x + a^T x \\geq b$ (greater-than with positive semidefinite Q)\n - $x^T Q x + a^T x = b$ (equality constraints)\n - Bilinear terms like $xy$\n\nConvex quadratic constraints define a convex feasible region (like the interior of an ellipse), while nonconvex constraints can create disconnected or non-convex regions. Solvers like Gurobi use spatial branch-and-bound to handle nonconvex cases." + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/linopy/common.py b/linopy/common.py index 9d530cc2..f27c51b9 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -876,9 +876,22 @@ def print_line( coeff_string = f"{coeff_string[0]} {coeff_string[1:]}" if isinstance(var, list): - var_string = "" - for name, coords in var: - if name is not None: + # Quadratic term - check if it's a squared term (same variable twice) + var_parts = [(name, coords) for name, coords in var if name is not None] + if len(var_parts) == 2: + (name1, coords1), (name2, coords2) = var_parts + coord_string1 = print_coord(coords1) + coord_string2 = print_coord(coords2) + if name1 == name2 and coord_string1 == coord_string2: + # Squared term: x² instead of x x + var_string = f" {name1}{coord_string1}²" + else: + # Cross term: x·y + var_string = f" {name1}{coord_string1}·{name2}{coord_string2}" + else: + # Fallback for other cases + var_string = "" + for name, coords in var_parts: coord_string = print_coord(coords) var_string += f" {name}{coord_string}" else: @@ -927,6 +940,21 @@ def print_line( def print_single_constraint(model: Any, label: int) -> str: + """ + Print a single linear constraint by its label. + + Parameters + ---------- + model : linopy.Model + The model containing the constraint. + label : int + The label of the constraint to print. + + Returns + ------- + str + Formatted string representation of the constraint. + """ constraints = model.constraints name, coord = constraints.get_label_position(label) @@ -941,6 +969,56 @@ def print_single_constraint(model: Any, label: int) -> str: return f"{name}{print_coord(coord)}: {expr} {sign} {rhs:.12g}" +def print_single_quadratic_constraint(model: Any, label: int) -> str: + """ + Print a single quadratic constraint by its label. + + Parameters + ---------- + model : linopy.Model + The model containing the constraint. + label : int + The label of the constraint to print. + + Returns + ------- + str + Formatted string representation of the constraint. + """ + qconstraints = model.quadratic_constraints + name, coord = qconstraints.get_label_position(label) + + qcon = model.quadratic_constraints[name] + quad_coeffs = qcon.quad_coeffs.sel(coord).values # shape: (qterm, factor) + quad_vars = qcon.quad_vars.sel(coord).values # shape: (qterm, factor) + lin_coeffs = qcon.lin_coeffs.sel(coord).values # shape: (term,) + lin_vars = qcon.lin_vars.sel(coord).values # shape: (term,) + sign = qcon.sign.sel(coord).item() + rhs = qcon.rhs.sel(coord).item() + + # Combine quadratic and linear terms for print_single_expression + # Quadratic terms: coeffs shape (qterm,), vars shape (2, qterm) + # Linear terms: coeffs shape (term,), vars shape (term,) -> need to expand to (2, term) + n_lterms = len(lin_vars) + + # Stack coefficients + coeffs = np.concatenate([quad_coeffs[:, 0], lin_coeffs]) + + # Stack vars: quad_vars is (qterm, 2), need (2, qterm); lin_vars needs -1 padding + quad_vars_t = quad_vars.T # shape (2, qterm) + lin_vars_expanded = np.stack( + [lin_vars, np.full(n_lterms, -1)], axis=0 + ) # shape (2, term) + vars_combined = np.concatenate( + [quad_vars_t, lin_vars_expanded], axis=1 + ) # shape (2, total) + + expr = print_single_expression(coeffs, vars_combined, 0, model) + sign_pretty = SIGNS_pretty[sign] + + return f"{name}{print_coord(coord)}: {expr} {sign_pretty} {rhs:.12g}" + + def has_optimized_model(func: Callable[..., Any]) -> Callable[..., Any]: """ Check if a reference model is set. diff --git a/linopy/constants.py b/linopy/constants.py index 3f6886ec..a71137fa 100644 --- a/linopy/constants.py +++ b/linopy/constants.py @@ -39,12 +39,14 @@ GROUP_DIM = "_group" FACTOR_DIM = "_factor" CONCAT_DIM = "_concat" +QTERM_DIM = "_qterm" # Quadratic term dimension for quadratic constraints HELPER_DIMS: list[str] = [ TERM_DIM, STACKED_TERM_DIM, GROUPED_TERM_DIM, FACTOR_DIM, CONCAT_DIM, + QTERM_DIM, ] @@ -207,6 +209,7 @@ class Solution: primal: pd.Series = field(default_factory=_pd_series_float) dual: pd.Series = field(default_factory=_pd_series_float) objective: float = field(default=np.nan) + qc_dual: pd.Series = field(default_factory=_pd_series_float) @dataclass @@ -224,8 +227,11 @@ def __repr__(self) -> str: "not available" if self.solver_model is None else "available" ) if self.solution is not None: + n_qc_duals = len(self.solution.qc_dual) + qc_dual_str = f", {n_qc_duals} qc_duals" if n_qc_duals > 0 else "" solution_string = ( - f"Solution: {len(self.solution.primal)} primals, {len(self.solution.dual)} duals\n" + f"Solution: {len(self.solution.primal)} primals, " + f"{len(self.solution.dual)} duals{qc_dual_str}\n" f"Objective: {self.solution.objective:.2e}\n" ) else: diff --git a/linopy/constraints.py b/linopy/constraints.py index a329891a..b7018344 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -48,6 +48,7 @@ print_coord, print_single_constraint, print_single_expression, + print_single_quadratic_constraint, replace_by_map, save_join, to_dataframe, @@ -56,9 +57,11 @@ from linopy.config import options from linopy.constants import ( EQUAL, + FACTOR_DIM, GREATER_EQUAL, HELPER_DIMS, LESS_EQUAL, + QTERM_DIM, TERM_DIM, SIGNS_pretty, ) @@ -1123,3 +1126,789 @@ def rhs(self) -> int | float | np.floating | np.integer: def to_constraint(self) -> Constraint: data = self.lhs.to_linexpr().data.assign(sign=self.sign, rhs=self.rhs) return Constraint(data=data, model=self.lhs.model) + + +QFILL_VALUE = { + "labels": -1, + "rhs": np.nan, + "lin_coeffs": 0, + "lin_vars": -1, + "quad_coeffs": 0, + "quad_vars": -1, + "sign": "=", +} + + +def qconwrap( + method: Callable, *default_args: Any, **new_default_kwargs: Any +) -> Callable: + @functools.wraps(method) + def _qconwrap( + con: QuadraticConstraint, *args: Any, **kwargs: Any + ) -> QuadraticConstraint: + for k, v in new_default_kwargs.items(): + kwargs.setdefault(k, v) + return con.__class__( + method(con.data, *default_args, *args, **kwargs), con.model, con.name + ) + + _qconwrap.__doc__ = f"Wrapper for the xarray {method.__qualname__} function for linopy.QuadraticConstraint" + if new_default_kwargs: + _qconwrap.__doc__ += f" with default arguments: {new_default_kwargs}" + + return _qconwrap + + +class QuadraticConstraint: + """ + A quadratic constraint of the form: x'Qx + a'x <= b (or >=, =). + + The QuadraticConstraint class stores quadratic constraints with both + quadratic and linear terms. It follows the same design patterns as the + linear Constraint class. + + Dataset structure: + { + 'quad_coeffs': DataArray[float] # shape: (..., _factor, _qterm) + 'quad_vars': DataArray[int] # shape: (..., _factor, _qterm) + 'lin_coeffs': DataArray[float] # shape: (..., _term) + 'lin_vars': DataArray[int] # shape: (..., _term) + 'sign': DataArray[str] # '=', '<=', '>=' + 'rhs': DataArray[float] # Right-hand side constant + 'labels': DataArray[int] # Constraint labels (-1 if masked) + 'dual': DataArray[float] # [OPTIONAL] Dual values (convex only) + } + """ + + __slots__ = ("_data", "_model", "_assigned") + + _fill_value = QFILL_VALUE + + def __init__( + self, + data: Dataset, + model: Model, + name: str = "", + skip_broadcast: bool = False, + ) -> None: + """ + Initialize the QuadraticConstraint. + + Parameters + ---------- + data : xarray.Dataset + Dataset containing the constraint data. + model : linopy.Model + Underlying model. + name : str + Name of the constraint. + skip_broadcast : bool + Skip broadcasting of data arrays. + """ + from linopy.model import Model + + if not isinstance(data, Dataset): + raise ValueError(f"data must be a Dataset, got {type(data)}") + + if not isinstance(model, Model): + raise ValueError(f"model must be a Model, got {type(model)}") + + # Check required fields + for attr in ( + "quad_coeffs", + "quad_vars", + "lin_coeffs", + "lin_vars", + "sign", + "rhs", + ): + if attr not in data: + raise ValueError(f"missing '{attr}' in data") + + data = data.assign_attrs(name=name) + + if not skip_broadcast: + (data,) = xr.broadcast(data, exclude=[TERM_DIM, QTERM_DIM, FACTOR_DIM]) + + self._assigned = "labels" in data + self._data = data + self._model = model + + def __getitem__( + self, selector: str | int | slice | list | tuple | dict + ) -> QuadraticConstraint: + """ + Get selection from the constraint. + """ + data = Dataset({k: self.data[k][selector] for k in self.data}, attrs=self.attrs) + return self.__class__(data, self.model, self.name) + + @property + def attrs(self) -> dict[str, Any]: + """Get the attributes of the constraint.""" + return self.data.attrs + + @property + def coords(self) -> DatasetCoordinates: + """Get the coordinates of the constraint.""" + return self.data.coords + + @property + def indexes(self) -> Indexes: + """Get the indexes of the constraint.""" + return self.data.indexes + + @property + def dims(self) -> Frozen[Hashable, int]: + """Get the dimensions of the constraint.""" + return self.data.dims + + @property + def sizes(self) -> Frozen[Hashable, int]: + """Get the sizes of the constraint.""" + return self.data.sizes + + @property + def nterm(self) -> int: + """Get the number of linear terms in the constraint.""" + return self.data.sizes.get(TERM_DIM, 0) + + @property + def nqterm(self) -> int: + """Get the number of quadratic terms in the constraint.""" + return self.data.sizes.get(QTERM_DIM, 0) + + @property + def ndim(self) -> int: + """Get the number of dimensions of the constraint.""" + return self.rhs.ndim + + @property + def shape(self) -> tuple[int, ...]: + """Get the shape of the constraint.""" + return self.rhs.shape + + @property + def size(self) -> int: + """Get the size of the constraint.""" + return self.rhs.size + + @property + def loc(self) -> LocIndexer: + return LocIndexer(self) + + @property + def data(self) -> Dataset: + """Get the underlying Dataset.""" + return self._data + + @property + def labels(self) -> DataArray: + """Get the labels of the constraint.""" + return self.data.get("labels", DataArray([])) + + @property + def model(self) -> Model: + """Get the model of the constraint.""" + return self._model + + @property + def name(self) -> str: + """Return the name of the constraint.""" + return self.attrs["name"] + + @property + def coord_dims(self) -> tuple[Hashable, ...]: + return tuple(k for k in self.dims if k not in HELPER_DIMS) + + @property + def coord_sizes(self) -> dict[Hashable, int]: + return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS} + + @property + def coord_names(self) -> list[str]: + """Get the names of the coordinates.""" + return get_dims_with_index_levels(self.data, self.coord_dims) + + @property + def is_assigned(self) -> bool: + return self._assigned + + @property + def type(self) -> str: + """Get the type of the constraint.""" + return ( + "QuadraticConstraint" + if self.is_assigned + else "QuadraticConstraint (unassigned)" + ) + + @property + def range(self) -> tuple[int, int]: + """Return the range of the constraint.""" + return self.data.attrs["label_range"] + + @property + def mask(self) -> DataArray | None: + """ + Get the mask of the constraint. + + The mask indicates on which coordinates the constraint is enabled + (True) and disabled (False). + """ + if self.is_assigned: + return (self.data.labels != QFILL_VALUE["labels"]).astype(bool) + return None + + @property + def quad_coeffs(self) -> DataArray: + """Get the quadratic coefficients of the constraint.""" + return self.data.quad_coeffs + + @property + def quad_vars(self) -> DataArray: + """Get the quadratic variables of the constraint.""" + return self.data.quad_vars + + @property + def lin_coeffs(self) -> DataArray: + """Get the linear coefficients of the constraint.""" + return self.data.lin_coeffs + + @property + def lin_vars(self) -> DataArray: + """Get the linear variables of the constraint.""" + return self.data.lin_vars + + @property + def lhs(self) -> expressions.QuadraticExpression: + """ + Get the left-hand-side quadratic expression of the constraint. + """ + # Reconstruct the QuadraticExpression from quad and lin parts + # QuadraticExpression stores vars with _factor dimension + quad_data = Dataset( + { + "coeffs": self.quad_coeffs.rename({QTERM_DIM: TERM_DIM}), + "vars": self.quad_vars.rename({QTERM_DIM: TERM_DIM}), + "const": xr.zeros_like(self.rhs), + } + ) + return expressions.QuadraticExpression(quad_data, self.model) + + @property + def sign(self) -> DataArray: + """Get the signs of the constraint.""" + return self.data.sign + + @sign.setter + @is_constant + def sign(self, value: SignLike) -> None: + value = maybe_replace_signs(DataArray(value)).broadcast_like(self.sign) + self._data = assign_multiindex_safe(self.data, sign=value) + + @property + def rhs(self) -> DataArray: + """Get the right hand side constants of the constraint.""" + return self.data.rhs + + @rhs.setter + def rhs(self, value: ConstantLike) -> None: + value = DataArray(value).broadcast_like(self.rhs) + self._data = assign_multiindex_safe(self.data, rhs=value) + + @property + @has_optimized_model + def dual(self) -> DataArray: + """ + Get the dual values of the constraint. + + Note: Dual values are only available for convex quadratic constraints. + """ + if "dual" not in self.data: + raise AttributeError( + "Underlying is optimized but does not have dual values stored." + ) + return self.data["dual"] + + @dual.setter + def dual(self, value: ConstantLike) -> None: + """Set the dual values of the constraint.""" + value = DataArray(value).broadcast_like(self.labels) + self._data = assign_multiindex_safe(self.data, dual=value) + + def __repr__(self) -> str: + """Print the quadratic constraint arrays.""" + max_lines = options["display_max_rows"] + dims = list(self.coord_sizes.keys()) + ndim = len(dims) + dim_names = self.coord_names + dim_sizes = list(self.coord_sizes.values()) + size = np.prod(dim_sizes) if dim_sizes else 1 + masked_entries = (~self.mask).sum().values if self.mask is not None else 0 + lines = [] + + header_string = f"{self.type} `{self.name}`" if self.name else f"{self.type}" + + if size > 1 or ndim > 0: + for indices in generate_indices_for_printout(dim_sizes, max_lines): + if indices is None: + lines.append("\t\t...") + else: + coord = [ + self.data.indexes[dims[i]][int(ind)] + for i, ind in enumerate(indices) + ] + if self.mask is None or self.mask.values[indices]: + line = self._format_single_constraint(indices, coord) + else: + line = print_coord(coord) + ": None" + lines.append(line) + lines = align_lines_by_delimiter(lines, list(SIGNS_pretty.values())) + + shape_str = ", ".join(f"{d}: {s}" for d, s in zip(dim_names, dim_sizes)) + mask_str = f" - {masked_entries} masked entries" if masked_entries else "" + underscore = "-" * (len(shape_str) + len(mask_str) + len(header_string) + 4) + lines.insert(0, f"{header_string} [{shape_str}]{mask_str}:\n{underscore}") + elif size == 1: + line = self._format_single_constraint((), None) + lines.append(f"{header_string}\n{'-' * len(header_string)}\n{line}") + else: + lines.append(f"{header_string}\n{'-' * len(header_string)}\n") + + return "\n".join(lines) + + def _format_single_constraint(self, indices: tuple, coord: list | None) -> str: + """Format a single constraint for display.""" + # Format quadratic terms + quad_parts: list[str] = [] + if indices: + qcoeffs = self.quad_coeffs.values[indices] + qvars = self.quad_vars.values[indices] + else: + qcoeffs = self.quad_coeffs.values + qvars = self.quad_vars.values + + # qvars has shape (_qterm, _factor), qcoeffs has shape (_qterm, _factor) + if qvars.ndim >= 2 and qvars.shape[-1] == 2: + for t in range(qvars.shape[0]): + v1, v2 = qvars[t, 0], qvars[t, 1] + c = qcoeffs[t, 0] if qcoeffs.ndim >= 2 else qcoeffs[t] + if v1 != -1 and v2 != -1 and c != 0: + v1_name = self.model.variables.get_label_position(v1) + v2_name = self.model.variables.get_label_position(v2) + if v1_name[0] is not None and v2_name[0] is not None: + v1_str = f"{v1_name[0]}{print_coord(list(v1_name[1].values()))}" + v2_str = f"{v2_name[0]}{print_coord(list(v2_name[1].values()))}" + sign = "+" if c >= 0 and quad_parts else "" + if v1 == v2: + quad_parts.append(f"{sign}{c} {v1_str}²") + else: + quad_parts.append(f"{sign}{c} {v1_str}·{v2_str}") + + # Format linear terms + lin_parts: list[str] = [] + if indices: + lcoeffs = self.lin_coeffs.values[indices] + lvars = self.lin_vars.values[indices] + else: + lcoeffs = self.lin_coeffs.values + lvars = self.lin_vars.values + + for t in range(len(lvars)): + v, c = lvars[t], lcoeffs[t] + if v != -1 and c != 0: + v_name = self.model.variables.get_label_position(v) + if v_name[0] is not None: + v_str = f"{v_name[0]}{print_coord(list(v_name[1].values()))}" + sign = "+" if c >= 0 and (quad_parts or lin_parts) else "" + lin_parts.append(f"{sign}{c} {v_str}") + + expr_string = " ".join(quad_parts + lin_parts) or "0" + sign = SIGNS_pretty[self.sign.values[indices] if indices else self.sign.item()] + rhs = self.rhs.values[indices] if indices else self.rhs.item() + + if coord is not None: + return print_coord(coord) + f": {expr_string} {sign} {rhs}" + return f"{expr_string} {sign} {rhs}" + + def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None: + """ + Print the quadratic constraint. + + Parameters + ---------- + display_max_rows : int + Maximum number of rows to be displayed. + display_max_terms : int + Maximum number of terms to be displayed. + """ + with options as opts: + opts.set_value( + display_max_rows=display_max_rows, display_max_terms=display_max_terms + ) + print(self) + + def __contains__(self, value: Any) -> bool: + return self.data.__contains__(value) + + @property + def flat(self) -> pd.DataFrame: + """ + Convert the quadratic constraint to a pandas DataFrame. + + Returns a long format DataFrame with columns for both quadratic + and linear terms. + """ + ds = self.data + + # Process quadratic terms + quad_df_list = [] + if QTERM_DIM in ds.quad_vars.dims: + quad_vars = ds.quad_vars.assign_coords( + {FACTOR_DIM: ["vars1", "vars2"]} + ).to_dataset(FACTOR_DIM) + quad_ds = ds[["quad_coeffs", "labels"]].rename({"quad_coeffs": "coeffs"}) + # Take first factor's coefficients (they're the same) + if FACTOR_DIM in quad_ds.coeffs.dims: + quad_ds["coeffs"] = quad_ds.coeffs.isel({FACTOR_DIM: 0}) + quad_ds = quad_ds.assign(quad_vars) + + def quad_mask_func(data: pd.DataFrame) -> pd.Series: + mask = ((data["vars1"] != -1) | (data["vars2"] != -1)) & ( + data["coeffs"] != 0 + ) + if "labels" in data: + mask &= data["labels"] != -1 + return mask + + quad_df = to_dataframe(quad_ds, mask_func=quad_mask_func) + if not quad_df.empty: + quad_df["is_quadratic"] = True + quad_df_list.append(quad_df) + + # Process linear terms + lin_df_list = [] + if TERM_DIM in ds.lin_vars.dims: + lin_ds = ds[["lin_coeffs", "lin_vars", "labels"]].rename( + {"lin_coeffs": "coeffs", "lin_vars": "vars"} + ) + + def lin_mask_func(data: pd.DataFrame) -> pd.Series: + mask = (data["vars"] != -1) & (data["coeffs"] != 0) + if "labels" in data: + mask &= data["labels"] != -1 + return mask + + lin_df = to_dataframe(lin_ds, mask_func=lin_mask_func) + if not lin_df.empty: + lin_df["is_quadratic"] = False + lin_df["vars1"] = -1 + lin_df["vars2"] = -1 + lin_df_list.append(lin_df) + + # Combine + dfs = quad_df_list + lin_df_list + if not dfs: + return pd.DataFrame( + columns=["labels", "coeffs", "vars", "vars1", "vars2", "is_quadratic"] + ) + + df = pd.concat(dfs, ignore_index=True) + return df + + def to_polars(self) -> pl.DataFrame: + """ + Convert the quadratic constraint to a polars DataFrame. + """ + df = self.flat + return pl.from_pandas(df) + + # Wrapped xarray functions + assign = qconwrap(Dataset.assign) + assign_multiindex_safe = qconwrap(assign_multiindex_safe) + assign_attrs = qconwrap(Dataset.assign_attrs) + assign_coords = qconwrap(Dataset.assign_coords) + broadcast_like = qconwrap(Dataset.broadcast_like) + chunk = qconwrap(Dataset.chunk) + drop_sel = qconwrap(Dataset.drop_sel) + drop_isel = qconwrap(Dataset.drop_isel) + expand_dims = qconwrap(Dataset.expand_dims) + sel = qconwrap(Dataset.sel) + isel = qconwrap(Dataset.isel) + shift = qconwrap(Dataset.shift) + swap_dims = qconwrap(Dataset.swap_dims) + set_index = qconwrap(Dataset.set_index) + reindex = qconwrap(Dataset.reindex, fill_value=_fill_value) + reindex_like = qconwrap(Dataset.reindex_like, fill_value=_fill_value) + rename = qconwrap(Dataset.rename) + rename_dims = qconwrap(Dataset.rename_dims) + roll = qconwrap(Dataset.roll) + stack = qconwrap(Dataset.stack) + unstack = qconwrap(Dataset.unstack) + iterate_slices = iterate_slices + + +@dataclass(repr=False) +class QuadraticConstraints: + """ + A container for storing multiple quadratic constraint arrays. + """ + + data: dict[str, QuadraticConstraint] + model: Model + + dataset_attrs = [ + "labels", + "quad_coeffs", + "quad_vars", + "lin_coeffs", + "lin_vars", + "sign", + "rhs", + ] + dataset_names = [ + "Labels", + "Quadratic coefficients", + "Quadratic variables", + "Linear coefficients", + "Linear variables", + "Signs", + "Right-hand-side constants", + ] + + def _formatted_names(self) -> dict[str, str]: + """Get a dictionary of formatted names to proper constraint names.""" + return {format_string_as_variable_name(n): n for n in self} + + def __repr__(self) -> str: + """Return a string representation of the quadratic constraints.""" + r = "linopy.model.QuadraticConstraints" + line = "-" * len(r) + r += f"\n{line}\n" + + for name, ds in self.items(): + coords = ( + " (" + ", ".join([str(c) for c in ds.coords.keys()]) + ")" + if ds.coords + else "" + ) + r += f" * {name}{coords}\n" + if not len(list(self)): + r += "\n" + return r + + @overload + def __getitem__(self, names: str) -> QuadraticConstraint: ... + + @overload + def __getitem__(self, names: list[str]) -> QuadraticConstraints: ... + + def __getitem__( + self, names: str | list[str] + ) -> QuadraticConstraint | QuadraticConstraints: + if isinstance(names, str): + return self.data[names] + return QuadraticConstraints( + {name: self.data[name] for name in names}, self.model + ) + + def __getattr__(self, name: str) -> QuadraticConstraint: + if name in self.data: + return self.data[name] + else: + if name in (formatted_names := self._formatted_names()): + return self.data[formatted_names[name]] + raise AttributeError( + f"QuadraticConstraints has no attribute `{name}` or the attribute is not accessible." + ) + + def __getstate__(self) -> dict: + return self.__dict__ + + def __setstate__(self, d: dict) -> None: + self.__dict__.update(d) + + def __dir__(self) -> list[str]: + base_attributes = list(super().__dir__()) + formatted_names = [ + n for n in self._formatted_names() if n not in base_attributes + ] + return base_attributes + formatted_names + + def __len__(self) -> int: + return self.data.__len__() + + def __iter__(self) -> Iterator[str]: + return self.data.__iter__() + + def items(self) -> ItemsView[str, QuadraticConstraint]: + return self.data.items() + + def _ipython_key_completions_(self) -> list[str]: + """Provide method for key-autocompletions in IPython.""" + return list(self) + + def add(self, constraint: QuadraticConstraint) -> None: + """Add a quadratic constraint to the container.""" + self.data[constraint.name] = constraint + + def remove(self, name: str) -> None: + """Remove quadratic constraint `name` from the container.""" + self.data.pop(name) + + @property + def labels(self) -> Dataset: + """Get the labels of all quadratic constraints.""" + return save_join( + *[v.labels.rename(k) for k, v in self.items()], + integer_dtype=True, + ) + + @property + def ncons(self) -> int: + """Get the number of quadratic constraints effectively used by the model.""" + if not len(self): + return 0 + return len(self.flat.labels.unique()) + + @property + def flat(self) -> pd.DataFrame: + """Convert all quadratic constraints to a single pandas DataFrame.""" + dfs = [self[k].flat for k in self] + if not len(dfs): + return pd.DataFrame( + columns=[ + "coeffs", + "vars", + "vars1", + "vars2", + "labels", + "key", + "is_quadratic", + ] + ) + df = pd.concat(dfs, ignore_index=True) + unique_labels = df.labels.unique() + map_labels = pd.Series(np.arange(len(unique_labels)), index=unique_labels) + df["key"] = df.labels.map(map_labels) + return df + + @property + def sign(self) -> Dataset: + """Get the signs of all quadratic constraints.""" + return save_join(*[v.sign.rename(k) for k, v in self.items()]) + + @property + def rhs(self) -> Dataset: + """Get the right-hand-side constants of all quadratic constraints.""" + return save_join(*[v.rhs.rename(k) for k, v in self.items()]) + + @property + def dual(self) -> Dataset: + """Get the dual values of all quadratic constraints.""" + try: + return save_join(*[v.dual.rename(k) for k, v in self.items()]) + except AttributeError: + return Dataset() + + def get_name_by_label(self, label: int | float) -> str: + """Get the constraint name containing the passed label.""" + if not isinstance(label, float | int) or label < 0: + raise ValueError("Label must be a positive number.") + for name, ds in self.items(): + if label in ds.labels: + return name + raise ValueError(f"No quadratic constraint found containing the label {label}.") + + def get_label_position( + self, values: int | ndarray + ) -> ( + tuple[str, dict] + | tuple[None, None] + | list[tuple[str, dict] | tuple[None, None]] + | list[list[tuple[str, dict] | tuple[None, None]]] + ): + """Get tuple of name and coordinate for constraint labels.""" + return get_label_position(self, values) + + def reset_dual(self) -> None: + """Reset the stored dual values of quadratic constraints.""" + for k, c in self.items(): + if "dual" in c: + c._data = c.data.drop_vars("dual") + + @property + def inequalities(self) -> QuadraticConstraints: + """ + Get the subset of quadratic constraints which are purely inequalities. + """ + return self[[n for n, s in self.items() if (s.sign != EQUAL).all()]] + + @property + def equalities(self) -> QuadraticConstraints: + """ + Get the subset of quadratic constraints which are purely equalities. + """ + return self[[n for n, s in self.items() if (s.sign == EQUAL).all()]] + + def sanitize_zeros(self) -> None: + """ + Filter out terms with zero and close-to-zero coefficient. + + For quadratic constraints, this filters both quadratic and linear terms. + """ + for name in self: + con = self[name] + # Filter linear terms + not_zero_lin = abs(con.lin_coeffs) > 1e-10 + con._data = assign_multiindex_safe( + con.data, + lin_vars=con.lin_vars.where(not_zero_lin, -1), + lin_coeffs=con.lin_coeffs.where(not_zero_lin), + ) + # Filter quadratic terms + not_zero_quad = abs(con.quad_coeffs) > 1e-10 + con._data = assign_multiindex_safe( + con.data, + quad_vars=con.quad_vars.where(not_zero_quad, -1), + quad_coeffs=con.quad_coeffs.where(not_zero_quad), + ) + + def sanitize_missings(self) -> None: + """ + Set constraint labels to -1 where all variables in the lhs are missing. + + For quadratic constraints, checks both quadratic and linear terms. + """ + for name in self: + con = self[name] + # Check if any quadratic term has valid variables + contains_non_missing_quad = (con.quad_vars != -1).any( + [QTERM_DIM, FACTOR_DIM] + ) + # Check if any linear term has valid variables + contains_non_missing_lin = (con.lin_vars != -1).any(TERM_DIM) + # Constraint is valid if it has either quadratic or linear terms + contains_non_missing = contains_non_missing_quad | contains_non_missing_lin + labels = con.labels.where(contains_non_missing, -1) + con._data = assign_multiindex_safe(con.data, labels=labels) + + def print_labels( + self, values: Sequence[int], display_max_terms: int | None = None + ) -> None: + """ + Print a selection of labels of the quadratic constraints. + + Parameters + ---------- + values : list, array-like + One dimensional array of constraint labels. + display_max_terms : int, optional + Maximum number of terms to display per constraint. + """ + with options as opts: + if display_max_terms is not None: + opts.set_value(display_max_terms=display_max_terms) + res = [print_single_quadratic_constraint(self.model, v) for v in values] + print("\n".join(res)) diff --git a/linopy/expressions.py b/linopy/expressions.py index d60c8be5..28bc166a 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -73,6 +73,7 @@ GROUPED_TERM_DIM, HELPER_DIMS, LESS_EQUAL, + QTERM_DIM, STACKED_TERM_DIM, TERM_DIM, ) @@ -86,7 +87,11 @@ ) if TYPE_CHECKING: - from linopy.constraints import AnonymousScalarConstraint, Constraint + from linopy.constraints import ( + AnonymousScalarConstraint, + Constraint, + QuadraticConstraint, + ) from linopy.model import Model from linopy.variables import ScalarVariable, Variable @@ -791,55 +796,6 @@ def sum( return res - def cumsum( - self, - dim: DimsLike | None = None, - *, - skipna: bool | None = None, - keep_attrs: bool | None = None, - **kwargs: Any, - ) -> LinearExpression: - """ - Cumulated sum along a given axis. - - Docstring and arguments are borrowed from `xarray.Dataset.cumsum` - - Parameters - ---------- - dim : str, Iterable of Hashable, "..." or None, default: None - Name of dimension[s] along which to apply ``cumsum``. For e.g. ``dim="x"`` - or ``dim=["x", "y"]``. If "..." or None, will reduce over all dimensions. - skipna : bool or None, optional - If True, skip missing values (as marked by NaN). By default, only - skips missing values for float dtypes; other dtypes either do not - have a sentinel missing value (int) or ``skipna=True`` has not been - implemented (object, datetime64 or timedelta64). - keep_attrs : bool or None, optional - If True, ``attrs`` will be copied from the original - object to the new one. If False, the new object will be - returned without attributes. - **kwargs : Any - Additional keyword arguments passed on to the appropriate array - function for calculating ``cumsum`` on this object's data. - These could include dask-specific kwargs like ``split_every``. - - Returns - ------- - linopy.expression.LinearExpression - """ - # Along every dimensions, we want to perform cumsum along, get the size of the - # dimension to pass that to self.rolling. - if not dim: - # If user did not specify a dimension to sum over, use all relevant - # dimensions - dim = self.coord_dims - if isinstance(dim, str): - dim = [dim] - elif isinstance(dim, EllipsisType) or dim is None: - dim = self.coord_dims - dim_dict = {dim_name: self.data.sizes[dim_name] for dim_name in dim} - return self.rolling(dim=dim_dict).sum(keep_attrs=keep_attrs, skipna=skipna) - def to_constraint( self, sign: SignLike, rhs: ConstantLike | VariableLike | ExpressionLike ) -> Constraint: @@ -986,74 +942,6 @@ def diff(self: GenericExpression, dim: str, n: int = 1) -> GenericExpression: """ return self - self.shift({dim: n}) - def groupby( - self, - group: DataFrame | Series | DataArray, - restore_coord_dims: bool | None = None, - **kwargs: Any, - ) -> LinearExpressionGroupby: - """ - Returns a LinearExpressionGroupBy object for performing grouped - operations. - - Docstring and arguments are borrowed from `xarray.Dataset.groupby` - - Parameters - ---------- - group : str, DataArray or IndexVariable - Array whose unique values should be used to group this array. If a - string, must be the name of a variable contained in this dataset. - restore_coord_dims : bool, optional - If True, also restore the dimension order of multi-dimensional - coordinates. - - Returns - ------- - grouped - A `LinearExpressionGroupBy` containing the xarray groups and ensuring - the correct return type. - """ - ds = self.data - kwargs = dict(restore_coord_dims=restore_coord_dims, **kwargs) - return LinearExpressionGroupby(ds, group, model=self.model, kwargs=kwargs) - - def rolling( - self, - dim: Mapping[Any, int] | None = None, - min_periods: int | None = None, - center: bool | Mapping[Any, bool] = False, - **window_kwargs: int, - ) -> LinearExpressionRolling: - """ - Rolling window object. - - Docstring and arguments are borrowed from `xarray.Dataset.rolling` - - Parameters - ---------- - dim : dict, optional - Mapping from the dimension name to create the rolling iterator - along (e.g. `time`) to its moving window size. - min_periods : int, default: None - Minimum number of observations in window required to have a value - (otherwise result is NA). The default, None, is equivalent to - setting min_periods equal to the size of the window. - center : bool or mapping, default: False - Set the labels at the center of the window. - **window_kwargs : optional - The keyword arguments form of ``dim``. - One of dim or window_kwargs must be provided. - - Returns - ------- - linopy.expression.LinearExpressionRolling - """ - ds = self.data - rolling = ds.rolling( - dim=dim, min_periods=min_periods, center=center, **window_kwargs - ) - return LinearExpressionRolling(rolling, model=self.model) - @property def nterm(self) -> int: """ @@ -1064,10 +952,9 @@ def nterm(self) -> int: @property def shape(self) -> tuple[int, ...]: """ - Get the total shape of the linear expression. + Get the shape of the expression (excluding term dimension). """ - assert self.vars.shape == self.coeffs.shape - return self.vars.shape + return tuple(v for k, v in self.sizes.items() if k != TERM_DIM) @property def size(self) -> int: @@ -1647,6 +1534,119 @@ def process_one( return merge(exprs, cls=cls) if len(exprs) > 1 else exprs[0] + def cumsum( + self, + dim: DimsLike | None = None, + *, + skipna: bool | None = None, + keep_attrs: bool | None = None, + **kwargs: Any, + ) -> LinearExpression: + """ + Cumulated sum along a given axis. + + Docstring and arguments are borrowed from `xarray.Dataset.cumsum` + + Parameters + ---------- + dim : str, Iterable of Hashable, "..." or None, default: None + Name of dimension[s] along which to apply ``cumsum``. For e.g. ``dim="x"`` + or ``dim=["x", "y"]``. If "..." or None, will reduce over all dimensions. + skipna : bool or None, optional + If True, skip missing values (as marked by NaN). By default, only + skips missing values for float dtypes; other dtypes either do not + have a sentinel missing value (int) or ``skipna=True`` has not been + implemented (object, datetime64 or timedelta64). + keep_attrs : bool or None, optional + If True, ``attrs`` will be copied from the original + object to the new one. If False, the new object will be + returned without attributes. + **kwargs : Any + Additional keyword arguments passed on to the appropriate array + function for calculating ``cumsum`` on this object's data. + These could include dask-specific kwargs like ``split_every``. + + Returns + ------- + linopy.expression.LinearExpression + """ + if not dim: + dim = self.coord_dims + if isinstance(dim, str): + dim = [dim] + elif isinstance(dim, EllipsisType) or dim is None: + dim = self.coord_dims + dim_dict = {dim_name: self.data.sizes[dim_name] for dim_name in dim} + return self.rolling(dim=dim_dict).sum(keep_attrs=keep_attrs, skipna=skipna) + + def groupby( + self, + group: DataFrame | Series | DataArray, + restore_coord_dims: bool | None = None, + **kwargs: Any, + ) -> LinearExpressionGroupby: + """ + Returns a LinearExpressionGroupBy object for performing grouped + operations. + + Docstring and arguments are borrowed from `xarray.Dataset.groupby` + + Parameters + ---------- + group : str, DataArray or IndexVariable + Array whose unique values should be used to group this array. If a + string, must be the name of a variable contained in this dataset. + restore_coord_dims : bool, optional + If True, also restore the dimension order of multi-dimensional + coordinates. + + Returns + ------- + grouped + A `LinearExpressionGroupBy` containing the xarray groups and ensuring + the correct return type. + """ + ds = self.data + kwargs = dict(restore_coord_dims=restore_coord_dims, **kwargs) + return LinearExpressionGroupby(ds, group, model=self.model, kwargs=kwargs) + + def rolling( + self, + dim: Mapping[Any, int] | None = None, + min_periods: int | None = None, + center: bool | Mapping[Any, bool] = False, + **window_kwargs: int, + ) -> LinearExpressionRolling: + """ + Rolling window object. + + Docstring and arguments are borrowed from `xarray.Dataset.rolling` + + Parameters + ---------- + dim : dict, optional + Mapping from the dimension name to create the rolling iterator + along (e.g. `time`) to its moving window size. + min_periods : int, default: None + Minimum number of observations in window required to have a value + (otherwise result is NA). The default, None, is equivalent to + setting min_periods equal to the size of the window. + center : bool or mapping, default: False + Set the labels at the center of the window. + **window_kwargs : optional + The keyword arguments form of ``dim``. + One of dim or window_kwargs must be provided. + + Returns + ------- + linopy.expression.LinearExpressionRolling + """ + ds = self.data + rolling = ds.rolling( + dim=dim, min_periods=min_periods, center=center, **window_kwargs + ) + return LinearExpressionRolling(rolling, model=self.model) + class QuadraticExpression(BaseExpression): """ @@ -1684,9 +1684,11 @@ def type(self) -> str: @property def shape(self) -> tuple[int, ...]: - # TODO Implement this - raise NotImplementedError( - f"{self.__class__.__name__} does not support shape property" + """ + Get the shape of the expression (excluding term dimensions). + """ + return tuple( + v for k, v in self.sizes.items() if k not in (TERM_DIM, FACTOR_DIM) ) def __mul__(self, other: SideLike) -> QuadraticExpression: @@ -1802,11 +1804,84 @@ def solution(self) -> DataArray: sol = (self.coeffs * vals.prod(FACTOR_DIM)).sum(TERM_DIM) + self.const return sol.rename("solution") - def to_constraint(self, sign: SignLike, rhs: SideLike) -> NotImplementedType: - raise NotImplementedError( - "Quadratic expressions cannot be used in constraints." + def to_constraint(self, sign: SignLike, rhs: ConstantLike) -> QuadraticConstraint: # type: ignore[override] + """ + Convert a quadratic expression to a quadratic constraint. + + Parameters + ---------- + sign : str + Constraint sense: '<=', '>=', or '=' + rhs : float or array-like + Right-hand side constant + + Returns + ------- + QuadraticConstraint + + Examples + -------- + >>> from linopy import Model + >>> m = Model() + >>> x = m.add_variables(name="x") + >>> y = m.add_variables(name="y") + >>> qc = (x**2 + y**2).to_constraint("<=", 25) # x² + y² <= 25 + """ + # Move rhs to left-hand side to get: quadexpr - rhs {sign} 0 + all_to_lhs = self - rhs + + # Separate quadratic and linear terms + # QuadraticExpression has vars with shape (..., _factor, _term) + # where _factor has size 2 (vars1, vars2) + vars_data = all_to_lhs.vars + coeffs_data = all_to_lhs.coeffs + + # Identify linear terms: where one of the factors is -1 (missing) + is_linear = (vars_data.isel({FACTOR_DIM: 0}) == -1) | ( + vars_data.isel({FACTOR_DIM: 1}) == -1 ) + # Extract quadratic terms (neither var is -1) + is_quadratic = ~is_linear + + # Get quadratic parts + # Note: coeffs only has _term dim, vars has both _factor and _term + quad_vars = vars_data.where(is_quadratic, -1) + quad_coeffs = coeffs_data.where(is_quadratic, 0) + + # Expand quad_coeffs to have _factor dimension (for consistency with quad_vars) + quad_coeffs = quad_coeffs.expand_dims({FACTOR_DIM: 2}, axis=-2) + + # Rename TERM_DIM to QTERM_DIM for quadratic terms + quad_vars = quad_vars.rename({TERM_DIM: QTERM_DIM}) + quad_coeffs = quad_coeffs.rename({TERM_DIM: QTERM_DIM}) + + # Get linear parts - extract the non-missing variable + lin_vars_factor0 = vars_data.isel({FACTOR_DIM: 0}) + lin_vars_factor1 = vars_data.isel({FACTOR_DIM: 1}) + # For linear terms, one factor is -1, use the other + lin_vars = xr.where( + is_linear & (lin_vars_factor0 != -1), + lin_vars_factor0, + xr.where(is_linear & (lin_vars_factor1 != -1), lin_vars_factor1, -1), + ) + # Get linear coefficients (coeffs doesn't have _factor dim) + lin_coeffs = coeffs_data.where(is_linear, 0) + + # Build the constraint data + data = Dataset( + { + "quad_coeffs": quad_coeffs, + "quad_vars": quad_vars, + "lin_coeffs": lin_coeffs, + "lin_vars": lin_vars, + "sign": sign, + "rhs": -all_to_lhs.const, # Move constant to RHS + } + ) + + return constraints.QuadraticConstraint(data, model=self.model) + @property def flat(self) -> DataFrame: """ @@ -1878,6 +1953,193 @@ def to_matrix(self) -> scipy.sparse._csc.csc_matrix: nvars = self.model.shape[1] return csc_matrix((data, (row, col)), shape=(nvars, nvars)) + def groupby( + self, + group: Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame, + restore_coord_dims: bool = True, + **kwargs: Any, + ) -> QuadraticExpressionGroupby: + """ + GroupBy operation for QuadraticExpression. + + Returns a QuadraticExpressionGroupby object for performing grouped + operations. + """ + ds = self.data + kwargs_dict = dict(restore_coord_dims=restore_coord_dims, **kwargs) + return QuadraticExpressionGroupby( + ds, group, model=self.model, kwargs=kwargs_dict + ) + + def rolling( + self, + dim: Mapping[Any, int] | None = None, + min_periods: int | None = None, + center: bool | Mapping[Any, bool] = False, + **window_kwargs: int, + ) -> QuadraticExpressionRolling: + """ + Rolling window object for QuadraticExpression. + + Returns a QuadraticExpressionRolling object for performing rolling + operations. + """ + ds = self.data + rolling = ds.rolling( + dim=dim, min_periods=min_periods, center=center, **window_kwargs + ) + return QuadraticExpressionRolling(rolling, model=self.model) + + def cumsum( + self, + dim: DimsLike | None = None, + *, + skipna: bool | None = None, + keep_attrs: bool | None = None, + **kwargs: Any, + ) -> QuadraticExpression: + """ + Cumulated sum along a given axis. + + Returns + ------- + QuadraticExpression + """ + if not dim: + dim = self.coord_dims + if isinstance(dim, str): + dim = [dim] + elif isinstance(dim, EllipsisType) or dim is None: + dim = self.coord_dims + dim_dict = {dim_name: self.data.sizes[dim_name] for dim_name in dim} + return self.rolling(dim=dim_dict).sum(keep_attrs=keep_attrs, skipna=skipna) + + +@dataclass +@forward_as_properties(groupby=["dims", "groups"]) +class QuadraticExpressionGroupby: + """ + GroupBy object specialized to grouping QuadraticExpression objects. + """ + + data: xr.Dataset + group: Hashable | DataArray | IndexVariable | pd.Series | pd.DataFrame + model: Any + kwargs: Mapping[str, Any] = field(default_factory=dict) + + @property + def groupby(self) -> xarray.core.groupby.DatasetGroupBy: + """ + Groups the data using the specified group and kwargs. + """ + if isinstance(self.group, pd.DataFrame): + raise ValueError( + "Grouping by a DataFrame only supported for `sum` operation with `use_fallback=False`." + ) + if isinstance(self.group, pd.Series): + group_name = self.group.name or "group" + group = DataArray(self.group, name=group_name) + else: + group = self.group # type: ignore + + return self.data.groupby(group=group, **self.kwargs) + + def map( + self, + func: Callable[..., Dataset], + shortcut: bool = False, + args: tuple[()] = (), + **kwargs: Any, + ) -> QuadraticExpression: + """ + Apply a specified function to the groupby object. + """ + return QuadraticExpression( + self.groupby.map(func, shortcut=shortcut, args=args, **kwargs), self.model + ) + + def sum(self, use_fallback: bool = False, **kwargs: Any) -> QuadraticExpression: + """ + Sum the groupby object. + """ + non_fallback_types = (pd.Series, pd.DataFrame, xr.DataArray) + if isinstance(self.group, non_fallback_types) and not use_fallback: + group: pd.Series | pd.DataFrame | xr.DataArray = self.group + if isinstance(group, pd.DataFrame): + final_group_name = "group" + else: + final_group_name = getattr(group, "name", "group") or "group" + + if isinstance(group, DataArray): + group = group.to_pandas() + + int_map = None + if isinstance(group, pd.DataFrame): + index_name = group.index.name + group = group.reindex(self.data.indexes[group.index.name]) + group.index.name = index_name + int_map = get_index_map(*group.values.T) + orig_group = group + group = group.apply(tuple, axis=1).map(int_map) + + assert isinstance(group, pd.Series) + group_dim = group.index.name + + arrays = [group, group.groupby(group).cumcount()] + idx = pd.MultiIndex.from_arrays(arrays, names=[GROUP_DIM, GROUPED_TERM_DIM]) + new_coords = Coordinates.from_pandas_multiindex(idx, group_dim) + coords = self.data.indexes[group_dim] + names_to_drop = [coords.name] + if isinstance(coords, pd.MultiIndex): + names_to_drop += list(coords.names) + ds = self.data.drop_vars(names_to_drop).assign_coords(new_coords) + ds = ds.unstack(group_dim, fill_value=QuadraticExpression._fill_value) + ds = LinearExpression._sum(ds, dim=GROUPED_TERM_DIM) + + if int_map is not None: + index = ds.indexes[GROUP_DIM].map({v: k for k, v in int_map.items()}) + index.names = [str(col) for col in orig_group.columns] + index.name = GROUP_DIM + new_coords = Coordinates.from_pandas_multiindex(index, GROUP_DIM) + ds = xr.Dataset(ds.assign_coords(new_coords)) + + ds = ds.rename({GROUP_DIM: final_group_name}) + return QuadraticExpression(ds, self.model) + + def func(ds: Dataset) -> Dataset: + ds = LinearExpression._sum(ds, str(self.groupby._group_dim)) + ds = ds.assign_coords({TERM_DIM: np.arange(len(ds._term))}) + return ds + + return self.map(func, **kwargs, shortcut=True) + + def roll(self, **kwargs: Any) -> QuadraticExpression: + """ + Roll the groupby object. + """ + return self.map(Dataset.roll, **kwargs) + + +@dataclass +@forward_as_properties(rolling=["center", "dim", "obj", "rollings", "window"]) +class QuadraticExpressionRolling: + """ + Rolling object specialized to rolling QuadraticExpression objects. + """ + + rolling: DatasetRolling + model: Any + + def sum(self, **kwargs: Any) -> QuadraticExpression: + data = self.rolling.construct("_rolling_term", keep_attrs=True) + ds = ( + data[["coeffs", "vars"]] + .rename({TERM_DIM: STACKED_TERM_DIM}) + .stack({TERM_DIM: [STACKED_TERM_DIM, "_rolling_term"]}, create_index=False) + ) + ds = assign_multiindex_safe(ds, const=data.const.sum("_rolling_term")) + return QuadraticExpression(ds, self.model) + def as_expression( obj: Any, model: Model, **kwargs: Any diff --git a/linopy/io.py b/linopy/io.py index 7065adbb..74969aee 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -8,7 +8,7 @@ import logging import shutil import time -from collections.abc import Callable +from collections.abc import Callable, Iterable from io import BufferedWriter from pathlib import Path from tempfile import TemporaryDirectory @@ -51,11 +51,10 @@ def clean_name(name: str) -> str: coord_sanitizer = str.maketrans("[,]", "(,)", " ") -def print_coord(coord: str) -> str: - from linopy.common import print_coord +def print_coord(coord: dict[str, Any] | Iterable[Any]) -> str: + from linopy.common import print_coord as _print_coord - coord = print_coord(coord).translate(coord_sanitizer) - return coord + return _print_coord(coord).translate(coord_sanitizer) def get_printers_scalar( @@ -177,10 +176,11 @@ def objective_to_file( f.write(f"{sense}\n\nobj:\n\n".encode()) df = m.objective.to_polars() - if m.is_linear: + if m.objective.is_linear: objective_write_linear_terms(f, df, print_variable) - elif m.is_quadratic: + else: + # Quadratic objective linear_terms = df.filter(pl.col("vars1").eq(-1) | pl.col("vars2").eq(-1)) linear_terms = linear_terms.with_columns( pl.when(pl.col("vars1").eq(-1)) @@ -421,6 +421,127 @@ def constraints_to_file( # formatted.sink_csv(f, **kwargs) +def quadratic_constraints_to_file( + m: Model, + f: BufferedWriter, + progress: bool = False, + explicit_coordinate_names: bool = False, + write_section_header: bool = True, +) -> None: + """ + Write out quadratic constraints of a model to an LP file. + + LP format for quadratic constraints (Gurobi/CPLEX style): + qc0: 3.1 x + 4.5 y + [ x ^ 2 + 2 x * y + 3 y ^ 2 ] <= 10 + """ + if not len(m.quadratic_constraints): + return + + # Write "s.t." section header if needed (when there are no linear constraints) + if write_section_header and not len(m.constraints): + f.write(b"\n\ns.t.\n\n") + + print_variable, _ = get_printers_scalar( + m, explicit_coordinate_names=explicit_coordinate_names + ) + + if progress: + logger.info("Writing quadratic constraints.") + + names = list(m.quadratic_constraints) + if progress: + names = tqdm( + names, + desc="Writing quadratic constraints.", + colour=TQDM_COLOR, + ) + + for name in names: + qcon = m.quadratic_constraints[name] + df = qcon.to_polars() + + if df.height == 0: + continue + + # Build label -> (sign, rhs) lookup for multi-dimensional constraints + labels_flat = qcon.labels.values.ravel() + signs_flat = np.broadcast_to(qcon.sign.values, qcon.labels.shape).ravel() + rhs_flat = np.broadcast_to(qcon.rhs.values, qcon.labels.shape).ravel() + label_metadata = { + int(lab): (str(signs_flat[i]), float(rhs_flat[i])) + for i, lab in enumerate(labels_flat) + if lab != -1 + } + + # Get unique labels (constraint indices) + labels = df["labels"].unique().to_list() + + for label in labels: + label_df = df.filter(pl.col("labels") == label) + label_int = int(label) + if label_int not in label_metadata: + msg = ( + f"Label {label_int} from flat representation not found in " + f"constraint metadata for '{name}'. This indicates a mismatch " + "between the flat DataFrame and constraint labels." + ) + raise ValueError(msg) + sign, rhs = label_metadata[label_int] + + # Start constraint line with label + constraint_name = clean_name(name) + if explicit_coordinate_names: + label_pos = m.quadratic_constraints.get_label_position(label) + if label_pos and label_pos[0] is not None: + con_name, coord = label_pos + assert isinstance(con_name, str) + assert isinstance(coord, dict) + coord_str = print_coord(coord) + constraint_name = f"{clean_name(con_name)}{coord_str}#{label}" + else: + constraint_name = f"{constraint_name}#{label}" + else: + constraint_name = f"qc{label}" + + f.write(f"{constraint_name}:\n".encode()) + + # Write linear terms first + linear_terms = label_df.filter(~pl.col("is_quadratic")) + for row in linear_terms.iter_rows(named=True): + coeff = row["coeffs"] + var_idx = int(row["vars"]) + if var_idx >= 0: + var_name = print_variable(var_idx) + sign_char = "+" if coeff >= 0 else "" + f.write(f"{sign_char}{coeff} {var_name}\n".encode()) + + # Write quadratic terms in brackets + quad_terms = label_df.filter(pl.col("is_quadratic")) + if quad_terms.height > 0: + f.write(b"+ [\n") + for i, row in enumerate(quad_terms.iter_rows(named=True)): + coeff = row["coeffs"] + var1 = int(row["vars1"]) + var2 = int(row["vars2"]) + var1_name = print_variable(var1) + var2_name = print_variable(var2) + + sign_char = "+" if coeff >= 0 or i == 0 else "" + if var1 == var2: + # Squared term: x ^ 2 + f.write(f"{sign_char}{coeff} {var1_name} ^ 2\n".encode()) + else: + # Cross term: x * y + f.write( + f"{sign_char}{coeff} {var1_name} * {var2_name}\n".encode() + ) + f.write(b"]\n") + + # Write comparison and RHS + rhs_sign = "+" if rhs >= 0 else "" + f.write(f"{sign} {rhs_sign}{rhs}\n".encode()) + + def to_lp_file( m: Model, fn: Path, @@ -442,6 +563,12 @@ def to_lp_file( slice_size=slice_size, explicit_coordinate_names=explicit_coordinate_names, ) + quadratic_constraints_to_file( + m, + f=f, + progress=progress, + explicit_coordinate_names=explicit_coordinate_names, + ) bounds_to_file( m, f=f, @@ -505,15 +632,23 @@ def to_file( ) elif io_api == "mps": - if "highs" not in solvers.available_solvers: - raise RuntimeError( - "Package highspy not installed. This is required to exporting to MPS file." - ) - - # Use very fast highspy implementation - # Might be replaced by custom writer, however needs C/Rust bindings for performance - h = m.to_highspy(explicit_coordinate_names=explicit_coordinate_names) - h.writeModel(str(fn)) + if m.has_quadratic_constraints: + # MPS with quadratic constraints requires Gurobi + if "gurobi" not in solvers.available_solvers: + raise RuntimeError( + "Package Gurobipy not installed. This is requiredd for MPS export with quadratic constraints. " + "Use LP format instead" + ) + gm = m.to_gurobipy(explicit_coordinate_names=explicit_coordinate_names) + gm.write(str(fn)) + else: + # Use fast HiGHS implementation for models without QC + if "highs" not in solvers.available_solvers: + raise RuntimeError( + "Package highspy not installed. This is required for exporting to MPS file." + ) + h = m.to_highspy(explicit_coordinate_names=explicit_coordinate_names) + h.writeModel(str(fn)) else: raise ValueError( f"Invalid io_api '{io_api}'. Choose from 'lp', 'lp-polars' or 'mps'." @@ -558,7 +693,12 @@ def to_mosek( labels = np.vectorize(print_variable)(M.vlabels).astype(object) task.generatevarnames( - np.arange(0, len(labels)), "%0", [len(labels)], None, [0], list(labels) + np.arange(0, len(labels), dtype=np.int32), + "%0", + [len(labels)], + None, + [0], + list(labels), ) ## Variables @@ -617,7 +757,12 @@ def to_mosek( if M.A is not None: A = M.A.tocsr() task.putarowslice( - 0, m.ncons, A.indptr[:-1], A.indptr[1:], A.indices, A.data + 0, + m.ncons, + A.indptr[:-1].astype(np.int64), + A.indptr[1:].astype(np.int64), + A.indices.astype(np.int32), + A.data.astype(np.float64), ) task.putconboundslice(0, m.ncons, bkc, blc, buc) @@ -631,6 +776,65 @@ def to_mosek( task.putobjsense(mosek.objsense.maximize) else: task.putobjsense(mosek.objsense.minimize) + + ## Quadratic Constraints + if len(m.quadratic_constraints): + # Get the number of quadratic constraints + n_qcons = len(M.qclabels) + + # Append quadratic constraints to the task + # In MOSEK, quadratic constraints are added to regular constraints + # with quadratic terms via putqconk + qc_start_idx = m.ncons # Start after linear constraints + task.appendcons(n_qcons) + + # Get matrices for QC + Qc_list = M.Qc + qc_linear = M.qc_linear + qc_sense = M.qc_sense + qc_rhs = M.qc_rhs + + for i, label in enumerate(M.qclabels): + con_idx = qc_start_idx + i + + # Set constraint name + task.putconname(con_idx, f"qc{label}") + + # Set constraint bound based on sense + # Note: For boundkey.up, bound_lo is ignored; for boundkey.lo, bound_up is ignored + # We use -inf/+inf to match linear constraint handling pattern + sense = qc_sense[i] + rhs = qc_rhs[i] + if sense == "<=": + bk = mosek.boundkey.up + bound_lo, bound_up = -np.inf, rhs + elif sense == ">=": + bk = mosek.boundkey.lo + bound_lo, bound_up = rhs, np.inf + else: # "=" + bk = mosek.boundkey.fx + bound_lo, bound_up = rhs, rhs + task.putconbound(con_idx, bk, bound_lo, bound_up) + + # Add linear terms if any + if qc_linear is not None: + row = qc_linear.getrow(i).tocoo() + if row.nnz > 0: + task.putarow(con_idx, list(row.col), list(row.data)) + + # Add quadratic terms + # MOSEK expects lower triangular part only + Q = Qc_list[i] + if Q.nnz > 0: + # Get lower triangular part (MOSEK requirement) + Q_lower = tril(Q).tocoo() + # MOSEK uses 0.5 * x'Qx convention, and our Q matrix is already + # built with doubled diagonal terms for this convention. + # So we pass Q directly without dividing. + task.putqconk( + con_idx, list(Q_lower.row), list(Q_lower.col), list(Q_lower.data) + ) + return task @@ -670,7 +874,7 @@ def to_gurobipy( kwargs["vtype"] = M.vtypes x = model.addMVar(M.vlabels.shape, M.lb, M.ub, name=list(names), **kwargs) - if m.is_quadratic: + if not m.objective.is_linear: model.setObjective(0.5 * x.T @ M.Q @ x + M.c @ x) # type: ignore else: model.setObjective(M.c @ x) @@ -683,6 +887,72 @@ def to_gurobipy( c = model.addMConstr(M.A, x, M.sense, M.b) # type: ignore c.setAttr("ConstrName", list(names)) # type: ignore + # Add quadratic constraints + if len(m.quadratic_constraints): + for name in m.quadratic_constraints: + qcon = m.quadratic_constraints[name] + df = qcon.to_polars() + + # Build label -> (sign, rhs) lookup for multi-dimensional constraints + labels_flat = qcon.labels.values.ravel() + signs_flat = np.broadcast_to(qcon.sign.values, qcon.labels.shape).ravel() + rhs_flat = np.broadcast_to(qcon.rhs.values, qcon.labels.shape).ravel() + label_metadata = { + int(lab): (str(signs_flat[i]), float(rhs_flat[i])) + for i, lab in enumerate(labels_flat) + if lab != -1 + } + + # Build QuadExpr for each constraint label + for label in df["labels"].unique().to_list(): + label_df = df.filter(pl.col("labels") == label) + + # Build the quadratic expression + qexpr = gurobipy.QuadExpr() + + # Add linear terms + linear_terms = label_df.filter(~pl.col("is_quadratic")) + for row in linear_terms.iter_rows(named=True): + coeff = row["coeffs"] + var_idx = int(row["vars"]) + if var_idx >= 0: + qexpr.addTerms(coeff, x[var_idx].item()) + + # Add quadratic terms + quad_terms = label_df.filter(pl.col("is_quadratic")) + for row in quad_terms.iter_rows(named=True): + coeff = row["coeffs"] + var1 = int(row["vars1"]) + var2 = int(row["vars2"]) + qexpr.addTerms(coeff, x[var1].item(), x[var2].item()) + + # Get sign and rhs for this specific label + sign, rhs = label_metadata[int(label)] + + # Map sign to gurobipy sense + if sign == "<=": + sense = gurobipy.GRB.LESS_EQUAL + elif sign == ">=": + sense = gurobipy.GRB.GREATER_EQUAL + else: + sense = gurobipy.GRB.EQUAL + + # Determine constraint name + if explicit_coordinate_names: + label_pos = m.quadratic_constraints.get_label_position(label) + if label_pos and label_pos[0] is not None: + con_name, coord = label_pos + assert isinstance(con_name, str) + assert isinstance(coord, dict) + coord_str = print_coord(coord) + qc_name = f"{clean_name(con_name)}{coord_str}#{label}" + else: + qc_name = f"{clean_name(name)}#{label}" + else: + qc_name = f"qc{label}" + + model.addQConstr(qexpr, sense, rhs, name=qc_name) + model.update() return model @@ -706,6 +976,12 @@ def to_highspy(m: Model, explicit_coordinate_names: bool = False) -> Highs: """ import highspy + if m.has_quadratic_constraints: + raise ValueError( + "HiGHS does not support quadratic constraints. " + "Use a solver that supports QCP: gurobi, cplex, mosek, xpress, copt, scip." + ) + print_variable, print_constraint = get_printers_scalar( m, explicit_coordinate_names=explicit_coordinate_names ) @@ -917,6 +1193,10 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: with_prefix(con.data, f"constraints-{name}") for name, con in m.constraints.items() ] + qcons = [ + with_prefix(qcon.data, f"quadratic_constraints-{name}") + for name, qcon in m.quadratic_constraints.items() + ] objective = m.objective.data objective = objective.assign_attrs(sense=m.objective.sense) if m.objective.value is not None: @@ -925,7 +1205,7 @@ def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: params = [with_prefix(m.parameters, "parameters")] scalars = {k: getattr(m, k) for k in m.scalar_attrs} - ds = xr.merge(vars + cons + obj + params, combine_attrs="drop_conflicts") + ds = xr.merge(vars + cons + qcons + obj + params, combine_attrs="drop_conflicts") ds = ds.assign_attrs(scalars) ds.attrs = non_bool_dict(ds.attrs) @@ -950,6 +1230,7 @@ def read_netcdf(path: Path | str, **kwargs: Any) -> Model: ------- m : linopy.Model """ + from linopy.constraints import QuadraticConstraint, QuadraticConstraints from linopy.model import ( Constraint, Constraints, @@ -1002,7 +1283,7 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: m._variables = Variables(variables, m) - cons = [str(k) for k in ds if str(k).startswith("constraints")] + cons = [str(k) for k in ds if str(k).startswith("constraints-")] con_names = list({str(k).rsplit("-", 1)[0] for k in cons}) constraints = {} for k in sorted(con_names): @@ -1010,6 +1291,16 @@ def get_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: constraints[name] = Constraint(get_prefix(ds, k), m, name) m._constraints = Constraints(constraints, m) + qcons = [str(k) for k in ds if str(k).startswith("quadratic_constraints-")] + qcon_names = list({str(k).rsplit("-", 1)[0] for k in qcons}) + quadratic_constraints = {} + for k in sorted(qcon_names): + name = remove_prefix(k, "quadratic_constraints") + quadratic_constraints[name] = QuadraticConstraint( + get_prefix(ds, k), m, name, skip_broadcast=True + ) + m._quadratic_constraints = QuadraticConstraints(quadratic_constraints, m) + objective = get_prefix(ds, "objective") m.objective = Objective( LinearExpression(objective, m), m, objective.attrs.pop("sense") diff --git a/linopy/matrices.py b/linopy/matrices.py index a55bb0bd..0f35cd5b 100644 --- a/linopy/matrices.py +++ b/linopy/matrices.py @@ -177,3 +177,210 @@ def Q(self) -> csc_matrix | None: if not isinstance(expr, expressions.QuadraticExpression): return None return expr.to_matrix()[self.vlabels][:, self.vlabels] + + # Quadratic constraint accessors + + @cached_property + def flat_qcons(self) -> pd.DataFrame: + """Flat DataFrame of all quadratic constraints.""" + m = self._parent + return m.quadratic_constraints.flat + + @property + def qclabels(self) -> ndarray: + """Vector of labels of all non-missing quadratic constraints.""" + df: pd.DataFrame = self.flat_qcons + if df.empty: + return np.array([], dtype=int) + return np.sort(df["labels"].unique()) + + @property + def qc_sense(self) -> ndarray: + """Vector of senses of all non-missing quadratic constraints.""" + m = self._parent + if not len(m.quadratic_constraints): + return np.array([], dtype=" ndarray: + """Vector of right-hand-sides of all non-missing quadratic constraints.""" + m = self._parent + if not len(m.quadratic_constraints): + return np.array([], dtype=float) + + labels = self.qclabels + rhs = np.empty(len(labels), dtype=float) + # Track which labels have been set to detect duplicates + seen_labels: set[int] = set() + + for name in m.quadratic_constraints: + qcon = m.quadratic_constraints[name] + qc_labels = qcon.labels.values.ravel() + qc_rhs = np.broadcast_to(qcon.rhs.values, qcon.labels.shape).ravel() + for i, lab in enumerate(qc_labels): + if lab != -1: + assert lab not in seen_labels, ( + f"Duplicate quadratic constraint label {lab} found. " + "Labels must be unique across all quadratic constraints." + ) + seen_labels.add(lab) + idx = np.searchsorted(labels, lab) + rhs[idx] = float(qc_rhs[i]) + + return rhs + + @property + def Qc(self) -> list[csc_matrix]: + """ + List of Q matrices for quadratic constraints. + + Returns a list where each element is a sparse matrix representing the + quadratic terms of one constraint. The matrix follows the convention + x'Qx, where Q is symmetric with doubled diagonal terms. + + Notes + ----- + This method assumes that flat_qcons stores each quadratic term exactly + once (i.e., for off-diagonal terms, only one of (i,j) or (j,i) appears). + The method then constructs a symmetric Q matrix by: + - Doubling diagonal terms (x_i^2 coefficient becomes 2*coeff) + - Adding both (i,j) and (j,i) entries for off-diagonal terms + + If flat_qcons semantics change to store already-symmetric entries, + this logic would need adjustment to avoid double-counting. + """ + m = self._parent + if not len(m.quadratic_constraints): + return [] + + df = self.flat_qcons + labels = self.qclabels + n_vars = len(self.vlabels) + + # Build variable label to index mapping + var_map = pd.Series(index=self.vlabels, data=np.arange(n_vars)) + + matrices = [] + for label in labels: + label_df = df[(df["labels"] == label) & df["is_quadratic"]] + + if label_df.empty: + # No quadratic terms - empty matrix + matrices.append(csc_matrix((n_vars, n_vars))) + continue + + rows = [] + cols = [] + data = [] + + for _, row in label_df.iterrows(): + var1 = int(row["vars1"]) + var2 = int(row["vars2"]) + coeff = row["coeffs"] + + if var1 < 0 or var2 < 0: + continue + + # Map to matrix indices + i = var_map.get(var1, -1) + j = var_map.get(var2, -1) + + if i < 0 or j < 0: + continue + + if i == j: + # Diagonal term - double it for x'Qx convention + rows.append(i) + cols.append(j) + data.append(coeff * 2) + else: + # Off-diagonal - add symmetric entries + rows.extend([i, j]) + cols.extend([j, i]) + data.extend([coeff, coeff]) + + Q = csc_matrix((data, (rows, cols)), shape=(n_vars, n_vars)) + matrices.append(Q) + + return matrices + + @property + def qc_linear(self) -> csc_matrix | None: + """ + Matrix of linear coefficients for quadratic constraints. + + Returns a sparse matrix of shape (n_qconstraints, n_variables) where + each row contains the linear coefficients for one quadratic constraint. + """ + m = self._parent + if not len(m.quadratic_constraints): + return None + + df = self.flat_qcons + labels = self.qclabels + n_cons = len(labels) + n_vars = len(self.vlabels) + + if n_cons == 0 or n_vars == 0: + return csc_matrix((n_cons, n_vars)) + + # Build variable label to index mapping + var_map = pd.Series(index=self.vlabels, data=np.arange(n_vars)) + + # Build constraint label to index mapping + con_map = pd.Series(index=labels, data=np.arange(n_cons)) + + # Filter to linear terms only + # Note: 'vars' column may not exist if there are no linear terms + if "vars" not in df.columns: + return csc_matrix((n_cons, n_vars)) + + linear_df = df[~df["is_quadratic"] & (df["vars"] >= 0)] + + if linear_df.empty: + return csc_matrix((n_cons, n_vars)) + + rows = [] + cols = [] + data = [] + skipped_rows = 0 + + for _, row in linear_df.iterrows(): + con_idx = con_map.get(row["labels"], -1) + var_idx = var_map.get(int(row["vars"]), -1) + + if con_idx >= 0 and var_idx >= 0: + rows.append(con_idx) + cols.append(var_idx) + data.append(row["coeffs"]) + else: + skipped_rows += 1 + + assert skipped_rows == 0, ( + f"Skipped {skipped_rows} linear term(s) in qc_linear due to " + "missing variable or constraint labels. This indicates a mismatch " + "between flat_qcons and the global variable/constraint indexing." + ) + + return csc_matrix((data, (rows, cols)), shape=(n_cons, n_vars)) diff --git a/linopy/model.py b/linopy/model.py index 3982b84d..b96a8094 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -42,7 +42,13 @@ ModelStatus, TerminationCondition, ) -from linopy.constraints import AnonymousScalarConstraint, Constraint, Constraints +from linopy.constraints import ( + AnonymousScalarConstraint, + Constraint, + Constraints, + QuadraticConstraint, + QuadraticConstraints, +) from linopy.expressions import ( LinearExpression, QuadraticExpression, @@ -63,6 +69,7 @@ IO_APIS, NO_SOLUTION_FILE_SOLVERS, available_solvers, + quadratic_constraint_solvers, quadratic_solvers, ) from linopy.types import ( @@ -98,6 +105,7 @@ class Model: solver_name: str _variables: Variables _constraints: Constraints + _quadratic_constraints: QuadraticConstraints _objective: Objective _parameters: Dataset _solution: Dataset @@ -118,6 +126,7 @@ class Model: # containers "_variables", "_constraints", + "_quadratic_constraints", "_objective", "_parameters", "_solution", @@ -128,8 +137,10 @@ class Model: # TODO: move counters to Variables and Constraints class "_xCounter", "_cCounter", + "_qcCounter", "_varnameCounter", "_connameCounter", + "_qconnameCounter", "_blocks", # TODO: check if these should not be mutable "_chunk", @@ -171,6 +182,9 @@ def __init__( """ self._variables: Variables = Variables({}, model=self) self._constraints: Constraints = Constraints({}, model=self) + self._quadratic_constraints: QuadraticConstraints = QuadraticConstraints( + {}, model=self + ) self._objective: Objective = Objective(LinearExpression(None, self), self) self._parameters: Dataset = Dataset() @@ -178,8 +192,10 @@ def __init__( self._termination_condition: str = "" self._xCounter: int = 0 self._cCounter: int = 0 + self._qcCounter: int = 0 self._varnameCounter: int = 0 self._connameCounter: int = 0 + self._qconnameCounter: int = 0 self._blocks: DataArray | None = None self._chunk: T_Chunks = chunk @@ -204,6 +220,13 @@ def constraints(self) -> Constraints: """ return self._constraints + @property + def quadratic_constraints(self) -> QuadraticConstraints: + """ + Quadratic constraints assigned to the model. + """ + return self._quadratic_constraints + @property def objective(self) -> Objective: """ @@ -688,6 +711,120 @@ def add_constraints( self.constraints.add(constraint) return constraint + def add_quadratic_constraints( + self, + lhs: QuadraticExpression | LinearExpression | QuadraticConstraint | Callable, + sign: SignLike | None = None, + rhs: ConstantLike | None = None, + name: str | None = None, + coords: Sequence[Sequence | pd.Index | DataArray] | Mapping | None = None, + mask: MaskLike | None = None, + ) -> QuadraticConstraint: + """ + Add a quadratic constraint to the model. + + Quadratic constraints are of the form: x'Qx + a'x <= b (or >=, =) + + Parameters + ---------- + lhs : linopy.QuadraticExpression, linopy.QuadraticConstraint, or callable + Left hand side of the constraint(s) or a full quadratic constraint. + If a QuadraticExpression is passed, `sign` and `rhs` must be provided. + If a callable is passed, it is called for every combination of + coordinates given in `coords`. + sign : str or array_like, optional + Relation between the lhs and rhs: '=', '>=', or '<='. + rhs : float or array_like, optional + Right hand side constant(s) of the constraint. + name : str, optional + Reference name for the constraint. Default generates names like + "qcon0", "qcon1", etc. + coords : list or xarray.Coordinates, optional + Coordinates for the constraint array. Only used when lhs is a callable. + mask : array_like, optional + Boolean mask indicating which constraints to include. + + Returns + ------- + linopy.QuadraticConstraint + The quadratic constraint added to the model. + + Examples + -------- + >>> m = Model() + >>> x = m.add_variables(name="x") + >>> y = m.add_variables(name="y") + >>> qc = m.add_quadratic_constraints(x**2 + y**2, "<=", 25, name="circle") + """ + if name in list(self.quadratic_constraints): + raise ValueError(f"Quadratic constraint '{name}' already assigned to model") + elif name is None: + name = f"qcon{self._qconnameCounter}" + self._qconnameCounter += 1 + + if sign is not None: + from linopy.common import maybe_replace_signs + + sign = maybe_replace_signs(as_dataarray(sign)) + + if isinstance(lhs, QuadraticExpression): + if sign is None or rhs is None: + raise ValueError( + "Arguments `sign` and `rhs` must be provided when lhs is a QuadraticExpression." + ) + data = lhs.to_constraint(sign, rhs).data + elif isinstance(lhs, QuadraticConstraint): + if sign is not None or rhs is not None: + raise ValueError( + "Arguments `sign` and `rhs` cannot be provided when lhs is a QuadraticConstraint." + ) + data = lhs.data + elif callable(lhs): + raise NotImplementedError( + "Rule-based quadratic constraint creation is not yet implemented." + ) + else: + raise ValueError( + f"Invalid type for `lhs` ({type(lhs)}). Expected QuadraticExpression, " + "QuadraticConstraint, or callable." + ) + + # Ensure helper dimensions are not set as coordinates + if drop_dims := set(HELPER_DIMS).intersection(data.coords): + data = data.drop_vars(drop_dims) + + data["labels"] = -1 + from linopy.constants import FACTOR_DIM, QTERM_DIM + + (data,) = xr.broadcast(data, exclude=[TERM_DIM, QTERM_DIM, FACTOR_DIM]) + + if mask is not None: + mask = as_dataarray(mask).astype(bool) + assert set(mask.dims).issubset(data.dims), ( + "Dimensions of mask not a subset of resulting labels dimensions." + ) + + self.check_force_dim_names(data) + + start = self._qcCounter + end = start + data.labels.size + data.labels.values = np.arange(start, end).reshape(data.labels.shape) + self._qcCounter += data.labels.size + + if mask is not None: + data.labels.values = data.labels.where(mask, -1).values + + data = data.assign_attrs(label_range=(start, end), name=name) + + if self.chunk: + data = data.chunk(self.chunk) + + constraint = QuadraticConstraint( + data, name=name, model=self, skip_broadcast=True + ) + self.quadratic_constraints.add(constraint) + return constraint + def add_objective( self, expr: Variable @@ -809,14 +946,30 @@ def integers(self) -> Variables: @property def is_linear(self) -> bool: - return self.objective.is_linear + return self.objective.is_linear and not self.has_quadratic_constraints @property def is_quadratic(self) -> bool: - return self.objective.is_quadratic + return self.objective.is_quadratic or self.has_quadratic_constraints + + @property + def has_quadratic_constraints(self) -> bool: + """Return True if the model has any quadratic constraints.""" + return len(self.quadratic_constraints) > 0 @property def type(self) -> str: + """ + Return the problem type string. + + Returns one of: + - LP: Linear Program + - QP: Quadratic Program (quadratic objective, linear constraints) + - QCP: Quadratically Constrained Program (linear objective, quadratic constraints) + - QCQP: Quadratically Constrained Quadratic Program + - ILP/IQP/IQCP/IQCQP: Integer versions + - MILP/MIQP/MIQCP/MIQCQP: Mixed-Integer versions + """ if (len(self.binaries) or len(self.integers)) and len(self.continuous): variable_type = "MI" elif len(self.binaries) or len(self.integers): @@ -824,9 +977,14 @@ def type(self) -> str: else: variable_type = "" - objective_type = "Q" if self.is_quadratic else "L" + # Determine constraint type + has_qc = self.has_quadratic_constraints + constraint_type = "QC" if has_qc else "" - return f"{variable_type}{objective_type}P" + # Determine objective type + objective_type = "Q" if self.objective.is_quadratic else "L" + + return f"{variable_type}{constraint_type}{objective_type}P" @property def nvars(self) -> int: @@ -1213,6 +1371,15 @@ def solve( f"Solver {solver_name} does not support quadratic problems." ) + if ( + self.has_quadratic_constraints + and solver_name not in quadratic_constraint_solvers + ): + raise ValueError( + f"Solver {solver_name} does not support quadratic constraints. " + f"Use one of: {quadratic_constraint_solvers}" + ) + try: solver_class = getattr(solvers, f"{solvers.SolverName(solver_name).name}") # initialize the solver as object of solver subclass @@ -1294,6 +1461,20 @@ def solve( vals = dual.reindex(idx).values.reshape(con.labels.shape) con.dual = xr.DataArray(vals, con.labels.coords) + # Assign quadratic constraint dual values + if not result.solution.qc_dual.empty: + qc_dual = result.solution.qc_dual.copy() + qc_dual = set_int_index(qc_dual) + qc_dual.loc[-1] = nan + + for name, qcon in self.quadratic_constraints.items(): + idx = np.ravel(qcon.labels) + try: + vals = qc_dual[idx].values.reshape(qcon.labels.shape) + except KeyError: + vals = qc_dual.reindex(idx).values.reshape(qcon.labels.shape) + qcon.dual = xr.DataArray(vals, qcon.labels.coords) + return result.status.status.value, result.status.termination_condition.value def compute_infeasibilities(self) -> list[int]: diff --git a/linopy/solvers.py b/linopy/solvers.py index a121a2b5..1ce54958 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -50,6 +50,21 @@ "mindopt", ] +# Solvers that support quadratic constraints (QCP/QCQP) via LP file format +# Note: HiGHS does NOT support quadratic constraints +# Note: SCIP supports QC but its LP file reading doesn't work atm +# Note: XPRESS support convex QC but are excluded due to licensing/config complexity with QCQP +QUADRATIC_CONSTRAINT_SOLVERS = [ + "gurobi", + "mosek", + "copt", +] + +# Solvers that support nonconvex quadratic constraints (bilinear terms, >= constraints) +NONCONVEX_QUADRATIC_CONSTRAINT_SOLVERS = [ + "gurobi", +] + # Solvers that don't need a solution file when keep_files=False NO_SOLUTION_FILE_SOLVERS = [ "xpress", @@ -146,6 +161,12 @@ class xpress_Namespaces: # type: ignore[no-redef] pass quadratic_solvers = [s for s in QUADRATIC_SOLVERS if s in available_solvers] +quadratic_constraint_solvers = [ + s for s in QUADRATIC_CONSTRAINT_SOLVERS if s in available_solvers +] +nonconvex_quadratic_constraint_solvers = [ + s for s in NONCONVEX_QUADRATIC_CONSTRAINT_SOLVERS if s in available_solvers +] logger = logging.getLogger(__name__) @@ -1181,10 +1202,20 @@ def get_solver_solution() -> Solution: logger.warning("Dual values of MILP couldn't be parsed") dual = pd.Series(dtype=float) - return Solution(sol, dual, objective) + # Retrieve quadratic constraint dual values + qc_dual = pd.Series(dtype=float) + try: + qcs = m.getQConstrs() + if qcs: + qc_dual = pd.Series({qc.QCName: qc.QCPi for qc in qcs}, dtype=float) + except (AttributeError, gurobipy.GurobiError): + # QCPi not available (non-convex or QCPDual=0) + pass + + return Solution(sol, dual, objective, qc_dual) solution = self.safe_get_solution(status=status, func=get_solver_solution) - solution = solution = maybe_adjust_objective_sign(solution, io_api, sense) + solution = maybe_adjust_objective_sign(solution, io_api, sense) return Result(status, solution, m) diff --git a/test/test_common.py b/test/test_common.py index 0ec933bf..d2f55538 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -708,6 +708,6 @@ def test_align(x: Variable, u: Variable) -> None: # noqa: F811 expr = 20 * x x_obs, expr_obs, alpha_obs = align(x, expr, alpha) assert x_obs.shape == alpha_obs.shape == (1,) - assert expr_obs.shape == (1, 1) # _term dim + assert expr_obs.shape == (1,) # _term dim excluded from shape assert isinstance(expr_obs, LinearExpression) assert_linequal(expr_obs, expr.loc[[1]]) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 2551c203..ae44e28e 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -317,7 +317,7 @@ def test_linear_expression_with_constant_multiplication( obs = expr * pd.Series([1, 2, 3], index=pd.RangeIndex(3, name="new_dim")) assert isinstance(obs, LinearExpression) - assert obs.shape == (2, 3, 1) + assert obs.shape == (2, 3) def test_linear_expression_multi_indexed(u: Variable) -> None: diff --git a/test/test_quadratic_constraint.py b/test/test_quadratic_constraint.py new file mode 100644 index 00000000..f74009c4 --- /dev/null +++ b/test/test_quadratic_constraint.py @@ -0,0 +1,1232 @@ +#!/usr/bin/env python3 +""" +Tests for quadratic constraints. +""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import numpy as np +import pandas as pd +import polars as pl +import pytest + +import linopy +from linopy import Model +from linopy.constraints import QuadraticConstraint +from linopy.solvers import ( + available_solvers, + nonconvex_quadratic_constraint_solvers, + quadratic_constraint_solvers, +) + +# Build parameter list: (solver, io_api) for QC-capable solvers +qc_solver_params: list[tuple[str, str]] = [] +for solver in quadratic_constraint_solvers: + if solver in available_solvers: + qc_solver_params.append((solver, "lp")) + if solver in ["gurobi", "mosek"]: + qc_solver_params.append((solver, "direct")) + +# Build parameter list for nonconvex QC tests +nonconvex_qc_solver_params: list[tuple[str, str]] = [] +for solver in nonconvex_quadratic_constraint_solvers: + if solver in available_solvers: + nonconvex_qc_solver_params.append((solver, "lp")) + if solver == "gurobi": + nonconvex_qc_solver_params.append((solver, "direct")) + + +@pytest.fixture +def m() -> Model: + m = Model() + m.add_variables(lower=0, name="x") + m.add_variables(lower=0, name="y") + return m + + +@pytest.fixture +def x(m: Model) -> linopy.Variable: + return m.variables["x"] + + +@pytest.fixture +def y(m: Model) -> linopy.Variable: + return m.variables["y"] + + +class TestQuadraticConstraintCreation: + """Tests for quadratic constraint creation.""" + + def test_create_simple_quadratic_constraint( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test creating a simple quadratic constraint.""" + qexpr = x * x + y * y + qcon = m.add_quadratic_constraints(qexpr, "<=", 100, name="qc1") + + assert isinstance(qcon, QuadraticConstraint) + assert qcon.name == "qc1" + assert str(qcon.sign.values) == "<=" + assert float(qcon.rhs.values) == 100.0 + + def test_create_mixed_quadratic_constraint( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test creating a quadratic constraint with both linear and quadratic terms.""" + qexpr = x * x + 2 * x * y + y * y + 3 * x + 4 * y + qcon = m.add_quadratic_constraints(qexpr, "<=", 100, name="mixed") + + assert isinstance(qcon, QuadraticConstraint) + # Check repr works + repr_str = repr(qcon) + assert "x²" in repr_str or "x^2" in repr_str or "x·x" in repr_str + + def test_create_cross_product_constraint( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test creating a constraint with cross product term.""" + qexpr = x * y + qcon = m.add_quadratic_constraints(qexpr, "<=", 10, name="cross") + + assert isinstance(qcon, QuadraticConstraint) + + def test_create_constraint_with_different_signs( + self, m: Model, x: linopy.Variable + ) -> None: + """Test creating constraints with different comparison operators.""" + qexpr = x * x + + qcon_le = m.add_quadratic_constraints(qexpr, "<=", 100, name="qc_le") + assert str(qcon_le.sign.values) == "<=" + + qcon_ge = m.add_quadratic_constraints(qexpr, ">=", 1, name="qc_ge") + assert str(qcon_ge.sign.values) == ">=" + + qcon_eq = m.add_quadratic_constraints(qexpr, "==", 50, name="qc_eq") + assert str(qcon_eq.sign.values) == "=" + + def test_create_constraint_with_negative_rhs( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test creating a constraint where terms move to create negative RHS.""" + qexpr = x * x - 100 + qcon = m.add_quadratic_constraints(qexpr, "<=", 0, name="neg_rhs") + + assert isinstance(qcon, QuadraticConstraint) + + +class TestQuadraticConstraintsContainer: + """Tests for the QuadraticConstraints container class.""" + + def test_empty_quadratic_constraints(self) -> None: + """Test that a new model has empty quadratic constraints.""" + m = Model() + assert len(m.quadratic_constraints) == 0 + assert list(m.quadratic_constraints) == [] + + def test_add_single_constraint( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test adding a single quadratic constraint.""" + qexpr = x * x + y * y + m.add_quadratic_constraints(qexpr, "<=", 100, name="qc1") + + assert len(m.quadratic_constraints) == 1 + assert "qc1" in m.quadratic_constraints + assert m.quadratic_constraints["qc1"].name == "qc1" + + def test_add_multiple_constraints( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test adding multiple quadratic constraints.""" + m.add_quadratic_constraints(x * x, "<=", 100, name="qc1") + m.add_quadratic_constraints(y * y, "<=", 50, name="qc2") + m.add_quadratic_constraints(x * y, "<=", 25, name="qc3") + + assert len(m.quadratic_constraints) == 3 + assert set(m.quadratic_constraints) == {"qc1", "qc2", "qc3"} + + def test_remove_constraint( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test removing a quadratic constraint.""" + m.add_quadratic_constraints(x * x, "<=", 100, name="qc1") + m.add_quadratic_constraints(y * y, "<=", 50, name="qc2") + + assert len(m.quadratic_constraints) == 2 + + m.quadratic_constraints.remove("qc1") + + assert len(m.quadratic_constraints) == 1 + assert "qc1" not in m.quadratic_constraints + assert "qc2" in m.quadratic_constraints + + +class TestModelTypeDetection: + """Tests for model type detection with quadratic constraints.""" + + def test_model_type_qclp( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that adding quadratic constraints changes model type to QCLP.""" + # Add linear objective + m.add_objective(x + y) + + # Add quadratic constraint + m.add_quadratic_constraints(x * x + y * y, "<=", 100, name="qc1") + + assert m.type == "QCLP" + assert m.has_quadratic_constraints + + def test_model_type_qcqp( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test QCQP type detection with quadratic objective and constraints.""" + # Add quadratic objective + m.add_objective(x * x + y * y) + + # Add quadratic constraint + m.add_quadratic_constraints(x * y, "<=", 10, name="qc1") + + assert m.type == "QCQP" + + def test_model_type_miqclp(self, x: linopy.Variable, y: linopy.Variable) -> None: + """Test MIQCLP type with integers and quadratic constraints.""" + m = Model() + x = m.add_variables(lower=0, upper=10, integer=True, name="x") + y = m.add_variables(lower=0, name="y") + + m.add_objective(x + y) + m.add_quadratic_constraints(y * y, "<=", 100, name="qc1") + + assert "MI" in m.type + assert "QC" in m.type + + +class TestQuadraticConstraintProperties: + """Tests for QuadraticConstraint properties and methods.""" + + def test_flat_property( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test the flat property returns a DataFrame.""" + qexpr = x * x + 2 * x * y + y * y + qcon = m.add_quadratic_constraints(qexpr, "<=", 100, name="qc1") + + flat_df = qcon.flat + assert isinstance(flat_df, pd.DataFrame) + assert "coeffs" in flat_df.columns + + def test_to_polars(self, m: Model, x: linopy.Variable, y: linopy.Variable) -> None: + """Test the to_polars method.""" + qexpr = x * x + y * y + qcon = m.add_quadratic_constraints(qexpr, "<=", 100, name="qc1") + + df = qcon.to_polars() + assert isinstance(df, pl.DataFrame) + assert "coeffs" in df.columns + assert "is_quadratic" in df.columns + + def test_ncons_property( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test the ncons property of QuadraticConstraints container.""" + m.add_quadratic_constraints(x * x, "<=", 100, name="qc1") + m.add_quadratic_constraints(y * y, "<=", 50, name="qc2") + + assert m.quadratic_constraints.ncons == 2 + + def test_labels_property( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test the labels property.""" + qcon = m.add_quadratic_constraints(x * x + y * y, "<=", 100, name="qc1") + + labels = qcon.labels + assert labels is not None + + +class TestLPFileExport: + """Tests for LP file export with quadratic constraints.""" + + def test_lp_file_with_quadratic_constraint( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that quadratic constraints are written to LP files.""" + m.add_objective(x + y) + m.add_constraints(x + y <= 10, name="linear_c") + m.add_quadratic_constraints(x * x + y * y, "<=", 100, name="qc1") + + with tempfile.NamedTemporaryFile(mode="w", suffix=".lp", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn, progress=False) + content = fn.read_text() + + # Check that quadratic constraint is in the file + assert "qc0" in content or "qc1" in content + assert "^ 2" in content # Squared term + assert "<=" in content + + # Clean up + fn.unlink() + + def test_lp_file_with_mixed_constraint( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test LP export with mixed linear/quadratic constraint.""" + m.add_objective(x + y) + qexpr = x * x + 2 * x * y + y * y + 3 * x + 4 * y + m.add_quadratic_constraints(qexpr, "<=", 100, name="mixed") + + with tempfile.NamedTemporaryFile(mode="w", suffix=".lp", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn, progress=False) + content = fn.read_text() + + # Check for both linear and quadratic terms + assert "[" in content # Opening bracket for quadratic section + assert "]" in content # Closing bracket + assert "^ 2" in content or "* x" in content # Quadratic terms + + # Clean up + fn.unlink() + + def test_lp_file_with_multidimensional_constraint(self) -> None: + """Test LP export with multi-dimensional quadratic constraints.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + y = m.add_variables(lower=0, coords=[range(3)], name="y") + + m.add_objective((x + y).sum()) + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circles") + + with tempfile.NamedTemporaryFile(mode="w", suffix=".lp", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn, progress=False) + content = fn.read_text() + + # Should have 3 quadratic constraints (qc0, qc1, qc2) + assert "qc0:" in content + assert "qc1:" in content + assert "qc2:" in content + + # Clean up + fn.unlink() + + def test_mps_export_with_quadratic_constraints( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that MPS export works with quadratic constraints (using Gurobi).""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available for MPS export with QC") + + m.add_objective(x + y) + m.add_quadratic_constraints(x * x + y * y, "<=", 100, name="qc1") + + with tempfile.NamedTemporaryFile(mode="w", suffix=".mps", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn, progress=False) + content = fn.read_text() + + # Check that QCMATRIX section is present + assert "QCMATRIX" in content + assert "qc" in content.lower() + + fn.unlink(missing_ok=True) + + +class TestFileRoundtrip: + """Tests for export/import/solve roundtrip with quadratic constraints.""" + + def test_lp_roundtrip_solve_compare(self) -> None: + """Test LP export -> read with Gurobi -> solve -> compare solutions.""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available") + + # Create model with QC + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_objective(x + 2 * y, sense="max") + + # Solve directly + m.solve("gurobi", io_api="direct") + direct_x = float(m.solution["x"].values) + direct_y = float(m.solution["y"].values) + direct_obj = m.objective.value + assert direct_obj is not None + + # Export to LP, read back with gurobipy, solve + import gurobipy + + with tempfile.NamedTemporaryFile(suffix=".lp", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn) + gm = gurobipy.read(str(fn)) + gm.optimize() + + var_x = gm.getVarByName("x0") + var_y = gm.getVarByName("x1") + assert var_x is not None and var_y is not None + file_x = var_x.X + file_y = var_y.X + file_obj = gm.ObjVal + + fn.unlink() + + # Compare solutions + assert np.isclose(direct_x, file_x, atol=0.01) + assert np.isclose(direct_y, file_y, atol=0.01) + assert file_obj is not None and np.isclose(direct_obj, file_obj, atol=0.01) + + def test_mps_roundtrip_solve_compare(self) -> None: + """Test MPS export -> read with Gurobi -> solve -> compare solutions.""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available") + + # Create model with QC + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_objective(x + 2 * y, sense="max") + + # Solve directly + m.solve("gurobi", io_api="direct") + direct_x = float(m.solution["x"].values) + direct_y = float(m.solution["y"].values) + direct_obj = m.objective.value + assert direct_obj is not None + + # Export to MPS, read back with gurobipy, solve + import gurobipy + + with tempfile.NamedTemporaryFile(suffix=".mps", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn) + gm = gurobipy.read(str(fn)) + gm.optimize() + + var_x = gm.getVarByName("x0") + var_y = gm.getVarByName("x1") + assert var_x is not None and var_y is not None + file_x = var_x.X + file_y = var_y.X + file_obj = gm.ObjVal + + fn.unlink() + + # Compare solutions + assert np.isclose(direct_x, file_x, atol=0.01) + assert np.isclose(direct_y, file_y, atol=0.01) + assert file_obj is not None and np.isclose(direct_obj, file_obj, atol=0.01) + + def test_lp_roundtrip_multidim_qc(self) -> None: + """Test LP roundtrip with multi-dimensional quadratic constraints.""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available") + + # Create model with multi-dim QC + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + y = m.add_variables(lower=0, coords=[range(3)], name="y") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circles") + m.add_objective((x + 2 * y).sum(), sense="max") + + # Solve directly + m.solve("gurobi", io_api="direct") + direct_obj = m.objective.value + assert direct_obj is not None + + # Export to LP, read back, solve + import gurobipy + + with tempfile.NamedTemporaryFile(suffix=".lp", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn) + gm = gurobipy.read(str(fn)) + gm.optimize() + + file_obj = gm.ObjVal + fn.unlink() + + # Compare objective values (3 independent problems, each with obj ≈ 11.18) + assert file_obj is not None and np.isclose(direct_obj, file_obj, atol=0.05) + assert np.isclose(direct_obj, 3 * 11.18, atol=0.1) + + def test_lp_roundtrip_mixed_linear_qc(self) -> None: + """Test LP roundtrip with both linear and quadratic constraints.""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available") + + # Create model with both linear and quadratic constraints + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_constraints(x + y <= 10, name="linear") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_objective(x + 2 * y, sense="max") + + # Solve directly + m.solve("gurobi", io_api="direct") + direct_obj = m.objective.value + assert direct_obj is not None + + # Export to LP, read back, solve + import gurobipy + + with tempfile.NamedTemporaryFile(suffix=".lp", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn) + gm = gurobipy.read(str(fn)) + gm.optimize() + + file_obj = gm.ObjVal + fn.unlink() + + # Compare objective values + assert file_obj is not None and np.isclose(direct_obj, file_obj, atol=0.01) + + def test_mps_roundtrip_with_linear_terms_in_qc(self) -> None: + """Test MPS roundtrip with QC that has linear terms.""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available") + + # Create model: min x s.t. (x-1)² <= 0 + m = Model() + x = m.add_variables(lower=0, name="x") + m.add_quadratic_constraints(x * x - 2 * x + 1, "<=", 0, name="qc") + m.add_objective(x, sense="min") + + # Solve directly + m.solve("gurobi", io_api="direct") + direct_x = float(m.solution["x"].values) + + # Export to MPS, read back, solve + import gurobipy + + with tempfile.NamedTemporaryFile(suffix=".mps", delete=False) as f: + fn = Path(f.name) + + m.to_file(fn) + gm = gurobipy.read(str(fn)) + gm.optimize() + + var_x = gm.getVarByName("x0") + assert var_x is not None + file_x = var_x.X + fn.unlink() + + # Solution should be x = 1 + assert np.isclose(direct_x, 1.0, atol=0.01) + assert np.isclose(file_x, 1.0, atol=0.01) + + +class TestSolverValidation: + """Tests for solver validation with quadratic constraints.""" + + def test_highs_rejects_quadratic_constraints( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that HiGHS raises an error for quadratic constraints.""" + if "highs" not in linopy.available_solvers: + pytest.skip("HiGHS not available") + + m.add_objective(x + y) # Linear objective + m.add_quadratic_constraints(x * x + y * y, "<=", 100, name="qc1") + + # HiGHS supports QP (quadratic objective) but not QCP (quadratic constraints) + from linopy.solvers import quadratic_constraint_solvers + + if "highs" not in quadratic_constraint_solvers: + with pytest.raises( + ValueError, match="does not support quadratic constraints" + ): + m.solve(solver_name="highs") + + def test_highs_accepts_quadratic_objective( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that HiGHS accepts quadratic objectives (but not QC).""" + if "highs" not in linopy.available_solvers: + pytest.skip("HiGHS not available") + + # Quadratic objective, no quadratic constraints + m.add_objective(x * x + y * y) + m.add_constraints(x + y >= 1, name="c1") + + # This should work - HiGHS supports QP + status, _ = m.solve(solver_name="highs") + assert status == "ok" + + def test_supported_solver_accepts_quadratic_constraints( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that supported solvers accept quadratic constraints.""" + from linopy.solvers import quadratic_constraint_solvers + + # Find a solver that supports QC + available_qc_solvers = [ + s for s in quadratic_constraint_solvers if s in linopy.available_solvers + ] + if not available_qc_solvers: + pytest.skip("No QC-supporting solver available") + + solver = available_qc_solvers[0] + + m.add_objective(x + y, sense="max") + m.add_constraints(x + y <= 10, name="budget") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + + # Should succeed + status, _ = m.solve(solver_name=solver) + assert status == "ok" + + def test_is_quadratic_with_qc_only( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that is_quadratic is True when only QC are present.""" + m.add_objective(x + y) # Linear objective + m.add_quadratic_constraints(x * x, "<=", 10, name="qc") + + assert m.has_quadratic_constraints is True + assert m.objective.is_quadratic is False + assert m.is_quadratic is True # True because of QC + + def test_is_quadratic_with_quadratic_objective_only( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that is_quadratic is True when only quadratic objective.""" + m.add_objective(x * x + y * y) # Quadratic objective + m.add_constraints(x + y <= 10, name="c") # Linear constraint + + assert m.has_quadratic_constraints is False + assert m.objective.is_quadratic is True + assert m.is_quadratic is True # True because of objective + + +class TestQuadraticConstraintRepr: + """Tests for QuadraticConstraint string representations.""" + + def test_repr_simple( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test repr of simple quadratic constraint.""" + qcon = m.add_quadratic_constraints(x * x, "<=", 100, name="simple") + repr_str = repr(qcon) + + assert "QuadraticConstraint" in repr_str + assert "simple" in repr_str + assert "100" in repr_str + + def test_repr_with_cross_term( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test repr includes cross product terms.""" + qcon = m.add_quadratic_constraints(x * y, "<=", 50, name="cross") + repr_str = repr(qcon) + + assert "QuadraticConstraint" in repr_str + # Should show cross product (x·y or similar) + + def test_quadratic_constraints_container_repr( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test repr of QuadraticConstraints container.""" + m.add_quadratic_constraints(x * x, "<=", 100, name="qc1") + m.add_quadratic_constraints(y * y, "<=", 50, name="qc2") + + repr_str = repr(m.quadratic_constraints) + assert "QuadraticConstraints" in repr_str + assert "qc1" in repr_str or "2" in repr_str # Either name or count + + def test_empty_container_repr(self) -> None: + """Test repr of empty QuadraticConstraints.""" + m = Model() + repr_str = repr(m.quadratic_constraints) + assert "QuadraticConstraints" in repr_str + + +class TestMatrixAccessor: + """Tests for matrix accessor with quadratic constraints.""" + + def test_qclabels(self, m: Model, x: linopy.Variable, y: linopy.Variable) -> None: + """Test qclabels property.""" + m.add_objective(x + y) + m.add_quadratic_constraints(x * x, "<=", 25, name="qc1") + m.add_quadratic_constraints(y * y, ">=", 10, name="qc2") + + labels = m.matrices.qclabels + assert len(labels) == 2 + assert labels[0] == 0 + assert labels[1] == 1 + + def test_qc_sense(self, m: Model, x: linopy.Variable, y: linopy.Variable) -> None: + """Test qc_sense property.""" + m.add_objective(x + y) + m.add_quadratic_constraints(x * x, "<=", 25, name="qc1") + m.add_quadratic_constraints(y * y, ">=", 10, name="qc2") + + senses = m.matrices.qc_sense + assert len(senses) == 2 + assert senses[0] == "<=" + assert senses[1] == ">=" + + def test_qc_rhs(self, m: Model, x: linopy.Variable, y: linopy.Variable) -> None: + """Test qc_rhs property.""" + m.add_objective(x + y) + m.add_quadratic_constraints(x * x, "<=", 25, name="qc1") + m.add_quadratic_constraints(y * y, ">=", 10, name="qc2") + + rhs = m.matrices.qc_rhs + assert len(rhs) == 2 + assert rhs[0] == 25.0 + assert rhs[1] == 10.0 + + def test_Qc_matrices( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test Qc property returns Q matrices.""" + m.add_objective(x + y) + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + + Qc = m.matrices.Qc + assert len(Qc) == 1 + Q = Qc[0].toarray() + + # Should be symmetric with doubled diagonal + assert Q[0, 0] == 2.0 # x^2 coefficient doubled + assert Q[1, 1] == 2.0 # y^2 coefficient doubled + + def test_Qc_cross_terms( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test Qc with cross product terms.""" + m.add_objective(x + y) + m.add_quadratic_constraints(x * y, "<=", 10, name="cross") + + Qc = m.matrices.Qc + Q = Qc[0].toarray() + + # Cross term should be symmetric + assert Q[0, 1] == 1.0 + assert Q[1, 0] == 1.0 + assert Q[0, 0] == 0.0 # No x^2 term + assert Q[1, 1] == 0.0 # No y^2 term + + def test_qc_linear(self, m: Model, x: linopy.Variable, y: linopy.Variable) -> None: + """Test qc_linear property.""" + m.add_objective(x + y) + m.add_quadratic_constraints(x * x + 3 * x + 4 * y, "<=", 25, name="mixed") + + A = m.matrices.qc_linear + assert A is not None + assert A.shape == (1, 2) # 1 constraint, 2 variables + + A_dense = A.toarray() + assert A_dense[0, 0] == 3.0 # coefficient of x + assert A_dense[0, 1] == 4.0 # coefficient of y + + def test_empty_quadratic_constraints(self, m: Model) -> None: + """Test matrix accessors with no quadratic constraints.""" + m.add_objective(m.variables["x"]) + + assert len(m.matrices.qclabels) == 0 + assert len(m.matrices.qc_sense) == 0 + assert len(m.matrices.qc_rhs) == 0 + assert len(m.matrices.Qc) == 0 + assert m.matrices.qc_linear is None + + +class TestNetCDFSerialization: + """Tests for netCDF serialization of quadratic constraints.""" + + def test_netcdf_roundtrip( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test saving and loading a model with quadratic constraints.""" + m.add_objective(x + y) + m.add_constraints(x + y <= 10, name="linear") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_quadratic_constraints(x * y, "<=", 10, name="mixed") + + with tempfile.NamedTemporaryFile(suffix=".nc", delete=False) as f: + fn = Path(f.name) + + m.to_netcdf(fn) + m2 = linopy.read_netcdf(fn) + + # Check quadratic constraints were loaded + assert len(m2.quadratic_constraints) == 2 + assert "circle" in m2.quadratic_constraints + assert "mixed" in m2.quadratic_constraints + assert m2.type == "QCLP" + + # Check constraint properties preserved + assert float(m2.quadratic_constraints["circle"].rhs.values) == 25.0 + assert str(m2.quadratic_constraints["circle"].sign.values) == "<=" + + fn.unlink() + + def test_netcdf_roundtrip_multidimensional(self) -> None: + """Test netCDF roundtrip with multi-dimensional quadratic constraints.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + y = m.add_variables(lower=0, coords=[range(3)], name="y") + + m.add_objective((x + y).sum()) + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circles") + + with tempfile.NamedTemporaryFile(suffix=".nc", delete=False) as f: + fn = Path(f.name) + + m.to_netcdf(fn) + m2 = linopy.read_netcdf(fn) + + # Check constraint shape preserved + assert m2.quadratic_constraints["circles"].shape == (3,) + assert len(m2.quadratic_constraints["circles"].labels.values.ravel()) == 3 + + fn.unlink() + + +class TestMOSEKExport: + """Tests for MOSEK direct API export with quadratic constraints.""" + + def test_to_mosek_with_quadratic_constraints( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that to_mosek works with quadratic constraints.""" + if "mosek" not in linopy.available_solvers: + pytest.skip("MOSEK not available") + + m.add_constraints(x + y <= 8, name="budget") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_objective(x + 2 * y, sense="max") + + from linopy.io import to_mosek + + task = to_mosek(m) + # If we got here without error, the export worked + assert task is not None + + def test_to_mosek_multidimensional(self) -> None: + """Test MOSEK export with multi-dimensional quadratic constraints.""" + if "mosek" not in linopy.available_solvers: + pytest.skip("MOSEK not available") + + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + y = m.add_variables(lower=0, coords=[range(3)], name="y") + + m.add_constraints(x + y <= 8, name="budget") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circles") + m.add_objective((x + 2 * y).sum(), sense="max") + + from linopy.io import to_mosek + + task = to_mosek(m) + assert task is not None + + +class TestDualValues: + """Tests for dual value retrieval for quadratic constraints.""" + + def test_qc_dual_with_gurobi( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test that dual values can be retrieved for convex QC with Gurobi.""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available") + + m.add_constraints(x + y <= 8, name="budget") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_objective(x + 2 * y, sense="max") + + # Solve with QCPDual enabled + m.solve(solver_name="gurobi", QCPDual=1) + + # Check dual values exist + dual = m.quadratic_constraints["circle"].dual + assert dual is not None + assert not dual.isnull().all() + # Dual should be positive for binding <= constraint + assert float(dual.values) > 0 + + def test_qc_dual_multidimensional(self) -> None: + """Test dual values for multi-dimensional quadratic constraints.""" + if "gurobi" not in linopy.available_solvers: + pytest.skip("Gurobi not available") + + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + y = m.add_variables(lower=0, coords=[range(3)], name="y") + + m.add_constraints(x + y <= 8, name="budget") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circles") + m.add_objective((x + 2 * y).sum(), sense="max") + + m.solve(solver_name="gurobi", QCPDual=1) + + dual = m.quadratic_constraints["circles"].dual + assert dual.shape == (3,) + assert not dual.isnull().all() + + +# ============================================================================ +# Fixtures for solver correctness tests +# ============================================================================ + + +@pytest.fixture +def qc_circle_model() -> Model: + """ + Model: max x + 2y s.t. x² + y² <= 25, x,y >= 0 + + This is a convex QCP. The optimal point is where the gradient of the + objective (1, 2) is parallel to the gradient of the constraint (2x, 2y). + Solution: x = 1/√5 * 5 ≈ 2.236, y = 2/√5 * 5 ≈ 4.472, obj ≈ 11.18 + """ + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_objective(x + 2 * y, sense="max") + return m + + +@pytest.fixture +def qc_multidim_model() -> Model: + """ + Multi-dimensional model: 3 independent circle constraints. + max sum(x + 2y) s.t. x[i]² + y[i]² <= 25 for each i + Each dimension has same solution as qc_circle_model. + """ + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + y = m.add_variables(lower=0, coords=[range(3)], name="y") + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circles") + m.add_objective((x + 2 * y).sum(), sense="max") + return m + + +@pytest.fixture +def qc_mixed_model() -> Model: + """ + QC with both quadratic and linear terms. + min x s.t. x² - 2x + 1 <= 0, x >= 0 + This is (x-1)² <= 0, so x = 1 exactly. + """ + m = Model() + x = m.add_variables(lower=0, name="x") + m.add_quadratic_constraints(x * x - 2 * x + 1, "<=", 0, name="qc") + m.add_objective(x, sense="min") + return m + + +@pytest.fixture +def qc_cross_terms_model() -> Model: + """ + Model with cross product constraint: xy <= 4. + max x + y s.t. xy <= 4, x,y >= 0, x <= 4, y <= 4 + + This is a NONCONVEX bilinear constraint. The optimal solutions are + corners like (4, 1) or (1, 4) with objective value 5 and xy = 4. + """ + m = Model() + x = m.add_variables(lower=0, upper=4, name="x") + y = m.add_variables(lower=0, upper=4, name="y") + m.add_quadratic_constraints(x * y, "<=", 4, name="cross") + m.add_objective(x + y, sense="max") + return m + + +@pytest.fixture +def qc_geq_model() -> Model: + """ + Greater-than quadratic constraint: x² + y² >= 4. + min x + y s.t. x² + y² >= 4, x,y >= 0 + + This is NONCONVEX. The optimal solution is at an extreme point on + the constraint boundary: either x=0,y=2 or x=2,y=0, giving obj=2. + """ + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_quadratic_constraints(x * x + y * y, ">=", 4, name="circle_geq") + m.add_objective(x + y, sense="min") + return m + + +@pytest.fixture +def qc_equality_model() -> Model: + """ + Equality quadratic constraint: x² + y² = 25. + max x + 2y s.t. x² + y² = 25, x,y >= 0 + Same solution as qc_circle_model since constraint is binding. + """ + m = Model() + x = m.add_variables(lower=0, name="x") + y = m.add_variables(lower=0, name="y") + m.add_quadratic_constraints(x * x + y * y, "=", 25, name="circle_eq") + m.add_objective(x + 2 * y, sense="max") + return m + + +# ============================================================================ +# Solver correctness tests +# ============================================================================ + + +@pytest.mark.skipif(len(qc_solver_params) == 0, reason="No QC solver available") +class TestQuadraticConstraintSolving: + """Tests that verify QC solutions are mathematically correct.""" + + @pytest.mark.parametrize("solver,io_api", qc_solver_params) + def test_qc_circle_solution( + self, qc_circle_model: Model, solver: str, io_api: str + ) -> None: + """Test basic convex QC produces correct solution.""" + status, condition = qc_circle_model.solve(solver, io_api=io_api) + assert status == "ok" + assert condition == "optimal" + + # Expected: x = 5/√5 ≈ 2.236, y = 10/√5 ≈ 4.472, obj ≈ 11.18 + x_val = float(qc_circle_model.solution["x"].values) + y_val = float(qc_circle_model.solution["y"].values) + obj_val = qc_circle_model.objective.value + assert obj_val is not None + + assert np.isclose(x_val, 2.236, atol=0.01) + assert np.isclose(y_val, 4.472, atol=0.01) + assert np.isclose(obj_val, 11.18, atol=0.01) + + @pytest.mark.parametrize("solver,io_api", qc_solver_params) + def test_qc_multidim_solution( + self, qc_multidim_model: Model, solver: str, io_api: str + ) -> None: + """Test multi-dimensional QC with broadcasting.""" + status, condition = qc_multidim_model.solve(solver, io_api=io_api) + assert status == "ok" + assert condition == "optimal" + + # Each dimension should have same solution + x_vals = qc_multidim_model.solution["x"].values + y_vals = qc_multidim_model.solution["y"].values + obj_val = qc_multidim_model.objective.value + assert obj_val is not None + + assert np.allclose(x_vals, 2.236, atol=0.01) + assert np.allclose(y_vals, 4.472, atol=0.01) + assert np.isclose(obj_val, 3 * 11.18, atol=0.05) # 3x single solution + + @pytest.mark.parametrize("solver,io_api", nonconvex_qc_solver_params) + def test_qc_mixed_linear_quad( + self, qc_mixed_model: Model, solver: str, io_api: str + ) -> None: + """Test QC with both quadratic and linear terms (nonconvex).""" + status, condition = qc_mixed_model.solve(solver, io_api=io_api) + assert status == "ok" + assert condition == "optimal" + + # (x-1)² <= 0 means x = 1 exactly + x_val = float(qc_mixed_model.solution["x"].values) + assert np.isclose(x_val, 1.0, atol=0.01) + + @pytest.mark.parametrize("solver,io_api", nonconvex_qc_solver_params) + def test_qc_cross_terms( + self, qc_cross_terms_model: Model, solver: str, io_api: str + ) -> None: + """Test QC with cross product terms (xy) - nonconvex bilinear.""" + status, condition = qc_cross_terms_model.solve(solver, io_api=io_api) + assert status == "ok" + assert condition == "optimal" + + # Nonconvex - verify constraint satisfaction rather than exact values + # Optimal is x+y = 5 with xy = 4 (e.g., x=4,y=1 or x=1,y=4) + x_val = float(qc_cross_terms_model.solution["x"].values) + y_val = float(qc_cross_terms_model.solution["y"].values) + obj_val = qc_cross_terms_model.objective.value + assert obj_val is not None + + # Verify constraint is satisfied + assert x_val * y_val <= 4.0 + 0.01 + # Verify optimal objective value + assert np.isclose(obj_val, 5.0, atol=0.01) + + @pytest.mark.parametrize("solver,io_api", nonconvex_qc_solver_params) + def test_qc_geq_constraint( + self, qc_geq_model: Model, solver: str, io_api: str + ) -> None: + """Test >= quadratic constraint - nonconvex.""" + status, condition = qc_geq_model.solve(solver, io_api=io_api) + assert status == "ok" + assert condition == "optimal" + + # min x+y s.t. x²+y² >= 4, x,y >= 0 + # Optimal: either (0,2) or (2,0) with obj=2 + x_val = float(qc_geq_model.solution["x"].values) + y_val = float(qc_geq_model.solution["y"].values) + obj_val = qc_geq_model.objective.value + assert obj_val is not None + + # Verify constraint is satisfied + assert x_val**2 + y_val**2 >= 4.0 - 0.01 + # Verify optimal objective value + assert np.isclose(obj_val, 2.0, atol=0.01) + + @pytest.mark.parametrize("solver,io_api", nonconvex_qc_solver_params) + def test_qc_equality_constraint( + self, qc_equality_model: Model, solver: str, io_api: str + ) -> None: + """Test = quadratic constraint - nonconvex equality.""" + status, condition = qc_equality_model.solve(solver, io_api=io_api) + assert status == "ok" + assert condition == "optimal" + + # Same as circle model since constraint is binding + x_val = float(qc_equality_model.solution["x"].values) + y_val = float(qc_equality_model.solution["y"].values) + obj_val = qc_equality_model.objective.value + assert obj_val is not None + + # Verify constraint is satisfied (x² + y² = 25) + assert np.isclose(x_val**2 + y_val**2, 25.0, atol=0.1) + # Verify optimal solution + assert np.isclose(x_val, 2.236, atol=0.01) + assert np.isclose(y_val, 4.472, atol=0.01) + assert np.isclose(obj_val, 11.18, atol=0.01) + + +class TestQuadraticConstraintsContainerMethods: + """Tests for QuadraticConstraints container methods.""" + + def test_inequalities_property( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test inequalities property returns only inequality constraints.""" + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="qc_le") + m.add_quadratic_constraints(x * x, ">=", 1, name="qc_ge") + m.add_quadratic_constraints(x * y, "=", 4, name="qc_eq") + + inequalities = m.quadratic_constraints.inequalities + assert len(inequalities) == 2 + assert "qc_le" in inequalities + assert "qc_ge" in inequalities + assert "qc_eq" not in inequalities + + def test_equalities_property( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test equalities property returns only equality constraints.""" + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="qc_le") + m.add_quadratic_constraints(x * x, ">=", 1, name="qc_ge") + m.add_quadratic_constraints(x * y, "=", 4, name="qc_eq") + + equalities = m.quadratic_constraints.equalities + assert len(equalities) == 1 + assert "qc_eq" in equalities + assert "qc_le" not in equalities + assert "qc_ge" not in equalities + + def test_inequalities_empty(self, m: Model, x: linopy.Variable) -> None: + """Test inequalities when all constraints are equalities.""" + m.add_quadratic_constraints(x * x, "=", 1, name="qc_eq") + + inequalities = m.quadratic_constraints.inequalities + assert len(inequalities) == 0 + + def test_equalities_empty(self, m: Model, x: linopy.Variable) -> None: + """Test equalities when all constraints are inequalities.""" + m.add_quadratic_constraints(x * x, "<=", 1, name="qc_le") + + equalities = m.quadratic_constraints.equalities + assert len(equalities) == 0 + + def test_sanitize_zeros( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test sanitize_zeros filters out zero coefficients.""" + # Add a constraint then manually set some coefficients to zero + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="qc") + + # Check that sanitize_zeros runs without error + m.quadratic_constraints.sanitize_zeros() + + # The constraint should still exist + assert "qc" in m.quadratic_constraints + + def test_sanitize_missings( + self, m: Model, x: linopy.Variable, y: linopy.Variable + ) -> None: + """Test sanitize_missings runs without error.""" + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="qc") + + # Check that sanitize_missings runs without error + m.quadratic_constraints.sanitize_missings() + + # The constraint should still exist + assert "qc" in m.quadratic_constraints + + def test_print_labels( + self, + m: Model, + x: linopy.Variable, + y: linopy.Variable, + capsys: pytest.CaptureFixture[str], + ) -> None: + """Test print_labels outputs constraint information.""" + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circle") + m.add_quadratic_constraints(x * y, ">=", 1, name="bilinear") + + m.quadratic_constraints.print_labels([0, 1]) + + captured = capsys.readouterr() + assert "circle" in captured.out + assert "bilinear" in captured.out + assert "≤" in captured.out or "<=" in captured.out + assert "≥" in captured.out or ">=" in captured.out + + def test_print_labels_multidimensional( + self, capsys: pytest.CaptureFixture[str] + ) -> None: + """Test print_labels with multi-dimensional constraints.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + y = m.add_variables(lower=0, coords=[range(3)], name="y") + + m.add_quadratic_constraints(x * x + y * y, "<=", 25, name="circles") + + m.quadratic_constraints.print_labels([0, 1, 2]) + + captured = capsys.readouterr() + assert "circles[0]" in captured.out + assert "circles[1]" in captured.out + assert "circles[2]" in captured.out + assert "x[0]²" in captured.out or "x[0]^2" in captured.out + + def test_inequalities_multidimensional(self) -> None: + """Test inequalities with multi-dimensional constraints.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(3)], name="x") + + m.add_quadratic_constraints(x * x, "<=", 25, name="qc_le") + m.add_quadratic_constraints(x * x, "=", 1, name="qc_eq") + + inequalities = m.quadratic_constraints.inequalities + assert len(inequalities) == 1 + assert "qc_le" in inequalities + # Each constraint has 3 elements (one per coordinate) + assert inequalities["qc_le"].size == 3 diff --git a/test/test_quadratic_expression.py b/test/test_quadratic_expression.py index f5f86c35..118107da 100644 --- a/test/test_quadratic_expression.py +++ b/test/test_quadratic_expression.py @@ -329,8 +329,10 @@ def test_matrices_matrix_mixed_linear_and_quadratic( def test_quadratic_to_constraint(x: Variable, y: Variable) -> None: - with pytest.raises(NotImplementedError): - x * y <= 10 + from linopy.constraints import QuadraticConstraint + + con = x * y <= 10 + assert isinstance(con, QuadraticConstraint) def test_power_of_three(x: Variable) -> None: @@ -344,3 +346,165 @@ def test_power_of_three(x: Variable) -> None: x**3 with pytest.raises(TypeError): (x * x) * (x * x) + + +class TestQuadraticExpressionOperations: + """Tests for QuadraticExpression groupby, rolling, and cumsum operations.""" + + def test_quadratic_expression_cumsum(self) -> None: + """Test cumsum returns QuadraticExpression.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(5)], name="x") + y = m.add_variables(lower=0, coords=[range(5)], name="y") + + qexpr = x * x + y * y + + # Test cumsum + result = qexpr.cumsum() + assert isinstance(result, QuadraticExpression) + # Cumsum should preserve shape + assert result.shape == qexpr.shape + + def test_quadratic_expression_cumsum_with_dim(self) -> None: + """Test cumsum with specific dimension.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(3), range(4)], name="x") + + qexpr = x * x + + # Test cumsum along first dimension + result = qexpr.cumsum(dim="dim_0") + assert isinstance(result, QuadraticExpression) + assert result.shape == qexpr.shape + + def test_quadratic_expression_groupby_sum(self) -> None: + """Test groupby().sum() returns QuadraticExpression.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(6)], name="x") + y = m.add_variables(lower=0, coords=[range(6)], name="y") + + qexpr = x * x + y * y + + # Create a grouping - index name must match dimension name + index = pd.Index(range(6), name="dim_0") + groups = pd.Series([0, 0, 1, 1, 2, 2], index=index, name="group") + + # Test groupby sum + result = qexpr.groupby(groups).sum() + assert isinstance(result, QuadraticExpression) + # Should have 3 groups + assert result.shape == (3,) + + def test_quadratic_expression_rolling_sum(self) -> None: + """Test rolling().sum() returns QuadraticExpression.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(5)], name="x") + y = m.add_variables(lower=0, coords=[range(5)], name="y") + + qexpr = x * x + y * y + + # Test rolling sum with window of 2 + result = qexpr.rolling(dim_0=2).sum() + assert isinstance(result, QuadraticExpression) + # Rolling preserves shape + assert result.shape == qexpr.shape + + def test_quadratic_expression_rolling_different_window(self) -> None: + """Test rolling with different window sizes.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(10)], name="x") + + qexpr = x * x + + # Test with window of 3 + result = qexpr.rolling(dim_0=3).sum() + assert isinstance(result, QuadraticExpression) + assert result.shape == (10,) + + def test_quadratic_expression_groupby_with_dataarray(self) -> None: + """Test groupby with DataArray group.""" + import xarray as xr + + m = Model() + x = m.add_variables(lower=0, coords=[range(6)], name="x") + + qexpr = x * x + + # Create a DataArray for grouping + groups = xr.DataArray([0, 0, 1, 1, 2, 2], dims=["dim_0"], name="group") + + result = qexpr.groupby(groups).sum() + assert isinstance(result, QuadraticExpression) + assert result.shape == (3,) + + def test_quadratic_expression_shape_property(self) -> None: + """Test shape property for QuadraticExpression.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(3), range(4)], name="x") + y = m.add_variables(lower=0, coords=[range(3), range(4)], name="y") + + qexpr = x * x + y * y + + # Shape should exclude TERM_DIM and FACTOR_DIM + assert qexpr.shape == (3, 4) + + def test_quadratic_expression_cumsum_preserves_model(self) -> None: + """Test that cumsum preserves the model reference.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(5)], name="x") + + qexpr = x * x + result = qexpr.cumsum() + + assert result.model is m + + def test_quadratic_expression_groupby_preserves_model(self) -> None: + """Test that groupby().sum() preserves the model reference.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(6)], name="x") + + qexpr = x * x + index = pd.Index(range(6), name="dim_0") + groups = pd.Series([0, 0, 1, 1, 2, 2], index=index, name="group") + + result = qexpr.groupby(groups).sum() + assert result.model is m + + def test_quadratic_expression_rolling_preserves_model(self) -> None: + """Test that rolling().sum() preserves the model reference.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(5)], name="x") + + qexpr = x * x + result = qexpr.rolling(dim_0=2).sum() + + assert result.model is m + + def test_quadratic_expression_cross_product_cumsum(self) -> None: + """Test cumsum with cross product terms.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(4)], name="x") + y = m.add_variables(lower=0, coords=[range(4)], name="y") + + # Cross product creates off-diagonal terms + qexpr = x * y + + result = qexpr.cumsum() + assert isinstance(result, QuadraticExpression) + assert result.shape == (4,) + + def test_quadratic_expression_mixed_terms_groupby(self) -> None: + """Test groupby with both quadratic and linear terms.""" + m = Model() + x = m.add_variables(lower=0, coords=[range(6)], name="x") + y = m.add_variables(lower=0, coords=[range(6)], name="y") + + # Mixed quadratic and linear terms + qexpr = x * x + 2 * x + y + + index = pd.Index(range(6), name="dim_0") + groups = pd.Series([0, 0, 1, 1, 2, 2], index=index, name="group") + result = qexpr.groupby(groups).sum() + + assert isinstance(result, QuadraticExpression) + assert result.shape == (3,)