add to_polars improvements

Author: Robbie Muir

linopy/expressions.py

@@ -441,6 +441,11 @@ def __repr__(self) -> str:
 
         return "\n".join(lines)
 
+    @property
+    def is_constant(self) -> bool:
+        """This is true if the expression contains no variables"""
+        return self.data["coeffs"].values.size == 0
+
     def print(self, display_max_rows: int = 20, display_max_terms: int = 20) -> None:
         """
         Print the linear expression.
@@ -1439,12 +1444,18 @@ def to_polars(self) -> pl.DataFrame:
 
         The resulting DataFrame represents a long table format of the all
         non-masked expressions with non-zero coefficients. It contains the
-        columns `coeffs`, `vars`.
+        columns `coeffs`, `vars`, `const`. The coeffs and vars columns will be null if the expression is constant.
 
         Returns
         -------
         df : polars.DataFrame
         """
+        if self.is_constant:
+            df = pl.DataFrame(
+                {"const": self.data["const"].values.reshape(-1)}
+            ).with_columns(pl.lit(None).alias("coeffs"), pl.lit(None).alias("vars"))
+            return df.select(["vars", "coeffs", "const"])
+
         df = to_polars(self.data)
         df = filter_nulls_polars(df)
         df = group_terms_polars(df)
@@ -1855,12 +1866,22 @@ def to_polars(self, **kwargs: Any) -> pl.DataFrame:
 
         The resulting DataFrame represents a long table format of the all
         non-masked expressions with non-zero coefficients. It contains the
-        columns `coeffs`, `vars`.
+        columns `vars1`, `vars2`, `coeffs`, `const`. If the expression is constant, the `vars1` and `vars2` and `coeffs` columns will be null.
 
         Returns
         -------
         df : polars.DataFrame
         """
+        if self.is_constant:
+            df = pl.DataFrame(
+                {"const": self.data["const"].values.reshape(-1)}
+            ).with_columns(
+                pl.lit(None).alias("coeffs"),
+                pl.lit(None).alias("vars1"),
+                pl.lit(None).alias("vars2"),
+            )
+            return df.select(["vars1", "vars2", "coeffs", "const"])
+
         vars = self.data.vars.assign_coords(
             {FACTOR_DIM: ["vars1", "vars2"]}
         ).to_dataset(FACTOR_DIM)

test/test_linear_expression.py

@@ -1125,17 +1125,45 @@ def test_linear_expression_from_tuples_bad_calls(
 
 def test_linear_expression_from_constant_scalar(m: Model) -> None:
     expr = LinearExpression.from_constant(model=m, constant=10)
+    assert expr.is_constant
     assert isinstance(expr, LinearExpression)
     assert (expr.const == 10).all()
 
 
-def test_linear_expression_from_constant_array(m: Model) -> None:
+def test_linear_expression_from_constant_1D(m: Model) -> None:
     arr = pd.Series(index=pd.Index([0, 1], name="t"), data=[10, 20])
     expr = LinearExpression.from_constant(model=m, constant=arr)
     assert isinstance(expr, LinearExpression)
     assert list(expr.coords.keys())[0] == "t"
     assert expr.nterm == 0
     assert (expr.const.values == [10, 20]).all()
+    assert expr.is_constant
+
+    exp_polars = expr.to_polars()
+    assert exp_polars.columns == ["const", "coeffs", "vars"]
+    assert exp_polars["const"].to_list() == [10, 20]
+    assert exp_polars["coeffs"].to_list() == [None, None]
+    assert exp_polars["vars"].to_list() == [None, None]
+
+
+def test_linear_expression_two_dimensional_from_constant_2D(m: Model) -> None:
+    index_a = pd.Index([0, 1], name="a")
+    index_b = pd.Index([0, 1, 2], name="b")
+    arr = np.array([[10, 20, 30], [40, 50, 60]])
+    const = xr.DataArray(data=arr, coords=[index_a, index_b])
+
+    le_variable = m.add_variables(name="var", coords=[index_a, index_b]) * 1 + const
+    assert not le_variable.is_constant
+    le_const = LinearExpression.from_constant(model=m, constant=const)
+    assert le_const.is_constant
+
+    var_pol = le_variable.to_polars()
+    const_pol = le_const.to_polars()
+    assert var_pol.shape == const_pol.shape
+    assert var_pol.columns == const_pol.columns
+    assert all(const_pol["const"] == var_pol["const"])
+    assert all(const_pol["coeffs"].is_null())
+    assert all(const_pol["vars"].is_null())
 
 
 def test_linear_expression_sanitize(x: Variable, y: Variable, z: Variable) -> None:

test/test_quadratic_expression.py

@@ -287,7 +287,7 @@ def test_quadratic_expression_flat(x: Variable, y: Variable) -> None:
     assert len(expr.flat) == 2
 
 
-def test_linear_expression_to_polars(x: Variable, y: Variable) -> None:
+def test_quadratic_expression_to_polars(x: Variable, y: Variable) -> None:
     expr = x * y + x + 5
     df = expr.to_polars()
     assert isinstance(df, pl.DataFrame)
@@ -296,6 +296,22 @@ def test_linear_expression_to_polars(x: Variable, y: Variable) -> None:
     assert len(df) == expr.nterm * 2
 
 
+def test_quadratic_expression_constant_to_polars() -> None:
+    m = Model()
+    arr = pd.Series(index=pd.Index([0, 1], name="t"), data=[10, 20])
+    lin_expr = LinearExpression.from_constant(model=m, constant=arr)
+    quad_expr = lin_expr.to_quadexpr()
+
+    assert quad_expr.is_constant
+    df = quad_expr.to_polars()
+    assert isinstance(df, pl.DataFrame)
+    assert df.columns == ["vars1", "vars2", "coeffs", "const"]
+    assert all(df["vars1"].is_null())
+    assert all(df["vars2"].is_null())
+    assert all(df["coeffs"].is_null())
+    assert all(arr.to_numpy() == df["const"].to_numpy())
+
+
 def test_quadratic_expression_to_matrix(model: Model, x: Variable, y: Variable) -> None:
     expr: QuadraticExpression = x * y + x + 5  # type: ignore