From 959ce51809ddb8c5ae80a5b50ce89744c2a22d0a Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 21 Sep 2026 07:17:09 +0000 Subject: [PATCH] feat(language): a where counts the coordinates a predicate admits, and reads one at a neighbour The pin follows math-spec #592, merged onto #589 so the stack keeps the assumptions this branch's base consumes: `count(, over=)` answers a number and `shift(, along=, offset=)` answers a predicate, on both lanes. The relational lane reads the counted predicate over its own product and joins the total back, so a coordinate it admits nowhere counts zero rather than losing its row; a shifted predicate carries its truth set forward along the axis, and a position the move runs off has no row to join, which is the false the language gives a vacated coordinate. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01TWu2WSipjg5T3khNLfTXDK --- docs/about/linopy.md | 2 + pyproject.toml | 2 +- src/lpspec/linopy/where.py | 8 ++ .../relational/engines/polars/predicates.py | 81 ++++++++++- tests/test_resolution_parity.py | 7 + tests/test_where_count.py | 136 ++++++++++++++++++ uv.lock | 6 +- 7 files changed, 232 insertions(+), 10 deletions(-) create mode 100644 tests/test_where_count.py diff --git a/docs/about/linopy.md b/docs/about/linopy.md index 4240496c4..a6b7d1d05 100644 --- a/docs/about/linopy.md +++ b/docs/about/linopy.md @@ -114,6 +114,8 @@ in `linopy/builder.py`, one section per group below. | `defined(x)` | `Model.variables['x'].labels != -1`, linopy's own marker for an absent slot | | a comparison | the Python comparison operators element-wise, absence reading as false | | a comparison of expressions | each side evaluated as a constant position is, and compared element-wise; a side with no value at a coordinate reads as false there, and under a sum, a window or a `+` an absent term is one fewer | +| `count(, over=)` | the predicate summed over that dimension, which reduces it away, and the total compared against the literal | +| `shift(, along=, offset=)` | `DataArray.shift` with `fill_value=False`, so a vacated position is false and no `edge=` is read | Absence has no single row. It is positional: a missing parameter row is zero in a coefficient, an error in `bounds:`, and false in a `where` operand. `linopy/absence.py` holds all four spellings, and the builder calls them diff --git a/pyproject.toml b/pyproject.toml index acdc33fcb..de8256ac8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ # A direct reference until math-spec publishes, which is a thing PyPI refuses # (RELEASING.md) — the same wall the `linopy` extra is behind, with the same # exit: an ordinary floor the day there is a release to floor against. - "math-spec @ git+https://github.com/energy-models/math-spec@5cb06abf3401d816fa371246f54760b56fd74bbb", + "math-spec @ git+https://github.com/energy-models/math-spec@2a9d0d599d3f5fd60acddec991412ff58b05110f", # the relational engine: every frame a model is built from, and the parquet # reader that fills them. Arrow-backed, so a result frame exports the # PyCapsule protocol and any consumer takes it without this package diff --git a/src/lpspec/linopy/where.py b/src/lpspec/linopy/where.py index 5537cfd52..5d9c2b005 100644 --- a/src/lpspec/linopy/where.py +++ b/src/lpspec/linopy/where.py @@ -99,6 +99,14 @@ def evaluate(child: program.Predicate) -> xr.DataArray: if isinstance(node, program.VariableDefined): return absence.present(ctx.model, node.name) + if isinstance(node, program.CountComparison): + counted = evaluate(node.predicate.root).sum(dim=node.over) + return _PREDICATE_OPS[node.op](counted, node.value).astype(bool) + + if isinstance(node, program.TranslatedPredicate): + moved = evaluate(node.operand.root).shift({node.along: node.offset}, fill_value=False) + return moved.astype(bool) + if isinstance(node, program.ExpressionComparison): from lpspec.linopy.builder import _eval # mask ↔ expression recursion diff --git a/src/lpspec/relational/engines/polars/predicates.py b/src/lpspec/relational/engines/polars/predicates.py index bdaa22a6f..bacea9139 100644 --- a/src/lpspec/relational/engines/polars/predicates.py +++ b/src/lpspec/relational/engines/polars/predicates.py @@ -6,11 +6,12 @@ built first and the frame read after. A closed vocabulary of its own — comparisons against a parameter, a dimension -label, a position along a dimension, a relation, arithmetic over parameters, -and the three connectives. It -takes the :class:`~lpspec.relational.engines.polars.scope.Scope` as an -argument and holds nothing. :func:`masked` is the product a declaration is -instantiated over, cut by its mask: the one place the two meet. +label, a position along a dimension, a relation, arithmetic over parameters, a +count of what a predicate admits, a predicate read at a neighbour, and the three +connectives. It takes the +:class:`~lpspec.relational.engines.polars.scope.Scope` as an argument and holds +nothing. :func:`masked` is the product a declaration is instantiated over, cut +by its mask: the one place the two meet. :class:`Carrier` lives here too, and the bounds walk imports it: both walks that read parameters build an expression over columns they are joining on as @@ -30,7 +31,7 @@ if TYPE_CHECKING: import datetime - from collections.abc import Callable + from collections.abc import Callable, Iterable from polars._typing import JoinStrategy @@ -77,6 +78,18 @@ def _defined(col: pl.Expr, dtype: program.ParameterDtype) -> pl.Expr: return col.is_not_null() & col.is_finite() +#: Column names the count and the shift work in, inside their own frames: the +#: count, a dimension's ordinal, the label the move lands on, and the truth a +#: moved coordinate carries. +_COUNT, _ORD, _MOVED, _TRUE = '__count__', '__ord__', '__moved__', '__true__' + + +def _in_declaration_order(scope: Scope, dims: Iterable[str]) -> tuple[str, ...]: + """*dims* in the program's declaration order, which is the order a product is built in.""" + read = set(dims) + return tuple(d for d in scope.program.dimensions if d in read) + + def masked(scope: Scope, dims: tuple[str, ...], where: program.Mask | None) -> pl.LazyFrame: """The masked coordinate product over *dims*. @@ -229,7 +242,63 @@ def attach(frame: pl.LazyFrame, alias: str) -> pl.LazyFrame: return carrier.once(f'__where expression {expression!r}__', attach) + def join_count(p: program.CountComparison) -> str: + """One column: how many coordinates along ``over`` the counted predicate admits, per coordinate of ``dims``. + + The counted predicate is read over its own product — ``over`` plus what + is left — so the count is taken where the predicate is true and joined + back onto the frame. A coordinate the predicate admits nowhere has no + row to join and counts zero, which is the answer rather than a gap. + """ + for dim in p.dims: + refuse_outside_frame(f"count over '{p.over}' reading dimension '{dim}'", dim) + counted = _in_declaration_order(scope, (p.over, *p.dims)) + admitted = masked(scope, counted, p.predicate) + totals = ( + admitted.group_by(list(p.dims)).agg(pl.len().alias(_COUNT)) + if p.dims + else admitted.select(pl.len().alias(_COUNT)) + ) + return carrier.once( + f'__where count {p.over} {p.predicate!r}__', + lambda f, alias: join_on(f, totals.rename({_COUNT: alias}), p.dims, 'left'), + ) + + def join_translated(p: program.TranslatedPredicate) -> str: + """One boolean column: the operand's truth set moved *offset* along the axis, null where the translation vacates. + + The operand is read over its own product and the coordinates it admits + are carried forward by ``offset`` positions in the axis' own order. A + position the move runs off the end of has no row to carry, so it joins + null and the caller reads it as false — which is what a missing row + means in any mask, and why the predicate form states no ``edge``. + """ + read = _in_declaration_order(scope, p.dims) + for dim in read: + refuse_outside_frame(f"shift along '{p.along}' reading dimension '{dim}'", dim) + axis = scope.data.dimensions[p.along].lazy() + at = axis.select(pl.col('val').alias(p.along), pl.col('ord').alias(_ORD)) + landing = axis.select(pl.col('val').alias(_MOVED), pl.col('ord').alias(_ORD)) + moved = ( + masked(scope, read, p.operand) + .select(*read) + .join(at, on=p.along, how='inner') + .with_columns((pl.col(_ORD) + p.offset).alias(_ORD)) + .join(landing, on=_ORD, how='inner') + .drop(_ORD, p.along) + .rename({_MOVED: p.along}) + .with_columns(pl.lit(value=True).alias(_TRUE)) + ) + return carrier.once( + f'__where shift {p.along} {p.offset} {p.operand!r}__', + lambda f, alias: join_on(f, moved.rename({_TRUE: alias}), read, 'left'), + ) + def walk(p: program.Predicate) -> pl.Expr: + if isinstance(p, program.CountComparison): + return _COLUMN_COMPARISONS[p.op](pl.col(join_count(p)).fill_null(0), pl.lit(p.value)) + if isinstance(p, program.TranslatedPredicate): + return falsy_if_null(pl.col(join_translated(p))) if isinstance(p, program.ExpressionComparison): left, right = (_side(p.left, p.dims), _side(p.right, p.dims)) return falsy_if_null(_COLUMN_COMPARISONS[p.op](pl.col(left), pl.col(right))) diff --git a/tests/test_resolution_parity.py b/tests/test_resolution_parity.py index 4c99b2298..30cfcde57 100644 --- a/tests/test_resolution_parity.py +++ b/tests/test_resolution_parity.py @@ -90,6 +90,13 @@ def test_both_lanes_refuse_a_comparison_that_carries_no_variable(tmp_path, dispa #: carries no dimension. The eager lane read the result as an array and got #: a plain bool, which has no `fillna`. '1 > 0', + #: A count reduces the dim it counts along away, so both of these carry no + #: dimension at all: one number against a literal. + 'count(p_max > 0, over=generator) >= 2', + 'count(load > 0, over=snapshot) >= 1', + #: A predicate read one coordinate back, false where the translation + #: vacates — so this is true at the first snapshot alone. + 'load AND NOT shift(load, along=snapshot, offset=1)', ] #: Predicates this sweep cannot host, with where they are checked instead. The diff --git a/tests/test_where_count.py b/tests/test_where_count.py new file mode 100644 index 000000000..fa76b54ba --- /dev/null +++ b/tests/test_where_count.py @@ -0,0 +1,136 @@ +"""A ``where`` that counts what a predicate admits, and one that reads a predicate at a neighbour. + +Two calls read a predicate where every other operator reads arithmetic +(math-spec's expressions page, *a predicate is an operand*): +``count(, over=)`` answers a number, and +``shift(, along=, offset=)`` answers a predicate. Both lanes build +the same model from either, which is what these assert — the coordinates built, +not just the count of them, because a predicate that inverted its sense would +keep the complement. +""" + +from __future__ import annotations + +import polars as pl +import pytest + +from tests.conftest import DISPATCH_SPEC, override +from tests.differential import differential +from tests.oracle import pd # skips the module without the [linopy] extra + + +def _spec(where: str) -> dict: + return override( + DISPATCH_SPEC, + **{'parameters.offer': {'dims': ['snapshot', 'generator']}, 'variables.p.where': where}, + ) + + +def _sources(offer: pl.DataFrame) -> dict: + return { + 'p_max': pd.Series({'wind': 100.0, 'gas': 200.0}), + 'cost': pd.Series({'wind': 0.0, 'gas': 50.0}), + 'load': pd.Series([80.0] * 4, index=pd.RangeIndex(4, name='snapshot')), + 'snapshot': pd.RangeIndex(4, name='snapshot'), + 'generator': ['wind', 'gas'], + 'offer': offer, + } + + +def _offer(**per_generator: list[float | None]) -> pl.DataFrame: + """One tidy frame from ``generator=[value per snapshot]``, a ``None`` standing for no row at all.""" + rows = [ + (t, name, value) + for name, values in per_generator.items() + for t, value in enumerate(values) + if value is not None + ] + return pl.DataFrame( + {'snapshot': [t for t, _, _ in rows], 'generator': [g for _, g, _ in rows], 'value': [v for _, _, v in rows]}, + schema={'snapshot': pl.Int64, 'generator': pl.String, 'value': pl.Float64}, + ) + + +def _built(where: str, offer: pl.DataFrame) -> list[tuple[int, str]]: + """``p``'s coordinates under *where*, both lanes agreeing on the model first.""" + with differential(_spec(where), _sources(offer)) as agreed: + rows = agreed.result.primal('p').to_dicts() + return sorted((row['snapshot'], row['generator']) for row in rows) + + +def test_a_count_is_one_number_per_coordinate_the_predicate_keeps(): + """``over`` is reduced away, so the count speaks about each generator without the file saying "each".""" + offer = _offer(wind=[10.0, 10.0, 10.0, 10.0], gas=[10.0, 0.0, 0.0, 0.0]) + built = _built('count(offer > 0, over=snapshot) >= 3', offer) + + assert built == [(0, 'wind'), (1, 'wind'), (2, 'wind'), (3, 'wind')], ( + 'wind offers at three snapshots or more and gas at one, so the mask keeps wind at every snapshot' + ) + + +def test_a_coordinate_the_predicate_admits_nowhere_counts_zero(): + """A generator with no row at all is a count of zero, which is the answer rather than a gap that drops the row.""" + offer = _offer(wind=[None, None, None, None], gas=[None, None, None, None]) + built = _built('count(offer > 0, over=snapshot) == 0', offer) + + assert built == [(t, g) for t in range(4) for g in ('gas', 'wind')], ( + 'neither generator offers anywhere, so both count zero and the mask keeps every coordinate' + ) + + +def test_a_shift_over_a_predicate_is_false_where_the_translation_vacates(): + """``offset=1`` reads one coordinate back, and the first coordinate reads false rather than wrapping or filling.""" + offer = _offer(wind=[150.0, 50.0, 50.0, 50.0], gas=[10.0, 10.0, 10.0, 10.0]) + built = _built('NOT shift(offer > 100, along=snapshot, offset=1)', offer) + + assert built == [(0, 'gas'), (0, 'wind'), (1, 'gas'), (2, 'gas'), (2, 'wind'), (3, 'gas'), (3, 'wind')], ( + 'only wind at snapshot 1 is dropped: the predicate holds at snapshot 0, and snapshot 0 itself ' + 'reads false because the translation vacates there' + ) + + +def test_a_count_over_a_translated_predicate_counts_the_runs_a_mask_starts(): + """The shape #582 needs a curve to state: a coordinate the predicate admits whose neighbour it does not.""" + offer = _offer(wind=[10.0, 10.0, 10.0, 10.0], gas=[0.0, 10.0, 0.0, 10.0]) + built = _built('count(offer > 0 AND NOT shift(offer > 0, along=snapshot, offset=1), over=snapshot) == 1', offer) + + assert built == [(0, 'wind'), (1, 'wind'), (2, 'wind'), (3, 'wind')], ( + "wind's offers are one consecutive run and gas's are two, so only wind admits exactly one run start" + ) + + +def test_a_count_reads_a_dimension_the_frame_does_not_span(): + """The count reduces its own dim away, so a declaration need not span it — it is the frame's other dims that must match.""" + offer = _offer(wind=[10.0, 10.0, 10.0, 10.0], gas=[10.0, 10.0, 10.0, 10.0]) + spec = override( + DISPATCH_SPEC, + **{ + 'parameters.offer': {'dims': ['snapshot', 'generator']}, + 'constraints.balance.where': 'count(offer > 0, over=generator) >= 2', + }, + ) + with differential(spec, _sources(offer)) as agreed: + rows = agreed.result.dual('balance').to_dicts() + + assert sorted(row['snapshot'] for row in rows) == [0, 1, 2, 3], ( + 'both generators offer at every snapshot, so the balance row stands at all four — the count ' + 'over generator asks nothing of the row it masks' + ) + + +@pytest.mark.parametrize( + ('where', 'kept'), + [ + pytest.param('count(offer > 0, over=snapshot) > 2', ['wind'], id='strictly-more-than'), + pytest.param('count(offer > 0, over=snapshot) == 1', ['gas'], id='exactly-one'), + pytest.param('count(offer > 0, over=snapshot) != 1', ['wind'], id='not-one'), + pytest.param('count(offer > 0, over=snapshot) <= 1', ['gas'], id='at-most-one'), + ], +) +def test_a_count_is_compared_by_the_operator_the_file_wrote(where, kept): + """Every relation a where may write, over the same two counts — four and one.""" + offer = _offer(wind=[10.0, 10.0, 10.0, 10.0], gas=[10.0, 0.0, 0.0, 0.0]) + built = _built(where, offer) + + assert sorted({g for _, g in built}) == kept, f'where: {where!r} kept the wrong generators' + assert len(built) == 4, 'whichever generator the count keeps, it is kept at every snapshot' diff --git a/uv.lock b/uv.lock index ce4f553d4..912f240f7 100644 --- a/uv.lock +++ b/uv.lock @@ -471,7 +471,7 @@ requires-dist = [ { name = "gurobipy", marker = "extra == 'gurobi'", specifier = ">=11" }, { name = "highspy", specifier = ">=1.13" }, { name = "linopy", marker = "extra == 'linopy'", git = "https://github.com/PyPSA/linopy?rev=master" }, - { name = "math-spec", git = "https://github.com/energy-models/math-spec?rev=5cb06abf3401d816fa371246f54760b56fd74bbb" }, + { name = "math-spec", git = "https://github.com/energy-models/math-spec?rev=2a9d0d599d3f5fd60acddec991412ff58b05110f" }, { name = "numpy", specifier = ">=1.26" }, { name = "pandas", marker = "extra == 'linopy'", specifier = ">=2.2" }, { name = "polars", specifier = ">=1.30" }, @@ -610,8 +610,8 @@ wheels = [ [[package]] name = "math-spec" -version = "0.0.0a108.dev17+g5cb06abf3" -source = { git = "https://github.com/energy-models/math-spec?rev=5cb06abf3401d816fa371246f54760b56fd74bbb#5cb06abf3401d816fa371246f54760b56fd74bbb" } +version = "0.0.0a108.dev20+g2a9d0d599" +source = { git = "https://github.com/energy-models/math-spec?rev=2a9d0d599d3f5fd60acddec991412ff58b05110f#2a9d0d599d3f5fd60acddec991412ff58b05110f" } dependencies = [ { name = "pydantic" }, { name = "pyparsing" },