Fix warning due to timezone aware pandas index

doc/release_notes.rst

@@ -3,6 +3,7 @@ Release Notes
 
 .. Upcoming Version
 
+* Fix warning when multiplying variables with pd.Series containing time-zone aware index
 * 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/common.py

@@ -18,6 +18,7 @@
 import numpy as np
 import pandas as pd
 import polars as pl
+import xarray as xr
 from numpy import arange, signedinteger
 from xarray import DataArray, Dataset, apply_ufunc, broadcast
 from xarray import align as xr_align
@@ -41,6 +42,9 @@
     from linopy.variables import Variable
 
 
+class CoordAlignWarning(UserWarning): ...
+
+
 def set_int_index(series: pd.Series) -> pd.Series:
     """
     Convert string index to int index.
@@ -124,6 +128,21 @@ def get_from_iterable(lst: DimsLike | None, index: int) -> Any | None:
     return lst[index] if 0 <= index < len(lst) else None
 
 
+def try_to_convert_to_pd_datetime_index(
+    coord: xr.DataArray | Sequence | pd.Index | Any,
+) -> pd.DatetimeIndex | xr.DataArray | Sequence | pd.Index | Any:
+    if isinstance(coord, pd.DatetimeIndex):
+        return coord
+    try:
+        if isinstance(coord, xr.DataArray):
+            index = coord.to_index()
+            assert isinstance(index, pd.DatetimeIndex)
+            return index
+        return pd.DatetimeIndex(coord)
+    except Exception:
+        return coord
+
+
 def pandas_to_dataarray(
     arr: pd.DataFrame | pd.Series,
     coords: CoordsLike | None = None,
@@ -164,7 +183,10 @@ def pandas_to_dataarray(
         shared_dims = set(pandas_coords.keys()) & set(coords.keys())
         non_aligned = []
         for dim in shared_dims:
+            pd_coord = pandas_coords[dim]
             coord = coords[dim]
+            if isinstance(pd_coord, pd.DatetimeIndex):
+                coord = try_to_convert_to_pd_datetime_index(coord)
             if not isinstance(coord, pd.Index):
                 coord = pd.Index(coord)
             if not pandas_coords[dim].equals(coord):
@@ -174,7 +196,8 @@ def pandas_to_dataarray(
                 f"coords for dimension(s) {non_aligned} is not aligned with the pandas object. "
                 "Previously, the indexes of the pandas were ignored and overwritten in "
                 "these cases. Now, the pandas object's coordinates are taken considered"
-                " for alignment."
+                " for alignment.",
+                CoordAlignWarning,
             )
 
     return DataArray(arr, coords=None, dims=dims, **kwargs)
@@ -454,7 +477,7 @@ def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset:
     except ValueError:
         warn(
             "Coordinates across variables not equal. Perform outer join.",
-            UserWarning,
+            CoordAlignWarning,
         )
         arrs = xr_align(*dataarrays, join="outer")
         if integer_dtype:

pyproject.toml

@@ -70,6 +70,7 @@ dev = [
     "types-requests",
     "gurobipy",
     "highspy",
+    "types-pytz"
 ]
 solvers = [
     "gurobipy",

test/test_common.py

@@ -5,17 +5,21 @@
 @author: fabian
 """
 
+from datetime import datetime
+
 import numpy as np
 import pandas as pd
 import polars as pl
 import pytest
 import xarray as xr
+from pytz import UTC
 from test_linear_expression import m, u, x  # noqa: F401
 from xarray import DataArray
 from xarray.testing.assertions import assert_equal
 
 from linopy import LinearExpression, Variable
 from linopy.common import (
+    CoordAlignWarning,
     align,
     as_dataarray,
     assign_multiindex_safe,
@@ -72,6 +76,36 @@ def test_as_dataarray_with_series_dims_priority() -> None:
     assert list(da.coords[target_dim].values) == target_index
 
 
+def test_as_datarray_with_tz_aware_series_index() -> None:
+    time_index = pd.date_range(
+        start=datetime(2025, 1, 1),
+        freq="15min",
+        periods=4,
+        tz=UTC,
+        name="time",
+    )
+    other_index = pd.Index(name="time", data=[0, 1, 2, 3])
+
+    panda_series = pd.Series(index=time_index, data=1.0)
+
+    data_array = xr.DataArray(data=[0, 1, 2, 3], coords=[time_index])
+    result = as_dataarray(arr=panda_series, coords=data_array.coords)
+    assert time_index.equals(result.coords["time"].to_index())
+
+    data_array = xr.DataArray(data=[0, 1, 2, 3], coords=[other_index])
+    with pytest.warns(CoordAlignWarning):
+        result = as_dataarray(arr=panda_series, coords=data_array.coords)
+    assert time_index.equals(result.coords["time"].to_index())
+
+    coords = {"time": time_index}
+    result = as_dataarray(arr=panda_series, coords=coords)
+    assert time_index.equals(result.coords["time"].to_index())
+
+    coords = {"time": [0, 1, 2, 3]}
+    result = as_dataarray(arr=panda_series, coords=coords)
+    assert time_index.equals(result.coords["time"].to_index())
+
+
 def test_as_dataarray_with_series_dims_subset() -> None:
     target_dim = "dim_0"
     target_index = ["a", "b", "c"]
@@ -98,7 +132,7 @@ def test_as_dataarray_with_series_override_coords() -> None:
     target_dim = "dim_0"
     target_index = ["a", "b", "c"]
     s = pd.Series([1, 2, 3], index=target_index)
-    with pytest.warns(UserWarning):
+    with pytest.warns(CoordAlignWarning):
         da = as_dataarray(s, coords=[[1, 2, 3]])
     assert isinstance(da, DataArray)
     assert da.dims == (target_dim,)
@@ -217,7 +251,7 @@ def test_as_dataarray_dataframe_override_coords() -> None:
     target_index = ["a", "b"]
     target_columns = ["A", "B"]
     df = pd.DataFrame([[1, 2], [3, 4]], index=target_index, columns=target_columns)
-    with pytest.warns(UserWarning):
+    with pytest.warns(CoordAlignWarning):
         da = as_dataarray(df, coords=[[1, 2], [2, 3]])
     assert isinstance(da, DataArray)
     assert da.dims == target_dims

test/test_constraint.py

@@ -12,7 +12,6 @@
 import polars as pl
 import pytest
 import xarray as xr
-import xarray.core
 from xarray.testing import assert_equal
 
 import linopy

test/test_variables.py

@@ -3,15 +3,20 @@
 This module aims at testing the correct behavior of the Variables class.
 """
 
+import warnings
+from datetime import datetime
+
 import numpy as np
 import pandas as pd
 import pytest
 import xarray as xr
 import xarray.core.indexes
 import xarray.core.utils
+from pytz import UTC
 
 import linopy
 from linopy import Model
+from linopy.common import CoordAlignWarning
 from linopy.testing import assert_varequal
 from linopy.variables import ScalarVariable
 
@@ -122,3 +127,24 @@ def test_scalar_variable(m: Model) -> None:
     x = ScalarVariable(label=0, model=m)
     assert isinstance(x, ScalarVariable)
     assert x.__rmul__(x) is NotImplemented  # type: ignore
+
+
+def test_timezone_alignment_with_multiplication() -> None:
+    utc_index = pd.date_range(
+        start=datetime(2025, 1, 1),
+        freq="15min",
+        periods=4,
+        tz=UTC,
+        name="time",
+    )
+    model = Model()
+    series1 = pd.Series(index=utc_index, data=1.0)
+    var1 = model.add_variables(coords=[utc_index], name="var1")
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", CoordAlignWarning)
+        expr = var1 * series1
+
+    index: pd.DatetimeIndex = expr.coords["time"].to_index()
+    assert index.equals(utc_index)
+    assert index.tzinfo is UTC