Custom multivariate Dirichlet priors in pymc3 - pymc3

Question

Is it possible create custom multivariate distributions in pymc3? In the following, I have tried to create a linear transformation of a Dirichlet distribution. All variants on this have returned numerous errors, perhaps to do with theano data types? Any help would be gratefully appreciated.
import numpy as np
import pymc3 as pymc
import theano.tensor as tt
# data
n = 5
prior_params = np.ones(n - 1) / (n - 1)
mx = np.array([[0.25 , 0.5 , 0.75 , 1. ],
[0.25 , 0.333, 0.25 , 0. ],
[0.25 , 0.167, 0. , 0. ],
[0.25 , 0. , 0. , 0. ]])
# Note that the matrix mx takes the unit simplex into the unit simplex.
# custom log-liklihood
def generate_function(mx, prior_params):
def log_trunc_dir(x):
return pymc.Dirichlet.dist(a=prior_params).logp(mx.dot(x.T)).eval()
return log_trunc_dir
#model
with pymc.Model() as simple_model:
x = pymc.Dirichlet('x', a=np.ones(n - 1))
q = pymc.DensityDist('q', generate_function(mx, prior_params), observed={'x': x})

Answer 1

Thanks to significant help from the PyMC3 development community, I can post the
following working example of a customised Dirichlet prior in PyMC3.
import pymc3 as pm
import numpy as np
import scipy.special as special
import theano.tensor as tt
import matplotlib.pyplot as plt
n = 4
with pm.Model() as model:
prior = np.ones(n) / n
def dirich_logpdf(value=prior):
return -n * special.gammaln(1/n) + (-1 + 1/n) * tt.log(value).sum()
stick = pm.distributions.transforms.StickBreaking()
probs = pm.DensityDist('probs', dirich_logpdf, shape=n,
testval=np.array(prior), transform=stick)
data = np.array([5, 7, 1, 0])
sfs_obs = pm.Multinomial('sfs_obs', n=np.sum(data), p=probs, observed=data)
with model:
step = pm.Metropolis()
trace = pm.sample(100000, tune=10000, step=step)
print('MLE = ', data / np.sum(data))
print(pm.summary(trace))
pm.traceplot(trace, [probs])
plt.show()