File size: 11,354 Bytes
d06fa57
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
"""Safe parser for the SRHarness symbolic model language."""

from __future__ import annotations

import ast
from collections.abc import Mapping
from typing import Any

from .analysis import fold_constants
from .desugar import desugar
from .expression import (
    Aggregate,
    Binary,
    Expression,
    Function,
    Gather,
    GroupedParameter,
    Index,
    Indexed,
    Number,
    Parameter,
    Reduction,
    RelationLift,
    Symbol,
    Unary,
)


FUNCTIONS = {
    "abs", "arccos", "arcsin", "arctan", "cos", "cosh", "cot", "csc", "exp", "inv",
    "log", "log10", "log2", "sec", "sech", "sigmoid", "sign", "sin", "sinh", "sqrt", "tan",
    "tanh", "pow2", "pow3",
}
FUNCTION_ALIASES = {
    "acos": "arccos",
    "asin": "arcsin",
    "atan": "arctan",
    "lg": "log10",
    "ln": "log",
}
BINARY_FUNCTIONS = {"max", "min"}


class ExpressionParser(ast.NodeVisitor):
    """Convert a restricted Python expression AST into engine nodes."""

    def __init__(self, symbols: Mapping[str, Any] | None = None):
        self.symbols = {}
        for name, value in (symbols or {}).items():
            if isinstance(value, Expression):
                self.symbols[name] = value
            elif isinstance(value, bool) or not isinstance(value, (int, float)):
                raise TypeError(f"Symbol override {name!r} must be an expression or number.")
            else:
                self.symbols[name] = Number(value)

    def parse(self, source: str) -> Expression:
        """Parse input into the canonical representation.

        Args:
            source: Source text to parse.

        Returns:
            The parsed, desugared, and constant-folded expression.
        """
        try:
            tree = ast.parse(source, mode="eval")
        except SyntaxError as error:
            raise SyntaxError(f"Invalid symbolic expression: {error.msg}.") from error
        return fold_constants(desugar(self.visit(tree.body)))

    def generic_visit(self, node: ast.AST):
        """Reject syntax that is outside the symbolic language.

        Args:
            node: Expression or syntax-tree node.
        """
        raise ValueError(f"Unsupported syntax: {type(node).__name__}.")

    def visit_Constant(self, node: ast.Constant) -> Expression:
        """Convert a numeric literal into a number node.

        Args:
            node: Expression or syntax-tree node.

        Returns:
            A numeric expression node.
        """
        if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
            raise ValueError(f"Only numerical literals are expressions, got {node.value!r}.")
        return Number(node.value)

    def visit_Name(self, node: ast.Name) -> Expression:
        """Resolve a name to a supplied or newly created symbol.

        Args:
            node: Expression or syntax-tree node.

        Returns:
            The expression associated with the name.
        """
        return self.symbols.setdefault(node.id, Symbol(node.id))

    def visit_UnaryOp(self, node: ast.UnaryOp) -> Expression:
        """Convert a supported unary operator.

        Args:
            node: Expression or syntax-tree node.

        Returns:
            The converted operand or unary expression.
        """
        if isinstance(node.op, ast.USub):
            return Unary("-", self.visit(node.operand))
        if isinstance(node.op, ast.UAdd):
            return self.visit(node.operand)
        raise ValueError(f"Unsupported unary operator: {type(node.op).__name__}.")

    def visit_BinOp(self, node: ast.BinOp) -> Expression:
        """Convert a supported binary arithmetic operator.

        Args:
            node: Expression or syntax-tree node.

        Returns:
            A binary expression node.
        """
        operators = {
            ast.Add: "+", ast.Sub: "-", ast.Mult: "*", ast.Div: "/", ast.Pow: "**"
        }
        operator = operators.get(type(node.op))
        if operator is None:
            raise ValueError(f"Unsupported binary operator: {type(node.op).__name__}.")
        return Binary(operator, self.visit(node.left), self.visit(node.right))

    def visit_Subscript(self, node: ast.Subscript) -> Expression:
        """Convert symbolic indexing such as ``x[i]``.

        Args:
            node: Expression or syntax-tree node.

        Returns:
            An indexed expression node.
        """
        base = self.visit(node.value)
        return Indexed(base, self._indices(node.slice))

    def visit_Call(self, node: ast.Call) -> Expression:
        """Convert an approved symbolic function or reduction call.

        Args:
            node: Expression or syntax-tree node.

        Returns:
            The corresponding symbolic expression.
        """
        if isinstance(node.func, ast.Subscript) and self._name(node.func.value) == "sum":
            if node.keywords or len(node.args) not in {1, 2}:
                raise ValueError(
                    "sum[index](...) expects an operand, optionally preceded by a relation."
                )
            arguments = [self.visit(argument) for argument in node.args]
            relation, operand = (None, arguments[0]) if len(arguments) == 1 else arguments
            return Reduction(self._indices(node.func.slice), operand, relation)

        source_name = self._name(node.func)
        name = FUNCTION_ALIASES.get(source_name, source_name)
        if name in FUNCTIONS:
            if node.keywords or len(node.args) != 1:
                raise ValueError(f"{source_name}(...) expects exactly one positional argument.")
            return Function(name, (self.visit(node.args[0]),))
        if name in BINARY_FUNCTIONS:
            if node.keywords or len(node.args) != 2:
                raise ValueError(f"{name}(...) expects exactly two positional arguments.")
            return Function(name, tuple(self.visit(argument) for argument in node.args))
        if name == "delay":
            if node.keywords or len(node.args) != 2:
                raise ValueError("delay(value, delta) expects exactly two arguments.")
            return Function(name, tuple(self.visit(argument) for argument in node.args))
        if name == "gather":
            if node.keywords or len(node.args) != 2:
                raise ValueError("gather(relation, expression) expects exactly two arguments.")
            return Gather(self.visit(node.args[0]), self.visit(node.args[1]))
        if name == "Number":
            if node.keywords or len(node.args) != 1:
                raise ValueError("Number(...) expects exactly one numerical literal.")
            return self.visit(node.args[0])
        if name == "param":
            return self._parameter(node)
        if name == "grouped_param":
            return self._grouped_parameter(node)
        if name == "aggr":
            if node.keywords or len(node.args) not in {1, 2}:
                raise ValueError(
                    "aggr expects aggr(relation * expression) or "
                    "aggr(relation, expression)."
                )
            if len(node.args) == 2:
                return Aggregate(self.visit(node.args[0]), self.visit(node.args[1]))
            relation, operand = _split_aggregation(self.visit(node.args[0]))
            return Aggregate(relation, operand)
        if name in {"targ", "sour"}:
            if node.keywords or len(node.args) not in {1, 2}:
                raise ValueError(f"{name} expects one value, optionally preceded by a relation.")
            arguments = [self.visit(argument) for argument in node.args]
            relation, operand = (None, arguments[0]) if len(arguments) == 1 else arguments
            return RelationLift("target" if name == "targ" else "source", operand, relation)
        raise ValueError(f"Unknown symbolic function: {name!r}.")

    def _parameter(self, node: ast.Call) -> Parameter:
        if len(node.args) != 1 or not isinstance(node.args[0], ast.Constant):
            raise ValueError("param(...) requires a literal string name.")
        name = node.args[0].value
        if not isinstance(name, str) or not name:
            raise ValueError("Parameter names must be non-empty strings.")
        keywords = self._literal_keywords(node, {"value"})
        value = keywords.get("value")
        if value is not None and (isinstance(value, bool) or not isinstance(value, (int, float))):
            raise ValueError("A parameter initial value must be numerical.")
        return Parameter(name, value)

    def _grouped_parameter(self, node: ast.Call) -> GroupedParameter:
        if len(node.args) != 1:
            raise ValueError("grouped_param(...) requires one grouping expression.")
        keywords = self._literal_keywords(node, {"name", "value", "default"})
        name = keywords.get("name")
        if name is not None and (not isinstance(name, str) or not name):
            raise ValueError("A grouped parameter name must be a non-empty string.")
        value = keywords.get("value")
        if value is not None and not isinstance(value, dict):
            raise ValueError("grouped_param value must be a category-to-number dictionary.")
        default = keywords.get("default")
        return GroupedParameter(self.visit(node.args[0]), name=name, value=value, default=default)

    @staticmethod
    def _name(node: ast.AST) -> str:
        if not isinstance(node, ast.Name):
            raise ValueError("Only direct symbolic function calls are allowed.")
        return node.id

    @staticmethod
    def _indices(node: ast.AST) -> tuple[Index, ...]:
        nodes = node.elts if isinstance(node, ast.Tuple) else (node,)
        if not nodes or not all(isinstance(item, ast.Name) for item in nodes):
            raise ValueError("Indices must be identifiers such as i, j, or k.")
        return tuple(Index(item.id) for item in nodes)

    @staticmethod
    def _literal_keywords(node: ast.Call, allowed: set[str]) -> dict[str, Any]:
        result = {}
        for keyword in node.keywords:
            if keyword.arg not in allowed:
                raise ValueError(f"Unexpected keyword argument: {keyword.arg!r}.")
            result[keyword.arg] = ast.literal_eval(keyword.value)
        return result


def parse(
    source: str,
    symbols: Mapping[str, Any] | None = None,
) -> Expression:
    """Parse *source* without using ``eval`` or executing user code.

    Args:
        source: Source text to parse.
        symbols: Optional predefined symbols or numeric constants.

    Returns:
        The parsed canonical expression.
    """
    return ExpressionParser(symbols).parse(source)


def _split_aggregation(expression: Expression) -> tuple[Expression, Expression]:
    factors = _multiplication_factors(expression)
    if len(factors) < 2 or not isinstance(factors[0], Symbol):
        raise ValueError(
            "aggr(relation * expression) requires the relation symbol as its first factor."
        )
    operand = factors[1]
    for factor in factors[2:]:
        operand = Binary("*", operand, factor)
    return factors[0], operand


def _multiplication_factors(expression: Expression) -> list[Expression]:
    if isinstance(expression, Binary) and expression.operator == "*":
        return _multiplication_factors(expression.left) + _multiplication_factors(expression.right)
    return [expression]