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

// 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, n_classes] int<lower=0> vocs;
    array[n_children] int<lower=1> corpus;

    real<lower=0> 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,n_classes] int vtc;
    array [n_clips,n_classes] int truth;

    int<lower=1> n_validation;

    // actual speech rates
    int<lower=1> n_rates;
    array [n_rates,n_classes] int<lower=0> speech_rates;
    array [n_rates] int group_corpus;
    array [n_rates] real<lower=0> durations;
}

parameters {
    matrix<lower=0> [n_recs, n_classes] truth_vocs;
    array [n_children] matrix<lower=0,upper=1>[n_classes,n_classes] actual_confusion_baseline;

    // confusion parameters
    matrix<lower=0,upper=1>[n_classes,n_classes] mus;
    matrix<lower=1>[n_classes,n_classes] etas;
    array [n_groups] matrix<lower=0,upper=1>[n_classes,n_classes] group_confusion;

    array [n_corpora] matrix[n_classes,n_classes] corpus_bias;
    matrix<lower=0>[n_classes,n_classes] corpus_sigma;

    // speech rates
    matrix<lower=1>[n_classes,n_corpora] speech_rate_alpha;
    matrix<lower=0>[n_classes,n_corpora] speech_rate_mu;
    matrix<lower=0> [n_classes,n_rates] speech_rate;
}

transformed parameters {
    array [n_children] matrix<lower=0,upper=1>[n_classes,n_classes] actual_confusion;

    for (c in 1:n_children) {
        actual_confusion[c] = inv_logit(logit(actual_confusion_baseline[c])+corpus_bias[corpus[c]]);
    }
}

model {
    // actual model
    vector [4] expect;
    vector [4] sd;

    for (k in 1:n_recs) {
        expect = rep_vector(0, 4);
        sd = rep_vector(0, 4);

        for (i in 1:n_classes) {
            for (j in 1:n_classes) {
                expect[i] += truth_vocs[k,j]*actual_confusion[children[k],j,i];
                sd[i] += truth_vocs[k,j]*actual_confusion[children[k],j,i]*(1-actual_confusion[children[k],j,i]);
            }
        }
        vocs[k,:] ~ normal(expect, sqrt(sd));
    }

    for (c in 1:n_children) {
        for (i in 1:n_classes) {
            actual_confusion_baseline[c,i] ~ beta_proportion(mus[i,:], etas[i,:]);
        }
    }

    for (k in 1:n_recs) {
        truth_vocs[k,:] ~ gamma(
            speech_rate_alpha[:,corpus[children[k]]],
            (speech_rate_alpha[:,corpus[children[k]]]./speech_rate_mu[:,corpus[children[k]]])/1000/duration
        );
    }

    // speaker confusion model
    for (k in n_validation:n_clips) {
        for (i in 1:n_classes) {
            for (j in 1:n_classes) {
                vtc[k,i,j] ~ binomial(
                    truth[k,j], inv_logit(logit(group_confusion[group[k],j,i]) + corpus_bias[conf_corpus[k],j,i])
                );
            }
        }
    }

    for (i in 1:n_classes) {
        for (j in 1:n_classes) {
            if (i==j) {
                mus[i,j] ~ beta(2,1);
            }
            else {
                mus[i,j] ~ beta(1,2);
            }
            etas[i,j] ~ pareto(1,1.5);
        }
    }

    for (c in 1:n_groups) {
        for (i in 1:n_classes) {
            for (j in 1:n_classes) {
                group_confusion[c,i,j] ~ beta_proportion(mus[i,j], etas[i,j]);
            }
        }
    }

    for (i in 1:n_classes) {
        for (j in 1:n_classes) {
            for (c in 1:n_corpora) {
                corpus_bias[c,j,i] ~ normal(0, corpus_sigma[j,i]);
            }
            corpus_sigma[j,i] ~ normal(0, 1);
        }
    }

    // speech rates
    for (i in 1:n_classes) {
        speech_rate_alpha[i,:] ~ normal(1, 1);
        speech_rate_mu[i,:] ~ exponential(2);
    }

    for (g in 1:n_rates) {
        for (i in 1:n_classes) {
            speech_rate[i,g] ~ gamma(
                speech_rate_alpha[i,group_corpus[g]],
                (speech_rate_alpha[i,group_corpus[g]]/speech_rate_mu[i,group_corpus[g]])/1000
            );
            speech_rates[g,i] ~ poisson(speech_rate[i,g]*durations[g]);
        }
    }
}