Дискретная байесовская сеть о вероятности тензорного потока, Edward2 и Python

Question

У меня есть простая байесовская сеть:

        state
      /       \
     /         \
    V           V  
signal_1    signal_2

со случайными переменными «состояние», «сигнал_1» и «сигнал_2» с соответствующими дискретными значениями: Val (состояние) = {0, 1, 2, 3}, Val (signal_1) = {0, 1} и Val (signal_2) = {0, 1}.

У меня есть распределение предельной вероятности P (состояние) и распределение условной вероятности P (signal_1 |состояние) и P (сигнал_2 | состояние) в виде таблиц.

Совместная вероятность P (состояние, сигнал_1, сигнал_2) равна P (состояние) * P (сигнал_1 | состояние) * P (сигнал_2 | состояние) и log P (состояние, сигнал_1, сигнал_2) = log P(состояние) + log P (сигнал_1 | состояние) + log P (сигнал_2 | состояние).

Я пытаюсь построить эту модель в TensorFlow Вероятность:

Вариант 1:

import tensorflow as tf
tfd = tfp.distributions

def joint_log_prob(state, signal_1, signal_2):
        state_prob = tf.constant([0.5, 0.1, 0.3, 0.1])
        signal_1_prob = tf.constant([[0.03, 0.97], [0.05, 0.95], [0.20, 0.80], [0.97, 0.03]])
        signal_2_prob = tf.constant([[0.03, 0.97], [0.97, 0.03], [0.97, 0.03], [0.99, 0.01]])

        state_dist = tfd.Categorical(probs=state_prob)
        signal_1_dist = tfd.Categorical(probs=signal_1_prob[state]) # state=0,1,2,3
        signal_2_dist = tfd.Categorical(probs=signal_2_prob[state]) # state=0,1,2,3

        return state_dist.log_prob(state) + signal_1_dist.log_prob(signal_1) + signal_2_dist.log_prob(signal_2)

Моими доказательствами являются, например, «signal_1 = 1» и «signal_2 = 0», и я хочу получить вероятность «состояние = 0», т.е. P (состояние = 0 | сигнал_1 = 1, сигнал_2= 0)

Я определяю функцию:

def unnormalized_log_posterior(state):
    return joint_log_prob(state, 1, 0)

Вариант 2 с помощью Edward2:

import tensorflow as tf
from tensorflow_probability import edward2 as ed

def model():
    state_prob = tf.constant([0.5, 0.1, 0.3, 0.1])
    signal_1_prob = tf.constant([[0.03, 0.97], [0.05, 0.95], [0.20, 0.80], [0.97, 0.03]])
    signal_2_prob = tf.constant([[0.03, 0.97], [0.97, 0.03], [0.97, 0.03], [0.99, 0.01]])

    state = ed.Categorical(probs=state_prob, name='state')
    signal_1 = ed.Categorical(probs=signal_1_prob[state], name='signal_1')
    signal_2 = ed.Categorical(probs=signal_2_prob[state], name='signal_2')

    return [state, signal_1, signal_2] # what does the function have to return? [state, signal_1, signal_2]?

Затем я получаю объединенную функцию логарифмической вероятности

log_joint = ed.make_log_joint_fn(model) # joint log-probability function

, который можно использовать с помощью лямбда-выражения:

lambda x: log_joint(state=x, signal_1=1, signal_2=0)

Правильно ли я определил модель?

Вопрос в том, как мне продолжать использовать MCMC? Какой метод MCMC я могу использовать для этой дискретной модели? Может кто-нибудь показать код MCMC для моей модели? Как я могу сказать алгоритму MCMC, чтобы он брал выборки для "состояния" только из набора {0, 1, 2, 3}?

Спасибо!