import numpy as np
import matplotlib.pyplot as plt

from beast.physicsmodel.prior_weights_stars import compute_metallicity_prior_weights

fig, ax = plt.subplots()

# met grid with linear spacing
mets = np.linspace(0.004, 0.03)

met_prior_models = [{"name": "flat"},]

for mp_mod in met_prior_models:
    ax.plot(mets, compute_metallicity_prior_weights(mets, mp_mod), label=mp_mod["name"])

ax.set_ylabel("probability")
ax.set_xlabel("metallicity")
ax.legend(loc="best")
plt.tight_layout()
plt.show()