Getting max treedepth iterations for parallelized but not standard model

Modeling
, ,
4

In case it’s of use to anyone, following these other posts 1 2 I modified the reduce_sum code as follows, to compute the linear predictor inside reduce_sum, and the model is now on par with map-rect, with much less coding complexity:

functions {
  real partial_sum(array[] real Response_slice, int start, int end, vector Stimulation,
                   array[] int neuronId, vector neuronBaseline, vector neuronEffect,
                   real sigmaResidual) {
      vector[end-start+1] mu = neuronBaseline[neuronId[start:end]] +
        Stimulation[start:end] .* neuronEffect[neuronId[start:end]];
    return normal_lpdf(Response_slice | mu, sigmaResidual);
  }
}

data {
  int<lower=1> nNeurons;  // Number of groups
  int<lower=1> N; // Number of total datapoints
  vector[N] Stimulation; // Stimulation indicator
  array[N] real Response; // Response data
  array[N] int<lower=1> neuronId; // Neuron indices for each observation
  int grainsize; // Number of observations per thread
}

parameters {
  real meanBaseline;      // Average baseline FR
  real meanStimEffect;      // Increment of FR with stimulation
  vector[nNeurons] neuronBaseline;  // Neuron-specific baseline FR
  vector[nNeurons] neuronEffect;  // Neuron-specific effect of stimulation
  real<lower=0> sigmaBaseline;     // Standard deviation of baseline FR
  real<lower=0> sigmaEffect;     // Standard deviation of stimulation effect
  real<lower=0> sigmaResidual;  // Standard deviation of residuals
}

model {
  // Priors
  meanBaseline ~ normal(0, 30);
  sigmaBaseline ~ normal(0, 5);
  neuronBaseline ~ normal(meanBaseline, sigmaBaseline);
  meanStimEffect ~ normal(0, 30);
  sigmaEffect ~ normal(0, 5);
  neuronEffect ~ normal(meanStimEffect, sigmaEffect);
  sigmaResidual ~ normal(0, 5);
  target += reduce_sum(partial_sum, Response, grainsize, Stimulation, neuronId,
    neuronBaseline, neuronEffect, sigmaResidual);
}
3 Likes