address deprecations and issue with multiindex deserialization

linopy/constraints.py

@@ -153,7 +153,7 @@ def name(self):
 
     @property
     def coord_dims(self):
-        return {k: self.data.dims[k] for k in self.dims if k not in HELPER_DIMS}
+        return {k for k in self.dims if k not in HELPER_DIMS}
 
     @property
     def is_assigned(self):

linopy/expressions.py

@@ -119,7 +119,6 @@ def map(self, func, shortcut=False, args=(), **kwargs):
         LinearExpression
             The result of applying the function to the groupby object.
         """
-
         return LinearExpression(
             self.groupby.map(func, shortcut=shortcut, args=args, **kwargs), self.model
         )
@@ -191,7 +190,7 @@ def func(ds):
             ds = ds.assign_coords({TERM_DIM: np.arange(len(ds._term))})
             return ds
 
-        return self.map(func, **kwargs)
+        return self.map(func, **kwargs, shortcut=True)
 
     def roll(self, **kwargs):
         """
@@ -321,7 +320,7 @@ def __init__(self, data, model):
         # ensure helper dimensions are not set as coordinates
         if drop_dims := set(HELPER_DIMS).intersection(data.coords):
             # TODO: add a warning here, routines should be safe against this
-            data = data.drop(drop_dims)
+            data = data.drop_vars(drop_dims)
 
         if not isinstance(model, Model):
             raise ValueError("model must be an instance of linopy.Model")
@@ -336,7 +335,7 @@ def __repr__(self):
         max_lines = options["display_max_rows"]
         dims = list(self.coord_dims)
         ndim = len(dims)
-        dim_sizes = list(self.coord_dims.values())
+        dim_sizes = list(self.coord_sizes.values())
         size = np.prod(dim_sizes)  # that the number of theoretical printouts
         masked_entries = self.mask.sum().values if self.mask is not None else 0
         lines = []
@@ -512,11 +511,15 @@ def model(self):
     @property
     def dims(self):
         # do explicitly sort as in vars (same as in coeffs)
-        return {k: self.data.dims[k] for k in self.vars.dims}
+        return self.vars.dims
 
     @property
     def coord_dims(self):
-        return {k: self.data.dims[k] for k in self.dims if k not in HELPER_DIMS}
+        return {k for k in self.dims if k not in HELPER_DIMS}
+
+    @property
+    def coord_sizes(self):
+        return {k: v for k, v in self.sizes.items() if k not in HELPER_DIMS}
 
     @property
     def vars(self):
@@ -637,7 +640,7 @@ def cumsum(
         if isinstance(dim, str):
             # Make sure, single mentioned dimensions is handled correctly.
             dim = [dim]
-        dim_dict = {dim_name: self.data.dims[dim_name] for dim_name in dim}
+        dim_dict = {dim_name: self.data.sizes[dim_name] for dim_name in dim}
         return self.rolling(dim=dim_dict).sum(keep_attrs=keep_attrs, skipna=skipna)
 
     @classmethod
@@ -1378,7 +1381,7 @@ def merge(*exprs, dim=TERM_DIM, cls=LinearExpression, **kwargs):
 
     if cls in linopy_types and dim in HELPER_DIMS:
         coord_dims = [
-            {k: v for k, v in e.dims.items() if k not in HELPER_DIMS} for e in exprs
+            {k: v for k, v in e.sizes.items() if k not in HELPER_DIMS} for e in exprs
         ]
         override = check_common_keys_values(coord_dims)
     else:

linopy/io.py

@@ -654,7 +654,7 @@ def get_prefix(ds, prefix):
         for dim in ds.dims:
             for name in ds.attrs.get(f"{dim}_multiindex", []):
                 multiindexes.append(prefix + "-" + name)
-        ds = ds.drop(set(ds.coords) - set(ds.dims) - set(multiindexes))
+        ds = ds.drop_vars(set(ds.coords) - set(ds.dims) - set(multiindexes))
         to_rename = set([*ds.dims, *ds.coords, *ds])
         ds = ds.rename({d: remove_prefix(d, prefix) for d in to_rename})
         ds.attrs = {

linopy/model.py

@@ -565,7 +565,7 @@ def assert_sign_rhs_not_None(lhs, sign, rhs):
         # ensure helper dimensions are not set as coordinates
         if drop_dims := set(HELPER_DIMS).intersection(data.coords):
             # TODO: add a warning here, routines should be safe against this
-            data = data.drop(drop_dims)
+            data = data.drop_vars(drop_dims)
 
         if mask is not None:
             mask = as_dataarray(mask).astype(bool)

test/test_io.py

@@ -91,6 +91,11 @@ def test_model_to_netcdf_with_dash_names(model_with_dash_names, tmp_path):
     assert_model_equal(m, p)
 
 
+# skip it xarray version is 2024.01.0 due to issue https://github.com/pydata/xarray/issues/8628
+@pytest.mark.skipif(
+    xr.__version__ == "2024.1.0",
+    reason="xarray version 2024.1.0 has a bug with MultiIndex deserialize",
+)
 def test_model_to_netcdf_with_multiindex(model_with_multiindex, tmp_path):
     m = model_with_multiindex
     fn = tmp_path / "test.nc"

test/test_quadratic_expression.py

@@ -35,13 +35,13 @@ def z(model):
 def test_quadratic_expression_from_variables_multiplication(x, y):
     expr = x * y
     assert isinstance(expr, QuadraticExpression)
-    assert expr.data.dims[FACTOR_DIM] == 2
+    assert expr.data.sizes[FACTOR_DIM] == 2
 
 
 def test_quadratic_expression_from_linexpr_multiplication(x, y):
     expr = (10 * x + y) * y
     assert isinstance(expr, QuadraticExpression)
-    assert expr.data.dims[FACTOR_DIM] == 2
+    assert expr.data.sizes[FACTOR_DIM] == 2
     assert expr.nterm == 2