computation_model.py

"""Renderer-independent computation store and typed information flow."""
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

import numpy as np

from math_expression import Expr, operation_expression, snapshot
from backward_rules import ScalarLocalRule, TensorBackwardRule
from backward_explanation import BackwardExplanation, explain_step, explain_stop
from scalar_rules import scalar_rule_from_record


class BackwardEdgeType(Enum):
    ADJOINT_INTRODUCTION = 'adjoint_introduction'
    ADJOINT_PROPAGATION = 'adjoint_propagation'
    CONTRIBUTION = 'contribution'


@dataclass
class QuantityState:
    uid: str
    name: str
    primal_value: Any
    adjoint_value: Any
    differentiable: bool
    is_input: bool
    shape: tuple
    is_constant: bool = False


@dataclass
class Operation:
    quantity: QuantityState
    operands: list[QuantityState]
    kind: str
    function: str | None = None
    exponent: float | None = None
    index: int | None = None
    formula: Expr = field(init=False)
    stop_explanation: BackwardExplanation | None = None

    def __post_init__(self):
        self.formula = operation_expression(self)


@dataclass
class BackwardStep:
    uid: str
    source_quantity: QuantityState
    destination_quantity: QuantityState
    incoming_adjoint: Any
    operation: Operation
    rule: ScalarLocalRule | TensorBackwardRule
    local_value: Any
    contribution: Any
    operand_occurrence: int
    explanation: BackwardExplanation | None = None


@dataclass
class Accumulation:
    uid: str
    destination_quantity: QuantityState
    contributions: list[BackwardStep]
    resulting_adjoint: Any


@dataclass
class BackwardEdge:
    source: str
    destination: str
    kind: BackwardEdgeType
    value: Any
    quantity: QuantityState | None = None
    # A single contribution is allowed to collapse into an introduction edge.
    contribution_step: BackwardStep | None = None
    seed: bool = False


@dataclass
class OutputState:
    uid: str
    quantity: QuantityState


@dataclass
class ForwardEdge:
    source: str
    destination: str
    value: Any
    operand_occurrence: int | None = None


@dataclass
class ComputationModel:
    quantities: list[QuantityState]
    operations: list[Operation]
    root: QuantityState
    output_name: str
    formula: str
    forward_only: bool
    output: OutputState | None = None
    steps: list[BackwardStep] = field(default_factory=list)
    accumulations: list[Accumulation] = field(default_factory=list)
    forward_edges: list[ForwardEdge] = field(default_factory=list)
    backward_edges: list[BackwardEdge] = field(default_factory=list)


def build_computation_model(graph, engine, output_name='out', forward_only=False):
    """Adapt numerical engine records once; formatting does not infer edge roles."""
    if engine not in {'scalar', 'tensor'}:
        raise ValueError('Unknown numerical engine')
    quantities, operations, steps = [], [], []
    by_id = {}
    for node in graph.nodes:
        differentiable = not node.is_constant if node.is_input else any(by_id[a.uid].differentiable for a in node.args)
        q = QuantityState(node.uid, node.name, snapshot(node.val),
                          None if forward_only or not differentiable or not node.live else snapshot(node.grad),
                          differentiable, node.is_input, np.shape(node.val),
                          is_constant=getattr(node, 'is_constant', False))
        by_id[q.uid] = q
        quantities.append(q)
        if node.is_input:
            continue
        if engine == 'scalar':
            names = {'Add': 'add', 'Mul': 'mul', 'Pow': 'pow'}
            func = node.expr.func.__name__
            operation = Operation(q, [by_id[a.uid] for a in node.args], names.get(func, 'func'), func)
        else:
            operation = Operation(q, [by_id[a.uid] for a in node.args], node.op_type,
                                  node.func_name, node.pow_exp, node.pick_index)
        operations.append(operation)
        is_nograd = func == 'nograd' if engine == 'scalar' else node.func_name == 'nograd'
        if is_nograd and not forward_only and q.differentiable:
            operation.stop_explanation = explain_stop(operation)
        if forward_only or not getattr(node, 'live', True) or is_nograd:
            continue
        records = node.local_derivatives if engine == 'scalar' else node.backward_rules
        for record in records:
            occurrence = record['operand_occurrence']
            destination = by_id[record['arg_node'].uid]
            if not destination.differentiable:
                continue
            rule = (scalar_rule_from_record(operation, occurrence, record)
                    if engine == 'scalar' else record['rule'])
            local_value = record['local_val'] if engine == 'scalar' else record.get('local_val')
            step = BackwardStep(
                f'der_{node.uid}_{occurrence}', q, destination, snapshot(node.grad),
                operation, rule, snapshot(local_value),
                snapshot(record['grad_contrib'] if engine == 'scalar' else record['contrib_val']), occurrence,
            )
            values = {a.uid: a.primal_value for a in operation.operands}
            values[q.uid] = q.primal_value
            values['adjoint:' + q.uid] = step.incoming_adjoint
            step.explanation = explain_step(step, values)
            steps.append(step)
    root_q = by_id[graph.root.uid]
    output_state = OutputState(f'output_{root_q.uid}', root_q)
    model = ComputationModel(quantities, operations, root_q, output_name,
                         str(graph.orig_expr) if engine == 'scalar' else graph.orig_expr_str,
                         forward_only, output_state, steps)
    for operation in operations:
        for i, operand in enumerate(operation.operands):
            model.forward_edges.append(ForwardEdge(operand.uid, operation.quantity.uid, operand.primal_value, i))
    # The final operation/result sends its primal result forward to the output presentation node
    model.forward_edges.append(ForwardEdge(model.root.uid, model.output.uid, model.root.primal_value, None))
    if forward_only or not model.root.differentiable:
        return model
    model.backward_edges.append(BackwardEdge(model.output.uid, model.root.uid,
        BackwardEdgeType.ADJOINT_INTRODUCTION, model.root.adjoint_value, model.root, seed=True))
    contributions = {q.uid: [] for q in quantities}
    for step in steps:
        contributions[step.destination_quantity.uid].append(step)
        model.backward_edges.append(BackwardEdge(step.source_quantity.uid, step.uid,
            BackwardEdgeType.ADJOINT_PROPAGATION, step.incoming_adjoint, step.source_quantity))
    for quantity in quantities:
        incoming = contributions[quantity.uid]
        if len(incoming) > 1:
            accumulation = Accumulation('sum_' + quantity.uid, quantity, incoming, quantity.adjoint_value)
            model.accumulations.append(accumulation)
            for step in incoming:
                model.backward_edges.append(BackwardEdge(step.uid, accumulation.uid,
                    BackwardEdgeType.CONTRIBUTION, step.contribution, contribution_step=step))
            model.backward_edges.append(BackwardEdge(accumulation.uid, quantity.uid,
                BackwardEdgeType.ADJOINT_INTRODUCTION, quantity.adjoint_value, quantity))
        elif incoming:
            step = incoming[0]
            model.backward_edges.append(BackwardEdge(step.uid, quantity.uid,
                BackwardEdgeType.ADJOINT_INTRODUCTION, quantity.adjoint_value, quantity, step))
    return model


def tensor_operation(node):
    """Numerical tensor-node adapter; rule factories depend only on this record."""
    def quantity(n):
        return QuantityState(n.uid, n.name, n.val, n.grad, not n.is_constant,
                             n.is_input, n.shape, n.is_constant)
    return Operation(quantity(node), [quantity(a) for a in node.args],
                     node.op_type, node.func_name, node.pow_exp, node.pick_index)