import numpy as np
import matplotlib.pyplot as plt

from shone.chemistry import FastchemWrapper

pressure = np.geomspace(1e-6, 10, 15)  # [bar]
temperature = 2300 * (pressure / 0.1) ** 0.1  # [K]

chem = FastchemWrapper(
    temperature, pressure,
    # these are in linear (not log) units:
    metallicity=1,  # M/H
    c_to_o_ratio=1  # C/O
)
species_table = chem.get_species()
vmr = chem.vmr()

species = ['H2O1', 'O1Ti1', 'e-']
indices = species_table.loc[species]['index']
names = species_table.loc[species]['name']

ax = plt.gca()
ax.loglog(vmr[:, indices], pressure, label=names)
ax.legend(loc='lower left')
ax.invert_yaxis()
ax.set(
    xlabel='Volume mixing ratio',
    ylabel='Pressure [bar]'
)