open-standards/wai/quantum-cross-simulator/provider/seeded_statevector_provider.py

#!/usr/bin/env python3
"""A provider of the open provider protocol (wai/extensions/quantum-providers.md) around a
third-party state-vector simulator, run as a local process: one JSON message a line on standard
input, one JSON answer a line on standard output (§2.2).

    wai-quantum providers add sim --kind process -- python3 -I seeded_statevector_provider.py
    wai-quantum providers test sim
    wai-quantum run bell.qasm --provider sim --shots 1000 --json > job.json
    wai-quantum verify job.json

It is the template to copy for a device of your own. Everything the protocol asks of a provider
is here, and the device is one class, `Device`: replace its `run` with a call to your device and
its `DESCRIPTION` with what your device is, and the rest stays as it is.

The device here is the pinned simulator the corpus compares against (env/requirements.lock),
seeded, so the same job gives the same counts. A program arrives in the device's own gates with
every angle an exact fraction of a turn; each gate is built from the simulator's own gate of that
meaning, and nothing is transpiled, so the program the client compiled and checked is the one
that runs. Nothing here takes or needs a credential.
"""
import json
import math
import sys

PROTOCOL = "wai.quantum.provider/1"
BASES = {"rzsxxcz": ["rz", "sx", "x", "cz"], "prxcz": ["prx", "cz"], "gpizz": ["gpi", "gpi2", "zz"]}
MAX_SHOTS = 1_000_000


class Refusal(Exception):
    def __init__(self, code, message):
        super().__init__(message)
        self.code, self.message = code, message


class Device:
    """The simulator. Replace `run` (and DESCRIPTION) to bring another device."""

    NAME = "seeded-statevector"
    QUBITS = 24
    BASIS = "rzsxxcz"
    DESCRIPTION = "a third-party state-vector simulator in double precision (Qiskit Aer, the version pinned in the corpus's lock), sampled with the job's seed"  # allow-name: the simulator this template wraps, named as data

    def __init__(self):
        import qiskit_aer

        self.version = qiskit_aer.__version__

    def run(self, program, shots, seed):
        """Counts keyed by the program's classical bits, the highest bit first."""
        from qiskit import QuantumCircuit
        from qiskit_aer import AerSimulator

        used = sorted({q for g in program["gates"] for q in g["qubits"]} | {m["qubit"] for m in program["measure"]})
        at = {q: i for i, q in enumerate(used)}
        qc = QuantumCircuit(len(used), program["bits"])
        for g in program["gates"]:
            q = [at[x] for x in g["qubits"]]
            t = [2 * math.pi * m / (1 << k) for m, k in g.get("turns", [])]
            name = g["gate"]
            # Each gate by its definition in the protocol (§4), from the simulator's gate of that meaning.
            if name == "rz":
                qc.rz(t[0], q[0])  # exp(-iπt·Z) is Rz(2πt)
            elif name == "sx":
                qc.sx(q[0])
            elif name == "x":
                qc.x(q[0])
            elif name == "cz":
                qc.cz(q[0], q[1])
            elif name in ("prx", "gpi", "gpi2"):
                theta, phi = {"prx": (t[0] if t else 0, t[1] if len(t) > 1 else 0), "gpi": (math.pi, t[0]), "gpi2": (math.pi / 2, t[0])}[name]
                qc.rz(-phi, q[0])
                qc.rx(theta, q[0])
                qc.rz(phi, q[0])
            elif name == "zz":
                qc.rzz(t[0], q[0], q[1])  # exp(-iπt·Z⊗Z) is Rzz(2πt)
        for m in program["measure"]:
            qc.measure(at[m["qubit"]], m["bit"])
        sim = AerSimulator(method="statevector", seed_simulator=seed)
        counts = sim.run(qc, shots=shots).result().get_counts()
        return {k.replace(" ", ""): int(v) for k, v in counts.items()}


def check_program(m, device):
    """The program and its shots, refused as the protocol says (§4) when they are not one this
    device runs."""
    p = m.get("program")
    if not isinstance(p, dict) or p.get("format") != "wai.native/1":
        raise Refusal("invalid", "a program's format is wai.native/1")
    basis = p.get("basis")
    if basis not in BASES:
        raise Refusal("invalid", "a program's basis is rzsxxcz, prxcz or gpizz")
    bits = p.get("bits")
    if not isinstance(bits, int) or isinstance(bits, bool) or not 1 <= bits <= 64:
        raise Refusal("invalid", "a program's bits count its classical bits, 1 to 64")
    shots = m.get("shots")
    if not isinstance(shots, int) or isinstance(shots, bool) or not 1 <= shots <= MAX_SHOTS:
        raise Refusal("invalid", "%r shots; this device takes 1 to %d" % (shots, MAX_SHOTS))
    arity = {"rz": (1, 1), "sx": (1, 0), "x": (1, 0), "cz": (2, 0), "prx": (1, 2), "gpi": (1, 1), "gpi2": (1, 1), "zz": (2, 1)}
    gates = p.get("gates")
    if not isinstance(gates, list):
        raise Refusal("invalid", "a program's gates are a list")
    for i, g in enumerate(gates):
        name = g.get("gate") if isinstance(g, dict) else None
        if name not in BASES[basis]:
            raise Refusal("invalid", "gate %d: %r is not a gate of the basis %s" % (i, name, basis))
        qs, ts = g.get("qubits"), g.get("turns", [])
        nq, nt = arity[name]
        if not isinstance(qs, list) or len(qs) != nq or not all(isinstance(q, int) and q >= 0 for q in qs) or len(set(qs)) != nq:
            raise Refusal("invalid", "gate %d: the wrong qubits for %s" % (i, name))
        if not isinstance(ts, list) or len(ts) != nt:
            raise Refusal("invalid", "gate %d: the wrong turns for %s" % (i, name))
        for t in ts:
            if not (isinstance(t, list) and len(t) == 2 and all(isinstance(x, int) and x >= 0 for x in t) and t[1] <= 32 and t[0] < (1 << t[1])):
                raise Refusal("invalid", "gate %d: %r is not m/2^k of a turn with m < 2^k, k <= 32" % (i, t))
    measure = p.get("measure")
    if not isinstance(measure, list):
        raise Refusal("invalid", "a program's measurements are a list")
    read = {}
    for x in measure:
        b, q = x.get("bit"), x.get("qubit")
        if not isinstance(b, int) or not 0 <= b < bits or not isinstance(q, int) or q < 0:
            raise Refusal("invalid", "%r: a measurement reads a qubit into one of the program's bits" % (x,))
        if b in read:
            raise Refusal("invalid", "bit %d is measured twice" % b)
        read[b] = q
    if sorted(read) != list(range(bits)):
        raise Refusal("invalid", "bit %d measures no qubit" % min(set(range(bits)) - set(read)))
    if len(set(read.values())) != len(read):
        raise Refusal("invalid", "a qubit is measured into two bits")
    # Well formed; now, can this device run it?
    if basis != device.BASIS:
        raise Refusal("refused", "this device runs %s, not %s" % (" ".join(BASES[device.BASIS]), " ".join(BASES[basis])))
    if any(q >= device.QUBITS for g in gates for q in g["qubits"]) or any(q >= device.QUBITS for q in read.values()):
        raise Refusal("refused", "a qubit this device does not have: it has %d" % device.QUBITS)
    return p, shots


class Provider:
    def __init__(self, device):
        self.device = device
        self.jobs = {}
        self.next = 1

    def describe(self):
        d = self.device
        return {
            "ok": True,
            "protocol": PROTOCOL,
            "provider": d.NAME,
            "device": "%s-%d" % (d.NAME, d.QUBITS),
            "provenance": "simulator",
            "simulator": "%s, version %s" % (d.DESCRIPTION, d.version),
            "qubits": d.QUBITS,
            "native_gates": BASES[d.BASIS],
            "connectivity": "all",
            "calibration": {"id": "%s-%s" % (d.NAME, d.version)},
            "billing": {"free": True},
            "max_shots": MAX_SHOTS,
            "seeded": True,
        }

    def handle(self, m):
        if not isinstance(m, dict):
            raise Refusal("invalid", "a message is a JSON object")
        if "protocol" not in m:
            raise Refusal("invalid", "a message names its protocol")
        if m["protocol"] != PROTOCOL:
            raise Refusal("unsupported_version", "this provider speaks %s, not %s" % (PROTOCOL, m["protocol"]))
        op = m.get("op")
        if op == "describe":
            return self.describe()
        if op == "estimate":
            check_program(m, self.device)
            return {"ok": True, "free": True, "seconds": 0}
        if op == "submit":
            p, shots = check_program(m, self.device)
            seed = m.get("seed", 1)
            if not isinstance(seed, int) or seed < 0:
                raise Refusal("invalid", "a seed is a whole number")
            try:
                counts = self.device.run(p, shots, seed)
            except Exception as e:  # noqa: BLE001 - the device's failure, reported as the protocol's
                raise Refusal("internal", "the simulator failed: %s" % type(e).__name__)
            job = "sv-%d" % self.next
            self.next += 1
            self.jobs[job] = {"counts": counts, "shots": shots}
            return {"ok": True, "job": job}
        if op in ("status", "result", "cancel"):
            job = m.get("job")
            if not isinstance(job, str):
                raise Refusal("invalid", "name the `job`")
            if job not in self.jobs:
                raise Refusal("not_found", "no job %s" % job)
            if op == "result":
                j = self.jobs[job]
                return {"ok": True, "counts": j["counts"], "shots": j["shots"]}
            # A job here finishes when it is submitted, so it is always done.
            return {"ok": True, "status": "done"}
        if op is None:
            raise Refusal("invalid", "a message names its op")
        raise Refusal("unsupported", "no operation %r; the operations are describe, estimate, submit, status, result and cancel" % (op,))


def main():
    provider = Provider(Device())
    for line in sys.stdin:
        if not line.strip():
            continue
        try:
            answer = provider.handle(json.loads(line))
        except Refusal as r:
            answer = {"ok": False, "error": {"code": r.code, "message": r.message}}
        except ValueError:
            answer = {"ok": False, "error": {"code": "invalid", "message": "not JSON"}}
        sys.stdout.write(json.dumps(answer) + "\n")
        sys.stdout.flush()


if __name__ == "__main__":
    main()