From ce2475d0a00a9eec9eaeb8faa30bd5d5f8c454b2 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 10 Nov 2023 11:34:36 +0100 Subject: [PATCH 1/5] io: fix dashes in names and objective storage, support multiindexes repr: fix single list-like indexed representation --- doc/release_notes.rst | 9 +++- linopy/common.py | 1 + linopy/constraints.py | 3 +- linopy/expressions.py | 3 +- linopy/io.py | 71 ++++++++++++++++++++++------- linopy/testing.py | 21 +++++++++ test/test_io.py | 96 +++++++++++++++++++++++++++++---------- test/test_model.py | 17 +++++++ test/test_repr.py | 101 +++++++++++++++++++++++++----------------- 9 files changed, 237 insertions(+), 85 deletions(-) diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 2ecb044d..71328c0c 100644 --- a/doc/release_notes.rst +++ b/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 ------------- diff --git a/linopy/common.py b/linopy/common.py index f851d493..496be869 100644 --- a/linopy/common.py +++ b/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] diff --git a/linopy/constraints.py b/linopy/constraints.py index 3551cd16..6aca342b 100644 --- a/linopy/constraints.py +++ b/linopy/constraints.py @@ -165,6 +165,7 @@ 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 @@ -172,7 +173,7 @@ def __repr__(self): 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...") diff --git a/linopy/expressions.py b/linopy/expressions.py index 6619eb92..dea3a677 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -330,6 +330,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()) 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 @@ -337,7 +338,7 @@ def __repr__(self): 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...") diff --git a/linopy/io.py b/linopy/io.py index 4536a9dd..b32ad775 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -298,7 +298,11 @@ def to_file(m, fn, integer_label="general"): 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 to_netcdf(m, *args, **kwargs): """ 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 @@ def with_prefix(ds, prefix): 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 @@ def read_netcdf(path, **kwargs): 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)) diff --git a/linopy/testing.py b/linopy/testing.py index 4d727e54..51a6e97e 100644 --- a/linopy/testing.py +++ b/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 diff --git a/test/test_io.py b/test/test_io.py index 780ff57f..7c5cc099 100644 --- a/test/test_io.py +++ b/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): diff --git a/test/test_model.py b/test/test_model.py index 59755af8..ab3ee3b9 100644 --- a/test/test_model.py +++ b/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) diff --git a/test/test_repr.py b/test/test_repr.py index e3d0b994..a3f63a41 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -76,63 +76,82 @@ cu_masked = m.add_constraints(cu_, name="cu_masked", mask=xr.full_like(u.labels, False)) -def test_variable_repr(): - for var in [u, v, x, y, z, a, b, c, d, e, f]: - repr(var) +variables = [u, v, x, y, z, a, b, c, d, e, f] +expressions = [lu, lv, lx, ly, lz, la, lb, lc, ld, lav, luc, lq, lq2, lq3] +anonymous_constraints = [cu_, cv_, cx_, cy_, cz_, ca_, cb_, cc_, cd_, cav_, cuc_] +constraints = [cu, cv, cx, cy, cz, ca, cb, cc, cd, cav, cuc, cu_masked] -def test_scalar_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[coord]) +@pytest.mark.parametrize("var", variables) +def test_variable_repr(var): + repr(var) -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]) +@pytest.mark.parametrize("var", variables) +def test_scalar_variable_repr(var): + coord = tuple(var.indexes[c][0] for c in var.dims) + repr(var[coord]) -def test_linear_expression_repr(): - for expr in [lu, lv, lx, ly, lz, la, lb, lc, ld, lav, luc, lq, lq2, lq3]: - repr(expr) +@pytest.mark.parametrize("var", variables) +def test_single_variable_repr(var): + coord = tuple(var.indexes[c][0] for c in var.dims) + repr(var.loc[coord]) + + +@pytest.mark.parametrize("expr", expressions) +def test_linear_expression_repr(expr): + repr(expr) def test_linear_expression_long(): repr(x.sum()) -def test_scalar_linear_expression_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(1 * var[coord]) +@pytest.mark.parametrize("var", variables) +def test_scalar_linear_expression_repr(var): + coord = tuple(var.indexes[c][0] for c in var.dims) + repr(1 * var[coord]) + +@pytest.mark.parametrize("var", variables) +def test_single_linear_repr(var): + coord = tuple(var.indexes[c][0] for c in var.dims) + repr(1 * var.loc[coord]) -def test_single_linear_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(1 * var.loc[coord]) +@pytest.mark.parametrize("var", variables) +def test_single_array_linear_repr(var): + coord = [[var.indexes[c][0]] for c in var.dims] + repr(1 * var.loc[*coord]) -def test_anonymous_constraint_repr(): - for con in [cu_, cv_, cx_, cy_, cz_, ca_, cb_, cc_, cd_, cav_, cuc_]: - repr(con) + +@pytest.mark.parametrize("con", anonymous_constraints) +def test_anonymous_constraint_repr(con): + repr(con) def test_scalar_constraint_repr(): repr(u[0, 0] >= 0) -def test_single_constraint_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] == 0) - repr(1 * var.loc[coord] - var.loc[coord] == 0) +@pytest.mark.parametrize("var", variables) +def test_single_constraint_repr(var): + coord = tuple(var.indexes[c][0] for c in var.dims) + repr(var.loc[coord] == 0) + repr(1 * var.loc[coord] - var.loc[coord] == 0) + + +@pytest.mark.parametrize("var", variables) +def test_single_array_constraint_repr(var): + coord = [[var.indexes[c][0]] for c in var.dims] + repr(var.loc[*coord] == 0) + repr(1 * var.loc[*coord] - var.loc[*coord] == 0) -def test_constraint_repr(): - for con in [cu, cv, cx, cy, cz, ca, cb, cc, cd, cav, cuc, cu_masked]: - repr(con) +@pytest.mark.parametrize("con", constraints) +def test_constraint_repr(con): + repr(con) def test_empty_repr(): @@ -141,15 +160,15 @@ def test_empty_repr(): repr(lu.sel(dim_0=[]) >= 0) -def test_print_options(): - for o in [v, lv, cv_, cv]: - default_repr = repr(o) - with options as opts: - opts.set_value(display_max_rows=20) - longer_repr = repr(o) - assert len(default_repr) < len(longer_repr) +@pytest.mark.parametrize("obj", [v, lv, cv_, cv]) +def test_print_options(obj): + default_repr = repr(obj) + with options as opts: + opts.set_value(display_max_rows=20) + longer_repr = repr(obj) + assert len(default_repr) < len(longer_repr) - longer_repr = o.print(display_max_rows=20) + longer_repr = obj.print(display_max_rows=20) def test_print_labels(): From b14b57cc05ebd1efc4f1520d6415bd0b5665286b Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 10 Nov 2023 12:23:26 +0100 Subject: [PATCH 2/5] test_repr: flake8 syntax --- test/test_repr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_repr.py b/test/test_repr.py index a3f63a41..74a7981a 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -123,7 +123,7 @@ def test_single_linear_repr(var): @pytest.mark.parametrize("var", variables) def test_single_array_linear_repr(var): coord = [[var.indexes[c][0]] for c in var.dims] - repr(1 * var.loc[*coord]) + repr(1 * var.loc[*coord]) # noqa: E999 @pytest.mark.parametrize("con", anonymous_constraints) From ca2c95271ab45ba664b54794dd8ab1c21fe4b1f5 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 10 Nov 2023 12:28:24 +0100 Subject: [PATCH 3/5] test_repr: follow up --- test/test_repr.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_repr.py b/test/test_repr.py index 74a7981a..66c9cb9f 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -122,8 +122,8 @@ def test_single_linear_repr(var): @pytest.mark.parametrize("var", variables) def test_single_array_linear_repr(var): - coord = [[var.indexes[c][0]] for c in var.dims] - repr(1 * var.loc[*coord]) # noqa: E999 + coord = {c: [var.indexes[c][0]] for c in var.dims} + repr(1 * var.sel(coord)) @pytest.mark.parametrize("con", anonymous_constraints) From 56b81e81d953c2a2af53c7d5f9561cf288ed4dd4 Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 10 Nov 2023 13:00:33 +0100 Subject: [PATCH 4/5] follow up --- test/test_repr.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_repr.py b/test/test_repr.py index 66c9cb9f..0fe95329 100644 --- a/test/test_repr.py +++ b/test/test_repr.py @@ -144,9 +144,9 @@ def test_single_constraint_repr(var): @pytest.mark.parametrize("var", variables) def test_single_array_constraint_repr(var): - coord = [[var.indexes[c][0]] for c in var.dims] - repr(var.loc[*coord] == 0) - repr(1 * var.loc[*coord] - var.loc[*coord] == 0) + coord = {c: [var.indexes[c][0]] for c in var.dims} + repr(var.sel(coord) == 0) + repr(1 * var.sel(coord) - var.sel(coord) == 0) @pytest.mark.parametrize("con", constraints) From 2ac42c59371100ea2cb9ed72c74907cc641b3cca Mon Sep 17 00:00:00 2001 From: Fabian Date: Fri, 10 Nov 2023 13:28:45 +0100 Subject: [PATCH 5/5] setup: add netcdf4 for backend --- setup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.py b/setup.py index 85055536..37c1ec06 100644 --- a/setup.py +++ b/setup.py @@ -67,6 +67,7 @@ "dev": [ "pytest", "pytest-cov", + "netcdf4", "pre-commit", "paramiko", "gurobipy",