feat(language): a where counts the coordinates a predicate admits, and reads one at a neighbour

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(<predicate>, over=)` | the predicate summed over that dimension, which reduces it away, and the total compared against the literal |
+| `shift(<predicate>, 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

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

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
 

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)))

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

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(<predicate>, over=<dim>)`` answers a number, and
+``shift(<predicate>, 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'