open-standards/wai/wai-quantum-py/tests/test_widgets.py

"""The widgets: the data each draws (the engine's answers, nothing else), their
static fallbacks, their mime bundle, and their front ends run on that data."""

from __future__ import annotations

import json
import subprocess
import sys
import xml.etree.ElementTree as ET

import pytest

import wai_quantum as wq
from wai_quantum import widgets
from conftest import BELL, HERE, node

DEMO = wq.Key.demo()
SVG = "{http://www.w3.org/2000/svg}"


def test_circuit_data_is_the_engines_drawing():
    d = widgets.circuit_data(BELL)
    assert d["svg"] == wq.draw(BELL, format="svg") and d["text"] == wq.draw(BELL) and d["layout"] == wq.layout(BELL)
    assert d["hands"] is None and d["key"] == []
    c = wq.Circuit(BELL)
    assert c.data == d and c._repr_svg_() == d["svg"] and d["svg"] in c._repr_html_()
    assert repr(c) == "Circuit(2 qubits, 4 operations in 3 columns)"
    assert getattr(c, "selected", None) is None


def test_circuit_marks_each_gate_with_its_hand_from_a_history():
    src = "OPENQASM 3.0;\nqubit[2] q;\nh q[0];\ncx q[0], q[1];\n"
    k = wq.Key.generate()
    h = wq.history("start", {"source": src, "by": {"kind": "origin", "name": "example"}, "at": 1}, key=k)
    h = wq.history("propose", {"history": h["history"], "parent": "v0", "proposed": src + "x q[1];\n", "accepted": [0], "by": {"kind": "assistant"}, "at": 2}, key=k)
    d = widgets.circuit_data(src + "x q[1];\n", history=h["history"])
    assert d["hands"] == ["origin", "origin", "assistant"]
    assert d["key"] == [{"hand": "assistant", "label": "the assistant", "gates": 1}, {"hand": "origin", "label": "where it came from", "gates": 2}]
    root = ET.fromstring(d["svg"])
    marked = {g.get("data-op"): g.get("class") for g in root.iter(SVG + "g") if g.get("data-op")}
    assert marked == {"0": "qc-item qc-by-origin", "1": "qc-item qc-by-origin", "2": "qc-item qc-by-assistant"}
    with pytest.raises(ValueError):
        widgets.circuit_data(src, hands=["person"] * 2, history=h["history"])


def test_result_data_is_the_jobs_counts_method_accuracy_and_receipt_check():
    job = wq.run(BELL, shots=1000, seed=3, key=DEMO)
    d = widgets.result_data(job)
    assert d["counts"] == job.counts and list(d["counts"]) == sorted(job.counts, key=lambda k: -job.counts[k])
    assert (d["method"], d["accuracy"], d["shots"], d["distinct_outcomes"], d["omitted_shots"]) == ("stabilizer", "exact", 1000, 2, 0)
    assert d["receipt"]["valid"] is True and d["receipt"]["demo"] is True and d["receipt"]["signer"] == DEMO.did
    r = wq.Result(job)
    assert r.data == d and "stabilizer" in r._repr_html_() and "verified" in r._repr_html_()
    # The histogram draws one bar per outcome.
    html = r._repr_html_()
    svg = html[html.index("<svg") : html.index("</svg>") + 6]
    assert len(ET.fromstring(svg).findall(SVG + "rect")) == 2
    # A job's own card is the same drawing.
    assert job._repr_html_() == html
    forged = json.loads(json.dumps(job))
    forged["result"]["counts"]["00"] += 1
    assert widgets.result_data(forged)["receipt"]["valid"] is False and "INVALID" in wq.Result(forged)._repr_html_()


def test_result_data_carries_expectations_with_their_errors_and_amplitudes():
    job = wq.run("OPENQASM 3.0; qubit[2] q; h q[0]; ry(0.3) q[1];", observables=["Z1", "X0"], method="pauli", key=DEMO)
    d = widgets.result_data(job)
    assert [e["observable"] for e in d["expectations"]] == ["Z1", "X0"] and all("bound" in e for e in d["expectations"])
    assert "(bound)" in wq.Result(job)._repr_html_()
    amp = wq.run("OPENQASM 3.0; qubit[1] q; h q[0];", amplitudes=["0", "1"], key=DEMO)
    assert [a["bits"] for a in widgets.result_data(amp)["amplitudes"]] == ["0", "1"]


def test_receipt_data_is_the_engines_check():
    k = wq.Key.generate()
    job = wq.run(BELL, shots=10, key=k)
    d = widgets.receipt_data(job)
    bare_result = {x: v for x, v in job["result"].items() if x != "receipt"}
    assert d["valid"] is True and d["key"] == k.public_key and d["signer"] == k.did and d["demo"] is False
    assert d["digest"] == wq.digest(job["request"], bare_result) and d["alg"] == "ed25519" and d["canon"] == "rfc8785"
    w = wq.Receipt(job)
    assert w.valid and "verified" in w._repr_html_() and k.did in w._repr_html_()
    forged = json.loads(json.dumps(job))
    forged["request"]["shots"] = 11
    assert not wq.Receipt(forged).valid and "INVALID" in wq.Receipt(forged)._repr_html_()
    with pytest.raises(ValueError):
        widgets.receipt_data({"request": {}, "result": {}})


def test_bloch_data_is_the_engines_vectors():
    d = widgets.bloch_data("OPENQASM 3.0; qubit[2] q; h q[0]; ry(0.5) q[1];", key=DEMO)
    assert d["method"] == "statevector" and d["accuracy"] == "exact" and "job" not in d
    q0, q1 = d["qubits"]
    assert (q0["x"], q0["z"]) == pytest.approx((1, 0), abs=1e-12) and q1["x"] == pytest.approx(0.479425538604203, abs=1e-12)
    b = wq.Bloch("OPENQASM 3.0; qubit[2] q; h q[0]; ry(0.5) q[1];", key=DEMO)
    assert b.data == d and b._repr_html_().count("<svg") == 2
    assert widgets.bloch_data(wq.bloch("OPENQASM 3.0; qubit[1] q;", key=DEMO)) == widgets.bloch_data("OPENQASM 3.0; qubit[1] q;", key=DEMO)


def test_the_static_bloch_projection_puts_the_poles_and_the_equator_where_they_belong():
    from wai_quantum._render import project

    assert project((0, 0, 1), -35, 20)[1] == pytest.approx(0.9396926, abs=1e-6)
    assert project((0, 0, 1), 0, 0) == pytest.approx((0, 1))
    assert project((1, 0, 0), 0, 0) == pytest.approx((1, 0))
    # Seen from the equator straight on, +y points at the reader.
    assert project((0, 1, 0), 0, 0) == pytest.approx((0, 0))


@pytest.mark.skipif(not widgets.HAVE_WIDGETS, reason="the widgets extra is not installed")
def test_each_widget_offers_its_view_and_its_static_html():
    from IPython.core.formatters import DisplayFormatter

    job = wq.run(BELL, shots=100, key=DEMO)
    for w in (wq.Circuit(BELL), wq.Result(job), wq.Receipt(job), wq.Bloch(BELL, key=DEMO)):
        bundle, _ = DisplayFormatter().format(w)
        assert "application/vnd.jupyter.widget-view+json" in bundle, type(w)
        assert "text/html" in bundle and bundle["text/html"] == w._repr_html_()
        assert "export default" in w._esm and w.data
    c = wq.Circuit(BELL)
    c.selected = 1
    assert c.selected == 1


def test_the_widgets_render_without_their_library():
    code = (
        "import sys; sys.modules['anywidget'] = None\n"
        "import wai_quantum as wq\n"
        "from wai_quantum import widgets\n"
        "assert not widgets.HAVE_WIDGETS\n"
        f"bell = {BELL!r}\n"
        "job = wq.run(bell, shots=50, key=wq.Key.demo())\n"
        "parts = [wq.Circuit(bell)._repr_html_(), wq.Result(job)._repr_html_(), wq.Receipt(job)._repr_html_(), wq.Bloch(bell, key=wq.Key.demo())._repr_html_()]\n"
        "assert all(parts) and 'verified' in parts[2]\n"
        "print('ok')\n"
    )
    p = subprocess.run([sys.executable, "-I", "-c", code], capture_output=True, text=True, timeout=120)
    assert p.returncode == 0 and p.stdout.strip() == "ok", p.stderr


def test_the_front_ends_draw_the_data_they_are_given(tmp_path):
    exe = node()
    job = wq.run(BELL, shots=1000, seed=3, key=DEMO)
    forged = json.loads(json.dumps(job))
    forged["result"]["counts"]["00"] += 1
    payloads = {
        "circuit": widgets.circuit_data(BELL),
        "circuit_hands": widgets.circuit_data("OPENQASM 3.0; qubit[2] q; h q[0]; cx q[0], q[1];", hands=["person", "assistant"]),
        "result": widgets.result_data(job),
        "receipt": widgets.receipt_data(job),
        "receipt_forged": widgets.receipt_data(forged),
        "bloch": widgets.bloch_data("OPENQASM 3.0; qubit[2] q; h q[0]; ry(0.5) q[1];", key=DEMO),
    }
    f = tmp_path / "payloads.json"
    f.write_text(json.dumps(payloads))
    p = subprocess.run([exe, str(HERE / "front_ends.mjs"), str(f)], capture_output=True, text=True, timeout=120)
    assert p.returncode == 0, p.stdout + p.stderr
    assert "front ends: 4 widgets," in p.stdout, p.stdout