Какой степпер Монте-Карло pymc3 я могу использовать для пользовательского категориального распределения?

Я работаю над реализацией скрытых цепей Маркова в pymc3. Я продвинулся довольно далеко в реализации скрытых состояний. Ниже я показываю простую марковскую цепь с двумя состояниями:

import numpy as np
import pymc3 as pm
import theano.tensor as tt

# Markov chain sample with 2 states that was created
# to have prob 0->1 = 0.1 and prob 1->0 = 0.3
sample = np.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
   0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0,
   0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
   0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0,
   1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 1,
   0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 0],
dtype=np.uint8)

Сейчас я определяю класс, описывающий состояния. В качестве входных данных мне нужно знать вероятность P1 перейти из состояния 0 в состояние 1 и P2 перейти из 1-> 0. Мне также нужно знать вероятность PA для первого состояния, равного 0.

class HMMStates(pm.Discrete):
    """
    Hidden Markov Model States
    Parameters
    ----------
    P1 : tensor
        probability to remain in state 1
    P2 : tensor
        probability to move from state 2 to state 1

    """

    def __init__(self, PA=None, P1=None, P2=None,
                 *args, **kwargs):
        super(HMMStates, self).__init__(*args, **kwargs)
        self.PA = PA
        self.P1 = P1
        self.P2 = P2
        self.mean = 0.
        self.mode = tt.cast(0,dtype='int64')

    def logp(self, x):
        PA = self.PA
        P1 = self.P1
        P2 = self.P2

        # now we need to create an array with probabilities
        # so that for x=A: PA=P1, PB=(1-P1)
        # and for x=B: PA=P2, PB=(1-P2)

        choice = tt.stack((P1,P2))
        P = choice[x[:-1]]

        x_i = x[1:]

        ou_like = pm.Categorical.dist(P).logp(x_i)
        return pm.Categorical.dist(PA).logp(x[0]) + tt.sum(ou_like)

Я очень горжусь продвинутыми трюками ниндзя индексирования, которым я научился в группе theano google. Вы также можете реализовать то же самое с помощью tt.switch. В чем я не был слишком уверен, так это в self.mode. Я просто дал ему 0, чтобы избежать ошибки тестового значения. Вот как использовать класс в модели, которая проверяет, работает ли он. В этом случае состояние не скрыто, а наблюдается.

with pm.Model() as model:
    # 2 state model
    # P1 is probablility to stay in state 1
    # P2 is probability to move from state 2 to state 1
    P1 = pm.Dirichlet('P1', a=np.ones(2))
    P2 = pm.Dirichlet('P2', a=np.ones(2))

    PA = pm.Deterministic('PA',P2/(P2+1-P1))

    states = HMMStates('states',PA,P1,P2, observed=sample)

    start = pm.find_MAP()

    trace = pm.sample(5000, start=start)

вывод хорошо воспроизводит данные. В следующей модели я покажу проблему. Здесь я наблюдаю не состояния напрямую, а состояние с добавлением некоторого гауссовского шума (таким образом, скрытое состояние). Если вы запустите модель со степпером Metropolis, произойдет сбой с ошибкой индекса, которую я проследил до проблем, связанных с использованием степпера Metropolis в категориальных распределениях. К сожалению, единственным степпером, который можно было бы применить к моему классу, является степпер CategoricalGibbsMetropolis, но он отказывается работать с моим классом, поскольку явно не является категориальным распределением.

gauss_sample = sample*1.0 + 0.1*np.random.randn(len(sample))
from scipy import optimize
with pm.Model() as model2:
    # 2 state model
    # P1 is probablility to stay in state 1
    # P2 is probability to move from state 2 to state 1
    P1 = pm.Dirichlet('P1', a=np.ones(2))
    P2 = pm.Dirichlet('P2', a=np.ones(2))

    S = pm.InverseGamma('S',alpha=2.1, beta=1.1)

    PA = pm.Deterministic('PA',P2/(P2+1-P1))

    states = HMMStates('states',PA,P1,P2, shape=len(gauss_sample))

    emission = pm.Normal('emission',
                         mu=tt.cast(states,dtype='float64'),
                         sd=S,
                         observed = gauss_sample)

    start2 = pm.find_MAP(fmin=optimize.fmin_powell)
    step1 = pm.Metropolis(vars=[P1, P2, S, PA, emission])
    step2 = pm.ElemwiseCategorical(vars=[states], values=[0,1])
    trace2 = pm.sample(10000, start=start, step=[step1,step2])

ElemwiseCategorical запускает его, но не присваивает правильное значение моим состояниям. Состояния либо все 0, либо все 1 с.

Как я могу указать ElemwiseCategorial назначить вектор состояний 1 и 0, или, альтернативно, как я могу заставить CategorialGibbsMetropolis распознать мое распределение как категориальное. Это должно быть распространенной проблемой с пользовательскими дистрибутивами.