File size: 4,403 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
"""Expression simplification and structural analysis."""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

import numpy as np

from .evaluation import grouped_parameter_key
from .expression import Binary, Expression, Function, GroupedParameter, Number, Parameter, Unary
from .tree import iter_preorder, transform


def fold_constants(expression: Expression) -> Expression:
    """Evaluate closed numerical subexpressions without reordering terms.

    Args:
        expression: Symbolic expression to process.

    Returns:
        A simplified expression with constant-only branches evaluated.
    """

    def fold(node: Expression) -> Expression:
        if isinstance(node, Unary) and isinstance(node.operand, Number):
            return Number(-node.operand.value)
        if (
            isinstance(node, Binary)
            and isinstance(node.left, Number)
            and isinstance(node.right, Number)
        ):
            operations = {
                "+": lambda left, right: left + right,
                "-": lambda left, right: left - right,
                "*": lambda left, right: left * right,
            }
            if node.operator not in operations:
                return node
            try:
                value = operations[node.operator](node.left.value, node.right.value)
            except (ArithmeticError, OverflowError, ValueError):
                return node
            return Number(value) if _finite_scalar(value) else node
        return node

    return transform(expression, fold)


def count_parameters(
    expression: Expression,
    values: Mapping[str, Any] | None = None,
    *,
    parameters: Mapping[str, Any] | None = None,
) -> int:
    """Count independent fitted values represented by an expression.

    Repeated named parameters count once. A grouped parameter counts once per
    known category. When categories are unavailable, it counts as one
    unresolved parameter family.

    Args:
        expression: Symbolic expression to process.
        values: Values keyed by symbol name.
        parameters: Fitted parameter values keyed by parameter name.

    Returns:
        The number of independent scalar parameter values.
    """
    values = dict(values or {})
    parameters = dict(parameters or {})
    named: set[str] = set()
    grouped: dict[str, GroupedParameter] = {}
    for node in iter_preorder(expression):
        if isinstance(node, Parameter):
            named.add(node.name)
        elif isinstance(node, GroupedParameter):
            grouped.setdefault(grouped_parameter_key(node), node)

    count = len(named)
    for key, node in grouped.items():
        if isinstance(parameters.get(key), Mapping):
            count += len(parameters[key])
            continue
        try:
            labels = np.asarray(node.by.evaluate(values), dtype=object).reshape(-1)
        except (KeyError, ValueError):
            count += len(node.value) if isinstance(node.value, Mapping) else 1
        else:
            count += len({_hashable_label(label) for label in labels})
    return count


def parameter_values(expression: Expression) -> dict[str, Any]:
    """Collect parameter values already bound into an expression."""
    values: dict[str, Any] = {}
    for node in iter_preorder(expression):
        if isinstance(node, Parameter) and node.value is not None:
            values[node.name] = float(node.value)
        elif isinstance(node, GroupedParameter) and node.value is not None:
            values[grouped_parameter_key(node)] = dict(node.value)
    return values


def unbound_parameters(expression: Expression) -> list[str]:
    """Return stable names for parameter nodes that have no fitted value."""
    names = []
    for node in iter_preorder(expression):
        if isinstance(node, Parameter) and node.value is None:
            names.append(node.name)
        elif isinstance(node, GroupedParameter) and node.value is None:
            names.append(grouped_parameter_key(node))
    return list(dict.fromkeys(names))


def _finite_scalar(value: Any) -> bool:
    array = np.asarray(value)
    return array.ndim == 0 and bool(np.isfinite(array))


def _hashable_label(value: Any) -> Any:
    value = value.item() if isinstance(value, np.generic) else value
    try:
        hash(value)
    except TypeError:
        return repr(value)
    return value