variances.hpp

Go to the documentation of this file.

#ifndef TATAMI_STATS_VARS_HPP
#define TATAMI_STATS_VARS_HPP
 
#include "tatami/tatami.hpp"
#include "utils.hpp"
 
#include <vector>
#include <cmath>
#include <numeric>
#include <limits>
#include <algorithm>
 
namespace tatami_stats {
 
namespace variances {
 
struct Options {
    bool skip_nan = false;
 
    int num_threads = 1;
};
 
namespace internal {
 
template<typename Output_ = double, typename Value_, typename Index_ >
void add_welford(Output_& mean, Output_& sumsq, Value_ value, Index_ count) {
    Output_ delta = value - mean;
    mean += delta / count;
    sumsq += delta * (value - mean);
}
 
template<typename Output_ = double, typename Index_ >
void add_welford_zeros(Output_& mean, Output_& sumsq, Index_ num_nonzero, Index_ num_all) {
    auto ratio = static_cast<Output_>(num_nonzero) / static_cast<Output_>(num_all);
    sumsq += mean * mean * ratio * (num_all - num_nonzero);
    mean *= ratio;
}
 
// Avoid problems from interactions between constexpr/lambda/std::conditional. 
template<typename Index_>
struct MockVector {
    MockVector(size_t) {}
    Index_& operator[](size_t) { return out; }
    Index_ out = 0;
};
 
}
template<typename Output_ = double, typename Value_, typename Index_ >
std::pair<Output_, Output_> direct(const Value_* value, Index_ num_nonzero, Index_ num_all, bool skip_nan) {
    Output_ mean = 0;
    Index_ lost = 0;
 
    ::tatami_stats::internal::nanable_ifelse<Value_>(
        skip_nan,
        [&]() {
            auto copy = value;
            for (Index_ i = 0; i < num_nonzero; ++i, ++copy) {
                auto val = *copy;
                if (std::isnan(val)) {
                    ++lost;
                } else {
                    mean += val;
                }
            }
        },
        [&]() {
            auto copy = value;
            for (Index_ i = 0; i < num_nonzero; ++i, ++copy) {
                mean += *copy;
            }
        }
    );
 
    auto count = num_all - lost;
    mean /= count;
 
    Output_ var = 0;
    ::tatami_stats::internal::nanable_ifelse<Value_>(
        skip_nan,
        [&]() {
            for (Index_ i = 0; i < num_nonzero; ++i) {
                auto val = value[i];
                if (!std::isnan(val)) {
                    auto delta = static_cast<Output_>(val) - mean;
                    var += delta * delta;
                }
            }
        },
        [&]() {
            for (Index_ i = 0; i < num_nonzero; ++i) {
                auto delta = static_cast<Output_>(value[i]) - mean;
                var += delta * delta;
            }
        }
    );
 
    if (num_nonzero < num_all) {
        var += static_cast<Output_>(num_all - num_nonzero) * mean * mean;
    }
 
    if (count == 0) {
        return std::make_pair(std::numeric_limits<Output_>::quiet_NaN(), std::numeric_limits<Output_>::quiet_NaN());
    } else if (count == 1) {
        return std::make_pair(mean, std::numeric_limits<Output_>::quiet_NaN());
    } else {
        return std::make_pair(mean, var / (count - 1));
    }
}
 
template<typename Output_ = double, typename Value_, typename Index_ >
std::pair<Output_, Output_> direct(const Value_* ptr, Index_ num, bool skip_nan) {
    return direct<Output_>(ptr, num, num, skip_nan);
}
 
template<typename Output_, typename Value_, typename Index_>
class RunningDense {
public:
    RunningDense(Index_ num, Output_* mean, Output_* variance, bool skip_nan) : 
        my_num(num), my_mean(mean), my_variance(variance), my_skip_nan(skip_nan), my_ok_count(skip_nan ? num : 0) {}
 
    void add(const Value_* ptr) {
        ::tatami_stats::internal::nanable_ifelse<Value_>(
            my_skip_nan,
            [&]() {
                for (Index_ i = 0; i < my_num; ++i, ++ptr) {
                    auto val = *ptr;
                    if (!std::isnan(val)) {
                        internal::add_welford(my_mean[i], my_variance[i], val, ++(my_ok_count[i]));
                    }
                }
            },
            [&]() {
                ++my_count;
                for (Index_ i = 0; i < my_num; ++i, ++ptr) {
                    internal::add_welford(my_mean[i], my_variance[i], *ptr, my_count);
                }
            }
        );
    }
 
    void finish() {
        ::tatami_stats::internal::nanable_ifelse<Value_>(
            my_skip_nan,
            [&]() {
                for (Index_ i = 0; i < my_num; ++i) {
                    auto ct = my_ok_count[i];
                    if (ct < 2) {
                        my_variance[i] = std::numeric_limits<Output_>::quiet_NaN();
                        if (ct == 0) {
                            my_mean[i] = std::numeric_limits<Output_>::quiet_NaN();
                        }
                    } else {
                        my_variance[i] /= ct - 1;
                    }
                }
            },
            [&]() {
                if (my_count < 2) {
                    std::fill_n(my_variance, my_num, std::numeric_limits<Output_>::quiet_NaN());
                    if (my_count == 0) {
                        std::fill_n(my_mean, my_num, std::numeric_limits<Output_>::quiet_NaN());
                    }
                } else {
                    for (Index_ i = 0; i < my_num; ++i) {
                        my_variance[i] /= my_count - 1;
                    }
                }
            }
        );
    }
 
private:
    Index_ my_num;
    Output_* my_mean;
    Output_* my_variance;
    bool my_skip_nan;
    Index_ my_count = 0;
    typename std::conditional<std::numeric_limits<Value_>::has_quiet_NaN, std::vector<Index_>, internal::MockVector<Index_> >::type my_ok_count;
};
 
template<typename Output_, typename Value_, typename Index_>
class RunningSparse {
public:
    RunningSparse(Index_ num, Output_* mean, Output_* variance, bool skip_nan, Index_ subtract = 0) : 
        my_num(num), my_mean(mean), my_variance(variance), my_nonzero(num), my_skip_nan(skip_nan), my_subtract(subtract), my_nan(skip_nan ? num : 0) {}
 
    void add(const Value_* value, const Index_* index, Index_ number) {
        ++my_count;
 
        ::tatami_stats::internal::nanable_ifelse<Value_>(
            my_skip_nan,
            [&]() {
                for (Index_ i = 0; i < number; ++i) {
                    auto val = value[i];
                    auto ri = index[i] - my_subtract;
                    if (std::isnan(val)) {
                        ++my_nan[ri];
                    } else {
                        internal::add_welford(my_mean[ri], my_variance[ri], val, ++(my_nonzero[ri]));
                    }
                }
            },
            [&]() {
                for (Index_ i = 0; i < number; ++i) {
                    auto ri = index[i] - my_subtract;
                    internal::add_welford(my_mean[ri], my_variance[ri], value[i], ++(my_nonzero[ri]));
                }
            }
        );
    }
 
    void finish() {
        ::tatami_stats::internal::nanable_ifelse<Value_>(
            my_skip_nan,
            [&]() {
                for (Index_ i = 0; i < my_num; ++i) {
                    auto& curM = my_mean[i];
                    auto& curV = my_variance[i];
                    Index_ ct = my_count - my_nan[i];
 
                    if (ct < 2) {
                        curV = std::numeric_limits<Output_>::quiet_NaN();
                        if (ct == 0) {
                            curM = std::numeric_limits<Output_>::quiet_NaN();
                        }
                    } else {
                        internal::add_welford_zeros(curM, curV, my_nonzero[i], ct);
                        curV /= ct - 1;
                    }
                }
            },
            [&]() {
                if (my_count < 2) {
                    std::fill_n(my_variance, my_num, std::numeric_limits<Output_>::quiet_NaN());
                    if (my_count == 0) {
                        std::fill_n(my_mean, my_num, std::numeric_limits<Output_>::quiet_NaN());
                    }
                } else {
                    for (Index_ i = 0; i < my_num; ++i) {
                        auto& var = my_variance[i];
                        internal::add_welford_zeros(my_mean[i], var, my_nonzero[i], my_count);
                        var /= my_count - 1;
                    }
                }
            }
        );
    }
 
private:
    Index_ my_num;
    Output_* my_mean;
    Output_* my_variance;
    std::vector<Index_> my_nonzero;
    bool my_skip_nan;
    Index_ my_subtract;
    Index_ my_count = 0;
    typename std::conditional<std::numeric_limits<Value_>::has_quiet_NaN, std::vector<Index_>, internal::MockVector<Index_> >::type my_nan;
};
 
template<typename Value_, typename Index_, typename Output_>
void apply(bool row, const tatami::Matrix<Value_, Index_>* p, Output_* output, const Options& vopt) {
    auto dim = (row ? p->nrow() : p->ncol());
    auto otherdim = (row ? p->ncol() : p->nrow());
    const bool direct = p->prefer_rows() == row;
 
    if (p->sparse()) {
        if (direct) {
            tatami::Options opt;
            opt.sparse_extract_index = false;
            tatami::parallelize([&](int, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<true>(p, row, s, l);
                std::vector<Value_> vbuffer(otherdim);
                for (Index_ x = 0; x < l; ++x) {
                    auto out = ext->fetch(vbuffer.data(), NULL);
                    output[x + s] = variances::direct<Output_>(out.value, out.number, otherdim, vopt.skip_nan).second;
                }
            }, dim, vopt.num_threads);
 
        } else {
            tatami::parallelize([&](int thread, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<true>(p, !row, static_cast<Index_>(0), otherdim, s, l);
                std::vector<Value_> vbuffer(l);
                std::vector<Index_> ibuffer(l);
 
                std::vector<Output_> running_means(l);
                LocalOutputBuffer<Output_> local_output(thread, s, l, output);
                variances::RunningSparse<Output_, Value_, Index_> runner(l, running_means.data(), local_output.data(), vopt.skip_nan, s);
 
                for (Index_ x = 0; x < otherdim; ++x) {
                    auto out = ext->fetch(vbuffer.data(), ibuffer.data());
                    runner.add(out.value, out.index, out.number);
                }
                runner.finish();
 
                local_output.transfer(); 
            }, dim, vopt.num_threads);
        }
 
    } else {
        if (direct) {
            tatami::parallelize([&](int, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<false>(p, row, s, l);
                std::vector<Value_> buffer(otherdim);
                for (Index_ x = 0; x < l; ++x) {
                    auto out = ext->fetch(buffer.data());
                    output[x + s] = variances::direct<Output_>(out, otherdim, vopt.skip_nan).second;
                }
            }, dim, vopt.num_threads);
 
        } else {
            tatami::parallelize([&](int thread, Index_ s, Index_ l) {
                auto ext = tatami::consecutive_extractor<false>(p, !row, static_cast<Index_>(0), otherdim, s, l);
                std::vector<Value_> buffer(l);
 
                std::vector<Output_> running_means(l);
                LocalOutputBuffer<Output_> local_output(thread, s, l, output);
                variances::RunningDense<Output_, Value_, Index_> runner(l, running_means.data(), local_output.data(), vopt.skip_nan);
 
                for (Index_ x = 0; x < otherdim; ++x) {
                    runner.add(ext->fetch(buffer.data()));
                }
                runner.finish();
 
                local_output.transfer(); 
            }, dim, vopt.num_threads);
        }
    }
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_column(const tatami::Matrix<Value_, Index_>* p, const Options& vopt) {
    std::vector<Output_> output(p->ncol());
    apply(false, p, output.data(), vopt);
    return output;
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_column(const tatami::Matrix<Value_, Index_>* p) {
    Options vopt;
    return by_column(p, vopt);
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_row(const tatami::Matrix<Value_, Index_>* p, const Options& vopt) {
    std::vector<Output_> output(p->nrow());
    apply(true, p, output.data(), vopt);
    return output;
}
 
template<typename Output_ = double, typename Value_, typename Index_>
std::vector<Output_> by_row(const tatami::Matrix<Value_, Index_>* p) {
    Options vopt;
    return by_row(p, vopt);
}
 
}
 
}
 
#endif