fix: xarray new rolling namespace (#432)

* fix: xarray new namespace
fix: assign copt as solver only if env can be solved

* fix: ruff complaining about unused variables

* fix: remove doc test from hard imported xarray doc

* fix type annotations

* fix type annotations

* follow up

* fix: version comparison for highspy

* follow up

Author: FabianHofmann

linopy/common.py

@@ -19,7 +19,6 @@
 import pandas as pd
 import polars as pl
 from numpy import arange, signedinteger
-from pandas.util._decorators import doc
 from xarray import DataArray, Dataset, apply_ufunc, broadcast
 from xarray import align as xr_align
 from xarray.core import dtypes, indexing
@@ -990,20 +989,73 @@ def check_common_keys_values(list_of_dicts: list[dict[str, Any]]) -> bool:
     return all(len({d[k] for d in list_of_dicts if k in d}) == 1 for k in common_keys)
 
 
-@doc(xr_align)
 def align(
     *objects: LinearExpression | Variable | T_Alignable,
     join: JoinOptions = "inner",
     copy: bool = True,
-    indexes=None,
+    indexes: Any = None,
     exclude: str | Iterable[Hashable] = frozenset(),
-    fill_value=dtypes.NA,
+    fill_value: Any = dtypes.NA,
 ) -> tuple[LinearExpression | Variable | T_Alignable, ...]:
+    """
+    Given any number of Variables, Expressions, Dataset and/or DataArray objects,
+    returns new objects with aligned indexes and dimension sizes.
+
+    Array from the aligned objects are suitable as input to mathematical
+    operators, because along each dimension they have the same index and size.
+
+    Missing values (if ``join != 'inner'``) are filled with ``fill_value``.
+    The default fill value is NaN.
+
+    This functions essentially wraps the xarray function
+    :py:func:`xarray.align`.
+
+    Parameters
+    ----------
+    *objects : Variable, LinearExpression, Dataset or DataArray
+        Objects to align.
+    join : {"outer", "inner", "left", "right", "exact", "override"}, optional
+        Method for joining the indexes of the passed objects along each
+        dimension:
+
+        - "outer": use the union of object indexes
+        - "inner": use the intersection of object indexes
+        - "left": use indexes from the first object with each dimension
+        - "right": use indexes from the last object with each dimension
+        - "exact": instead of aligning, raise `ValueError` when indexes to be
+        aligned are not equal
+        - "override": if indexes are of same size, rewrite indexes to be
+        those of the first object with that dimension. Indexes for the same
+        dimension must have the same size in all objects.
+
+    copy : bool, default: True
+        If ``copy=True``, data in the return values is always copied. If
+        ``copy=False`` and reindexing is unnecessary, or can be performed with
+        only slice operations, then the output may share memory with the input.
+        In either case, new xarray objects are always returned.
+    indexes : dict-like, optional
+        Any indexes explicitly provided with the `indexes` argument should be
+        used in preference to the aligned indexes.
+    exclude : str, iterable of hashable or None, optional
+        Dimensions that must be excluded from alignment
+    fill_value : scalar or dict-like, optional
+        Value to use for newly missing values. If a dict-like, maps
+        variable names to fill values. Use a data array's name to
+        refer to its values.
+
+    Returns
+    -------
+    aligned : tuple of DataArray or Dataset
+        Tuple of objects with the same type as `*objects` with aligned
+        coordinates.
+
+
+    """
     from linopy.expressions import LinearExpression
     from linopy.variables import Variable
 
-    finisher = []
-    das = []
+    finisher: list[partial[Any] | Callable[[Any], Any]] = []
+    das: list[Any] = []
     for obj in objects:
         if isinstance(obj, LinearExpression):
             finisher.append(partial(obj.__class__, model=obj.model))

linopy/expressions.py

@@ -24,17 +24,23 @@
 import scipy
 import xarray as xr
 import xarray.core.groupby
-import xarray.core.rolling
 from numpy import array, nan, ndarray
 from pandas.core.frame import DataFrame
 from pandas.core.series import Series
 from scipy.sparse import csc_matrix
 from xarray import Coordinates, DataArray, Dataset, IndexVariable
 from xarray.core.coordinates import DataArrayCoordinates, DatasetCoordinates
 from xarray.core.indexes import Indexes
-from xarray.core.rolling import DatasetRolling
 from xarray.core.utils import Frozen
 
+try:
+    # resolve breaking change in xarray 2025.03.0
+    import xarray.computation.rolling
+    from xarray.computation.rolling import DatasetRolling
+except ImportError:
+    import xarray.core.rolling
+    from xarray.core.rolling import DatasetRolling  # type: ignore
+
 from linopy import constraints, variables
 from linopy.common import (
     EmptyDeprecationWrapper,

linopy/solvers.py

@@ -20,6 +20,7 @@
 
 import numpy as np
 import pandas as pd
+from packaging.version import parse as parse_version
 
 from linopy.constants import (
     Result,
@@ -62,7 +63,8 @@
     available_solvers.append("highs")
     from importlib.metadata import version
 
-    if version("highspy") < "1.7.1":
+    if parse_version(version("highspy")) < parse_version("1.7.1"):
+        # Fallback if parse_version is not available or version string is invalid
         _new_highspy_mps_layout = False
     else:
         _new_highspy_mps_layout = True
@@ -111,7 +113,11 @@
 with contextlib.suppress(ModuleNotFoundError):
     import coptpy
 
-    available_solvers.append("copt")
+    try:
+        coptpy.Envr()
+        available_solvers.append("copt")
+    except coptpy.CoptError:
+        pass
 
 quadratic_solvers = [s for s in QUADRATIC_SOLVERS if s in available_solvers]
 logger = logging.getLogger(__name__)

test/test_common.py

@@ -13,7 +13,7 @@
 from xarray import DataArray
 from xarray.testing.assertions import assert_equal
 
-from linopy import Model, Variable
+from linopy import LinearExpression, Variable
 from linopy.common import (
     align,
     as_dataarray,
@@ -651,7 +651,7 @@ def test_get_dims_with_index_levels() -> None:
     assert get_dims_with_index_levels(ds5) == []
 
 
-def test_align(m: Model, x: Variable, u: Variable) -> None:
+def test_align(x: Variable, u: Variable) -> None:  # noqa: F811
     alpha = xr.DataArray([1, 2], [[1, 2]])
     beta = xr.DataArray(
         [1, 2, 3],
@@ -667,18 +667,21 @@ def test_align(m: Model, x: Variable, u: Variable) -> None:
 
     # inner join
     x_obs, alpha_obs = align(x, alpha)
+    assert isinstance(x_obs, Variable)
     assert x_obs.shape == alpha_obs.shape == (1,)
     assert_varequal(x_obs, x.loc[[1]])
 
     # left-join
     x_obs, alpha_obs = align(x, alpha, join="left")
     assert x_obs.shape == alpha_obs.shape == (2,)
+    assert isinstance(x_obs, Variable)
     assert_varequal(x_obs, x)
     assert_equal(alpha_obs, DataArray([np.nan, 1], [[0, 1]]))
 
     # multiindex
     beta_obs, u_obs = align(beta, u)
     assert u_obs.shape == beta_obs.shape == (2,)
+    assert isinstance(u_obs, Variable)
     assert_varequal(u_obs, u.loc[[(1, "b"), (2, "b")]])
     assert_equal(beta_obs, beta.loc[[(1, "b"), (2, "b")]])
 
@@ -687,4 +690,5 @@ def test_align(m: Model, x: Variable, u: Variable) -> None:
     x_obs, expr_obs, alpha_obs = align(x, expr, alpha)
     assert x_obs.shape == alpha_obs.shape == (1,)
     assert expr_obs.shape == (1, 1)  # _term dim
+    assert isinstance(expr_obs, LinearExpression)
     assert_linequal(expr_obs, expr.loc[[1]])