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

functions {
    real confusion_model_lpmf(array[] int group,
        int start, int end,
        int n_classes,
        array[,] int algo,
        array[,] int truth,
        array[] real clip_duration,
        array[] matrix lambda,
        matrix omega//,
        //array[] vector lambda_fp,
    ) {
        real ll = 0;
        vector [4] bp;

        vector[8192] log_contrib_comb;
        int n = size(log_contrib_comb);

        for (k in start:end) {
            for (i in 1:n_classes) {
                log_contrib_comb[:n] = rep_vector(0, n);
                n = 1;

                for (chi in 0:(truth[k,1]>0?max(truth[k,1], algo[k,i]):0)) {
                    bp[1] = truth[k,1]==0?0:neg_binomial_lpmf(chi | truth[k,1]*lambda[group[k-start+1],1,i]/(omega[1,i]-1), 1/(omega[1,i]-1));

                    for (och in 0:(truth[k,2]>0?max(truth[k,2], algo[k,i]-chi):0)) {
                        bp[2] = truth[k,2]==0?0:neg_binomial_lpmf(och | truth[k,2]*lambda[group[k-start+1],2,i]/(omega[2,i]-1), 1/(omega[2,i]-1));

                        for (fem in 0:(truth[k,3]>0?max(truth[k,3], algo[k,i]-chi-och):0)) {
                            bp[3] = truth[k,3]==0?0:neg_binomial_lpmf(fem | truth[k,3]*lambda[group[k-start+1],3,i]/(omega[3,i]-1), 1/(omega[3,i]-1));

                            for (mal in 0:(truth[k,4]>0?max(truth[k,4], algo[k,i]-chi-och-fem):0)) {
                                bp[4] = truth[k,4]==0?0:neg_binomial_lpmf(mal | truth[k,4]*lambda[group[k-start+1],4,i]/(omega[4,i]-1), 1/(omega[4,i]-1));

                                int delta = algo[k,i] - (mal+fem+och+chi);
                                // if (delta >= 0) {
                                //     log_contrib_comb[n] += sum(bp);
                                //     log_contrib_comb[n] += poisson_lpmf(
                                //         delta | lambda_fp[group[k-start+1],i]*clip_duration[k]
                                //     );
                                //     n = n+1;
                                // }
                                if (delta==0) {
                                    log_contrib_comb[n] += sum(bp);
                                    n = n+1;
                                }
                            }
                        }
                    }
                }
                if (n>1) {
                    ll += log_sum_exp(log_contrib_comb[1:n-1]);
                }
            }
        }
        return ll;
    }
}

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

    // analysis data block
    int<lower=1> n_recs;
    int<lower=1> n_children;

    array[n_recs] int<lower=1> children;
    array[n_recs] real<lower=1> age;
    array[n_recs] int<lower=-1> siblings;
    array[n_recs, n_classes] int<lower=0> vocs;
    array[n_children] int<lower=1> corpus;

    real<lower=0> recs_duration;

    // 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 {
    vector<lower=0>[n_groups] recording_age;
    array[n_children] int<lower=-1> child_siblings;
    int no_siblings = 0;
    int has_siblings = 0;

    array [n_groups,n_classes] int group_truth;

    for (c in 1:n_clips) {
        recording_age[group[c]] = clip_age[c];
    }

    for (k in 1:n_recs) {
        child_siblings[children[k]] = siblings[k];
    }

    for (c in 1:n_children) {
        if (child_siblings[c] == 0) {
            no_siblings += 1;
        }
        else if (child_siblings[c] > 0) {
            has_siblings += 1;
        }
    }

    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 {
    // confusion parameters
    // confusion matrix
    matrix<lower=0>[n_classes,n_classes] alphas;
    matrix<lower=0>[n_classes,n_classes] mus;
    matrix<lower=1>[n_classes,n_classes] omega;
    array [n_groups] matrix<lower=0>[n_classes,n_classes] lambda;

    // false positives
    //vector<lower=0>[n_classes] alphas_fp;
    //vector<lower=0>[n_classes] mus_fp;
    //array [n_groups] vector<lower=0>[n_classes] lambda_fp;
}

model {
    target += reduce_sum(
        confusion_model_lpmf, group, n_clips%/%(threads*4),
        n_classes,
        algo_total, truth_total, clip_duration,
        lambda, omega//, lambda_fp
    );

    //mus_fp ~ exponential(1);
    //alphas_fp ~ gamma(2, 1);

    for (i in 1:n_classes) {
        //lambda_fp[:,i] ~ gamma(alphas_fp[i], alphas_fp[i]/mus_fp[i]);
        omega[i,:] ~ pareto(1, 2);

        for (j in 1:n_classes) {
            mus[i,j] ~ exponential(2);
            alphas[i,j] ~ inv_gamma(1, 1);
            for (c in 1:n_groups) {
                lambda[c,i,j] ~ gamma(alphas[i,j], alphas[i,j]/mus[i,j]);
            }
        }
    }
}

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) {
        for (i in 1:n_classes) {
            vector[n_classes] mu_given_lambda = lambda[g,:,i].*to_vector(group_truth[g,:]);
            vector[n_classes] mu = to_vector(gamma_rng(alphas[:,i],alphas[:,i]./mus[:,i])).*to_vector(group_truth[g,:]);

            sim_vocs_given_lambda[g,i] = 0;
            sim_vocs[g,i] = 0;
            for (j in 1:n_classes) {
                if (group_truth[g,j] > 0) {
                    sim_vocs_given_lambda[g,i] += neg_binomial_rng(mu_given_lambda[j]/(omega[j,i]-1), 1/(omega[j,i]-1));
                    sim_vocs[g,i] += neg_binomial_rng(mu[j]/(omega[j,i]-1), 1/(omega[j,i]-1));
                }
            }
        }
    }
}