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

1"""Utility functions""" 

2 

3import jax.numpy as jnp 

4import jax.scipy as jsp 

5from jax import lax 

6from scipy import constants 

7 

8 

9def factorial_approx(n): 

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

11 

12 

13# physics utils 

14 

15# ---------------- 

16 

17 

18def harm_osc_wavefunction(n, x, l_osc): 

19 r""" 

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

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

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

23 

24 Parameters 

25 ---------- 

26 n: 

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

28 x: 

29 coordinate(s) where wave function is evaluated 

30 l_osc: 

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

32 

33 Returns 

34 ------- 

35 value of harmonic oscillator wave function 

36 """ 

37 return harm_osc_wavefunctions(n + 1, x, l_osc)[n] 

38 

39 

40def harm_osc_wavefunctions(num_levels, x, l_osc): 

41 """Evaluate the first ``num_levels`` normalized oscillator wavefunctions.""" 

42 if num_levels < 1: 

43 raise ValueError("num_levels must be positive") 

44 

45 coordinate = 2 * jnp.pi * jnp.asarray(x) / l_osc 

46 psi0 = jnp.exp(-(coordinate**2) / 2) / jnp.sqrt(l_osc * jnp.sqrt(jnp.pi)) 

47 if num_levels == 1: 

48 return psi0[None] 

49 

50 psi1 = jnp.sqrt(2.0) * coordinate * psi0 

51 

52 def next_level(carry, level): 

53 previous, current = carry 

54 following = ( 

55 jnp.sqrt(2.0 / (level + 1)) * coordinate * current 

56 - jnp.sqrt(level / (level + 1)) * previous 

57 ) 

58 return (current, following), following 

59 

60 _, remaining = lax.scan( 

61 next_level, 

62 (psi0, psi1), 

63 jnp.arange(1, num_levels - 1), 

64 ) 

65 return jnp.concatenate((psi0[None], psi1[None], remaining), axis=0) 

66 

67 

68def calculate_lambda_over_four_resonator_zpf(freq, impedance): 

69 expected_Z0 = impedance # Ohms 

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

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

72 ) 

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

74 desired_E_L = freq**2 / desired_E_C / 8 

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

76 return storage_q_zpf, desired_E_C, desired_E_L