def _base_dist(self, mu: TensorLike, alpha: TensorLike, *args, **kwargs):
total_count = mu + alpha
probs = mu / (mu + alpha)
return tfd.NegativeBinomial(total_count=total_count,
probs=probs,
*args,
**kwargs)
def _init_distribution(conditions, **kwargs):
return tfd.Mixture(
cat=tfd.Categorical(
probs=[1.0 - conditions["psi"], conditions["psi"]]),
components=[
tfd.Deterministic(loc=tf.zeros_like(conditions["mu"])),
tfd.NegativeBinomial(
total_count=conditions["alpha"],
probs=(conditions["mu"]) /
(conditions["mu"] + conditions["alpha"]),
),
],
**kwargs,
)
def _base_dist(self, *args, **kwargs):
"""
Zero-inflated negative binomial base distribution.
A ZeroInflatedNegativeBinomial is a mixture between a deterministic
distribution and a NegativeBinomial distribution.
"""
mix = kwargs.pop("mix")
return tfd.Mixture(
cat=tfd.Categorical(probs=[mix, 1.0 - mix]),
components=[
tfd.Deterministic(0.0),
tfd.NegativeBinomial(*args, **kwargs)
],
name="ZeroInflatedNegativeBinomial",
)
def _init_distribution(conditions):
mu, alpha = conditions["mu"], conditions["alpha"]
return tfd.NegativeBinomial(total_count=mu + alpha,
probs=mu / (mu + alpha))
def _init_distribution(conditions):
total_count, probs = conditions["total_count"], conditions["probs"]
return tfd.NegativeBinomial(total_count=total_count, probs=probs)
rate=tf.exp(theta["log_lambda"])),
pi=theta["pi"])
},
"negative binomial": {
"parameters": {
"p": {
"support": [0, 1],
"activation function": sigmoid
},
"log_r": {
"support": [-10, 10],
"activation function": identity
}
},
"class":
lambda theta: tensorflow_distributions.NegativeBinomial(
total_count=tf.exp(theta["log_r"]), probs=theta["p"])
},
"zero-inflated negative binomial": {
"parameters": {
"pi": {
"support": [0, 1],
"activation function": sigmoid
},
"p": {
"support": [0, 1],
"activation function": sigmoid
},
"log_r": {
"support": [-10, 10],
"activation function": identity
}
)
},
"negative binomial": {
"parameters": {
"p": {
"support": [0, 1],
"activation function": sigmoid
},
"log_r": {
"support": [-10, 10],
"activation function": identity
}
},
"class": lambda theta: tensorflow_distributions.NegativeBinomial(
total_count = tf.exp(theta["log_r"]),
probs = theta["p"]
)
},
"zero-inflated negative binomial": {
"parameters": {
"pi": {
"support": [0, 1],
"activation function": sigmoid
},
"p": {
"support": [0, 1],
"activation function": sigmoid
},
"log_r": {
"support": [-10, 10],