First attempt at a hierarchical model with custom likelihood

Hello, I am trying to implement a simple hierarchical model in pymc3. My data is a population of n=1\dots N
different time series. Each time series \mathbf F^{(n)} is modeled with a Gaussian likelihood with a diagonal covariance matrix (which is known), and the mean function is a non-linear function of 5 parameters:

\boldsymbol\theta^{(n)}=\{\Delta F^{(n)},F_b^{(n)},t_0^{(n)} t_\textrm{eff}^{(n)},t_E^{(n)}\}

The joint posterior for all time series parameters and the hyperparameters of the model is given by:

p\left(\{\boldsymbol\theta_n\},\mu,\sigma\lvert\{\mathbf F_n\}, \{\boldsymbol\sigma_n\}\right) \ \propto \prod_{n=1}^N\mathcal{L}\left(\mathbf F^{(n)},\boldsymbol\sigma^{(n)}\lvert \boldsymbol\theta^{(n)}\right)\,p\left(\boldsymbol\theta^{(n)}\lvert \mu, \sigma\right)\,p(\mu,\sigma)

where \{\sigma_n\} are the diagonal elements of the covariance matrix for the n-th time series, and

p\left(\boldsymbol\theta^{(n)}\lvert \mu, \sigma\right)= p\left(\Delta F^{(n)}\right)\,p\left(F_b^{(n)}\right)\, p\left(t_0^{(n)}\right) \,p\left(t_\textrm{eff}^{(n)}\lvert t_E^{(n)}\right)p\left(t_E^{(n)}\lvert \mu, \sigma\right)

I’ve attempted to implement this model in pymc3 as

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

model = pm.Model()

with model:    
    n_population = len(lightcurves)
    
    BoundedNormal = pm.Bound(pm.Normal, lower=0.0) # Half-Normal dist

    # Hyperpriors
    mu_tE = BoundedNormal('mu_tE', mu=100, sd=200.)
    sig_tE = BoundedNormal('sig_tE', mu=20, sd=100.)
        
    # Priors 
    tE = pm.Lognormal('tE', mu=mu_tE, sd=sig_tE, shape=n_population)
    DeltaF = BoundedNormal('DeltaF', mu=[np.max(lc[:, 1]) for lc in lightcurves], 
                           sd=1., shape=n_population)
    Fb = pm.Normal('Fb', mu=0., sd=0.1, shape=n_population)
    t0 = pm.Uniform('t0', [np.min(lc[:, 0]) for lc in lightcurves], 
                    [np.max(lc[:, 0]) for lc in lightcurves], shape=n_population)
    # p(teff|tE)
    def conditional_prior(value):
        teff = T.cast(value, 'float64')
        return -T.log(tE) - (teff/tE)**2/1**2 
    
    teff = pm.DensityDist('teff', conditional_prior, shape=n_population)
    
    # Likelihood 
    def ln_likelihood(lightcurves):
        lnL = T.cast(0., 'float64')

        for n in range(n_population):
            lightcurve = T.cast(lightcurves[n], 'float64')
            t = lightcurve[:, 0]
            F = lightcurve[:, 1]
            sig_F = lightcurve[:, 2]
        
            # Mean model
            u0 = teff/tE
            u = T.sqrt(u0[n]**2 + ((t - t0[n])/tE[n])**2)
            A = lambda u: (u**2 + 2)/(u*T.sqrt(u**2 + 4))

            mu_F = DeltaF[n]*(A(u) - 1)/(A(u0[n]) - 1) + Fb[n]
            
            # Gaussian likelihood
            lnL += T.sum((-(mu_F - F)**2/sig_F**2))
        
        return lnL
    
    likelihood = pm.DensityDist('likelihood', ln_likelihood, observed={'lightcurves':lightcurves})

model.check_test_point()
model.vars

where lightcurves is a list containing N time series, each with a different length.

I have no idea if this is implemented correctly. I am getting an error: ValueError: setting an array element with a sequence..

You need to cast n within the forloop of ln_likelihood into theano using theano.shared (or there is a theano range as well you can try). But more generally I think you need to rewrite the ln_likelihood using theano.scan to make this efficient.