fit_variance_trend.hpp

Go to the documentation of this file.

#ifndef SCRAN_VARIANCES_FIT_VARIANCE_TREND_H
#define SCRAN_VARIANCES_FIT_VARIANCE_TREND_H
 
#include <algorithm>
#include <vector>
#include <array>
#include <cstddef>
 
#include "WeightedLowess/WeightedLowess.hpp"
#include "sanisizer/sanisizer.hpp"
 
namespace scran_variances {
 
struct FitVarianceTrendOptions {
    double minimum_mean = 0.1;
 
    bool mean_filter = true;
 
    bool transform = true;
 
    double span = 0.3;
 
    bool use_minimum_width = false;
 
    double minimum_width = 1;
 
    int minimum_window_count = 200;
 
    int num_threads = 1;
};
 
template<typename Float_>
struct FitVarianceTrendWorkspace {
    WeightedLowess::SortBy sorter;
 
    std::vector<uint8_t> sort_workspace;
 
    std::vector<Float_> xbuffer, ybuffer;
};
 
template<typename Float_>
void fit_variance_trend(std::size_t n, const Float_* mean, const Float_* variance, Float_* fitted, Float_* residuals, FitVarianceTrendWorkspace<Float_>& workspace, const FitVarianceTrendOptions& options) {
    auto& xbuffer = workspace.xbuffer;
    xbuffer.resize(sanisizer::cast<decltype(xbuffer.size())>(n));
    auto& ybuffer = workspace.ybuffer;
    ybuffer.resize(sanisizer::cast<decltype(ybuffer.size())>(n));
 
    auto quad = [](Float_ x) -> Float_ {
        return x * x * x * x;
    };
 
    std::size_t counter = 0;
    Float_ min_mean = options.minimum_mean;
    for (decltype(n) i = 0; i < n; ++i) {
        if (!options.mean_filter || mean[i] >= min_mean) {
            xbuffer[counter] = mean[i];
            if (options.transform) {
                ybuffer[counter] = std::pow(variance[i], 0.25); // Using the same quarter-root transform that limma::voom uses.
            } else {
                ybuffer[counter] = variance[i];
            }
            ++counter;
        }
    }
 
    if (counter < 2) {
        throw std::runtime_error("not enough observations above the minimum mean");
    }
 
    auto& sorter = workspace.sorter;
    sorter.set(counter, xbuffer.data());
    auto& work = workspace.sort_workspace;
    sorter.permute(std::array<Float_*, 2>{ xbuffer.data(), ybuffer.data() }, work);
 
    WeightedLowess::Options<Float_> smooth_opt;
    if (options.use_minimum_width) {
        smooth_opt.span = options.minimum_window_count;
        smooth_opt.span_as_proportion = false;
        smooth_opt.minimum_width = options.minimum_width;
    } else {
        smooth_opt.span = options.span;
    }
    smooth_opt.num_threads = options.num_threads;
 
    // Using the residual array to store the robustness weights as a placeholder;
    // we'll be overwriting this later.
    WeightedLowess::compute(counter, xbuffer.data(), ybuffer.data(), fitted, residuals, smooth_opt);
 
    // Determining the left edge before we unpermute.
    Float_ left_x = xbuffer[0];
    Float_ left_fitted = (options.transform ? quad(fitted[0]) : fitted[0]);
 
    sorter.unpermute(fitted, work);
 
    // Walking backwards to shift the elements back to their original position
    // (i.e., before filtering on the mean) on the same array. We need to walk
    // backwards to ensure that writing to the original position on this array
    // doesn't clobber the first 'counter' positions containing the fitted
    // values, at least not until each value is shifted to its original place.
    for (auto i = n; i > 0; --i) {
        auto j = i - 1;
        if (!options.mean_filter || mean[j] >= min_mean) {
            --counter;
            fitted[j] = (options.transform ? quad(fitted[counter]) : fitted[counter]);
        } else {
            fitted[j] = mean[j] / left_x * left_fitted; // draw a y = x line to the origin from the left of the fitted trend.
        }
    }
 
    for (decltype(n) i = 0; i < n; ++i) {
        residuals[i] = variance[i] - fitted[i];
    }
    return;
}
 
template<typename Float_>
struct FitVarianceTrendResults {
    FitVarianceTrendResults() {}
 
    FitVarianceTrendResults(std::size_t n) :
        fitted(sanisizer::cast<decltype(fitted.size())>(n)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        ),
        residuals(sanisizer::cast<decltype(residuals.size())>(n)
#ifdef SCRAN_VARIANCES_TEST_INIT
            , SCRAN_VARIANCES_TEST_INIT
#endif
        )
    {}
    std::vector<Float_> fitted;
 
    std::vector<Float_> residuals;
};
 
template<typename Float_>
FitVarianceTrendResults<Float_> fit_variance_trend(std::size_t n, const Float_* mean, const Float_* variance, const FitVarianceTrendOptions& options) {
    FitVarianceTrendResults<Float_> output(n);
    FitVarianceTrendWorkspace<Float_> work;
    fit_variance_trend(n, mean, variance, output.fitted.data(), output.residuals.data(), work, options);
    return output;
}
 
}
 
#endif