dynamic_model.stan

functions {
    // For sum to 0. From here: https://discourse.mc-stan.org/t/test-soft-vs-hard-sum-to-zero-constrain-choosing-the-right-prior-for-soft-constrain/3884/31

    vector Q_sum_to_zero_QR(int N) {
        vector [2*N] Q_r;

        for(i in 1:N) {
          Q_r[i] = -sqrt((N-i)/(N-i+1.0));
          Q_r[i+N] = inv_sqrt((N-i) * (N-i+1));
        }
        return Q_r;
      }

    vector sum_to_zero_QR(vector x_raw, vector Q_r) {
        int N = num_elements(x_raw) + 1;
        vector [N] x;
        real x_aux = 0;

        for(i in 1:N-1){
          x[i] = x_aux + x_raw[i] * Q_r[i];
          x_aux = x_aux + x_raw[i] * Q_r[i+N];
        }
        x[N] = x_aux;
        return x;
      }
}

data {
    real<lower=0> nu;  // number of degree of freedom for evolution-prior
    real<lower=0> scale_sigma;
    int<lower=1> N;  // sample wave_size
    int<lower=1> n_waves;
    int<lower=1,upper=n_waves> wave[N];


    int<lower=1,upper=2> gender[N];
    int start_gender;



    int<lower=1,upper=2> question_1[N];

    int<lower=1,upper=2> question_2[N];


    int<lower = 0> p;  // the number of factor covariates
    int<lower = 1> n_covariate[p];  // the number of distinct levels of each covariate
}

transformed data {
    int<lower = 0> n_change = n_waves - 1;


    vector[2*2] Qr_gender = Q_sum_to_zero_QR(2);
    real x_sum0_sigma_gender = inv_sqrt(1 - inv(2));


    int<lower = 1> sum_n_covariate = sum(n_covariate);  // the sum of the number of distinct levels of each covariate after stacking them all

    // X is a one-hot encoding of all the covariates
    // rows should sum to p
    matrix<lower = 0, upper = 1>[N, sum_n_covariate * n_waves] X = rep_matrix(0, N, sum_n_covariate * n_waves);

    // check variables
    vector[N] sum_rows_X = X * rep_vector(1, sum_n_covariate * n_waves);

    // Define end indexes of levels in the stacked covariate (start indices are data)

    int end_gender = start_gender + 2 - 1;


    // Define X, the design matrix
    for (i in 1:N) {
        int w = wave[i];

        int l_gender = start_gender - 1 + gender[i];
        int m_gender = l_gender + (w - 1) * sum_n_covariate; // idx within this row of X
        X[i, m_gender] = 1;

    }

    // Run Checks
    sum_rows_X = X * rep_vector(1, sum_n_covariate * n_waves);
    if (min(sum_rows_X) != max(sum_rows_X) || min(sum_rows_X) != p) reject("X is configured incorrectly. The rows should sum to ", p);
}


parameters {

    positive_ordered[2-1] c_question_1;
    vector<lower=0>[n_waves-1] d_raw_question_1;
    real <lower=0> d_sigma_question_1;
    real <lower=0, upper=1> d_phi_question_1;

    positive_ordered[2-1] c_question_2;
    vector<lower=0>[n_waves-1] d_raw_question_2;
    real <lower=0> d_sigma_question_2;
    real <lower=0, upper=1> d_phi_question_2;



    real <lower=0> d_0_question_2;
    real <lower=0> d_mu_question_2;



    vector[2 - 1] beta_0_sum0_gender;
    vector[2] beta_raw_gender [n_waves-1];
    vector[2] mu_dyn_gender;
    real <lower=0> sigma_dyn_gender;
    real <lower=0, upper=1> phi_dyn_gender;

}

transformed parameters {


    vector[2] beta_0_gender = sum_to_zero_QR(beta_0_sum0_gender, Qr_gender);
    vector[2] beta_dyn_gender [n_waves];
    beta_dyn_gender[1] = mu_dyn_gender + beta_0_gender * sigma_dyn_gender * scale_sigma / sqrt(1 - phi_dyn_gender * phi_dyn_gender);



    vector[n_waves] d_question_1;

    vector[n_waves] d_question_2;


    real d_mu_question_1;
    d_mu_question_1 = 1;
    d_question_1[1] = 1;


    d_question_2[1] = d_mu_question_2 + d_0_question_2 * d_sigma_question_2 * scale_sigma / sqrt(1 - d_phi_question_2 * d_phi_question_2);


    for (k in 2:n_waves) {

        beta_dyn_gender[k] = mu_dyn_gender + phi_dyn_gender * (beta_dyn_gender[k-1] - mu_dyn_gender) + beta_raw_gender[k-1] * sigma_dyn_gender * scale_sigma;



        d_question_1[k] = d_mu_question_1 + d_phi_question_1 * (d_question_1[k-1] - d_mu_question_1) + d_raw_question_1[k-1] * d_sigma_question_1 * scale_sigma;

        d_question_2[k] = d_mu_question_2 + d_phi_question_2 * (d_question_2[k-1] - d_mu_question_2) + d_raw_question_2[k-1] * d_sigma_question_2 * scale_sigma;

    }
}

model {

    vector[n_waves * sum_n_covariate] b_question_1;

    vector[n_waves * sum_n_covariate] b_question_2;


    for (w in 1:n_waves) {
        vector[sum_n_covariate] b_wave;


        b_wave[start_gender:end_gender] = beta_dyn_gender[w];



        b_question_1[(w - 1) * sum_n_covariate + 1: w * sum_n_covariate] = b_wave*d_question_1[w];

        b_question_2[(w - 1) * sum_n_covariate + 1: w * sum_n_covariate] = b_wave*d_question_2[w];

    }


    for (k in 1:(2-1)) {
        c_question_1[k] ~ normal(k, 2-1);
    }

    for (k in 1:(2-1)) {
        c_question_2[k] ~ normal(k, 2-1);
    }


    for (k in 1:(n_waves-1)) {

        beta_raw_gender[k] ~ student_t(nu, 0, 1);


        d_raw_question_1[k] ~ student_t(nu, 0, 1);

        d_raw_question_2[k] ~ student_t(nu, 0, 1);

    }


    mu_dyn_gender ~ normal(0, 0.5);
    sigma_dyn_gender ~ normal(0, 1);
    phi_dyn_gender ~ beta(10, 1);
    beta_0_sum0_gender ~ normal(0, x_sum0_sigma_gender);



    d_0_question_2 ~ student_t(nu, 0, 1);
    d_mu_question_2 ~ normal(1, 1);



    d_sigma_question_1 ~ normal(0, 1);
    d_phi_question_1 ~ beta(10, 1);
    target += ordered_logistic_glm_lpmf(question_1 | X, b_question_1, c_question_1);

    d_sigma_question_2 ~ normal(0, 1);
    d_phi_question_2 ~ beta(10, 1);
    target += ordered_logistic_glm_lpmf(question_2 | X, b_question_2, c_question_2);

}