mathspec 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
mathspec/program.py ADDED
@@ -0,0 +1,1514 @@
1
+ # SPDX-FileCopyrightText: mathspec Contributors
2
+ #
3
+ # SPDX-License-Identifier: MIT
4
+
5
+ """The program: what a file declares, with names resolved and shapes fixed.
6
+
7
+ The second public state, and the one a consumer reads. A :class:`Program` is
8
+ the file typed, section for section: every declaration it makes, with names
9
+ resolved, shapes fixed and every rule decidable without data checked, and no
10
+ data at all. Lowering, as a :class:`~mathspec.model.Spec` loads, is the only
11
+ thing that builds one, so nothing here re-checks a hand-built one.
12
+
13
+ Node and declaration classes are matched with ``isinstance``. The rules a
14
+ node's structure does not show is :func:`children`; the questions over the walk
15
+ are :func:`walk_regions`, :func:`walk` and the filters beside them. A
16
+ resolved ``where`` arrives as a :class:`Mask`. Frozen dataclasses only — no
17
+ execution logic, and nothing imported from a consumer. How a consumer reads
18
+ one: ``docs/reference/reading.md``.
19
+ """
20
+
21
+ from __future__ import annotations
22
+
23
+ from collections.abc import Mapping
24
+ from dataclasses import dataclass, fields, replace
25
+ from functools import cached_property
26
+ from typing import TYPE_CHECKING, Literal, assert_never
27
+
28
+ from mathspec._expression_parser import ComparisonOperator
29
+ from mathspec._sealed import Sealed
30
+ from mathspec.errors import did_you_mean
31
+
32
+ if TYPE_CHECKING:
33
+ import datetime
34
+ from collections.abc import Iterator
35
+
36
+
37
+ #: What ``mathspec.program`` promises a consumer, sorted.
38
+ __all__ = [
39
+ 'Add',
40
+ 'And',
41
+ 'Assumption',
42
+ 'BooleanLiteral',
43
+ 'Cases',
44
+ 'Connective',
45
+ 'Constant',
46
+ 'ConstraintDeclaration',
47
+ 'ConstraintSense',
48
+ 'CountComparison',
49
+ 'DeclaredDtype',
50
+ 'DimensionComparison',
51
+ 'DimensionDeclaration',
52
+ 'DimensionDtype',
53
+ 'DimensionPosition',
54
+ 'Direction',
55
+ 'Divide',
56
+ 'Dual',
57
+ 'Expression',
58
+ 'ExpressionComparison',
59
+ 'ExpressionDeclaration',
60
+ 'Footprint',
61
+ 'GroupSum',
62
+ 'Link',
63
+ 'Mask',
64
+ 'Multiply',
65
+ 'Named',
66
+ 'Negate',
67
+ 'Not',
68
+ 'ObjectiveDeclaration',
69
+ 'ObjectiveSense',
70
+ 'Or',
71
+ 'Parameter',
72
+ 'ParameterComparison',
73
+ 'ParameterDeclaration',
74
+ 'ParameterDefined',
75
+ 'ParameterDtype',
76
+ 'Partition',
77
+ 'PiecewiseDeclaration',
78
+ 'PiecewiseMethod',
79
+ 'Power',
80
+ 'Predicate',
81
+ 'PredicateOperator',
82
+ 'Program',
83
+ 'Pullback',
84
+ 'PulledBackPredicate',
85
+ 'QuadraticPosition',
86
+ 'Reach',
87
+ 'Region',
88
+ 'RelationComparison',
89
+ 'RelationDeclaration',
90
+ 'RelationDefined',
91
+ 'RelationPairComparison',
92
+ 'Separability',
93
+ 'SosDeclaration',
94
+ 'SosType',
95
+ 'Sum',
96
+ 'Translate',
97
+ 'TranslatedPredicate',
98
+ 'TypedPredicate',
99
+ 'Variable',
100
+ 'VariableAbsence',
101
+ 'VariableDeclaration',
102
+ 'VariableDefined',
103
+ 'VariableDomain',
104
+ 'WindowSum',
105
+ 'assumption_message',
106
+ 'carries_variable',
107
+ 'children',
108
+ 'is_quadratic',
109
+ 'parameters_of',
110
+ 'variables_of',
111
+ 'walk',
112
+ 'walk_regions',
113
+ 'where_children',
114
+ ]
115
+
116
+
117
+ ConstraintSense = ComparisonOperator
118
+
119
+ #: Where a degree-2 product may stand in the math a solver sees. An objective
120
+ #: and a constraint take ``variable * variable``; a bound and a ``piecewise:``
121
+ #: link are read affinely (``mathspec.degree``), so those are the two.
122
+ QuadraticPosition = Literal['objective', 'constraint']
123
+
124
+ #: The dtype a dimension index may declare (the declaration rules), and what
125
+ #: its labels are. ``datetime`` is a dimension's alone — labels on a timeline
126
+ #: order and compare, where a *value* of that type is a moment nothing
127
+ #: computes with.
128
+ DimensionDtype = Literal['float', 'int', 'str', 'datetime']
129
+
130
+ #: The dtype a parameter may declare (the declaration rules), and what its bound
131
+ #: column must be. ``bool`` is a parameter's alone — a value column may be a
132
+ #: flag a mask reads, where a label set of two members is a dimension nothing
133
+ #: indexes by.
134
+ ParameterDtype = Literal['float', 'int', 'bool', 'str']
135
+
136
+ #: What a *name* a where comparison tests may be — a parameter's dtype or a
137
+ #: dimension's, since a relation's is its target's. The union rather than either
138
+ #: half, because a mask names all three kinds and reads the dtype the same way.
139
+ DeclaredDtype = ParameterDtype | DimensionDtype
140
+
141
+ #: The domain a variable may declare.
142
+ VariableDomain = Literal['continuous', 'integer', 'binary']
143
+
144
+ #: What a masked variable's non-existence *means* where it does not exist.
145
+ #: ``undefined`` is the absence rules' default — a term carrying it takes its
146
+ #: row. ``zero`` says the quantity *is* zero there, so the term contributes
147
+ #: nothing and the row stands.
148
+ VariableAbsence = Literal['undefined', 'zero']
149
+
150
+ #: Which way an objective is optimised (the declaration rules).
151
+ ObjectiveSense = Literal['minimize', 'maximize']
152
+
153
+ #: The order of special ordered set.
154
+ SosType = Literal[1, 2]
155
+
156
+ #: How a ``piecewise:`` block restricts its interpolation weights. Kept in step
157
+ #: with :data:`~mathspec.model.PIECEWISE_METHODS`, which says what each one
158
+ #: emits, by ``tests/test_schema.py``.
159
+ PiecewiseMethod = Literal['adjacency', 'sos2', 'convex', 'lp']
160
+
161
+
162
+ # --------------------------------------------------------------------------
163
+ # Affine expressions
164
+ # --------------------------------------------------------------------------
165
+
166
+
167
+ @dataclass(frozen=True)
168
+ class Constant:
169
+ """A scalar constant."""
170
+
171
+ value: float
172
+
173
+
174
+ @dataclass(frozen=True)
175
+ class Parameter:
176
+ """A parameter reference — contributes to the constant part."""
177
+
178
+ name: str
179
+
180
+
181
+ @dataclass(frozen=True)
182
+ class Variable:
183
+ """A variable reference — one term per existing variable row."""
184
+
185
+ name: str
186
+
187
+
188
+ @dataclass(frozen=True)
189
+ class Dual:
190
+ """A constraint's dual — its shadow price, read after the solve.
191
+
192
+ Stands only under an :class:`ExpressionDeclaration` the math never reads:
193
+ the loader refuses ``dual()`` anywhere a solver ingests. One value per
194
+ coordinate of the named constraint's own ``dims`` frame: the leaf reshapes
195
+ nothing, like a parameter.
196
+ """
197
+
198
+ constraint: str
199
+
200
+
201
+ @dataclass(frozen=True)
202
+ class Negate:
203
+ operand: Expression
204
+
205
+
206
+ @dataclass(frozen=True)
207
+ class Add:
208
+ left: Expression
209
+ right: Expression
210
+
211
+
212
+ @dataclass(frozen=True)
213
+ class Multiply:
214
+ """Product of two operands.
215
+
216
+ Affine where at least one factor is variable-free; degree 2 where neither
217
+ is, which ``mathspec.degree`` admits in a :data:`QuadraticPosition` alone.
218
+ """
219
+
220
+ left: Expression
221
+ right: Expression
222
+
223
+
224
+ @dataclass(frozen=True)
225
+ class Power:
226
+ """``base ** exponent``, both variable-free wherever the math reads it.
227
+
228
+ The language refuses a variable anywhere under it (``mathspec.degree``),
229
+ so in the program a solver sees it is degree 0 and folds to one number per
230
+ coordinate like any other parameter arithmetic.
231
+ """
232
+
233
+ base: Expression
234
+ exponent: Expression
235
+
236
+
237
+ @dataclass(frozen=True)
238
+ class Divide:
239
+ """Quotient ``numerator / divisor``, the divisor variable-free wherever the math reads it (``mathspec.degree``)."""
240
+
241
+ numerator: Expression
242
+ divisor: Expression
243
+
244
+
245
+ @dataclass(frozen=True)
246
+ class Sum:
247
+ """Sum ``operand`` over the named dims, removing them from the result."""
248
+
249
+ operand: Expression
250
+ over: tuple[str, ...]
251
+
252
+
253
+ @dataclass(frozen=True)
254
+ class GroupSum:
255
+ """Sum ``operand`` through a relation: the dims ``direction`` consumes go, the dims it produces arrive, the dims it joins on stay.
256
+
257
+ The join keys on the consumed columns and every joined column, and the
258
+ operand carries every dim consumed or joined on.
259
+ """
260
+
261
+ operand: Expression
262
+ direction: Direction
263
+
264
+
265
+ @dataclass(frozen=True)
266
+ class Pullback:
267
+ """Read ``operand`` through a relation — the adjoint of :class:`GroupSum`.
268
+
269
+ The dims ``direction`` consumes go and the dims it produces arrive, one
270
+ value per coordinate because the read takes value columns at a key the
271
+ result fixes, which the loader checks. The join fans out, many
272
+ produced tuples sharing one consumed tuple — at each coordinate of the
273
+ joined columns, which the operand carries and the result keeps.
274
+ """
275
+
276
+ operand: Expression
277
+ direction: Direction
278
+
279
+
280
+ @dataclass(frozen=True)
281
+ class Translate:
282
+ """Re-index along one dimension: the result at *t* is ``operand`` at *t - offset*.
283
+
284
+ ``wrap`` is ``edge='wrap'`` in the file: periodic, and stated on every
285
+ node. ``fill`` is what an acyclic shift leaves behind: ``None`` leaves the
286
+ vacated positions absent, so the row drops; a number makes them present
287
+ and contribute it. Always ``None`` under ``wrap``.
288
+
289
+ ``offset`` is an integer, or the name of an integer parameter that does
290
+ not depend on ``along`` and carries its sign in the values.
291
+
292
+ ``partition`` is a relation with a key column over ``along``
293
+ (:class:`Partition`), and the translation then happens inside each group
294
+ its ``within=`` columns make: the neighbour is the one before in the same
295
+ group, the edge is the group's, and a wrap closes each group onto itself.
296
+ A coordinate the relation sends nowhere reaches nothing.
297
+ """
298
+
299
+ operand: Expression
300
+ along: str
301
+ offset: int | str
302
+ wrap: bool
303
+ fill: float | None = None
304
+ partition: Partition | None = None
305
+
306
+
307
+ @dataclass(frozen=True)
308
+ class WindowSum:
309
+ """Sum ``operand`` over a trailing window along one dimension.
310
+
311
+ The result at *t* is the sum of the operand at every position from
312
+ *t - width + 1* through *t*, so a width of 1 is the operand itself. The
313
+ dimension survives: this replicates terms onto the positions that can see
314
+ them rather than reducing anything away.
315
+
316
+ ``width`` is a whole number, or the name of an integer parameter when the
317
+ window differs per entity — a minimum up time, a rolling budget, a delivery
318
+ horizon. A named width may not depend on the dimension being summed over.
319
+
320
+ ``wrap`` says whether the window reaches around the start of the axis
321
+ instead of stopping short at it, and is stated on every node.
322
+
323
+ ``partition`` names a relation over that dimension, and the window then stops
324
+ at each group's edge. Positions are counted inside the group, so a
325
+ coordinate the relation places nowhere reaches nothing — not even itself.
326
+ """
327
+
328
+ operand: Expression
329
+ along: str
330
+ width: int | str
331
+ wrap: bool
332
+ partition: Partition | None = None
333
+
334
+
335
+ @dataclass(frozen=True)
336
+ class Region:
337
+ """One region of a :class:`Cases`: where it applies, and the value there.
338
+
339
+ ``when`` is stated on every region; the one the file wrote as
340
+ ``otherwise:`` carries the negation of the others.
341
+ """
342
+
343
+ when: Mask
344
+ value: Expression
345
+
346
+
347
+ @dataclass(frozen=True)
348
+ class Cases:
349
+ """A value defined by region — exactly one region applies at each coordinate.
350
+
351
+ The regions are disjoint and total, so a consumer adds them rather than
352
+ ranking them. Not a shape operator: every region spans the dims the
353
+ expression does.
354
+ """
355
+
356
+ regions: tuple[Region, ...]
357
+
358
+
359
+ @dataclass(frozen=True)
360
+ class Named:
361
+ """A use of an ``expressions:`` entry, standing where its name was written, with the entry's body under it.
362
+
363
+ Its value is its body's: a consumer building rows steps through it, as
364
+ :func:`children` does. It is kept as a node rather than written in so the
365
+ typesetter can print the symbol where the name stood and define it once.
366
+ Every use of one entry holds the one node resolution built for it, whose
367
+ :attr:`body` is the :attr:`ExpressionDeclaration.expression` of that entry.
368
+ """
369
+
370
+ name: str
371
+ body: Expression
372
+
373
+
374
+ #: Every expression node, as one type — what a walk takes. The set is
375
+ #: *closed*: nothing registers into it, so a consumer that walks it ends in
376
+ #: ``assert_never`` and a node added without a branch is a type error at the
377
+ #: site that must grow one, rather than a ``LanguageError`` raised at the first
378
+ #: model that uses it. The degree rules (``mathspec.degree``) hold on every
379
+ #: tree the math reads — :attr:`Program.roots`, a bound, and a named expression
380
+ #: that is ``in_math`` — affine but where a :class:`QuadraticPosition` admits a
381
+ #: :class:`Multiply` of two variable-carrying operands. A
382
+ #: :class:`ExpressionDeclaration` the math never reads is held to none of them.
383
+ #: No node records which tree it stands in.
384
+ Expression = (
385
+ Constant
386
+ | Parameter
387
+ | Variable
388
+ | Dual
389
+ | Negate
390
+ | Add
391
+ | Multiply
392
+ | Power
393
+ | Divide
394
+ | Sum
395
+ | GroupSum
396
+ | Pullback
397
+ | Translate
398
+ | WindowSum
399
+ | Cases
400
+ | Named
401
+ )
402
+
403
+
404
+ def children(expression: Expression) -> tuple[Expression, ...]:
405
+ """The sub-expressions of *expression* — what every walk recurses through."""
406
+ if isinstance(expression, Named):
407
+ return (expression.body,)
408
+ if isinstance(expression, Negate):
409
+ return (expression.operand,)
410
+ if isinstance(expression, (Add, Multiply)):
411
+ return (expression.left, expression.right)
412
+ if isinstance(expression, Divide):
413
+ return (expression.numerator, expression.divisor)
414
+ if isinstance(expression, Power):
415
+ return (expression.base, expression.exponent)
416
+ if isinstance(expression, (Sum, GroupSum, Pullback, Translate, WindowSum)):
417
+ return (expression.operand,)
418
+ if isinstance(expression, Cases):
419
+ return tuple(region.value for region in expression.regions)
420
+ if isinstance(expression, (Constant, Parameter, Variable, Dual)):
421
+ return ()
422
+ assert_never(expression)
423
+
424
+
425
+ # --------------------------------------------------------------------------
426
+ # Declarations
427
+ # --------------------------------------------------------------------------
428
+
429
+
430
+ @dataclass(frozen=True)
431
+ class RelationDeclaration:
432
+ """One declared relation: a table over its ``columns``, single-valued per ``key``.
433
+
434
+ ``columns`` maps each role to its dimension in the order the table
435
+ carries them, the key's roles first; ``key`` is the roles a row is
436
+ identified by, and :attr:`values` the rest — every role is a key role for
437
+ a bare relation, which is one with no value columns. Every value is
438
+ checked when the data is attached to be a label of its column's dimension, and the table to
439
+ have one row per key tuple — which keeps a mistyped label from silently
440
+ dropping its terms in the join that places them, and is what lets ``at``
441
+ read one value.
442
+ """
443
+
444
+ columns: tuple[tuple[str, str], ...]
445
+ key: tuple[str, ...]
446
+ description: str | None = None
447
+
448
+ @property
449
+ def roles(self) -> tuple[str, ...]:
450
+ return tuple(role for role, _ in self.columns)
451
+
452
+ @property
453
+ def dims(self) -> tuple[str, ...]:
454
+ return tuple(dim for _, dim in self.columns)
455
+
456
+ @property
457
+ def values(self) -> tuple[str, ...]:
458
+ """The roles the key determines."""
459
+ return tuple(role for role in self.roles if role not in self.key)
460
+
461
+ @cached_property
462
+ def _dim_of(self) -> Mapping[str, str]:
463
+ """Each role's dimension, built once: :meth:`dim` is called per role inside loops over roles."""
464
+ return dict(self.columns)
465
+
466
+ def dim(self, role: str) -> str:
467
+ return self._dim_of[role]
468
+
469
+
470
+ @dataclass(frozen=True)
471
+ class Direction:
472
+ """One relation as one call reads it — which columns are consumed, which produced, which joined on.
473
+
474
+ The declaration fixes no direction; the call does, and this is the one it
475
+ named. ``name`` is the relation's, as :attr:`Program.relations` keys it.
476
+ ``consumed``, ``produced`` and ``joined`` are *roles* — column names of
477
+ ``relation``, which maps every role to its dimension and names the key.
478
+ ``joined`` is the key roles the call did not name (every role, for a bare
479
+ relation): the join keys on them, and a value role left unnamed is not
480
+ read.
481
+ """
482
+
483
+ name: str
484
+ relation: RelationDeclaration
485
+ consumed: tuple[str, ...]
486
+ produced: tuple[str, ...]
487
+ joined: tuple[str, ...]
488
+
489
+ def dim(self, role: str) -> str:
490
+ """The dimension *role* ranges over."""
491
+ return self.relation.dim(role)
492
+
493
+ @property
494
+ def consumed_dims(self) -> tuple[str, ...]:
495
+ return tuple(self.dim(role) for role in self.consumed)
496
+
497
+ @property
498
+ def produced_dims(self) -> tuple[str, ...]:
499
+ return tuple(self.dim(role) for role in self.produced)
500
+
501
+ @property
502
+ def joined_dims(self) -> tuple[str, ...]:
503
+ return tuple(self.dim(role) for role in self.joined)
504
+
505
+
506
+ @dataclass(frozen=True)
507
+ class Partition:
508
+ """One relation as a partition steps along it — the key column stepped along, the group columns, and the key columns joined on.
509
+
510
+ ``name`` is the relation's, as :attr:`Program.relations` keys it.
511
+ ``along``, ``group`` and ``joined`` are *roles* — column names of
512
+ ``relation``, which maps every role to its dimension and names the key.
513
+ ``along`` is the one key column over the dimension stepped along, and
514
+ the frame keeps it. ``group`` is the value columns ``within=`` named,
515
+ read at the row's key. ``joined`` is the other key columns, whose
516
+ dimensions the frame carries. Nothing is consumed and nothing is
517
+ produced: the frame does not change.
518
+ """
519
+
520
+ name: str
521
+ relation: RelationDeclaration
522
+ along: str
523
+ group: tuple[str, ...]
524
+ joined: tuple[str, ...]
525
+
526
+ def dim(self, role: str) -> str:
527
+ """The dimension *role* ranges over."""
528
+ return self.relation.dim(role)
529
+
530
+ @property
531
+ def along_dim(self) -> str:
532
+ return self.dim(self.along)
533
+
534
+ @property
535
+ def joined_dims(self) -> tuple[str, ...]:
536
+ return tuple(self.dim(role) for role in self.joined)
537
+
538
+
539
+ @dataclass(frozen=True)
540
+ class DimensionDeclaration:
541
+ """A dimension, as the file declares it."""
542
+
543
+ #: What the labels are, as the file declares them. A dimension is read from
544
+ #: whatever table carries it, so the declared type is what that column is
545
+ #: checked against — the same claim ``ParameterDeclaration.dtype`` makes
546
+ #: about a value column, one axis over.
547
+ dtype: DimensionDtype = 'str'
548
+ description: str | None = None
549
+
550
+
551
+ @dataclass(frozen=True)
552
+ class Assumption:
553
+ """A predicate the file states of its data, under the name it wrote in ``assumptions:``.
554
+
555
+ ``predicate`` is true at every coordinate of its frame — the product of
556
+ every dim the two masks name — that ``where`` admits, a missing row
557
+ reading as false as it does in any mask. Nothing here is decidable at
558
+ load: both sides are the data's, which is why the consumer attaching it
559
+ checks.
560
+ """
561
+
562
+ predicate: Mask
563
+ where: Mask | None = None
564
+ #: What the file wrote under ``description:``, or the sentence a
565
+ #: ``piecewise:`` method implies. The refusal trails it: the names alone
566
+ #: say which columns are wrong, and not why the rule is there.
567
+ description: str | None = None
568
+
569
+
570
+ def assumption_message(name: str, assumption: Assumption) -> str:
571
+ """The sentence a consumer raises when the data attached to *assumption*, called *name*, fails it.
572
+
573
+ The language's own wording, so every consumer refuses in the same words;
574
+ a consumer appends the coordinates it saw. Where the file wrote a
575
+ ``description:``, or a ``piecewise:`` method implied one, it trails the
576
+ sentence: the names say which columns are wrong, and the description says
577
+ why the rule is there.
578
+ """
579
+ read = ', '.join(f"'{n}'" for n in sorted(assumption.predicate.names_read))
580
+ sentence = f"assumption '{name}' does not hold for the data attached to {read}"
581
+ return f'{sentence} — {assumption.description}' if assumption.description else sentence
582
+
583
+
584
+ @dataclass(frozen=True)
585
+ class ParameterDeclaration:
586
+ """Shape declaration; data is attached at execution time by name.
587
+
588
+ ``dtype`` is what the declaration claims the values are, and a consumer
589
+ attaching data refuses a column that is not it — so the *declaration* is
590
+ what is read, rather than whatever the column happens to hold.
591
+ """
592
+
593
+ dims: tuple[str, ...]
594
+ dtype: ParameterDtype = 'float'
595
+ description: str | None = None
596
+
597
+
598
+ @dataclass(frozen=True)
599
+ class VariableDeclaration:
600
+ dims: tuple[str, ...]
601
+ where: Mask | None = None
602
+ #: A number or a parameter, or ``None`` where that side is open. What stands
603
+ #: for an open side in a solve is the consumer's to choose.
604
+ lower: Expression | None = None
605
+ #: As :attr:`lower`, for the other side.
606
+ upper: Expression | None = None
607
+ domain: VariableDomain = 'continuous'
608
+ absence: VariableAbsence = 'undefined'
609
+ description: str | None = None
610
+
611
+
612
+ @dataclass(frozen=True)
613
+ class ConstraintDeclaration:
614
+ """``lhs sense rhs`` for each coord combination of ``dims``.
615
+
616
+ Either side may carry variables and constants alike; which side a
617
+ consumer gathers them onto is its own arrangement and not stated here.
618
+ ``where`` masks out coord combinations (row absence, like variables).
619
+ """
620
+
621
+ dims: tuple[str, ...]
622
+ lhs: Expression
623
+ sense: ConstraintSense
624
+ rhs: Expression
625
+ where: Mask | None = None
626
+ description: str | None = None
627
+
628
+
629
+ @dataclass(frozen=True)
630
+ class SosDeclaration:
631
+ """One special-ordered set per coordinate of the variable's ``dims`` minus ``along``.
632
+
633
+ The only declaration that adds neither a column nor a row: it names
634
+ columns a consumer already has and says what may be nonzero among them. Which
635
+ dims those are is the variable's own ``dims`` and is read from it: a
636
+ copy here would be a second home for a fact
637
+ (:attr:`Program.variables`).
638
+ """
639
+
640
+ variable: str
641
+ along: str
642
+ sos_type: SosType
643
+ description: str | None = None
644
+
645
+
646
+ @dataclass(frozen=True)
647
+ class ObjectiveDeclaration:
648
+ """Objective — scalar, every reduction in it one the file wrote."""
649
+
650
+ sense: ObjectiveSense
651
+ expression: Expression
652
+ description: str | None = None
653
+
654
+
655
+ @dataclass(frozen=True)
656
+ class ExpressionDeclaration:
657
+ """A named quantity — one the math reads, or one only read back after a solve.
658
+
659
+ ``in_math`` where the objective or a constraint reads it, directly or
660
+ through another entry or a macro; its body then stands inside
661
+ :attr:`Program.roots` and is held to the degree rules where it is
662
+ read. Otherwise nothing a solver sees contains it: it is a reported
663
+ quantity, its body held to no degree, the one place a :class:`Dual` may
664
+ stand. A bound and a ``where`` name no entry, so neither decides this.
665
+ """
666
+
667
+ expression: Expression
668
+ #: The frame the entry is read over: the ``dims:`` a cased entry
669
+ #: declares, or the dims a plain entry's body carries.
670
+ dims: tuple[str, ...]
671
+ in_math: bool
672
+ description: str | None = None
673
+
674
+
675
+ @dataclass(frozen=True)
676
+ class Link:
677
+ """One link of a ``piecewise:`` block: an expression tied to the breakpoints a values parameter holds.
678
+
679
+ ``sign`` is ``'=='`` where the link is pinned to the curve, and one side
680
+ of it where the link is bounded by the curve instead.
681
+ """
682
+
683
+ expression: Expression
684
+ values: str
685
+ sign: ConstraintSense = '=='
686
+
687
+
688
+ @dataclass(frozen=True)
689
+ class PiecewiseDeclaration:
690
+ """A ``piecewise:`` block as the curve it states, which :meth:`~mathspec.model.Spec.expand` writes out as rows.
691
+
692
+ A program of a model that still declares one carries it here, typed; a
693
+ program of the expanded model carries the rows instead, under
694
+ :attr:`Program.variables` and :attr:`Program.constraints`, and what the
695
+ method assumes of the breakpoints under :attr:`Program.assumptions`. A
696
+ consumer building rows takes the expanded model.
697
+
698
+ Attributes:
699
+ over: The breakpoint dimension.
700
+ links: The links, in the order the file wrote them.
701
+ method: How the weights are restricted.
702
+ activity: The binary the weights sum to, or ``None`` where they sum
703
+ to 1.
704
+ points: The parameter saying how far each curve runs, or ``None``.
705
+ frame: The dimensions the block builds one curve per coordinate of,
706
+ in declaration order.
707
+ description: What the file wrote under ``description:``, or ``None``.
708
+ """
709
+
710
+ over: str
711
+ links: tuple[Link, ...]
712
+ method: PiecewiseMethod
713
+ frame: tuple[str, ...]
714
+ activity: str | None = None
715
+ points: str | None = None
716
+ description: str | None = None
717
+
718
+ @property
719
+ def nominated(self) -> str | None:
720
+ """The block's own values parameter ``points:`` names, so the mask is derived from it — or ``None``."""
721
+ return self.points if self.points in {link.values for link in self.links} else None
722
+
723
+ @property
724
+ def curve(self) -> tuple[Link, Link]:
725
+ """The two links as ``(x, y)``, the bounded one last. Two-link blocks only."""
726
+ x, y = self.links
727
+ return (y, x) if x.sign != '==' else (x, y)
728
+
729
+
730
+ @dataclass(frozen=True)
731
+ class Footprint:
732
+ """Which of the language's constructs one program uses.
733
+
734
+ A subset, never the whole: an empty field says this program does not use
735
+ the construct.
736
+
737
+ Attributes:
738
+ quadratic: Each position a product of two variable-carrying operands
739
+ stands in; empty is affine throughout.
740
+ domains: Every domain declared.
741
+ sos_types: The order of each special-ordered set declared.
742
+ kinds: Every expression node kind that appears.
743
+ """
744
+
745
+ quadratic: frozenset[QuadraticPosition]
746
+ domains: frozenset[VariableDomain]
747
+ sos_types: frozenset[SosType]
748
+ kinds: frozenset[type[Expression]]
749
+
750
+
751
+ @dataclass(frozen=True)
752
+ class Reach:
753
+ """One read along an axis whose distance only data can say.
754
+
755
+ Attributes:
756
+ label: The declaration reading, as the lowering's messages label it.
757
+ name: The parameter or relation that says how far.
758
+ kind: An ``offset`` is a parameter's values, which
759
+ :meth:`Separability.resolved` folds in; a ``partition`` and a
760
+ ``coordinate`` are a relation's groups, which it does not.
761
+ """
762
+
763
+ label: str
764
+ name: str
765
+ kind: Literal['offset', 'partition', 'coordinate']
766
+
767
+
768
+ @dataclass(frozen=True)
769
+ class Separability:
770
+ """What building one dimension a window at a time asks of a driver, and what it would break.
771
+
772
+ A rolling-horizon or myopic driver cuts an axis into windows and builds
773
+ each on its own, and a decomposition cuts the same axis and solves each
774
+ piece on its own. What the program can say is whether every row it builds
775
+ is then complete inside one window: how far a row reads ahead along the
776
+ axis, which declarations tie the axis together so that no window holds
777
+ them, and — the same fact read as a set — which rows and columns are left
778
+ over as the border every window shares. It cannot say whether the windowed
779
+ answer is the one a whole-horizon solve would give — a store carried over
780
+ one row windows cleanly, and a rolling solve of it is still a different
781
+ answer — which is the driver's design and not the model's.
782
+
783
+ What a row reads *behind* is not reported. A window starts where the
784
+ driver puts it, and what its first rows meet there is the edge policy:
785
+ the opening state a rolling horizon seeds, and the driver's to carry.
786
+
787
+ Attributes:
788
+ dimension: The axis asked about.
789
+ ahead: Coordinates a window must see after its last row for every row
790
+ it builds to be complete — what a negative ``shift`` reads. ``0``
791
+ is pointwise; a ``shift`` of ``-2`` is ``2``.
792
+ coupled: Each declaration that ties the axis together, to what ties it
793
+ and the one modelling change that would not: a sum over the axis
794
+ in a constraint, a grouping that consumes it, a wrapped
795
+ translation, a set. No window satisfies these, and no rewrite here
796
+ would keep the model's meaning, so the remedy is named rather than
797
+ applied.
798
+ undecided: Each read along the axis whose reach only data can say —
799
+ a named offset, a partition whose groups a window may cut, a read
800
+ through a relation at a coordinate the data chooses.
801
+ :meth:`resolved` folds a parameter's values in.
802
+ restarts: Each declaration counting a position along the axis, which a
803
+ window restarts at its first row. Whether that is wanted — a seed
804
+ once per window, or once per horizon — is the modeller's, so it is
805
+ reported rather than refused.
806
+ linking_rows: Each constraint no one window holds whole, in declaration
807
+ order: one the axis does not index, whose row stands in every
808
+ window, and one :attr:`coupled` names. A reach the data decides is
809
+ not one, so a row waiting on :attr:`undecided` may span two windows.
810
+ linking_columns: Each variable the axis does not index, in declaration
811
+ order, whose column every window reads. A decomposition calls a
812
+ window a block, and with :attr:`linking_rows` this is the border of
813
+ a bordered block-diagonal form cut along the axis, whole where
814
+ nothing is :attr:`undecided` and no set runs through it. A set
815
+ couples the axis without building a row, so it stands in neither
816
+ field. The form is exactly that where :attr:`ahead` is ``0``: a
817
+ positive lookahead is neighbouring blocks overlapping by that much.
818
+ """
819
+
820
+ dimension: str
821
+ ahead: int
822
+ coupled: Mapping[str, str]
823
+ undecided: tuple[Reach, ...]
824
+ restarts: Mapping[str, str]
825
+ linking_rows: tuple[str, ...]
826
+ linking_columns: tuple[str, ...]
827
+
828
+ @property
829
+ def windowable(self) -> bool:
830
+ """Whether every row builds complete inside a window looking :attr:`ahead` past its last row.
831
+
832
+ ``False`` while a reach is :attr:`undecided`, which a driver holding
833
+ the data may resolve; :attr:`restarts` do not count against it.
834
+ """
835
+ return not self.coupled and not self.undecided
836
+
837
+ def resolved(self, least: Mapping[str, int]) -> Separability:
838
+ """The same verdict with each named offset folded into :attr:`ahead`.
839
+
840
+ A driver holding the data reads the least value of each parameter an
841
+ :attr:`undecided` reach names and hands it here, so the rule that
842
+ turns a value into a reach — a negative offset reads ahead by that
843
+ much, a positive one reads behind and asks nothing — has one home.
844
+
845
+ Args:
846
+ least: Parameter name to the least of its values. A reach through
847
+ a relation — a partition, a coordinate — cannot be folded this
848
+ way and stays undecided, as does a parameter left out.
849
+
850
+ Raises:
851
+ KeyError: A name no undecided reach along this axis waits on.
852
+ """
853
+ waiting = {reach.name for reach in self.undecided if reach.kind == 'offset'}
854
+ for name in least:
855
+ if name not in waiting:
856
+ raise KeyError(
857
+ f"'{name}' is not a parameter an undecided reach along '{self.dimension}' waits on. "
858
+ + did_you_mean(name, sorted(waiting))
859
+ )
860
+ folded = {reach for reach in self.undecided if reach.kind == 'offset' and reach.name in least}
861
+ ahead = max([self.ahead, *(-least[reach.name] for reach in folded)])
862
+ return replace(self, ahead=ahead, undecided=tuple(r for r in self.undecided if r not in folded))
863
+
864
+
865
+ @dataclass(frozen=True, kw_only=True)
866
+ class Program:
867
+ """A complete declarative description of a mathematical program, with no data in it.
868
+
869
+ Every group of declarations is keyed by the name the file wrote, in the
870
+ order it wrote them, and is read-only: the mappings are wrapped at
871
+ construction, so a consumer cannot rewrite what another consumer reads.
872
+ A whole program is not hashable — the declarations and expression nodes
873
+ inside it are, which is what dedup and memoisation ask for.
874
+ """
875
+
876
+ parameters: Mapping[str, ParameterDeclaration]
877
+ variables: Mapping[str, VariableDeclaration]
878
+ constraints: Mapping[str, ConstraintDeclaration]
879
+ #: ``None`` where the file declares no objective — a feasibility problem,
880
+ #: whose answer is whether the constraints can be met at all.
881
+ objective: ObjectiveDeclaration | None
882
+ dimensions: Mapping[str, DimensionDeclaration] = Sealed({})
883
+ relations: Mapping[str, RelationDeclaration] = Sealed({})
884
+ sos: Mapping[str, SosDeclaration] = Sealed({})
885
+ #: Each ``piecewise:`` block the model still declares, as the curve it
886
+ #: states; empty on a program of a model whose curves are written out.
887
+ piecewise: Mapping[str, PiecewiseDeclaration] = Sealed({})
888
+ #: What the data has to satisfy for the answer to mean anything, by the
889
+ #: name a refusal quotes: every ``assumptions:`` entry the file wrote, then
890
+ #: what each ``piecewise:`` block's method assumes of its breakpoints. The
891
+ #: language decides none of it, so the consumer attaching the data checks
892
+ #: each and refuses with :func:`assumption_message`.
893
+ assumptions: Mapping[str, Assumption] = Sealed({})
894
+ #: Declared ``expressions:``, each saying whether the math reads it. None
895
+ #: builds a row of its own — one the math reads stands as a :class:`Named`
896
+ #: where it is read — but all are lowered with the program, so a file whose
897
+ #: named expression is outside the language is refused by every verb that
898
+ #: reads the file rather than only by the one that reads the expression.
899
+ expressions: Mapping[str, ExpressionDeclaration] = Sealed({})
900
+ #: What the file as a whole is, as its ``description:`` says.
901
+ description: str | None = None
902
+
903
+ def __post_init__(self) -> None:
904
+ """Seal every group, so a program handed out cannot be written to."""
905
+ for f in fields(self):
906
+ group = getattr(self, f.name)
907
+ if isinstance(group, Mapping):
908
+ object.__setattr__(self, f.name, Sealed(group))
909
+
910
+ def _by_position(self) -> Iterator[tuple[QuadraticPosition, tuple[Expression, ...]]]:
911
+ """The row-building expressions, grouped by the position they stand in."""
912
+ yield 'objective', (self.objective.expression,) if self.objective is not None else ()
913
+ yield 'constraint', tuple(side for c in self.constraints.values() for side in (c.lhs, c.rhs))
914
+
915
+ def relations_of(self, dimension: str) -> Mapping[str, RelationDeclaration]:
916
+ """The relations with a column over *dimension*, by name."""
917
+ return Sealed({n: lk for n, lk in self.relations.items() if dimension in lk.dims})
918
+
919
+ @property
920
+ def roots(self) -> tuple[Expression, ...]:
921
+ """Every tree a row is built from — the objective and both sides of each constraint.
922
+
923
+ An :attr:`expressions` entry builds no row and is not among them. Nor is
924
+ a curve still under :attr:`piecewise`: it is not a row until
925
+ :meth:`~mathspec.model.Spec.expand` writes it out, and its rows are in
926
+ the program of the expansion.
927
+ """
928
+ return tuple(e for _, group in self._by_position() for e in group)
929
+
930
+ @cached_property
931
+ def footprint(self) -> Footprint:
932
+ """Which constructs this program uses — walked once, then held.
933
+
934
+ It answers for the rows this program holds. A curve still under
935
+ :attr:`piecewise` is not counted, so a ``sos2`` curve adds no set order
936
+ here; ask the program of ``spec.expand('piecewise')`` for its rows.
937
+ """
938
+ return Footprint(
939
+ quadratic=frozenset(
940
+ position for position, group in self._by_position() if any(is_quadratic(e) for e in group)
941
+ ),
942
+ domains=frozenset(v.domain for v in self.variables.values()),
943
+ sos_types=frozenset(s.sos_type for s in self.sos.values()),
944
+ kinds=frozenset(type(node) for node in walk(*self.roots)),
945
+ )
946
+
947
+ @cached_property
948
+ def separability(self) -> Mapping[str, Separability]:
949
+ """Every axis, to what building it a window at a time asks and what it would break.
950
+
951
+ The locality :doc:`the ceiling </about/ceiling>` argues in — pointwise,
952
+ bounded halo, global — asked about the axes rather than about the
953
+ operators, so a driver may know before it cuts a horizon whether every
954
+ row it builds is complete inside some window.
955
+
956
+ **A reduction means opposite things by position**, which is the whole of
957
+ the care: in a constraint a sum over the axis ties every window to every
958
+ other, and in the objective it is additively separable, an objective
959
+ being a sum already.
960
+
961
+ Every declared dimension has an entry, an axis nothing mentions being
962
+ trivially windowable. Walked once and held, like :attr:`footprint` and
963
+ for the same reason — a program cannot change after construction — and
964
+ answering for every axis costs what answering for one did, every
965
+ construct that ties an axis naming the axis it ties (#248).
966
+
967
+ It answers for the rows this program holds, as :attr:`footprint` does.
968
+ A curve still under :attr:`piecewise` ties nothing here, although its
969
+ rows sum over its breakpoint dimension; ask the program of
970
+ ``spec.expand('piecewise')``.
971
+ """
972
+ from mathspec.separability import separabilities
973
+
974
+ return Sealed(separabilities(self))
975
+
976
+
977
+ # --------------------------------------------------------------------------
978
+ # Walks, and the questions asked through them
979
+ # --------------------------------------------------------------------------
980
+
981
+
982
+ def walk_regions(*expressions: Expression) -> Iterator[tuple[Expression, tuple[Mask, ...]]]:
983
+ """Every node under *expressions*, each with the regions it stands inside, outermost first.
984
+
985
+ The traversal every *question* about a program is a filter of — which names
986
+ it mentions, whether a variable stands under it, which divisions it
987
+ contains, which rows a piece owes data at. One generator rather than that
988
+ five-line recursion once per question: how a program is traversed is one
989
+ fact, so a node kind :func:`children` learns to descend into reaches every
990
+ caller at once rather than the callers that remembered.
991
+
992
+ The regions are the ``when`` of every :class:`Cases` region the node's
993
+ value stands under, the outermost first, which is the order the masks
994
+ conjoin in. A node outside any ``cases:`` block carries the empty tuple,
995
+ and a ``Cases`` node carries only the regions above it, not its own. The
996
+ tuple rather than one conjoined mask: what a consumer does with the
997
+ regions is its own, and the conjunction is one ``&`` away.
998
+ """
999
+ yield from _walk_regions(expressions, ())
1000
+
1001
+
1002
+ def _walk_regions(
1003
+ expressions: tuple[Expression, ...], above: tuple[Mask, ...]
1004
+ ) -> Iterator[tuple[Expression, tuple[Mask, ...]]]:
1005
+ """The recursion under :func:`walk_regions`, with the regions above *expressions* carried down.
1006
+
1007
+ A ``Cases`` descends by its regions rather than by :func:`children`, because
1008
+ only the region pairs a value with its ``when``; every other node kind
1009
+ descends by :func:`children`, which stays the one home of what sits under
1010
+ a node.
1011
+ """
1012
+ for expression in expressions:
1013
+ yield expression, above
1014
+ if isinstance(expression, Cases):
1015
+ for region in expression.regions:
1016
+ yield from _walk_regions((region.value,), (*above, region.when))
1017
+ else:
1018
+ yield from _walk_regions(children(expression), above)
1019
+
1020
+
1021
+ def walk(*expressions: Expression) -> Iterator[Expression]:
1022
+ """Every node under *expressions*, each expression itself included, parents first.
1023
+
1024
+ :func:`walk_regions` with the regions dropped, for the questions that do
1025
+ not ask where a node stands.
1026
+ """
1027
+ return (node for node, _ in walk_regions(*expressions))
1028
+
1029
+
1030
+ def is_quadratic(expression: Expression) -> bool:
1031
+ """Whether *expression* contains a product of two variable-carrying operands.
1032
+
1033
+ A structural question over the program, and unrelated consumers ask it —
1034
+ what a solver must support, which declarations to build last, whether this
1035
+ form can be represented at all — so it is answered once here beside the
1036
+ other walks rather than once per consumer in its own terms.
1037
+
1038
+ Whether a degree *may be written* is the language's verdict, and this is
1039
+ not a second opinion on it: by the time a program exists the question is
1040
+ which shape the expression has, and the program is what is in hand to
1041
+ answer it.
1042
+ """
1043
+ return any(
1044
+ isinstance(node, Multiply) and all(carries_variable(side) for side in (node.left, node.right))
1045
+ for node in walk(expression)
1046
+ )
1047
+
1048
+
1049
+ def carries_variable(expression: Expression) -> bool:
1050
+ """Whether a variable appears anywhere under *expression*."""
1051
+ return any(isinstance(node, Variable) for node in walk(expression))
1052
+
1053
+
1054
+ def parameters_of(*expressions: Expression) -> frozenset[str]:
1055
+ """Every parameter named anywhere under *expressions*."""
1056
+ return frozenset(node.name for node in walk(*expressions) if isinstance(node, Parameter))
1057
+
1058
+
1059
+ def variables_of(*expressions: Expression) -> frozenset[str]:
1060
+ """Every variable named anywhere under *expressions*."""
1061
+ return frozenset(node.name for node in walk(*expressions) if isinstance(node, Variable))
1062
+
1063
+
1064
+ # ---------------------------------------------------------------------------
1065
+ # the resolved where vocabulary
1066
+ # ---------------------------------------------------------------------------
1067
+
1068
+
1069
+ PredicateOperator = Literal['<=', '>=', '==', '!=', '<', '>']
1070
+
1071
+
1072
+ @dataclass(frozen=True)
1073
+ class BooleanLiteral:
1074
+ value: bool
1075
+
1076
+
1077
+ @dataclass(frozen=True)
1078
+ class ParameterDefined:
1079
+ """True wherever the named parameter is non-null and finite.
1080
+
1081
+ ``dims`` is the parameter's own, copied off the declaration during
1082
+ resolution; every leaf below that names a declaration carries its dims
1083
+ (or ``over``) the same way.
1084
+ """
1085
+
1086
+ name: str
1087
+ dims: tuple[str, ...]
1088
+
1089
+
1090
+ @dataclass(frozen=True)
1091
+ class VariableDefined:
1092
+ """True at the coordinates where the named variable exists."""
1093
+
1094
+ name: str
1095
+ dims: tuple[str, ...]
1096
+
1097
+
1098
+ @dataclass(frozen=True)
1099
+ class ParameterComparison:
1100
+ """Compare a parameter against a literal, element-wise."""
1101
+
1102
+ name: str
1103
+ op: PredicateOperator
1104
+ value: float | str
1105
+ dims: tuple[str, ...]
1106
+
1107
+
1108
+ @dataclass(frozen=True)
1109
+ class ExpressionComparison:
1110
+ """Compare two variable-free expressions, coordinate by coordinate — ``p_min <= 0.5 * p_max``.
1111
+
1112
+ ``dims`` is every dim either side carries. A side whose value is absent at
1113
+ a coordinate — a parameter row missing, a translation that vacated it —
1114
+ makes the comparison false there, as a null does in every other
1115
+ comparison; under a summing operator the absent term is one fewer.
1116
+ """
1117
+
1118
+ left: Expression
1119
+ op: PredicateOperator
1120
+ right: Expression
1121
+ dims: tuple[str, ...]
1122
+
1123
+
1124
+ @dataclass(frozen=True)
1125
+ class DimensionComparison:
1126
+ """Compare a dimension's own coordinates against a literal."""
1127
+
1128
+ name: str
1129
+ op: PredicateOperator
1130
+ value: float | str | datetime.date
1131
+
1132
+
1133
+ @dataclass(frozen=True)
1134
+ class DimensionPosition:
1135
+ """Compare where a row sits along a dimension against a position — ``position(snapshot) == 0``.
1136
+
1137
+ Both sides are integers, negative counting from the end. With a
1138
+ ``partition`` the position is counted within each group the relation makes
1139
+ (:class:`Partition`), whose joined columns' dimensions the frame carries.
1140
+ """
1141
+
1142
+ name: str
1143
+ op: PredicateOperator
1144
+ position: int
1145
+ partition: Partition | None = None
1146
+
1147
+
1148
+ @dataclass(frozen=True)
1149
+ class RelationComparison:
1150
+ """Compare one value column of a keyed relation against a literal — ``period_of == 2030``.
1151
+
1152
+ ``column`` is the role read, and ``dims`` the dimensions of the key
1153
+ columns: the leaf is read at them, one value per coordinate.
1154
+ """
1155
+
1156
+ name: str
1157
+ column: str
1158
+ op: PredicateOperator
1159
+ value: float | str | datetime.date
1160
+ dims: tuple[str, ...] = ()
1161
+
1162
+
1163
+ @dataclass(frozen=True)
1164
+ class RelationPairComparison:
1165
+ """Compare a value column of one keyed relation with one of another — ``from_bus != to_bus`` — row by row on the key.
1166
+
1167
+ Both keys are over the same ``dims``, and the two columns are over one
1168
+ dimension, so a match is possible at all.
1169
+ """
1170
+
1171
+ name: str
1172
+ column: str
1173
+ other: str
1174
+ other_column: str
1175
+ op: PredicateOperator
1176
+ dims: tuple[str, ...] = ()
1177
+
1178
+
1179
+ @dataclass(frozen=True)
1180
+ class RelationDefined:
1181
+ """True where the relation has a row at the frame's coordinates.
1182
+
1183
+ ``dims`` is what the frame supplies: the key's dimensions, whose row is
1184
+ then the one the key finds — every column's for a bare relation, where a
1185
+ row is the whole tuple.
1186
+ """
1187
+
1188
+ name: str
1189
+ dims: tuple[str, ...] = ()
1190
+
1191
+
1192
+ @dataclass(frozen=True)
1193
+ class CountComparison:
1194
+ """How many coordinates *predicate* admits along *over*, against a literal — ``count(points, over=bp) >= 2``.
1195
+
1196
+ The count is one number per coordinate of ``dims``, which is every dim
1197
+ *predicate* reads minus *over*, so a claim about each curve is written
1198
+ without saying "each curve". A predicate a leaf reads arrives as a
1199
+ :class:`Mask`, where a connective's operand is a bare :data:`Predicate`:
1200
+ a walk recurses through the second and stops at the first.
1201
+ """
1202
+
1203
+ predicate: Mask
1204
+ over: str
1205
+ op: PredicateOperator
1206
+ value: float
1207
+ dims: tuple[str, ...]
1208
+
1209
+
1210
+ @dataclass(frozen=True)
1211
+ class TranslatedPredicate:
1212
+ """*operand* read at a neighbouring coordinate — ``shift(points, along=bp, offset=1)``.
1213
+
1214
+ False where the translation vacates, and there is no ``edge=`` to state.
1215
+ The arithmetic translation needs one because no number is neutral and
1216
+ inventing one changes the answer; false is what a missing row already
1217
+ means in a mask, so the predicate form has the value the language already
1218
+ gives it.
1219
+ """
1220
+
1221
+ operand: Mask
1222
+ along: str
1223
+ offset: int
1224
+ dims: tuple[str, ...]
1225
+
1226
+
1227
+ @dataclass(frozen=True)
1228
+ class PulledBackPredicate:
1229
+ """*operand* read through a relation — ``at(has_curve, by=converter_of, over=converter, into=flow)``.
1230
+
1231
+ True at a coordinate where the relation has a row and *operand* holds at
1232
+ the coordinate that row reads. False where the relation has no row, which
1233
+ is what a missing row already means in a mask. The dims ``direction``
1234
+ consumes go and the dims it produces arrive, as :class:`Pullback`'s do.
1235
+ """
1236
+
1237
+ operand: Mask
1238
+ direction: Direction
1239
+ dims: tuple[str, ...]
1240
+
1241
+
1242
+ @dataclass(frozen=True)
1243
+ class Not:
1244
+ operand: Predicate
1245
+
1246
+
1247
+ @dataclass(frozen=True)
1248
+ class And:
1249
+ left: Predicate
1250
+ right: Predicate
1251
+
1252
+
1253
+ @dataclass(frozen=True)
1254
+ class Or:
1255
+ left: Predicate
1256
+ right: Predicate
1257
+
1258
+
1259
+ #: Every predicate resolution has typed: it names a declaration and the kind is
1260
+ #: settled. Resolution passes these straight through, having nothing left to
1261
+ #: decide about them.
1262
+ TypedPredicate = (
1263
+ ParameterComparison
1264
+ | ExpressionComparison
1265
+ | ParameterDefined
1266
+ | VariableDefined
1267
+ | DimensionComparison
1268
+ | DimensionPosition
1269
+ | RelationComparison
1270
+ | RelationPairComparison
1271
+ | RelationDefined
1272
+ | CountComparison
1273
+ | TranslatedPredicate
1274
+ | PulledBackPredicate
1275
+ )
1276
+
1277
+ #: The boolean connectives — the only where nodes carrying other where nodes,
1278
+ #: and so the only place a walk over a predicate recurses. The grammar builds
1279
+ #: these classes directly, over leaves still unresolved, so a pre-resolution
1280
+ #: tree shares them — the transient impurity resolution normalizes away.
1281
+ Connective = Not | And | Or
1282
+
1283
+ #: Every resolved predicate node. The parser's ``Unresolved*`` nodes are not members: they live with the
1284
+ #: grammar in :mod:`mathspec._where_parser`, and resolution rewrites them away
1285
+ #: before anything here is asked.
1286
+ Predicate = BooleanLiteral | TypedPredicate | Connective
1287
+
1288
+
1289
+ def where_children(where: Predicate) -> tuple[Predicate, ...]:
1290
+ """The predicates under *where* — a connective's operands, and nothing under a leaf.
1291
+
1292
+ What every walk over a predicate recurses through, as :func:`children` is
1293
+ for an expression. A leaf has nothing under it whether or not it is
1294
+ resolved, so the grammar measures its own output with this too.
1295
+ """
1296
+ if isinstance(where, Not):
1297
+ return (where.operand,)
1298
+ if isinstance(where, (And, Or)):
1299
+ return (where.left, where.right)
1300
+ return ()
1301
+
1302
+
1303
+ def _atoms(where: Predicate) -> Iterator[TypedPredicate]:
1304
+ """Every node in *where* that reads a declaration, connectives removed.
1305
+
1306
+ A boolean literal yields nothing.
1307
+
1308
+ Raises:
1309
+ AssertionError: An unresolved node reached the walk.
1310
+ """
1311
+ if isinstance(where, TypedPredicate):
1312
+ yield where
1313
+ elif isinstance(where, BooleanLiteral | Connective):
1314
+ for child in where_children(where):
1315
+ yield from _atoms(child)
1316
+ else:
1317
+ msg = f'{type(where).__name__} reached a predicate walk unresolved.'
1318
+ raise AssertionError(msg)
1319
+
1320
+
1321
+ def _atom_dims(atom: TypedPredicate) -> frozenset[str]:
1322
+ """One leaf's dims — the rule :attr:`Mask.dims` is the union of.
1323
+
1324
+ A parameter or variable leaf carries its own dims off the declaration; a
1325
+ comparison on a dimension is read through that dimension, and a relation
1326
+ through the dimensions of the columns it is read at — its key for a
1327
+ comparison, every column for a bare existence.
1328
+ Separate from the union because the load-time frame check reports per
1329
+ leaf. Closed by ``assert_never``: a predicate node added without a reading
1330
+ is a type error here, at the one place that has to grow a branch, rather
1331
+ than a wrong dim set at the first model to use it.
1332
+ """
1333
+ match atom:
1334
+ case (
1335
+ ParameterComparison()
1336
+ | ExpressionComparison()
1337
+ | ParameterDefined()
1338
+ | VariableDefined()
1339
+ | CountComparison()
1340
+ | TranslatedPredicate()
1341
+ | RelationComparison()
1342
+ | RelationPairComparison()
1343
+ | RelationDefined()
1344
+ | PulledBackPredicate()
1345
+ ):
1346
+ return frozenset(atom.dims)
1347
+ case DimensionComparison():
1348
+ return frozenset({atom.name})
1349
+ case DimensionPosition():
1350
+ return frozenset({atom.name, *(atom.partition.joined_dims if atom.partition is not None else ())})
1351
+ case _:
1352
+ assert_never(atom)
1353
+
1354
+
1355
+ def _atom_names(atom: TypedPredicate) -> frozenset[str]:
1356
+ """One leaf's declarations, its dimension apart — the rule :attr:`Mask.names_read` is the union of.
1357
+
1358
+ A comparison on a dimension names no declaration — a coordinate is not
1359
+ data to feed — a relation pair names both maps it compares, and a
1360
+ comparison of expressions names every parameter and relation its sides
1361
+ read, answered on the program's form of it since only a program mask is
1362
+ asked. ``assert_never``-closed for the reason :func:`_atom_dims` is: a predicate
1363
+ node added without a reading is a type error at this one branch rather
1364
+ than a name silently dropped at the first model to use it.
1365
+ """
1366
+ match atom:
1367
+ case ParameterComparison() | ParameterDefined() | VariableDefined() | RelationComparison() | RelationDefined():
1368
+ return frozenset({atom.name})
1369
+ case RelationPairComparison():
1370
+ return frozenset({atom.name, atom.other})
1371
+ case ExpressionComparison():
1372
+ return _names_under(atom.left, atom.right)
1373
+ case CountComparison():
1374
+ return atom.predicate.names_read
1375
+ case TranslatedPredicate():
1376
+ return atom.operand.names_read
1377
+ case PulledBackPredicate():
1378
+ return atom.operand.names_read | {atom.direction.name}
1379
+ case DimensionComparison() | DimensionPosition():
1380
+ return frozenset()
1381
+ case _:
1382
+ assert_never(atom)
1383
+
1384
+
1385
+ def _names_under(*expressions: Expression) -> frozenset[str]:
1386
+ """Every parameter and relation the data has to supply for *expressions* — what a mask's ``names_read`` promises.
1387
+
1388
+ :func:`parameters_of` alone misses the data an operator reads beside its
1389
+ operand: the relation a grouping or a pullback reads through, the one a
1390
+ translation or a window is partitioned by, the parameter a named offset or
1391
+ width is read from, and whatever decides which region of a cased value
1392
+ applies.
1393
+ """
1394
+ names: set[str] = set(parameters_of(*expressions))
1395
+ for node in walk(*expressions):
1396
+ if isinstance(node, Cases):
1397
+ names.update(*(region.when.names_read for region in node.regions))
1398
+ elif isinstance(node, (GroupSum, Pullback)):
1399
+ names.add(node.direction.name)
1400
+ elif isinstance(node, (Translate, WindowSum)):
1401
+ if node.partition is not None:
1402
+ names.add(node.partition.name)
1403
+ amount = node.offset if isinstance(node, Translate) else node.width
1404
+ if isinstance(amount, str):
1405
+ names.add(amount)
1406
+ return frozenset(names)
1407
+
1408
+
1409
+ def _conjuncts(where: Predicate) -> tuple[Predicate, ...]:
1410
+ """The flatten rule behind :attr:`Mask.conjuncts` — the one home of the split.
1411
+
1412
+ ``a AND b AND c`` gives three, and a predicate that is not an ``AND`` gives
1413
+ itself. The walk stops at the first node that is not an ``AND``: the
1414
+ conjuncts of ``a AND (b OR c)`` are ``a`` and ``b OR c``, and of
1415
+ ``NOT (a AND b)`` the single ``NOT`` — neither an ``OR`` nor a ``NOT`` is a
1416
+ claim the predicate makes on its own, so neither is split.
1417
+ """
1418
+ if isinstance(where, And):
1419
+ return _conjuncts(where.left) + _conjuncts(where.right)
1420
+ return (where,)
1421
+
1422
+
1423
+ def _fold(node: Predicate) -> Predicate:
1424
+ """*node* with every connective a literal or a double negation decides evaluated away.
1425
+
1426
+ ``X AND True`` is ``X``, ``X OR True`` is every row, ``X AND False`` is
1427
+ none, ``NOT True`` is ``False`` and ``NOT NOT X`` is ``X``. What survives
1428
+ is a predicate over data, or the one literal the whole mask reduces to —
1429
+ the invariant :class:`Mask` applies at construction, so it holds wherever
1430
+ a mask is built.
1431
+ """
1432
+ if isinstance(node, Not):
1433
+ operand = _fold(node.operand)
1434
+ if isinstance(operand, BooleanLiteral):
1435
+ return BooleanLiteral(not operand.value)
1436
+ if isinstance(operand, Not):
1437
+ return operand.operand
1438
+ return Not(operand)
1439
+ if isinstance(node, And):
1440
+ left, right = _fold(node.left), _fold(node.right)
1441
+ if isinstance(left, BooleanLiteral):
1442
+ return right if left.value else left
1443
+ if isinstance(right, BooleanLiteral):
1444
+ return left if right.value else right
1445
+ return And(left, right)
1446
+ if isinstance(node, Or):
1447
+ left, right = _fold(node.left), _fold(node.right)
1448
+ if isinstance(left, BooleanLiteral):
1449
+ return left if left.value else right
1450
+ if isinstance(right, BooleanLiteral):
1451
+ return right if right.value else left
1452
+ return Or(left, right)
1453
+ return node
1454
+
1455
+
1456
+ @dataclass(frozen=True)
1457
+ class Mask:
1458
+ """A resolved ``where`` and the questions the language answers about it.
1459
+
1460
+ ``root`` is the predicate a consumer dispatches on with ``isinstance``;
1461
+ every question below is derived from it. Construction folds, so a boolean
1462
+ literal stands at the root or nowhere in it, and refuses an unresolved
1463
+ tree.
1464
+
1465
+ Attributes:
1466
+ root: The resolved predicate the mask restricts rows by, folded.
1467
+ """
1468
+
1469
+ root: Predicate
1470
+
1471
+ def __post_init__(self) -> None:
1472
+ object.__setattr__(self, 'root', _fold(self.root))
1473
+ _ = self.atoms # the walk is the refusal, and runs after the fold
1474
+
1475
+ @cached_property
1476
+ def atoms(self) -> tuple[TypedPredicate, ...]:
1477
+ """The mask's leaves, connectives removed — the one walk the other questions read.
1478
+
1479
+ Held rather than re-walked: construction takes this walk anyway, to
1480
+ refuse an unresolved tree, and a mask cannot change afterwards.
1481
+ """
1482
+ return tuple(_atoms(self.root))
1483
+
1484
+ @property
1485
+ def conjuncts(self) -> tuple[Predicate, ...]:
1486
+ """The predicates the mask joins with ``AND`` — its ``AND`` spine flattened, stopping at an ``OR`` or a ``NOT``."""
1487
+ return _conjuncts(self.root)
1488
+
1489
+ @property
1490
+ def names_read(self) -> frozenset[str]:
1491
+ """The parameters, relations and variables the mask names."""
1492
+ return frozenset(name for atom in self.atoms for name in _atom_names(atom))
1493
+
1494
+ @property
1495
+ def dims(self) -> frozenset[str]:
1496
+ """The dims the mask is read at — the union of what each leaf carries.
1497
+
1498
+ Empty for a mask over nothing but literals. Read off the leaves, which
1499
+ resolution stamped with their declarations' dims, so a predicate built
1500
+ from resolved pieces answers exactly as a declaration's own does.
1501
+ """
1502
+ return frozenset(dim for atom in self.atoms for dim in _atom_dims(atom))
1503
+
1504
+ def __invert__(self) -> Mask:
1505
+ """The mask admitting exactly the rows this one refuses — construction folds a double negation or a literal flip."""
1506
+ return Mask(Not(self.root))
1507
+
1508
+ def __and__(self, other: Mask) -> Mask:
1509
+ """Both masks at once — construction absorbs a literal side rather than burying it."""
1510
+ return Mask(And(self.root, other.root))
1511
+
1512
+ def __or__(self, other: Mask) -> Mask:
1513
+ """Either mask — construction absorbs a literal side rather than burying it."""
1514
+ return Mask(Or(self.root, other.root))