Coverage for jaxquantum/circuits/library/sbs/device.py: 90%
216 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Shared cat and GKP sBs device simulations."""
3from __future__ import annotations
5from dataclasses import dataclass
6from typing import Callable
8import jax
9import jax.numpy as jnp
10import numpy as np
12import jaxquantum as jqt
13import jaxquantum.codes as jqcodes
14from jaxquantum.circuits.library.qubit import Rx, Ry
16from .core import (
17 SBSCDGeometry,
18 SBSNoise,
19 SBSProtocol,
20 build_sbs_cd_geometry,
21 build_sbs_half_round,
22 simulate_sbs,
23)
24from .parameters import GKP_JULY30
26__all__ = (
27 "ERROR_CHANNELS",
28 "DeviceParameters",
29 "CAT_DEVICE",
30 "CAT_MEASURED_DEVICE",
31 "GKP_DEVICE",
32 "GKP_JULY30_DEVICE",
33 "GKP_LEGACY_DEVICE",
34 "DecayResult",
35 "ErrorBudget",
36 "round_time",
37 "cat_protocol",
38 "gkp_protocol",
39 "gkp_displacements",
40 "prepare_gkp_protocol",
41 "prepare_cat_protocol",
42 "cat_problem",
43 "gkp_problem",
44 "fit_decay",
45 "simulate_decay",
46 "simulate_decay_variants",
47 "compute_error_budget",
48 "gkp_error_budget",
49 "cat_error_budget",
50)
53ERROR_CHANNELS = (
54 "storage_t1",
55 "storage_tphi",
56 "qubit_t1",
57 "qubit_tphi",
58 "qubit_t1_cd",
59 "reset",
60)
63@dataclass(frozen=True)
64class DeviceParameters:
65 storage_t1: float
66 storage_tphi: float
67 storage_nbar: float
68 qubit_t1: float
69 qubit_t1_cd: float | tuple[float, float, float]
70 qubit_tphi: float
71 qubit_excited_population: float
72 cd_durations: tuple[float, float, float]
73 rotation_durations: tuple[float, float, float, float]
74 reset_duration: float
75 reset_error: float = 0.0
76 reset_chi: float = 0.0
77 extra_storage_duration: float = 0.0
78 qubit_cd_excited_population: float | tuple[float, float, float] | None = None
81CAT_DEVICE = DeviceParameters(
82 storage_t1=100e-6,
83 storage_tphi=0.87e-3,
84 storage_nbar=0.0,
85 qubit_t1=400e-6,
86 qubit_t1_cd=400e-6,
87 qubit_tphi=1 / (1 / 90e-6 - 1 / (2 * 400e-6)),
88 qubit_excited_population=0.0,
89 cd_durations=(1.088e-6, 2.688e-6, 1.088e-6),
90 rotation_durations=(144e-9,) * 4,
91 reset_duration=324e-9,
92 extra_storage_duration=28e-9,
93)
95# July 30 measured values used in the cat-sBs lifetime audit.
96CAT_MEASURED_DEVICE = DeviceParameters(
97 storage_t1=90.96363465452683e-6,
98 storage_tphi=1.080e-3,
99 storage_nbar=0.0,
100 qubit_t1=438.2766324140162e-6,
101 qubit_t1_cd=438.2766324140162e-6,
102 qubit_tphi=1 / (1 / 51.56110574223823e-6 - 1 / (2 * 438.2766324140162e-6)),
103 qubit_excited_population=0.0,
104 cd_durations=(1.088e-6, 3.088e-6, 1.088e-6),
105 rotation_durations=(144e-9,) * 4,
106 reset_duration=480e-9,
107 extra_storage_duration=28e-9,
108)
111GKP_DEVICE = DeviceParameters(
112 storage_t1=606e-6,
113 storage_tphi=24e-3,
114 storage_nbar=0.0,
115 qubit_t1=280e-6,
116 qubit_t1_cd=280e-6,
117 qubit_tphi=1 / (1 / 238e-6 - 1 / (2 * 280e-6)),
118 qubit_excited_population=0.04,
119 cd_durations=(470e-9, 676e-9, 230e-9),
120 rotation_durations=(0.0, 32e-9, 32e-9, 0.0),
121 reset_duration=2.380e-6,
122 reset_error=0.01,
123)
125GKP_JULY30_DEVICE = DeviceParameters(**GKP_JULY30["device"])
128GKP_LEGACY_DEVICE = DeviceParameters(
129 storage_t1=90e-6,
130 storage_tphi=1e-3,
131 storage_nbar=0.00702828004369955,
132 qubit_t1=200e-6,
133 qubit_t1_cd=30e-6,
134 qubit_tphi=30e-6,
135 qubit_excited_population=0.430562654241043,
136 cd_durations=(400e-9, 1.2e-6, 400e-9),
137 rotation_durations=(0.0, 144e-9, 144e-9, 0.0),
138 reset_duration=236e-9,
139 reset_error=0.1,
140 reset_chi=2 * np.pi * 30e3,
141)
144@dataclass
145class DecayResult:
146 lifetime: float
147 rate: float
148 r2: float
149 contrast: np.ndarray
150 trace_error: float
151 hermiticity_error: float
152 minimum_eigenvalue: float
155@dataclass
156class ErrorBudget:
157 baseline: DecayResult
158 all_on: DecayResult
159 isolated: dict[str, DecayResult]
160 without: dict[str, DecayResult]
161 isolated_increments: dict[str, float]
162 context_increments: dict[str, float]
163 interaction_rate: float
165 @property
166 def ranking(self):
167 return sorted(
168 self.context_increments,
169 key=self.context_increments.get,
170 reverse=True,
171 )
173 def summary(self):
174 return {
175 "baseline_rate": self.baseline.rate,
176 "all_on_rate": self.all_on.rate,
177 "isolated_increments": self.isolated_increments,
178 "context_increments": self.context_increments,
179 "interaction_rate": self.interaction_rate,
180 "ranking": self.ranking,
181 }
184def round_time(device: DeviceParameters, half_rounds=1):
185 return half_rounds * (
186 sum(device.cd_durations)
187 + sum(device.rotation_durations)
188 + device.reset_duration
189 + device.extra_storage_duration
190 )
193def _enabled(value, name, enabled):
194 return value if name in enabled else jnp.inf
197def _noise(device, enabled):
198 return SBSNoise(
199 oscillator_t1=_enabled(
200 device.storage_t1,
201 "storage_t1",
202 enabled,
203 ),
204 oscillator_tphi=_enabled(
205 device.storage_tphi,
206 "storage_tphi",
207 enabled,
208 ),
209 oscillator_nbar=device.storage_nbar,
210 qubit_t1=_enabled(device.qubit_t1, "qubit_t1", enabled),
211 qubit_t1_cd=_enabled(
212 device.qubit_t1_cd,
213 "qubit_t1_cd",
214 enabled,
215 ),
216 qubit_tphi=_enabled(
217 device.qubit_tphi,
218 "qubit_tphi",
219 enabled,
220 ),
221 qubit_excited_population=device.qubit_excited_population,
222 qubit_cd_excited_population=device.qubit_cd_excited_population,
223 reset_error=device.reset_error if "reset" in enabled else 0.0,
224 reset_chi=device.reset_chi,
225 )
228def cat_protocol(
229 dimension,
230 nbar,
231 *,
232 delta=0.6,
233 ratio=3.125,
234 device=CAT_DEVICE,
235 enabled=ERROR_CHANNELS,
236 microsteps=1,
237 jump_samples=4,
238 max_loss=8,
239 cd_geometry=None,
240 alternate_cd_direction=False,
241):
242 """Build the nominal cat sBs measurement round."""
243 alpha = jnp.sqrt(nbar)
244 small = jnp.pi * delta**2 / (4 * alpha)
245 displacements = (small, -1j * jnp.pi / (2 * alpha), ratio * small)
246 rotations = (
247 Ry(jnp.pi / 2).U.data,
248 Rx(-jnp.pi / 2).U.data,
249 Rx(-jnp.pi / 2).U.data,
250 Ry(jnp.pi / 2).U.data,
251 )
253 def build(values, geometry):
254 return build_sbs_half_round(
255 dimension,
256 values,
257 rotations,
258 device.cd_durations,
259 device.rotation_durations,
260 device.reset_duration,
261 _noise(device, set(enabled)),
262 microsteps=microsteps,
263 jump_samples=jump_samples,
264 max_loss=max_loss,
265 storage_placement="lumped",
266 extra_storage_duration=device.extra_storage_duration,
267 reset_qubit_duration=0.0,
268 cd_geometry=geometry,
269 )
271 forward = build(displacements, cd_geometry)
272 if not alternate_cd_direction:
273 return (forward,)
274 reverse_geometry = SBSCDGeometry(
275 jnp.swapaxes(forward.cd.displacements.conj(), -1, -2),
276 jnp.swapaxes(forward.cd.jump_displacements.conj(), -1, -2),
277 )
278 reverse = build(tuple(-value for value in displacements), reverse_geometry)
279 return SBSProtocol((forward,), (reverse,))
282def gkp_protocol(
283 dimension,
284 *,
285 delta=0.428,
286 small_ratio=1.083,
287 small_displacement_scales=(1.0, 1.0),
288 big_displacement=None,
289 epsilon_model="sinh",
290 final_storage_rotation=0.0,
291 alternate_cd_direction=False,
292 length_scale=1.0,
293 device=GKP_DEVICE,
294 enabled=ERROR_CHANNELS,
295 microsteps=1,
296 jump_samples=4,
297 max_loss=8,
298 max_reset=12,
299 cd_geometries=None,
300):
301 """Build the two complementary GKP sBs half-rounds."""
302 z_displacements, x_displacements = gkp_displacements(
303 delta,
304 small_ratio,
305 small_displacement_scales,
306 big_displacement,
307 epsilon_model,
308 length_scale,
309 )
310 rotations = (
311 Ry(jnp.pi / 2).U.data,
312 Rx(-jnp.pi / 2).U.data,
313 Rx(jnp.pi / 2).U.data,
314 Ry(-jnp.pi / 2).U.data,
315 )
316 noise = _noise(device, set(enabled))
317 if cd_geometries is None:
318 cd_geometries = (None, None)
320 def build(displacements, geometry):
321 half_round = build_sbs_half_round(
322 dimension,
323 displacements,
324 rotations,
325 device.cd_durations,
326 device.rotation_durations,
327 device.reset_duration,
328 noise,
329 microsteps=microsteps,
330 jump_samples=jump_samples,
331 max_loss=max_loss,
332 max_reset=max_reset,
333 storage_placement="segment",
334 extra_storage_duration=device.extra_storage_duration,
335 cd_geometry=geometry,
336 )
337 if final_storage_rotation:
338 phase = jnp.exp(
339 -1j * final_storage_rotation * jnp.arange(dimension)
340 )
341 phase_factor = phase[:, None] * phase.conj()[None, :]
342 half_round = half_round._replace(
343 reset=half_round.reset._replace(
344 phase_factor=phase_factor * half_round.reset.phase_factor,
345 )
346 )
347 return half_round
349 forward = (
350 build(z_displacements, cd_geometries[0]),
351 build(x_displacements, cd_geometries[1]),
352 )
353 if not alternate_cd_direction:
354 return forward
355 reverse_geometries = tuple(
356 SBSCDGeometry(
357 jnp.swapaxes(half_round.cd.displacements.conj(), -1, -2),
358 jnp.swapaxes(half_round.cd.jump_displacements.conj(), -1, -2),
359 )
360 for half_round in forward
361 )
362 reverse = tuple(
363 build(
364 tuple(-value for value in displacements),
365 geometry,
366 )
367 for displacements, geometry in zip(
368 (z_displacements, x_displacements),
369 reverse_geometries,
370 )
371 )
372 return SBSProtocol(forward, reverse)
375def gkp_displacements(
376 delta,
377 small_ratio,
378 small_displacement_scales,
379 big_displacement,
380 epsilon_model,
381 length_scale,
382):
383 """Return complementary Z- and X-stabilizer displacements."""
384 length = jnp.sqrt(2 * jnp.pi) * length_scale
385 if epsilon_model == "sinh":
386 epsilon = jnp.sinh(delta**2) * length
387 elif epsilon_model == "quadratic":
388 epsilon = delta**2 * length
389 else:
390 raise ValueError("epsilon_model must be 'sinh' or 'quadratic'")
391 scales = jnp.asarray(small_displacement_scales)
392 small = (epsilon / 2 * scales[0], small_ratio * epsilon / 2 * scales[1])
393 big = length if big_displacement is None else big_displacement
394 return (
395 (small[0], -1j * big, small[1]),
396 (1j * small[0], big, 1j * small[1]),
397 )
400def prepare_gkp_protocol(
401 dimension,
402 *,
403 delta=0.428,
404 small_ratio=1.083,
405 small_displacement_scales=(1.0, 1.0),
406 big_displacement=None,
407 epsilon_model="sinh",
408 final_storage_rotation=0.0,
409 alternate_cd_direction=False,
410 length_scale=1.0,
411 device=GKP_DEVICE,
412 microsteps=1,
413 jump_samples=4,
414 max_loss=8,
415 max_reset=12,
416):
417 """Return an error-channel builder with shared pulse geometry."""
418 displacements = gkp_displacements(
419 delta,
420 small_ratio,
421 small_displacement_scales,
422 big_displacement,
423 epsilon_model,
424 length_scale,
425 )
426 geometries = tuple(
427 build_sbs_cd_geometry(
428 dimension,
429 values,
430 microsteps=microsteps,
431 jump_samples=jump_samples,
432 )
433 for values in displacements
434 )
436 def build(enabled):
437 return gkp_protocol(
438 dimension,
439 delta=delta,
440 small_ratio=small_ratio,
441 small_displacement_scales=small_displacement_scales,
442 big_displacement=big_displacement,
443 epsilon_model=epsilon_model,
444 final_storage_rotation=final_storage_rotation,
445 alternate_cd_direction=alternate_cd_direction,
446 length_scale=length_scale,
447 device=device,
448 enabled=enabled,
449 microsteps=microsteps,
450 jump_samples=jump_samples,
451 max_loss=max_loss,
452 max_reset=max_reset,
453 cd_geometries=geometries,
454 )
456 return build
459def prepare_cat_protocol(
460 dimension,
461 nbar,
462 *,
463 delta=0.6,
464 ratio=3.125,
465 device=CAT_DEVICE,
466 microsteps=1,
467 jump_samples=4,
468 max_loss=8,
469 alternate_cd_direction=False,
470):
471 """Return an error-channel builder with shared pulse geometry."""
472 alpha = jnp.sqrt(nbar)
473 small = jnp.pi * delta**2 / (4 * alpha)
474 displacements = (small, -1j * jnp.pi / (2 * alpha), ratio * small)
475 geometry = build_sbs_cd_geometry(
476 dimension,
477 displacements,
478 microsteps=microsteps,
479 jump_samples=jump_samples,
480 )
482 def build(enabled):
483 return cat_protocol(
484 dimension,
485 nbar,
486 delta=delta,
487 ratio=ratio,
488 device=device,
489 enabled=enabled,
490 microsteps=microsteps,
491 jump_samples=jump_samples,
492 max_loss=max_loss,
493 cd_geometry=geometry,
494 alternate_cd_direction=alternate_cd_direction,
495 )
497 return build
500def cat_problem(dimension, nbar, kind="bit"):
501 """Return the two cat states and observable defining a contrast."""
502 alpha = jnp.sqrt(nbar)
503 plus = jqt.displace(dimension, alpha) @ jqt.basis(dimension, 0)
504 minus = jqt.displace(dimension, -alpha) @ jqt.basis(dimension, 0)
505 if kind == "bit":
506 q = (jqt.destroy(dimension) + jqt.create(dimension)).data / jnp.sqrt(2)
507 values, vectors = jnp.linalg.eigh(q)
508 observable = (vectors * jnp.sign(values)) @ vectors.conj().T
509 states = (plus, minus)
510 elif kind == "phase":
511 states = (jqt.unit(plus + minus), jqt.unit(plus - minus))
512 observable = jnp.diag((-1.0) ** jnp.arange(dimension))
513 else:
514 raise ValueError("kind must be 'bit' or 'phase'")
515 density = jnp.stack([(state @ state.dag()).data for state in states])
516 return density, jnp.broadcast_to(observable, density.shape)
519def gkp_problem(dimension, delta=0.428, kind="x"):
520 """Return the two GKP states and observable defining a logical contrast."""
521 kind = kind.lower()
522 if kind not in ("x", "z"):
523 raise ValueError("kind must be 'x' or 'z'")
524 code = jqcodes.GKPQubit({"delta": delta, "N": dimension})
525 states = (code.basis[f"-{kind}"], code.basis[f"+{kind}"])
526 density = jnp.stack([(state @ state.dag()).data for state in states])
527 observable = code.common_gates[f"{kind.upper()}_0"].data
528 return density, jnp.broadcast_to(observable, density.shape)
531def fit_decay(times, contrast, start=4, floor=1e-10):
532 values = np.abs(np.asarray(contrast))
533 times = np.asarray(times)
534 valid = np.isfinite(values) & (values > floor) & (np.arange(values.size) >= start)
535 if valid.sum() < 3:
536 raise ValueError("at least three finite decay samples are required")
537 x = times[valid]
538 y = np.log(values[valid])
539 slope, intercept = np.polyfit(x, y, 1)
540 fitted = intercept + slope * x
541 residual = np.sum((y - fitted) ** 2)
542 total = np.sum((y - y.mean()) ** 2)
543 r2 = 1.0 if total == 0 else 1 - residual / total
544 lifetime = np.inf if slope >= 0 else -1 / slope
545 return lifetime, max(0.0, -slope), r2
548def simulate_decay(
549 initial_states,
550 observables,
551 half_rounds,
552 cycles,
553 cycle_time,
554 *,
555 fit_start=4,
556 fit_floor=1e-10,
557):
558 final, values, _ = simulate_sbs(
559 initial_states,
560 observables,
561 half_rounds,
562 cycles,
563 )
564 final, values = jax.device_get((final, values))
565 return _analyze_decay(
566 final,
567 values,
568 cycles,
569 cycle_time,
570 fit_start,
571 fit_floor,
572 )
575def _analyze_decay(final, values, cycles, cycle_time, fit_start, fit_floor):
576 contrast = np.asarray(values[:, 0] - values[:, 1])
577 times = np.arange(cycles + 1) * cycle_time
578 lifetime, rate, r2 = fit_decay(
579 times,
580 contrast,
581 start=fit_start,
582 floor=fit_floor,
583 )
584 traces = np.trace(final, axis1=-2, axis2=-1)
585 hermiticity = np.max(np.abs(final - final.conj().swapaxes(-1, -2)))
586 minimum_eigenvalue = np.min(np.linalg.eigvalsh(final))
587 return DecayResult(
588 lifetime=lifetime,
589 rate=rate,
590 r2=r2,
591 contrast=contrast,
592 trace_error=float(np.max(np.abs(traces - 1))),
593 hermiticity_error=float(hermiticity),
594 minimum_eigenvalue=float(minimum_eigenvalue),
595 )
598def _stack_protocols(protocols):
599 flattened = [jax.tree.flatten(protocol) for protocol in protocols]
600 structure = flattened[0][1]
601 if any(item[1] != structure for item in flattened[1:]):
602 raise ValueError("protocol variants must have equal structures")
603 leaves = []
604 axes = []
605 for items in zip(*(item[0] for item in flattened)):
606 if all(item is items[0] for item in items[1:]):
607 leaves.append(items[0])
608 axes.append(None)
609 else:
610 leaves.append(jnp.stack(items))
611 axes.append(0)
612 return (
613 jax.tree.unflatten(structure, leaves),
614 jax.tree.unflatten(structure, axes),
615 )
618def simulate_decay_variants(
619 initial_states,
620 observables,
621 protocols,
622 cycles,
623 cycle_time,
624 *,
625 fit_start=4,
626 fit_floor=1e-10,
627):
628 """Simulate equal-structure protocol variants in one mapped call."""
629 if len(protocols) == 1:
630 return [
631 simulate_decay(
632 initial_states,
633 observables,
634 protocols[0],
635 cycles,
636 cycle_time,
637 fit_start=fit_start,
638 fit_floor=fit_floor,
639 )
640 ]
641 protocols, axes = _stack_protocols(protocols)
642 final, values, _ = jax.vmap(
643 lambda rounds: simulate_sbs(
644 initial_states,
645 observables,
646 rounds,
647 cycles,
648 ),
649 in_axes=(axes,),
650 )(protocols)
651 final, values = jax.device_get((final, values))
652 return [
653 _analyze_decay(
654 result,
655 samples,
656 cycles,
657 cycle_time,
658 fit_start,
659 fit_floor,
660 )
661 for result, samples in zip(final, values)
662 ]
665def compute_error_budget(
666 protocol: Callable[[set[str]], tuple],
667 initial_states,
668 observables,
669 cycles,
670 cycle_time,
671 *,
672 channels=ERROR_CHANNELS,
673 fit_start=4,
674 fit_floor=1e-10,
675 batched=True,
676):
677 """Return baseline-subtracted and all-on-context channel budgets."""
678 channels = tuple(channels)
680 def run(enabled):
681 return simulate_decay(
682 initial_states,
683 observables,
684 protocol(set(enabled)),
685 cycles,
686 cycle_time,
687 fit_start=fit_start,
688 fit_floor=fit_floor,
689 )
691 variants = [
692 (),
693 channels,
694 *((channel,) for channel in channels),
695 *(set(channels) - {channel} for channel in channels),
696 ]
697 if batched:
698 results = simulate_decay_variants(
699 initial_states,
700 observables,
701 [protocol(set(enabled)) for enabled in variants],
702 cycles,
703 cycle_time,
704 fit_start=fit_start,
705 fit_floor=fit_floor,
706 )
707 else:
708 results = [run(enabled) for enabled in variants]
709 baseline, all_on = results[:2]
710 split = 2 + len(channels)
711 isolated = dict(zip(channels, results[2:split]))
712 without = dict(zip(channels, results[split:]))
713 isolated_increments = {
714 channel: isolated[channel].rate - baseline.rate for channel in channels
715 }
716 context_increments = {
717 channel: all_on.rate - without[channel].rate for channel in channels
718 }
719 interaction = all_on.rate - baseline.rate - sum(isolated_increments.values())
720 return ErrorBudget(
721 baseline=baseline,
722 all_on=all_on,
723 isolated=isolated,
724 without=without,
725 isolated_increments=isolated_increments,
726 context_increments=context_increments,
727 interaction_rate=interaction,
728 )
731def gkp_error_budget(
732 *,
733 dimension=60,
734 delta=0.428,
735 state_delta=None,
736 kind="x",
737 small_ratio=1.083,
738 small_displacement_scales=(1.0, 1.0),
739 big_displacement=None,
740 epsilon_model="sinh",
741 final_storage_rotation=0.0,
742 alternate_cd_direction=False,
743 cycles=512,
744 microsteps=1,
745 device=GKP_DEVICE,
746 fit_start=4,
747 fit_floor=1e-10,
748 max_loss=8,
749 max_reset=12,
750):
751 initial, observables = gkp_problem(
752 dimension,
753 delta if state_delta is None else state_delta,
754 kind,
755 )
756 protocol = prepare_gkp_protocol(
757 dimension,
758 delta=delta,
759 small_ratio=small_ratio,
760 small_displacement_scales=small_displacement_scales,
761 big_displacement=big_displacement,
762 epsilon_model=epsilon_model,
763 final_storage_rotation=final_storage_rotation,
764 alternate_cd_direction=alternate_cd_direction,
765 device=device,
766 microsteps=microsteps,
767 max_loss=max_loss,
768 max_reset=max_reset,
769 )
770 return compute_error_budget(
771 protocol,
772 initial,
773 observables,
774 cycles,
775 round_time(device, 2),
776 fit_start=fit_start,
777 fit_floor=fit_floor,
778 )
781def cat_error_budget(
782 nbar,
783 *,
784 dimension=48,
785 kind="bit",
786 cycles=480,
787 microsteps=1,
788 device=CAT_DEVICE,
789 fit_start=4,
790 fit_floor=1e-10,
791 max_loss=8,
792):
793 initial, observables = cat_problem(dimension, nbar, kind)
794 protocol = prepare_cat_protocol(
795 dimension,
796 nbar,
797 device=device,
798 microsteps=microsteps,
799 max_loss=max_loss,
800 )
801 return compute_error_budget(
802 protocol,
803 initial,
804 observables,
805 cycles,
806 round_time(device),
807 fit_start=fit_start,
808 fit_floor=fit_floor,
809 )