From 61ad751f1fad681d65cae4cab995cc04cb950f2c Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Wed, 10 Dec 2025 21:15:25 +0100 Subject: [PATCH 01/15] minor changes --- doc/release_notes.rst | 1 + linopy/common.py | 15 +++++++++++++++ test/test_constraint.py | 1 - test/test_variables.py | 23 +++++++++++++++++++++++ 4 files changed, 39 insertions(+), 1 deletion(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 5ca5ecc7..1c9ec17c 100644 --- a/doc/release_notes.rst +++ b/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 diff --git a/linopy/common.py b/linopy/common.py index f9474d3a..3f5b24f6 100644 --- a/linopy/common.py +++ b/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 @@ -124,6 +125,17 @@ 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: + try: + if isinstance(coord, xr.DataArray): + return coord.to_index() + return pd.DatetimeIndex(coord) # type: ignore + except Exception: + return coord + + def pandas_to_dataarray( arr: pd.DataFrame | pd.Series, coords: CoordsLike | None = None, @@ -164,7 +176,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): diff --git a/test/test_constraint.py b/test/test_constraint.py index 716fff40..55500f7b 100644 --- a/test/test_constraint.py +++ b/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 diff --git a/test/test_variables.py b/test/test_variables.py index 3984b091..98feffac 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -3,6 +3,9 @@ This module aims at testing the correct behavior of the Variables class. """ +import warnings +from datetime import UTC, datetime + import numpy as np import pandas as pd import pytest @@ -122,3 +125,23 @@ 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") + expr = var1 * series1 + index: pd.DatetimeIndex = expr.coords["time"].to_index() # type: ignore + assert index.equals(utc_index) + assert index.tzinfo is UTC From 3a10b12f25ad65904c1c4f02548bf384555c654d Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Wed, 10 Dec 2025 21:23:25 +0100 Subject: [PATCH 02/15] fix test --- linopy/common.py | 8 ++++++-- test/test_variables.py | 7 +++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 3f5b24f6..4ff4f15c 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -42,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. @@ -189,7 +192,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) @@ -469,7 +473,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: diff --git a/test/test_variables.py b/test/test_variables.py index 98feffac..c64adde0 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -4,7 +4,7 @@ """ import warnings -from datetime import UTC, datetime +from datetime import datetime import numpy as np import pandas as pd @@ -12,9 +12,11 @@ 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 @@ -139,9 +141,10 @@ def test_timezone_alignment_with_multiplication() -> None: series1 = pd.Series(index=utc_index, data=1.0) var1 = model.add_variables(coords=[utc_index], name="var1") - with warnings.catch_warnings(): + with warnings.catch_warnings(category=CoordAlignWarning): warnings.simplefilter("error") expr = var1 * series1 + index: pd.DatetimeIndex = expr.coords["time"].to_index() # type: ignore assert index.equals(utc_index) assert index.tzinfo is UTC From af58a0d0e60009701f733cd100927ce1c7f5f67b Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Wed, 10 Dec 2025 21:24:33 +0100 Subject: [PATCH 03/15] fix typing --- linopy/common.py | 2 +- test/test_variables.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 4ff4f15c..45b6e7b5 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -134,7 +134,7 @@ def try_to_convert_to_pd_datetime_index( try: if isinstance(coord, xr.DataArray): return coord.to_index() - return pd.DatetimeIndex(coord) # type: ignore + return pd.DatetimeIndex(coord) except Exception: return coord diff --git a/test/test_variables.py b/test/test_variables.py index c64adde0..0bfc63ae 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -145,6 +145,6 @@ def test_timezone_alignment_with_multiplication() -> None: warnings.simplefilter("error") expr = var1 * series1 - index: pd.DatetimeIndex = expr.coords["time"].to_index() # type: ignore + index: pd.DatetimeIndex = expr.coords["time"].to_index() assert index.equals(utc_index) assert index.tzinfo is UTC From b1c5ce1aeed195b63d04e34758e8d0043d92f963 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Wed, 10 Dec 2025 21:39:53 +0100 Subject: [PATCH 04/15] add 3.10 support --- test/test_variables.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/test/test_variables.py b/test/test_variables.py index 0bfc63ae..783d7d78 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -3,6 +3,7 @@ This module aims at testing the correct behavior of the Variables class. """ +import sys import warnings from datetime import datetime @@ -141,9 +142,14 @@ def test_timezone_alignment_with_multiplication() -> None: series1 = pd.Series(index=utc_index, data=1.0) var1 = model.add_variables(coords=[utc_index], name="var1") - with warnings.catch_warnings(category=CoordAlignWarning): - warnings.simplefilter("error") - expr = var1 * series1 + if sys.version_info >= (3, 11): + with warnings.catch_warnings(category=CoordAlignWarning): + warnings.simplefilter("error") + expr = var1 * series1 + else: + with warnings.catch_warnings(): + warnings.simplefilter("error") + expr = var1 * series1 index: pd.DatetimeIndex = expr.coords["time"].to_index() assert index.equals(utc_index) From 59adfafe278a7bc2601f5c0981dd564037de5997 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Wed, 10 Dec 2025 22:00:53 +0100 Subject: [PATCH 05/15] update workflow --- .github/workflows/test.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index ad501ed7..bc5fb996 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -124,4 +124,5 @@ jobs: - name: Run type checker (mypy) run: | + mypy --install-types mypy . From eba373f99f5a8b8cfdf7e60afa9ad54ef37aad8f Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Wed, 10 Dec 2025 22:02:40 +0100 Subject: [PATCH 06/15] fix types for mypy --- .github/workflows/test.yml | 1 - pyproject.toml | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index bc5fb996..ad501ed7 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -124,5 +124,4 @@ jobs: - name: Run type checker (mypy) run: | - mypy --install-types mypy . diff --git a/pyproject.toml b/pyproject.toml index b5105230..c089da13 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,6 +70,7 @@ dev = [ "types-requests", "gurobipy", "highspy", + "types-pytz" ] solvers = [ "gurobipy", From 22ca6d9f18f1e2a35c8ce41bf8f5ff5183171123 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Wed, 10 Dec 2025 22:17:30 +0100 Subject: [PATCH 07/15] fix 3.10 issue maybe --- test/test_common.py | 5 +++-- test/test_variables.py | 13 ++++--------- 2 files changed, 7 insertions(+), 11 deletions(-) diff --git a/test/test_common.py b/test/test_common.py index 0ec933bf..6d379748 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -16,6 +16,7 @@ from linopy import LinearExpression, Variable from linopy.common import ( + CoordAlignWarning, align, as_dataarray, assign_multiindex_safe, @@ -98,7 +99,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 +218,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 diff --git a/test/test_variables.py b/test/test_variables.py index 783d7d78..92a2f081 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -3,7 +3,6 @@ This module aims at testing the correct behavior of the Variables class. """ -import sys import warnings from datetime import datetime @@ -142,14 +141,10 @@ def test_timezone_alignment_with_multiplication() -> None: series1 = pd.Series(index=utc_index, data=1.0) var1 = model.add_variables(coords=[utc_index], name="var1") - if sys.version_info >= (3, 11): - with warnings.catch_warnings(category=CoordAlignWarning): - warnings.simplefilter("error") - expr = var1 * series1 - else: - with warnings.catch_warnings(): - warnings.simplefilter("error") - expr = var1 * series1 + # TODO increase coverage for datarray when coords are not dataarray + with warnings.catch_warnings(): + warnings.simplefilter("error", CoordAlignWarning) + expr = var1 * series1 index: pd.DatetimeIndex = expr.coords["time"].to_index() assert index.equals(utc_index) From a29b6b28eac309baf2cf8e08d7d0a8923e902975 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Thu, 11 Dec 2025 09:19:59 +0100 Subject: [PATCH 08/15] increase test coverage --- linopy/common.py | 2 ++ test/test_common.py | 26 +++++++++++++++++++++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index 45b6e7b5..50cafb99 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -131,6 +131,8 @@ def get_from_iterable(lst: DimsLike | None, index: int) -> Any | 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): return coord.to_index() diff --git a/test/test_common.py b/test/test_common.py index 6d379748..57b49500 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -5,12 +5,14 @@ @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 test_linear_expression import m, u, x # noqa: F401 +from pytz import UTC from xarray import DataArray from xarray.testing.assertions import assert_equal @@ -73,6 +75,28 @@ 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()) + + def test_as_dataarray_with_series_dims_subset() -> None: target_dim = "dim_0" target_index = ["a", "b", "c"] From 49856c706cf917c42e62277031ba16fe8c33a743 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Thu, 11 Dec 2025 09:36:52 +0100 Subject: [PATCH 09/15] fix test --- test/test_common.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_common.py b/test/test_common.py index 57b49500..462f9c22 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -13,6 +13,7 @@ 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 e07cceb747f286f2f359635a8e99978d68a10782 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Thu, 11 Dec 2025 11:23:58 +0100 Subject: [PATCH 10/15] improve test coverage --- linopy/common.py | 3 ++- test/test_common.py | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index 50cafb99..318cd30e 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -135,7 +135,8 @@ def try_to_convert_to_pd_datetime_index( return coord try: if isinstance(coord, xr.DataArray): - return coord.to_index() + index = coord.to_index() + assert isinstance(index, pd.DatetimeIndex) return pd.DatetimeIndex(coord) except Exception: return coord diff --git a/test/test_common.py b/test/test_common.py index 462f9c22..c95b8dc8 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -97,6 +97,14 @@ def test_as_datarray_with_tz_aware_series_index() -> None: 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" From e71fadb24a0ddd0a601ea20124522d52fcbff3cc Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Thu, 11 Dec 2025 12:01:27 +0100 Subject: [PATCH 11/15] fix test --- linopy/common.py | 1 + test/test_variables.py | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/linopy/common.py b/linopy/common.py index 318cd30e..e6176836 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -137,6 +137,7 @@ def try_to_convert_to_pd_datetime_index( if isinstance(coord, xr.DataArray): index = coord.to_index() assert isinstance(index, pd.DatetimeIndex) + return index return pd.DatetimeIndex(coord) except Exception: return coord diff --git a/test/test_variables.py b/test/test_variables.py index 92a2f081..f26d5690 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -141,7 +141,6 @@ def test_timezone_alignment_with_multiplication() -> None: series1 = pd.Series(index=utc_index, data=1.0) var1 = model.add_variables(coords=[utc_index], name="var1") - # TODO increase coverage for datarray when coords are not dataarray with warnings.catch_warnings(): warnings.simplefilter("error", CoordAlignWarning) expr = var1 * series1 From 7728270447505f7e0d564b4af9218f003ab22f26 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Sun, 14 Dec 2025 13:55:37 +0100 Subject: [PATCH 12/15] formatting --- linopy/common.py | 42 ++++++++++++++++++++++++++++++- linopy/expressions.py | 5 +++- linopy/variables.py | 2 ++ test/test_common.py | 32 ++++++++++++++++++++++- test/test_linear_expression.py | 28 +++++++++++++++++++++ test/test_quadratic_expression.py | 29 +++++++++++++++++++++ test/test_variables.py | 26 ++++++++++++++++++- 7 files changed, 160 insertions(+), 4 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index e6176836..aa15c0e9 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -12,7 +12,7 @@ from collections.abc import Callable, Generator, Hashable, Iterable, Sequence from functools import partial, reduce, wraps from pathlib import Path -from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload +from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar, overload from warnings import warn import numpy as np @@ -45,6 +45,44 @@ class CoordAlignWarning(UserWarning): ... +class TimezoneAlignError(ValueError): ... + + +P = ParamSpec("P") +R = TypeVar("R") + + +class CatchDatetimeTypeError: + """Context manager that catches datetime-related TypeErrors and re-raises as TimezoneAlignError.""" + + def __enter__(self) -> CatchDatetimeTypeError: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> bool: + if exc_type is TypeError and exc_val is not None: + if "Cannot interpret 'datetime" in str(exc_val): + raise TimezoneAlignError( + "Timezone information across datetime coordinates not aligned." + ) from exc_val + return False + + +def catch_datetime_type_error_and_re_raise(func: Callable[P, R]) -> Callable[P, R]: + """Decorator that catches datetime-related TypeErrors and re-raises as TimezoneAlignError.""" + + @wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + with CatchDatetimeTypeError(): + return func(*args, **kwargs) + + return wrapper + + def set_int_index(series: pd.Series) -> pd.Series: """ Convert string index to int index. @@ -468,6 +506,7 @@ def group_terms_polars(df: pl.DataFrame) -> pl.DataFrame: return df +@catch_datetime_type_error_and_re_raise def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset: """ Join multiple xarray Dataarray's to a Dataset and warn if coordinates are not equal. @@ -485,6 +524,7 @@ def save_join(*dataarrays: DataArray, integer_dtype: bool = False) -> Dataset: return Dataset({ds.name: ds for ds in arrs}) +@catch_datetime_type_error_and_re_raise def assign_multiindex_safe(ds: Dataset, **fields: Any) -> Dataset: """ Assign a field to a xarray Dataset while being safe against warnings about multiindex corruption. diff --git a/linopy/expressions.py b/linopy/expressions.py index d60c8be5..61945d8c 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -47,6 +47,7 @@ LocIndexer, as_dataarray, assign_multiindex_safe, + catch_datetime_type_error_and_re_raise, check_common_keys_values, check_has_nulls, check_has_nulls_polars, @@ -498,6 +499,7 @@ def __neg__(self: GenericExpression) -> GenericExpression: """ return self.assign_multiindex_safe(coeffs=-self.coeffs, const=-self.const) + @catch_datetime_type_error_and_re_raise def _multiply_by_linear_expression( self, other: LinearExpression | ScalarLinearExpression ) -> QuadraticExpression: @@ -519,6 +521,7 @@ def _multiply_by_linear_expression( res = res + self.reset_const() * other.const return res + @catch_datetime_type_error_and_re_raise def _multiply_by_constant( self: GenericExpression, other: ConstantLike ) -> GenericExpression: @@ -1361,7 +1364,7 @@ def __mul__( return self._multiply_by_linear_expression(other) else: return self._multiply_by_constant(other) - except TypeError: + except AssertionError: return NotImplemented def __pow__(self, other: int) -> QuadraticExpression: diff --git a/linopy/variables.py b/linopy/variables.py index 396f165f..230b3413 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -34,6 +34,7 @@ LocIndexer, as_dataarray, assign_multiindex_safe, + catch_datetime_type_error_and_re_raise, check_has_nulls, check_has_nulls_polars, filter_nulls_polars, @@ -287,6 +288,7 @@ def loc(self) -> LocIndexer: def to_pandas(self) -> pd.Series: return self.labels.to_pandas() + @catch_datetime_type_error_and_re_raise def to_linexpr( self, coefficient: ConstantLike = 1, diff --git a/test/test_common.py b/test/test_common.py index c95b8dc8..00500084 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -13,7 +13,6 @@ 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 @@ -106,6 +105,37 @@ def test_as_datarray_with_tz_aware_series_index() -> None: assert time_index.equals(result.coords["time"].to_index()) +def test_as_datarray_with_tz_aware_dataframe_columns_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]) + + index = pd.Index([0, 1, 2, 3], name="x") + pandas_df = pd.DataFrame(index=index, columns=time_index, data=1.0) + + data_array = xr.DataArray(data=[0, 1, 2, 3], coords=[time_index]) + result = as_dataarray(arr=pandas_df, 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=pandas_df, coords=data_array.coords) + assert time_index.equals(result.coords["time"].to_index()) + + coords = {"time": time_index} + result = as_dataarray(arr=pandas_df, coords=coords) + assert time_index.equals(result.coords["time"].to_index()) + + coords = {"time": [0, 1, 2, 3]} + result = as_dataarray(arr=pandas_df, 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"] diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 2551c203..2a8a3fa0 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -7,14 +7,18 @@ from __future__ import annotations +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 xarray.testing import assert_equal from linopy import LinearExpression, Model, QuadraticExpression, Variable, merge +from linopy.common import TimezoneAlignError from linopy.constants import HELPER_DIMS, TERM_DIM from linopy.expressions import ScalarLinearExpression from linopy.testing import assert_linequal, assert_quadequal @@ -1191,3 +1195,27 @@ def test_cumsum(m: Model, multiple: float) -> None: expr = m.variables["x"] + m.variables["y"] cumsum = (multiple * expr).cumsum() cumsum.nterm == 2 + + +def test_timezone_alignment_failure() -> None: + utc_index = pd.date_range( + start=datetime(2025, 1, 1), + freq="15min", + periods=4, + tz=UTC, + name="time", + ) + tz_naive_index = pd.date_range( + start=datetime(2025, 1, 1), + freq="15min", + periods=4, + tz=None, + name="time", + ) + model = Model() + series1 = pd.Series(index=tz_naive_index, data=1.0) + expr = model.add_variables(coords=[utc_index], name="var1") * 1.0 + + with pytest.raises(TimezoneAlignError): + # We expect to get a useful error (TimezoneAlignError) instead of a not implemented error falsely claiming that we cannot multiply these types together + _ = expr * series1 diff --git a/test/test_quadratic_expression.py b/test/test_quadratic_expression.py index f5f86c35..5bd0457f 100644 --- a/test/test_quadratic_expression.py +++ b/test/test_quadratic_expression.py @@ -1,13 +1,17 @@ #!/usr/bin/env python3 +from datetime import datetime + import numpy as np import pandas as pd import polars as pl import pytest +from pytz import UTC from scipy.sparse import csc_matrix from xarray import DataArray from linopy import Model, Variable, merge +from linopy.common import TimezoneAlignError from linopy.constants import FACTOR_DIM, TERM_DIM from linopy.expressions import LinearExpression, QuadraticExpression from linopy.testing import assert_quadequal @@ -344,3 +348,28 @@ def test_power_of_three(x: Variable) -> None: x**3 with pytest.raises(TypeError): (x * x) * (x * x) + + +def test_timezone_alignment_failure() -> None: + utc_index = pd.date_range( + start=datetime(2025, 1, 1), + freq="15min", + periods=4, + tz=UTC, + name="time", + ) + tz_naive_index = pd.date_range( + start=datetime(2025, 1, 1), + freq="15min", + periods=4, + tz=None, + name="time", + ) + model = Model() + series1 = pd.Series(index=tz_naive_index, data=1.0) + var = model.add_variables(coords=[utc_index], name="var1") + expr = var * var + + with pytest.raises(TimezoneAlignError): + # We expect to get a useful error (TimezoneAlignError) instead of a not implemented error falsely claiming that we cannot multiply these types together + _ = expr * series1 diff --git a/test/test_variables.py b/test/test_variables.py index f26d5690..aed5219c 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -16,7 +16,7 @@ import linopy from linopy import Model -from linopy.common import CoordAlignWarning +from linopy.common import CoordAlignWarning, TimezoneAlignError from linopy.testing import assert_varequal from linopy.variables import ScalarVariable @@ -148,3 +148,27 @@ def test_timezone_alignment_with_multiplication() -> None: index: pd.DatetimeIndex = expr.coords["time"].to_index() assert index.equals(utc_index) assert index.tzinfo is UTC + + +def test_timezone_alignment_failure() -> None: + utc_index = pd.date_range( + start=datetime(2025, 1, 1), + freq="15min", + periods=4, + tz=UTC, + name="time", + ) + tz_naive_index = pd.date_range( + start=datetime(2025, 1, 1), + freq="15min", + periods=4, + tz=None, + name="time", + ) + model = Model() + series1 = pd.Series(index=tz_naive_index, data=1.0) + var1 = model.add_variables(coords=[utc_index], name="var1") + + with pytest.raises(TimezoneAlignError): + # We expect to get a useful error (TimezoneAlignError) instead of a not implemented error falsely claiming that we cannot multiply these types together + _ = var1 * series1 From 1d4ca3b5652b146e4986623ba319c44982575746 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Sun, 14 Dec 2025 14:02:11 +0100 Subject: [PATCH 13/15] fixed typing --- linopy/common.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index aa15c0e9..dd6a19cd 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -12,7 +12,7 @@ from collections.abc import Callable, Generator, Hashable, Iterable, Sequence from functools import partial, reduce, wraps from pathlib import Path -from typing import TYPE_CHECKING, Any, Generic, ParamSpec, TypeVar, overload +from typing import TYPE_CHECKING, Any, Generic, Literal, ParamSpec, TypeVar, overload from warnings import warn import numpy as np @@ -63,7 +63,7 @@ def __exit__( exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any, - ) -> bool: + ) -> Literal[False]: if exc_type is TypeError and exc_val is not None: if "Cannot interpret 'datetime" in str(exc_val): raise TimezoneAlignError( @@ -78,7 +78,8 @@ def catch_datetime_type_error_and_re_raise(func: Callable[P, R]) -> Callable[P, @wraps(func) def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: with CatchDatetimeTypeError(): - return func(*args, **kwargs) + result = func(*args, **kwargs) + return result return wrapper From 2d67883bf2e538a26e9fedb6bc61980a20a00352 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Sun, 14 Dec 2025 14:03:32 +0100 Subject: [PATCH 14/15] fix tests --- linopy/expressions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index 61945d8c..b3360bc3 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -1364,7 +1364,7 @@ def __mul__( return self._multiply_by_linear_expression(other) else: return self._multiply_by_constant(other) - except AssertionError: + except TypeError: return NotImplemented def __pow__(self, other: int) -> QuadraticExpression: From bd5f7b428c4b9a3c0eeedd54dbc759d7e48b0853 Mon Sep 17 00:00:00 2001 From: Robbie Muir Date: Sun, 14 Dec 2025 14:26:15 +0100 Subject: [PATCH 15/15] formatting --- test/test_common.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_common.py b/test/test_common.py index 00500084..76b7a492 100644 --- a/test/test_common.py +++ b/test/test_common.py @@ -13,6 +13,7 @@ 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