Hierarchical model in Tensorflow probability using JointDistributionSequential

Question

I'm trying to understand how to implement the following model in Tensorflow probability.

1. An angle, theta, has uniform prior probability in the range [-pi / 2, +pi / 2];
2. A direction-flip probability, beta, has uniform prior probability in range [0, 1];
3. theta' is set to either:
   • theta' = theta + pi with probability beta; or
   • theta' = theta with probability (1 - beta);
4. Concentration c has a HalfCauchy prior probability; and
5. An observation, alpha is drawn from the von Mises distribution, centred on theta' with concentration c.

So far what I've tried is

import tensorflow_probability as tfp
import numpy as np
tfd = tfp.distributions

model = tfd.JointDistributionSequential(
    [
        tfd.Uniform(-np.pi / 2, +np.pi / 2, name='theta'), # theta
        tfd.Uniform(0.0, 1.0, name='beta'), # beta
        tfd.HalfCauchy(loc=0, scale=1), # c
        lambda c, beta, theta: tfd.VonMises(
            loc=theta + np.pi * tfd.Binomial(probs=beta),
            concentration=c,
            name='observed'
        ), # Observation, alpha
    ]
)

Calling this gives an error on the binomial part: TypeError: __init__() missing 1 required positional argument: 'total_count'. What am I doing wrong?

Updated 2020-03-17

Latest code is as follows. I'm still trying to find out how to implement part (3) of my model, i.e. flip the direction of my angle, theta, by adding pi with probability beta. Any help on this would be appreciated! What I have so far doesn't work because I can't multiply the Bernoulli object by a float.

model = tfd.JointDistributionSequential(
    [
        tfd.Uniform(-np.pi / 2, +np.pi / 2, name='theta'), # theta
        tfd.Uniform(0.0, 1.0, name='beta'), # beta
        tfd.HalfCauchy(loc=0, scale=1), # c
        lambda c, beta, theta: tfd.VonMises(
            loc=theta + np.pi * tfd.Bernoulli(probs=beta, dtype=tf.float32),
            concentration=c,
            name='observed'
        ), # Observation, alpha
    ]
)

Answer 1

The problem of multiplying a number by a distribution in the loc parameter calculation can be fixed by sampling from the Bernoulli distribution i.e.

...
loc=theta + np.pi*tfd.Bernoulli(probs=beta, dtype=tf.float32).sample(),
...

This allows sampling from the joint distribution but I'm not sure that it's correct.

Another approach is to pull out the flip random variable and scale using a bijector i.e.

tpb = tfp.bijectors
model = tfd.JointDistributionSequential(
[
    tfd.Uniform(-np.pi / 2, +np.pi / 2, name='theta'), # theta
    tfd.Uniform(0.0, 1.0, name='beta'), # beta
    lambda beta, theta: tfb.Scale(np.pi)(tfd.Bernoulli(probs=beta, dtype=tf.float32)), #flip
    tfd.HalfCauchy(loc=0, scale=1), # c
    lambda c, flip, beta, theta: tfd.VonMises(
        loc=theta + flip ,
        concentration=c,
        name='observed'
    ), # Observation, alpha
]

)

This also allows sampling from the joint dictribution and has the advantage of being able to see when flips happen.