joint_log_probability.py

# We calculate the joint log probability.
# It answers the big and central question in each Markov chain simulation: What's the probability
# that this data occurs together with these distribution parameters?
@tf.function(input_signature=5 * (tf.TensorSpec(shape=[], dtype=tf.float32),))
def joint_log_prob(total_rock, total_paper, total_scissors, p_rock, p_paper):
    '''
    Joint log probability of data occuring together with given parameters.

    :param total_rock: number of rock occurences
    :param total_paper: number of paper occurences
    :param total_scissors: number of scissors occurences
    :param p_rock: probability of rock
    :param p_paper: probability of paper
    '''
    total_count = total_rock + total_paper + total_scissors 
    # Start by defining a prior for the rock / stone / scissor outcomes.
    # The three parameters are proportional to the prior probabilities of rock, paper, and scissors, respectively.
    # We want to start with a situation in which all combinations of probabilities for 
    # rock/paper/scissors are equally likely.
    # Dirichlet is the analogue of having a beta prior in the case of only two options instead
    # of three. In particular, Dirichlet([1, 1]) is a uniform distribution.
    rv_rsp_prior = tfd.Dirichlet([1., 1., 1.], name='rv_rsp_prior')  
    rv_rsp_outcomes = tfd.Multinomial(total_count=total_count, probs=[p_rock, p_paper, 
                                      tf.constant(1., tf.float32) - p_rock - p_paper],
                                      name='rv_rsp_outcomes')

    # calculate prior log probability of parameters
    joint_log_prob = rv_rsp_prior.log_prob([p_rock, p_paper, tf.constant(1., tf.float32) - p_rock - p_paper])
    
    # add log likelihood of data
    joint_log_prob = tf.add(joint_log_prob, rv_rsp_outcomes.log_prob([total_rock, total_paper, total_scissors]))

    return joint_log_prob