diff --git a/doc/release_notes.rst b/doc/release_notes.rst index aa05eda8..079330eb 100644 --- a/doc/release_notes.rst +++ b/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 diff --git a/linopy/common.py b/linopy/common.py index 6804ac1e..9d530cc2 100644 --- a/linopy/common.py +++ b/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( diff --git a/linopy/expressions.py b/linopy/expressions.py index cff9b0cd..d60c8be5 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -98,6 +98,7 @@ pd.Series, pd.DataFrame, np.ndarray, + pl.Series, ) diff --git a/linopy/types.py b/linopy/types.py index d5276dbb..68e5c307 100644 --- a/linopy/types.py +++ b/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"] diff --git a/linopy/variables.py b/linopy/variables.py index 2a929515..0013ff75 100644 --- a/linopy/variables.py +++ b/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. diff --git a/test/test_common.py b/test/test_common.py index 85059487..0ec933bf 100644 --- a/test/test_common.py +++ b/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] diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 438dcab5..2551c203 100644 --- a/test/test_linear_expression.py +++ b/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)