zinb_example.py

from __future__ import print_function, division, absolute_import

import tensorflow as tf
tf.enable_eager_execution()

from tensorflow_probability.python.distributions import NegativeBinomial
from zero_inflated import ZeroInflated

nb = NegativeBinomial(total_count=10, probs=0.5)
zfnb = ZeroInflated(dist=nb, pi=0.9)

# denoised values
print(nb.mean())

# matching values of imputed data
print(zfnb.mean())
print(zfnb.log_prob(10))

# Note: zfnb.sample is not working recently