formatting

Author: RobbieKiwi

doc/release_notes.rst

@@ -3,6 +3,7 @@ Release Notes
 
 .. Upcoming Version
 
+* Add simplify method to LinearExpression to combine duplicate terms
 * Fix compatibility for xpress versions below 9.6 (regression)
 * Performance: Up to 50x faster ``repr()`` for variables/constraints via O(log n) label lookup and direct numpy indexing
 * Performance: Up to 46x faster ``ncons`` property by replacing ``.flat.labels.unique()`` with direct counting

linopy/constants.py

@@ -39,12 +39,14 @@
 GROUP_DIM = "_group"
 FACTOR_DIM = "_factor"
 CONCAT_DIM = "_concat"
+CV_DIM = "_cv"
 HELPER_DIMS: list[str] = [
     TERM_DIM,
     STACKED_TERM_DIM,
     GROUPED_TERM_DIM,
     FACTOR_DIM,
     CONCAT_DIM,
+    CV_DIM,
 ]
 
 

linopy/expressions.py

@@ -66,6 +66,7 @@
 )
 from linopy.config import options
 from linopy.constants import (
+    CV_DIM,
     EQUAL,
     FACTOR_DIM,
     GREATER_EQUAL,
@@ -1473,63 +1474,73 @@ def simplify(self) -> LinearExpression:
         """
 
         def _simplify_row(vars_row: np.ndarray, coeffs_row: np.ndarray) -> np.ndarray:
-            """Simplify a single row by grouping vars and summing coefficients.
-            
+            """
+            Simplify a single row by grouping vars and summing coefficients.
+
             Returns a 2D array of shape (2, input_len) where first row is vars, second is coeffs.
             """
             input_len = len(vars_row)
-            
+
             # Filter out invalid entries
             mask = (vars_row != -1) & (coeffs_row != 0) & ~np.isnan(coeffs_row)
             valid_vars = vars_row[mask]
             valid_coeffs = coeffs_row[mask]
 
             if len(valid_vars) == 0:
                 # Return arrays filled with -1 and 0.0, same length as input
-                return np.vstack([
-                    np.full(input_len, -1, dtype=float),
-                    np.zeros(input_len, dtype=float)
-                ])
+                return np.vstack(
+                    [
+                        np.full(input_len, -1, dtype=float),
+                        np.zeros(input_len, dtype=float),
+                    ]
+                )
 
             # Use bincount to sum coefficients for each variable ID efficiently
             max_var = int(valid_vars.max())
-            summed = np.bincount(valid_vars, weights=valid_coeffs, minlength=max_var + 1)
+            summed = np.bincount(
+                valid_vars, weights=valid_coeffs, minlength=max_var + 1
+            )
 
             # Get non-zero entries
             unique_vars = np.where(summed != 0)[0]
             unique_coeffs = summed[unique_vars]
-            
+
             # Pad to match input length
             result_vars = np.full(input_len, -1, dtype=float)
             result_coeffs = np.zeros(input_len, dtype=float)
-            
+
             n_unique = len(unique_vars)
             result_vars[:n_unique] = unique_vars
             result_coeffs[:n_unique] = unique_coeffs
 
             return np.vstack([result_vars, result_coeffs])
 
-        # Stack vars and coeffs, apply simplification once, then unstack
-        combined = xr.apply_ufunc(
+        # Coeffs and vars have dimensions (.., TERM_DIM)
+        # A row-wise operation is applied over the .. dimensions on both coeffs and vars, which are stacked together over a new "CV_DIM" dimension
+        combined: xr.DataArray = xr.apply_ufunc(
             _simplify_row,
             self.vars,
             self.coeffs,
             input_core_dims=[[TERM_DIM], [TERM_DIM]],
-            output_core_dims=[["_field", TERM_DIM]],
+            output_core_dims=[[CV_DIM, TERM_DIM]],
             vectorize=True,
         )
-        
+        # Combined has dimensions (.., CV_DIM, TERM_DIM)
+
+        # Drop terms where all vars are -1 (i.e., empty terms across all positions)
+        vars = combined.isel({CV_DIM: 0}).astype(int)
+        non_empty_terms = (vars != -1).any(dim=[d for d in vars.dims if d != TERM_DIM])
+        combined = combined.isel({TERM_DIM: non_empty_terms})
+
         # Extract vars and coeffs from the combined result
-        vars_simplified = combined.isel(_field=0).astype(int)
-        coeffs_simplified = combined.isel(_field=1)
+        vars = combined.isel({CV_DIM: 0}).astype(int)
+        coeffs = combined.isel({CV_DIM: 1})
 
         # Create new dataset with simplified data
         new_data = self.data.copy()
-        new_data = assign_multiindex_safe(
-            new_data, vars=vars_simplified, coeffs=coeffs_simplified
-        )
+        new_data = assign_multiindex_safe(new_data, vars=vars, coeffs=coeffs)
 
-        return LinearExpression(new_data, self.model).densify_terms()
+        return LinearExpression(new_data, self.model)
 
     @classmethod
     def _from_scalarexpression_list(

test/test_linear_expression.py

@@ -1195,23 +1195,29 @@ def test_cumsum(m: Model, multiple: float) -> None:
 
 def test_simplify_basic(x: Variable) -> None:
     """Test basic simplification with duplicate terms."""
-    expr = 2 * x + 3 * x
+    expr = 2 * x + 3 * x + 1 * x
     simplified = expr.simplify()
     assert simplified.nterm == 1, f"Expected 1 term, got {simplified.nterm}"
 
+    x_len = len(x.coords["dim_0"])
     # Check that the coefficient is 5
     coeffs: np.ndarray = simplified.coeffs.values
-    assert len(coeffs) == 1, f"Expected 1 valid coefficient, got {len(coeffs)}"
-    assert all(coeffs == 5.0), f"Expected coefficient 5.0, got {coeffs[0]}"
+    assert len(coeffs) == x_len, f"Expected {x_len} coefficients, got {len(coeffs)}"
+    assert all(coeffs == 6.0), f"Expected coefficient 5.0, got {coeffs[0]}"
 
 
-def test_simplify_array(x: Variable) -> None:
-    """Test simplification with array variables."""
-    # Create expression with duplicate terms
+def test_simplify_multiple_dimensions() -> None:
+    model = Model()
+    a_index = pd.Index([0, 1, 2, 3], name="a")
+    b_index = pd.Index([0, 1, 2], name="b")
+    coords = [a_index, b_index]
+    x = model.add_variables(name="x", coords=coords)
+
     expr = 2 * x + 3 * x + x
     # Simplify
     simplified = expr.simplify()
     assert simplified.nterm == 1, f"Expected 1 term, got {simplified.nterm}"
+    assert simplified.ndim == 2, f"Expected 2 dimensions, got {simplified.ndim}"
     assert all(simplified.coeffs.values == 6), (
         f"Expected coefficients of 6, got {simplified.coeffs.values}"
     )
@@ -1227,9 +1233,7 @@ def test_simplify_with_different_variables(x: Variable, y: Variable) -> None:
     # Should have 2 terms (one for x with coeff 5, one for y with coeff 4)
     assert simplified.nterm == 2, f"Expected 2 terms, got {simplified.nterm}"
 
-    coeffs: np.ndarray = simplified.coeffs.values
-    assert len(coeffs) == 2, f"Expected 2 valid coefficients, got {len(coeffs)}"
-    # Check that coefficients are 5 and 4 (in some order)
+    coeffs: list[float] = simplified.coeffs.values.flatten().tolist()
     assert set(coeffs) == {5.0, 4.0}, (
         f"Expected coefficients {{5.0, 4.0}}, got {set(coeffs)}"
     )