fix: let sum_back stop at each group's edge, as its checks already assumed

docs/examples/operators.md

@@ -427,6 +427,42 @@ objective: { sense: minimize, expression: sum(on) }
 ```
 
 $\sum_{h' \in \mathcal{H} \thinspace:\thinspace 0 \le h \ominus h' < \mathrm{min\_up}} \mathit{started}_{u,h'} \le \mathit{on}_{u,h} \qquad \forall\thinspace u \in \mathcal{U},\enspace h \in \mathcal{H}$
+
+### `sum_back(array, over=dim, within=n, by=lookup)`
+
+`examples/operators/sum_back_partitioned.yaml`
+
+```yaml
+description: >-
+  A window that stops at each group's edge: representative days are separate
+  samples rather than consecutive hours, so a window must not reach across the
+  seam between two of them.
+
+dimensions:
+  unit: { dtype: str }
+  hour: { dtype: int }
+  day: { dtype: str }
+
+lookups:
+  day_of: { over: hour, into: day }
+
+variables:
+  started:
+    foreach: [unit, hour]
+    domain: binary
+  on:
+    foreach: [unit, hour]
+    domain: binary
+
+constraints:
+  stays_up_inside_its_day:
+    foreach: [unit, hour]
+    expression: sum_back(started, over=hour, within=3, by=day_of) <= on
+
+objective: { sense: minimize, expression: sum(on) }
+```
+
+$\sum_{h' \in \mathcal{H} \thinspace:\thinspace 0 \le h -^{\mathrm{day\_of}(h)} h' < 3} \mathit{started}_{u,h'} \le \mathit{on}_{u,h} \qquad \forall\thinspace u \in \mathcal{U},\enspace h \in \mathcal{H}$
 <!-- gallery:end -->
 
 Regenerate with `pixi run python -m tools.gallery`.

docs/reference/language/operators.md

@@ -360,6 +360,7 @@ language is rendered the same way, on one page:
 | `sum_back(array, over=dim, within=n)` | $\sum_{h' \in \mathcal{H} \thinspace:\thinspace 0 \le h - h' < 3} \mathit{started}_{u,h'} \le \mathit{on}_{u,h} \qquad \forall\thinspace u \in \mathcal{U},\enspace h \in \mathcal{H}$ |
 | `sum_back(array, over=dim, within=p)` | $\sum_{h' \in \mathcal{H} \thinspace:\thinspace 0 \le h - h' < \mathrm{min\_up}} \mathit{started}_{u,h'} \le \mathit{on}_{u,h} \qquad \forall\thinspace u \in \mathcal{U},\enspace h \in \mathcal{H}$ |
 | `sum_back(array, over=dim, within=p, edge='wrap')` | $\sum_{h' \in \mathcal{H} \thinspace:\thinspace 0 \le h \ominus h' < \mathrm{min\_up}} \mathit{started}_{u,h'} \le \mathit{on}_{u,h} \qquad \forall\thinspace u \in \mathcal{U},\enspace h \in \mathcal{H}$ |
+| `sum_back(array, over=dim, within=n, by=lookup)` | $\sum_{h' \in \mathcal{H} \thinspace:\thinspace 0 \le h -^{\mathrm{day\_of}(h)} h' < 3} \mathit{started}_{u,h'} \le \mathit{on}_{u,h} \qquad \forall\thinspace u \in \mathcal{U},\enspace h \in \mathcal{H}$ |
 
 $t \ominus k$ denotes cyclic translation: index $t-k$ taken modulo the size of the dimension (`roll`). Plain $t-k$ (`shift`) has no wraparound — terms translated past the edge are simply absent.
 

docs/reference/notation.md

@@ -288,6 +288,18 @@ history:
 
 $$\sum_{t' \in \mathcal{T} \thinspace:\thinspace 0 \le t \ominus t' < \mathrm{min\_up}} \mathit{on}_{t',g} \le \mathit{units}_{g} \qquad \forall\thinspace t \in \mathcal{T},\enspace g \in \mathcal{G}$$
 
+#### `seasonal_window`
+
+a window partitioned by a lookup: the group rides on the operator
+
+```yaml
+seasonal_window:
+  foreach: [snapshot, generator]
+  expression: sum_back(on, over=snapshot, within=3, by=season_of) <= units
+```
+
+$$\sum_{t' \in \mathcal{T} \thinspace:\thinspace 0 \le t -^{\mathrm{season\_of}(t)} t' < 3} \mathit{on}_{t',g} \le \mathit{units}_{g} \qquad \forall\thinspace t \in \mathcal{T},\enspace g \in \mathcal{G}$$
+
 #### `pullback`
 
 at(), which re-indexes through a lookup instead of an offset

examples/operators/sum_back_partitioned.yaml

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: math-spec Contributors
+#
+# SPDX-License-Identifier: MIT
+
+description: >-
+  A window that stops at each group's edge: representative days are separate
+  samples rather than consecutive hours, so a window must not reach across the
+  seam between two of them.
+
+dimensions:
+  unit: { dtype: str }
+  hour: { dtype: int }
+  day: { dtype: str }
+
+lookups:
+  day_of: { over: hour, into: day }
+
+variables:
+  started:
+    foreach: [unit, hour]
+    domain: binary
+  on:
+    foreach: [unit, hour]
+    domain: binary
+
+constraints:
+  stays_up_inside_its_day:
+    foreach: [unit, hour]
+    expression: sum_back(started, over=hour, within=3, by=day_of) <= on
+
+objective: { sense: minimize, expression: sum(on) }

src/math_spec/operators.py

@@ -97,10 +97,12 @@ def optional(self) -> frozenset[str]:
     ),
     'sum_back': Builtin(
         1,
-        "sum_back(<expr>, over=<dim>, within=<n|parameter>[, edge='wrap'])",
+        "sum_back(<expr>, over=<dim>, within=<n|parameter>[, edge='wrap'][, by=<lookup>])",
         dimension_kwargs=('over',),
+        lookup_kwargs=('by',),
         required_value_kwargs=('within',),
         edge_kwargs=('edge',),
+        optional_kwargs=('by',),
     ),
     'shift': Builtin(
         1,

src/math_spec/typesetting/walk.py

@@ -373,7 +373,13 @@ def _call(self, node: FunctionCallNode, ctx: _Context) -> tuple[str, int]:
             step = _Step(1, 'wrap' if isinstance(node.kwargs.get('edge'), EdgeNode) else 'plain')
             self.policies.add(step.policy)
             source = f'{self.symbols.index[over.name]}{PRIME}'
-            lag = f'{ctx.subscript(over.name)} {self.translation(step)} {source}'
+            # a partition rides the operator here exactly as it does on a leaf
+            # translation, and for the same reason: what the group changes is
+            # where the axis ends, not which coordinate is being written
+            lag = (
+                f'{ctx.subscript(over.name)} '
+                f'{self.translation(step, self._group(node.kwargs.get("by"), over.name))} {source}'
+            )
             domain = (
                 f'{source} {self.op("in")} {self.symbols.set[over.name]} {self.op("such_that")} '
                 f'0 {self.op("le")} {lag} {self.op("lt")} {self._width(node.kwargs["within"])}'
@@ -406,6 +412,18 @@ def _call(self, node: FunctionCallNode, ctx: _Context) -> tuple[str, int]:
             domain = self.format.joined([self.membership(d) for d in dims], '')
         return self.format.summation(domain, self.reduction_body(node.args[0], ctx)), _PRECEDENCE['+']
 
+    def _group(self, by: ArithmeticNode | None, dim: str) -> str:
+        """A window's ``by=`` as the superscript its operator carries.
+
+        The bare index, not the subscript in force: the group is a property of
+        the row being written, and a window whose operand is itself translated
+        still asks which group *that row* is in.
+        """
+        if by is None:
+            return ''
+        assert isinstance(by, LookupNode)
+        return self.format.apply(self.format.upright(by.names[0]), self.symbols.index[dim])
+
     def _width(self, node: ArithmeticNode) -> str:
         """``sum_back``'s ``within=``: a number, or a parameter's own symbol.
 

tests/typesetting/golden/latex.out

@@ -81,6 +81,7 @@
 \text{held\_in\_season} && p_{t,g} & \le p_{t \boxminus_{0}^{\mathrm{season\_of}(t)} 1,g} && \forall\, t \in \mathcal{T},\ g \in \mathcal{G} \\
 \text{window} && \sum_{t' \in \mathcal{T} \,:\, 0 \le t - t' < 3} \mathit{on}_{t',g} & \le \mathit{units}_{g} && \forall\, t \in \mathcal{T},\ g \in \mathcal{G} \\
 \text{history} && \sum_{t' \in \mathcal{T} \,:\, 0 \le t \ominus t' < \mathrm{min\_up}} \mathit{on}_{t',g} & \le \mathit{units}_{g} && \forall\, t \in \mathcal{T},\ g \in \mathcal{G} \\
+\text{seasonal\_window} && \sum_{t' \in \mathcal{T} \,:\, 0 \le t -^{\mathrm{season\_of}(t)} t' < 3} \mathit{on}_{t',g} & \le \mathit{units}_{g} && \forall\, t \in \mathcal{T},\ g \in \mathcal{G} \\
 \text{pullback} && \mathit{spill}_{t} & \le \mathrm{zone\_cap}_{\mathrm{zone\_of}(b)} && \forall\, t \in \mathcal{T},\ b \in \mathcal{B} \\
 \text{grouped\_twice} && \sum_{g \in \mathcal{G} \,:\, \mathrm{gen\_bus}(g) = b \wedge \mathrm{gen\_tech}(g) = e} p_{t,g} & \le \mathrm{tech\_cap}_{b,e} && \forall\, t \in \mathcal{T},\ b \in \mathcal{B},\ e \in \mathcal{E} \\
 \text{pulled\_back\_twice} && \mathit{units}_{g} & \le \mathrm{tech\_cap}_{\mathrm{gen\_bus}(g),\mathrm{gen\_tech}(g)} && \forall\, g \in \mathcal{G} \\

tests/typesetting/golden/markdown.out

@@ -114,6 +114,10 @@ $$\sum_{t' \in \mathcal{T} \thinspace:\thinspace 0 \le t - t' < 3} \mathit{on}_{
 
 $$\sum_{t' \in \mathcal{T} \thinspace:\thinspace 0 \le t \ominus t' < \mathrm{min\_up}} \mathit{on}_{t',g} \le \mathit{units}_{g} \qquad \forall\thinspace t \in \mathcal{T},\enspace g \in \mathcal{G}$$
 
+**`seasonal_window`**
+
+$$\sum_{t' \in \mathcal{T} \thinspace:\thinspace 0 \le t -^{\mathrm{season\_of}(t)} t' < 3} \mathit{on}_{t',g} \le \mathit{units}_{g} \qquad \forall\thinspace t \in \mathcal{T},\enspace g \in \mathcal{G}$$
+
 **`pullback`**
 
 $$\mathit{spill}_{t} \le \mathrm{zone\_cap}_{\mathrm{zone\_of}(b)} \qquad \forall\thinspace t \in \mathcal{T},\enspace b \in \mathcal{B}$$

tests/typesetting/golden/model.yaml

@@ -124,6 +124,9 @@ constraints:
   history: # the same window, its width in the data and its edge wrapped
     foreach: [snapshot, generator]
     expression: sum_back(on, over=snapshot, within=min_up, edge='wrap') <= units
+  seasonal_window: # a window partitioned by a lookup: the group rides on the operator
+    foreach: [snapshot, generator]
+    expression: sum_back(on, over=snapshot, within=3, by=season_of) <= units
   pullback: # at(), which re-indexes through a lookup instead of an offset
     foreach: [snapshot, bus]
     expression: spill <= at(zone_cap, by=zone_of)

tests/typesetting/golden/typst.out

@@ -70,6 +70,7 @@ $ upright("balance") & sum_(g in cal(G) colon upright("gen_bus")(g) = b) p_(t,g)
   upright("held_in_season") & p_(t,g) & <= p_(t minus.square_(0)^(upright("season_of")(t)) 1,g) & forall t in cal(T), g in cal(G) \
   upright("window") & sum_(t' in cal(T) colon 0 <= t - t' < 3) italic("on")_(t',g) & <= italic("units")_(g) & forall t in cal(T), g in cal(G) \
   upright("history") & sum_(t' in cal(T) colon 0 <= t minus.o t' < upright("min_up")) italic("on")_(t',g) & <= italic("units")_(g) & forall t in cal(T), g in cal(G) \
+  upright("seasonal_window") & sum_(t' in cal(T) colon 0 <= t -^(upright("season_of")(t)) t' < 3) italic("on")_(t',g) & <= italic("units")_(g) & forall t in cal(T), g in cal(G) \
   upright("pullback") & italic("spill")_(t) & <= upright("zone_cap")_(upright("zone_of")(b)) & forall t in cal(T), b in cal(B) \
   upright("grouped_twice") & sum_(g in cal(G) colon upright("gen_bus")(g) = b and upright("gen_tech")(g) = e) p_(t,g) & <= upright("tech_cap")_(b,e) & forall t in cal(T), b in cal(B), e in cal(E) \
   upright("pulled_back_twice") & italic("units")_(g) & <= upright("tech_cap")_(upright("gen_bus")(g),upright("gen_tech")(g)) & forall g in cal(G) \

tools/spec_math.py

@@ -57,6 +57,7 @@
     'sum_back(array, over=dim, within=n)': 'sum_back',
     'sum_back(array, over=dim, within=p)': 'sum_back_by_parameter',
     "sum_back(array, over=dim, within=p, edge='wrap')": 'sum_back_wrap',
+    'sum_back(array, over=dim, within=n, by=lookup)': 'sum_back_partitioned',
 }