Coverage for qpdk / models / __init__.py: 97%

1"""Model definitions for qpdk.""" 

2 

3# ruff: noqa: E402 

4 

5import jax 

6import sax 

7 

8from qpdk.models.capacitor import ( 

9 interdigital_capacitor, 

10 plate_capacitor, 

11) 

12from qpdk.models.inductor import ( 

13 lumped_element_resonator, 

14 meander_inductor, 

15 meander_inductor_inductance_analytical, 

16) 

17from qpdk.models.junction import josephson_junction, squid_junction 

18 

19sax.set_port_naming_strategy("optical") 

20 

21from qpdk.models.constants import ( 

22 DEFAULT_FREQUENCY, 

23) 

24from qpdk.models.couplers import ( 

25 coupler_ring, 

26 coupler_straight, 

27 cpw_cpw_coupling_capacitance, 

28) 

29from qpdk.models.cpw import ( 

30 cpw_epsilon_eff, 

31 cpw_parameters, 

32 cpw_thickness_correction, 

33 cpw_z0, 

34 cpw_z0_from_cross_section, 

35 microstrip_epsilon_eff, 

36 microstrip_thickness_correction, 

37 microstrip_z0, 

38 propagation_constant, 

39 transmission_line_s_params, 

40) 

41from qpdk.models.generic import ( 

42 admittance, 

43 capacitor, 

44 electrical_open, 

45 electrical_short, 

46 electrical_short_2_port, 

47 gamma_0_load, 

48 impedance, 

49 inductor, 

50 lc_resonator, 

51 lc_resonator_coupled, 

52 open, 

53 series_impedance, 

54 short, 

55 short_2_port, 

56 shunt_admittance, 

57) 

58from qpdk.models.perturbation import ( 

59 dispersive_shift, 

60 dispersive_shift_to_coupling, 

61 ej_ec_to_frequency_and_anharmonicity, 

62 measurement_induced_dephasing, 

63 purcell_decay_rate, 

64 resonator_linewidth_from_q, 

65) 

66from qpdk.models.qubit import ( 

67 coupling_strength_to_capacitance, 

68 double_island_transmon, 

69 double_island_transmon_with_bbox, 

70 double_island_transmon_with_resonator, 

71 double_pad_transmon, 

72 double_pad_transmon_with_bbox, 

73 double_pad_transmon_with_resonator, 

74 ec_to_capacitance, 

75 ej_to_inductance, 

76 el_to_inductance, 

77 flipmon, 

78 flipmon_with_bbox, 

79 flipmon_with_resonator, 

80 fluxonium, 

81 fluxonium_coupled, 

82 fluxonium_with_bbox, 

83 fluxonium_with_resonator, 

84 qubit_with_resonator, 

85 shunted_transmon, 

86 transmon_coupled, 

87 transmon_with_resonator, 

88 xmon_transmon, 

89) 

90from qpdk.models.resonator import ( 

91 quarter_wave_resonator_coupled, 

92 resonator, 

93 resonator_coupled, 

94 resonator_frequency, 

95 resonator_half_wave, 

96 resonator_quarter_wave, 

97) 

98from qpdk.models.unimon import ( 

99 el_to_arm_inductance, 

100 unimon_coupled, 

101 unimon_energies, 

102 unimon_frequency_and_anharmonicity, 

103 unimon_hamiltonian, 

104) 

105from qpdk.models.waveguides import ( 

106 airbridge, 

107 bend_circular, 

108 bend_euler, 

109 bend_s, 

110 indium_bump, 

111 launcher, 

112 nxn, 

113 rectangle, 

114 straight, 

115 straight_double_open, 

116 straight_microstrip, 

117 straight_open, 

118 straight_shorted, 

119 taper_cross_section, 

120 tee, 

121 tsv, 

122) 

123 

124__all__ = [ 

125 "DEFAULT_FREQUENCY", 

126 "admittance", 

127 "airbridge", 

128 "bend_circular", 

129 "bend_euler", 

130 "bend_s", 

131 "capacitor", 

132 "coupler_ring", 

133 "coupler_straight", 

134 "coupling_strength_to_capacitance", 

135 "cpw_cpw_coupling_capacitance", 

136 "cpw_epsilon_eff", 

137 "cpw_parameters", 

138 "cpw_thickness_correction", 

139 "cpw_z0", 

140 "cpw_z0_from_cross_section", 

141 "dispersive_shift", 

142 "dispersive_shift_to_coupling", 

143 "double_island_transmon", 

144 "double_island_transmon_with_bbox", 

145 "double_island_transmon_with_resonator", 

146 "double_pad_transmon", 

147 "double_pad_transmon_with_bbox", 

148 "double_pad_transmon_with_resonator", 

149 "ec_to_capacitance", 

150 "ej_ec_to_frequency_and_anharmonicity", 

151 "ej_to_inductance", 

152 "el_to_arm_inductance", 

153 "el_to_inductance", 

154 "electrical_open", 

155 "electrical_short", 

156 "electrical_short_2_port", 

157 "flipmon", 

158 "flipmon_with_bbox", 

159 "flipmon_with_resonator", 

160 "fluxonium", 

161 "fluxonium_coupled", 

162 "fluxonium_with_bbox", 

163 "fluxonium_with_resonator", 

164 "gamma_0_load", 

165 "impedance", 

166 "indium_bump", 

167 "inductor", 

168 "interdigital_capacitor", 

169 "josephson_junction", 

170 "launcher", 

171 "lc_resonator", 

172 "lc_resonator_coupled", 

173 "lumped_element_resonator", 

174 "meander_inductor", 

175 "meander_inductor_inductance_analytical", 

176 "measurement_induced_dephasing", 

177 "microstrip_epsilon_eff", 

178 "microstrip_thickness_correction", 

179 "microstrip_z0", 

180 "models", 

181 "nxn", 

182 "open", 

183 "plate_capacitor", 

184 "propagation_constant", 

185 "purcell_decay_rate", 

186 "quarter_wave_resonator_coupled", 

187 "qubit_with_resonator", 

188 "rectangle", 

189 "resonator", 

190 "resonator_coupled", 

191 "resonator_frequency", 

192 "resonator_half_wave", 

193 "resonator_linewidth_from_q", 

194 "resonator_quarter_wave", 

195 "series_impedance", 

196 "short", 

197 "short_2_port", 

198 "shunt_admittance", 

199 "shunted_transmon", 

200 "squid_junction", 

201 "straight", 

202 "straight_double_open", 

203 "straight_microstrip", 

204 "straight_open", 

205 "straight_shorted", 

206 "taper_cross_section", 

207 "tee", 

208 "transmission_line_s_params", 

209 "transmon_coupled", 

210 "transmon_with_resonator", 

211 "tsv", 

212 "unimon_coupled", 

213 "unimon_energies", 

214 "unimon_frequency_and_anharmonicity", 

215 "unimon_hamiltonian", 

216 "xmon_transmon", 

217] 

218 

219 

220def _is_sax_model(obj: object) -> bool: 

221 """Check if an object is a SAX model function.""" 

222 if not callable(obj): 

223 return False 

224 

225 # Skip functions that return jax.Array or tuple (these are Hamiltonian models) 

226 if hasattr(obj, "__annotations__") and "return" in obj.__annotations__: 

227 ret = obj.__annotations__["return"] 

228 # If it returns jax.Array or a tuple of floats, it's not a SAX S-parameter model 

229 if ret in {jax.Array, "jax.Array"}: 

230 return False 

231 # Match names like tuple[float, float] 

232 if hasattr(ret, "__name__") and ret.__name__ == "tuple": 

233 return False 

234 if str(ret).startswith("tuple") or str(ret).startswith("Tuple"): 

235 return False 

236 

237 # Check if return type is sax.SType or sax.SDict 

238 for target in [obj, getattr(obj, "__wrapped__", None)]: 

239 if target is None: 

240 continue 

241 if hasattr(target, "__annotations__") and "return" in target.__annotations__: 

242 ret = target.__annotations__["return"] 

243 if hasattr(ret, "__name__") and ret.__name__ in {"SType", "SDict"}: 

244 return True 

245 # Also check identity for standard cases 

246 if ret in {sax.SType, sax.SDict}: 

247 return True 

248 return sax.try_into[sax.Model](obj) is not None 

249 

250 

251models = { 

252 k: v 

253 for k, v in globals().items() 

254 if k in __all__ and k != "models" and _is_sax_model(v) 

255}