import pandas as pd
import numpy as np
# working for SUEWS v2019a: supy_driver >= 2019a2
from supy_driver import atmmoiststab_module as atm

# create z/L range
ser_zl = pd.Series(np.arange(-5, 2, .2))

# base DataFrame with z/L and scheme columns
df_zl_psi = pd.DataFrame({i: ser_zl
                          for i in [0, 2, 3, 4]}).rename({
                              0: '$z/L$'
                          }, axis=1).set_index(r'$z/L$')
# momentum
df_zl_psi_mom = df_zl_psi.apply(lambda x: x.apply(lambda zl: atm.stab_psi_mom(
    x.name, zl, zl)))
# heat and vapour
df_zl_psi_heat = df_zl_psi.apply(lambda x: x.apply(
    lambda zl: atm.stab_psi_heat(x.name, zl, zl)))

# plot Psi_m
ax = df_zl_psi_mom.plot(logy=False)
_ = ax.set_ylabel(r'$\Psi_m$')

# plot Psi_h
ax = df_zl_psi_heat.plot(logy=False)
_ = ax.set_ylabel(r'$\Psi_h$')