Using pm.DensityDist and customized likelihood with a black-box function

Questions
8

Hi, I implemented the tt.Op extension and it works with the metropolis sampler. However, I’m a bit stuck on getting the gradient to work (so I can use NUTS). It returns an error saying the object returned the wrong number of gradient terms. For my case, the inputs are 3 vector variables, so total number of free parameter is 3*len(ABC). As a test I set them to just 3 scalars, so the input to the Op class has a length of 3 and gradient also has a length of 3, so I’m not sure why Theano is complaining.

### Log likelihood class with gradient ####
class LogLikelihood(theano.tensor.Op):
    #itypes = [tt.dvector, tt.dvector, tt.dvector]
    itypes = [tt.scalar, tt.dscalar, tt.dscalar]
    otypes = [tt.fscalar]

    def __init__(self, fm):
        self.fm = fm
        self.grad = LoglikeGrad(self.fm)
        super(LogLikelihood, self).__init__()

    def perform(self,node,inputs,outputs):

        A, B, C = inputs
        model = np.copy(self.fm.model)
        model = np.array([A,B,C]).T
        self.fm.setModel(model)
        likelihood = self.fm.getOF()
        outputs[0][0] = np.array(likelihood)

    def grad(self, inputs, g):
        A, B, C = inputs
        return [g[0]*self.grad(A,B,C)]

### Gradient class  ###
class LoglikeGrad(theano.tensor.Op):
    #itypes = [tt.dvector, tt.dvector, tt.dvector]
    itypes = [tt.scalar, tt.dscalar, tt.dscalar]
    otypes = [tt.fvector]
    
    def __init__(self, fm):
        self.fm = fm
        super(LoglikeGrad, self).__init__()

    def perform(self,node,inputs,outputs):
        A, B, C = inputs
        model = np.copy(self.fm.model)
        model = np.array([A,B,C]).T
        self.fm.setModel(model)
        self.fm.buildFrechet()
        grads = self.fm.einSum()
        outputs[0][0] = grads

### Instantiate logLike tt.Op class
fm = myOwnClass(data,ABC_init)
logP = LogLikelihood(fm)

## Set up probablistic model
with pm.Model():
    
    A = pm.Normal('A', mu=ABC_init[-1,1], sd=200.0)
    B = pm.Normal('B', mu=ABC_init[-1,2], sd=100.0)
    C = pm.Normal('C', mu=ABC_init[-1,3], sd=0.1)
    
    # Custom likelihood
    likelihood = pm.DensityDist('likelihood', lambda A,B,C: logP(A,B,C), observed={'A':A,'B':B,'C':C})
    
    # Sample from posterior
    trace = pm.sample(draws=5000)