import numpy as np
import matplotlib.pyplot as plt
import astropy.units as u

from beast.physicsmodel.priormodel import PriorDistanceModel

fig, ax = plt.subplots()

# met grid with linear spacing
dists = np.arange(50., 70, 0.1) * 1e3

distance_prior_models = [
  {"name": "flat"},
  {"name": "absexponential",
    "dist0": 60.0*u.kpc,
    "tau": 5.*u.kpc,
    "amp": 1.0}
  ]

for dp_mod in distance_prior_models:
    pmod = PriorDistanceModel(dp_mod)
    ax.plot(dists, pmod(dists), label=dp_mod["name"])

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