sbs

Compatibility imports for the original sBs circuit API.

SBSNoise dataclass

Coherence parameters in the same time unit as the sequence.

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

@dataclass(frozen=True)
class SBSNoise:
    """Coherence parameters in the same time unit as the sequence."""

    oscillator_t1: float | None = None
    oscillator_tphi: float | None = None
    oscillator_nbar: float = 0.0
    qubit_t1: float | None = None
    qubit_t1_cd: float | Sequence[float] | None = None
    qubit_tphi: float | None = None
    qubit_excited_population: float = 0.0
    reset_error: float = 0.0
    reset_chi: float = 0.0
    qubit_cd_excited_population: float | Sequence[float] | None = None

SBSProtocol

Bases: NamedTuple

sBs rounds with an optional alternating sequence.

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

class SBSProtocol(NamedTuple):
    """sBs rounds with an optional alternating sequence."""

    rounds: tuple[SBSHalfRound, ...]
    alternate_rounds: tuple[SBSHalfRound, ...] | None = None

apply_sbs_half_round(joint, ops)

Apply one prepared sBs half-round to a joint density matrix.

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

def apply_sbs_half_round(joint, ops: SBSHalfRound):
    """Apply one prepared sBs half-round to a joint density matrix."""

    def segment(index, state):
        state = _apply_qubit_unitary(state, ops.rotations[index])
        state = _apply_noise(state, ops.rotation_noise, index)

        def apply_cd(inverse):
            def step(_, value):
                value = _apply_noisy_cd(value, ops.cd, index, inverse=inverse)
                return _apply_noise(value, ops.cd_noise, index)

            return step

        state = jax.lax.fori_loop(0, ops.microsteps, apply_cd(False), state)
        state = _apply_qubit_unitary(state, ops.echoes[index])
        return jax.lax.fori_loop(0, ops.microsteps, apply_cd(True), state)

    if jax.default_backend() == "cpu":
        for index in range(3):
            joint = segment(index, joint)
    else:
        joint = jax.lax.fori_loop(0, 3, segment, joint)

    joint = _apply_qubit_unitary(joint, ops.rotations[3])
    joint = _apply_noise(joint, ops.rotation_noise, 3)
    joint = _apply_reset(joint, ops.reset)
    return _apply_noise(joint, ops.reset_noise, 0)

build_sbs_cd_geometry(dimension, displacements, *, microsteps=1, jump_samples=4)

Precompute the displacement matrices shared by noise variants.

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

def build_sbs_cd_geometry(
    dimension,
    displacements,
    *,
    microsteps=1,
    jump_samples=4,
):
    """Precompute the displacement matrices shared by noise variants."""
    if microsteps < 1 or jump_samples < 1:
        raise ValueError("microsteps and jump_samples must be positive")
    times = (jnp.arange(jump_samples) + 0.5) / jump_samples

    def build(beta):
        beta = beta / (2 * microsteps)
        displacement = displace(dimension, beta / 2).data
        jumps = jax.vmap(
            lambda time: displace(
                dimension,
                beta * (2 * time - 1) / 2,
            ).data
        )(times)
        return displacement, jumps

    displacement, jumps = zip(*(build(beta) for beta in displacements))
    return SBSCDGeometry(
        displacements=jnp.stack(displacement),
        jump_displacements=jnp.stack(jumps),
    )

build_sbs_half_round(dimension, displacements, rotations, cd_durations, rotation_durations, reset_duration, noise, *, echoes=None, microsteps=1, jump_samples=4, max_loss=8, max_reset=10, storage_placement='segment', extra_storage_duration=0.0, reset_qubit_duration=None, cd_geometry=None)

Build one three-ECD sBs half-round.

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

def build_sbs_half_round(
    dimension: int,
    displacements: Sequence[complex],
    rotations: Sequence[jax.Array],
    cd_durations: Sequence[float],
    rotation_durations: Sequence[float],
    reset_duration: float,
    noise: SBSNoise,
    *,
    echoes: Sequence[jax.Array] | None = None,
    microsteps: int = 1,
    jump_samples: int = 4,
    max_loss: int = 8,
    max_reset: int = 10,
    storage_placement: str = "segment",
    extra_storage_duration: float = 0.0,
    reset_qubit_duration: float | None = None,
    cd_geometry: SBSCDGeometry | None = None,
):
    """Build one three-ECD sBs half-round."""
    if len(displacements) != 3 or len(cd_durations) != 3:
        raise ValueError("sBs requires three displacements and CD durations")
    if len(rotations) != 4 or len(rotation_durations) != 4:
        raise ValueError("sBs requires four rotations and rotation durations")
    if microsteps < 1 or jump_samples < 1:
        raise ValueError("microsteps and jump_samples must be positive")
    if max_loss < 0:
        raise ValueError("max_loss must be non-negative")
    if max_reset < 2:
        raise ValueError("max_reset must be at least two")
    if storage_placement not in {"segment", "lumped"}:
        raise ValueError("storage_placement must be 'segment' or 'lumped'")

    cd_durations = jnp.asarray(cd_durations)
    rotation_durations = jnp.asarray(rotation_durations)
    if (
        bool(jnp.any(cd_durations < 0))
        or bool(jnp.any(rotation_durations < 0))
        or reset_duration < 0
        or extra_storage_duration < 0
    ):
        raise ValueError("durations must be nonnegative")
    reset_qubit_duration = (
        reset_duration if reset_qubit_duration is None else reset_qubit_duration
    )
    if reset_qubit_duration < 0:
        raise ValueError("durations must be nonnegative")

    echoes = (
        [Rx(jnp.pi).U.data] * 3
        if echoes is None
        else [getattr(echo, "data", echo) for echo in echoes]
    )
    if len(echoes) != 3:
        raise ValueError("echoes must contain three qubit rotations")

    cd_t1 = noise.qubit_t1 if noise.qubit_t1_cd is None else noise.qubit_t1_cd
    substep_durations = cd_durations / (2 * microsteps)
    if cd_geometry is None:
        cd_geometry = build_sbs_cd_geometry(
            dimension,
            displacements,
            microsteps=microsteps,
            jump_samples=jump_samples,
        )
    cd_population = (
        noise.qubit_excited_population
        if noise.qubit_cd_excited_population is None
        else noise.qubit_cd_excited_population
    )
    cd = _build_cd_ops(
        cd_geometry,
        substep_durations,
        cd_t1,
        cd_population,
    )

    if storage_placement == "segment":
        rotation_storage = rotation_durations
        cd_storage = substep_durations
        reset_storage = jnp.asarray([reset_duration + extra_storage_duration])
    else:
        rotation_storage = jnp.zeros_like(rotation_durations)
        cd_storage = jnp.zeros_like(substep_durations)
        reset_storage = jnp.asarray(
            [
                rotation_durations.sum()
                + cd_durations.sum()
                + reset_duration
                + extra_storage_duration
            ]
        )

    cd_noise = SBSNoise(
        oscillator_t1=noise.oscillator_t1,
        oscillator_tphi=noise.oscillator_tphi,
        oscillator_nbar=noise.oscillator_nbar,
        qubit_tphi=noise.qubit_tphi,
        qubit_excited_population=noise.qubit_excited_population,
    )
    transfer_factor, excited_factor = _dephasing_reset_factors(
        dimension,
        noise.reset_error,
        reset_duration,
        noise.reset_chi,
        max_reset,
    )
    return SBSHalfRound(
        rotations=jnp.stack(
            [getattr(rotation, "data", rotation) for rotation in rotations]
        ),
        cd=cd,
        echoes=jnp.stack(echoes),
        rotation_noise=_noise_ops(
            dimension,
            rotation_storage,
            rotation_durations,
            noise,
            max_loss,
        ),
        cd_noise=_noise_ops(
            dimension,
            cd_storage,
            substep_durations,
            cd_noise,
            max_loss,
        ),
        reset=SBSResetOps(
            transfer_factor,
            excited_factor,
            jnp.asarray(1, dtype=transfer_factor.dtype),
        ),
        reset_noise=_noise_ops(
            dimension,
            reset_storage,
            jnp.asarray([reset_qubit_duration]),
            noise,
            max_loss,
        ),
        microsteps=microsteps,
    )

noisy_cd_kraus(dimension, beta, duration, t1, excited_population=0.0, jump_samples=4)

Return a thermally damped CD with midpoint-sampled jump times.

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

@partial(jax.jit, static_argnames=("dimension", "jump_samples"))
def noisy_cd_kraus(
    dimension,
    beta,
    duration,
    t1,
    excited_population=0.0,
    jump_samples=4,
):
    """Return a thermally damped CD with midpoint-sampled jump times."""
    if jump_samples < 1:
        raise ValueError("jump_samples must be positive")
    probability = _probability(duration, t1)
    excited_population = jnp.asarray(excited_population)
    p_down = probability * (1 - excited_population)
    p_up = probability * excited_population
    ground = basis(2, 0).data
    excited = basis(2, 1).data
    Pg = jnp.outer(ground, ground.conj())
    Pe = jnp.outer(excited, excited.conj())
    identity = jnp.eye(dimension)

    no_jump = (
        jnp.kron(jnp.sqrt(1 - p_up) * Pg + jnp.sqrt(1 - p_down) * Pe, identity)
        @ CD(dimension, beta).U.data
    )
    times = (jnp.arange(jump_samples) + 0.5) / jump_samples
    displacements = jax.vmap(lambda time: CD(dimension, beta * (2 * time - 1)).U.data)(
        times
    )
    # JAXQuantum names |g><e| ``sigmap`` and |e><g| ``sigmam``.
    lower = jnp.kron(sigmap().data, identity)
    raise_ = jnp.kron(sigmam().data, identity)
    relaxation = jnp.sqrt(p_down / jump_samples) * jnp.einsum(
        "ij,kjl->kil", lower, displacements
    )
    excitation = jnp.sqrt(p_up / jump_samples) * jnp.einsum(
        "ij,kjl->kil", raise_, displacements
    )
    return jnp.concatenate((no_jump[None], relaxation, excitation))

oscillator_state(joint)

Trace the ancilla from a joint density matrix.

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

def oscillator_state(joint):
    """Trace the ancilla from a joint density matrix."""
    return jnp.trace(joint, axis1=-4, axis2=-2)

simulate_sbs(initial_states, observables, half_rounds, cycles)

Evolve batched oscillator states and retain only observables.

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

@partial(jax.jit, static_argnames=("cycles",))
def simulate_sbs(initial_states, observables, half_rounds, cycles):
    """Evolve batched oscillator states and retain only observables."""
    if cycles < 0:
        raise ValueError("cycles must be non-negative")
    if isinstance(half_rounds, SBSProtocol):
        rounds = half_rounds.rounds
        alternate_rounds = half_rounds.alternate_rounds
    else:
        rounds, alternate_rounds = half_rounds, None
    ground = jnp.array([[1.0, 0.0], [0.0, 0.0]])
    joint = jnp.einsum("ij,...mn->...imjn", ground, initial_states)

    def expectation(state):
        reduced = oscillator_state(state)
        return jnp.einsum("...ij,...ji->...", observables, reduced).real

    initial = expectation(joint)

    def apply_round(state, selected):
        for half_round in selected:
            state = apply_sbs_half_round(state, half_round)
        return state

    def step(state, _):
        state = apply_round(state, rounds)
        return state, expectation(state)

    if alternate_rounds is None:
        joint, values = jax.lax.scan(step, joint, None, length=cycles)
    else:

        def alternate_step(state, cycle):
            state = jax.lax.cond(
                cycle % 2 == 0,
                lambda value: apply_round(value, rounds),
                lambda value: apply_round(value, alternate_rounds),
                state,
            )
            return state, expectation(state)

        joint, values = jax.lax.scan(
            alternate_step,
            joint,
            jnp.arange(cycles),
        )
    return (
        oscillator_state(joint),
        jnp.concatenate(
            (initial[None], values),
        ),
        joint,
    )