import numpy as np
import matplotlib.pyplot as plt

from shone.opacity import Opacity

# load the tiny opacity archive:
tiny_opacity = Opacity.load_demo_species('H2O')

# get a jitted interpolator:
interp_opacity = tiny_opacity.get_interpolator()

# get opacity at several temperatures:
wavelength = np.geomspace(0.6, 5, 500)  # [µm]
temperature = np.geomspace(100, 3000, 5)  # [K]
pressure = np.ones_like(temperature)  # [bar]

kappa = interp_opacity(wavelength, temperature, pressure)

n = len(temperature)
ax = plt.gca()

for i in range(n):
    color = plt.cm.plasma(i / n)
    label = f"{temperature[i]:.0f} K"
    ax.semilogy(wavelength, kappa[i], label=label, color=color)

ax.legend(title='Temperature', loc='lower right', framealpha=1)
ax.set(
    xlabel='Wavelength [µm]',
    ylabel='Opacity [cm$^2$ g$^{-1}$]',
    title="Demo opacity: H$_2$O"
)