Coverage for jaxquantum/circuits/library/sbs/core.py: 97%
245 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Functional sBs circuit primitives."""
3from __future__ import annotations
5from dataclasses import dataclass
6from functools import partial
7from typing import NamedTuple, Sequence
9import jax
10import jax.numpy as jnp
12from jaxquantum.circuits.channels import apply_kraus_map, apply_shifted_channel
13from jaxquantum.circuits.library.oscillator import (
14 CD,
15 _amp_damp_coefficients,
16 _apply_dephasing_reset_factors,
17 _dephasing_reset_factors,
18 _thermal_coefficients,
19)
20from jaxquantum.circuits.library.qubit import Rx
21from jaxquantum.core.operators import basis, displace, sigmam, sigmap
23__all__ = (
24 "SBSNoise",
25 "SBSNoiseOps",
26 "SBSCDGeometry",
27 "SBSCDOps",
28 "SBSResetOps",
29 "SBSHalfRound",
30 "SBSProtocol",
31 "noisy_cd_kraus",
32 "build_sbs_cd_geometry",
33 "build_sbs_half_round",
34 "apply_sbs_half_round",
35 "oscillator_state",
36 "simulate_sbs",
37)
40@dataclass(frozen=True)
41class SBSNoise:
42 """Coherence parameters in the same time unit as the sequence."""
44 oscillator_t1: float | None = None
45 oscillator_tphi: float | None = None
46 oscillator_nbar: float = 0.0
47 qubit_t1: float | None = None
48 qubit_t1_cd: float | Sequence[float] | None = None
49 qubit_tphi: float | None = None
50 qubit_excited_population: float = 0.0
51 reset_error: float = 0.0
52 reset_chi: float = 0.0
53 qubit_cd_excited_population: float | Sequence[float] | None = None
56class SBSNoiseOps(NamedTuple):
57 oscillator_coefficients: jax.Array
58 oscillator_shifts: jax.Array
59 oscillator_dephasing: jax.Array
60 qubit_probability: jax.Array
61 qubit_excited_population: jax.Array
62 qubit_dephasing: jax.Array
65class SBSCDGeometry(NamedTuple):
66 displacements: jax.Array
67 jump_displacements: jax.Array
70class SBSCDOps(NamedTuple):
71 displacements: jax.Array
72 jump_displacements: jax.Array
73 relaxation_probability: jax.Array
74 excitation_probability: jax.Array
77class SBSResetOps(NamedTuple):
78 transfer_factor: jax.Array
79 excited_factor: jax.Array
80 phase_factor: jax.Array
83class SBSHalfRound(NamedTuple):
84 rotations: jax.Array
85 cd: SBSCDOps
86 echoes: jax.Array
87 rotation_noise: SBSNoiseOps
88 cd_noise: SBSNoiseOps
89 reset: SBSResetOps
90 reset_noise: SBSNoiseOps
91 microsteps: int
94class SBSProtocol(NamedTuple):
95 """sBs rounds with an optional alternating sequence."""
97 rounds: tuple[SBSHalfRound, ...]
98 alternate_rounds: tuple[SBSHalfRound, ...] | None = None
101def _lifetime(value):
102 return jnp.inf if value is None else jnp.asarray(value)
105def _probability(duration, lifetime):
106 return -jnp.expm1(-jnp.asarray(duration) / _lifetime(lifetime))
109def _noise_ops(
110 dimension,
111 oscillator_durations,
112 qubit_durations,
113 noise,
114 max_loss,
115):
116 oscillator_durations = jnp.atleast_1d(jnp.asarray(oscillator_durations))
117 qubit_durations = jnp.atleast_1d(jnp.asarray(qubit_durations))
118 if oscillator_durations.shape != qubit_durations.shape:
119 raise ValueError("oscillator and qubit durations must have equal shapes")
121 loss_probability = _probability(
122 oscillator_durations,
123 noise.oscillator_t1,
124 )
125 if noise.oscillator_nbar == 0:
126 coefficients = _amp_damp_coefficients(
127 dimension,
128 loss_probability,
129 max_loss,
130 )
131 shifts = jnp.arange(coefficients.shape[-2])
132 else:
133 coefficients, gains, losses = _thermal_coefficients(
134 dimension,
135 loss_probability,
136 noise.oscillator_nbar,
137 max_loss,
138 )
139 shifts = losses - gains
140 indices = jnp.arange(dimension)
141 delta = indices[:, None] - indices[None, :]
142 oscillator_dephasing = jnp.exp(
143 -oscillator_durations[:, None, None]
144 / _lifetime(noise.oscillator_tphi)
145 * delta**2
146 )
147 return SBSNoiseOps(
148 oscillator_coefficients=coefficients,
149 oscillator_shifts=shifts,
150 oscillator_dephasing=oscillator_dephasing,
151 qubit_probability=_probability(qubit_durations, noise.qubit_t1),
152 qubit_excited_population=jnp.broadcast_to(
153 jnp.asarray(noise.qubit_excited_population),
154 qubit_durations.shape,
155 ),
156 qubit_dephasing=jnp.exp(-qubit_durations / _lifetime(noise.qubit_tphi)),
157 )
160@partial(jax.jit, static_argnames=("dimension", "jump_samples"))
161def noisy_cd_kraus(
162 dimension,
163 beta,
164 duration,
165 t1,
166 excited_population=0.0,
167 jump_samples=4,
168):
169 """Return a thermally damped CD with midpoint-sampled jump times."""
170 if jump_samples < 1:
171 raise ValueError("jump_samples must be positive")
172 probability = _probability(duration, t1)
173 excited_population = jnp.asarray(excited_population)
174 p_down = probability * (1 - excited_population)
175 p_up = probability * excited_population
176 ground = basis(2, 0).data
177 excited = basis(2, 1).data
178 Pg = jnp.outer(ground, ground.conj())
179 Pe = jnp.outer(excited, excited.conj())
180 identity = jnp.eye(dimension)
182 no_jump = (
183 jnp.kron(jnp.sqrt(1 - p_up) * Pg + jnp.sqrt(1 - p_down) * Pe, identity)
184 @ CD(dimension, beta).U.data
185 )
186 times = (jnp.arange(jump_samples) + 0.5) / jump_samples
187 displacements = jax.vmap(lambda time: CD(dimension, beta * (2 * time - 1)).U.data)(
188 times
189 )
190 # JAXQuantum names |g><e| ``sigmap`` and |e><g| ``sigmam``.
191 lower = jnp.kron(sigmap().data, identity)
192 raise_ = jnp.kron(sigmam().data, identity)
193 relaxation = jnp.sqrt(p_down / jump_samples) * jnp.einsum(
194 "ij,kjl->kil", lower, displacements
195 )
196 excitation = jnp.sqrt(p_up / jump_samples) * jnp.einsum(
197 "ij,kjl->kil", raise_, displacements
198 )
199 return jnp.concatenate((no_jump[None], relaxation, excitation))
202def build_sbs_cd_geometry(
203 dimension,
204 displacements,
205 *,
206 microsteps=1,
207 jump_samples=4,
208):
209 """Precompute the displacement matrices shared by noise variants."""
210 if microsteps < 1 or jump_samples < 1:
211 raise ValueError("microsteps and jump_samples must be positive")
212 times = (jnp.arange(jump_samples) + 0.5) / jump_samples
214 def build(beta):
215 beta = beta / (2 * microsteps)
216 displacement = displace(dimension, beta / 2).data
217 jumps = jax.vmap(
218 lambda time: displace(
219 dimension,
220 beta * (2 * time - 1) / 2,
221 ).data
222 )(times)
223 return displacement, jumps
225 displacement, jumps = zip(*(build(beta) for beta in displacements))
226 return SBSCDGeometry(
227 displacements=jnp.stack(displacement),
228 jump_displacements=jnp.stack(jumps),
229 )
232def _build_cd_ops(geometry, durations, t1, excited_population):
233 probabilities = _probability(durations, t1)
234 excited_population = jnp.asarray(excited_population)
235 return SBSCDOps(
236 displacements=geometry.displacements,
237 jump_displacements=geometry.jump_displacements,
238 relaxation_probability=probabilities * (1 - excited_population),
239 excitation_probability=probabilities * excited_population,
240 )
243def build_sbs_half_round(
244 dimension: int,
245 displacements: Sequence[complex],
246 rotations: Sequence[jax.Array],
247 cd_durations: Sequence[float],
248 rotation_durations: Sequence[float],
249 reset_duration: float,
250 noise: SBSNoise,
251 *,
252 echoes: Sequence[jax.Array] | None = None,
253 microsteps: int = 1,
254 jump_samples: int = 4,
255 max_loss: int = 8,
256 max_reset: int = 10,
257 storage_placement: str = "segment",
258 extra_storage_duration: float = 0.0,
259 reset_qubit_duration: float | None = None,
260 cd_geometry: SBSCDGeometry | None = None,
261):
262 """Build one three-ECD sBs half-round."""
263 if len(displacements) != 3 or len(cd_durations) != 3:
264 raise ValueError("sBs requires three displacements and CD durations")
265 if len(rotations) != 4 or len(rotation_durations) != 4:
266 raise ValueError("sBs requires four rotations and rotation durations")
267 if microsteps < 1 or jump_samples < 1:
268 raise ValueError("microsteps and jump_samples must be positive")
269 if max_loss < 0:
270 raise ValueError("max_loss must be non-negative")
271 if max_reset < 2:
272 raise ValueError("max_reset must be at least two")
273 if storage_placement not in {"segment", "lumped"}:
274 raise ValueError("storage_placement must be 'segment' or 'lumped'")
276 cd_durations = jnp.asarray(cd_durations)
277 rotation_durations = jnp.asarray(rotation_durations)
278 if (
279 bool(jnp.any(cd_durations < 0))
280 or bool(jnp.any(rotation_durations < 0))
281 or reset_duration < 0
282 or extra_storage_duration < 0
283 ):
284 raise ValueError("durations must be nonnegative")
285 reset_qubit_duration = (
286 reset_duration if reset_qubit_duration is None else reset_qubit_duration
287 )
288 if reset_qubit_duration < 0:
289 raise ValueError("durations must be nonnegative")
291 echoes = (
292 [Rx(jnp.pi).U.data] * 3
293 if echoes is None
294 else [getattr(echo, "data", echo) for echo in echoes]
295 )
296 if len(echoes) != 3:
297 raise ValueError("echoes must contain three qubit rotations")
299 cd_t1 = noise.qubit_t1 if noise.qubit_t1_cd is None else noise.qubit_t1_cd
300 substep_durations = cd_durations / (2 * microsteps)
301 if cd_geometry is None:
302 cd_geometry = build_sbs_cd_geometry(
303 dimension,
304 displacements,
305 microsteps=microsteps,
306 jump_samples=jump_samples,
307 )
308 cd_population = (
309 noise.qubit_excited_population
310 if noise.qubit_cd_excited_population is None
311 else noise.qubit_cd_excited_population
312 )
313 cd = _build_cd_ops(
314 cd_geometry,
315 substep_durations,
316 cd_t1,
317 cd_population,
318 )
320 if storage_placement == "segment":
321 rotation_storage = rotation_durations
322 cd_storage = substep_durations
323 reset_storage = jnp.asarray([reset_duration + extra_storage_duration])
324 else:
325 rotation_storage = jnp.zeros_like(rotation_durations)
326 cd_storage = jnp.zeros_like(substep_durations)
327 reset_storage = jnp.asarray(
328 [
329 rotation_durations.sum()
330 + cd_durations.sum()
331 + reset_duration
332 + extra_storage_duration
333 ]
334 )
336 cd_noise = SBSNoise(
337 oscillator_t1=noise.oscillator_t1,
338 oscillator_tphi=noise.oscillator_tphi,
339 oscillator_nbar=noise.oscillator_nbar,
340 qubit_tphi=noise.qubit_tphi,
341 qubit_excited_population=noise.qubit_excited_population,
342 )
343 transfer_factor, excited_factor = _dephasing_reset_factors(
344 dimension,
345 noise.reset_error,
346 reset_duration,
347 noise.reset_chi,
348 max_reset,
349 )
350 return SBSHalfRound(
351 rotations=jnp.stack(
352 [getattr(rotation, "data", rotation) for rotation in rotations]
353 ),
354 cd=cd,
355 echoes=jnp.stack(echoes),
356 rotation_noise=_noise_ops(
357 dimension,
358 rotation_storage,
359 rotation_durations,
360 noise,
361 max_loss,
362 ),
363 cd_noise=_noise_ops(
364 dimension,
365 cd_storage,
366 substep_durations,
367 cd_noise,
368 max_loss,
369 ),
370 reset=SBSResetOps(
371 transfer_factor,
372 excited_factor,
373 jnp.asarray(1, dtype=transfer_factor.dtype),
374 ),
375 reset_noise=_noise_ops(
376 dimension,
377 reset_storage,
378 jnp.asarray([reset_qubit_duration]),
379 noise,
380 max_loss,
381 ),
382 microsteps=microsteps,
383 )
386def _apply_joint_kraus(joint, kraus):
387 dimension = joint.shape[-1]
388 flat = joint.reshape(joint.shape[:-4] + (2 * dimension,) * 2)
389 flat = apply_kraus_map(kraus, flat)
390 return flat.reshape(flat.shape[:-2] + (2, dimension, 2, dimension))
393def _apply_reset(joint, reset):
394 dimension = joint.shape[-1]
395 flat = joint.reshape(joint.shape[:-4] + (2 * dimension,) * 2)
396 flat = _apply_dephasing_reset_factors(
397 flat,
398 reset.transfer_factor,
399 reset.excited_factor,
400 dimension,
401 reset.phase_factor,
402 )
403 return flat.reshape(flat.shape[:-2] + (2, dimension, 2, dimension))
406def _apply_noisy_cd(joint, cd: SBSCDOps, index, inverse=False):
407 displacement = cd.displacements[index]
408 jumps = cd.jump_displacements[index]
409 if inverse:
410 displacement = jnp.swapaxes(displacement.conj(), -1, -2)
411 jumps = jnp.swapaxes(jumps.conj(), -1, -2)
412 conditional = jnp.stack((displacement, jnp.swapaxes(displacement.conj(), -1, -2)))
413 no_jump = jnp.einsum(
414 "ami,...aibj,bnj->...ambn",
415 conditional,
416 joint,
417 conditional.conj(),
418 )
419 relaxation = cd.relaxation_probability[index]
420 excitation = cd.excitation_probability[index]
421 scales = jnp.sqrt(jnp.stack((1 - excitation, 1 - relaxation)))
422 no_jump = no_jump * scales[:, None, None, None] * scales[None, None, :, None]
424 ground = joint[..., 0, :, 0, :]
425 excited = joint[..., 1, :, 1, :]
427 def transform(unitary, state):
428 return unitary @ state @ jnp.swapaxes(unitary.conj(), -1, -2)
430 down = jax.vmap(
431 lambda unitary: transform(
432 jnp.swapaxes(unitary.conj(), -1, -2),
433 excited,
434 )
435 )(jumps).mean(axis=0)
436 up = jax.vmap(lambda unitary: transform(unitary, ground))(jumps).mean(axis=0)
437 no_jump = no_jump.at[..., 0, :, 0, :].add(relaxation * down)
438 return no_jump.at[..., 1, :, 1, :].add(excitation * up)
441def _apply_qubit_unitary(joint, unitary):
442 return jnp.einsum(
443 "ai,...imjn,cj->...amcn",
444 unitary,
445 joint,
446 unitary.conj(),
447 )
450def _apply_noise(joint, noise, index):
451 moved = jnp.moveaxis(joint, (-3, -1), (-2, -1))
452 moved = apply_shifted_channel(
453 moved,
454 {
455 "_coefficients": noise.oscillator_coefficients[index],
456 "_shifts": noise.oscillator_shifts,
457 },
458 )
459 joint = jnp.moveaxis(moved, (-2, -1), (-3, -1))
460 joint = joint * noise.oscillator_dephasing[index][None, :, None, :]
462 probability = noise.qubit_probability[index]
463 excited_population = noise.qubit_excited_population[index]
464 ground = joint[..., 0, :, 0, :]
465 excited = joint[..., 1, :, 1, :]
466 coherence = jnp.sqrt(1 - probability) * noise.qubit_dephasing[index]
467 output_ground = (1 - excited_population * probability) * ground + (
468 1 - excited_population
469 ) * probability * excited
470 output_excited = (
471 excited_population * probability * ground
472 + (1 - (1 - excited_population) * probability) * excited
473 )
474 ge = coherence * joint[..., 0, :, 1, :]
475 eg = coherence * joint[..., 1, :, 0, :]
476 return jnp.stack(
477 (
478 jnp.stack((output_ground, ge), axis=-2),
479 jnp.stack((eg, output_excited), axis=-2),
480 ),
481 axis=-4,
482 )
485def apply_sbs_half_round(joint, ops: SBSHalfRound):
486 """Apply one prepared sBs half-round to a joint density matrix."""
488 def segment(index, state):
489 state = _apply_qubit_unitary(state, ops.rotations[index])
490 state = _apply_noise(state, ops.rotation_noise, index)
492 def apply_cd(inverse):
493 def step(_, value):
494 value = _apply_noisy_cd(value, ops.cd, index, inverse=inverse)
495 return _apply_noise(value, ops.cd_noise, index)
497 return step
499 state = jax.lax.fori_loop(0, ops.microsteps, apply_cd(False), state)
500 state = _apply_qubit_unitary(state, ops.echoes[index])
501 return jax.lax.fori_loop(0, ops.microsteps, apply_cd(True), state)
503 if jax.default_backend() == "cpu":
504 for index in range(3):
505 joint = segment(index, joint)
506 else:
507 joint = jax.lax.fori_loop(0, 3, segment, joint)
509 joint = _apply_qubit_unitary(joint, ops.rotations[3])
510 joint = _apply_noise(joint, ops.rotation_noise, 3)
511 joint = _apply_reset(joint, ops.reset)
512 return _apply_noise(joint, ops.reset_noise, 0)
515def oscillator_state(joint):
516 """Trace the ancilla from a joint density matrix."""
517 return jnp.trace(joint, axis1=-4, axis2=-2)
520@partial(jax.jit, static_argnames=("cycles",))
521def simulate_sbs(initial_states, observables, half_rounds, cycles):
522 """Evolve batched oscillator states and retain only observables."""
523 if cycles < 0:
524 raise ValueError("cycles must be non-negative")
525 if isinstance(half_rounds, SBSProtocol):
526 rounds = half_rounds.rounds
527 alternate_rounds = half_rounds.alternate_rounds
528 else:
529 rounds, alternate_rounds = half_rounds, None
530 ground = jnp.array([[1.0, 0.0], [0.0, 0.0]])
531 joint = jnp.einsum("ij,...mn->...imjn", ground, initial_states)
533 def expectation(state):
534 reduced = oscillator_state(state)
535 return jnp.einsum("...ij,...ji->...", observables, reduced).real
537 initial = expectation(joint)
539 def apply_round(state, selected):
540 for half_round in selected:
541 state = apply_sbs_half_round(state, half_round)
542 return state
544 def step(state, _):
545 state = apply_round(state, rounds)
546 return state, expectation(state)
548 if alternate_rounds is None:
549 joint, values = jax.lax.scan(step, joint, None, length=cycles)
550 else:
552 def alternate_step(state, cycle):
553 state = jax.lax.cond(
554 cycle % 2 == 0,
555 lambda value: apply_round(value, rounds),
556 lambda value: apply_round(value, alternate_rounds),
557 state,
558 )
559 return state, expectation(state)
561 joint, values = jax.lax.scan(
562 alternate_step,
563 joint,
564 jnp.arange(cycles),
565 )
566 return (
567 oscillator_state(joint),
568 jnp.concatenate(
569 (initial[None], values),
570 ),
571 joint,
572 )