arborist/bench/scripts/demo_plot.py

"""End-to-end demo for ``function-sampled@v1`` (#000030 Phase 7).

Closes the loop on opencompletion's ``activity24-math-plot.yaml``:
SymPy expression → quantized integer-vector signature (canonical
bytes) → optional matplotlib PNG (a downstream view of the
canonical evidence).

The canonical bytes are the proof: two function-sampled outputs
that match byte-for-byte represent the same function on the same
quantization grid. The PNG is just a human-readable rendering of
those same bytes — different DPIs / colormaps / axis spans don't
change identity.

Usage::

    python -m bench.scripts.demo_plot \\
        --expr 'sin(x)' --x-min 0 --x-max 6.283 --n-samples 200 --dv 0.01 \\
        --png /tmp/sin.png

    # Or via Make (PNG optional):
    make demo-plot Q='sin(x)'
    make demo-plot Q='x**2 + 2*x + 1' PNG=/tmp/parabola.png

The matplotlib import is gated — the canonical bytes always print
to stdout; the PNG is written only when ``--png`` is provided AND
matplotlib is importable.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import sys
from pathlib import Path


def _build_payload(
    expr: str, x: str, x_min: float, x_max: float,
    n_samples: int, dv: float,
) -> bytes:
    return json.dumps({
        "f": expr, "x": x,
        "x_min": x_min, "x_max": x_max,
        "n_samples": n_samples, "dv": dv,
    }).encode("utf-8")


def _maybe_render_png(
    *, expr: str, x_var: str, x_min: float, x_max: float,
    n_samples: int, png_path: Path,
) -> None:
    """Render a PNG via matplotlib for human-readable inspection.
    No-op if matplotlib isn't installed."""
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        import sympy as sp
    except ImportError:
        print(
            "warning: matplotlib + sympy required for PNG render; "
            "skipping --png",
            file=sys.stderr,
        )
        return
    x_sym = sp.Symbol(x_var)
    f_expr = sp.sympify(expr, locals={x_var: x_sym})
    f = sp.lambdify(x_sym, f_expr, "math")
    step = (x_max - x_min) / (n_samples - 1)
    xs = [x_min + i * step for i in range(n_samples)]
    ys = []
    for xi in xs:
        try:
            yi = f(xi)
            if isinstance(yi, complex):
                yi = yi.real
            ys.append(float(yi))
        except (ValueError, ZeroDivisionError, OverflowError):
            ys.append(float("nan"))
    plt.figure(figsize=(8, 5))
    plt.plot(xs, ys, linewidth=2, color="#3366cc")
    plt.title(f"y = {expr}")
    plt.xlabel(x_var)
    plt.ylabel("y")
    plt.grid(True, alpha=0.3)
    png_path.parent.mkdir(parents=True, exist_ok=True)
    plt.tight_layout()
    plt.savefig(png_path, dpi=120)
    plt.close()


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument("--expr", required=True, help="SymPy expression in the variable --x")
    p.add_argument("--x", default="x", help="independent-variable name (default 'x')")
    p.add_argument("--x-min", dest="x_min", type=float, default=-3.14159)
    p.add_argument("--x-max", dest="x_max", type=float, default=3.14159)
    p.add_argument("--n-samples", dest="n_samples", type=int, default=200)
    p.add_argument("--dv", type=float, default=0.01,
                   help="value-quantization step (smaller = tighter)")
    p.add_argument("--png", type=Path, default=None,
                   help="optional PNG output path (requires matplotlib)")
    args = p.parse_args(argv)

    from arborist.pi_star import PiStarError, get
    try:
        ps = get("function-sampled@v1")
    except KeyError:
        print(
            "error: function-sampled@v1 not registered (sympy missing?); "
            "install with `pip install arborist[math]`",
            file=sys.stderr,
        )
        return 2

    payload = _build_payload(
        args.expr, args.x, args.x_min, args.x_max,
        args.n_samples, args.dv,
    )
    try:
        canonical_bytes = ps.canonicalize(payload)
    except PiStarError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1

    digest = hashlib.sha256(canonical_bytes).hexdigest()

    out = {
        "expr": args.expr,
        "x": args.x,
        "grid": {
            "x_min": args.x_min,
            "x_max": args.x_max,
            "n_samples": args.n_samples,
            "dv": args.dv,
        },
        "canonical_bytes_sha256": digest,
        "canonical_bytes_preview": canonical_bytes.decode("utf-8")[:200] + (
            "…" if len(canonical_bytes) > 200 else ""
        ),
        "canonical_bytes_total_chars": len(canonical_bytes),
        "pi_star_ref": "function-sampled@v1",
    }

    if args.png is not None:
        _maybe_render_png(
            expr=args.expr, x_var=args.x,
            x_min=args.x_min, x_max=args.x_max,
            n_samples=args.n_samples,
            png_path=args.png,
        )
        out["png_path"] = str(args.png)

    print(json.dumps(out, indent=2, ensure_ascii=False))
    return 0


if __name__ == "__main__":
    sys.exit(main())