Compute a correlation matrix for pairs of variables in a ROOT TNtuple.

tntuple_correlation.h

/** Get a value from a TNtuple by event id and name */
float get_ntuple_entry(TNtuple* nt, int i, std::string field) {
    float v;
    nt->SetBranchAddress(field.c_str(), &v);
    nt->GetEvent(i);
    nt->ResetBranchAddresses();
    return v;
}


/**
 * Build a correlation matrix for a TNtuple.
 *
 * Creates a matrix with Pearson product-moment correlation coefficients
 * computed between pairs of variables in a TNtuple. The matrix expressed
 * as a vector of length (entries x entries). Only the upper half is set.
 *
 * \param nt The source TNtuple
 */
std::vector<float> get_correlation_matrix(TNtuple* nt) {
    int nentries = nt->GetEntries();

    // get list of branch names
    std::vector<std::string> names;
    for (size_t i=0; i<nt->GetListOfBranches()->GetEntries(); i++) {
        names.push_back(nt->GetListOfBranches()->At(i)->GetName());
    }

    std::vector<float> matrix(names.size() * names.size());

    // convert the ntuple to a vector, calculating means as we go
    std::vector<float> table(names.size() * nentries);
    std::vector<float> means(names.size(), 0);
    for (size_t i=0; i<nentries; i++) {
        for (size_t j=0; j<names.size(); j++) {
            float v = get_ntuple_entry(nt, i, names[j]);
            table[j + i * names.size()] = v;
            means[j] += v;
        }
    }

    // sums to means
    for (size_t i=0; i<names.size(); i++) {
        means[i] /= nentries;
    }

    // compute correlations
    for (size_t i=0; i<names.size(); i++) {
        for (size_t j=i; j<names.size(); j++) {
            float t = 0;
            float dx2 = 0;
            float dy2 = 0;
            for (int k=0; k<nentries; k++) {
                float x1 = table[k * names.size() + i] - means[i];
                float x2 = table[k * names.size() + j] - means[j];
                t += x1 * x2;
                dx2 += x1 * x1;
                dy2 += x2 * x2;
            }
            matrix[i * names.size() + j] = t / TMath::Sqrt(dx2 * dy2);
        }
    }

    return matrix;
}