import matplotlib.pyplot as plt
from pymc_marketing.mmm.events import HalfGaussianBasis
from pymc_extras.prior import Prior
half_gaussian = HalfGaussianBasis(
    priors={
        "sigma": Prior("Gamma", mu=[3, 4], sigma=1, dims="event"),
    }
)
coords = {"event": ["PyData-Berlin", "PyCon-Finland"]}
prior = half_gaussian.sample_prior(coords=coords)
curve = half_gaussian.sample_curve(prior)
fig, axes = half_gaussian.plot_curve(
    curve, subplot_kwargs={"figsize": (6, 3), "sharey": True}
)
for ax in axes:
    ax.set_xlabel("")
plt.show()