device

Shared cat and GKP sBs device simulations.

cat_problem(dimension, nbar, kind='bit')

Return the two cat states and observable defining a contrast.

Source code in jaxquantum/circuits/library/sbs/device.py

def cat_problem(dimension, nbar, kind="bit"):
    """Return the two cat states and observable defining a contrast."""
    alpha = jnp.sqrt(nbar)
    plus = jqt.displace(dimension, alpha) @ jqt.basis(dimension, 0)
    minus = jqt.displace(dimension, -alpha) @ jqt.basis(dimension, 0)
    if kind == "bit":
        q = (jqt.destroy(dimension) + jqt.create(dimension)).data / jnp.sqrt(2)
        values, vectors = jnp.linalg.eigh(q)
        observable = (vectors * jnp.sign(values)) @ vectors.conj().T
        states = (plus, minus)
    elif kind == "phase":
        states = (jqt.unit(plus + minus), jqt.unit(plus - minus))
        observable = jnp.diag((-1.0) ** jnp.arange(dimension))
    else:
        raise ValueError("kind must be 'bit' or 'phase'")
    density = jnp.stack([(state @ state.dag()).data for state in states])
    return density, jnp.broadcast_to(observable, density.shape)

cat_protocol(dimension, nbar, *, delta=0.6, ratio=3.125, device=CAT_DEVICE, enabled=ERROR_CHANNELS, microsteps=1, jump_samples=4, max_loss=8, cd_geometry=None, alternate_cd_direction=False)

Build the nominal cat sBs measurement round.

Source code in jaxquantum/circuits/library/sbs/device.py

def cat_protocol(
    dimension,
    nbar,
    *,
    delta=0.6,
    ratio=3.125,
    device=CAT_DEVICE,
    enabled=ERROR_CHANNELS,
    microsteps=1,
    jump_samples=4,
    max_loss=8,
    cd_geometry=None,
    alternate_cd_direction=False,
):
    """Build the nominal cat sBs measurement round."""
    alpha = jnp.sqrt(nbar)
    small = jnp.pi * delta**2 / (4 * alpha)
    displacements = (small, -1j * jnp.pi / (2 * alpha), ratio * small)
    rotations = (
        Ry(jnp.pi / 2).U.data,
        Rx(-jnp.pi / 2).U.data,
        Rx(-jnp.pi / 2).U.data,
        Ry(jnp.pi / 2).U.data,
    )

    def build(values, geometry):
        return build_sbs_half_round(
            dimension,
            values,
            rotations,
            device.cd_durations,
            device.rotation_durations,
            device.reset_duration,
            _noise(device, set(enabled)),
            microsteps=microsteps,
            jump_samples=jump_samples,
            max_loss=max_loss,
            storage_placement="lumped",
            extra_storage_duration=device.extra_storage_duration,
            reset_qubit_duration=0.0,
            cd_geometry=geometry,
        )

    forward = build(displacements, cd_geometry)
    if not alternate_cd_direction:
        return (forward,)
    reverse_geometry = SBSCDGeometry(
        jnp.swapaxes(forward.cd.displacements.conj(), -1, -2),
        jnp.swapaxes(forward.cd.jump_displacements.conj(), -1, -2),
    )
    reverse = build(tuple(-value for value in displacements), reverse_geometry)
    return SBSProtocol((forward,), (reverse,))

compute_error_budget(protocol, initial_states, observables, cycles, cycle_time, *, channels=ERROR_CHANNELS, fit_start=4, fit_floor=1e-10, batched=True)

Return baseline-subtracted and all-on-context channel budgets.

Source code in jaxquantum/circuits/library/sbs/device.py

def compute_error_budget(
    protocol: Callable[[set[str]], tuple],
    initial_states,
    observables,
    cycles,
    cycle_time,
    *,
    channels=ERROR_CHANNELS,
    fit_start=4,
    fit_floor=1e-10,
    batched=True,
):
    """Return baseline-subtracted and all-on-context channel budgets."""
    channels = tuple(channels)

    def run(enabled):
        return simulate_decay(
            initial_states,
            observables,
            protocol(set(enabled)),
            cycles,
            cycle_time,
            fit_start=fit_start,
            fit_floor=fit_floor,
        )

    variants = [
        (),
        channels,
        *((channel,) for channel in channels),
        *(set(channels) - {channel} for channel in channels),
    ]
    if batched:
        results = simulate_decay_variants(
            initial_states,
            observables,
            [protocol(set(enabled)) for enabled in variants],
            cycles,
            cycle_time,
            fit_start=fit_start,
            fit_floor=fit_floor,
        )
    else:
        results = [run(enabled) for enabled in variants]
    baseline, all_on = results[:2]
    split = 2 + len(channels)
    isolated = dict(zip(channels, results[2:split]))
    without = dict(zip(channels, results[split:]))
    isolated_increments = {
        channel: isolated[channel].rate - baseline.rate for channel in channels
    }
    context_increments = {
        channel: all_on.rate - without[channel].rate for channel in channels
    }
    interaction = all_on.rate - baseline.rate - sum(isolated_increments.values())
    return ErrorBudget(
        baseline=baseline,
        all_on=all_on,
        isolated=isolated,
        without=without,
        isolated_increments=isolated_increments,
        context_increments=context_increments,
        interaction_rate=interaction,
    )

gkp_displacements(delta, small_ratio, small_displacement_scales, big_displacement, epsilon_model, length_scale)

Return complementary Z- and X-stabilizer displacements.

Source code in jaxquantum/circuits/library/sbs/device.py

def gkp_displacements(
    delta,
    small_ratio,
    small_displacement_scales,
    big_displacement,
    epsilon_model,
    length_scale,
):
    """Return complementary Z- and X-stabilizer displacements."""
    length = jnp.sqrt(2 * jnp.pi) * length_scale
    if epsilon_model == "sinh":
        epsilon = jnp.sinh(delta**2) * length
    elif epsilon_model == "quadratic":
        epsilon = delta**2 * length
    else:
        raise ValueError("epsilon_model must be 'sinh' or 'quadratic'")
    scales = jnp.asarray(small_displacement_scales)
    small = (epsilon / 2 * scales[0], small_ratio * epsilon / 2 * scales[1])
    big = length if big_displacement is None else big_displacement
    return (
        (small[0], -1j * big, small[1]),
        (1j * small[0], big, 1j * small[1]),
    )

gkp_problem(dimension, delta=0.428, kind='x')

Return the two GKP states and observable defining a logical contrast.

Source code in jaxquantum/circuits/library/sbs/device.py

def gkp_problem(dimension, delta=0.428, kind="x"):
    """Return the two GKP states and observable defining a logical contrast."""
    kind = kind.lower()
    if kind not in ("x", "z"):
        raise ValueError("kind must be 'x' or 'z'")
    code = jqcodes.GKPQubit({"delta": delta, "N": dimension})
    states = (code.basis[f"-{kind}"], code.basis[f"+{kind}"])
    density = jnp.stack([(state @ state.dag()).data for state in states])
    observable = code.common_gates[f"{kind.upper()}_0"].data
    return density, jnp.broadcast_to(observable, density.shape)

gkp_protocol(dimension, *, delta=0.428, small_ratio=1.083, small_displacement_scales=(1.0, 1.0), big_displacement=None, epsilon_model='sinh', final_storage_rotation=0.0, alternate_cd_direction=False, length_scale=1.0, device=GKP_DEVICE, enabled=ERROR_CHANNELS, microsteps=1, jump_samples=4, max_loss=8, max_reset=12, cd_geometries=None)

Build the two complementary GKP sBs half-rounds.

Source code in jaxquantum/circuits/library/sbs/device.py

def gkp_protocol(
    dimension,
    *,
    delta=0.428,
    small_ratio=1.083,
    small_displacement_scales=(1.0, 1.0),
    big_displacement=None,
    epsilon_model="sinh",
    final_storage_rotation=0.0,
    alternate_cd_direction=False,
    length_scale=1.0,
    device=GKP_DEVICE,
    enabled=ERROR_CHANNELS,
    microsteps=1,
    jump_samples=4,
    max_loss=8,
    max_reset=12,
    cd_geometries=None,
):
    """Build the two complementary GKP sBs half-rounds."""
    z_displacements, x_displacements = gkp_displacements(
        delta,
        small_ratio,
        small_displacement_scales,
        big_displacement,
        epsilon_model,
        length_scale,
    )
    rotations = (
        Ry(jnp.pi / 2).U.data,
        Rx(-jnp.pi / 2).U.data,
        Rx(jnp.pi / 2).U.data,
        Ry(-jnp.pi / 2).U.data,
    )
    noise = _noise(device, set(enabled))
    if cd_geometries is None:
        cd_geometries = (None, None)

    def build(displacements, geometry):
        half_round = build_sbs_half_round(
            dimension,
            displacements,
            rotations,
            device.cd_durations,
            device.rotation_durations,
            device.reset_duration,
            noise,
            microsteps=microsteps,
            jump_samples=jump_samples,
            max_loss=max_loss,
            max_reset=max_reset,
            storage_placement="segment",
            extra_storage_duration=device.extra_storage_duration,
            cd_geometry=geometry,
        )
        if final_storage_rotation:
            phase = jnp.exp(
                -1j * final_storage_rotation * jnp.arange(dimension)
            )
            phase_factor = phase[:, None] * phase.conj()[None, :]
            half_round = half_round._replace(
                reset=half_round.reset._replace(
                    phase_factor=phase_factor * half_round.reset.phase_factor,
                )
            )
        return half_round

    forward = (
        build(z_displacements, cd_geometries[0]),
        build(x_displacements, cd_geometries[1]),
    )
    if not alternate_cd_direction:
        return forward
    reverse_geometries = tuple(
        SBSCDGeometry(
            jnp.swapaxes(half_round.cd.displacements.conj(), -1, -2),
            jnp.swapaxes(half_round.cd.jump_displacements.conj(), -1, -2),
        )
        for half_round in forward
    )
    reverse = tuple(
        build(
            tuple(-value for value in displacements),
            geometry,
        )
        for displacements, geometry in zip(
            (z_displacements, x_displacements),
            reverse_geometries,
        )
    )
    return SBSProtocol(forward, reverse)

prepare_cat_protocol(dimension, nbar, *, delta=0.6, ratio=3.125, device=CAT_DEVICE, microsteps=1, jump_samples=4, max_loss=8, alternate_cd_direction=False)

Return an error-channel builder with shared pulse geometry.

Source code in jaxquantum/circuits/library/sbs/device.py

def prepare_cat_protocol(
    dimension,
    nbar,
    *,
    delta=0.6,
    ratio=3.125,
    device=CAT_DEVICE,
    microsteps=1,
    jump_samples=4,
    max_loss=8,
    alternate_cd_direction=False,
):
    """Return an error-channel builder with shared pulse geometry."""
    alpha = jnp.sqrt(nbar)
    small = jnp.pi * delta**2 / (4 * alpha)
    displacements = (small, -1j * jnp.pi / (2 * alpha), ratio * small)
    geometry = build_sbs_cd_geometry(
        dimension,
        displacements,
        microsteps=microsteps,
        jump_samples=jump_samples,
    )

    def build(enabled):
        return cat_protocol(
            dimension,
            nbar,
            delta=delta,
            ratio=ratio,
            device=device,
            enabled=enabled,
            microsteps=microsteps,
            jump_samples=jump_samples,
            max_loss=max_loss,
            cd_geometry=geometry,
            alternate_cd_direction=alternate_cd_direction,
        )

    return build

prepare_gkp_protocol(dimension, *, delta=0.428, small_ratio=1.083, small_displacement_scales=(1.0, 1.0), big_displacement=None, epsilon_model='sinh', final_storage_rotation=0.0, alternate_cd_direction=False, length_scale=1.0, device=GKP_DEVICE, microsteps=1, jump_samples=4, max_loss=8, max_reset=12)

Return an error-channel builder with shared pulse geometry.

Source code in jaxquantum/circuits/library/sbs/device.py

def prepare_gkp_protocol(
    dimension,
    *,
    delta=0.428,
    small_ratio=1.083,
    small_displacement_scales=(1.0, 1.0),
    big_displacement=None,
    epsilon_model="sinh",
    final_storage_rotation=0.0,
    alternate_cd_direction=False,
    length_scale=1.0,
    device=GKP_DEVICE,
    microsteps=1,
    jump_samples=4,
    max_loss=8,
    max_reset=12,
):
    """Return an error-channel builder with shared pulse geometry."""
    displacements = gkp_displacements(
        delta,
        small_ratio,
        small_displacement_scales,
        big_displacement,
        epsilon_model,
        length_scale,
    )
    geometries = tuple(
        build_sbs_cd_geometry(
            dimension,
            values,
            microsteps=microsteps,
            jump_samples=jump_samples,
        )
        for values in displacements
    )

    def build(enabled):
        return gkp_protocol(
            dimension,
            delta=delta,
            small_ratio=small_ratio,
            small_displacement_scales=small_displacement_scales,
            big_displacement=big_displacement,
            epsilon_model=epsilon_model,
            final_storage_rotation=final_storage_rotation,
            alternate_cd_direction=alternate_cd_direction,
            length_scale=length_scale,
            device=device,
            enabled=enabled,
            microsteps=microsteps,
            jump_samples=jump_samples,
            max_loss=max_loss,
            max_reset=max_reset,
            cd_geometries=geometries,
        )

    return build

simulate_decay_variants(initial_states, observables, protocols, cycles, cycle_time, *, fit_start=4, fit_floor=1e-10)

Simulate equal-structure protocol variants in one mapped call.

Source code in jaxquantum/circuits/library/sbs/device.py

def simulate_decay_variants(
    initial_states,
    observables,
    protocols,
    cycles,
    cycle_time,
    *,
    fit_start=4,
    fit_floor=1e-10,
):
    """Simulate equal-structure protocol variants in one mapped call."""
    if len(protocols) == 1:
        return [
            simulate_decay(
                initial_states,
                observables,
                protocols[0],
                cycles,
                cycle_time,
                fit_start=fit_start,
                fit_floor=fit_floor,
            )
        ]
    protocols, axes = _stack_protocols(protocols)
    final, values, _ = jax.vmap(
        lambda rounds: simulate_sbs(
            initial_states,
            observables,
            rounds,
            cycles,
        ),
        in_axes=(axes,),
    )(protocols)
    final, values = jax.device_get((final, values))
    return [
        _analyze_decay(
            result,
            samples,
            cycles,
            cycle_time,
            fit_start,
            fit_floor,
        )
        for result, samples in zip(final, values)
    ]