import matplotlib.pyplot as plt
from shone.opacity import Opacity, generate_synthetic_opacity

generate_synthetic_opacity()

# load the one opacity file:
opacity = Opacity.load_species_from_name('synthetic')

# get a jitted 3D interpolator over wavelength, temperature, pressure:
interp_opacity = opacity.get_interpolator()

wavelength = np.linspace(1, 5, 500)  # [µm]
temperature = np.array([200, 400, 600, 800])  # [K]
pressure = np.ones_like(temperature)  # [bar]

# interpolate for a range of wavelengths at one pressure and temperature:
temperature = np.array([200, 400, 600, 800])
label = [f"{t} K" for t in temperature]
example_opacity = interp_opacity(wavelength, temperature, pressure)

kappa_cloud = 5e-2  # [cm2/g]

plt.figure()
plt.semilogy(wavelength, example_opacity.T, label=label)
plt.semilogy(wavelength, kappa_cloud * np.ones_like(wavelength), ls='--', label="Cloud")
plt.legend()
plt.gca().set(
    xlabel='Wavelength [µm]',
    ylabel='Opacity, $\kappa$ [cm$^2$ g$^{-1}$]'
)