Skip to content

Gate metrics: rlquantopt/jx/metrics.py

The reward and GRAPE both use the paper's cost (eq. 2):

\[ J_T = 1 - \frac{w_C\, C + w_U\, U}{w_C + w_U}, \qquad w_C = 1,\; w_U = 3, \]

where \(C\) is the gate concurrence (1 for a perfect entangler) and \(U = \operatorname{Tr}(G^\dagger G)/4\) is the unitarity (1 when nothing leaks out of the computational subspace). The step reward is \(-\log_{10} J_T\), so a reward of 4 means \(J_T = 10^{-4}\).

Weyl-chamber coordinates in closed form

The concurrence comes from the Weyl-chamber coordinates \((c_1, c_2, c_3)\), which v1 computes with weylchamber.c1c2c3 (Childs et al., PRA 68, 052311). That algorithm needs the eigenvalues of \(U \tilde U / \sqrt{\det U}\) with \(\tilde U = (\sigma_y\otimes\sigma_y) U^T (\sigma_y\otimes\sigma_y)\). A general 4×4 eig is not available on GPU in JAX, and it is awkward to differentiate.

Our gate is block diagonal, \(G = \operatorname{diag}(g_{00}, V, e)\) with \(V = \begin{pmatrix} a & b \\ c & d \end{pmatrix}\). Working out \(\tilde G\) gives \(\operatorname{diag}\big(e, \begin{pmatrix} d & b \\ c & a\end{pmatrix}, g_{00}\big)\), so

\[ G\tilde G = \operatorname{diag}\Big(g_{00} e,\; \begin{pmatrix} ad+bc & 2ab \\ 2cd & ad+bc \end{pmatrix},\; g_{00} e\Big), \]

whose eigenvalues are \(g_{00} e\) (twice) and \(ad + bc \pm 2\sqrt{abcd}\). The rest of the algorithm (phases, sorting, the \(n = \mathrm{round}(\sum S)\) shift, mapping into the chamber) follows weylchamber line by line:

def weyl_eigenvalues(G):
    g00, a, b, c, d, e = G[0, 0], G[1, 1], G[1, 2], G[2, 1], G[2, 2], G[3, 3]
    p = a * d + b * c
    r = 2 * jnp.sqrt(a * b * c * d)
    ev = jnp.stack([g00 * e, p + r, p - r, g00 * e])
    det = g00 * e * (a * d - b * c)
    return ev / jnp.sqrt(det)


def c1c2c3(G, digits=WEYL_DIGITS):
    """Weyl chamber coordinates (units of π) of a block-diagonal 4x4 gate, as weylchamber.c1c2c3.

    ``digits=None`` skips the final rounding, which otherwise has zero gradient (use it for GRAPE).
    """
    two_S = jnp.angle(weyl_eigenvalues(G)) / jnp.pi
    two_S = jnp.where(two_S <= -0.5, two_S + 2.0, two_S)
    S = jnp.sort(two_S / 2.0)[::-1]
    n = jnp.round(jnp.sum(S)).astype(int)
    S = S - (jnp.arange(4) < n)
    S = jnp.roll(S, -n)
    c1, c2, c3 = _M @ S[:3]
    c1 = jnp.where(c3 < 0, 1 - c1, c1)
    c3 = jnp.abs(c3)
    c = jnp.stack([c1, c2, c3])
    return c if digits is None else jnp.round(c, digits) + 0.0
  • digits=8 reproduces weylchamber's rounding exactly. That rounding has zero gradient, so GRAPE calls it with digits=None.
  • \(\sqrt{abcd}\) has an infinite slope when an off-diagonal element is exactly 0 (for example at \(u = 0\)); GRAPE replaces the resulting NaN gradients with 0.
  • Inside the perfect-entangler polyhedron the concurrence is exactly 1 and flat, so the gradient there comes only from the unitarity.

Concurrence, unitarity and J_T

def concurrence(c):
    """Gate concurrence from Weyl coordinates, as weylchamber.concurrence."""
    c1, c2, c3 = c
    in_pe = (c1 + c2 >= 0.5) & (c1 - c2 <= 0.5) & (c2 + c3 <= 0.5)
    rolled = jnp.roll(c, 1)
    m = jnp.concatenate([c - rolled, c + rolled])
    return jnp.where(in_pe, 1.0, jnp.max(jnp.abs(jnp.sin(jnp.pi * m))))


def unitarity(G):
    """U = Tr(G†G)/4: 1 when no population leaks out of the computational subspace."""
    return jnp.sum(jnp.abs(G) ** 2) / 4


def cost_JT(G, concurrence_weight=1.0, unitarity_weight=3.0, digits=WEYL_DIGITS):
    """J_T = 1 - (w_c C + w_u U)/(w_c + w_u); paper eq. 2 with the default 1:3 weights."""
    C = concurrence(c1c2c3(G, digits))
    U = unitarity(G)
    K = concurrence_weight + unitarity_weight
    return 1 - (concurrence_weight * C + unitarity_weight * U) / K, C, U

Why the concurrence can look noisy

When \(|b|, |c| \ll 1\) (early in a pulse), the eigenvalue phases are ill-conditioned and the concurrence changes by ~1e-5 for tiny changes of the state. The v1-equivalence test therefore compares concurrence at 3e-4 while amplitudes are compared at 1e-5. On exact gates the metric agrees with weylchamber to 1e-8.

The perfect-entangler target and its named-gate alternatives are discussed in Named gates and a richer model.

Named-gate fidelity with free Z corrections

Option B of Named gates, implemented for √iSWAP (idea i07). Virtual Z rotations cost nothing on hardware, so the fidelity is taken after the best Z corrections on each qubit before and after the gate. For the block-diagonal gates of this model only three combinations of the four phases change \(|\mathrm{Tr}\,M|\); they are found by a 10³ grid search and eight modified-Newton steps. The first term of the average gate fidelity is the unitarity, so leakage is penalised without a separate term.

\[ F = \max_{\text{Z phases}} \frac{\operatorname{Tr}(M M^\dagger) + |\operatorname{Tr} M|^2}{20}, \qquad M = V^\dagger D_\text{after}\, P\, D_\text{before} \]
# Named-gate fidelity (idea i06). Target in the basis |00>, |01>, |10>, |11> (|q0 q1>).
SQRT_ISWAP = jnp.array([[1, 0, 0, 0],
                        [0, 1 / jnp.sqrt(2), 1j / jnp.sqrt(2), 0],
                        [0, 1j / jnp.sqrt(2), 1 / jnp.sqrt(2), 0],
                        [0, 0, 0, 1]], dtype=jnp.complex128)
_GRID = jnp.stack(jnp.meshgrid(*[jnp.linspace(-jnp.pi, jnp.pi, 10, endpoint=False)] * 3, indexing="ij"), -1).reshape(-1, 3)


def physical_gate(G):
    """<b_i|U|b_j> from ``physics.realised_gate``, which follows v1 and returns its conjugate transpose."""
    return G.conj().T


def _z_trace(P, V, th):
    """Tr(V† D_after P D_before) for block-diagonal P and V, as a function of the three phase combinations
    (x, y, z) = (a1 + b1, a1 + b0, a0 + b1) that the Z corrections D = diag(1, e^{i a1}, e^{i a0}, e^{i(a0 + a1)})
    before (b) and after (a) the gate can set; the global phase does not enter |Tr|."""
    t = jnp.conj(V) * P
    x, y, z = th[..., 0], th[..., 1], th[..., 2]
    e = lambda p: jnp.exp(1j * p)
    return (t[0, 0] + t[1, 1] * e(x) + t[1, 2] * e(y) + t[2, 1] * e(z) + t[2, 2] * e(y + z - x)
            + t[3, 3] * e(y + z))


def best_z_phases(P, V):
    """Z-correction phases maximising |Tr|: best point of a 10^3 grid, then 8 modified-Newton steps."""
    f = lambda th: jnp.abs(_z_trace(P, V, th)) ** 2
    th = _GRID[jnp.argmax(jax.vmap(f)(_GRID))]

    def newton(th, _):
        g, H = jax.grad(f)(th), jax.hessian(f)(th)
        w, Q = jnp.linalg.eigh(H)
        w = jnp.minimum(w, -1e-6 * (1 + jnp.abs(w).max()))     # ascent direction even off the concave region
        return th - Q @ ((Q.T @ g) / w), None

    th, _ = jax.lax.scan(newton, th, None, 8)
    return th


def fidelity_free_z(G, V=SQRT_ISWAP):
    """Average gate fidelity to V after the best virtual-Z corrections, leakage included
    (Pedersen et al., Phys. Lett. A 367, 47 (2007)): F = (Tr(M M†) + |Tr M|^2) / 20, M = V† D_a P D_b.
    The phases are optimised on a stopped-gradient copy: at the optimum dF/dphase = 0, so the gradient
    with respect to the pulse is exact. Returns (F, U) with U the unitarity Tr(P†P)/4."""
    P = physical_gate(G)
    th = jax.lax.stop_gradient(best_z_phases(jax.lax.stop_gradient(P), V))
    U = jnp.sum(jnp.abs(P) ** 2) / 4
    F = (4 * U + jnp.abs(_z_trace(P, V, th)) ** 2) / 20
    return F, U


def cost(G, objective="pe", concurrence_weight=1.0, unitarity_weight=3.0, digits=WEYL_DIGITS):
    """(cost, C, U) for the objective: ``pe`` = the paper's J_T; ``sqrt_iswap`` = 1 - F to sqrt(iSWAP), free Z."""
    if objective == "pe":
        return cost_JT(G, concurrence_weight, unitarity_weight, digits)
    if objective == "sqrt_iswap":
        F, U = fidelity_free_z(G, SQRT_ISWAP)
        return 1 - F, concurrence(c1c2c3(G, digits)), U
    raise ValueError(f"unknown objective {objective!r}")

tests/test_jx_fidelity.py checks \(F\) against a brute-force optimisation over the four phases (agreement to 1e-9), unit fidelity for √iSWAP and √iSWAP† under arbitrary Z phases, and the pulse gradient against finite differences. The pulses optimised for the perfect-entangler objective sit 0.6-1.8 % away from √iSWAP: they land near Weyl point (0.25, 0.25, 0.06-0.09), √iSWAP plus a small conditional phase.