import jax.numpy as jnp
from qpdk.models import tsv_transition_double_sided

freq = jnp.linspace(1e9, 10e9, 501)
freq_ghz = freq / 1e9
s_model = tsv_transition_double_sided(f=freq)

for (pout, pin), sij in s_model.items():
   if pout >= pin:
       plt.plot(
           freq_ghz,
           20 * jnp.log10(jnp.abs(jnp.broadcast_to(sij, freq.shape))),
           label=f"S[{pout}, {pin}]"
       )
plt.xlabel("Frequency [GHz]", fontsize=12)
plt.ylabel("Magnitude [dB]", fontsize=12)
plt.grid(True)
plt.legend()
plt.show()