Allow users to define custom priors

[drbenvincent]

High level goals

• The API should include default priors and not require the user to define their own priors.
• But the user should be able to customise their own priors.
• Users should be able to ask the experiments what the default priors are, and use that as part of a workflow in going from default priors to customising priors for specific situations.
• Use of custom priors should be documented not just in the API docs, but we should also have a number of examples sprinkled in to multiple (existing) example notebooks.

Implementation

We can perhaps learn from pymc-marketing which started off allowing users to specify priors by providing dicts but which has moved to having a prior class (see pymc-labs/pymc-marketing#759). We may find that dicts are fine for our purposes, but there could be advantages of going down the path of using classes. (Tagging @wd60622)

[drbenvincent]

See the discussion here pymc-devs/pymc#7416 on whether to port the Prior class from pymc-marketing into the main pymc-repo. This would be amazing because we could add this new functionality with very little change to CausalPy itself.

[williambdean]

Based on the discussion, it seems that the Prior class is most likely best suited for pymc-experimental. However, I don't have a timeline for that just yet.

The code is a single file so it would be easy to copy and tweak to your liking.

With access to the Prior class, the custom priors doesn't seems very difficult since the dims are already being specified in the classes.
Some things to consider:

• pass the configuration into the class keeps consistency with previous API. I've been using dict[str, Prior] mainly but using keywords is alsso good way to restrict to only the allowed variables. A check for dict keys is also possible.
• It's possible to add some additional checks to ensure that the custom prior will work with the model.
  • dim name works. dist = Prior("Normal", mu=0, sigma=1, dims="coeffs"); if dist.dims != ("coeffs", ): raise ...
  • restrict the type of the distribution. dist = Prior("Normal", mu=0, sigma=1, dims="coeffs"); if dist.distribution != "Normal": raise ...
  • etc

Based on the old API, it might look like this:

from pymc_marketing.prior import Prior
from causalpy.pymc_models import WeightedSumFitter

priors = {
    "beta": Prior("Dirichlet", a=[1, 2, 3, 4], dims="coeffs"), 
    "sigma": Prior("Gamma", mu=1.5, sigma=0.25),
}
model = WeightedSumFitter(priors=priors)

# Check the priors
print(model.priors)

model.fit(X, y)

[drbenvincent]

Note to self: Ideally we wait for pymc-devs/pymc-extras#448 (maybe done by @williambdean) then we use that :)