This is fantastic, thank you!
Initial runs playing with your code throws me an error if you supply a single threshold value in the priors:
I’ve had a dig at fixing this line, but cant spot what it should be changed to:
threshold = model.parameters["threshold"].prior
From what I can tell the parameters just dont exist in the model in this way. Do we need to extend the idata to contain the prior and take from there?
I’ve also attempted to edit your example into using a negative binomial as the base distribution as thats what I’m using. It seems to be working (aside the issue above) but I’ve done something wrong in the family definition as I now need to specify a prior for alpha every run, rather than letting the back end in Bambi grab a default for me. I’m also in a bit of a kerfuffle about handling the survival function given the difference in parameterisation between pymc and numpy.
Any help is greatly appreciated.
Mike
import numpy as np
import pandas as pd
import pymc as pm
import xarray as xr
from scipy.stats import nbinom
import bambi as bmb
class NegativeBinomialWithThreshold(pm.NegativeBinomial):
"""A NegativeBinomial distribution that records ``response > threshold`` as a deterministic."""
##Modified from fantastic forum post here: https://discourse.pymc.io/t/bambi-count-data-and-thresholds/17926
rv_op = pm.NegativeBinomial.rv_op
threshold_name = "TestThreshold"
def __new__(cls, name, mu, threshold, *args, **kwargs):
dims = kwargs.get("dims")
response = super().__new__(cls, name, mu, threshold, *args, **kwargs)
pm.Deterministic(cls.threshold_name, response > threshold, dims=dims)
return response
@classmethod
def dist(cls, mu, threshold, *args, **kwargs):
"""Return the unchanged NegativeBinomial distribution used for prediction and log-likelihood."""
del threshold
return super().dist( mu, *args, **kwargs)
likelihood = bmb.Likelihood(
name="NegativeBinomialWithThreshold",
params=["mu", "threshold","alpha"],
parent="mu",
dist=NegativeBinomialWithThreshold,
)
custom_negativebinomial = bmb.Family(
name="custom_negativebinomial",
likelihood=likelihood,
link={"mu": "log", "threshold": "identity"},
)
custom_negativebinomial.set_default_priors({"threshold": 5})
def probability_above_threshold(model, idata):
"""Return an ``interpret`` transform for ``P(Y > threshold | mu)``.
This works with either a constant threshold or posterior draws from a threshold prior.
"""
if "threshold" in idata.posterior:
threshold = idata.posterior["threshold"]
else:
threshold = model.parameters["threshold"].prior
alpha=idata.posterior["alpha"]
def transform(mu):
# For a count Y, Y > t is equivalent to Y > floor(t).
return xr.apply_ufunc(nbinom.sf, np.floor(threshold), mu,1.0/(1.0+alpha*mu))
#since scipi uses k n and p and we use mu and alpha need to transform the parameters here.
##This doesnt feel correct to me though. n=1.0/alpha in place of the mu?
return transform
def run_example(draws=500, tune=500):
"""Fit the custom family and obtain predictions with ``bmb.interpret``.
Returns
-------
dict
The data, fitted model, posterior draws, prediction summaries, and plots.
"""
rng = np.random.default_rng(1234)
x = rng.normal(size=150)
data = pd.DataFrame({"x": x, "y": rng.negative_binomial(len(x),0.5+x/1500)})
# Alternatives for the auxiliary parameter are:
#
# - Omit "threshold" to use the family's default constant, 5.
# - Pass another constant with priors={"threshold": 7}.
# - Pass a prior, as below, to propagate uncertainty about the threshold.
#Due to edits alpha prior is currently required. Needs investigating.
threshold_prior = bmb.Prior("Normal", mu=5, sigma=0.75)
alpha_prior=bmb.Prior("Exponential", lam=1)
model = bmb.Model(
"y ~ x",
data=data,
family=custom_negativebinomial,
priors={"threshold": threshold_prior, "alpha":alpha_prior},
)
idata = model.fit(
draws=draws,
tune=tune,
random_seed=1234,
)
x_grid = np.linspace(data["x"].min(), data["x"].max(), 30)
conditional = {"x": x_grid}
# Expected Poisson count, E[Y | x] = mu.
mean_predictions = bmb.interpret.predictions(
model,
idata,
conditional=conditional,
)
# Posterior predictive counts.
count_predictions = bmb.interpret.predictions(
model,
idata,
conditional=conditional,
target="y",
)
# Transform the posterior expected count into P(Y > threshold | x), using the Poisson survival
# function. This propagates uncertainty in both mu and threshold.
probability_transform = probability_above_threshold(model, idata)
over_threshold_predictions = bmb.interpret.predictions(
model,
idata,
conditional=conditional,
transforms={"y": probability_transform},
)
mean_plot = bmb.interpret.plot_predictions(
model,
idata,
conditional=conditional,
)
over_threshold_plot = bmb.interpret.plot_predictions(
model,
idata,
conditional=conditional,
transforms={"y": probability_transform},
)
return {
"data": data,
"model": model,
"idata": idata,
"mean_predictions": mean_predictions,
"count_predictions": count_predictions,
"over_threshold_predictions": over_threshold_predictions,
"mean_plot": mean_plot,
"over_threshold_plot": over_threshold_plot,
}
if __name__ == "__main__":
example = run_example()
print(example["mean_predictions"].summary)
print(example["over_threshold_predictions"].summary)
example["mean_plot"].show()
example["over_threshold_plot"].show()
