Probability with TensorFlow#

While TensorFlow offers some support for statistical inference, TensorFlow-Probability is very strong at this and provides MCMC methods, probability distributions and more.

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import zfit
from zfit import z

Distributions#

There is a whole collection of different distributions to be found in TFP. They have a minimal and well designed interface, which is similar to the SciPy distributions.

tfd = tfp.distributions
cauchy = tfd.Cauchy(loc=1., scale=10.)
sample = cauchy.sample(10)
cauchy.prob(sample)
<tf.Tensor: shape=(10,), dtype=float32, numpy=
array([0.01464302, 0.00850127, 0.02412639, 0.02055683, 0.02073182,
       0.03043594, 0.03124683, 0.00549949, 0.00826702, 0.02922291],
      dtype=float32)>

Mixtures of PDFs#

TensorFlow-Probability also supports creating mixtures of different distributions.

mix = 0.3
mix_gauss_cauchy = tfd.Mixture(
  cat=tfd.Categorical(probs=[mix, 1.-mix]),
  components=[
    cauchy,
    tfd.Normal(loc=+1., scale=0.5),
])
sample_mixed = mix_gauss_cauchy.sample(10)
mix_gauss_cauchy.prob(sample_mixed)
<tf.Tensor: shape=(10,), dtype=float32, numpy=
array([3.7576836e-01, 5.5732805e-01, 4.8493055e-01, 2.4973764e-01,
       5.4685426e-01, 7.1766470e-03, 3.8791910e-01, 9.9788621e-02,
       1.7146308e-04, 2.9484937e-03], dtype=float32)>

Joint distributions#

Furthermore, joint distributions of multiple variables are supported.

joint = tfd.JointDistributionNamed(dict(
    c=             tfd.Cauchy(loc=10., scale=1.),
    n=             tfd.Normal(loc=0, scale=2.),
    m=lambda n, c: tfd.Normal(loc=n, scale=c),
))
sample_joint = joint.sample(10)
sample_joint
{'n': <tf.Tensor: shape=(10,), dtype=float32, numpy=
 array([-0.5834539 , -0.14766367, -0.6320378 , -2.3588283 , -1.189753  ,
        -0.09000712,  1.9659712 , -0.8830064 , -1.8737171 , -1.6860838 ],
       dtype=float32)>,
 'c': <tf.Tensor: shape=(10,), dtype=float32, numpy=
 array([ 6.942772 , 11.169306 , 12.133305 , 10.727866 ,  8.620009 ,
         8.341726 , 10.108339 , 10.023032 ,  7.0252953, 10.317402 ],
       dtype=float32)>,
 'm': <tf.Tensor: shape=(10,), dtype=float32, numpy=
 array([ -4.6998634 , -10.244301  ,  -2.2449243 ,  -6.4879284 ,
          0.99081635,  -5.7633195 ,   9.00574   ,   7.6700816 ,
         -5.720553  ,  -3.9813764 ], dtype=float32)>}
joint.prob(sample_joint)
<tf.Tensor: shape=(10,), dtype=float32, numpy=
array([0.00028346, 0.00063495, 0.00035462, 0.00071494, 0.000821  ,
       0.00064193, 0.00119881, 0.00159205, 0.00020318, 0.00152513],
      dtype=float32)>

How TFP compares to zfit#

TensorFlow-Probability offers a great choice of distributions to build a model. The flexibility in terms of vectorization and parametrization is larger than in zfit. However, they only provide models with analytically known CDFs and lack any numerical normalization or sampling methods. This excludes any more sophisticated model, convolutions and more.

Internally, zfit simply wraps TFP distributions for certain implementations, such as the Gauss. There is also a standard wrapper, WrapDistribution, that allows to easily wrap any TFP distribution and use it in zfit.