Merge pull request #88 from PyPSA/improve-repr-follow-up

Improve repr follow up

Author: FabianHofmann

.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'

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
 --------------

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,17 +89,23 @@ 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:
         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):

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}"
 

linopy/expressions.py

@@ -166,15 +166,24 @@ 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. For "
+                "consistency all dimensions must have coordinates."
+            )
 
         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")
@@ -205,9 +214,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 +230,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
@@ -452,6 +461,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 +817,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 +853,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 +906,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)

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 = [

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: