Coverage for jaxquantum/circuits/library/oscillator.py: 92%
218 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Oscillator gates."""
3from jaxquantum.core.operators import (displace, basis, destroy, create, num)
4from jaxquantum.circuits.gates import Gate
5from jaxquantum.circuits.channels import (
6 apply_elementwise_channel,
7 apply_shifted_channel,
8)
9from jax.scipy.special import gammaln
10import jax.numpy as jnp
11from jaxquantum import Qarray
12from jaxquantum.utils import hermgauss
13from functools import partial
14import jax
16def diag_expm(diag_matrix):
17 """Computes expm of a diagonal matrix efficiently (O(N) instead of O(N^3))."""
18 # Extract diagonal, exponentiate elements, put back on diagonal
19 return jnp.diag(jnp.exp(jnp.diagonal(diag_matrix)))
23def D(N, alpha, ts=None, c_ops=None):
24 """Displacement gate.
26 Args:
27 N: Hilbert space dimension.
28 alpha: Displacement amplitude.
29 ts: Optional time array for hamiltonian simulation.
30 c_ops: Optional collapse operators.
32 Returns:
33 Displacement gate.
34 """
35 gen_Ht = None
36 if ts is not None:
37 delta_t = ts[-1] - ts[0]
38 amp = 1j * alpha / delta_t
39 a = destroy(N)
40 gen_Ht = lambda params: (lambda t: jnp.conj(amp) * a + amp * a.dag())
42 return Gate.create(
43 N,
44 name="D",
45 params={"alpha": alpha},
46 gen_U=lambda params: displace(N, params["alpha"]),
47 gen_Ht=gen_Ht,
48 ts=ts,
49 gen_c_ops=lambda params: Qarray.from_list([]) if c_ops is None else c_ops,
50 num_modes=1,
51 )
54def _conditional_displacement(N, beta, echoed=False):
55 displacement = displace(N, beta / 2).data
56 inverse = jnp.swapaxes(jnp.conj(displacement), -1, -2)
57 blocks = jnp.zeros(
58 displacement.shape[:-2] + (2, N, 2, N),
59 dtype=displacement.dtype,
60 )
61 if echoed:
62 blocks = blocks.at[..., 1, :, 0, :].set(displacement)
63 blocks = blocks.at[..., 0, :, 1, :].set(inverse)
64 else:
65 blocks = blocks.at[..., 0, :, 0, :].set(displacement)
66 blocks = blocks.at[..., 1, :, 1, :].set(inverse)
67 return Qarray.create(blocks.reshape(blocks.shape[:-4] + (2 * N, 2 * N)),
68 dims=[[2, N], [2, N]])
71def CD(N, beta, ts=None):
72 """Conditional displacement gate.
74 Args:
75 N: Hilbert space dimension.
76 beta: Conditional displacement amplitude.
77 ts: Optional time sequence for hamiltonian simulation.
79 Returns:
80 Conditional displacement gate.
81 """
82 gen_Ht = None
83 if ts is not None:
84 g = basis(2, 0)
85 e = basis(2, 1)
86 gg = g @ g.dag()
87 ee = e @ e.dag()
88 delta_t = ts[-1] - ts[0]
89 amp = 1j * beta / delta_t / 2
90 a = destroy(N)
91 gen_Ht = lambda params: lambda t: (
92 gg
93 ^ (jnp.conj(amp) * a + amp * a.dag()) + ee
94 ^ (jnp.conj(-amp) * a + (-amp) * a.dag())
95 )
97 return Gate.create(
98 [2, N],
99 name="CD",
100 params={"beta": beta},
101 gen_U=lambda params: _conditional_displacement(N, params["beta"]),
102 gen_Ht=gen_Ht,
103 ts=ts,
104 num_modes=2,
105 )
108def ECD(N, beta, ts=None):
109 """Echoed conditional displacement gate.
111 Args:
112 N: Hilbert space dimension.
113 beta: Conditional displacement amplitude.
114 ts: Optional time sequence for hamiltonian simulation.
116 Returns:
117 Echoed conditional displacement gate.
118 """
119 return Gate.create(
120 [2, N],
121 name="ECD",
122 params={"beta": beta},
123 gen_U=lambda params: _conditional_displacement(
124 N,
125 params["beta"],
126 echoed=True,
127 ),
128 gen_Ht=None,
129 ts=ts,
130 num_modes=2,
131 )
133def CR(N, theta):
134 """Conditional rotation gate.
136 Args:
137 N: Hilbert space dimension.
138 theta: Conditional rotation angle.
140 Returns:
141 Conditional rotation gate.
142 """
143 g = basis(2, 0)
144 e = basis(2, 1)
146 gg = g @ g.dag()
147 ee = e @ e.dag()
150 return Gate.create(
151 [2, N],
152 name="CR",
153 params={"theta": theta},
154 gen_U=lambda params: (gg ^ (-1.j*theta/2*create(N)@destroy(N)).expm())
155 + (ee ^ (1.j*theta/2*create(N)@destroy(N)).expm()),
156 num_modes=2,
157 )
160# --- 2. Optimized Kernels (Using diag_expm) ---
162@partial(jax.jit, static_argnames=["N", "max_l"])
163def _Amp_Damp_Kraus_Map_JIT(N, err_prob, max_l):
164 coefficients = _amp_damp_coefficients(
165 N,
166 err_prob,
167 max_l,
168 truncate=False,
169 )
170 indices = jnp.arange(N)
171 sources = jnp.minimum(
172 indices + jnp.arange(coefficients.shape[-2])[:, None],
173 N - 1,
174 )
175 return _shifted_kraus_maps(coefficients, sources)
178def _amp_damp_coefficients(
179 dimension,
180 probability,
181 max_l,
182 truncate=True,
183):
184 probability = jnp.asarray(probability)
185 indices = jnp.arange(dimension)
186 order = min(max_l, dimension - 1) if truncate else max_l
187 losses = jnp.arange(order + 1)
188 retention = jnp.sqrt(1.0 - probability)
189 log_binomial = (
190 gammaln(indices + losses[:, None] + 1)
191 - gammaln(indices + 1)
192 - gammaln(losses[:, None] + 1)
193 )
194 coefficients = (
195 jnp.exp(0.5 * log_binomial)
196 * jnp.power(probability[..., None, None], 0.5 * losses[:, None])
197 * jnp.power(retention[..., None, None], indices)
198 )
199 return jnp.where(
200 indices + losses[:, None] < dimension,
201 coefficients,
202 0,
203 )
206def _shifted_kraus_maps(coefficients, sources):
207 coefficients = jnp.moveaxis(coefficients, -2, 0)
208 indices = jnp.arange(coefficients.shape[-1])
210 def build(coefficient, source):
211 shape = coefficient.shape[:-1] + (coefficient.shape[-1],) * 2
212 return jnp.zeros(shape, coefficient.dtype).at[..., indices, source].set(
213 coefficient
214 )
216 return jax.vmap(build)(coefficients, sources)
219def _amp_damp_apply(rho, params, dimension, max_l):
220 coefficients = _amp_damp_coefficients(
221 dimension,
222 params["err_prob"],
223 max_l,
224 )
225 return apply_shifted_channel(
226 rho,
227 {
228 "_coefficients": coefficients,
229 "_shifts": jnp.arange(coefficients.shape[-2]),
230 },
231 )
234def Amp_Damp(N, err_prob, max_l):
235 if max_l < 0:
236 raise ValueError("max_l must be non-negative")
238 def kmap(params):
239 data = _Amp_Damp_Kraus_Map_JIT(
240 params["N"],
241 params["err_prob"],
242 params["max_l"],
243 )
244 return Qarray.create(data, dims=[[N], [N]], bdims=data.shape[:-2])
246 return Gate.create(
247 N,
248 name="Amp_Damp",
249 params={"err_prob": err_prob, "max_l": max_l, "N": N},
250 gen_KM=kmap,
251 channel_apply=partial(
252 _amp_damp_apply,
253 dimension=N,
254 max_l=max_l,
255 ),
256 lazy_kraus=True,
257 num_modes=1,
258 )
261@partial(jax.jit, static_argnames=["N", "max_l"])
262def _Amp_Gain_Kraus_Map_JIT(N, err_prob, max_l):
263 coefficients = _amp_gain_coefficients(
264 N,
265 err_prob,
266 max_l,
267 truncate=False,
268 )
269 indices = jnp.arange(N)
270 sources = jnp.maximum(
271 indices - jnp.arange(coefficients.shape[-2])[:, None],
272 0,
273 )
274 return _shifted_kraus_maps(coefficients, sources)
277def _amp_gain_coefficients(
278 dimension,
279 probability,
280 max_l,
281 truncate=True,
282):
283 probability = jnp.asarray(probability)
284 indices = jnp.arange(dimension)
285 order = min(max_l, dimension - 1) if truncate else max_l
286 gains = jnp.arange(order + 1)
287 source = indices - gains[:, None]
288 log_binomial = (
289 gammaln(indices + 1)
290 - gammaln(jnp.maximum(source, 0) + 1)
291 - gammaln(gains[:, None] + 1)
292 )
293 coefficients = (
294 jnp.exp(0.5 * log_binomial)
295 * jnp.power(probability[..., None, None], 0.5 * gains[:, None])
296 * jnp.power(
297 jnp.sqrt(1.0 - probability)[..., None, None],
298 jnp.maximum(source, 0),
299 )
300 )
301 return jnp.where(source >= 0, coefficients, 0)
304def _amp_gain_apply(rho, params, dimension, max_l):
305 coefficients = _amp_gain_coefficients(
306 dimension,
307 params["err_prob"],
308 max_l,
309 )
310 return apply_shifted_channel(
311 rho,
312 {
313 "_coefficients": coefficients,
314 "_shifts": -jnp.arange(coefficients.shape[-2]),
315 },
316 )
319def Amp_Gain(N, err_prob, max_l):
320 if max_l < 0:
321 raise ValueError("max_l must be non-negative")
323 def kmap(params):
324 data = _Amp_Gain_Kraus_Map_JIT(
325 params["N"],
326 params["err_prob"],
327 params["max_l"],
328 )
329 return Qarray.create(data, dims=[[N], [N]], bdims=data.shape[:-2])
331 return Gate.create(
332 N,
333 name="Amp_Gain",
334 params={"err_prob": err_prob, "max_l": max_l, "N": N},
335 gen_KM=kmap,
336 channel_apply=partial(
337 _amp_gain_apply,
338 dimension=N,
339 max_l=max_l,
340 ),
341 lazy_kraus=True,
342 num_modes=1,
343 )
346@partial(jax.jit, static_argnames=["N", "max_l"])
347def _Thermal_Ch_Kraus_Map_JIT(N, err_prob, n_bar, max_l):
348 coefficients, gains, losses = _thermal_coefficients(
349 N,
350 err_prob,
351 n_bar,
352 max_l,
353 truncate=False,
354 )
355 indices = jnp.arange(N)
356 sources = jnp.clip(
357 indices + losses[:, None] - gains[:, None],
358 0,
359 N - 1,
360 )
361 return _shifted_kraus_maps(coefficients, sources)
364def _thermal_coefficients(
365 dimension,
366 probability,
367 n_bar,
368 max_l,
369 truncate=True,
370):
371 probability = jnp.asarray(probability)
372 n_bar = jnp.asarray(n_bar)
373 order = (
374 min(max_l, dimension - 1) + 1
375 if truncate
376 else max_l + 1
377 )
378 pair_indices = jnp.arange(order**2)
379 gains = pair_indices // order
380 losses = pair_indices % order
381 output = jnp.arange(dimension)
382 source = output + losses[:, None] - gains[:, None]
383 valid = (source >= 0) & (source < dimension) & (
384 output + losses[:, None] < dimension
385 )
386 source_safe = jnp.clip(source, 0, dimension - 1)
387 prefactor = jnp.sqrt(
388 jnp.power(
389 probability[..., None, None] * (1.0 + n_bar[..., None, None]),
390 losses[:, None],
391 )
392 * jnp.power(
393 probability[..., None, None] * n_bar[..., None, None],
394 gains[:, None],
395 )
396 / (
397 jnp.exp(gammaln(losses[:, None] + 1))
398 * jnp.exp(gammaln(gains[:, None] + 1))
399 )
400 )
401 ratio = jnp.exp(
402 gammaln(output + losses[:, None] + 1)
403 - 0.5
404 * (
405 gammaln(source_safe + 1)
406 + gammaln(output + 1)
407 )
408 )
409 coefficients = (
410 prefactor
411 * ratio
412 * jnp.power(
413 jnp.sqrt(1.0 - probability)[..., None, None],
414 output,
415 )
416 )
417 return jnp.where(valid, coefficients, 0), gains, losses
420def _thermal_apply(rho, params, dimension, max_l):
421 coefficients, gains, losses = _thermal_coefficients(
422 dimension,
423 params["err_prob"],
424 params["n_bar"],
425 max_l,
426 )
427 return apply_shifted_channel(
428 rho,
429 {
430 "_coefficients": coefficients,
431 "_shifts": losses - gains,
432 },
433 )
436def Thermal_Ch(N, err_prob, n_bar, max_l):
437 if max_l < 0:
438 raise ValueError("max_l must be non-negative")
440 def kmap(params):
441 data = _Thermal_Ch_Kraus_Map_JIT(
442 params["N"],
443 params["err_prob"],
444 params["n_bar"],
445 params["max_l"],
446 )
447 return Qarray.create(data, dims=[[N], [N]], bdims=data.shape[:-2])
449 return Gate.create(
450 N,
451 name="Thermal_Ch",
452 params={
453 "err_prob": err_prob,
454 "n_bar": n_bar,
455 "max_l": max_l,
456 "N": N,
457 },
458 gen_KM=kmap,
459 channel_apply=partial(
460 _thermal_apply,
461 dimension=N,
462 max_l=max_l,
463 ),
464 lazy_kraus=True,
465 num_modes=1,
466 )
469@partial(jax.jit, static_argnames=["N", "max_l"])
470def _Dephasing_Ch_Kraus_Map_JIT(N, ws, phis, max_l):
471 def compute_op(w, phi):
472 diagonal = jnp.exp(1.0j * phi[..., None] * jnp.arange(N))
473 return jnp.sqrt(w) * diagonal[..., None, :] * jnp.eye(N)
474 return jax.vmap(compute_op)(ws, phis)
477def _dephasing_factor(dimension, probability, nodes, weights):
478 indices = jnp.arange(dimension)
479 delta = indices[:, None] - indices[None, :]
480 phases = jnp.sqrt(2 * jnp.asarray(probability)[..., None]) * nodes
481 return jnp.sum(
482 weights[:, None, None]
483 * jnp.exp(1.0j * phases[..., :, None, None] * delta),
484 axis=-3,
485 )
488def _dephasing_apply(rho, params, dimension):
489 return apply_elementwise_channel(
490 rho,
491 {
492 "_factor": _dephasing_factor(
493 dimension,
494 params["err_prob"],
495 params["_nodes"],
496 params["_weights"],
497 )
498 },
499 )
502def Dephasing_Ch(N, err_prob, max_l):
503 if max_l < 1:
504 raise ValueError("max_l must be positive")
505 xs, ws_raw = hermgauss(max_l)
506 ws = 1/jnp.sqrt(jnp.pi)*ws_raw
508 def kmap(params):
509 phases = (
510 jnp.sqrt(2 * params["err_prob"])[..., None]
511 * params["_nodes"]
512 )
513 data = _Dephasing_Ch_Kraus_Map_JIT(
514 params["N"],
515 params["_weights"],
516 jnp.moveaxis(phases, -1, 0),
517 params["max_l"],
518 )
519 return Qarray.create(data, dims=[[N], [N]], bdims=data.shape[:-2])
521 return Gate.create(
522 N,
523 name="Dephasing_Ch",
524 params={
525 "err_prob": err_prob,
526 "max_l": max_l,
527 "N": N,
528 "_nodes": xs,
529 "_weights": ws,
530 },
531 gen_KM=kmap,
532 channel_apply=partial(_dephasing_apply, dimension=N),
533 lazy_kraus=True,
534 num_modes=1,
535 )
538def selfKerr(N, K):
539 a = destroy(N)
540 return Gate.create(
541 N,
542 name="selfKerr",
543 params={"Kerr": K},
544 gen_U=lambda params: (-1.0j * K / 2 * (a.dag() @ a.dag() @ a @ a)).expm(),
545 num_modes=1,
546 )
549@partial(jax.jit, static_argnames=["N", "max_l"])
550def _Dephasing_Reset_Kraus_Map_JIT(N, p, t_rst, chi, max_l):
551 p = jnp.asarray(p)
552 g = basis(2, 0).data
553 e = basis(2, 1).data
554 gg = jnp.outer(g, jnp.conj(g))
555 ee = jnp.outer(e, jnp.conj(e))
556 ge = jnp.outer(g, jnp.conj(e))
558 n_op = num(N).data
559 I_N = jnp.eye(N)
561 exponents = jnp.arange(max_l - 1) / (max_l - 1)
562 raw_weights = jnp.power(p[..., None], exponents)
563 weights = (
564 (1 - p[..., None])
565 * raw_weights
566 / jnp.sum(raw_weights, axis=-1, keepdims=True)
567 )
569 def compute_op(l):
570 def branch_0(_):
571 matrix = jnp.kron(gg, I_N)
572 return jnp.broadcast_to(matrix, p.shape + matrix.shape)
574 def branch_1(_):
575 op_osc = diag_expm(-1.0j * chi * t_rst * n_op)
576 return jnp.sqrt(p)[..., None, None] * jnp.kron(ee, op_osc)
578 def branch_rest(_):
579 exponent = -1.0j * chi * t_rst * (l - 2) / (max_l - 1)
580 op_osc = diag_expm(exponent * n_op)
581 return (
582 jnp.sqrt(weights[..., l - 2])[..., None, None]
583 * jnp.kron(ge, op_osc)
584 )
586 return jax.lax.cond(
587 l == 0,
588 branch_0,
589 lambda _: jax.lax.cond(l == 1, branch_1, branch_rest, operand=None),
590 operand=None
591 )
593 ls = jnp.arange(max_l+1)
594 return jax.vmap(compute_op)(ls)
597def _dephasing_reset_factors(N, p, t_rst, chi, max_l):
598 indices = jnp.arange(N)
599 delta = indices[:, None] - indices[None, :]
600 transfer_indices = jnp.arange(2, max_l + 1)
601 p = jnp.asarray(p)
602 raw_weights = jnp.power(
603 p[..., None],
604 (transfer_indices - 2) / (max_l - 1),
605 )
606 weights = (
607 (1 - p[..., None])
608 * raw_weights
609 / jnp.sum(raw_weights, axis=-1, keepdims=True)
610 )
611 angles = chi * t_rst * (transfer_indices - 2) / (max_l - 1)
612 transfer = jnp.sum(
613 weights[..., :, None, None]
614 * jnp.exp(-1.0j * angles[:, None, None] * delta),
615 axis=-3,
616 )
617 excited = p[..., None, None] * jnp.exp(-1.0j * chi * t_rst * delta)
618 return transfer, excited
621def _apply_dephasing_reset_factors(
622 rho,
623 transfer_factor,
624 excited_factor,
625 dimension,
626 phase_factor=1,
627):
628 blocks = rho.reshape(rho.shape[:-2] + (2, dimension, 2, dimension))
629 ground = phase_factor * (
630 blocks[..., 0, :, 0, :] + transfer_factor * blocks[..., 1, :, 1, :]
631 )
632 excited = phase_factor * excited_factor * blocks[..., 1, :, 1, :]
633 zero = jnp.zeros_like(ground)
634 output = jnp.stack(
635 (
636 jnp.stack((ground, zero), axis=-2),
637 jnp.stack((zero, excited), axis=-2),
638 ),
639 axis=-4,
640 )
641 return output.reshape(output.shape[:-4] + (2 * dimension,) * 2)
644def _dephasing_reset_apply(rho, params, dimension, max_l):
645 """Apply structured ancilla reset directly to density-matrix blocks."""
646 transfer_factor, excited_factor = _dephasing_reset_factors(
647 dimension,
648 params["p"],
649 params["t_rst"],
650 params["chi"],
651 max_l,
652 )
653 return _apply_dephasing_reset_factors(
654 rho,
655 transfer_factor[..., None, :, :],
656 excited_factor[..., None, :, :],
657 dimension,
658 )
661def Dephasing_Reset(N, p, t_rst, chi, max_l):
662 if max_l < 2:
663 raise ValueError("max_l must be at least two")
665 def kmap(params):
666 data = _Dephasing_Reset_Kraus_Map_JIT(
667 params["N"],
668 params["p"],
669 params["t_rst"],
670 params["chi"],
671 params["max_l"],
672 )
673 return Qarray.create(
674 data,
675 dims=[[2, N], [2, N]],
676 bdims=data.shape[:-2],
677 )
679 return Gate.create(
680 [2, N],
681 name="Dephasing_Reset",
682 params={
683 "p": p,
684 "t_rst": t_rst,
685 "chi": chi,
686 "max_l": max_l,
687 "N": N,
688 },
689 gen_KM=kmap,
690 channel_apply=partial(
691 _dephasing_reset_apply,
692 dimension=N,
693 max_l=max_l,
694 ),
695 lazy_kraus=True,
696 num_modes=2,
697 )