From 8eeb91cab023d45a4ec5ab5df1a307849bb5ece4 Mon Sep 17 00:00:00 2001 From: Mayk Thewessen Date: Mon, 5 Oct 2026 18:28:14 +0200 Subject: [PATCH 1/2] Preserve CSR backing during coordinate concatenation Concatenate constant and coordinate metadata with xarray, remap COO rows into the output grid, and retain explicit zeros. Keep overlapping labels and unsupported operands on the existing dense fallback. Verify 18 targeted v1 cases and frozen constraints. Focused suite: 710 passed, 662 skipped. Full-suite comparison resolves 14 regression cases with no new failing test identities; 80 baseline failures/errors remain in the existing environment. Independently verify exact equality of a 168-hour BirdFlow LP and removal of its coordinate-merge densification warning. Ruff and read-only review pass. --- doc/release_notes.rst | 2 + linopy/expressions.py | 82 +++++++++++- test/test_sparse_coordinate_concat.py | 175 ++++++++++++++++++++++++++ 3 files changed, 258 insertions(+), 1 deletion(-) create mode 100644 test/test_sparse_coordinate_concat.py diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 42e8bee4..c1dc0421 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -4,6 +4,8 @@ Release Notes Upcoming Version ---------------- +* Concatenating ordinary sparse linear expressions along an existing coordinate dimension now preserves sparse backing for unique disjoint labels. Remaining axes and auxiliary coordinates retain xarray join semantics. Unsupported cases keep the existing dense fallback. + *Strict "v1" arithmetic semantics (opt-in)* diff --git a/linopy/expressions.py b/linopy/expressions.py index 083e673f..11ed7517 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -110,7 +110,13 @@ STACKED_TERM_DIM, TERM_DIM, ) -from linopy.csr import CSRLinearExpression, Grid, _aux_coords, _densify_notice +from linopy.csr import ( + CSRLinearExpression, + Grid, + _aux_coords, + _densify_notice, + coo_to_csr, +) from linopy.semantics import ( AbsentType, FillValueLike, @@ -3709,6 +3715,76 @@ def _aligned( return [p.reindexed(grid, fill) for p in csrs] +def _try_csr_concat( + exprs: Any, dim: str, join: JoinOptions | None +) -> LinearExpression | None: + """Stack disjoint sparse rows, leaving labelled metadata alignment to xarray.""" + if not all(type(e) is LinearExpression and e._csr is not None for e in exprs): + return None + csrs = [e._csr for e in exprs] + template = csrs[0] + dims = template.grid.dims + if dim not in dims or any( + e.model is not exprs[0].model + or set(p.grid.dims) != set(dims) + or not p.grid.is_unique + or any(isinstance(index, pd.MultiIndex) for index in p.grid.indexes.values()) + for e, p in zip(exprs, csrs) + ): + return None + labels = template.grid.indexes[dim].append([p.grid.indexes[dim] for p in csrs[1:]]) + if not labels.is_unique: + return None + metadata = [ + p.grid.to_dataset().assign(const=(p.grid.dims, p.const.reshape(p.grid.shape))) + for p in csrs + ] + if join is None: + enforce_merge_dims(metadata, concat_dim=dim, context=f"merge along dim {dim!r}") + enforce_aux_conflict(metadata, concat_dim=dim) + combined = xr.concat( + metadata, + dim=dim, + join=join or "outer", + coords="minimal", + compat="override", + fill_value={"const": np.nan}, + ).transpose(*dims) + grid = Grid.from_dataset(combined, dims) + rows: list[np.ndarray] = [] + cols: list[np.ndarray] = [] + values: list[np.ndarray] = [] + for p in csrs: + source_grid = p.grid + if join == "override": + source_grid = replace( + source_grid, + indexes={ + d: source_grid.indexes[d] if d == dim else grid.indexes[d] + for d in source_grid.dims + }, + ) + row_map, valid = grid.indexer(source_grid) + coo = p.csr.tocoo() + keep = valid[coo.coords[0]] + rows.append(row_map[coo.coords[0][keep]]) + cols.append(coo.coords[1][keep]) + values.append(coo.data[keep]) + csr = coo_to_csr( + np.concatenate(values), + np.concatenate(rows), + np.concatenate(cols), + (grid.size, max(p.csr.shape[1] for p in csrs)), + template.model, + ) + return LinearExpression._from_csr( + CSRLinearExpression( + csr, combined.const.to_numpy().reshape(-1), grid, template.model + ), + template.model, + ) + + def _try_csr_merge( exprs: Any, dim: str, @@ -3729,6 +3805,10 @@ def _try_csr_merge( """ if not any(type(e) is LinearExpression and e._csr is not None for e in exprs): return None + if dim != TERM_DIM and not kwargs: + result = _try_csr_concat(exprs, dim, join) + if result is not None: + return result if dim != TERM_DIM or kwargs: _densify_all( exprs, "merge along a coordinate dimension or with extra arguments" diff --git a/test/test_sparse_coordinate_concat.py b/test/test_sparse_coordinate_concat.py new file mode 100644 index 00000000..b0a07f41 --- /dev/null +++ b/test/test_sparse_coordinate_concat.py @@ -0,0 +1,175 @@ +""" +Coordinate concatenation contracts: input backing, row order, joins, absence, +explicit zeros, ragged terms, metadata, empty grids, and unsupported fallbacks. +Term merging already owns addition coverage; these tests stack disjoint rows. +""" + +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +import linopy +from linopy.csr import CSRLinearExpression +from linopy.expressions import LinearExpression + + +def parts(model: linopy.Model, buses: list[list[str]]) -> list[LinearExpression]: + result = [] + for i, labels in enumerate(buses): + coords = [pd.Index([2 - i], name="snapshot"), pd.Index(labels, name="bus")] + x = model.add_variables(coords=coords) + dense = x * (i + 1) + (i + 3) + dense = dense.assign_coords(season=("snapshot", [f"s{i}"])) + result.append( + LinearExpression._from_csr( + CSRLinearExpression.from_dense(dense.data, model), model + ) + ) + return result + + +@pytest.mark.v1 +@pytest.mark.parametrize( + "join", ["outer", "inner", "left", "right", "exact", "override"] +) +def test_coordinate_concat_preserves_sparse_inputs_and_rows(join: str) -> None: + model = linopy.Model(sparse=True) + labels = [["b", "a"], ["b", "a"]] if join == "exact" else [["b", "a"], ["c", "b"]] + sparse = parts(model, labels) + dense = [e._csr.to_dense() for e in sparse] + expected = linopy.merge(dense, dim="snapshot", join=join) + result = linopy.merge(sparse, dim="snapshot", join=join) + assert result.is_sparse and all(e.is_sparse for e in sparse) + actual = result._csr.to_dense() + xr.testing.assert_equal(actual.const, expected.const) + xr.testing.assert_equal(actual.coords.to_dataset(), expected.coords.to_dataset()) + np.testing.assert_array_equal( + result._csr.const.reshape(result._csr.grid.shape)[0], expected.const.values[0] + ) + for row, snapshot in enumerate([2, 1]): + for bus in result._csr.grid.indexes["bus"]: + cell = result._csr.cell((row, result._csr.grid.indexes["bus"].get_loc(bus))) + if join == "override" or bus in labels[row]: + np.testing.assert_array_equal(cell[0], [row + 1]) + source_position = ( + result._csr.grid.indexes["bus"].get_loc(bus) + if join == "override" + else labels[row].index(bus) + ) + np.testing.assert_array_equal(cell[1], [row * 2 + source_position]) + assert cell[2] == row + 3 + else: + assert np.isnan(cell[2]) and len(cell[0]) == 0 + + +@pytest.mark.v1 +@pytest.mark.parametrize("empty", ["snapshot", "bus", None]) +def test_ragged_concat_empty_absence_and_frozen_matrix(empty: str | None) -> None: + model = linopy.Model(sparse=True) + blocks = parts(model, [["a", "b"], ["a", "b"]]) + # A real zero term remains stored, while a masked cell remains absent. + first = blocks[0]._csr + from dataclasses import replace + + matrix = first.csr.copy() + matrix.data[0] = 0.0 + blocks[0] = LinearExpression._from_csr(replace(first, csr=matrix), model) + y = model.add_variables( + coords=[pd.Index([1], name="snapshot"), pd.Index(["a", "b"], name="bus")] + ) + second_dense = blocks[1]._csr.to_dense() + 5 * y + second = CSRLinearExpression.from_dense(second_dense.data, model) + blocks[1] = LinearExpression._from_csr( + second.with_const(np.array([np.nan, 9.0])), model + ) + # Reverse dimension order without expanding term rectangles. + blocks[1] = LinearExpression._from_csr( + blocks[1]._csr.reindexed(blocks[1]._csr.grid.reordered(("bus", "snapshot"))), + model, + ) + if empty is not None: + blocks[1] = blocks[1].isel({empty: slice(0, 0)}) + dense = [b._csr.to_dense() for b in blocks] + expected = linopy.merge(dense, dim="snapshot", join="outer") + actual = linopy.merge(blocks, dim="snapshot", join="outer") + assert actual.is_sparse and all(b.is_sparse for b in blocks) + xr.testing.assert_equal(actual._csr.to_dense().const, expected.const) + assert np.count_nonzero(actual._csr.csr.data == 0) == 1 + dense_csr = CSRLinearExpression.from_dense(expected.data, model) + np.testing.assert_array_equal(actual._csr.csr.toarray(), dense_csr.csr.toarray()) + np.testing.assert_array_equal(actual._csr.const, dense_csr.const) + # Frozen constraints exercise the solver-facing export, including masked rows. + left = model.add_constraints(actual == 0, name="sparse", freeze=True) + right = model.add_constraints(expected == 0, name="dense", freeze=True) + np.testing.assert_array_equal(left.rhs.values, right.rhs.values) + from linopy.testing import assert_conequal + + assert_conequal(left, right, strict=False) + + +@pytest.mark.v1 +@pytest.mark.parametrize("unsupported", ["overlap", "repeated", "mixed", "kwargs"]) +def test_unsupported_concat_keeps_observable_dense_fallback(unsupported: str) -> None: + from linopy.constants import PerformanceWarning + + model = linopy.Model(sparse=True) + blocks = parts(model, [["a"], ["a"]]) + kwargs = {} + if unsupported == "overlap": + blocks[1] = blocks[1].reindex(snapshot=[2]) + elif unsupported == "repeated": + blocks[0] = blocks[0].isel(snapshot=[0, 0]) + elif unsupported == "mixed": + blocks[1] = blocks[1]._csr.to_dense() + else: + kwargs = {"compat": "override", "coords": "minimal"} + dense = [b._csr.to_dense() if b.is_sparse else b for b in blocks] + expected = linopy.merge(dense, dim="snapshot", join="outer", **kwargs) + with linopy.options as options: + options.set_value(warn_on_densify=True) + with pytest.warns(PerformanceWarning, match="merge along a coordinate"): + actual = linopy.merge(blocks, dim="snapshot", join="outer", **kwargs) + assert not actual.is_sparse + xr.testing.assert_equal(actual.data, expected.data) + + +@pytest.mark.v1 +@pytest.mark.parametrize("size", [0, 1]) +def test_single_block_preserves_empty_grid_and_scalar_metadata(size: int) -> None: + model = linopy.Model(sparse=True) + block = parts(model, [["a", "b"]])[0].isel(snapshot=slice(0, size)) + from dataclasses import replace + + csr = block._csr + grid = replace(csr.grid, aux=csr.grid.aux | {"country": ((), np.array("NL"))}) + block = LinearExpression._from_csr(replace(csr, grid=grid), model) + expected = block._csr.to_dense() + actual = linopy.merge([block], dim="snapshot", join="outer") + assert actual.is_sparse and block.is_sparse + xr.testing.assert_equal(actual._csr.to_dense().data, expected.data) + + +@pytest.mark.v1 +@pytest.mark.parametrize("contract", ["exact", "override", "auxiliary"]) +def test_coordinate_alignment_errors_preserve_inputs(contract: str) -> None: + model = linopy.Model(sparse=True) + blocks = parts(model, [["a", "b"], ["b", "c"]]) + join = "outer" + if contract == "exact": + join = "exact" + elif contract == "override": + join = "override" + blocks[1] = blocks[1].isel(bus=[0]) + else: + for i, country in enumerate(["NL", "DE"]): + dense = blocks[i]._csr.to_dense().assign_coords(country=country) + blocks[i] = LinearExpression._from_csr( + CSRLinearExpression.from_dense(dense.data, model), model + ) + expected = [b._csr.to_dense() for b in blocks] + with pytest.raises(ValueError): + linopy.merge(expected, dim="snapshot", join=join) + with pytest.raises(ValueError): + linopy.merge(blocks, dim="snapshot", join=join) + assert all(b.is_sparse for b in blocks) From e3db8f403b664732a44ae9ad78598c601f1d08c0 Mon Sep 17 00:00:00 2001 From: Mayk Thewessen Date: Mon, 5 Oct 2026 21:06:08 +0200 Subject: [PATCH 2/2] Narrow sparse coordinate concatenation test types --- test/test_sparse_coordinate_concat.py | 86 +++++++++++++++++---------- 1 file changed, 53 insertions(+), 33 deletions(-) diff --git a/test/test_sparse_coordinate_concat.py b/test/test_sparse_coordinate_concat.py index b0a07f41..560acdd9 100644 --- a/test/test_sparse_coordinate_concat.py +++ b/test/test_sparse_coordinate_concat.py @@ -4,16 +4,25 @@ Term merging already owns addition coverage; these tests stack disjoint rows. """ +from typing import Any + import numpy as np import pandas as pd import pytest import xarray as xr +from xarray.core.types import JoinOptions import linopy from linopy.csr import CSRLinearExpression from linopy.expressions import LinearExpression +def backing(expression: LinearExpression) -> CSRLinearExpression: + """Require the sparse backing protected by these tests.""" + assert expression._csr is not None + return expression._csr + + def parts(model: linopy.Model, buses: list[list[str]]) -> list[LinearExpression]: result = [] for i, labels in enumerate(buses): @@ -33,27 +42,30 @@ def parts(model: linopy.Model, buses: list[list[str]]) -> list[LinearExpression] @pytest.mark.parametrize( "join", ["outer", "inner", "left", "right", "exact", "override"] ) -def test_coordinate_concat_preserves_sparse_inputs_and_rows(join: str) -> None: +def test_coordinate_concat_preserves_sparse_inputs_and_rows(join: JoinOptions) -> None: model = linopy.Model(sparse=True) labels = [["b", "a"], ["b", "a"]] if join == "exact" else [["b", "a"], ["c", "b"]] sparse = parts(model, labels) - dense = [e._csr.to_dense() for e in sparse] - expected = linopy.merge(dense, dim="snapshot", join=join) - result = linopy.merge(sparse, dim="snapshot", join=join) + dense = [backing(e).to_dense() for e in sparse] + expected = linopy.merge(dense, cls=LinearExpression, dim="snapshot", join=join) + result = linopy.merge(sparse, cls=LinearExpression, dim="snapshot", join=join) assert result.is_sparse and all(e.is_sparse for e in sparse) - actual = result._csr.to_dense() + actual = backing(result).to_dense() xr.testing.assert_equal(actual.const, expected.const) xr.testing.assert_equal(actual.coords.to_dataset(), expected.coords.to_dataset()) np.testing.assert_array_equal( - result._csr.const.reshape(result._csr.grid.shape)[0], expected.const.values[0] + backing(result).const.reshape(backing(result).grid.shape)[0], + expected.const.values[0], ) for row, snapshot in enumerate([2, 1]): - for bus in result._csr.grid.indexes["bus"]: - cell = result._csr.cell((row, result._csr.grid.indexes["bus"].get_loc(bus))) + for bus in backing(result).grid.indexes["bus"]: + cell = backing(result).cell( + (row, backing(result).grid.indexes["bus"].get_loc(bus)) + ) if join == "override" or bus in labels[row]: np.testing.assert_array_equal(cell[0], [row + 1]) source_position = ( - result._csr.grid.indexes["bus"].get_loc(bus) + backing(result).grid.indexes["bus"].get_loc(bus) if join == "override" else labels[row].index(bus) ) @@ -69,7 +81,7 @@ def test_ragged_concat_empty_absence_and_frozen_matrix(empty: str | None) -> Non model = linopy.Model(sparse=True) blocks = parts(model, [["a", "b"], ["a", "b"]]) # A real zero term remains stored, while a masked cell remains absent. - first = blocks[0]._csr + first = backing(blocks[0]) from dataclasses import replace matrix = first.csr.copy() @@ -78,27 +90,31 @@ def test_ragged_concat_empty_absence_and_frozen_matrix(empty: str | None) -> Non y = model.add_variables( coords=[pd.Index([1], name="snapshot"), pd.Index(["a", "b"], name="bus")] ) - second_dense = blocks[1]._csr.to_dense() + 5 * y + second_dense = backing(blocks[1]).to_dense() + 5 * y second = CSRLinearExpression.from_dense(second_dense.data, model) blocks[1] = LinearExpression._from_csr( second.with_const(np.array([np.nan, 9.0])), model ) # Reverse dimension order without expanding term rectangles. blocks[1] = LinearExpression._from_csr( - blocks[1]._csr.reindexed(blocks[1]._csr.grid.reordered(("bus", "snapshot"))), + backing(blocks[1]).reindexed( + backing(blocks[1]).grid.reordered(("bus", "snapshot")) + ), model, ) if empty is not None: blocks[1] = blocks[1].isel({empty: slice(0, 0)}) - dense = [b._csr.to_dense() for b in blocks] - expected = linopy.merge(dense, dim="snapshot", join="outer") - actual = linopy.merge(blocks, dim="snapshot", join="outer") + dense = [backing(b).to_dense() for b in blocks] + expected = linopy.merge(dense, cls=LinearExpression, dim="snapshot", join="outer") + actual = linopy.merge(blocks, cls=LinearExpression, dim="snapshot", join="outer") assert actual.is_sparse and all(b.is_sparse for b in blocks) - xr.testing.assert_equal(actual._csr.to_dense().const, expected.const) - assert np.count_nonzero(actual._csr.csr.data == 0) == 1 + xr.testing.assert_equal(backing(actual).to_dense().const, expected.const) + assert np.count_nonzero(backing(actual).csr.data == 0) == 1 dense_csr = CSRLinearExpression.from_dense(expected.data, model) - np.testing.assert_array_equal(actual._csr.csr.toarray(), dense_csr.csr.toarray()) - np.testing.assert_array_equal(actual._csr.const, dense_csr.const) + np.testing.assert_array_equal( + backing(actual).csr.toarray(), dense_csr.csr.toarray() + ) + np.testing.assert_array_equal(backing(actual).const, dense_csr.const) # Frozen constraints exercise the solver-facing export, including masked rows. left = model.add_constraints(actual == 0, name="sparse", freeze=True) right = model.add_constraints(expected == 0, name="dense", freeze=True) @@ -115,21 +131,25 @@ def test_unsupported_concat_keeps_observable_dense_fallback(unsupported: str) -> model = linopy.Model(sparse=True) blocks = parts(model, [["a"], ["a"]]) - kwargs = {} + kwargs: dict[str, Any] = {} if unsupported == "overlap": blocks[1] = blocks[1].reindex(snapshot=[2]) elif unsupported == "repeated": blocks[0] = blocks[0].isel(snapshot=[0, 0]) elif unsupported == "mixed": - blocks[1] = blocks[1]._csr.to_dense() + blocks[1] = backing(blocks[1]).to_dense() else: kwargs = {"compat": "override", "coords": "minimal"} - dense = [b._csr.to_dense() if b.is_sparse else b for b in blocks] - expected = linopy.merge(dense, dim="snapshot", join="outer", **kwargs) + dense = [backing(b).to_dense() if b.is_sparse else b for b in blocks] + expected = linopy.merge( + dense, cls=LinearExpression, dim="snapshot", join="outer", **kwargs + ) with linopy.options as options: options.set_value(warn_on_densify=True) with pytest.warns(PerformanceWarning, match="merge along a coordinate"): - actual = linopy.merge(blocks, dim="snapshot", join="outer", **kwargs) + actual = linopy.merge( + blocks, cls=LinearExpression, dim="snapshot", join="outer", **kwargs + ) assert not actual.is_sparse xr.testing.assert_equal(actual.data, expected.data) @@ -141,13 +161,13 @@ def test_single_block_preserves_empty_grid_and_scalar_metadata(size: int) -> Non block = parts(model, [["a", "b"]])[0].isel(snapshot=slice(0, size)) from dataclasses import replace - csr = block._csr + csr = backing(block) grid = replace(csr.grid, aux=csr.grid.aux | {"country": ((), np.array("NL"))}) block = LinearExpression._from_csr(replace(csr, grid=grid), model) - expected = block._csr.to_dense() - actual = linopy.merge([block], dim="snapshot", join="outer") + expected = backing(block).to_dense() + actual = linopy.merge([block], cls=LinearExpression, dim="snapshot", join="outer") assert actual.is_sparse and block.is_sparse - xr.testing.assert_equal(actual._csr.to_dense().data, expected.data) + xr.testing.assert_equal(backing(actual).to_dense().data, expected.data) @pytest.mark.v1 @@ -155,7 +175,7 @@ def test_single_block_preserves_empty_grid_and_scalar_metadata(size: int) -> Non def test_coordinate_alignment_errors_preserve_inputs(contract: str) -> None: model = linopy.Model(sparse=True) blocks = parts(model, [["a", "b"], ["b", "c"]]) - join = "outer" + join: JoinOptions = "outer" if contract == "exact": join = "exact" elif contract == "override": @@ -163,13 +183,13 @@ def test_coordinate_alignment_errors_preserve_inputs(contract: str) -> None: blocks[1] = blocks[1].isel(bus=[0]) else: for i, country in enumerate(["NL", "DE"]): - dense = blocks[i]._csr.to_dense().assign_coords(country=country) + dense = backing(blocks[i]).to_dense().assign_coords(country=country) blocks[i] = LinearExpression._from_csr( CSRLinearExpression.from_dense(dense.data, model), model ) - expected = [b._csr.to_dense() for b in blocks] + expected = [backing(b).to_dense() for b in blocks] with pytest.raises(ValueError): - linopy.merge(expected, dim="snapshot", join=join) + linopy.merge(expected, cls=LinearExpression, dim="snapshot", join=join) with pytest.raises(ValueError): - linopy.merge(blocks, dim="snapshot", join=join) + linopy.merge(blocks, cls=LinearExpression, dim="snapshot", join=join) assert all(b.is_sparse for b in blocks)