Feature/add polars input (#511)

* feat: add polars series for constant input

* add another test

* update doc

---------

Co-authored-by: Fabian Hofmann <[email protected]>

Author: RobbieKiwi

doc/release_notes.rst

@@ -9,6 +9,7 @@ Upcoming Version
 * Improved constraint equality check in `linopy.testing.assert_conequal` to less strict optionally
 * Minor bugfix for multiplying variables with numpy type constants
 * Harmonize dtypes before concatenation in lp file writing to avoid dtype mismatch errors. This error occurred when creating and storing models in netcdf format using windows machines and loading and solving them on linux machines.
+* Add option to use polars series as constant input
 * Fix expression merge to explicitly use outer join when combining expressions with disjoint coordinates for consistent behavior across xarray versions
 
 Version 0.5.6

linopy/common.py

@@ -254,6 +254,8 @@ def as_dataarray(
         arr = pandas_to_dataarray(arr, coords=coords, dims=dims, **kwargs)
     elif isinstance(arr, np.ndarray):
         arr = numpy_to_dataarray(arr, coords=coords, dims=dims, **kwargs)
+    elif isinstance(arr, pl.Series):
+        arr = numpy_to_dataarray(arr.to_numpy(), coords=coords, dims=dims, **kwargs)
     elif isinstance(arr, np.number):
         arr = DataArray(float(arr), coords=coords, dims=dims, **kwargs)
     elif isinstance(arr, int | float | str | bool | list):
@@ -269,6 +271,7 @@ def as_dataarray(
             pd.DataFrame,
             np.ndarray,
             DataArray,
+            pl.Series,
         ]
         supported_types_str = ", ".join([t.__name__ for t in supported_types])
         raise TypeError(

linopy/expressions.py

@@ -98,6 +98,7 @@
     pd.Series,
     pd.DataFrame,
     np.ndarray,
+    pl.Series,
 )
 
 

linopy/types.py

@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING, Union
 
 import numpy
-import numpy.typing
+import polars as pl
 from pandas import DataFrame, Index, Series
 from xarray import DataArray
 from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates
@@ -37,6 +37,7 @@
     DataArray,
     Series,
     DataFrame,
+    pl.Series,
 ]
 SignLike = Union[str, numpy.ndarray, DataArray, Series, DataFrame]  # noqa: UP007
 VariableLike = Union["ScalarVariable", "Variable"]

linopy/variables.py

@@ -287,12 +287,7 @@ def to_pandas(self) -> pd.Series:
 
     def to_linexpr(
         self,
-        coefficient: int
-        | float
-        | pd.Series
-        | pd.DataFrame
-        | np.ndarray
-        | DataArray = 1,
+        coefficient: ConstantLike = 1,
     ) -> expressions.LinearExpression:
         """
         Create a linear expression from the variables.

test/test_common.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import pandas as pd
+import polars as pl
 import pytest
 import xarray as xr
 from test_linear_expression import m, u, x  # noqa: F401
@@ -120,6 +121,16 @@ def test_as_dataarray_with_series_aligned_coords() -> None:
     assert list(da.coords[target_dim].values) == target_index
 
 
+def test_as_dataarray_with_pl_series_dims_default() -> None:
+    target_dim = "dim_0"
+    target_index = [0, 1, 2]
+    s = pl.Series([1, 2, 3])
+    da = as_dataarray(s)
+    assert isinstance(da, DataArray)
+    assert da.dims == (target_dim,)
+    assert list(da.coords[target_dim].values) == target_index
+
+
 def test_as_dataarray_dataframe_dims_default() -> None:
     target_dims = ("dim_0", "dim_1")
     target_index = [0, 1]

test/test_linear_expression.py

@@ -111,6 +111,13 @@ def test_linexpr_from_constant_dataarray(m: Model) -> None:
     assert expr.nterm == 0
 
 
+def test_linexpr_from_constant_pl_series(m: Model) -> None:
+    const = pl.Series([1, 2])
+    expr = LinearExpression(const, m)
+    assert (expr.const == const.to_numpy()).all()
+    assert expr.nterm == 0
+
+
 def test_linexpr_from_constant_pandas_series(m: Model) -> None:
     const = pd.Series([1, 2], index=pd.RangeIndex(2, name="dim_0"))
     expr = LinearExpression(const, m)