PYMC3: у NUTS есть трудности с выборкой из иерархической модели с нулевым раздувом

Question

Я пытаюсь повторить анализ данных из статьи Ричарда МакЭлрита, в которой он снабдил данные иерархической гамма-моделью с нулевым раздуванием. Данные об охоте возвращаются около 15000 Охотничьи поездки с о 150 охотники старше двадцати лет. Поскольку многие охотничьи поездки имеют нулевой доход, модель предполагает, что каждая поездка имеет pi вероятность нулевого возврата и 1 - pi вероятность положительных результатов, которые следуют за гамма-распределением с параметрами alpha а также beta,

Переменной предиктора является возраст, модель использует полином возраста (до 3-го порядка) для моделирования pi а также alpha, И так как 15000 поездки принадлежат 150 отдельные охотники, у каждого охотника есть свои коэффициенты, и все коэффициенты следуют общему многомерному нормальному распределению. Для получения подробной информации о модели, пожалуйста, обратитесь к следующему коду. Спецификация модели, похоже, в порядке, но NUTS испытывает затруднения с началом выборки: она дает только около 10 образцов через 20 минут, и пробоотборник просто остановился и сказал мне, что для завершения отбора проб потребуются сотни часов. Я хочу знать, что вызывает проблемы.

Обычный импорт

import pymc3 as pm
import numpy as np
from pymc3.distributions import Continuous, Gamma
import theano.tensor as tt

Данные могут быть получены из GitHub

n_trip = len(d)
n_hunter = len(d['hunter.id'].unique())
idx_hunter = d['hunter.id'].values

y = d['kg.meat'].values
age = d['age.s'].values
age2 = (d['age.s'].values)**2
age3 = (d['age.s'].values)**3

Функция плотности логарифмической вероятности для ноль-завышенной гаммы.

class ZeroInflatedGamma(Continuous):
    def __init__(self, alpha, beta, pi, *args, **kwargs):
        super(ZeroInflatedGamma, self).__init__(*args, **kwargs)
        self.alpha = alpha
        self.beta = beta
        self.pi = pi = tt.as_tensor_variable(pi)
        self.gamma = Gamma.dist(alpha, beta)

    def logp(self, value):
        return tt.switch(value > 0,
                         tt.log(1 - self.pi) + self.gamma.logp(value),
                         tt.log(self.pi))

Это матрица для индексации корреляционной матрицы до матрицы 9X9, LKJ до pymc3 задается как одномерный вектор

dim = 9
n_elem = dim * (dim - 1) / 2
tri_index = np.zeros([dim, dim], dtype=int)
tri_index[np.triu_indices(dim, k=1)] = np.arange(n_elem)
tri_index[np.triu_indices(dim, k=1)[::-1]] = np.arange(n_elem)

А вот и модель

with pm.Model() as Vary9_model:

    # hyper-priors
    mu_a = pm.Normal('mu_a', mu=0, sd=100, shape=9)
    sigma_a = pm.HalfCauchy('sigma_a', 5, shape=9)

    # build the covariance matrix
    C_triu = pm.LKJCorr('C_triu', n=2, p=9)    
    C = tt.fill_diagonal(C_triu[tri_index], 1)
    sigma_diag = tt.nlinalg.diag(sigma_a)
    cov = tt.nlinalg.matrix_dot(sigma_diag, C, sigma_diag)

    # priors for each hunter and all the linear components, 9 dimensional Gaussian  
    a = pm.MvNormal('a', mu=mu_a, cov=cov, shape=(n_hunter, 9))

    # linear function  
    mupi = a[:,0][idx_hunter] + a[:,1][idx_hunter] * age + a[:,2][idx_hunter] * age2 + a[:,3][idx_hunter] * age3
    mualpha = a[:,4][idx_hunter] + a[:,5][idx_hunter] * age + a[:,6][idx_hunter] * age2 + a[:,7][idx_hunter] * age3

    pi = pm.Deterministic('pi', pm.math.sigmoid(mupi))
    alpha = pm.Deterministic('alpha', pm.math.exp(mualpha))
    beta = pm.Deterministic('beta', pm.math.exp(a[:,8][idx_hunter]))

    y_obs = ZeroInflatedGamma('y_obs', alpha, beta, pi, observed=y)

    Vary9_trace = pm.sample(6000, njobs=2)

И это статус модели:

Auto-assigning NUTS sampler...
Initializing NUTS using advi...
Average ELBO = -28,366: 100%|██████████| 200000/200000 [15:36<00:00, 213.57it/s]
Finished [100%]: Average ELBO = -28,365
  0%|          | 22/6000 [15:51<63:49:25, 38.44s/it]

У меня есть некоторые мысли по поводу проблемы, но я не уверен, в чем может быть причина.

• Является ли девятимерный гауссов слишком трудным для выборки? Я ранее только моделировал перехваты для mualpha а также mupi как двумерный гауссов, он медленный, но работал (примерка модели заняла около 20 минут)

• это плотность вероятности, которая вызывает проблему? Я сам написал функцию плотности и не уверен, хорошо ли она работает. Я думаю, что функция плотности не дифференцируется в нуле, не вызовет ли это проблемы для пробоотборника гаек?

• потому что переменные предиктора сильно коррелированы? Компоненты линейной модели в этой модели являются полиномами от возраста до третьей степени, и, естественно, предикторы имеют высокую корреляцию.

А может это из-за чего-то еще?

В качестве примечания я попытался использовать Metropolis Сэмплер, мой компьютер исчерпал память, и цепи все еще не сходятся.

Answer 1

ZeroInflatedGamma выглядит хорошо. Функция плотности дифференцируема по отношению к пи, альфа и бета. Это все, что вам нужно для наблюдаемой переменной. Производные производные по стоимости нужны только в том случае, если вы пытаетесь оценить значения.

В реализации LKJCorr возникла проблема: https://github.com/pymc-devs/pymc3/pull/1863 Вы можете попробовать еще раз на master. К сожалению, pymc3 не поддерживает использование MVNormal и LKJCorr в разложенной параметризации cholesky. Это тоже может помочь. На github есть запрос на выполнение работ в данный момент: https://github.com/pymc-devs/pymc3/pull/1875

Чтобы улучшить сходимость, вы можете попробовать нецентрированную параметризацию для a, Нечто подобное

a_raw = pm.Normal('a_raw', shape=(9, n_hunter))
a = mu_a[None, :] + tt.dot(tt.slinalg.cholesky(cov), a_raw)

Конечно, это было бы быстрее, если бы у нас был этот холкий LKJCorr...