blocked.hpp

Go to the documentation of this file.

#ifndef MUMOSA_BLOCKED_HPP
#define MUMOSA_BLOCKED_HPP
 
#include <vector>
#include <algorithm>
#include <cstddef>
#include <optional>
 
#include "knncolle/knncolle.hpp"
#include "sanisizer/sanisizer.hpp"
#include "scran_blocks/scran_blocks.hpp"
#include "quickstats/quickstats.hpp"
 
#include "simple.hpp"
#include "utils.hpp"
 
namespace mumosa {
 
struct BlockedOptions {
    int num_neighbors = 20;
 
    scran_blocks::WeightPolicy block_weight_policy = scran_blocks::WeightPolicy::VARIABLE;
 
    scran_blocks::VariableWeightParameters variable_block_weight_parameters; 
 
    int num_threads = 1;
};
 
/*
 * We don't apply block-specific scaling factors as we don't want to alter the relative values within the same modality.
 * We shouldn't have to do it in the first place - as it's the same modality! - but more importantly, we could introduce spurious differences between blocks.
 * In the simplest case, two blocks have the same subpopulation structure but the number of cells is different.
 * We would get different distances in each block due to density, causing us to scale each block differently.
 * More generally, we could expect differences in subpopulation structure between blocks, leading to different distances even in the absence of any batch effects.
 * (Mind you, differences in subpopulation structure also interfere with accurate scaling between modalities,
 * but any errors in scaling modalities are much less obvious than those from scaling blocks.)
 */
 
template<typename Index_, typename Distance_>
std::pair<Distance_, Distance_> compute_distance_blocked(const std::vector<std::pair<Index_, Distance_*> >& blocks, const BlockedOptions& options) {
    const auto nblocks = blocks.size();
    auto block_weights = sanisizer::create<std::vector<Distance_> >(nblocks);
    scran_blocks::compute_weights(
        sanisizer::cast<std::size_t>(nblocks),
        [&](std::size_t b) -> Index_ { return blocks[b].first; },
        options.block_weight_policy,
        options.variable_block_weight_parameters,
        [&](std::size_t b, Distance_ w) -> void { block_weights[b] = w; }
    );
 
    const auto total_weight = [&]{
        quickstats::PairwiseSumWorkspace<Distance_> pswrk;
        quickstats::PairwiseSumOptions psopt;
        return quickstats::pairwise_sum(block_weights.size(), block_weights.data(), pswrk, psopt);
    }();
 
    auto outputs = sanisizer::create<std::vector<std::pair<Distance_, Distance_> > >(nblocks);
    knncolle::parallelize(options.num_threads, nblocks, [&](const int, I<decltype(nblocks)> start, I<decltype(nblocks)> length) -> void {
        for (I<decltype(nblocks)> b = start, bend = start + length; b < bend; ++b) {
            const auto curweight = block_weights[b];
            const auto curdist = compute_distance(blocks[b].first, blocks[b].second);
            outputs[b].first = curdist.first * curweight;
            outputs[b].second = curdist.second * curweight;
        }
    });
 
    std::pair<Distance_, Distance_> output{};
    for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
        output.first += outputs[b].first;
        output.second += outputs[b].second;
    }
 
    if (total_weight) {
        output.first /= total_weight;
        output.second /= total_weight;
    }
 
    return output;
}
 
template<typename Index_, typename Input_, typename Distance_>
std::pair<Distance_, Distance_> compute_distance_blocked(
    const std::vector<std::shared_ptr<const knncolle::Prebuilt<Index_, Input_, Distance_> > >& prebuilts,
    Distance_* const buffer,
    const BlockedOptions& options
) {
    const auto nblocks = prebuilts.size();
    std::size_t accumulated = 0;
    std::vector<std::pair<Index_, Distance_*> > blocks;
    blocks.reserve(nblocks);
 
    for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
        const auto nobs = prebuilts[b]->num_observations();
        const auto capped_k = knncolle::cap_k(options.num_neighbors, nobs);
 
        knncolle::parallelize(options.num_threads, nobs, [&](const int, const Index_ start, const Index_ length) -> void {
            const auto searcher = prebuilts[b]->initialize();
            std::vector<Distance_> cur_distances;
            for (Index_ i = start, end = start + length; i < end; ++i) {
                searcher->search(i, capped_k, NULL, &cur_distances);
                if (cur_distances.size()) {
                    buffer[accumulated + i] = cur_distances.back();
                } else {
                    buffer[accumulated + i] = 0; // i.e., only distance is that to itself.
                }
            }
        });
 
        blocks.emplace_back(nobs, buffer + accumulated);
        accumulated += nobs;
    }
 
    return compute_distance_blocked(blocks, options);
}
 
template<typename Index_, typename Input_, typename Block_, typename Distance_, class Matrix_ = knncolle::Matrix<Index_, Input_> >
std::pair<Distance_, Distance_> compute_distance_blocked(
    const std::size_t num_dim,
    const Index_ num_cells,
    const Input_* const data,
    const Block_* const blocks,
    const std::size_t num_blocks,
    const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& builder,
    Distance_* const buffer,
    const BlockedOptions& options
) {
    // Avoiding allocation of a temporary buffer if we're already dealing with contiguous blocks.
    auto block_details = sanisizer::create<std::vector<std::pair<Index_, Index_> > >(num_blocks);
    Index_ non_contiguous = 0;
    for (Index_ c = 0; c < num_cells; ++c) {
        auto& curblock = block_details[blocks[c]];
        if (curblock.second == 0) {
            curblock.first = c;
            curblock.second = 1;
        } else {
            non_contiguous += (c != curblock.first + curblock.second);
            ++curblock.second;
        }
    }
 
    const Input_* dataptr = data;
    std::optional<std::vector<Input_> > tmp_data;
    if (non_contiguous) {
        // Otherwise, we reorganize the data so that observations from the same batch are in a single block.
        Index_ accumulated = 0;
        auto offsets = sanisizer::create<std::vector<Index_> >(num_blocks);
        for (std::size_t b = 0; b < num_blocks; ++b) {
            offsets[b] = accumulated;
            block_details[b].first = accumulated;
            accumulated += block_details[b].second; // this won't overflow as we already know that num_cells fits in an Index_.
        }
 
        tmp_data.emplace(sanisizer::product<typename std::vector<Input_>::size_type>(num_dim, num_cells));
        for (Index_ c = 0; c < num_cells; ++c) {
            auto& off = offsets[blocks[c]];
            std::copy_n(
                data + sanisizer::product_unsafe<std::size_t>(c, num_dim),
                num_dim,
                tmp_data->data() + sanisizer::product_unsafe<std::size_t>(off, num_dim)
            );
            ++off;
        }
 
        dataptr = tmp_data->data();
    }
 
    auto prebuilts = sanisizer::create<std::vector<std::shared_ptr<const knncolle::Prebuilt<Index_, Input_, Distance_> > > >(num_blocks);
    knncolle::parallelize(options.num_threads, num_blocks, [&](const int, const std::size_t start, const std::size_t length) -> void {
        for (std::size_t b = start, end = start + length; b < end; ++b) {
            const auto sofar = block_details[b].first;
            const auto cursize = block_details[b].second;
            prebuilts[b] = builder.build_shared(knncolle::SimpleMatrix(num_dim, cursize, dataptr + sanisizer::product_unsafe<std::size_t>(sofar, num_dim)));
        }
    });
 
    return compute_distance_blocked(prebuilts, buffer, options);
}
 
}
 
#endif