from typing import Callable, Optional, Tuple from .parser import QASMCircuit from .gates import GATE_IDS from .compiler import QuantumTranspiler from .registry import HAS_JAX, NoiseModel, NoiseSpec if HAS_JAX: import jax import jax.numpy as jnp from .compiler import _compile_and_run_circuit_jit else: jnp = None #: gates that receive a value from theta — must match the n_params count #: below exactly, or theta's allocation order desyncs from the template's #: injection order. Kept identical to dashboard_core's own list (same #: engine, one source of truth). _PARAMETRIC_GATES = ('rx', 'ry', 'rz', 'u1', 'p', 'cp', 'crz') _TWO_QUBIT_GATES = ('cx', 'cy', 'cz', 'cp', 'crz', 'swap') def _require_jax(): if not HAS_JAX: raise ImportError( "circuit_to_energy_fn requires JAX. " "Install it with: pip install dense-evolution[jax]") def _build_template(circuit: QASMCircuit, n_qubits: int) -> "jnp.ndarray": """Builds the (n_ops, 4) float64 [g_id, q1, q2, sentinel] template that the energy function injects theta into (-1.0 in the param slot for gates whose value comes from theta, patched in via jnp.where inside a jax.lax.scan, never a Python float() — that would sever the JAX trace). Structural pass only: build (name, *qubits) tuples (no param values — QuantumTranspiler.transpile only inspects gate name/qubit-count, for ccx/swap decomposition), transpile once, then look up g_id per gate and mark parametric slots with the sentinel. ccx/toffoli decomposes into non-parametric gates only, so this never desyncs theta's order. circuit.ops qubits are always plain ints here (QASMCircuit is the package's own interchange type, produced by QASMParser.parse and by the Qiskit/PennyLane interop bridge alike) — no defensive unwrapping of framework-specific qubit/wire objects needed. """ tuples = [] for op in circuit.ops: name = str(op['name']).lower().strip() qubits = [int(q) for q in op.get('qubits', [])] if not qubits or any(q >= n_qubits for q in qubits): continue tuples.append((name, *qubits)) target = QuantumTranspiler.transpile(tuples) rows = [] for cmd in target: name = cmd[0].lower() if name not in GATE_IDS: continue g_id = float(GATE_IDS[name]) qubits = cmd[1:] sentinel = -1.0 if name in _PARAMETRIC_GATES else 0.0 if name in _TWO_QUBIT_GATES and len(qubits) >= 2: rows.append([g_id, float(qubits[0]), float(qubits[1]), sentinel]) elif qubits: rows.append([g_id, float(qubits[0]), 0.0, sentinel]) if not rows: return jnp.empty((0, 4), dtype=jnp.float64) return jnp.array(rows, dtype=jnp.float64) def circuit_to_energy_fn( circuit: QASMCircuit, n_qubits: int ) -> Tuple[Callable, int]: """ Convert a QASMCircuit into a JAX-differentiable energy function. circuit : QASMCircuit — from QASMParser.parse(qasm), or from the Qiskit/PennyLane interop bridge (from_qiskit/from_pennylane). Returns (energy_fn, n_params): energy_fn(theta, h_matrix, stato_zero=None, noise=None) -> (energy, statevector) is a pure JAX function, differentiable w.r.t. theta via jax.grad / jax.value_and_grad(energy_fn, argnums=0, has_aux=True). stato_zero defaults to |0...0> if not given. n_params is the number of parametric gates in the circuit, in the same order theta is injected — build theta as an array of that length. noise, when given, is a registry.NoiseSpec (a JAX PyTree) applied to the statevector right after the circuit and before the energy expectation value is computed — natively inside the same traced computation as theta, not as an external step the caller has to splice in around energy_fn themselves. Because NoiseSpec carries its own jax_key as a pytree leaf, the whole thing stays jit/grad/vmap-composable with no OS-entropy fallback and no external key-management workaround: noise = NoiseSpec(model='depolarizing', p=0.05, jax_key=jax.random.PRNGKey(0)) energy, sv = energy_fn(theta, h_matrix, noise=noise) This is the same engine dashboard_core.py's real VQE gradient uses internally (verified against finite differences, ~1e-11 agreement) — exposed here as public API so it's reachable without reading dashboard_core.py, and so circuits imported via from_qiskit/ from_pennylane (which are NOT differentiable on their own — see run_pennylane_circuit's docstring) have a real way to become differentiable instead of just a documented dead end. """ _require_jax() template = _build_template(circuit, n_qubits) n_params = sum(1 for op in circuit.ops if str(op['name']).lower().strip() in _PARAMETRIC_GATES) def energy_fn(theta, h_matrix, stato_zero: Optional["jnp.ndarray"] = None, noise: Optional["NoiseSpec"] = None): if stato_zero is None: stato_zero = jnp.zeros(2 ** n_qubits, dtype=jnp.complex128).at[0].set(1.0) if n_params == 0: # No parametric gates -> no sentinel (-1.0) rows in template, so # patch_and_apply below would never take its is_param branch. # Skip the scan entirely rather than index into an empty theta # array during tracing (n_params is a static Python int, fixed # at circuit_to_energy_fn() call time, so this branch is # resolved before any tracing happens — not a jax.lax.cond). sv = _compile_and_run_circuit_jit(stato_zero, template) else: def patch_and_apply(carry, op): idx = carry is_param = op[3] == -1.0 final_p = jnp.where(is_param, theta[idx], op[3]) next_idx = jnp.where(is_param, idx + jnp.int32(1), idx) return next_idx, jnp.array([op[0], op[1], op[2], final_p], dtype=jnp.float64) _, patched_ops = jax.lax.scan(patch_and_apply, jnp.int32(0), template) sv = _compile_and_run_circuit_jit(stato_zero, patched_ops) if noise is not None: sv = NoiseModel.apply_to_sv( sv, n_qubits, model=noise.model, p=noise.p, jax_key=noise.jax_key, qubits=list(noise.qubits) if noise.qubits else None, ) energy = jnp.real(jnp.vdot(sv, h_matrix @ sv)) return energy, sv return energy_fn, n_params