Repository: PyPSA/linopy
Pull request #5 (2021-10-06 15:20:25)

Masked arrays


Patch:
diff --git a/doc/release_notes.rst b/doc/release_notes.rst
index 12221be..b2f1679 100644
--- a/doc/release_notes.rst
+++ b/doc/release_notes.rst
@@ -4,9 +4,11 @@ Release Notes
 Upcoming Release
 ----------------
 
-* The "_term" dimension in the LinearExpression class is now stored without coordinates.
-* As soon as a LinearExpression is added to a model the "_term" dimension is rename to "{constraintname}_term" in order align the model better with the contained arrays and to avoid unnecessary nans.
-* `LinearExpression` now allows empty data as input.
+* The "_term" dimension in the ``LinearExpression`` class is now stored without coordinates.
+* As soon as a linear expression is added to a model the "_term" dimension is rename to "{constraintname}_term" in order align the model better with the contained arrays and to avoid unnecessary nans.
+* ``LinearExpression`` now allows empty data as input.
+* ``add_variables`` and ``add_constraints`` now accept a new argument ``mask``. The mask, which should be an boolean array, defines whether a variable/constraint is active (True) or should be ignored (False).
+* Missing values in the ``Model.variables`` and ``Model.constraints`` arrays are now indicated by -1. This circumvents changing the type from `int` to `float`.
 
 Version 0.0.3
 -------------
diff --git a/linopy/io.py b/linopy/io.py
index 8931643..2a907b0 100644
--- a/linopy/io.py
+++ b/linopy/io.py
@@ -4,7 +4,7 @@
 import logging
 import os
 import time
-from functools import reduce, partial
+from functools import partial, reduce
 
 import numpy as np
 import xarray as xr
@@ -49,7 +49,7 @@ def objective_to_file(m, f):
     coef = m.objective.coeffs
     var = m.objective.vars
 
-    nonnans = coef.notnull() & var.notnull()
+    nonnans = coef.notnull() & (var != -1)
     join = [to_float_str(coef), " x", to_int_str(var), "\n"]
     objective_str = join_str_arrays(join).where(nonnans, "")
     str_array_to_file(objective_str, f).compute()
@@ -66,12 +66,12 @@ def constraints_to_file(m, f):
 
     term_names = [f"{n}_term" for n in con]
 
-    nonnans = coef.notnull() & var.notnull()
+    nonnans = coef.notnull() & (var != -1)
     join = [to_float_str(coef), " x", to_int_str(var), "\n"]
     lhs_str = join_str_arrays(join).where(nonnans, "").reduce(np.sum, term_names)
     # .sum() does not work
 
-    nonnans = nonnans.any(term_names) & con.notnull() & sign.notnull() & rhs.notnull()
+    nonnans = nonnans.any(term_names) & (con != -1) & sign.notnull() & rhs.notnull()
 
     join = [
         "c",
@@ -94,7 +94,7 @@ def bounds_to_file(m, f):
     v = m.variables
     ub = m.variables_upper_bound
 
-    nonnans = lb.notnull() & v.notnull() & ub.notnull()
+    nonnans = lb.notnull() & ub.notnull() & (v != -1)
     join = [to_float_str(lb), " <= x", to_int_str(v), " <= ", to_float_str(ub), "\n"]
     bounds_str = join_str_arrays(join).where(nonnans, "")
     str_array_to_file(bounds_str, f).compute()
diff --git a/linopy/model.py b/linopy/model.py
index b36a079..5ecc749 100755
--- a/linopy/model.py
+++ b/linopy/model.py
@@ -102,17 +102,17 @@ class Model:
         """Get a model variable by the name."""
         return Variable(self.variables[key])
 
-    def _merge_inplace(self, attr, da, name):
+    def _merge_inplace(self, attr, da, name, **kwargs):
         """
         Assign a new variable to the dataset `attr` by merging.
 
         This takes care of all coordinate alignments, instead of a direct
         assignment like self.variables[name] = var
         """
-        ds = merge([getattr(self, attr), da.to_dataset(name=name)])
+        ds = merge([getattr(self, attr), da.to_dataset(name=name)], **kwargs)
         setattr(self, attr, ds)
 
-    def add_variables(self, lower=-inf, upper=inf, coords=None, name=None):
+    def add_variables(self, lower=-inf, upper=inf, coords=None, name=None, mask=None):
         """
         Assign a new, possibly multi-dimensional array of variables to the model.
 
@@ -135,6 +135,10 @@ class Model:
         name : str, optional
             Reference name of the added variables. The default None results in
             a name like "var1", "var2" etc.
+        mask : array_like, optional
+            Boolean mask with False values for variables which are skipped.
+            The shape of the mask has to match the shape the added variables.
+            Default is None.
 
         Raises
         ------
@@ -195,18 +199,21 @@ class Model:
         var = xr.DataArray(var, coords=broadcasted.coords)
         var = var.assign_attrs(name=name)
 
+        if mask is not None:
+            var = var.where(mask, -1)
+
         if self.chunk:
             lower = lower.chunk(self.chunk)
             upper = upper.chunk(self.chunk)
             var = var.chunk(self.chunk)
 
-        self._merge_inplace("variables", var, name)
+        self._merge_inplace("variables", var, name, fill_value=-1)
         self._merge_inplace("variables_lower_bound", lower, name)
         self._merge_inplace("variables_upper_bound", upper, name)
 
         return Variable(var)
 
-    def add_constraints(self, lhs, sign, rhs, name=None):
+    def add_constraints(self, lhs, sign, rhs, name=None, mask=None):
         """
         Assign a new, possibly multi-dimensional array of constraints to the model.
 
@@ -226,6 +233,10 @@ class Model:
         name : str, optional
             Reference name of the added constraints. The default None results
             results a name like "con1", "con2" etc.
+        mask : array_like, optional
+            Boolean mask with False values for variables which are skipped.
+            The shape of the mask has to match the shape the added variables.
+            Default is None.
 
 
         Returns
@@ -259,6 +270,9 @@ class Model:
         con = DataArray(con, coords=broadcasted.coords)
         con = con.assign_attrs(name=name)
 
+        if mask is not None:
+            con = con.where(mask, -1)
+
         lhs = lhs.rename({"_term": f"{name}_term"})
 
         if self.chunk:
@@ -268,7 +282,7 @@ class Model:
             con = con.chunk(self.chunk)
 
         # assign everything
-        self._merge_inplace("constraints", con, name)
+        self._merge_inplace("constraints", con, name, fill_value=-1)
         self._merge_inplace("constraints_lhs_coeffs", lhs.coeffs, name)
         self._merge_inplace("constraints_lhs_vars", lhs.vars, name)
         self._merge_inplace("constraints_sign", sign, name)
@@ -531,13 +545,13 @@ class Model:
         self.solver_model = res.pop("model", None)
         self.status = termination_condition
 
-        res["solution"].loc[np.nan] = np.nan
+        res["solution"].loc[-1] = np.nan
         for v in self.variables:
             idx = np.ravel(self.variables[v])
             sol = res["solution"][idx].values.reshape(self.variables[v].shape)
             self.solution[v] = xr.DataArray(sol, self.variables[v].coords)
 
-        res["dual"].loc[np.nan] = np.nan
+        res["dual"].loc[-1] = np.nan
         for c in self.constraints:
             idx = np.ravel(self.constraints[c])
             du = res["dual"][idx].values.reshape(self.constraints[c].shape)


Test patch:
diff --git a/test/test_io.py b/test/test_io.py
index d297bb8..545eb19 100644
--- a/test/test_io.py
+++ b/test/test_io.py
@@ -42,7 +42,7 @@ def test_str_arrays_with_nans():
     x = m.add_variables(4, pd.Series([8, 10]), name="x")
     # now expand the second dimension, expended values of x will be nan
     y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4], [5, 6]]), name="y")
-    assert not m["x"].notnull().all()
+    assert m["x"][-1].item() == -1
 
     da = to_int_str(m["x"])
     assert da.dtype == object
diff --git a/test/test_linear_expression.py b/test/test_linear_expression.py
index f50eb75..15ce555 100644
--- a/test/test_linear_expression.py
+++ b/test/test_linear_expression.py
@@ -59,15 +59,6 @@ def test_variable_to_linexpr():
     assert_equal(expr, m.linexpr((-1, "x"), (-8, "y")))
 
 
-def test_term_labels():
-    "Test that the _term dimension is named after the variables."
-    expr = 10 * x + y
-    other = m.linexpr((2, "y"), (1, "z"))
-
-    assert (expr._term == [0, 1]).all()
-    assert (other._term == [0, 1]).all()
-
-
 def test_add():
     expr = 10 * x + y
     other = 2 * y + z
diff --git a/test/test_model_creation.py b/test/test_model_creation.py
index 1630ae7..5f657ce 100755
--- a/test/test_model_creation.py
+++ b/test/test_model_creation.py
@@ -83,7 +83,7 @@ def test_add_variables_shape():
     assert m.variables.var9.shape == (20, 10)
     # var8 should now be aligned to new coords and contain 100 nans
     assert m.variables.var8.shape == (20, 10)
-    assert m.variables.var8.notnull().sum() == 100
+    assert (m.variables.var8 != -1).sum() == 100
 
     # setting with scalar and list
     with pytest.raises(ValueError):
@@ -94,6 +94,31 @@ def test_add_variables_shape():
         m.add_variables(lower, upper, name="var8")
 
 
+def test_masked_variables():
+    m = Model()
+
+    lower = pd.DataFrame(np.zeros((10, 10)))
+    upper = pd.Series(np.ones((10)))
+    mask = [True] * 5 + [False] * 5
+    m.add_variables(lower, upper, mask=mask)
+    assert m.variables.var0[-1, -1].item() == -1
+
+
+def test_variable_merging():
+    """
+    Test the merger of a variables with same dimension name but with
+    different lengths. Missing values should be filled up with -1.
+    """
+    m = Model()
+
+    upper = pd.Series(np.ones((10)))
+    m.add_variables(upper)
+
+    upper = pd.Series(np.ones((12)))
+    m.add_variables(upper)
+    assert m.variables.var0[-1].item() == -1
+
+
 def test_linexpr():
     m = Model()
 
@@ -125,6 +150,20 @@ def test_constraints():
     assert m.constraints_rhs.con0.dtype in (int, float)
 
 
+def test_masked_constraints():
+    m = Model()
+
+    lower = xr.DataArray(np.zeros((10, 10)), coords=[range(10), range(10)])
+    upper = xr.DataArray(np.ones((10, 10)), coords=[range(10), range(10)])
+    x = m.add_variables(lower, upper)
+    y = m.add_variables()
+
+    mask = [True] * 5 + [False] * 5
+    m.add_constraints(1 * x + 10 * y, "=", 0, mask=mask)
+    assert (m.constraints.con0[:, 0:5] != -1).all()
+    assert (m.constraints.con0[:, 5:10] == -1).all()
+
+
 def test_objective():
     m = Model()
 
diff --git a/test/test_optimization.py b/test/test_optimization.py
index 0451f50..ab4a596 100644
--- a/test/test_optimization.py
+++ b/test/test_optimization.py
@@ -62,6 +62,46 @@ def test_xpress():
     assert np.isclose(m.objective_value, 3.3)
 
 
+@pytest.mark.skipif("cplex" not in available_solvers, reason="Solver not available")
+def test_masked_variables_with_cplex():
+    m = Model()
+
+    lower = pd.Series(0, range(10))
+    x = m.add_variables(lower, name="x")
+    mask = [True] * 8 + [False, False]
+    y = m.add_variables(lower, name="y", mask=mask)
+
+    m.add_constraints(x + y, ">=", 10)
+
+    m.add_objective(2 * x + y)
+
+    m.solve("cbc")
+    assert m.solution.y[-2:].isnull().all()
+    assert m.solution.y[:-2].notnull().all()
+    assert m.solution.x.notnull().all()
+    assert (m.solution.x[-2:] == 10).all()
+
+
+@pytest.mark.skipif("cplex" not in available_solvers, reason="Solver not available")
+def test_masked_constraints_with_cplex():
+    m = Model()
+
+    lower = pd.Series(0, range(10))
+    x = m.add_variables(lower, name="x")
+    y = m.add_variables(lower, name="y")
+
+    mask = [True] * 8 + [False, False]
+    m.add_constraints(x + y, ">=", 10, mask=mask)
+    # for the last two entries only the following constraint will be active
+    m.add_constraints(x + y, ">=", 5)
+
+    m.add_objective(2 * x + y)
+
+    m.solve("cbc")
+    assert (m.solution.y[:-2] == 10).all()
+    assert (m.solution.y[-2:] == 5).all()
+
+
 def init_model_large():
     m = Model()
     time = pd.Index(range(10), name="time")
@@ -70,12 +110,12 @@ def init_model_large():
     y = m.add_variables(name="y", lower=0, coords=[time])
     factor = pd.Series(time, index=time)
 
-    m.add_constraints("Constraint1", 3 * x + 7 * y, ">=", 10 * factor)
-    m.add_constraints("Constraint2", 5 * x + 2 * y, ">=", 3 * factor)
+    m.add_constraints(3 * x + 7 * y, ">=", 10 * factor, name="Constraint1")
+    m.add_constraints(5 * x + 2 * y, ">=", 3 * factor, name="Constraint2")
 
     shifted = (1 * x).shift(time=1)
     lhs = (x - shifted).sel(time=time[1:])
-    m.add_constraints("Limited growth", lhs, "<=", 0.2)
+    m.add_constraints(lhs, "<=", 0.2, "Limited growth")
 
     m.add_objective((x + 2 * y).sum())
     m.solve()