speaker-confusion-model/code/models/cross_validation_multinomial.stan

functions {
    #include "blocks/confusion_model_multinomial.stan"
}

// TODO
// use speech rates to set priors on truth_vocs
data {
    int<lower=1> n_classes; // number of classes

    // speaker confusion data block
    int<lower=1> n_clips;   // number of clips
    int<lower=1> n_groups; // number of groups
    int<lower=1> n_corpora;
    array [n_clips] int group;
    array [n_clips] int conf_corpus;
    array [n_clips,n_classes] int<lower=0> algo_total; // algo vocs attributed to specific speakers
    array [n_clips,n_classes] int<lower=0> truth_total;
    array [n_clips] real<lower=0> clip_duration;
    array [n_clips] real<lower=0> clip_age;

    int<lower=0> n_validation;
    
    // parallel processing
    int<lower=1> threads;
}

transformed data {
    array [n_groups,n_classes] int group_truth;

    for (i in 1:n_classes) {
        for (g in 1:n_groups) {
            group_truth[g,i] = 0;
        }
        for (c in 1:n_clips) {
            group_truth[group[c],i] += truth_total[c,i];
        }
    } 
    
}

parameters {
    #include "blocks/confusion_model_parameters_multinomial.stan"
}

model {
    target += reduce_sum(
        confusion_model_lpdf, lambda, 1,
        n_classes, n_clips,
        algo_total, truth_total, group, clip_duration, clip_age
    );

    #include "blocks/confusion_model_priors_multinomial.stan"
}

generated quantities {
    array [n_groups,n_classes] int sim_vocs_given_lambda;
    array [n_groups,n_classes] int sim_vocs;

    for (g in 1:n_groups) {
        array [n_classes,n_classes+1] int n_given_lambda;
        array [n_classes,n_classes+1] int n;
        for (i in 1:n_classes) {
            if (group_truth[g,i]>0) {
                n_given_lambda[i] = multinomial_rng(lambda[g,i,:], group_truth[g,i]);
                n[i] = multinomial_rng(dirichlet_rng(mus[i]*etas[i]), group_truth[g,i]);
            }
            else {
                for (j in 1:n_classes+1) {
                    n_given_lambda[i,j] = 0;
                    n[i,j] = 0;
                }
            }
        }

        for (i in 1:n_classes) {
            sim_vocs_given_lambda[g,i] = sum(n_given_lambda[:,i]);
            sim_vocs[g,i] = sum(n[:,i]);
        }
    }
}