IO: multiindeces and general fixes

doc/release_notes.rst

@@ -1,8 +1,13 @@
 Release Notes
 =============
 
-.. Upcoming Release
-.. ----------------
+Upcoming Release
+----------------
+
+* The IO with NetCDF files was made more secure and fixed for some cases. In particular, variables and constraints with a dash in the name are now supported (as used by PyPSA). The object sense and value are now properly stored and retrieved from the netcdf file.
+* The IO with NetCDF file now supports multiindexed coordinates.
+* The representation of single indexed expressions and constraints with non-empty dimensions/coordinates was fixed, e.g. `x.loc[["a"]] > 0` where `x` has only one dimension. Therefore the representation now shows the coordinates.
+
 
 Version 0.3.1
 -------------

linopy/common.py

@@ -332,6 +332,7 @@ def print_line(expr, const):
         mask = v != -1
         c, v = c[mask], v[mask]
     else:
+        # case for quadratic expressions
         mask = (v != -1).any(0)
         c = c[mask]
         v = v[:, mask]

linopy/constraints.py

@@ -165,14 +165,15 @@ def __repr__(self):
         """
         max_lines = options["display_max_rows"]
         dims = list(self.dims)
+        ndim = len(self.coord_dims)
         dim_sizes = list(self.sizes.values())[:-1]
         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 = []
 
         header_string = f"{self.type} `{self.name}`" if self.name else f"{self.type}"
 
-        if size > 1:
+        if size > 1 or ndim > 0:
             for indices in generate_indices_for_printout(dim_sizes, max_lines):
                 if indices is None:
                     lines.append("\t\t...")

linopy/expressions.py

@@ -330,14 +330,15 @@ def __repr__(self):
         """
         max_lines = options["display_max_rows"]
         dims = list(self.coord_dims)
+        ndim = len(dims)
         dim_sizes = list(self.coord_dims.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 = []
 
         header_string = self.type
 
-        if size > 1:
+        if size > 1 or ndim > 0:
             for indices in generate_indices_for_printout(dim_sizes, max_lines):
                 if indices is None:
                     lines.append("\t\t...")

linopy/io.py

@@ -298,7 +298,11 @@
             bounds_to_file(m, f, log=log, batch_size=batch_size)
             binaries_to_file(m, f, log=log, batch_size=batch_size)
             integers_to_file(
-                m, f, log=log, batch_size=batch_size, integer_label=integer_label
+                m,
+                f,
+                log=log,
+                batch_size=batch_size,
+                integer_label=integer_label,
             )
             f.write("end\n")
 
@@ -572,8 +576,18 @@
     """
 
     def with_prefix(ds, prefix):
-        ds = ds.rename({d: f"{prefix}-{d}" for d in [*ds.dims, *ds]})
+        to_rename = set([*ds.dims, *ds.coords, *ds])
+        ds = ds.rename({d: f"{prefix}-{d}" for d in to_rename})
         ds.attrs = {f"{prefix}-{k}": v for k, v in ds.attrs.items()}
+
+        # Flatten multiindexes
+        for dim in ds.dims:
+            if isinstance(ds[dim].to_index(), pd.MultiIndex):
+                prefix_len = len(prefix) + 1  # leave original index level name
+                names = [n[prefix_len:] for n in ds[dim].to_index().names]
+                ds = ds.reset_index(dim)
+                ds.attrs[f"{dim}_multiindex"] = list(names)
+
         return ds
 
     vars = [
@@ -583,11 +597,15 @@
         with_prefix(con.data, f"constraints-{name}")
         for name, con in m.constraints.items()
     ]
-    obj = [with_prefix(m.objective.data, "objective")]
+    objective = m.objective.data
+    objective = objective.assign_attrs(sense=m.objective.sense)
+    if m.objective.value is not None:
+        objective = objective.assign_attrs(value=m.objective.value)
+    obj = [with_prefix(objective, "objective")]
     params = [with_prefix(m.parameters, "params")]
 
     scalars = {k: getattr(m, k) for k in m.scalar_attrs}
-    ds = xr.merge(vars + cons + obj + params)
+    ds = xr.merge(vars + cons + obj + params, combine_attrs="drop_conflicts")
     ds = ds.assign_attrs(scalars)
     ds.attrs = non_bool_dict(ds.attrs)
 
@@ -624,30 +642,51 @@
     m = Model()
     ds = xr.load_dataset(path, **kwargs)
 
+    def has_prefix(k, prefix):
+        return k.rsplit("-", 1)[0] == prefix
+
+    def remove_prefix(k, prefix):
+        return k[len(prefix) + 1 :]
+
     def get_prefix(ds, prefix):
-        ds = ds[[k for k in ds if k.startswith(prefix)]]
-        ds = ds.rename({d: d.split(prefix + "-", 1)[1] for d in [*ds.dims, *ds]})
+        ds = ds[[k for k in ds if has_prefix(k, prefix)]]
+        to_rename = set([*ds.dims, *ds.coords, *ds])
+        ds = ds.rename({d: remove_prefix(d, prefix) for d in to_rename})
         ds.attrs = {
-            k.split(prefix + "-", 1)[1]: v
+            remove_prefix(k, prefix): v
             for k, v in ds.attrs.items()
-            if k.startswith(prefix + "-")
+            if has_prefix(k, prefix)
         }
+
+        for dim in ds.dims:
+            if f"{dim}_multiindex" in ds.attrs:
+                names = ds.attrs.pop(f"{dim}_multiindex")
+                ds = ds.set_index({dim: names})
+
         return ds
 
-    vars = get_prefix(ds, "variables")
-    var_names = list({k.split("-", 1)[0] for k in vars})
-    variables = {k: Variable(get_prefix(vars, k), m, k) for k in var_names}
+    vars = [k for k in ds if k.startswith("variables")]
+    var_names = list({k.rsplit("-", 1)[0] for k in vars})
+    variables = {}
+    for k in var_names:
+        name = remove_prefix(k, "variables")
+        variables[name] = Variable(get_prefix(ds, k), m, name)
     m._variables = Variables(variables, m)
 
-    cons = get_prefix(ds, "constraints")
-    con_names = list({k.split("-", 1)[0] for k in cons})
-    constraints = {k: Constraint(get_prefix(cons, k), m, k) for k in con_names}
+    cons = [k for k in ds if k.startswith("constraints")]
+    con_names = list({k.rsplit("-", 1)[0] for k in cons})
+    constraints = {}
+    for k in con_names:
+        name = remove_prefix(k, "constraints")
+        constraints[name] = Constraint(get_prefix(ds, k), m, name)
     m._constraints = Constraints(constraints, m)
 
     objective = get_prefix(ds, "objective")
-    m._objective = LinearExpression(objective, m)
+    m.objective = LinearExpression(objective, m)
+    m.objective.sense = objective.attrs.pop("sense")
+    m.objective._value = objective.attrs.pop("value", None)
 
-    m._parameters = get_prefix(ds, "parameter")
+    m.parameters = get_prefix(ds, "parameter")
 
     for k in m.scalar_attrs:
         setattr(m, k, ds.attrs.get(k))

linopy/testing.py

@@ -18,3 +18,24 @@ def assert_linequal(a, b):
 def assert_conequal(a, b):
     """Assert that two constraints are equal."""
     return assert_equal(_con_unwrap(a), _con_unwrap(b))
+
+
+def assert_model_equal(a, b):
+    """Assert that two models are equal."""
+    for k in a.dataset_attrs:
+        assert_equal(getattr(a, k), getattr(b, k))
+
+    assert set(a.variables) == set(b.variables)
+    assert set(a.constraints) == set(b.constraints)
+
+    for v in a.variables:
+        assert_varequal(a.variables[v], b.variables[v])
+
+    for c in a.constraints:
+        assert_conequal(a.constraints[c], b.constraints[c])
+
+    assert_linequal(a.objective.expression, b.objective.expression)
+    assert a.objective.sense == b.objective.sense
+    assert a.objective.value == b.objective.value
+
+    assert a.type == b.type

setup.py

@@ -67,6 +67,7 @@
         "dev": [
             "pytest",
             "pytest-cov",
+            "netcdf4",
             "pre-commit",
             "paramiko",
             "gurobipy",

test/test_io.py

@@ -12,10 +12,11 @@
 from xarray.testing import assert_equal
 
 from linopy import LESS_EQUAL, Model, available_solvers, read_netcdf
+from linopy.testing import assert_model_equal
 
 
 @pytest.fixture
-def m():
+def model():
     m = Model()
 
     x = m.add_variables(4, pd.Series([8, 10]), name="x")
@@ -28,36 +29,83 @@ def m():
     return m
 
 
-def test_to_netcdf(m, tmp_path):
+@pytest.fixture
+def model_with_dash_names():
+    m = Model()
+
+    x = m.add_variables(4, pd.Series([8, 10]), name="x-var")
+    x = m.add_variables(4, pd.Series([8, 10]), name="x-var-2")
+    y = m.add_variables(0, pd.DataFrame([[1, 2], [3, 4]]), name="y-var")
+
+    m.add_constraints(x + y, LESS_EQUAL, 10, name="constraint-1")
+
+    m.add_objective(2 * x + 3 * y)
+
+    return m
+
+
+@pytest.fixture
+def model_with_multiindex():
+    m = Model()
+
+    index = pd.MultiIndex.from_tuples(
+        [(1, "a"), (1, "b"), (2, "a"), (2, "b")], names=["first", "second"]
+    )
+    x = m.add_variables(4, pd.Series([8, 10, 12, 14], index=index), name="x-var")
+    y = m.add_variables(
+        0, pd.DataFrame([[1, 2], [3, 4], [5, 6], [7, 8]], index=index), name="y-var"
+    )
+
+    m.add_constraints(x + y, LESS_EQUAL, 10, name="constraint-1")
+
+    m.add_objective(2 * x + 3 * y)
+
+    return m
+
+
+def test_model_to_netcdf(model, tmp_path):
+    m = model
+    fn = tmp_path / "test.nc"
+    m.to_netcdf(fn)
+    p = read_netcdf(fn)
+
+    assert_model_equal(m, p)
+
+
+def test_model_to_netcdf_with_sense(model, tmp_path):
+    m = model
+    m.objective.sense = "max"
     fn = tmp_path / "test.nc"
     m.to_netcdf(fn)
     p = read_netcdf(fn)
 
-    for k in m.scalar_attrs:
-        if k != "objective.value":
-            assert getattr(m, k) == getattr(p, k)
+    assert_model_equal(m, p)
+
 
-    for k in m.dataset_attrs:
-        assert_equal(getattr(m, k), getattr(p, k))
+def test_model_to_netcdf_with_dash_names(model_with_dash_names, tmp_path):
+    m = model_with_dash_names
+    fn = tmp_path / "test.nc"
+    m.to_netcdf(fn)
+    p = read_netcdf(fn)
 
-    assert set(m.variables) == set(p.variables)
-    assert set(m.constraints) == set(p.constraints)
+    assert_model_equal(m, p)
 
-    for v in m.variables:
-        assert_equal(m.variables[v].data, p.variables[v].data)
 
-    for c in m.constraints:
-        assert_equal(m.constraints[c].data, p.constraints[c].data)
+def test_model_to_netcdf_with_multiindex(model_with_multiindex, tmp_path):
+    m = model_with_multiindex
+    fn = tmp_path / "test.nc"
+    m.to_netcdf(fn)
+    p = read_netcdf(fn)
 
-    assert_equal(m.objective.data, p.objective.data)
+    assert_model_equal(m, p)
 
 
 @pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobipy not installed")
-def test_to_file_lp(m, tmp_path):
+def test_to_file_lp(model, tmp_path):
     import gurobipy
 
     fn = tmp_path / "test.lp"
-    m.to_file(fn)
+    model.to_file(fn)
 
     gurobipy.read(str(fn))
 
@@ -66,30 +114,30 @@ def test_to_file_lp(m, tmp_path):
     not {"gurobi", "highs"}.issubset(available_solvers),
     reason="Gurobipy of highspy not installed",
 )
-def test_to_file_mps(m, tmp_path):
+def test_to_file_mps(model, tmp_path):
     import gurobipy
 
     fn = tmp_path / "test.mps"
-    m.to_file(fn)
+    model.to_file(fn)
 
     gurobipy.read(str(fn))
 
 
 @pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobipy not installed")
-def test_to_file_invalid(m, tmp_path):
+def test_to_file_invalid(model, tmp_path):
     with pytest.raises(ValueError):
         fn = tmp_path / "test.failedtype"
-        m.to_file(fn)
+        model.to_file(fn)
 
 
 @pytest.mark.skipif("gurobi" not in available_solvers, reason="Gurobipy not installed")
-def test_to_gurobipy(m):
-    m.to_gurobipy()
+def test_to_gurobipy(model):
+    model.to_gurobipy()
 
 
 @pytest.mark.skipif("highs" not in available_solvers, reason="Highspy not installed")
-def test_to_highspy(m):
-    m.to_highspy()
+def test_to_highspy(model):
+    model.to_highspy()
 
 
 def test_to_blocks(tmp_path):

test/test_model.py

@@ -12,6 +12,7 @@
 import xarray as xr
 
 from linopy import EQUAL, Model
+from linopy.testing import assert_model_equal
 
 target_shape = (10, 10)
 
@@ -111,3 +112,19 @@ def test_remove_constraint():
     m.add_constraints(x, EQUAL, 0, name="x")
     m.remove_constraints("x")
     assert not len(m.constraints.labels)
+
+
+def test_assert_model_equal():
+    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, name="x")
+    y = m.add_variables(name="y")
+
+    m.add_constraints(1 * x + 10 * y, EQUAL, 0)
+
+    obj = (10 * x + 5 * y).sum()
+    m.add_objective(obj)
+
+    assert_model_equal(m, m)