test_math_renderer.py

"""Mathematical presentation checks against explicit expected equations."""
import numpy as np
import pytest
import sympy as sp

from autodiff_visualizer import ComputationGraph, safe_parse_expr
from tensor_engine import TensorComputationGraph
from math_expression import Expr, adjoint, literal, node
from math_renderer import MathRenderer
from backward_explanation import adjoint_equation


def graph(expr, feed):
    g = TensorComputationGraph().build_from_ast(expr)
    g.evaluate(feed)
    return g


def test_extrema_rule_shares_ties_for_vectors_and_matrices():
    for value in [np.array([1., 3., 3.]), np.array([[1., 3.], [3., 0.]])]:
        step = graph('max(v)', {'v': value}).to_computation_model().steps[0]
        upper, lower = MathRenderer('code').explanation(step.explanation)
        assert upper == 'c_out_to_v = out_bar * (mask(v == out) / sum(mask(v == out)))'
        assert 'Equal-share subgradient at ties.' in step.explanation.notes
        np.testing.assert_array_equal(step.contribution, (value == 3) / 2)
        assert MathRenderer('code').expression(literal(step.contribution)) in lower


def test_broadcast_rule_sums_contribution_and_keeps_scalar_unindexed():
    model = graph('sum(v / b)', {'v': np.array([2., 4.]), 'b': 2.}).to_computation_model()
    step = next(s for s in model.steps if s.destination_quantity.name == 'b')
    assert step.rule.application.kind == 'sum'
    upper, _ = MathRenderer().explanation(step.explanation)
    assert upper == 'c_n1→b = ∑(n̄1 ⊙ ((−v) / b^2))'
    assert 'Sum over broadcast entries.' in step.explanation.notes


def test_matrix_products_and_repeated_operands():
    model = graph('sum(W @ x)', {'W': np.ones((2, 3)), 'x': np.ones(3)}).to_computation_model()
    W, x = model.steps[:2]
    assert MathRenderer('code').explanation(W.explanation)[0] == 'c_n1_to_W = outer(n1_bar, x)'
    assert MathRenderer('code').explanation(x.explanation)[0] == 'c_n1_to_x = W.T @ n1_bar'
    assert MathRenderer(html=True).explanation(W.explanation)[0].endswith('x<sup>T</sup></span>')
    repeated = graph('sum(x / x)', {'x': np.array([2., 3.])}).to_computation_model()
    upper = [MathRenderer('code').explanation(s.explanation)[0] for s in repeated.steps[:2]]
    assert upper == ['c_n1_to_x_1 = n1_bar * (1 / x)', 'c_n1_to_x_2 = n1_bar * ((-x) / x**2)']


def test_zero_power_rule_does_not_display_undefined_reciprocal():
    model = graph('sum(v**0)', {'v': np.array([0., 2.])}).to_computation_model()
    upper, lower = MathRenderer('code').explanation(model.steps[0].explanation)
    assert 'zeros_like(v)' in upper
    assert '**(-1)' not in upper
    assert lower.endswith(' = [0, 0]')


def test_scalar_backward_result_and_simplified_power_derivative():
    cg = ComputationGraph().build(safe_parse_expr('x**2 * y'))
    cg.evaluate({sp.Symbol('x'): 2., sp.Symbol('y'): 2.})
    model = cg.to_computation_model()
    step = model.steps[0]
    upper, lower = MathRenderer('code').explanation(step.explanation)
    assert upper == 'diff(x**2, x) = 2 * x = 4'
    assert lower == 'c_n1_to_x = 2 * 4 = 8'
    assert 'x_bar = diff(out, x) = 8' in cg.to_mermaid()


@pytest.mark.parametrize('notation', ['code', 'mathbook'])
def test_adjoint_and_contribution_symbols(notation):
    q = graph('x*x', {'x': 3.}).to_computation_model().root
    renderer = MathRenderer(notation, html=True)
    expected = 'out_bar' if notation == 'code' else "<span style='text-decoration:overline'>out</span>"
    assert renderer.expression(adjoint(q)) == expected
    derivative = 'diff(out, out)' if notation == 'code' else '∂out/∂out'
    assert renderer.equation(adjoint_equation(q, 'out', 1)) == f'{expected} = {derivative} = 1'


def test_html_escapes_symbol_text_and_preserves_structure():
    expression = node('pow', Expr('ref', key='test', name='<x&y>'), literal(2))
    assert MathRenderer(html=True).expression(expression) == '&lt;x&amp;y&gt;<sup>2</sup>'
    assert MathRenderer().expression(expression) == '<x&y>^2'


def test_derivation_semantic_wrapping_invariants():
    # 1. Explicit inline layout produces a single line with '=' separation
    cg_scalar = ComputationGraph().build(safe_parse_expr('x * x'))
    cg_scalar.evaluate({sp.Symbol('x'): 3.})
    scalar_step = cg_scalar.to_computation_model().steps[0]
    scalar_inline = MathRenderer('code', html=True).application_derivation(scalar_step.explanation, layout='inline')
    assert 'class="derivation"' not in scalar_inline
    assert 'class="derivation-line"' not in scalar_inline
    assert scalar_inline == '<span class="math-stage" style="white-space:nowrap">c_out_to_x_1 = 1 * 3 = 3</span>'

    # Short scalar application + result stays inline under auto layout
    scalar_auto = MathRenderer('code', html=True).application_derivation(scalar_step.explanation, layout='auto')
    assert scalar_auto == scalar_inline

    # Explicit stacked layout breaks only between the two stages
    scalar_stacked = MathRenderer('code', html=True).application_derivation(scalar_step.explanation, layout='stacked')
    assert '<div class="derivation-line" style="white-space:nowrap">c_out_to_x_1 = 1 * 3</div>' in scalar_stacked
    assert '<div class="derivation-line derivation-continuation" style="white-space:nowrap">= 3</div>' in scalar_stacked

    # 2. Long tensor expression wraps into semantic derivation stages
    tg_tensor = TensorComputationGraph().build_from_ast('(softmax((K @ q) / 2**0.5) @ V)[0]')
    tg_tensor.evaluate({
        'K': np.array([[1., 0.], [0., 1.], [1., 1.]]),
        'q': np.array([0.5, 1.]),
        'V': np.array([[1., 0.], [0., 2.], [1., 1.]])
    })
    long_step = next(s for s in tg_tensor.to_computation_model().steps if len(s.explanation.application_derivation.stages) >= 2)
    long_rendered = MathRenderer('code', html=True).application_derivation(long_step.explanation, layout='auto')
    assert '<div class="derivation">' in long_rendered
    assert '<div class="derivation-line derivation-continuation" style="white-space:nowrap">= ' in long_rendered
    # Continuation lines must begin with '='
    assert '>=' in long_rendered

    # Negative assertions
    assert 'diff(out,' not in long_rendered
    assert 'δ' not in long_rendered
    # No empty continuation lines
    assert '>= </div>' not in long_rendered
    assert '>=</' not in long_rendered