diff --git a/src/math_spec/_expression_parser.py b/src/math_spec/_expression_parser.py index 76126006..bb404e82 100644 --- a/src/math_spec/_expression_parser.py +++ b/src/math_spec/_expression_parser.py @@ -132,24 +132,23 @@ def __str__(self) -> str: @dataclass(frozen=True) -class RelationNode: - """A resolved ``by=`` — the relation, and the use the call makes of it. - - ``use`` is a :class:`Direction` for ``sum`` and ``at``, and a - :class:`Partition` for ``shift`` and ``sum_back``. - ``dimensions`` is the fine side — what ``sum`` consumes and ``at`` - produces — and ``into`` the coarse dims, which ``sum`` produces and ``at`` - consumes. The roles joined on are the operand's to carry, and the operator - passes them through. - """ +class DirectionNode: + """A resolved ``by=`` on ``sum`` or ``at``: the relation, read in the :class:`Direction` the call names.""" - name: str - dimensions: tuple[str, ...] - into: tuple[str, ...] - use: Direction | Partition + direction: Direction def __str__(self) -> str: - return self.name + return self.direction.name + + +@dataclass(frozen=True) +class PartitionNode: + """A resolved ``by=`` on ``shift`` or ``sum_back``: the relation, as the :class:`Partition` the call steps inside.""" + + partition: Partition + + def __str__(self) -> str: + return self.partition.name @dataclass(frozen=True) @@ -284,7 +283,8 @@ def __str__(self) -> str: | ParameterNode | DualNode | DimensionNode - | RelationNode + | DirectionNode + | PartitionNode | EdgeNode | KeywordNode | UnaryOperatorNode @@ -338,7 +338,7 @@ def operand(node: ArithmeticNode) -> str: #: ``sum(x, along=d)``, ``sum(x, by=l)``, ``shift(..., edge='wrap')``. None of #: the three is data, so none may stand in arithmetic — which is why the passes #: that walk a value position refuse them together. -KwargNode = DimensionNode | RelationNode | EdgeNode +KwargNode = DimensionNode | DirectionNode | PartitionNode | EdgeNode #: What resolution rewrites away: a bare name, whose kind only the schema #: knows, and the two kwarg-only literals its kwarg consumes. Meeting one diff --git a/src/math_spec/advice.py b/src/math_spec/advice.py index 7458e466..f8e7a5b8 100644 --- a/src/math_spec/advice.py +++ b/src/math_spec/advice.py @@ -69,12 +69,11 @@ def _never_an_axis(program: Program) -> list[Advice]: def _produced_axes(program: Program) -> set[str]: """The axes the expressions create beyond what any declaration indexes. - ``sum(by=)`` lands on its target and ``at()`` spreads onto its fine dimension. + ``sum(by=)`` lands on its target and ``at()`` spreads onto its fine + dimension: either way, the dims the direction produces. """ axes: set[str] = set() for node in walk(*program.expressions): - if isinstance(node, GroupSum): - axes.update(node.into) - elif isinstance(node, At): - axes.update(node.over) + if isinstance(node, GroupSum | At): + axes.update(node.direction.produced_dims) return axes diff --git a/src/math_spec/dimensions.py b/src/math_spec/dimensions.py index 986d8135..6bec4fbf 100644 --- a/src/math_spec/dimensions.py +++ b/src/math_spec/dimensions.py @@ -23,6 +23,7 @@ ComparisonNode, DefinitionNode, DimensionNode, + DirectionNode, DualNode, EdgeNode, FunctionCallNode, @@ -30,7 +31,7 @@ NumberNode, ParameterNode, ParsedNode, - RelationNode, + PartitionNode, UnaryOperatorNode, UnresolvedNode, VariableNode, @@ -42,6 +43,7 @@ from math_spec.program import ( DimensionComparisonNode, DimensionPositionNode, + Direction, Mask, ParameterComparisonNode, ParameterDefinedNode, @@ -154,36 +156,33 @@ def _sum_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spec, conte ) return inner - {consumed.name} - assert isinstance(by, RelationNode) - if missing := sorted(set(by.dimensions) - inner): + assert isinstance(by, DirectionNode), 'resolution reads sum(by=) in a direction' + direction = by.direction + if missing := sorted(set(direction.consumed_dims) - inner): raise DimensionError( _not_carried( context, - f'sum(by={by}) consumes {missing}, the dims it reads from,', + f'sum(by={direction.name}) consumes {missing}, the dims it reads from,', inner, 'drop the sum, or fix the dim', ) ) - _check_lands_clear(f'sum(by={by})', set(by.into), set(by.dimensions), inner, context) - _check_joined(f'sum(by={by})', by, inner, context) - return (inner - set(by.dimensions)) | set(by.into) + return _read_dims(f'sum(by={direction.name})', direction, inner, context) def _at_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spec, context: str) -> frozenset[str]: """``at`` is the adjoint of ``sum(by=)``: it consumes the dims a sum produces and produces the ones it consumes.""" by = node.kwargs['by'] - assert isinstance(by, RelationNode) - absent = sorted(set(by.into) - inner) - if absent: + assert isinstance(by, DirectionNode), 'resolution reads at(by=) in a direction' + direction = by.direction + if absent := sorted(set(direction.consumed_dims) - inner): raise DimensionError( - f'{context}: at(by={by}) reads through ' + f'{context}: at(by={direction.name}) reads through ' f'{absent}, which the expression does not carry (dims ' f'{sorted(inner)}). A pullback needs the coarse dims to read *from* — ' f'sum is the direction that produces them.' ) - _check_lands_clear(f'at(by={by})', set(by.dimensions), set(by.into), inner, context) - _check_joined(f'at(by={by})', by, inner, context) - return (inner - set(by.into)) | set(by.dimensions) + return _read_dims(f'at(by={direction.name})', direction, inner, context) def _translation_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spec, context: str) -> frozenset[str]: @@ -202,22 +201,24 @@ def _translation_dims(node: FunctionCallNode, inner: frozenset[str], schema: Spe _check_named_amount(node, over.name, inner, schema, context) _check_amount_form(node, context) _check_edge(node, context) - partition = node.kwargs.get('by') - if partition is not None: - assert isinstance(partition, RelationNode) - _check_joined(f'{node.name}(along={over.name}, by={partition})', partition, inner, context) + by = node.kwargs.get('by') + if by is not None: + assert isinstance(by, PartitionNode), "resolution reads a translation's by= as a partition" + _check_joined(f'{node.name}(along={over.name}, by={by.partition.name})', by.partition, inner, context) return inner -def _check_lands_clear(call: str, produced: set[str], consumed: set[str], inner: frozenset[str], context: str) -> None: - """The dims a call lands on are its own to bring, so the operand does not already carry one. +def _read_dims(call: str, direction: Direction, inner: frozenset[str], context: str) -> frozenset[str]: + """The dims after a relation is read in *direction*: the consumed go, the produced arrive, the joined stay. - Where it does, the call would tie the operand's axis to the one it - produces rather than adding it, and it reads the same either way. - A relation into its own dimension is not that case: there the dim landed - on is the dim just consumed, so every factor is read at the coordinate - the sum runs over, and nothing is tied. + The dims a call lands on are its own to bring, so the operand does not + already carry one. Where it does, the call would tie the operand's axis to + the one it produces rather than adding it, and it reads the same either + way. A relation into its own dimension is not that case: there the dim + landed on is the dim just consumed, so every factor is read at the + coordinate the sum runs over, and nothing is tied. """ + consumed, produced = set(direction.consumed_dims), set(direction.produced_dims) if clash := sorted((produced & inner) - consumed): raise DimensionError( f'{context}: {call} lands on {clash}, which the expression already carries.\n' @@ -225,11 +226,17 @@ def _check_lands_clear(call: str, produced: set[str], consumed: set[str], inner: f'adds. Move the factor carrying {clash} outside the operator, or read to a column ' f'over another dimension.' ) + _check_joined(call, direction, inner, context) + return (inner - consumed) | produced + +def _check_joined(call: str, use: Direction | Partition, inner: frozenset[str], context: str) -> None: + """The columns a call joins on are read at their dimensions, so the operand carries every one, each once. -def _check_joined(call: str, by: RelationNode, inner: frozenset[str], context: str) -> None: - """The columns a call joins on are read at their dimensions, so the operand carries every one, each once.""" - use = by.use + A joined dimension the call also consumes is the same ambiguity as two + joined columns over one dimension: the operand's one coordinate would + have to be read as both. A partition consumes nothing. + """ dims = use.joined_dims if missing := sorted(set(dims) - inner): raise DimensionError( @@ -238,7 +245,8 @@ def _check_joined(call: str, by: RelationNode, inner: frozenset[str], context: s f'read between two of its columns and joined at the others — index the operand by them, or ' f'read it between different columns.' ) - if twice := sorted({d for d in dims if dims.count(d) > 1 or d in by.dimensions}): + consumed = use.consumed_dims if isinstance(use, Direction) else () + if twice := sorted({d for d in dims if dims.count(d) > 1 or d in consumed}): raise DimensionError( f"{context}: {call} joins '{use.name}' on {twice} through more than one column, and the operand " f'carries each dimension once. Read between different columns, or use a relation whose joined ' @@ -421,9 +429,10 @@ def _check_named_amount(node: FunctionCallNode, over: str, inner: frozenset[str] f'carries it. A named {words.noun} that varies over the axis it steps along is {words.varies} ' f"— declare '{amount.name}' over dims '{over}' is not one of." ) - partition = node.kwargs.get('by') - use = partition.use if isinstance(partition, RelationNode) else None - groups = frozenset(use.dim(v) for v in use.group) if isinstance(use, Partition) else frozenset() + by = node.kwargs.get('by') + groups = ( + frozenset(by.partition.dim(v) for v in by.partition.group) if isinstance(by, PartitionNode) else frozenset() + ) if stray := sorted(frozenset(declared.dims) - inner - groups): raise DimensionError( f'{context}: {node.name}({kwarg}={amount.name}) reads its {words.noun} at the coordinate it ' diff --git a/src/math_spec/lowering.py b/src/math_spec/lowering.py index 7d98fc1b..4f18f4bb 100644 --- a/src/math_spec/lowering.py +++ b/src/math_spec/lowering.py @@ -22,13 +22,14 @@ CasesNode, DefinitionNode, DimensionNode, + DirectionNode, DualNode, EdgeNode, FunctionCallNode, KwargNode, NumberNode, ParameterNode, - RelationNode, + PartitionNode, UnaryOperatorNode, UnresolvedNode, VariableNode, @@ -141,11 +142,7 @@ def lower_program(expanded: _ExpandedSpec) -> program.Program: dimensions = { dname: program.DimensionDeclaration( - tuple( - program.RelationDeclaration(lname, lk.pairs, lk.key_roles) - for lname, lk in expanded.relations.items() - if dname in lk.dims - ), + tuple(lk for lk in resolved.relations.values() if dname in lk.dims), ddef.dtype, ) for dname, ddef in expanded.dimensions.items() @@ -272,16 +269,14 @@ def sum(self, node: FunctionCallNode) -> program.ExpressionNode: consumed = node.kwargs['over'] assert isinstance(consumed, DimensionNode), 'resolution refuses a over= that is not a dimension' return program.Sum(operand, (consumed.name,)) - assert isinstance(by_node, RelationNode), 'resolution refuses a by= that is not a relation' - assert isinstance(by_node.use, program.Direction), 'resolution reads sum(by=) in a direction' - return program.GroupSum(operand, direction=by_node.use) + assert isinstance(by_node, DirectionNode), 'resolution reads sum(by=) in a direction' + return program.GroupSum(operand, direction=by_node.direction) def at(self, node: FunctionCallNode) -> program.ExpressionNode: """``at(x, by=relation)`` — the adjoint of :meth:`sum`'s ``by=`` form.""" by_node = node.kwargs['by'] - assert isinstance(by_node, RelationNode), 'resolution refuses a by= that is not a relation' - assert isinstance(by_node.use, program.Direction), 'resolution reads at(by=) in a direction' - return program.At(self.expr(node.args[0]), direction=by_node.use) + assert isinstance(by_node, DirectionNode), 'resolution reads at(by=) in a direction' + return program.At(self.expr(node.args[0]), direction=by_node.direction) def sum_back(self, node: FunctionCallNode) -> program.ExpressionNode: """``sum_back(x, along=d, window=w)`` — a trailing window along one dimension. @@ -353,9 +348,8 @@ def _partition_of(node: FunctionCallNode) -> program.Partition | None: by_node = node.kwargs.get('by') if by_node is None: return None - assert isinstance(by_node, RelationNode) - assert isinstance(by_node.use, program.Partition), "resolution reads a translation's by= as a partition" - return by_node.use + assert isinstance(by_node, PartitionNode), "resolution reads a translation's by= as a partition" + return by_node.partition def _bound_expression(value: float | str) -> program.ExpressionNode: diff --git a/src/math_spec/program.py b/src/math_spec/program.py index c9c80ffe..d718445e 100644 --- a/src/math_spec/program.py +++ b/src/math_spec/program.py @@ -266,70 +266,30 @@ class Sum(Expression): @dataclass(frozen=True) class GroupSum(Expression): - """Sum ``operand`` through a relation, consuming the dims ``over`` and producing ``into``. - - ``direction`` says which columns are consumed, which produced and which - joined on, and is the one fact the node holds: ``over`` is the dims it - consumes, ``into`` the dims it produces, and ``joined`` the dims it joins - on. The result replaces every dim in ``over`` with every dim in ``into`` - and keeps every dim in ``joined``. The join keys on the consumed columns - and every joined column. + """Sum ``operand`` through a relation: the dims ``direction`` consumes go, the dims it produces arrive, the dims it joins on stay. + + The join keys on the consumed columns and every joined column, and the + operand carries every dim consumed or joined on. """ operand: ExpressionNode direction: Direction - @property - def relation(self) -> str: - return self.direction.name - - @property - def over(self) -> tuple[str, ...]: - return self.direction.consumed_dims - - @property - def into(self) -> tuple[str, ...]: - return self.direction.produced_dims - - @property - def joined(self) -> tuple[str, ...]: - """The dims it joins on — the key columns neither consumed nor produced, which the operand carries.""" - return self.direction.joined_dims - @dataclass(frozen=True) class At(Expression): """Read ``operand`` through a relation — the adjoint of :class:`GroupSum`. - The same table read the other way: this consumes the dims in ``into`` - and produces the dims in ``over``, one value per coordinate because the - read takes value columns at a key the operand fixes - (``Direction.is_function_read``). The join fans out, many ``over`` tuples - sharing one ``into`` tuple — at each coordinate of the joined columns, - which the operand carries and the result keeps. As on - :class:`GroupSum`, ``direction`` is the fact and the rest are read off it. + The dims ``direction`` consumes go and the dims it produces arrive, one + value per coordinate because the read takes value columns at a key the + result fixes (``Direction.is_function_read``). The join fans out, many + produced tuples sharing one consumed tuple — at each coordinate of the + joined columns, which the operand carries and the result keeps. """ operand: ExpressionNode direction: Direction - @property - def relation(self) -> str: - return self.direction.name - - @property - def over(self) -> tuple[str, ...]: - return self.direction.produced_dims - - @property - def into(self) -> tuple[str, ...]: - return self.direction.consumed_dims - - @property - def joined(self) -> tuple[str, ...]: - """The dims it joins on — the key columns neither consumed nor produced, which the operand carries.""" - return self.direction.joined_dims - @dataclass(frozen=True) class Translate(Expression): @@ -492,7 +452,7 @@ class RelationDeclaration(NamedTuple): name: str columns: tuple[tuple[str, str], ...] - key: tuple[str, ...] = () + key: tuple[str, ...] @property def roles(self) -> tuple[str, ...]: @@ -531,18 +491,6 @@ class Direction(NamedTuple): def name(self) -> str: return self.relation.name - @property - def key(self) -> tuple[str, ...]: - return self.relation.key - - @property - def roles(self) -> tuple[str, ...]: - return self.relation.roles - - @property - def values(self) -> tuple[str, ...]: - return self.relation.values - def dim(self, role: str) -> str: """The dimension *role* is bound to.""" return self.relation.dim(role) @@ -562,7 +510,7 @@ def joined_dims(self) -> tuple[str, ...]: @property def is_function_read(self) -> bool: """Whether the read is one value per coordinate: the key lies inside what is fixed.""" - return bool(self.key) and set(self.key) <= {*self.joined, *self.produced} + return set(self.relation.key) <= {*self.joined, *self.produced} class Partition(NamedTuple): @@ -586,14 +534,6 @@ class Partition(NamedTuple): def name(self) -> str: return self.relation.name - @property - def key(self) -> tuple[str, ...]: - return self.relation.key - - @property - def values(self) -> tuple[str, ...]: - return self.relation.values - def dim(self, role: str) -> str: """The dimension *role* is bound to.""" return self.relation.dim(role) diff --git a/src/math_spec/resolution.py b/src/math_spec/resolution.py index d28f334c..20065626 100644 --- a/src/math_spec/resolution.py +++ b/src/math_spec/resolution.py @@ -25,6 +25,7 @@ ComparisonNode, DefinitionNode, DimensionNode, + DirectionNode, DualNode, EdgeNode, FunctionCallNode, @@ -35,7 +36,7 @@ NumberNode, ParameterNode, ParsedNode, - RelationNode, + PartitionNode, UnaryOperatorNode, VariableNode, case_context, @@ -161,10 +162,6 @@ def kind(self, name: str) -> DeclarationKind | None: return 'relation' return None - def shape_of(self, relation: str) -> RelationDeclaration: - """The columns and key of *relation*, as declared.""" - return self.relations[relation] - def unknown(self, name: str, context: str, *, allow_dims: bool, formals: Iterable[str] = ()) -> str: """The refusal for a *name* declared nowhere, listing what it could have been. @@ -224,12 +221,16 @@ class Resolved: constraints: Each constraint's comparison and ``where``. objective: The objective's expression, ``None`` where the file declares none. + relations: Each relation's columns and key, as declared — the one + copy, which every :class:`~math_spec.program.Direction` and + :class:`~math_spec.program.Partition` in the trees holds. """ expressions: dict[str, CasesNode | DefinitionNode] variables: dict[str, Mask | None] constraints: dict[str, ResolvedConstraint] objective: ArithmeticNode | None + relations: dict[str, RelationDeclaration] @cached_property def read_by_the_math(self) -> frozenset[str]: @@ -611,17 +612,11 @@ def _relation_ref( return value # the call shape refused it already, with the wording that names the rewrite over_dim = over.name if isinstance(over, NameNode | DimensionNode) else None partition = self._partition(name, operator, over_dim, named['within']) - if partition is None: - return value - return RelationNode(name, dimensions=(partition.along_dim,), into=(), use=partition) + return value if partition is None else PartitionNode(partition) if not ({'over', 'into'} <= set(named)): return value # the call shape refused it already, with the wording that names the rewrite direction = self._direction(name, operator, named['over'], named['into']) - if direction is None: - return value - fine = direction.produced_dims if operator == 'at' else direction.consumed_dims - coarse = direction.consumed_dims if operator == 'at' else direction.produced_dims - return RelationNode(name, dimensions=fine, into=coarse, use=direction) + return value if direction is None else DirectionNode(direction) def _role_name(self, value: ArithmeticNode, operator: str, key: str) -> tuple[str, ...] | None: """``over=`` or ``into=`` as the column names it must be — one bare name, or a bracketed list of them.""" @@ -651,7 +646,7 @@ def _direction( refused toward ``at``. """ ns, context = self.ns, self.context - shape = ns.shape_of(name) + shape = ns.relations[name] call = f'{operator}(by={name})' if not ( self._known_roles(name, call, from_roles, 'over') and self._known_roles(name, call, into_roles, 'into') @@ -673,13 +668,13 @@ def _direction( f'between columns over distinct dimensions.' ) return None - joined = tuple(r for r in (shape.key or shape.roles) if r not in from_roles and r not in into_roles) + joined = tuple(r for r in shape.key if r not in from_roles and r not in into_roles) direction = Direction(shape, from_roles, into_roles, joined) if not forward and not direction.is_function_read: self.errors.append( f"{context}: {call}: at reads one value per coordinate, and '{name}' is not single-valued in " - f'{list(from_roles)} at the columns the operand fixes ({[*into_roles, *joined]}) — its key is ' - f'{list(shape.key)}. Key the table by columns the read fixes, or read the other way.' + f'{list(from_roles)} at the columns the call lands on ({[*into_roles, *joined]}) — its key is ' + f'{list(shape.key)}. Key the table by the columns the call lands on, or read the other way.' ) return None if forward and direction.is_function_read: @@ -692,17 +687,17 @@ def _direction( return None return direction - def _known_roles(self, name: str, call: str, roles: tuple[str, ...] | None, kwarg: str) -> bool: + def _known_roles(self, name: str, call: str, roles: tuple[str, ...], kwarg: str) -> bool: """Whether every role *kwarg* names is a column of relation *name*, each once; the refusal otherwise.""" - shape = self.ns.shape_of(name) - for role in roles or (): + shape = self.ns.relations[name] + for role in roles: if role not in shape.roles: self.errors.append( f"{self.context}: {call}: {kwarg}={role} names no column of '{name}', whose columns are " f'{list(shape.roles)}.' ) return False - if roles is not None and len(set(roles)) < len(roles): + if len(set(roles)) < len(roles): self.errors.append(f'{self.context}: {call}: {kwarg}={list(roles)} names a column twice.') return False return True @@ -719,7 +714,7 @@ def _partition( ``within=`` names a column that is not a value column. """ context = self.context - shape = self.ns.shape_of(name) + shape = self.ns.relations[name] call = f'{operator}(by={name})' if along_dim is None or not self._known_roles(name, call, within_roles, 'within'): return None @@ -811,7 +806,7 @@ def _where_name(self, node: UnresolvedNameNode) -> WhereNode | UnresolvedWhereNo f'Remove it, or compare it: where: "{node.name} > 0".' ) case 'relation': - shape = ns.shape_of(node.name) + shape = ns.relations[node.name] dims = tuple(shape.dim(k) for k in shape.key) if len(set(dims)) < len(dims): self.errors.append( @@ -844,19 +839,15 @@ def _position(self, node: UnresolvedPositionNode) -> DimensionPositionNode | Unr return node if node.by is None: return DimensionPositionNode(node.dimension, node.op, node.position) - call = f'position({node.dimension}, by={node.by})' - if ns.kind(node.by) != 'relation': - self.errors.append( - f"{context}: '{call}' groups by '{node.by}', which is {_declared_as(ns, node.by)}. " - f'``by=`` takes a relation with a key column over that dimension. ' - f'{did_you_mean(node.by, ns.relations, label="Relations")}' - ) + if (problem := self._not_a_relation(node.by, 'position', 'by')) is not None: + self.errors.append(problem) return node + call = f'position({node.dimension}, by={node.by})' if node.into is None: self.errors.append( f'{context}: {call} leaves within= unsaid. {PARTITION_NAMES_ITS_GROUP} Write ' f"position({node.dimension}, by={node.by}, within=) — the value columns of '{node.by}' " - f'are {list(ns.shape_of(node.by).values)}.' + f'are {list(ns.relations[node.by].values)}.' ) return node partition = self._partition(node.by, 'position', node.dimension, node.into) @@ -880,7 +871,7 @@ def _comparison(self, node: UnresolvedComparisonNode) -> WhereNode | UnresolvedW if (refusal := _relation_pair_error(context, node, value, ns, left, right)) is not None: self.errors.append(refusal) return node - dims = tuple(ns.shape_of(left_name).dim(k) for k in ns.shape_of(left_name).key) + dims = tuple(ns.relations[left_name].dim(k) for k in ns.relations[left_name].key) return RelationPairComparisonNode(left_name, left, right_name, right, node.op, dims) self.errors.append(_declared_rhs_error(context, node, value, rhs_kind)) return node @@ -901,7 +892,7 @@ def _comparison(self, node: UnresolvedComparisonNode) -> WhereNode | UnresolvedW column = self._relation_column(left_name, left_column or None, node.name, node.op) if column is None: return node - dtype = ns.dtypes[ns.shape_of(left_name).dim(column)] + dtype = ns.dtypes[ns.relations[left_name].dim(column)] elif kind in ('parameter', 'dimension'): dtype = ns.dtypes[left_name] if dtype is not None: @@ -918,7 +909,7 @@ def _comparison(self, node: UnresolvedComparisonNode) -> WhereNode | UnresolvedW return DimensionComparisonNode(left_name, node.op, value) case 'relation': assert column is not None - shape = ns.shape_of(left_name) + shape = ns.relations[left_name] return RelationComparisonNode(left_name, column, node.op, value, tuple(shape.dim(k) for k in shape.key)) case 'variable': self.errors.append( @@ -936,7 +927,7 @@ def _relation_column(self, name: str, column: str | None, spelling: str, op: Pre column where there is exactly one. """ ns, context = self.ns, self.context - shape = ns.shape_of(name) + shape = ns.relations[name] if not shape.values: self.errors.append( f"{context}: '{spelling}' compares a column of '{name}', a bare relation — every column is in its " @@ -1116,7 +1107,7 @@ def _relation_pair_error( """ comparison = f"'{node.name} {node.op} {other}'" left_name, right_name = node.name.partition('.')[0], other.partition('.')[0] - ls, rs = ns.shape_of(left_name), ns.shape_of(right_name) + ls, rs = ns.relations[left_name], ns.relations[right_name] left_keys, right_keys = {ls.dim(k) for k in ls.key}, {rs.dim(k) for k in rs.key} if left_keys != right_keys: return ( diff --git a/src/math_spec/separability.py b/src/math_spec/separability.py index e51c4071..66307bbf 100644 --- a/src/math_spec/separability.py +++ b/src/math_spec/separability.py @@ -107,16 +107,16 @@ def waits_on(dimension: str, label: str, name: str, kind: Literal['offset', 'par f'sums over {dimension} — a rolling sum_back(window=n) windows, a total over the horizon does not', ) elif isinstance(node, GroupSum): - for dimension in node.over: + for dimension in node.direction.consumed_dims: report( 'coupled', dimension, label, - f'groups {dimension} into {", ".join(node.into)} — window that dimension instead, or cut only at the group edges', + f'groups {dimension} into {", ".join(node.direction.produced_dims)} — window that dimension instead, or cut only at the group edges', ) elif isinstance(node, At): - for dimension in node.into: - waits_on(dimension, label, node.relation, 'coordinate') + for dimension in node.direction.consumed_dims: + waits_on(dimension, label, node.direction.name, 'coordinate') elif isinstance(node, (Translate, Window)): dimension = node.dimension if node.wrap: diff --git a/src/math_spec/typesetting/walk.py b/src/math_spec/typesetting/walk.py index 6123b74c..d5a6684f 100644 --- a/src/math_spec/typesetting/walk.py +++ b/src/math_spec/typesetting/walk.py @@ -22,13 +22,14 @@ CasesNode, DefinitionNode, DimensionNode, + DirectionNode, DualNode, EdgeNode, FunctionCallNode, KwargNode, NumberNode, ParameterNode, - RelationNode, + PartitionNode, UnaryOperatorNode, UnresolvedNode, VariableNode, @@ -45,7 +46,6 @@ OrNode, ParameterComparisonNode, ParameterDefinedNode, - Partition, PredicateOperator, RelationComparisonNode, RelationDefinedNode, @@ -60,7 +60,6 @@ from collections.abc import Iterable, Mapping from math_spec.model import RelationBlock, SosBlock, _ExpandedSpec - from math_spec.program import RelationDeclaration from math_spec.typesetting.format import Format from math_spec.typesetting.symbols import Symbols @@ -276,53 +275,44 @@ def _translation(self, step: _Step) -> str: self.noticed.grouped = True return self.format.superscript(operator, step.within) - def _relation_read(self, relation: RelationDeclaration, at: Mapping[str, str], read: str) -> str: - """A relation's column *read* as a function at the columns *at* fixes: ``bus(g)``, ``zone_of(g, p)`` or ``ends.bus0(l)``. + def _relation_read(self, name: str, at: Mapping[str, str], read: str) -> str: + """Relation *name*'s column *read* as a function at the columns *at* fixes: ``bus(g)``, ``zone_of(g, p)`` or ``ends.bus0(l)``. *at* maps each key role to the index it is read at. The function is named after the relation alone where the key determines one column, and after the column read otherwise. """ - name = relation.name if len(relation.values) == 1 else f'{relation.name}.{read}' - return self.format.apply(self.format.upright(name), self.format.joined([at[k] for k in relation.key], '')) + lk = self.schema.relations[name] + function = name if len(lk.value_roles) == 1 else f'{name}.{read}' + return self.format.apply(self.format.upright(function), self.format.joined([at[k] for k in lk.key_roles], '')) - def _relation_row(self, name: str, key: list[str]) -> str: - """That relation *name* has a row at *key*, its key columns' indices in declared order. + def _relation_row(self, name: str, at: Mapping[str, str]) -> str: + """That relation *name* has a row at the key *at* fixes. A keyed table is a function, so the claim is that it is defined there: ``gen_zone(g, t) is defined``. A bare one is a set of rows, and every column of it is a key column, so the row is written out: ``(g, b) ∈ connection``. """ - if self.schema.relations[name].value_roles: - applied = self.format.apply(self.format.upright(name), self.format.joined(key, '')) - return f'{applied} {self.format.prose(" is defined")}' - row = self.format.parenthesise(self.format.joined(key, '')) - return f'{row} {self._op("in")} {self.format.upright(name)}' + lk = self.schema.relations[name] + key = self.format.joined([at[k] for k in lk.key_roles], '') + if lk.value_roles: + return f'{self.format.apply(self.format.upright(name), key)} {self.format.prose(" is defined")}' + return f'{self.format.parenthesise(key)} {self._op("in")} {self.format.upright(name)}' + + def _frame_key(self, name: str, ctx: _Context) -> dict[str, str]: + """Relation *name*'s key roles at the frame's own indices of their dimensions.""" + lk = self.schema.relations[name] + return {k: ctx.subscript(dict(lk.pairs)[k]) for k in lk.key_roles} def _value_read(self, name: str, column: str, ctx: _Context) -> str: """A keyed relation's value *column* read at the frame's own indices of its key: ``period_of(t)``.""" - lk = self.schema.relations[name] - keyed = self.format.joined([ctx.subscript(dict(lk.pairs)[k]) for k in lk.key_roles], '') - return self.format.apply(self._column(name, column, len(lk.value_roles) == 1), keyed) - - def _position_group(self, node: DimensionPositionNode, ctx: _Context) -> str: - """The group a grouped position counts within: the relation's group columns read at the row's key.""" - assert node.partition is not None - partition = node.partition - keyed = self.format.joined([ctx.subscript(partition.dim(k)) for k in partition.key], '') - single = len(partition.values) == 1 - reads = [self.format.apply(self._column(partition.name, column, single), keyed) for column in partition.group] - return self._tuple(reads) + return self._relation_read(name, self._frame_key(name, ctx), column) def _tuple(self, reads: list[str]) -> str: """Several reads as one group label: the read alone where there is one, a bracketed tuple otherwise.""" return reads[0] if len(reads) == 1 else self.format.parenthesise(self.format.joined(reads, '')) - def _column(self, name: str, column: str, single: bool) -> str: - """The function a keyed relation's value *column* is: the relation's own name where it has one value column.""" - return self.format.upright(name if single else f'{name}.{column}') - def _context(self, frame: Iterable[str] = ()) -> _Context: return _Context(self, bound=tuple(frame)) @@ -457,25 +447,24 @@ def _call(self, node: FunctionCallNode, ctx: _Context) -> tuple[str, int]: if node.name == 'at': by = node.kwargs['by'] - assert isinstance(by, RelationNode) + assert isinstance(by, DirectionNode) outer = ctx - direction = by.use - assert isinstance(direction, Direction) + direction = by.direction at = {r: outer.subscript(direction.dim(r)) for r in (*direction.produced, *direction.joined)} for read in direction.consumed: - ctx = ctx.pulled_back(direction.dim(read), self._relation_read(direction.relation, at, read)) + ctx = ctx.pulled_back(direction.dim(read), self._relation_read(direction.name, at, read)) return self._arithmetic(node.args[0], ctx) if (by := node.kwargs.get('by')) is not None: - assert isinstance(by, RelationNode) + assert isinstance(by, DirectionNode) + direction = by.direction dummies: dict[str, str] = {} inner = ctx - for d in by.dimensions: + for d in direction.consumed_dims: dummies[d], inner = inner.reducing(d) - assert isinstance(by.use, Direction) - conditions = list(self._grouping(by.use, dummies, ctx)) + conditions = list(self._grouping(direction, dummies, ctx)) domain = ( - f'{self.format.joined([self._membership(d, dummies[d]) for d in by.dimensions], "")} ' + f'{self.format.joined([self._membership(d, dummies[d]) for d in direction.consumed_dims], "")} ' f'{self._op("such_that")} {self.format.joined(conditions, self._op("and"))}' ) elif (consumed := node.kwargs.get('over')) is not None: @@ -504,10 +493,10 @@ def _grouping(self, direction: Direction, dummies: Mapping[str, str], ctx: _Cont **{r: dummies[direction.dim(r)] for r in direction.consumed}, **{r: ctx.subscript(direction.dim(r)) for r in (*direction.joined, *direction.produced)}, } - fixed = [r for r in direction.values if r in at] + fixed = [r for r in self.schema.relations[direction.name].value_roles if r in at] if not fixed: - return [self._relation_row(direction.name, [at[k] for k in direction.key])] - return [f'{self._relation_read(direction.relation, at, r)} {self._op("equal")} {at[r]}' for r in fixed] + return [self._relation_row(direction.name, at)] + return [f'{self._relation_read(direction.name, at, r)} {self._op("equal")} {at[r]}' for r in fixed] def _group(self, by: ArithmeticNode | None, dim: str) -> str: """A ``by=`` as the superscript its translation operator carries. @@ -518,11 +507,10 @@ def _group(self, by: ArithmeticNode | None, dim: str) -> str: """ if by is None: return '' - assert isinstance(by, RelationNode) - partition = by.use - assert isinstance(partition, Partition) + assert isinstance(by, PartitionNode) + partition = by.partition at = {r: self.symbols.index[partition.dim(r)] for r in (partition.along, *partition.joined)} - return self._tuple([self._relation_read(partition.relation, at, r) for r in partition.group]) + return self._tuple([self._relation_read(partition.name, at, r) for r in partition.group]) def _width(self, node: ArithmeticNode) -> str: """``sum_back``'s ``window=``: a number, or a parameter's own symbol. @@ -604,7 +592,11 @@ def _where(self, node: WhereNode, ctx: _Context) -> tuple[str, int]: ) if isinstance(node, DimensionPositionNode): - grouping = None if node.partition is None else self._position_group(node, ctx) + grouping = ( + None + if node.partition is None + else self._tuple([self._value_read(node.partition.name, c, ctx) for c in node.partition.group]) + ) place = self._position(ctx.subscript(node.name), grouping) ordinal = self._ordinal(node.name, node.position, grouping) return f'{place} {self._op(_PREDICATES[node.op])} {ordinal}', comparison @@ -619,8 +611,7 @@ def _where(self, node: WhereNode, ctx: _Context) -> tuple[str, int]: return f'{left} {self._op(_PREDICATES[node.op])} {right}', comparison if isinstance(node, RelationDefinedNode): - lk = self.schema.relations[node.name] - return self._relation_row(node.name, [ctx.subscript(dict(lk.pairs)[k]) for k in lk.key_roles]), comparison + return self._relation_row(node.name, self._frame_key(node.name, ctx)), comparison if isinstance(node, NotNode): return ( diff --git a/src/math_spec/validation.py b/src/math_spec/validation.py index c0fa5453..e292840e 100644 --- a/src/math_spec/validation.py +++ b/src/math_spec/validation.py @@ -165,7 +165,7 @@ def validate_expressions(schema: Spec) -> Resolved: if errors: raise SchemaError(_once(errors)) - resolved = Resolved(expressions, variables, constraints, objective) + resolved = Resolved(expressions, variables, constraints, objective, ns.relations) check_schema(schema, resolved) return resolved diff --git a/tests/test_dimensions.py b/tests/test_dimensions.py index e4dd5130..5258042b 100644 --- a/tests/test_dimensions.py +++ b/tests/test_dimensions.py @@ -290,6 +290,35 @@ def test_an_ill_dimensioned_expression_is_rejected(expr, match): _dims(expr) +@pytest.mark.parametrize( + ('expr', 'diag'), + [ + pytest.param( + 'sum(p, by=diag, over=k, into=z)', + {'key': {'k': 'generator', 'j': 'generator', 'z': 'zone'}}, + id='a-sum-consuming-a-column-over-the-dimension-it-joins-on', + ), + pytest.param( + 'at(load, by=diag, over=rep, into=generator)', + {'key': ['snapshot', 'generator'], 'values': {'rep': 'snapshot'}}, + id='a-read-consuming-a-column-over-the-dimension-it-joins-on', + ), + ], +) +def test_a_joined_column_is_not_also_consumed(expr, diag): + """The operand carries one coordinate per dimension, so a column consumed and a column joined on cannot share one. + + The `at` case passed: the check asked whether a joined dimension was + *produced*, which the landing check already refuses, and not whether it + was consumed. `at(load, by=diag, over=rep, into=generator)` then read + `rep` at the operand's snapshot and joined on the key's snapshot at the + same coordinate, and landed on `[bus, generator]` with the joined + dimension gone. + """ + with pytest.raises(DimensionError, match=r"joins 'diag' on \[.*\] through more than one column"): + _dims_with(expr, **{'relations.diag': diag}) + + def test_an_outer_product_is_legal_and_carries_both_dim_sets(): """Binary ops union. Requiring subset instead would reject the convex piecewise epigraph, which multiplies a per-segment slope by a per-snapshot diff --git a/tests/test_lowering.py b/tests/test_lowering.py index 31e5b47b..6420f866 100644 --- a/tests/test_lowering.py +++ b/tests/test_lowering.py @@ -553,12 +553,14 @@ def test_a_relation_lowers_with_the_direction_each_call_names(): 'a grouped sum names the column it consumes, the one it produces and the one it joins on' ) assert isinstance(zonal, GroupSum) - assert (zonal.over, zonal.into, zonal.joined, zonal.relation) == ( + assert (zonal.direction.consumed_dims, zonal.direction.produced_dims, zonal.direction.joined_dims) == ( ('generator',), ('zone',), ('snapshot',), - 'zone_of', ), 'the dims a consumer reads are read off the direction' + assert zonal.direction.relation is program.relations['zone_of'], ( + 'the direction holds the one declaration the program holds, not an equal copy built again' + ) assert program.constraints['history'].lhs == GroupSum( Variable('p'), direction=Direction(declared, ('snapshot',), ('zone',), ('generator',)) ), 'the same table read from its other key column' @@ -567,9 +569,11 @@ def test_a_relation_lowers_with_the_direction_each_call_names(): 'and its adjoint consumes the value column and produces the key column' ) assert isinstance(priced, At) - assert (priced.over, priced.into, priced.joined) == (('generator',), ('zone',), ('snapshot',)), ( - 'an at produces the fine dims, consumes the coarse, and joins on the rest of the key' - ) + assert (priced.direction.consumed_dims, priced.direction.produced_dims, priced.direction.joined_dims) == ( + ('zone',), + ('generator',), + ('snapshot',), + ), 'an at consumes the coarse dims, produces the fine, and joins on the rest of the key' p_where = program.variable('p').where assert p_where is not None assert [(type(a).__name__, a.dims) for a in p_where.atoms] == [ diff --git a/tests/test_parser.py b/tests/test_parser.py index 87fb4d02..e4cb1e77 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -23,6 +23,7 @@ ComparisonNode, DefinitionNode, DimensionNode, + DirectionNode, DualNode, EdgeNode, FunctionCallNode, @@ -30,7 +31,7 @@ NameNode, NumberNode, ParameterNode, - RelationNode, + PartitionNode, UnaryOperatorNode, VariableNode, parse_expression, @@ -48,6 +49,7 @@ Direction, NotNode, OrNode, + Partition, RelationDeclaration, _conjuncts, ) @@ -502,6 +504,9 @@ def test_a_node_prints_as_the_file_writes_it(text, printed): assert str(parse_expression(text)) == printed, 'the spelling is the one a file could be written with' +_ZONE_OF = RelationDeclaration('zone_of', (('u', 'unit'), ('zone', 'zone')), ('u',)) + + @pytest.mark.parametrize( ('node', 'printed'), [ @@ -510,9 +515,14 @@ def test_a_node_prints_as_the_file_writes_it(text, printed): pytest.param(DimensionNode('t'), 't', id='a-dimension'), pytest.param(DualNode('budget'), 'dual(budget)', id='a-dual'), pytest.param( - RelationNode('zone_of', ('u',), ('zone',), Direction(RelationDeclaration('zone_of', ()), (), (), ())), + DirectionNode(Direction(_ZONE_OF, ('u',), ('zone',), ())), + 'zone_of', + id='a-relation-read-in-a-direction', + ), + pytest.param( + PartitionNode(Partition(_ZONE_OF, 'u', ('zone',), ())), 'zone_of', - id='a-relation', + id='a-relation-stepped-along-as-a-partition', ), pytest.param(EdgeNode(), "'wrap'", id='a-resolved-edge'), pytest.param(DefinitionNode('headroom', NameNode('p')), 'headroom', id='a-named-expression-prints-its-name'), diff --git a/tests/test_validation.py b/tests/test_validation.py index 2bb3c379..ee4ce219 100644 --- a/tests/test_validation.py +++ b/tests/test_validation.py @@ -535,7 +535,10 @@ def test_it_resolves(self, mask: str, position: int, by: str | None): [ ('position(load) == 0', ["counts along a dimension's coordinates", "'load' is a parameter"]), ('position(nope) == 0', ["'nope' is not declared"]), - ('position(snapshot, by=load) == 0', ['groups by', '``by=`` takes a relation']), + ( + 'position(snapshot, by=load) == 0', + ['position(by=load) does not name a relation', "Declare it under 'relations:'"], + ), ( 'position(snapshot, by=starts_at, within=snapshot) == 0', ["no key column over 'snapshot'", "its key is ['period']"], @@ -787,7 +790,10 @@ class TestRulesDecidedWithoutData: 'relations.rel': {'key': ['g', 'h']}, 'objective': {'expression': 'sum(at(r, by=rel, over=h, into=g))'}, }, - ("at reads one value per coordinate, and 'rel' is not single-valued",), + ( + "at reads one value per coordinate, and 'rel' is not single-valued in ['h'] at the columns " + "the call lands on (['g'])", + ), id='at-through-a-bare-relation', ), pytest.param(