File size: 27,669 Bytes
2a87ede
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
import subprocess
import sys
import os
import time
import psutil
import platform
from typing import Optional, List, Dict, Any
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

try:
    import cupy as cp
    HAS_CUPY = True
except ImportError:
    HAS_CUPY = False

# JAX is now a mandatory dependency of dense_evolution (no numpy fallback
# detection) -- the numpy code paths below are kept as-is for reference/
# reuse, they are just never selected anymore.
import jax
import jax.numpy as jnp
HAS_JAX = True
jax.config.update("jax_enable_x64", True)


class QuantumHardwareRegistry:
    def __init__(self):
        self.processor = platform.processor()
        self.ram_total = psutil.virtual_memory().total / (1024**3)
        self.ram_avail = psutil.virtual_memory().available / (1024**3)
        self.has_cupy = HAS_CUPY
        self.has_jax = HAS_JAX
        self.has_gpu = self._detect_gpu()
        self.max_dense_qubits = self._get_qubit_limit()

    def _detect_gpu(self) -> bool:
        try:
            subprocess.check_output(['nvidia-smi'], stderr=subprocess.DEVNULL)
            return True
        except Exception:
            return False

    def _get_qubit_limit(self) -> int:
        if self.ram_total >= 50: return 28
        elif self.ram_total >= 12: return 24
        return 20

    def print_diagnostics(self):
        print(f"MAX_DENSE={self.max_dense_qubits}q | JAX={self.has_jax} | GPU={self.has_gpu}")


HARDWARE_REGISTRY = QuantumHardwareRegistry()
plt.style.use('dark_background')
matplotlib.rcParams.update({
    'figure.facecolor': '#010409',
    'axes.facecolor': '#0d1117',
    'axes.edgecolor': '#21262d',
    'grid.color': '#21262d',
    'font.family': 'monospace',
    'font.size': 9,
})


# ─────────────────────────────────────────────────────────────────────────────
# Internal helpers
# ─────────────────────────────────────────────────────────────────────────────

def _fresh_rng() -> np.random.Generator:
    """

    Create a hardware-entropy-seeded RNG.

    Combines os.urandom (CSPRNG) with a high-resolution nanosecond counter

    so two calls within the same microsecond still differ.

    """
    entropy_bytes = os.urandom(8)
    entropy_int   = int.from_bytes(entropy_bytes, byteorder='big')
    ns_counter    = time.perf_counter_ns() & 0xFFFF_FFFF_FFFF_FFFF
    seed          = (entropy_int ^ ns_counter) & 0xFFFF_FFFF_FFFF_FFFF
    return np.random.default_rng(seed)


def _qubit_index_pairs(dim: int, q: int):
    """

    Return (idx_0, idx_1) β€” two integer arrays of length dim/2 β€” where

    idx_0[i] has bit q == 0 and idx_1[i] = idx_0[i] | (1 << q).



    This is the correct and vectorised way to build qubit-pair indices.

    The original code used `xp.where()` which returns a *tuple*, then

    did `idx_1 = idx_0 | step` on that tuple β€” producing wrong indices

    for all models and making phaseflip look deterministic.

    """
    step   = 1 << q
    all_i  = np.arange(dim, dtype=np.intp)
    idx_0  = all_i[(all_i & step) == 0]          # shape: (dim//2,)
    idx_1  = idx_0 | step                          # shape: (dim//2,)
    return idx_0, idx_1


# ─────────────────────────────────────────────────────────────────────────────
# NoiseSpec β€” JAX PyTree wrapper for NoiseModel.apply_to_sv's parameters
# ─────────────────────────────────────────────────────────────────────────────

class NoiseSpec:
    """JAX PyTree wrapper around a NoiseModel configuration, so noise

    parameters can be threaded through jax.jit/jax.grad/jax.vmap natively

    -- e.g. as the `noise=` argument to circuit_to_energy_fn's energy_fn

    -- instead of being applied as an external, Python-side step around

    the already-traced circuit (the old way: build sv, exit the trace,

    call apply_to_sv separately).



    `model`/`qubits` are static (aux_data): they select which code path

    runs, not values to differentiate or batch over -- the same role

    `static_argnames` plays for a plain jax.jit function, but automatic

    here because it's part of the pytree structure. `p`/`jax_key` are

    pytree leaves (children): `p` can be a traced/differentiable value

    (e.g. optimizing noise strength itself), and `jax_key` flows through

    jit/vmap/scan the way any other JAX array does -- no external

    Python-level key management, no OS-entropy fallback (unlike

    apply_to_sv called standalone with jax_key=None), so a NoiseSpec's

    result is always reproducible from the key it was built with.



    `jax_key` is required (not Optional) -- the whole point of wiring

    noise into the traced computation this way is to remove the need for

    an external, ad-hoc key-management workaround; a caller who wants a

    fresh key per call should split one themselves (`jax.random.split`)

    and build a fresh NoiseSpec, the same as any other JAX-idiomatic

    stateless-key pattern.

    """
    def __init__(self, model: str, p, jax_key, qubits: Optional[List[int]] = None):
        self.model = model
        self.p = p
        self.jax_key = jax_key
        self.qubits = tuple(qubits) if qubits is not None else None

    def __repr__(self):
        return f"NoiseSpec(model={self.model!r}, p={self.p!r}, qubits={self.qubits!r})"

    def tree_flatten(self):
        children = (self.p, self.jax_key)
        aux_data = (self.model, self.qubits)
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        model, qubits = aux_data
        p, jax_key = children
        return cls(model, p, jax_key, qubits)


jax.tree_util.register_pytree_node(
    NoiseSpec, NoiseSpec.tree_flatten, NoiseSpec.tree_unflatten
)


# ─────────────────────────────────────────────────────────────────────────────
# NoiseModel
# ─────────────────────────────────────────────────────────────────────────────

class NoiseModel:
    """

    Stochastic single-qubit Kraus channels applied directly to a statevector.



    All channels are mathematically correct Kraus maps:

    - trace is preserved (normalisation enforced at the end)

    - phaseflip applies Z with probability p per qubit (non-deterministic)

    - amplitude_damping applies the correct K0/K1 Kraus operators

    - combined is a true worst-case NISQ mixture of all three Pauli errors

      plus amplitude damping



    Supported models

    ----------------

    'ideal'            identity β€” no modification

    'depolarizing'     {√(1-p)I, √(p/3)X, √(p/3)Y, √(p/3)Z}

    'bitflip'          {√(1-p)I, √p·X}

    'phaseflip'        {√(1-p)I, √pΒ·Z}    ← was broken, now fixed

    'amplitude_damping'{K0=diag(1,√(1-γ)), K1=[[0,√γ],[0,0]]}

    'combined'         depolarizing(p/2) + amplitude_damping(p/3), renormalised



    Every channel draws one fire/no-fire decision per qubit per shot

    (plus one Pauli choice for depolarizing/combined's depolarizing

    sub-step), applied identically across the whole statevector -- the

    same single-Pauli-per-qubit-per-shot convention STIM's

    DEPOLARIZE1(p) uses. Prior to v8.1.57, every channel instead drew

    2**(n-1) INDEPENDENT decisions per qubit per shot, one per amplitude

    pair (i.e. one per branch of the other n-1 qubits) -- inert on a

    product state, but on an entangled state it over-decohered any

    coherence-sensitive (off-diagonal) observable, up to hundreds of

    sigma vs the exact density-matrix Kraus-sum result on test cases

    (e.g. per-branch sampling dropped a measured value from 1.0 to 0.31

    at p=0.15 on one such test -- see the v8.1.57 changelog entry for

    the full reproduction).

    """

    MODELS = ['ideal', 'depolarizing', 'bitflip', 'phaseflip',
              'amplitude_damping', 'combined']

    @staticmethod
    def apply_to_sv(

        sv:       np.ndarray,

        n:        int,

        model:    str,

        p:        float,

        rng:      Optional[np.random.Generator] = None,

        qubits:   Optional[List[int]] = None,

        jax_key:  Optional[Any] = None,

    ) -> np.ndarray:
        """

        Apply a stochastic Kraus channel to statevector *sv* in-place

        (numpy path) or via functional updates (JAX path).



        Parameters

        ----------

        sv      : complex statevector of length 2**n

        n       : number of qubits

        model   : one of NoiseModel.MODELS

        p       : error probability (or damping rate Ξ³ for amplitude_damping)

        rng     : optional pre-seeded numpy Generator. Used directly when

                  *sv* is a NumPy array. When *sv* is a JAX array, `rng`

                  used to be silently ignored in favor of `jax_key` (or a

                  non-reproducible OS-entropy key if that was also None

                  -- issue #7); it is now used to *derive* a reproducible

                  jax_key (`rng.integers(...)` seeds `jax.random.PRNGKey`)

                  whenever `jax_key` isn't given explicitly, so seeding

                  `rng` has the effect a caller expects on both array

                  types instead of only on one of them.

        qubits  : subset of qubits to apply the channel to; defaults to all

        jax_key : optional JAX PRNGKey, only meaningful when *sv* is a JAX

                  array. Takes precedence over `rng` when both are given

                  (explicit key beats a derived one). Created from OS

                  entropy if neither `jax_key` nor `rng` is given.



        Returns

        -------

        Normalised statevector (same array type as input).

        """
        if model == 'ideal':
            return sv
        try:
            if p <= 0.0:
                return sv
        except jax.errors.TracerBoolConversionError:
            # p is a traced value (e.g. flowing through jax.jit/vmap/grad
            # as a NoiseSpec pytree leaf) -- can't early-exit on a Python
            # bool of it. Falling through is still correct: every channel
            # below already reduces to a no-op at p=0 (`fire = r < p` is
            # always False), this skips only the eager-mode optimization,
            # not correctness.
            pass

        is_jax = HAS_JAX and isinstance(sv, jnp.ndarray)
        dim    = len(sv)

        # ── RNG initialisation ────────────────────────────────────────
        if is_jax:
            if jax_key is not None:
                key = jax_key
            elif rng is not None:
                # Derive a reproducible JAX key from the caller's seeded
                # NumPy generator instead of silently ignoring it -- each
                # call advances `rng`'s state, so a fresh, identically-
                # seeded `rng` reproduces the exact same sequence of keys
                # across separate runs (same guarantee the NumPy path
                # already gives).
                key = jax.random.PRNGKey(int(rng.integers(0, 2**32 - 1)))
            else:
                seed_bytes = os.urandom(4)
                jax_seed   = int.from_bytes(seed_bytes, byteorder='big')
                jax_seed  ^= time.perf_counter_ns() & 0xFFFF_FFFF
                key = jax.random.PRNGKey(jax_seed)
        else:
            if rng is None:
                rng = _fresh_rng()

        target_qubits = qubits if qubits is not None else list(range(n))
        sv_out = sv  # JAX: functional; NumPy: will be modified in-place copy

        if not is_jax:
            sv_out = sv.copy()  # never mutate the caller's array

        for q in target_qubits:
            # ── correct index pair construction ───────────────────────
            idx_0, idx_1 = _qubit_index_pairs(dim, q)

            # ── channel application ───────────────────────────────────
            #
            # BUG FIX (all branches below): every channel used to draw
            # `half` = 2**(n-1) INDEPENDENT fire/no-fire (and, for
            # depolarizing, Pauli-choice) decisions per qubit per shot --
            # one per computational-basis amplitude pair, i.e. one per
            # branch of the OTHER n-1 qubits -- instead of ONE decision
            # per qubit per shot applied uniformly across the whole
            # statevector, which is what a correct Kraus-channel
            # unraveling requires (the convention STIM's DEPOLARIZE1(p)/
            # X_ERROR(p)/Z_ERROR(p) use: one Pauli draw per qubit per
            # shot, applied globally). A single-qubit noise event cannot
            # be correlated with what value the *other* qubits happen to
            # hold in superposition -- that correlation is exactly what
            # per-branch-independent sampling introduces.
            #
            # On a product state this was inert (only one branch has
            # nonzero amplitude). On an entangled/superposed state it
            # acts like measuring the other qubits in the computational
            # basis before deciding the error, over-decohering any
            # observable sensitive to coherence between branches --
            # confirmed for 'depolarizing' at 17-33 sigma vs an exact
            # density-matrix Kraus-sum reference on a GHZ-4 state's XXXX
            # expectation (equivalent to a true depolarizing channel at
            # p_eff ~2.2-2.5x nominal p), and for 'bitflip' even more
            # starkly: <XXXX> must be EXACTLY invariant under a bit-flip
            # channel (X commutes with the X observable) yet the old
            # per-branch sampling dropped it from 1.0 to 0.31 at p=0.15
            # (98-210 sigma). Fixed by drawing one scalar decision per
            # qubit per shot and applying it identically to every branch
            # (idx_0/idx_1 index the qubit's own two values; the fixed
            # code touches them as whole arrays, not element-by-element).
            if model == 'depolarizing':
                # Three equiprobable Pauli errors GIVEN that the channel
                # fired -- the choice AMONG X/Y/Z must be a uniform
                # 1-in-3 pick, independent of p, so ch's thresholds are
                # fixed at 1/3 and 2/3 (see historical note below).
                THIRD = 1.0 / 3.0
                if is_jax:
                    key, sk1, sk2 = jax.random.split(key, 3)
                    r  = jax.random.uniform(sk1, shape=(), minval=0.0, maxval=1.0)
                    ch = jax.random.uniform(sk2, shape=(), minval=0.0, maxval=1.0)
                    fire   = r < p
                    x_gate = fire & (ch < THIRD)
                    y_gate = fire & (ch >= THIRD) & (ch < 2.0 * THIRD)
                    z_gate = fire & (ch >= 2.0 * THIRD)
                    v0, v1 = sv_out[idx_0], sv_out[idx_1]
                    new_v0 = jnp.where(x_gate,  v1,
                             jnp.where(y_gate, -1j * v1, v0))
                    new_v1 = jnp.where(x_gate,  v0,
                             jnp.where(y_gate,  1j * v0,
                             jnp.where(z_gate, -v1, v1)))
                    sv_out = sv_out.at[idx_0].set(new_v0)
                    sv_out = sv_out.at[idx_1].set(new_v1)
                else:
                    r  = rng.random()
                    ch = rng.random()
                    if r < p:
                        if ch < THIRD:
                            sv_out[idx_0], sv_out[idx_1] = sv_out[idx_1].copy(), sv_out[idx_0].copy()
                        elif ch < 2.0 * THIRD:
                            v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                            sv_out[idx_0] = -1j * v1
                            sv_out[idx_1] =  1j * v0
                        else:
                            sv_out[idx_1] = -sv_out[idx_1]

            elif model == 'bitflip':
                # X gate applied with probability p, once per qubit per
                # shot, uniformly across the whole statevector.
                if is_jax:
                    key, subkey = jax.random.split(key)
                    r    = jax.random.uniform(subkey, shape=(), minval=0.0, maxval=1.0)
                    fire = r < p
                    v0, v1 = sv_out[idx_0], sv_out[idx_1]
                    sv_out = sv_out.at[idx_0].set(jnp.where(fire, v1, v0))
                    sv_out = sv_out.at[idx_1].set(jnp.where(fire, v0, v1))
                else:
                    r = rng.random()
                    if r < p:
                        sv_out[idx_0], sv_out[idx_1] = sv_out[idx_1].copy(), sv_out[idx_0].copy()

            elif model == 'phaseflip':
                # Z gate applied with probability p, once per qubit per
                # shot:
                # Z|0⟩ = |0⟩  β†’  no change to idx_0 amplitudes
                # Z|1⟩ = -|1⟩ β†’  negate idx_1 amplitudes when fired
                if is_jax:
                    key, subkey = jax.random.split(key)
                    r    = jax.random.uniform(subkey, shape=(), minval=0.0, maxval=1.0)
                    fire = r < p
                    v1     = sv_out[idx_1]
                    sv_out = sv_out.at[idx_1].set(jnp.where(fire, -v1, v1))
                else:
                    r = rng.random()
                    if r < p:
                        sv_out[idx_1] = -sv_out[idx_1]

            elif model == 'amplitude_damping':
                # K0 = [[1, 0], [0, √(1-Ξ³)]]  β€” no decay
                # K1 = [[0, √γ], [0, 0]]       β€” decay |1⟩ β†’ |0⟩
                #
                # Single-trajectory ("quantum jump") unraveling: ONE
                # decay/no-decay decision per qubit per shot, using the
                # Born-rule probability aggregated over the WHOLE
                # statevector, P(K1) = <psi|K1^dagger K1|psi> =
                # Ξ³ * Ξ£_i |v1[i]|^2 (sum over every branch of the other
                # n-1 qubits, not per branch -- see the per-branch-
                # sampling bug note above the 'depolarizing' branch).
                # Cross-checked against John Preskill's Ph219/CS219
                # Chapter 3 (Caltech lecture notes), whose own derivation
                # of this channel (system-environment isometry + partial
                # trace) gives exactly this K0/K1 pair. If it fires, K1
                # collapses the ENTIRE qubit-q=1 branch onto q=0,
                # preserving the other qubits' relative amplitudes
                # (v1[i]/norm, not flattened to unit phase per branch --
                # the previous per-branch fix only got this right for a
                # single isolated qubit, where there is only one branch
                # to preserve). Renormalizing by the GLOBAL P(K1) (or
                # 1-P(K1) for the no-decay branch) here, not a per-branch
                # norm, is what a correct sequential multi-qubit
                # trajectory needs -- the next qubit's Born-rule
                # probability must be computed on a properly
                # renormalized state.
                gamma = float(np.clip(p, 0.0, 1.0))
                if is_jax:
                    key, subkey = jax.random.split(key)
                    r = jax.random.uniform(subkey, shape=(), minval=0.0, maxval=1.0)
                    v0, v1 = sv_out[idx_0], sv_out[idx_1]
                    p1 = jnp.clip(gamma * jnp.sum(jnp.abs(v1) ** 2), 0.0, 1.0)
                    decay = r < p1
                    norm_decay    = jnp.sqrt(jnp.maximum(p1, 1e-15))
                    norm_no_decay = jnp.sqrt(jnp.maximum(1.0 - p1, 1e-15))
                    new_v0 = jnp.where(decay, v1 * jnp.sqrt(gamma) / norm_decay, v0 / norm_no_decay)
                    new_v1 = jnp.where(decay, 0.0 + 0j, v1 * jnp.sqrt(1.0 - gamma) / norm_no_decay)
                    sv_out = sv_out.at[idx_0].set(new_v0)
                    sv_out = sv_out.at[idx_1].set(new_v1)
                else:
                    r = rng.random()
                    v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                    p1 = float(np.clip(gamma * np.sum(np.abs(v1) ** 2), 0.0, 1.0))
                    if r < p1:
                        norm_decay = np.sqrt(max(p1, 1e-15))
                        sv_out[idx_0] = v1 * np.sqrt(gamma) / norm_decay
                        sv_out[idx_1] = 0.0
                    else:
                        norm_no_decay = np.sqrt(max(1.0 - p1, 1e-15))
                        sv_out[idx_0] = v0 / norm_no_decay
                        sv_out[idx_1] = v1 * np.sqrt(1.0 - gamma) / norm_no_decay

            elif model == 'combined':
                # Worst-case NISQ: depolarizing(p/2) then amplitude_damping(p/3)
                # applied sequentially on the same qubit, each as a single
                # per-qubit-per-shot decision (see the notes on the
                # 'depolarizing' and 'amplitude_damping' branches above --
                # this sub-channel used to share the same per-branch bug,
                # plus its own amplitude-damping sub-step never replaced
                # (only added to) the decayed |0> amplitude and never
                # renormalized the no-decay branch at all).
                p_dep   = p * 0.5
                p_damp  = p * 0.333333
                THIRD   = 1.0 / 3.0  # see the 'depolarizing' branch above for why
                                     # this must be fixed at 1/3, not p_dep/3

                # β€” depolarizing sub-channel β€”
                if is_jax:
                    key, sk1, sk2 = jax.random.split(key, 3)
                    r_dep = jax.random.uniform(sk1, shape=(), minval=0.0, maxval=1.0)
                    ch    = jax.random.uniform(sk2, shape=(), minval=0.0, maxval=1.0)
                    fire   = r_dep < p_dep
                    x_gate = fire & (ch < THIRD)
                    y_gate = fire & (ch >= THIRD) & (ch < 2.0 * THIRD)
                    z_gate = fire & (ch >= 2.0 * THIRD)
                    v0, v1 = sv_out[idx_0], sv_out[idx_1]
                    new_v0 = jnp.where(x_gate,  v1,
                             jnp.where(y_gate, -1j * v1, v0))
                    new_v1 = jnp.where(x_gate,  v0,
                             jnp.where(y_gate,  1j * v0,
                             jnp.where(z_gate, -v1, v1)))
                    sv_out = sv_out.at[idx_0].set(new_v0)
                    sv_out = sv_out.at[idx_1].set(new_v1)
                    # β€” amplitude_damping sub-channel, global Born rule β€”
                    key, sk3 = jax.random.split(key)
                    r_damp = jax.random.uniform(sk3, shape=(), minval=0.0, maxval=1.0)
                    v0, v1 = sv_out[idx_0], sv_out[idx_1]
                    p1 = jnp.clip(p_damp * jnp.sum(jnp.abs(v1) ** 2), 0.0, 1.0)
                    decay = r_damp < p1
                    norm_decay    = jnp.sqrt(jnp.maximum(p1, 1e-15))
                    norm_no_decay = jnp.sqrt(jnp.maximum(1.0 - p1, 1e-15))
                    new_v0d = jnp.where(decay, v1 * jnp.sqrt(p_damp) / norm_decay, v0 / norm_no_decay)
                    new_v1d = jnp.where(decay, 0.0 + 0j, v1 * jnp.sqrt(1.0 - p_damp) / norm_no_decay)
                    sv_out = sv_out.at[idx_0].set(new_v0d)
                    sv_out = sv_out.at[idx_1].set(new_v1d)
                else:
                    r_dep = rng.random()
                    ch    = rng.random()
                    if r_dep < p_dep:
                        if ch < THIRD:
                            sv_out[idx_0], sv_out[idx_1] = sv_out[idx_1].copy(), sv_out[idx_0].copy()
                        elif ch < 2.0 * THIRD:
                            v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                            sv_out[idx_0] = -1j * v1
                            sv_out[idx_1] =  1j * v0
                        else:
                            sv_out[idx_1] = -sv_out[idx_1]
                    r_damp = rng.random()
                    v0, v1 = sv_out[idx_0].copy(), sv_out[idx_1].copy()
                    p1 = float(np.clip(p_damp * np.sum(np.abs(v1) ** 2), 0.0, 1.0))
                    if r_damp < p1:
                        norm_decay = np.sqrt(max(p1, 1e-15))
                        sv_out[idx_0] = v1 * np.sqrt(p_damp) / norm_decay
                        sv_out[idx_1] = 0.0
                    else:
                        norm_no_decay = np.sqrt(max(1.0 - p1, 1e-15))
                        sv_out[idx_0] = v0 / norm_no_decay
                        sv_out[idx_1] = v1 * np.sqrt(1.0 - p_damp) / norm_no_decay

        # ── normalise ─────────────────────────────────────────────────
        if is_jax:
            norm = jnp.linalg.norm(sv_out)
            return sv_out / (norm + 1e-15)
        else:
            norm = np.linalg.norm(sv_out)
            return sv_out / (norm + 1e-15)

    @staticmethod
    def kraus_description(model: str) -> Dict:
        desc = {
            'ideal': {
                'kraus': 1,
                'formula': 'Kβ‚€ = I',
                'physical': 'No noise',
            },
            'depolarizing': {
                'kraus': 4,
                'formula': 'Kβ‚€=√(1-p)I  K₁=√(p/3)X  Kβ‚‚=√(p/3)Y  K₃=√(p/3)Z',
                'physical': 'Isotropic Pauli error β€” equiprobable X, Y, Z',
            },
            'bitflip': {
                'kraus': 2,
                'formula': 'Kβ‚€=√(1-p)I  K₁=√pΒ·X',
                'physical': 'Bit flip Οƒ_x with probability p',
            },
            'phaseflip': {
                'kraus': 2,
                'formula': 'Kβ‚€=√(1-p)I  K₁=√pΒ·Z',
                'physical': 'Pure dephasing Οƒ_z with probability p',
            },
            'amplitude_damping': {
                'kraus': 2,
                'formula': 'Kβ‚€=diag(1,√(1-Ξ³))  K₁=[[0,√γ],[0,0]]',
                'physical': 'T₁ energy relaxation |1βŸ©β†’|0⟩ with rate Ξ³',
            },
            'combined': {
                'kraus': 6,
                'formula': 'Depolarizing(p/2) ∘ AmplitudeDamping(p/3)',
                'physical': 'Worst-case NISQ: dephasing + relaxation',
            },
        }
        return desc.get(model, desc['ideal'])