From fe497beceb8ba8d8c7c34862b3235334c0a78c69 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 12 Jan 2023 09:32:08 +0100 Subject: [PATCH 1/6] raise UserWarning when adding unaligned variables --- doc/release_notes.rst | 1 + linopy/common.py | 12 ++- linopy/expressions.py | 8 +- linopy/model.py | 57 ++++++++---- linopy/variables.py | 151 +++++++++++++++++++++++-------- test/test_io.py | 13 ++- test/test_linear_expression.py | 18 ++-- test/test_repr.py | 22 ++++- test/test_variable.py | 14 ++- test/test_variable_assignment.py | 12 ++- test/test_variables.py | 9 +- 11 files changed, 227 insertions(+), 90 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 270e04b8..dd604a71 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -16,6 +16,7 @@ Upcoming Release * The classes Variable, LinearExpression, Constraint, ScalarVariable, ScalarLinearExpression and ScalarConstraint now require the model in the initialization (mostly internal code is affected). * The `eval` module was removed in favor of arithmetic operations on the classes `Variable`, `LinearExpression` and `Constraint`. * Solver options are now printed out in the console when solving a model. +* If a variable with indexes differing from the model internal indexes are assigned, linopy will raise a warning and align the variable to the model indexes. Version 0.0.15 -------------- diff --git a/linopy/common.py b/linopy/common.py index 7be53ab7..85ee2375 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -90,11 +90,15 @@ def print_coord(coord): return "" -def print_single_variable(lower, upper, var, vartype): - if vartype == "Binary Variable": - return f"\n {var}" +def print_single_variable(variable, name, coord, lower, upper): + if name in variable.model.variables._integer_variables: + bounds = "Z ⋂ " + f"[{lower},...,{upper}]" + elif name in variable.model.variables._binary_variables: + bounds = "{0, 1}" else: - return f"\n{lower} ≤ {var} ≤ {upper}" + bounds = f"[{lower}, {upper}]" + + return f"{name}{print_coord(coord)} ∈ {bounds}" def print_single_expression(c, v, model): diff --git a/linopy/expressions.py b/linopy/expressions.py index 18d51aed..8f65de9e 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -205,9 +205,9 @@ def __repr__(self): return f"LinearExpression:\n-----------------\n{expr_string}" # print only a few values - max_prints = 14 - split_at = max_prints // 2 - to_print = head_tail_range(nexprs, max_prints) + max_print = 14 + split_at = max_print // 2 + to_print = head_tail_range(nexprs, max_print) coords = self.unravel_coords(to_print) # loop over all values to print @@ -221,7 +221,7 @@ def __repr__(self): data_string += f"\n{coord_string}: {expr_string}" - if i == split_at - 1 and nexprs > max_prints: + if i == split_at - 1 and nexprs > max_print: data_string += "\n\t\t..." # create shape string diff --git a/linopy/model.py b/linopy/model.py index a1114d2e..f2424b69 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -8,6 +8,7 @@ import logging import os import re +import warnings from pathlib import Path from tempfile import NamedTemporaryFile, gettempdir @@ -406,18 +407,18 @@ def add_variables( >>> m = Model() >>> time = pd.RangeIndex(10, name="Time") >>> m.add_variables(lower=0, coords=[time], name="x") - Continuous Variable (Time: 10) - ------------------------------ - 0 ≤ x[0] ≤ inf - 0 ≤ x[1] ≤ inf - 0 ≤ x[2] ≤ inf - 0 ≤ x[3] ≤ inf - 0 ≤ x[4] ≤ inf - 0 ≤ x[5] ≤ inf - 0 ≤ x[6] ≤ inf - 0 ≤ x[7] ≤ inf - 0 ≤ x[8] ≤ inf - 0 ≤ x[9] ≤ inf + Variable (Time: 10) + ------------------- + [0]: x[0] ∈ [0, inf] + [1]: x[1] ∈ [0, inf] + [2]: x[2] ∈ [0, inf] + [3]: x[3] ∈ [0, inf] + [4]: x[4] ∈ [0, inf] + [5]: x[5] ∈ [0, inf] + [6]: x[6] ∈ [0, inf] + [7]: x[7] ∈ [0, inf] + [8]: x[8] ∈ [0, inf] + [9]: x[9] ∈ [0, inf] """ if name is None: name = "var" + str(self._varnameCounter) @@ -451,11 +452,37 @@ def add_variables( lower = DataArray(-inf, coords=coords, **kwargs) upper = DataArray(inf, coords=coords, **kwargs) - labels = DataArray(coords=coords).assign_attrs(binary=binary, integer=integer) + labels = DataArray(-2, coords=coords).assign_attrs( + binary=binary, integer=integer + ) + # ensure order of dims is the same lower = lower.transpose(*[d for d in labels.dims if d in lower.dims]) upper = upper.transpose(*[d for d in labels.dims if d in upper.dims]) + if mask is not None: + mask = DataArray(mask).astype(bool) + mask, _ = xr.align(mask, labels, join="right") + assert set(mask.dims).issubset( + labels.dims + ), "Dimensions of mask not a subset of resulting labels dimensions." + + # It is important to end up with monotonically increasing labels in the + # model's variables container as we use it for indirect indexing. + labels_reindexed = labels.reindex_like(self.variables.labels, fill_value=-1) + if not labels.equals(labels_reindexed): + warnings.warn( + f"Reindexing variable {name} to match existing coordinates.", + UserWarning, + ) + labels = labels_reindexed + lower = lower.reindex_like(labels) + upper = upper.reindex_like(labels) + if mask is None: + mask = labels != -1 + else: + mask = mask.reindex_like(labels_reindexed, fill_value=False) + self.check_force_dim_names(labels) start = self._xCounter @@ -464,10 +491,6 @@ def add_variables( self._xCounter += labels.size if mask is not None: - mask = DataArray(mask) - assert set(mask.dims).issubset( - labels.dims - ), "Dimensions of mask not a subset of resulting labels dimensions." labels = labels.where(mask, -1) if self.chunk: diff --git a/linopy/variables.py b/linopy/variables.py index 301379c4..4e49648a 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -24,6 +24,7 @@ _merge_inplace, forward_as_properties, has_optimized_model, + head_tail_range, is_constant, print_coord, print_single_variable, @@ -117,6 +118,8 @@ class Variable: __slots__ = ("_labels", "_model") __array_ufunc__ = None + _fill_value = -1 + def __init__(self, labels: DataArray, model: Any): """ Initialize the Constraint. @@ -181,52 +184,60 @@ def __repr__(self): """ Print the variable arrays. """ - lower = self.lower.values - upper = self.upper.values - # don't loop over all values if not necessary + # TODO: check that we can skip this if self.size == 1: - header = f"{self.type}\n{'-' * len(self.type)}" - lower = lower.item() - upper = upper.item() - data_string = print_single_variable(lower, upper, self.name, self.type) + header = f"Variable\n{'-' * len('Variable')}" + if self.lower.size > 1: + coord = {k: v for k, v in self.coords.items() if k in self.lower.dims} + lower = self.lower.sel(coord).item() + else: + lower = self.lower.item() + if self.upper.size > 1: + coord = {k: v for k, v in self.coords.items() if k in self.upper.dims} + upper = self.upper.sel(coord).item() + else: + upper = self.upper.item() + coord = [] + data_string = print_single_variable(self, self.name, coord, lower, upper) return f"{header}\n{data_string}" # print only a few values max_print = 14 split_at = max_print // 2 - to_print = np.flatnonzero(self.mask) - truncate = len(to_print) > max_print - if truncate: - to_print = np.hstack([to_print[:split_at], to_print[-split_at:]]) + to_print = head_tail_range(self.size, max_print) # create string, we use numpy to get the indexes - data_string = "" idx = np.unravel_index(to_print, self.shape) - indexes = np.stack(idx) coords = [self.indexes[self.dims[i]][idx[i]] for i in range(len(self.dims))] + coords = list(zip(*coords)) + labels = np.ravel(self.labels.values)[to_print] - # loop over all values to print - for i in range(len(to_print)): - # this is the index for the labels array - ix = tuple(indexes[..., i]) - # lower and upper bounds might only be defined for some dimensions - lix = tuple( - ix[i] for i in range(self.ndim) if self.dims[i] in self.lower.dims - ) - uix = tuple( - ix[i] for i in range(self.ndim) if self.dims[i] in self.upper.dims - ) + data_string = "" + for i, coord in enumerate(coords): - # create coordinate string - coord = [c[i] for c in coords] + label = labels[i] coord_string = print_coord(coord) - var_string = f"{self.name}{coord_string}" - data_string += print_single_variable( - lower[lix], upper[uix], var_string, self.type - ) - if i == split_at - 1 and truncate: + if label != -1: + vname, vcoord = self.model.variables.get_label_position(label) + + # get lower and upper bounds, which might have less dimensions + lcoord = [ + c for i, c in enumerate(vcoord) if self.dims[i] in self.lower.dims + ] + ucoord = [ + c for i, c in enumerate(vcoord) if self.dims[i] in self.upper.dims + ] + lower = self.lower.loc[tuple(lcoord)].item() + upper = self.upper.loc[tuple(ucoord)].item() + var_string = print_single_variable(self, vname, vcoord, lower, upper) + else: + var_string = "None" + + data_string += f"\n{coord_string}: {var_string}" + + if i == split_at - 1 and self.size > max_print: data_string += "\n\t\t..." # create shape string @@ -236,8 +247,8 @@ def __repr__(self): shape_string = f"({shape_string})" n_masked = (~self.mask).sum().item() mask_string = f" - {n_masked} masked entries" if n_masked else "" - header = f"{self.type} {shape_string}{mask_string}\n" + "-" * ( - len(self.type) + len(shape_string) + len(mask_string) + 1 + header = f"Variable {shape_string}{mask_string}\n" + "-" * ( + len("Variable") + len(shape_string) + len(mask_string) + 1 ) return f"{header}{data_string}" @@ -452,6 +463,14 @@ def type(self): else: return "Continuous Variable" + @classmethod + @property + def fill_value(self): + """ + Return the fill value of the variable. + """ + return self._fill_value + @property def mask(self): """ @@ -464,7 +483,7 @@ def mask(self): ------- xr.DataArray """ - return (self.labels != -1).astype(bool) + return (self.labels != self._fill_value).astype(bool) @property def upper(self): @@ -589,8 +608,64 @@ def where(self, cond, other=-1, **kwargs): ------- linopy.Variable """ + if isinstance(other, Variable): + other = other.labels + elif isinstance(other, ScalarVariable): + other = other.label return self.__class__(self.labels.where(cond, other, **kwargs), self.model) + def ffill(self, dim, limit=None): + """ + Forward fill the variable along a dimension. + + This operation call ``xarray.DataArray.ffill`` but ensures preserving + the linopy.Variable type. + + Parameters + ---------- + dim : str + Dimension over which to forward fill. + limit : int, optional + Maximum number of consecutive NaN values to forward fill. Must be greater than or equal to 0. + + Returns + ------- + linopy.Variable + """ + labels = ( + self.labels.where(self.labels != -1) + .ffill(dim, limit=limit) + .fillna(-1) + .astype(int) + ) + return self.__class__(labels, self.model) + + def bfill(self, dim, limit=None): + """ + Backward fill the variable along a dimension. + + This operation call ``xarray.DataArray.bfill`` but ensures preserving + the linopy.Variable type. + + Parameters + ---------- + dim : str + Dimension over which to backward fill. + limit : int, optional + Maximum number of consecutive NaN values to backward fill. Must be greater than or equal to 0. + + Returns + ------- + linopy.Variable + """ + labels = ( + self.labels.where(self.labels != -1) + .bfill(dim, limit=limit) + .fillna(-1) + .astype(int) + ) + return self.__class__(labels, self.model) + def sanitize(self): """ Sanitize variable by ensuring int dtype with fill value of -1. @@ -611,10 +686,6 @@ def equals(self, other): assign_coords = varwrap(DataArray.assign_coords) - astype = varwrap(DataArray.astype) - - bfill = varwrap(DataArray.bfill) - broadcast_like = varwrap(DataArray.broadcast_like) compute = varwrap(DataArray.compute) @@ -625,8 +696,6 @@ def equals(self, other): drop_isel = varwrap(DataArray.drop_isel) - ffill = varwrap(DataArray.ffill) - fillna = varwrap(DataArray.fillna) sel = varwrap(DataArray.sel) @@ -971,6 +1040,8 @@ def __init__(self, label: int, model: Any): self._model = model def __repr__(self) -> str: + if self.label == -1: + return "ScalarVariable: None" name, coord = self.model.variables.get_label_position(self.label) coord_string = print_coord(coord) return f"ScalarVariable: {name}{coord_string}" diff --git a/test/test_io.py b/test/test_io.py index 9f7331d1..565bcca8 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -20,7 +20,7 @@ def m(): m = Model() x = m.add_variables(4, pd.Series([8, 10]), name="x") - y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]), name="y") + y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4]]), name="y") m.add_constraints(x + y, LESS_EQUAL, 10) @@ -48,11 +48,14 @@ def test_str_arrays_with_nans(): m = Model() m.add_variables(4, pd.Series([8, 10]), name="x") - # now expand the second dimension, expended values of x will be nan - m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]), name="y") - assert m["x"].values[-1] == -1 + # now expand the second dimension, expanded values of x will be nan - da = int_to_str(m["x"].values) + with pytest.warns(UserWarning): + m.add_variables(0, pd.DataFrame([[1, 2]]), name="y") + + assert m["y"].values[-1, -1] == -1 + + da = int_to_str(m["y"].values) assert da.dtype == object diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 068145ee..04c2d312 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -433,23 +433,27 @@ def test_merge(x, y, z): def test_rolling_sum_variable_deprecated(v): - rolled = v.rolling_sum(dim_2=2) + with pytest.warns(DeprecationWarning): + rolled = v.rolling_sum(dim_2=2) assert rolled.nterm == 2 def test_linear_expression_rolling_sum_deprecated(x, v): - rolled = v.to_linexpr().rolling_sum(dim_2=2) + with pytest.warns(DeprecationWarning): + rolled = v.to_linexpr().rolling_sum(dim_2=2) assert rolled.nterm == 2 # multi-dimensional rolling with non-scalar _term dimension expr = 10 * x + v - rolled = expr.rolling_sum(dim_2=3) + with pytest.warns(DeprecationWarning): + rolled = expr.rolling_sum(dim_2=3) assert rolled.nterm == 6 def test_variable_groupby_sum_deprecated(v): groups = xr.DataArray([1] * 10 + [2] * 10, coords=v.coords) - grouped = v.groupby_sum(groups) + with pytest.warns(DeprecationWarning): + grouped = v.groupby_sum(groups) assert "group" in grouped.dims assert (grouped.data.group == [1, 2]).all() assert grouped.data._term.size == 10 @@ -457,14 +461,16 @@ def test_variable_groupby_sum_deprecated(v): def test_linear_expression_groupby_sum_deprecated(v): groups = xr.DataArray([1] * 10 + [2] * 10, coords=v.coords) - grouped = v.to_linexpr().groupby_sum(groups) + with pytest.warns(DeprecationWarning): + grouped = v.to_linexpr().groupby_sum(groups) assert "group" in grouped.dims assert (grouped.data.group == [1, 2]).all() assert grouped.data._term.size == 10 # now asymetric groups which result in different nterms groups = xr.DataArray([1] * 12 + [2] * 8, coords=v.coords) - grouped = v.to_linexpr().groupby_sum(groups) + with pytest.warns(DeprecationWarning): + grouped = v.to_linexpr().groupby_sum(groups) assert "group" in grouped.dims # first group must be full with vars assert (grouped.data.sel(group=1) > 0).all() diff --git a/test/test_repr.py b/test/test_repr.py index bd414cbf..53ce50a0 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -1,4 +1,6 @@ +import numpy as np import pandas as pd +import pytest import xarray as xr from linopy import Model @@ -21,6 +23,14 @@ c = m.add_variables(lower, upper, name="c", mask=c_mask) d = m.add_variables(0, 10, coords=[types], name="d") +# add variable with indexes to reindex +with pytest.warns(UserWarning): + e = m.add_variables(0, upper[5:], name="e") + + f_mask = np.full_like(upper[:5], True, dtype=bool) + f_mask[:3] = False + f = m.add_variables(0, upper[5:], name="f", mask=f_mask) + # create linear expression for each variable lu = 1 * u @@ -62,12 +72,20 @@ def test_variable_repr(): - for var in [u, v, x, y, z, a, b, c, d]: + for var in [u, v, x, y, z, a, b, c, d, e, f]: repr(var) def test_scalar_variable_repr(): - repr(u[0, 0]) + for var in [u, v, x, y, z, a, b, c, d]: + coord = tuple([var.indexes[c][0] for c in var.dims]) + repr(var[coord]) + + +def test_single_variable_repr(): + for var in [u, v, x, y, z, a, b, c, d]: + coord = tuple([var.indexes[c][0] for c in var.dims]) + repr(var.loc[coord]) def test_linear_expression_repr(): diff --git a/test/test_variable.py b/test/test_variable.py index 798dd90f..cdbc441f 100644 --- a/test/test_variable.py +++ b/test/test_variable.py @@ -142,8 +142,11 @@ def test_variable_shift(x): def test_variable_bfill(x): - result = x.bfill("first") - assert isinstance(result, linopy.variables.Variable) + x = x.where([False] * 4 + [True] * 6) + x = x.bfill("first") + assert isinstance(x, linopy.variables.Variable) + assert x.values[2] == x.values[4] + assert x.values[2] != x.values[5] def test_variable_broadcast_like(x): @@ -152,8 +155,11 @@ def test_variable_broadcast_like(x): def test_variable_ffill(x): - result = x.ffill("first") - assert isinstance(result, linopy.variables.Variable) + x = x.where([True] * 4 + [False] * 6) + x = x.ffill("first") + assert isinstance(x, linopy.variables.Variable) + assert x.values[9] == x.values[3] + assert x.values[3] != x.values[2] def test_variable_fillna(x): diff --git a/test/test_variable_assignment.py b/test/test_variable_assignment.py index 3dbaaef2..ecb46947 100644 --- a/test/test_variable_assignment.py +++ b/test/test_variable_assignment.py @@ -105,8 +105,8 @@ def test_variable_assignment_chunked(): def test_variable_assignment_different_coords(): - # set a variable with different set of coordinates, this should be properly - # merged + # set a variable with different set of coordinates + # since v0.1 new coordinates are reindexed to the old ones m = Model() lower = pd.DataFrame(np.zeros((10, 10))) upper = pd.Series(np.ones((10))) @@ -114,10 +114,12 @@ def test_variable_assignment_different_coords(): lower = pd.DataFrame(np.zeros((20, 10))) upper = pd.Series(np.ones((20))) - m.add_variables(lower, upper, name="y") - assert m.variables.labels.y.shape == (20, 10) + with pytest.warns(UserWarning): + m.add_variables(lower, upper, name="y") + + assert m.variables.labels.y.shape == (10, 10) # x should now be aligned to new coords and contain 100 nans - assert m.variables.labels.x.shape == (20, 10) + assert m.variables.labels.x.shape == (10, 10) assert (m.variables.labels.x != -1).sum() == 100 diff --git a/test/test_variables.py b/test/test_variables.py index 5c49d279..a678fb65 100644 --- a/test/test_variables.py +++ b/test/test_variables.py @@ -30,7 +30,8 @@ def test_variables_assignment_with_merge(): Test the merger of a variables with same dimension name but with different lengths. - Missing values should be filled up with -1. + New coordinates are aligned to the existing ones. Thus this should + raise a warning. """ m = Model() @@ -38,8 +39,10 @@ def test_variables_assignment_with_merge(): m.add_variables(upper) upper = pd.Series(np.ones((12))) - m.add_variables(upper) - assert m.variables.labels.var0[-1].item() == -1 + with pytest.warns(UserWarning): + m.add_variables(upper) + + assert m.variables.labels.var0[-1].item() != -1 def test_scalar_variables_name_counter(): From 33955206f8d03234a438b1c9dc66832e586efba6 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 12 Jan 2023 11:14:14 +0100 Subject: [PATCH 2/6] fix single variable repr, consolidate repr code between Variable and Constraint --- linopy/common.py | 7 +++++ linopy/constraints.py | 24 ++++++----------- linopy/variables.py | 60 +++++++++++++++++++------------------------ test/test_repr.py | 2 +- 4 files changed, 43 insertions(+), 50 deletions(-) diff --git a/linopy/common.py b/linopy/common.py index 85ee2375..8da23a31 100644 --- a/linopy/common.py +++ b/linopy/common.py @@ -75,6 +75,11 @@ def best_int(max_value): return t +def dictsel(d, keys): + "Reduce dictionary to keys that appear in selection." + return {k: v for k, v in d.items() if k in keys} + + def head_tail_range(stop, max_number_of_values=14): split_at = max_number_of_values // 2 if stop > max_number_of_values: @@ -84,6 +89,8 @@ def head_tail_range(stop, max_number_of_values=14): def print_coord(coord): + if isinstance(coord, dict): + coord = coord.values() if len(coord): return "[" + ", ".join([str(c) for c in coord]) + "]" else: diff --git a/linopy/constraints.py b/linopy/constraints.py index 826fb20b..b56bf935 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -22,6 +22,7 @@ from linopy import expressions, variables from linopy.common import ( _merge_inplace, + dictsel, forward_as_properties, has_optimized_model, is_constant, @@ -118,29 +119,20 @@ def __repr__(self): # create string, we use numpy to get the indexes data_string = "" idx = np.unravel_index(to_print, self.shape) - indexes = np.stack(idx) - coords = [self.indexes[self.dims[i]][idx[i]] for i in range(len(self.dims))] + coords = [self.indexes[dim][idx[i]] for i, dim in enumerate(self.dims)] # loop over all values to LinearExpression(print data_string = "" for i in range(len(to_print)): - # this is the index for the labels array - ix = tuple(indexes[..., i]) - # sign and rhs might only be defined for some dimensions - six = tuple( - ix[i] for i in range(self.ndim) if self.dims[i] in self.sign.dims - ) - rix = tuple( - ix[i] for i in range(self.ndim) if self.dims[i] in self.rhs.dims - ) - - coord = [c[i] for c in coords] + coord = {dim: c[i] for dim, c in zip(self.dims, coords)} coord_string = print_coord(coord) expr_string = print_single_expression( - self.coeffs.values[ix], self.vars.values[ix], self.lhs.model + self.coeffs.sel(coord).values, + self.vars.sel(coord).values, + self.lhs.model, ) - sign_string = f"{self.sign.values[six]}" - rhs_string = f"{self.rhs.values[rix]}" + sign_string = f"{self.sign.sel(**dictsel(coord, self.sign.dims)).item()}" + rhs_string = f"{self.rhs.sel(**dictsel(coord, self.sign.dims)).item()}" data_string += f"\n{self.name}{coord_string}: {expr_string} {sign_string} {rhs_string}" diff --git a/linopy/variables.py b/linopy/variables.py index 4e49648a..ff2b8f4c 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -22,6 +22,7 @@ import linopy.expressions as expressions from linopy.common import ( _merge_inplace, + dictsel, forward_as_properties, has_optimized_model, head_tail_range, @@ -185,19 +186,10 @@ def __repr__(self): Print the variable arrays. """ # don't loop over all values if not necessary - # TODO: check that we can skip this - if self.size == 1: + if not self.coords: header = f"Variable\n{'-' * len('Variable')}" - if self.lower.size > 1: - coord = {k: v for k, v in self.coords.items() if k in self.lower.dims} - lower = self.lower.sel(coord).item() - else: - lower = self.lower.item() - if self.upper.size > 1: - coord = {k: v for k, v in self.coords.items() if k in self.upper.dims} - upper = self.upper.sel(coord).item() - else: - upper = self.upper.item() + lower = self.lower.item() + upper = self.upper.item() coord = [] data_string = print_single_variable(self, self.name, coord, lower, upper) return f"{header}\n{data_string}" @@ -208,10 +200,16 @@ def __repr__(self): to_print = head_tail_range(self.size, max_print) # create string, we use numpy to get the indexes - idx = np.unravel_index(to_print, self.shape) - coords = [self.indexes[self.dims[i]][idx[i]] for i in range(len(self.dims))] - coords = list(zip(*coords)) - labels = np.ravel(self.labels.values)[to_print] + if self.shape: + idx = np.unravel_index(to_print, self.shape) + labels = np.ravel(self.labels.values)[to_print] + coords = [self.indexes[self.dims[i]][idx[i]] for i in range(len(self.dims))] + coords = list(zip(*coords)) + else: + # case a single variable was selected + idx = [0] + labels = np.ravel(self.labels.values) + coords = [[c.item() for c in self.coords.values()]] data_string = "" for i, coord in enumerate(coords): @@ -221,16 +219,8 @@ def __repr__(self): if label != -1: vname, vcoord = self.model.variables.get_label_position(label) - - # get lower and upper bounds, which might have less dimensions - lcoord = [ - c for i, c in enumerate(vcoord) if self.dims[i] in self.lower.dims - ] - ucoord = [ - c for i, c in enumerate(vcoord) if self.dims[i] in self.upper.dims - ] - lower = self.lower.loc[tuple(lcoord)].item() - upper = self.upper.loc[tuple(ucoord)].item() + lower = self.lower.sel(dictsel(vcoord, self.lower.dims)).item() + upper = self.upper.sel(dictsel(vcoord, self.upper.dims)).item() var_string = print_single_variable(self, vname, vcoord, lower, upper) else: var_string = "None" @@ -241,10 +231,13 @@ def __repr__(self): data_string += "\n\t\t..." # create shape string - shape_string = ", ".join( - [f"{self.dims[i]}: {self.shape[i]}" for i in range(self.ndim)] - ) - shape_string = f"({shape_string})" + if self.shape: + shape_string = ", ".join( + [f"{self.dims[i]}: {self.shape[i]}" for i in range(self.ndim)] + ) + shape_string = f"({shape_string})" + else: + shape_string = "" n_masked = (~self.mask).sum().item() mask_string = f" - {n_masked} masked entries" if n_masked else "" header = f"Variable {shape_string}{mask_string}\n" + "-" * ( @@ -916,9 +909,10 @@ def get_label_position(self, values): index = np.unravel_index(value - start, labels.shape) # Extract the coordinates from the indices - coord = [ - labels.indexes[dim][i] for dim, i in zip(labels.dims, index) - ] + coord = { + dim: labels.indexes[dim][i] + for dim, i in zip(labels.dims, index) + } # Add the name of the DataArray and the coordinates to the result list coords.append((name, coord)) diff --git a/test/test_repr.py b/test/test_repr.py index 53ce50a0..deed9f91 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -8,7 +8,7 @@ m = Model() lower = pd.Series(0, range(10)) -upper = pd.DataFrame(10, range(10), range(10)) +upper = pd.DataFrame(np.arange(10, 110).reshape(10, 10), range(10), range(10)) types = pd.Index(list("abcdefgh"), name="types") u = m.add_variables(0, upper, name="u") From 8412970189a5aa6bfa3b7fa6cc7e009f0eddd769 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 12 Jan 2023 15:18:00 +0100 Subject: [PATCH 3/6] fix rename function enable automatic for coefficient assignment in from_tuple if needed --- linopy/expressions.py | 32 +++++++++++++++++++++++--------- linopy/io.py | 3 +++ linopy/variables.py | 8 ++++++++ test/test_linear_expression.py | 14 ++++++++++++++ test/test_variable.py | 8 ++++++++ 5 files changed, 56 insertions(+), 9 deletions(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index 8f65de9e..91b3295c 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -166,15 +166,23 @@ def __init__(self, data, model): if not set(data).issuperset({"coeffs", "vars"}): raise ValueError("data must contain the variables 'coeffs' and 'vars'") + term_dims = [d for d in data.dims if d.endswith("_term")] + if not term_dims: + raise ValueError("data must contain one dimension ending with '_term'") + term_dim = term_dims[0] + # make sure that all non-helper dims have coordinates - if not all(dim in data.coords for dim in data.dims if not dim.startswith("_")): - raise ValueError("data must have coordinates for all non-helper dimensions") + missing_coords = set(data.dims) - set(data.coords) - {term_dim} + if missing_coords: + raise ValueError( + f"Dimensions {missing_coords} have no coordinates, please add them." + ) if np.issubdtype(data.vars, np.floating): data["vars"] = data.vars.fillna(-1).astype(int) (data,) = xr.broadcast(data) - data = data.transpose(..., "_term") + data = data.transpose(..., term_dim) if not isinstance(model, Model): raise ValueError("model must be an instance of linopy.Model") @@ -452,6 +460,11 @@ def from_tuples(cls, *tuples, chunk=None): c = DataArray(c, v.coords) else: c = as_dataarray(c) + # if a dimension is not in the coords, add it as a range index + for i, dim in enumerate(c.dims): + if dim not in c.coords: + c = c.assign_coords(**{dim: pd.RangeIndex(c.shape[i])}) + ds = Dataset({"coeffs": c, "vars": v}).expand_dims("_term") expr = cls(ds, model) exprs.append(expr) @@ -803,10 +816,6 @@ def sanitize(self): def equals(self, other: "LinearExpression"): return self.data.equals(_expr_unwrap(other)) - # TODO: make this return a LinearExpression (needs refactoring of __init__) - def rename(self, name_dict=None, **names) -> Dataset: - return self.data.rename(name_dict, **names) - def __iter__(self): return self.data.__iter__() @@ -843,6 +852,8 @@ def __iter__(self): reindex = exprwrap(Dataset.reindex, fill_value=_fill_value) + rename = exprwrap(Dataset.rename) + rename_dims = exprwrap(Dataset.rename_dims) roll = exprwrap(Dataset.roll) @@ -894,8 +905,11 @@ def merge(*exprs, dim="_term", cls=LinearExpression): model = exprs[0].model exprs = [e.data if isinstance(e, cls) else e for e in exprs] - if not all(len(expr._term) == len(exprs[0]._term) for expr in exprs[1:]): - exprs = [expr.assign_coords(_term=np.arange(len(expr._term))) for expr in exprs] + if cls == LinearExpression: + if not all(len(expr._term) == len(exprs[0]._term) for expr in exprs[1:]): + exprs = [ + expr.assign_coords(_term=np.arange(len(expr._term))) for expr in exprs + ] kwargs = dict(fill_value=cls._fill_value, coords="minimal", compat="override") ds = xr.concat(exprs, dim, **kwargs) diff --git a/linopy/io.py b/linopy/io.py index 9ee501ae..fd7f99c5 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -461,9 +461,12 @@ def to_netcdf(m, *args, **kwargs): **kwargs : TYPE Keyword arguments passed to ``xarray.Dataset.to_netcdf``. """ + from linopy.expressions import LinearExpression def get_and_rename(m, attr, prefix=""): ds = getattr(m, attr) + if isinstance(ds, LinearExpression): + ds = ds.data return ds.rename({v: prefix + attr + "-" + v for v in ds}) vars = [ diff --git a/linopy/variables.py b/linopy/variables.py index ff2b8f4c..57535a50 100644 --- a/linopy/variables.py +++ b/linopy/variables.py @@ -432,6 +432,14 @@ def labels(self): """ return self._labels + @property + def data(self): + """ + Get the data of the variable. + """ + # Needed for compatibility with linopy.merge + return self.labels + @property def model(self): """ diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index 04c2d312..b912284e 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -133,6 +133,9 @@ def test_linear_expression_with_multiplication(x): expr = np.array([1, 2]) * x assert isinstance(expr, LinearExpression) + expr = xr.DataArray(np.array([[1, 2], [2, 3]])) * x + assert isinstance(expr, LinearExpression) + def test_linear_expression_with_addition(m, x, y): expr = 10 * x + y @@ -429,6 +432,17 @@ def test_merge(x, y, z): assert res.sel(dim_1=0).vars[2].item() == -1 +def test_rename(x, y, z): + expr = 10 * x + y + z + renamed = expr.rename({"dim_0": "dim_5"}) + assert set(renamed.dims) == {"dim_1", "dim_5", "_term"} + assert renamed.nterm == 3 + + renamed = expr.rename({"dim_0": "dim_1", "dim_1": "dim_2"}) + assert set(renamed.dims) == {"dim_1", "dim_2", "_term"} + assert renamed.nterm == 3 + + # -------------------------------- deprecated -------------------------------- # diff --git a/test/test_variable.py b/test/test_variable.py index cdbc441f..5444ec35 100644 --- a/test/test_variable.py +++ b/test/test_variable.py @@ -134,6 +134,14 @@ def test_variable_where(x): assert isinstance(x, linopy.variables.Variable) assert x.values[9] == -1 + x = x.where([True] * 4 + [False] * 6, x[0]) + assert isinstance(x, linopy.variables.Variable) + assert x.values[9] == x[0].label + + x = x.where([True] * 4 + [False] * 6, x.loc[0]) + assert isinstance(x, linopy.variables.Variable) + assert x.values[9] == x[0].label + def test_variable_shift(x): x = x.shift(first=3) From e41672a2a31f4308bca73c4c23e870a72fe842d4 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 12 Jan 2023 15:32:14 +0100 Subject: [PATCH 4/6] improve test cov --- test/test_linear_expression.py | 4 ++++ test/test_variable.py | 10 +++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index b912284e..b5efba0d 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -80,6 +80,10 @@ def test_linexpr_with_wrong_data(m): with pytest.raises(ValueError): LinearExpression(data, None) + with pytest.raises(ValueError): + data = (1 * m["x"]).data.reset_index("dim_0") + LinearExpression(data, None) + def test_repr(m): expr = m.linexpr((10, "x"), (1, "y")) diff --git a/test/test_variable.py b/test/test_variable.py index 5444ec35..f7de5a72 100644 --- a/test/test_variable.py +++ b/test/test_variable.py @@ -39,6 +39,14 @@ def test_variable_repr(x): x.__repr__() +def test_variable_labels(x): + isinstance(x.labels, xr.DataArray) + + +def test_variable_data(x): + isinstance(x.data, xr.DataArray) + + def test_wrong_variable_init(m, x): with pytest.raises(ValueError): linopy.Variable(x.labels.values, m) @@ -132,7 +140,7 @@ def test_variable_sum(x): def test_variable_where(x): x = x.where([True] * 4 + [False] * 6) assert isinstance(x, linopy.variables.Variable) - assert x.values[9] == -1 + assert x.values[9] == x.fill_value x = x.where([True] * 4 + [False] * 6, x[0]) assert isinstance(x, linopy.variables.Variable) From 8cc8ea5341ca3d267d2c2e90555e0d8e0c3260c5 Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 12 Jan 2023 15:37:02 +0100 Subject: [PATCH 5/6] ci: switch to python 3.9 --- .github/workflows/CI.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml index 66043ed9..f804095e 100644 --- a/.github/workflows/CI.yaml +++ b/.github/workflows/CI.yaml @@ -16,10 +16,10 @@ jobs: steps: - uses: actions/checkout@v2 - - name: Set up Python 3.8 + - name: Set up Python 3.9 uses: actions/setup-python@v1 with: - python-version: 3.8 + python-version: 3.9 - name: Install ubuntu dependencies if: matrix.os == 'ubuntu-latest' From f727c8ed1d004ddf1ed4d0fee80f3a925ea185ba Mon Sep 17 00:00:00 2001 From: Fabian Date: Thu, 12 Jan 2023 15:51:24 +0100 Subject: [PATCH 6/6] add test for missing coords --- linopy/expressions.py | 3 ++- test/test_linear_expression.py | 8 ++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/linopy/expressions.py b/linopy/expressions.py index 91b3295c..f56d7cc8 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -175,7 +175,8 @@ def __init__(self, data, model): missing_coords = set(data.dims) - set(data.coords) - {term_dim} if missing_coords: raise ValueError( - f"Dimensions {missing_coords} have no coordinates, please add them." + f"Dimensions {missing_coords} have no coordinates. For " + "consistency all dimensions must have coordinates." ) if np.issubdtype(data.vars, np.floating): diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py index b5efba0d..da7bf212 100644 --- a/test/test_linear_expression.py +++ b/test/test_linear_expression.py @@ -80,9 +80,13 @@ def test_linexpr_with_wrong_data(m): with pytest.raises(ValueError): LinearExpression(data, None) + lhs = 1 * m["x"] + vars = xr.DataArray(lhs.vars.values, dims=["dim_0", "_term"]) + coeffs = xr.DataArray(lhs.coeffs.values, dims=["dim_0", "_term"]) + data = xr.Dataset({"vars": vars, "coeffs": coeffs}) with pytest.raises(ValueError): - data = (1 * m["x"]).data.reset_index("dim_0") - LinearExpression(data, None) + # test missing coords + LinearExpression(data, m) def test_repr(m):