import matplotlib.pyplot as plt
plt.clf()
from tf_pwa.utils import plot_particle_model
axis = plot_particle_model("BWR", {"mass": 0.1})
axis = plot_particle_model("BWR2", {"mass": 0.1}, axis=axis)
axis = plot_particle_model("BWR_below", {"mass": 0.1}, axis=axis)
axis = plot_particle_model("BWR_coupling", {"mass": 0.1}, axis=axis)
leg = axis[2].legend()