Declare priors using loop results in bracket nesting level exceeded maximum

@jessegrabowski: Thank you for your prompt reply :slight_smile:

I am interested in a branching process. It starts with a single particle that first accumulates y and then divides into two new particles. After the division, the y of the parent is divided equally between the two daughter particles.

A minimal model following this logic is given below:

import numpy as np
import pymc as pm


def get_keys(N_generation):
    """ Generate a binary Galton-Watson branching process (BP).

    The process starts with one particle 'R', upon dying it gives birth to two
    new particles 'R0' and 'R1', each behaving the same is its mother...

    """
    keys = []
    waiting_list = ['R']
    while waiting_list:
        key = waiting_list.pop(0)
        keys.append(key)
        for daughter in ['0', '1']:
            if len(key) < N_generation:
                waiting_list.append(key + daughter)
    return keys

def get_lifetime(key):
    """ Return lifetime for given particle (key).

    Assume that each particles lifetime is 1. So the first particle 'R' is
    observed from t=0 to t=1. It then give birth to 'R0' and 'R1', both living
    one time step, i.e. t=1 to t=2. And so on...

    """
    return np.linspace(len(key) - 1, len(key), 100)

def model(t, mu, y_0):
    """ Exponential growth model."""
    t_0 = np.min(t)
    return y_0 * np.exp(mu * (t - t_0))

def get_y_0(key, mu, y_0_R):
    """ Return y_0 for the likelihood.

    Note, as the initial condition is inherited, it depends on all growth rate
    of its ancestors and the intial condition 'y_0_R' of 'R'.

    """
    if key == 'R':
        y_0 = y_0_R
    else:
        mother = key[:-1]
        # y of mother at its last timepoint is splitted equally between
        # both daughters
        y_0 = 0.5 * model(t=get_lifetime(key=mother),
                          mu=mu[mother],
                          y_0=get_y_0(key=mother,
                                      mu=mu,
                                      y_0_R=y_0_R))[-1]
    return y_0

# get keys for BP with 7 generations
# Note, starting from 7 gen and higher, the code generates the error
keys = get_keys(N_generation=7)

# generate observed data
y_observed = {}
for key in keys:
    t = get_lifetime(key)

    # the initial value of the particle is inherited by the mother particle
    # in case of the mother assume y0 = 1
    if key == 'R':
        y_0 = 1
    else:
        mother = key[:-1]
        # Note, as keys are sorted, the mother entry always exists
        y_0 = 0.5 * y_observed[mother][-1]

    # for simplicity, assume that all particles follow the same exponential
    # growth model with normal noise and have the same true growth rate
    # mu = log(2) and sigma = 0.2
    y_observed[key] = np.random.normal(loc=model(t=t,
                                                 mu=np.log(2),
                                                 y_0=y_0),
                                       scale=0.2)


with pm.Model() as pymc_model:
    mu, y, y_data = {}, {}, {}
    # assume that sigma should be shared among all particles
    sigma = pm.Normal('sigma')
    y_0_R = pm.Uniform('y_0_R', 0, 2)

    #  add a rv (mu) for each particle
    for key in keys:
        mu[key] = pm.Uniform(f'mu_{key}', 0, 2)
    #  add the likelihood of the observation (y) for each particle
    for key in keys:
        # get the initial condition for the given particle
        y_0 = get_y_0(key, mu, y_0_R)

        # Expected value of outcome
        y_expected = model(t=get_lifetime(key),
                           mu=mu[key],
                           y_0=y_0)
        # Here, we create a 'pm.ConstantData' containers to hold the
        # observations, as thisa container is passed to the 'observed'
        y_data[key] = pm.ConstantData('y_data_' + key, y_observed[key])

        y[key] = pm.Normal(f'y_{key}',
                           mu=y_expected,
                           sigma=sigma,
                           observed=y_data[key])
    pymc_trace = pm.sample()