Coverage for jaxquantum/devices/common/utils.py: 0%

18 statements  

« prev     ^ index     » next       coverage.py v7.9.2, created at 2025-07-17 21:51 +0000

1"""Utility functions""" 

2 

3from scipy.special import pbdv 

4from scipy import constants 

5 

6import jax.scipy as jsp 

7import jax.numpy as jnp 

8 

9 

10def factorial_approx(n): 

11 return jsp.special.gamma(n + 1) 

12 

13 

14# physics utils 

15 

16# ---------------- 

17 

18 

19def harm_osc_wavefunction(n, x, l_osc): 

20 r""" 

21 Taken from scqubits... not jit-able 

22 

23 For given quantum number n=0,1,2,... return the value of the harmonic 

24 oscillator wave function :math:`\psi_n(x) = N H_n(x/l_{osc}) \exp(-x^2/2l_\text{ 

25 osc})`, N being the proper normalization factor. 

26 

27 Directly uses `scipy.special.pbdv` (implementation of the parabolic cylinder 

28 function) to mitigate numerical stability issues with the more commonly used 

29 expression in terms of a Gaussian and a Hermite polynomial factor. 

30 

31 Parameters 

32 ---------- 

33 n: 

34 index of wave function, n=0 is ground state 

35 x: 

36 coordinate(s) where wave function is evaluated 

37 l_osc: 

38 oscillator length, defined via <0|x^2|0> = l_osc^2/2 

39 

40 Returns 

41 ------- 

42 value of harmonic oscillator wave function 

43 """ 

44 x = 2 * jnp.pi * x 

45 result = pbdv(n, jnp.sqrt(2.0) * x / l_osc)[0] 

46 result = result / jnp.sqrt(l_osc * jnp.sqrt(jnp.pi) * factorial_approx(n)) 

47 return result 

48 

49 

50def calculate_lambda_over_four_resonator_zpf(freq, impedance): 

51 expected_Z0 = impedance # Ohms 

52 expected_E_L_over_E_C = (1 / (4 * expected_Z0)) ** 2 * ( 

53 constants.h**2 / (8 * constants.e**4) 

54 ) 

55 desired_E_C = jnp.sqrt(freq**2 / expected_E_L_over_E_C / 8) 

56 desired_E_L = freq**2 / desired_E_C / 8 

57 storage_q_zpf = (1 / 32 * desired_E_L / desired_E_C) ** (1 / 4) 

58 return storage_q_zpf, desired_E_C, desired_E_L