An application of spatial Gaussian process

functions {

  matrix kron_mvprod(matrix A, matrix B, matrix V)
  {
    return transpose(A*transpose(B*V));
  }

  matrix gp(int N_rows, int N_columns, real[] rows_idx, real[] columns_index,
  real delta0,
  real alpha_gp,
  real rho_gp1, real rho_gp2,
  matrix z1)
  {

    matrix[N_rows,N_columns] GP;

    matrix[N_rows, N_rows] K1;
    matrix[N_rows, N_rows] L_K1;

    matrix[N_columns, N_columns] K2;
    matrix[N_columns, N_columns] L_K2;

    for(i in 1:N_rows) {
      for(j in 1:N_rows) {
        K1[i, j] = alpha_gp * (1 + fabs(i - j) * sqrt(3) / rho_gp1) * exp(-fabs(i - j) * sqrt(3) / rho_gp1);
      }
    }
    K1 = K1 + diag_matrix(rep_vector(delta0, N_rows));

    for(i in 1:N_columns) {
      for(j in 1:N_columns) {
        K2[i, j] = alpha_gp * (1 + fabs(i - j) * sqrt(3) / rho_gp2) * exp(-fabs(i - j) * sqrt(3) / rho_gp2);
      }
    }
    K2 = K2 + diag_matrix(rep_vector(delta0, N_columns));

    L_K1 = cholesky_decompose(K1);
    L_K2 = cholesky_decompose(K2);

    GP = kron_mvprod(L_K2, L_K1, z1);

    return(GP);
  }

  real beta_neg_binomial_lpmf(int y, real size_param, real shape_param, real location_param) {
    // size is the size parameter of the negative binomial distribution
    // prob is the probability parameter of the the negative binomial distribution
    // shape the first shape parameter of the beta distribution
    // shape2 the second shape parameter of the beta distribution
    // location is the mean of the beta-negative-binomial parameter
    real value;
    real shape2 = (shape_param - 1) * location_param / size_param;

    // sanity checks
    if (size_param > 1E8) {
      value = negative_infinity();
    } else if (shape_param > 1E8) {
      value = negative_infinity();
    } else if (shape2 > 1E8) {
      value = negative_infinity();
    } else if (size_param <= 1E-8) {
      value = negative_infinity();
    } else if (shape_param <= 1E-8) {
      value = negative_infinity();
    } else if (shape2 <= 1E-8) {
      value = negative_infinity();
    } else {
      // numerator
      value = lgamma(y + size_param) - lgamma(y + 1) - lgamma(size_param);
      // and denominator
      value += lbeta(shape_param + size_param, shape2 + y) - lbeta(shape_param, shape2);
    }
    return value;
  }
}

data {
  // distance sampling
  int<lower = 1> n_det;
  int<lower = 1> n_level;
  real BEAUFORT[n_level];
  real<lower = 0.0> w; // truncation distance
  int<lower = 1> BEAUFORT_det[n_det];
  vector<lower = 0.0, upper = w>[n_det] DISTANCE;
  real<lower = 0.0> prior_scale_sigma;
  // density surface model
  int<lower = 1> n_seg;
  int<lower = 1> n_year;
  int<lower = 1> n_session;
  int<lower = 0> COUNT[n_seg];                    // response variable
  vector<lower = 0.0>[n_seg] EFFORT;              // offset
  int<lower = 1> BEAUFORT_seg[n_seg];
  int<lower = 1, upper = n_year> YEAR[n_seg];     // year
  int<lower = 1, upper = n_session> SESSION[n_seg];// session
  real prior_location_intercept;
  real<lower = 0.0> prior_scale_intercept;
  // spatial: BSplines to approximate GP
  int<lower = 1> n_cell;
  int<lower = 1> n_row;
  int<lower = 1> n_col;
  int<lower = 0, upper = n_cell> IDROW[n_seg];
  int<lower = 0, upper = n_cell> IDCOL[n_seg];
  int num_basis_rows;
  int num_basis_columns;
  matrix[num_basis_rows, n_row] BASIS_ROWS;
  matrix[num_basis_columns, n_col] BASIS_COLUMNS;
  real IDX_BASIS_ROWS[num_basis_rows];
  real IDX_BASIS_COLUMNS[num_basis_columns];
}

transformed data {
  real eps = 1e-9;
  real prior_scale_delta = log2() / 2;
}

parameters {
  vector[n_session] unscaled_intercept[n_year];
  // overdispersion
  real<lower = 0.0, upper = 1.0> inv_alpha;
  // temporal
  real<lower = -1.0, upper = 1.0> temporal_range[n_year];
  // spatial
  real<lower = 0.0> rho_1;
  real<lower = 0.0> rho_2;
  real<lower = 0.0> alpha_gp;
  matrix[num_basis_rows, num_basis_columns] eta[n_year, n_session];
  // distance sampling
  vector<lower = 0.0>[2] unscaled_sigma2;
  vector<lower = 0.0>[2] tau;
  real unscaled_delta;
}

transformed parameters {
  matrix[num_basis_rows,num_basis_columns] beta[n_year, n_session];
  matrix[n_row, n_col] latent[n_year, n_session];
  real r = 1 / (unscaled_sigma2[2] / tau[2]); // size parameter, https://github.com/stan-dev/rstanarm/issues/275
  real alpha = 1 / inv_alpha; // shape parameter > 1
  vector[n_session] intercept[n_year];
  vector[n_seg] log_lik;
  // distance sampling
  real delta = unscaled_delta * prior_scale_delta;
  vector[n_level] sigma;
  vector[n_level] esw;
  for(k in 1:n_level) {
    sigma[k] = prior_scale_sigma * sqrt(unscaled_sigma2[1] / tau[1]) * exp(delta * BEAUFORT[k]);
    esw[k] = (normal_cdf(w, 0.0, sigma[k]) - 0.5) * exp(-normal_lpdf(0.0| 0.0, sigma[k]));
  }
  // spatial effects
  for(k in 1:n_year) {
    intercept[k] = prior_location_intercept + prior_scale_intercept * unscaled_intercept[k];
    for(l in 1:n_session) {
      beta[k, l] = gp(num_basis_rows, num_basis_columns,
                      IDX_BASIS_ROWS, IDX_BASIS_COLUMNS,
                      eps, alpha_gp, rho_1,  rho_2,
                      eta[k, l]);
    }
    // temporal correlation
    latent[k, 1] = sqrt(1 - square(temporal_range[k])) * ((BASIS_ROWS') * beta[k, 1] * BASIS_COLUMNS); // initial condition
    for(l in 2:n_session) {
      latent[k, l] = temporal_range[k] * latent[k, l - 1] + ((BASIS_ROWS') * beta[k, l] * BASIS_COLUMNS);
    }
  }
  // linear predictor
  for(i in 1:n_seg) {
    log_lik[i] = beta_neg_binomial_lpmf(COUNT[i]| r, alpha, (2 * EFFORT[i] * esw[BEAUFORT_seg[i]]) * exp(intercept[YEAR[i], SESSION[i]] + latent[YEAR[i], SESSION[i], IDROW[i], IDCOL[i]]));
  }
}

model {
  rho_1 ~ inv_gamma(5.0, 5.0);
  rho_2 ~ inv_gamma(5.0, 5.0);
  alpha_gp ~ std_normal();
  for(i in 1:num_basis_rows) {
    for(j in 1:num_basis_columns) {
      for(k in 1:n_year) {
        for(l in 1:n_session) {
          eta[k, l, i, j] ~ std_normal();
        }
      }
    }
  }
  for(k in 1:n_year) {
    unscaled_intercept[k] ~ std_normal();
  }
  for(i in 1:n_seg) {
    target += log_lik[i];
  }
  // distance sampling
  unscaled_delta ~ normal(0.0, 1.0);
  unscaled_sigma2 ~ gamma(0.5, 1.0);
  tau ~ gamma(1.0, 1.0);
  for(j in 1:n_det) {
    target += normal_lpdf(DISTANCE[j]| 0.0, sigma[BEAUFORT_det[j]]);
  }
}