From 7e3cf9bb50d930679aaa4bbfbd9b51d677fb69da Mon Sep 17 00:00:00 2001 From: Fabian Date: Mon, 22 Jan 2024 18:29:38 +0100 Subject: [PATCH] address deprecations and issue with multiindex deserialization --- linopy/constraints.py | 2 +- linopy/expressions.py | 19 +++++++++++-------- linopy/io.py | 2 +- linopy/model.py | 2 +- test/test_io.py | 5 +++++ test/test_quadratic_expression.py | 4 ++-- 6 files changed, 21 insertions(+), 13 deletions(-) diff --git a/linopy/constraints.py b/linopy/constraints.py index 6aca342b..491f8d81 100644 --- a/linopy/constraints.py +++ b/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): diff --git a/linopy/expressions.py b/linopy/expressions.py index 3f1ed56a..f138917b 100644 --- a/linopy/expressions.py +++ b/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: diff --git a/linopy/io.py b/linopy/io.py index c52cf66e..1168697c 100644 --- a/linopy/io.py +++ b/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 = { diff --git a/linopy/model.py b/linopy/model.py index f6aad0ac..c887ffb4 100644 --- a/linopy/model.py +++ b/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) diff --git a/test/test_io.py b/test/test_io.py index 7c5cc099..04053acc 100644 --- a/test/test_io.py +++ b/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" diff --git a/test/test_quadratic_expression.py b/test/test_quadratic_expression.py index 92ed8b60..a0ccc8aa 100644 --- a/test/test_quadratic_expression.py +++ b/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