compact_graph.py

"""Compact presentation projection; all numerical values come from the store."""
from dataclasses import dataclass
from math_expression import Expr, ref


@dataclass(frozen=True)
class CompactNode:
    uid: str
    formula: Expr | None
    value: object
    adjoint: object
    role: str
    name: str


@dataclass(frozen=True)
class CompactEdge:
    source: str
    destination: str
    value: object
    role: str
    occurrence: int | None = None


@dataclass(frozen=True)
class CompactGraph:
    nodes: tuple[CompactNode, ...]
    edges: tuple[CompactEdge, ...]


def compact_graph(model, backward=True):
    operations = {op.quantity.uid: op for op in model.operations}
    nodes = []
    for q in model.quantities:
        op = operations.get(q.uid)
        nodes.append(CompactNode(q.uid, op.formula if op else ref(q), q.primal_value,
                                 q.adjoint_value, 'operation' if op else ('constant' if q.is_constant else 'input'), q.name))
    if model.output:
        nodes.append(CompactNode(model.output.uid, None, None, None, 'output', ''))
    edges = [CompactEdge(e.source, e.destination, e.value, 'forward', e.operand_occurrence)
             for e in model.forward_edges]
    if backward and not model.forward_only:
        edges.extend(CompactEdge(s.source_quantity.uid, s.destination_quantity.uid,
                                 s.contribution, 'backward', s.operand_occurrence) for s in model.steps)
        edges.extend(CompactEdge(e.source, e.destination, e.value, 'seed')
                     for e in model.backward_edges if e.seed)
    return CompactGraph(tuple(nodes), tuple(edges))