Coverage for jaxquantum/devices/superconducting/flux_base.py: 32%

1"""Flux base device.""" 

2 

3from abc import abstractmethod 

4 

5import jax.numpy as jnp 

6import matplotlib.pyplot as plt 

7from flax import struct 

8from jax import config 

9 

10from jaxquantum.devices.base.base import BasisTypes, Device 

11from jaxquantum.devices.common.utils import harm_osc_wavefunctions 

12 

13config.update("jax_enable_x64", True) 

14 

15 

16@struct.dataclass 

17class FluxDevice(Device): 

18 @abstractmethod 

19 def phi_zpf(self): 

20 """Return Phase ZPF.""" 

21 

22 def _calculate_wavefunctions_fock(self, phi_vals): 

23 """Calculate wavefunctions at phi_exts.""" 

24 phi_osc = self.phi_zpf() * jnp.sqrt(2) # length of oscillator 

25 basis_functions = harm_osc_wavefunctions( 

26 self.N_pre_diag, 

27 phi_vals, 

28 jnp.real(phi_osc), 

29 ) 

30 return self.get_vec_data_in_H_eigenbasis(basis_functions) 

31 

32 def _calculate_wavefunctions_charge(self, phi_vals): 

33 phi_vals = jnp.array(phi_vals) 

34 n_labels = jnp.diag(self.original_ops["n"].data) 

35 basis_functions = jnp.exp( 

36 -2j * jnp.pi * n_labels[:, None] * phi_vals 

37 ) / jnp.sqrt(2 * jnp.pi) 

38 wavefunctions = self.get_vec_data_in_H_eigenbasis( 

39 basis_functions, 

40 ) 

41 correction = jnp.power(1j, jnp.arange(wavefunctions.shape[-2])) 

42 return correction[:, None] * wavefunctions 

43 

44 @abstractmethod 

45 def potential(self, phi): 

46 """Return potential energy as a function of phi.""" 

47 

48 def plot_wavefunctions(self, phi_vals, max_n=None, which=None, ax=None, mode="abs", ylim=None, y_scale_factor=1, zero_potential=False, wavefunction_color=None): 

49 if self.basis == BasisTypes.fock: 

50 _calculate_wavefunctions = self._calculate_wavefunctions_fock 

51 elif self.basis == BasisTypes.charge: 

52 _calculate_wavefunctions = self._calculate_wavefunctions_charge 

53 else: 

54 raise NotImplementedError( 

55 f"The {self.basis} is not yet supported for plotting wavefunctions." 

56 ) 

57 

58 """Plot wavefunctions at phi_exts.""" 

59 wavefunctions = _calculate_wavefunctions(phi_vals) 

60 energy_levels = self.eig_systems["vals"][: self.N] 

61 

62 potential = self.potential(phi_vals) 

63 

64 min_potential = 0 if not zero_potential else jnp.min(potential) 

65 if ax is None: 

66 fig, ax = plt.subplots(1, 1, figsize=(3.5, 2.5), dpi=1000) 

67 else: 

68 fig = ax.get_figure() 

69 

70 min_val = None 

71 max_val = None 

72 

73 assert max_n is None or which is None, "Can't specify both max_n and which" 

74 

75 max_n = self.N if max_n is None else max_n 

76 levels = range(max_n) if which is None else which 

77 

78 for n in levels: 

79 if mode == "abs": 

80 wf_vals = jnp.abs(wavefunctions[n, :]) ** 2 

81 elif mode == "real": 

82 wf_vals = wavefunctions[n, :].real 

83 elif mode == "imag": 

84 wf_vals = wavefunctions[n, :].imag 

85 

86 wf_vals += energy_levels[n] 

87 curr_min_val = min(wf_vals) 

88 curr_max_val = max(wf_vals) 

89 

90 if min_val is None or curr_min_val < min_val: 

91 min_val = curr_min_val 

92 

93 if max_val is None or curr_max_val > max_val: 

94 max_val = curr_max_val 

95 

96 extra_kwargs = {} 

97 if wavefunction_color is not None: 

98 if isinstance(wavefunction_color, list): 

99 extra_kwargs["color"] = wavefunction_color[n] 

100 else: 

101 extra_kwargs["color"] = wavefunction_color 

102 

103 ax.plot( 

104 phi_vals, (wf_vals - min_potential)*y_scale_factor, label=f"$|${n}$\\rangle$", linestyle="-", linewidth=1, **extra_kwargs 

105 ) 

106 

107 ax.fill_between(phi_vals, (energy_levels[n] - min_potential)*y_scale_factor, (wf_vals - min_potential)*y_scale_factor, alpha=0.5, **extra_kwargs) 

108 

109 ax.plot( 

110 phi_vals, 

111 (potential - min_potential)*y_scale_factor, 

112 label="potential", 

113 color="black", 

114 linestyle="-", 

115 linewidth=1, 

116 ) 

117 

118 ylim = ylim if ylim is not None else [jnp.min(jnp.array([min_val - 1 - min_potential, jnp.min(potential) - min_potential]))*y_scale_factor, (max_val + 1 - min_potential)*y_scale_factor] 

119 ax.set_ylim(ylim) 

120 ax.set_xlabel(r"$\varphi/2\pi$") 

121 ax.set_ylabel(r"Energy [GHz]") 

122 

123 if mode == "abs": 

124 title_str = r"$|\psi_n(\Phi)|^2$" 

125 elif mode == "real": 

126 title_str = r"Re($\psi_n(\Phi)$)" 

127 elif mode == "imag": 

128 title_str = r"Im($\psi_n(\Phi)$)" 

129 

130 ax.set_title(f"{title_str}") 

131 

132 ax.legend(fontsize='xx-small') 

133 fig.tight_layout() 

134 

135 return ax