import numpy as np
import matplotlib.pyplot as plt

from beast.physicsmodel.prior_weights_stars import compute_distance_prior_weights

fig, ax = plt.subplots()

# met grid with linear spacing
dists = np.linspace(8e6, 9e6)

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

for mp_mod in distance_prior_models:
    ax.plot(dists, compute_distance_prior_weights(dists, mp_mod), label=mp_mod["name"])

ax.set_ylabel("probability")
ax.set_xlabel("distance [pc]")
ax.legend(loc="best")
plt.tight_layout()
plt.show()