dense_evolution/compiler.py

import numpy as np
from typing import List, Tuple
from dataclasses import dataclass

# ── optional JAX import (same pattern as registry) ────────────────────────────
try:
    from .registry import HAS_JAX
except ImportError:
    try:
        import jax
        HAS_JAX = True
    except ImportError:
        HAS_JAX = False

if HAS_JAX:
    import jax
    import jax.numpy as jnp
    jax.config.update("jax_enable_x64", True)

# ─────────────────────────────────────────────────────────────────────────────
# Gate-ID encoding (shared between _apply_gate_fast_step and the beast-mode
# command builder in the dashboard).
#
#  0  I (identity)       7  T
#  1  H                  8  Tdg
#  2  X                  9  Rx(θ)
#  3  Y                 10  Ry(θ)
#  4  Z                 11  Rz(θ)
#  5  S                 12  Phase / P(θ) / U1(θ)
#  6  Sdg               13  SX (√X)
#                       ──  2-qubit gates ──
#                       20  CX / CNOT
#                       21  CZ
#                       22  CP(θ)              — phase on |11⟩ only
#                       23  SWAP (never dispatched here — always decomposed
#                           into 3xCX by QuantumTranspiler.transpile first)
#                       24  CY
#                       25  CRZ(θ)             — phase on target conditioned
#                           on control, distinct from CP: NOT the same gate,
#                           do not merge into 22 (see apply_crz below)
# ─────────────────────────────────────────────────────────────────────────────

if HAS_JAX:

    @jax.jit
    def _apply_gate_fast_step(sv: "jnp.ndarray",
                               operation: "jnp.ndarray"):
        """
        Apply a single quantum gate to statevector *sv*.

        Parameters
        ----------
        sv        : complex128 statevector of shape (2**n,)
        operation : float64 array [g_id, q1, q2, param]
                    g_id  — gate identifier (see table above)
                    q1    — target qubit (1-qubit gates) or control qubit (2-qubit)
                    q2    — target qubit for 2-qubit gates; unused for 1-qubit
                    param — rotation angle in radians; 0.0 for non-parametric gates

        Returns
        -------
        (new_sv, None)  — compatible with jax.lax.scan
        """
        g_id  = operation[0].astype(jnp.int32)
        q1    = operation[1].astype(jnp.int32)
        q2    = operation[2].astype(jnp.int32)
        param = operation[3]
        dim   = sv.shape[0]                    # static (shape is concrete even under jit)
        n_qubits = dim.bit_length() - 1        # static; dim == 2**n_qubits always

        # Every constant below must match sv's own dtype (complex64 if
        # use_float32=True, complex128 otherwise), not a dtype hardcoded to
        # complex128 — a fixed dtype here made this function only traceable
        # for complex128 input: apply_cp (inside do_2q) multiplied by an
        # exp_pos hardcoded to complex128, while lax.cond's other branch
        # (identity, `lambda s: s`) preserved sv's real dtype, so any
        # complex64 circuit hit "cond branches must have equal output
        # types" immediately at trace time — even for circuits with only
        # 1-qubit gates, since do_2q's body (and this cond) is traced
        # unconditionally as part of the is_1q/is_2q dispatch below.
        sv_dtype = sv.dtype
        inv2    = jnp.asarray(1.0 / jnp.sqrt(2.0), dtype=sv_dtype)
        half_p  = param * jnp.float64(0.5)
        cos_p   = jnp.cos(half_p).astype(sv_dtype)
        sin_p   = jnp.sin(half_p).astype(sv_dtype)
        exp_pos = jnp.exp( 1j * param).astype(sv_dtype)
        exp_neg = jnp.exp(-1j * param).astype(sv_dtype)
        exp_ph4 = jnp.exp( 1j * jnp.pi / 4.0).astype(sv_dtype)
        exp_mh4 = jnp.exp(-1j * jnp.pi / 4.0).astype(sv_dtype)
        # half-angle phases for CRZ(θ) — same convention as the 1-qubit
        # RZ(θ) switch entry below, named here since apply_crz (do_2q) needs
        # them inside a conditional function rather than an inline literal.
        exp_pos_half = jnp.exp( 1j * half_p).astype(sv_dtype)
        exp_neg_half = jnp.exp(-1j * half_p).astype(sv_dtype)

        # ── 1-qubit gate matrix selection via lax.switch ──────────────
        # Index must be in [0, 13]; anything outside is clamped to 0 (I).
        safe_gid = jnp.clip(g_id, 0, 13)

        g_1q = jax.lax.switch(
            safe_gid,
            [
                # 0  I
                lambda _: jnp.eye(2, dtype=sv_dtype),
                # 1  H
                lambda _: jnp.array(
                    [[inv2,  inv2],
                     [inv2, -inv2]], dtype=sv_dtype),
                # 2  X
                lambda _: jnp.array(
                    [[0.0+0j, 1.0+0j],
                     [1.0+0j, 0.0+0j]], dtype=sv_dtype),
                # 3  Y
                lambda _: jnp.array(
                    [[0.0+0j, -1j],
                     [1j,      0.0+0j]], dtype=sv_dtype),
                # 4  Z
                lambda _: jnp.array(
                    [[1.0+0j,  0.0+0j],
                     [0.0+0j, -1.0+0j]], dtype=sv_dtype),
                # 5  S
                lambda _: jnp.array(
                    [[1.0+0j, 0.0+0j],
                     [0.0+0j, 1j    ]], dtype=sv_dtype),
                # 6  Sdg
                lambda _: jnp.array(
                    [[1.0+0j, 0.0+0j],
                     [0.0+0j, -1j   ]], dtype=sv_dtype),
                # 7  T
                lambda _: jnp.array(
                    [[1.0+0j, 0.0+0j],
                     [0.0+0j, exp_ph4]], dtype=sv_dtype),
                # 8  Tdg
                lambda _: jnp.array(
                    [[1.0+0j, 0.0+0j],
                     [0.0+0j, exp_mh4]], dtype=sv_dtype),
                # 9  Rx(θ)  = [[cos θ/2, -i sin θ/2], [-i sin θ/2, cos θ/2]]
                lambda _: jnp.array(
                    [[cos_p,      -1j * sin_p],
                     [-1j * sin_p, cos_p     ]], dtype=sv_dtype),
                # 10  Ry(θ) = [[cos θ/2, -sin θ/2], [sin θ/2, cos θ/2]]
                lambda _: jnp.array(
                    [[cos_p,  -sin_p],
                     [sin_p,   cos_p]], dtype=sv_dtype),
                # 11  Rz(θ) = [[e^{-iθ/2}, 0], [0, e^{iθ/2}]]
                lambda _: jnp.array(
                    [[jnp.exp(-1j * half_p), 0.0+0j            ],
                     [0.0+0j,                jnp.exp(1j * half_p)]],
                    dtype=sv_dtype),
                # 12  Phase / P(θ) / U1(θ) = [[1, 0], [0, e^{iθ}]]
                lambda _: jnp.array(
                    [[1.0+0j, 0.0+0j],
                     [0.0+0j, exp_pos]], dtype=sv_dtype),
                # 13  SX (√X) = 0.5 * [[1+i, 1-i], [1-i, 1+i]]
                lambda _: jnp.array(
                    [[0.5+0.5j, 0.5-0.5j],
                     [0.5-0.5j, 0.5+0.5j]], dtype=sv_dtype),
            ],
            operand=None,
        )

        # ── 1-qubit application ────────────────────────────────────────
        def do_1q(_sv):
            # MSB-first, matching the rest of the simulator (simulator.py's
            # own docstring/_qubit_stride_pairs: qubit 0 = most significant
            # bit, phys = n_qubits - 1 - qubit) — this used to be a raw
            # `1 << q1` (LSB-first), silently disagreeing with run_circuit()/
            # apply_gate_1q()/measure() on every multi-qubit circuit that
            # isn't symmetric under qubit reversal (verified: X on qubit 0
            # in a 3-qubit register gave index 1 here vs index 4 via
            # run_circuit() before this fix — see CHANGELOG).
            phys      = jnp.int64(n_qubits - 1) - q1.astype(jnp.int64)
            stride    = jnp.int64(1) << phys
            idx_full  = jnp.arange(dim, dtype=jnp.int64)
            mask_0    = (idx_full & stride) == 0

            # idx_0: indices where qubit q1 == 0
            # idx_1: corresponding |1⟩ partners
            # We build them without xp.where-tuple confusion:
            # any index i has its pair at i ^ stride.
            # For i in |0⟩ slots (mask_0): partner = i | stride = i ^ stride
            # For i in |1⟩ slots (¬mask_0): partner = i ^ stride  (clears bit)
            # We process all indices simultaneously using the |0⟩ slot's amplitude.
            idx_pair  = idx_full ^ stride          # each element's partner
            amp_self  = _sv[idx_full]              # a[i]
            amp_pair  = _sv[idx_pair]              # a[i ^ stride]

            # When mask_0: amp_self = a_0, amp_pair = a_1
            # new_0 = g00*a_0 + g01*a_1
            # new_1 = g10*a_0 + g11*a_1
            g00 = g_1q[0, 0];  g01 = g_1q[0, 1]
            g10 = g_1q[1, 0];  g11 = g_1q[1, 1]

            new_when_0 = g00 * amp_self + g01 * amp_pair   # result for |0⟩ slot
            new_when_1 = g10 * amp_pair + g11 * amp_self   # result for |1⟩ slot
            # NOTE: for |1⟩ slots, amp_pair is the |0⟩ amplitude and
            # amp_self is the |1⟩ amplitude — roles are swapped.

            return jnp.where(mask_0, new_when_0, new_when_1)

        # ── 2-qubit application ───────────────────────────────────────
        def do_2q(_sv):
            # Same MSB-first fix as do_1q above — ctrl/trgt are physical bit
            # positions (n_qubits-1-qubit), not raw qubit indices.
            ctrl     = jnp.int64(n_qubits - 1) - q1.astype(jnp.int64)
            trgt     = jnp.int64(n_qubits - 1) - q2.astype(jnp.int64)
            idx_full = jnp.arange(dim, dtype=jnp.int64)

            ctrl_bit_set = (idx_full & (jnp.int64(1) << ctrl)) != 0
            trgt_bit_set = (idx_full & (jnp.int64(1) << trgt)) != 0

            # CX: flip target bit when control is set
            def apply_cx(__sv):
                partner = idx_full ^ (jnp.int64(1) << trgt)
                swapped = __sv[partner]
                return jnp.where(ctrl_bit_set, swapped, __sv)

            # CZ: negate amplitude when both control and target bits are set
            def apply_cz(__sv):
                both_set = ctrl_bit_set & trgt_bit_set
                return jnp.where(both_set, -__sv, __sv)

            # CP(θ): phase kick e^{iθ} on |11⟩ component only — NOT the same
            # gate as CRZ below (see apply_crz docstring for the distinction).
            def apply_cp(__sv):
                both_set = ctrl_bit_set & trgt_bit_set
                return jnp.where(both_set, exp_pos * __sv, __sv)

            # CY: apply Y to target when control is set. Same index-pair
            # structure as apply_cx (swap target=0/target=1 partners), but
            # with Y's phase factors instead of a plain swap: Y|0⟩=i|1⟩,
            # Y|1⟩=-i|0⟩ — the new amplitude at each slot is ±i times its
            # PARTNER's amplitude (matches do_1q's own Y row-formula, since
            # Y's diagonal entries are both 0).
            def apply_cy(__sv):
                partner = idx_full ^ (jnp.int64(1) << trgt)
                swapped = __sv[partner]
                phase   = jnp.where(trgt_bit_set, 1j, -1j).astype(sv_dtype)
                return jnp.where(ctrl_bit_set, phase * swapped, __sv)

            # CRZ(θ): apply RZ(θ) to target when control is set — a DIAGONAL
            # phase conditioned on the TARGET's own bit value (e^{-iθ/2} if
            # target=0, e^{+iθ/2} if target=1), not on "both bits set" like
            # CP above. Deliberately a separate function/id from apply_cp:
            # CP and CRZ are mathematically different gates that used to be
            # silently merged under one id before this fix (CRZ wasn't
            # reachable at all — see CHANGELOG).
            def apply_crz(__sv):
                phase = jnp.where(trgt_bit_set, exp_pos_half, exp_neg_half)
                return jnp.where(ctrl_bit_set, phase * __sv, __sv)

            # SWAP: exchange amplitudes of ctrl-bit and trgt-bit positions
            def apply_swap(__sv):
                # Standard SWAP = CX(c,t) · CX(t,c) · CX(c,t)
                # Computed directly: for each (ctrl=0,trgt=1) pair with
                # the other bits identical, swap the two amplitudes.
                only_ctrl = ( ctrl_bit_set & ~trgt_bit_set)
                only_trgt = (~ctrl_bit_set &  trgt_bit_set)
                swap_mask = only_ctrl | only_trgt
                partner   = idx_full ^ (jnp.int64(1) << ctrl) ^ (jnp.int64(1) << trgt)
                return jnp.where(swap_mask, __sv[partner], __sv)

            # Dispatch on g_id: 20=CX, 21=CZ, 22=CP, 23=SWAP, 24=CY, 25=CRZ
            is_cx   = g_id == 20
            is_cz   = g_id == 21
            is_cp   = g_id == 22
            is_cy   = g_id == 24
            is_crz  = g_id == 25
            # is_swap = g_id == 23 (default branch — never actually reached,
            # see comment at the top of this file: QuantumTranspiler always
            # decomposes 'swap' into 3xCX before a gate name reaches here)

            after_cx   = jax.lax.cond(is_cx,   apply_cx,   lambda s: s, _sv)
            after_cz   = jax.lax.cond(is_cz,   apply_cz,   lambda s: s, _sv)
            after_cp   = jax.lax.cond(is_cp,   apply_cp,   lambda s: s, _sv)
            after_cy   = jax.lax.cond(is_cy,   apply_cy,   lambda s: s, _sv)
            after_crz  = jax.lax.cond(is_crz,  apply_crz,  lambda s: s, _sv)
            after_swap = apply_swap(_sv)

            # Pick the right result
            result = jnp.where(is_cx,  after_cx,
                     jnp.where(is_cz,  after_cz,
                     jnp.where(is_cp,  after_cp,
                     jnp.where(is_cy,  after_cy,
                     jnp.where(is_crz, after_crz,
                                       after_swap)))))
            return result

        # ── branch on 1-qubit vs 2-qubit ─────────────────────────────
        # g_id <= 12 → 1-qubit;  g_id >= 20 → 2-qubit.
        # Both branches must have identical output dtypes — enforced here by
        # casting both outputs to sv_dtype (sv's own dtype, complex64 or
        # complex128 depending on use_float32), not a dtype hardcoded to
        # complex128 (see note above: forcing complex128 here regardless of
        # input dtype would also break jax.lax.scan's carry-type check on
        # the very next iteration when sv started out as complex64).
        is_1q = g_id <= 13
        new_sv = jax.lax.cond(
            is_1q,
            lambda s: do_1q(s).astype(sv_dtype),
            lambda s: do_2q(s).astype(sv_dtype),
            sv,
        )
        return new_sv, None


    def _run_circuit_scan_core(state_vector: "jnp.ndarray",
                                compiled_ops:  "jnp.ndarray") -> "jnp.ndarray":
        """
        Execute a pre-compiled gate sequence on *state_vector* via jax.lax.scan.

        Parameters
        ----------
        state_vector : complex128 array of shape (2**n,)
        compiled_ops : float64 array of shape (n_gates, 4)
                       each row = [g_id, q1, q2, param]

        Returns
        -------
        Final statevector after all gates.
        """
        final_sv, _ = jax.lax.scan(_apply_gate_fast_step, state_vector, compiled_ops)
        return final_sv

    #: Plain (non-donating) wrapper. Required by callers that reuse the
    #: same state_vector buffer across multiple calls — run_parametric_batch_jit
    #: (simulator.py), where it's a vmap-broadcast closure variable shared
    #: across every batch lane, and circuit_to_energy_fn's energy_fn
    #: (autodiff.py), where the caller's VQE loop passes the same stato_zero
    #: to every epoch. Donating argument 0 there would make JAX raise on the
    #: second use (the buffer is deleted after a donated call) — not a
    #: silent bug, but a real break, verified by tracing every call site
    #: before adding donation anywhere.
    _compile_and_run_circuit_jit = jax.jit(_run_circuit_scan_core)

    #: Donating wrapper — safe ONLY where the caller immediately rebinds its
    #: reference to the input buffer and never reads the old one again.
    #: Currently just DenseSVSimulator.run_circuit_jit_beast_mode's
    #: `self.sv = ...` pattern (verified: no code path anywhere keeps a
    #: stale reference to self.sv across that call, including chunked/
    #: repeated invocations via run_circuit_with_chunking, and separate
    #: DenseSVSimulator instances never share a buffer).
    _compile_and_run_circuit_jit_donated = jax.jit(
        _run_circuit_scan_core, donate_argnums=(0,))


# ─────────────────────────────────────────────────────────────────────────────
# QuantumTranspiler
# ─────────────────────────────────────────────────────────────────────────────

class QuantumTranspiler:
    """
    Gate-level transpiler: decomposes non-native gates into the native
    {H, T, Tdg, CX, CZ} basis and performs basic circuit optimisations.
    """

    @staticmethod
    def decompose_toffoli(c1: int, c2: int, t: int) -> List[Tuple]:
        """
        Decompose CCX (Toffoli) into 15 native gates using the
        standard T/Tdg/CX Barenco decomposition.

        Gate count: 6 CX + 7 single-qubit (H, T, Tdg) = 15 total.
        """
        return [
            ('h',   t),
            ('cx',  c2, t),  ('tdg', t),
            ('cx',  c1, t),  ('t',   t),
            ('cx',  c2, t),  ('tdg', t),
            ('cx',  c1, t),
            ('t',   c2),     ('t',   t),
            ('cx',  c1, c2), ('h',   t),
            ('t',   c1),     ('tdg', c2),
            ('cx',  c1, c2),
        ]

    @staticmethod
    def decompose_swap(q1: int, q2: int) -> List[Tuple]:
        """Decompose SWAP into 3 CX gates."""
        return [('cx', q1, q2), ('cx', q2, q1), ('cx', q1, q2)]

    @classmethod
    def transpile(cls, circuit: List[Tuple]) -> List[Tuple]:
        """
        Expand CCX → 15 native gates and SWAP → 3 CX.
        All other gates are passed through unchanged.

        Parameters
        ----------
        circuit : list of tuples (gate_name, qubit, ...)

        Returns
        -------
        Expanded circuit as a list of tuples.
        """
        out: List[Tuple] = []
        for cmd in circuit:
            name = cmd[0].lower()            # BUG FIX: was cmd.lower() on a tuple
            if name == 'ccx':
                out.extend(cls.decompose_toffoli(*cmd[1:4]))
            elif name == 'swap':
                out.extend(cls.decompose_swap(*cmd[1:3]))
            else:
                out.append(cmd)
        return out