blocked.hpp

Go to the documentation of this file.

#ifndef MUMOSA_BLOCKED_HPP
#define MUMOSA_BLOCKED_HPP
 
#include <vector>
#include <algorithm>
#include <cstddef>
 
#include "knncolle/knncolle.hpp"
#include "sanisizer/sanisizer.hpp"
#include "scran_blocks/scran_blocks.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;
};
 
template<typename Distance_>
struct BlockedWorkspace {
    std::vector<Distance_> weights;
    Distance_ total_weight;
 
    std::vector<Distance_> distance_buffer;
};
 
template<typename Distance_, typename Index_>
BlockedWorkspace<Distance_> create_workspace(const std::vector<Index_>& block_sizes, const BlockedOptions& options) {
    BlockedWorkspace<Distance_> output;
    output.weights = scran_blocks::compute_weights<Distance_>(block_sizes, options.block_weight_policy, options.variable_block_weight_parameters);
    output.total_weight = std::accumulate(output.weights.begin(), output.weights.end(), static_cast<Distance_>(0));
 
    Index_ max_size = 0;
    if (block_sizes.size()) {
        max_size = *std::max_element(block_sizes.begin(), block_sizes.end());
    }
    sanisizer::resize(output.distance_buffer, max_size);
 
    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,
    BlockedWorkspace<Distance_>& workspace,
    const BlockedOptions& options
) {
    Options simple_opt;
    simple_opt.num_neighbors = options.num_neighbors;
    simple_opt.num_threads = options.num_threads;
 
    std::pair<Distance_, Distance_> output(0, 0);
 
    const auto nblocks = prebuilts.size();
    for (I<decltype(nblocks)> b = 0; b < nblocks; ++b) {
        const auto curweight = workspace.weights[b];
        const auto& pbptr = prebuilts[b];
        if (curweight && pbptr && pbptr->num_observations()) {
            const auto curdist = compute_distance(*pbptr, workspace.distance_buffer.data(), simple_opt);
            output.first += curdist.first * curweight;
            output.second += curdist.second * curweight;
        }
    }
 
    if (workspace.total_weight) {
        output.first /= workspace.total_weight;
        output.second /= workspace.total_weight;
    }
 
    return output;
}
 
template<typename Index_, typename Input_, typename Distance_, class Matrix_ = knncolle::Matrix<Index_, Input_> >
std::vector<std::shared_ptr<const knncolle::Prebuilt<Index_, Input_, Distance_> > > build_blocked_indices(
    const std::size_t num_dim,
    const std::vector<Index_> block_sizes,
    const Input_* const data,
    const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& builder
) {
    const auto num_blocks = block_sizes.size();
    auto prebuilts = sanisizer::create<std::vector<std::shared_ptr<const knncolle::Prebuilt<Index_, Input_, Distance_> > > >(num_blocks);
 
    Index_ sofar = 0;
    for (I<decltype(num_blocks)> b = 0; b < num_blocks; ++b) {
        const auto cursize = block_sizes[b];
        if (cursize) {
            prebuilts[b] = builder.build_shared(knncolle::SimpleMatrix(num_dim, cursize, data + sanisizer::product_unsafe<std::size_t>(sofar, num_dim)));
        }
        sofar += cursize;
    }
 
    return prebuilts;
}
 
template<typename Index_, typename Input_, typename Distance_, class Matrix_ = knncolle::Matrix<Index_, Input_> >
std::pair<Distance_, Distance_> compute_distance_blocked(
    const std::size_t num_dim,
    const std::vector<Index_>& block_sizes,
    const Input_* const data,
    const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& builder,
    const BlockedOptions& options
) {
    const auto prebuilts = build_blocked_indices(num_dim, block_sizes, data, builder);
    auto workspace = create_workspace<Distance_>(block_sizes, options);
    return compute_distance_blocked(prebuilts, workspace, options);
}
 
template<typename Index_, typename Block_>
class BlockedIndicesFactory {
private:
    Index_ my_num_cells;
    const Block_* my_block;
    Block_ my_num_blocks = 0;
    std::vector<Index_> my_block_sizes;
 
    std::vector<std::pair<Index_, Index_> > my_contigs;
    Index_ my_non_contig_total = 0;
    std::vector<Index_> my_non_contig_offsets;
 
public:
    BlockedIndicesFactory(
        const Index_ num_cells,
        const Block_* block
    ) :
        my_num_cells(num_cells),
        my_block(block)
    {
        if (num_cells) {
            my_num_blocks = sanisizer::sum<Block_>(1, *std::max_element(block, block + num_cells));
        }
 
        sanisizer::resize(my_block_sizes, my_num_blocks);
        auto block_non_contig = sanisizer::create<std::vector<char> >(my_num_blocks);
        auto& block_ends = my_non_contig_offsets; // repurposing the offset vector to store the end of each contiguous block.
        sanisizer::resize(block_ends, my_num_blocks);
 
        for (Index_ c = 0; c < my_num_cells; ++c) {
            const auto curb = my_block[c];
            my_block_sizes[curb] += 1;
 
            auto& nc = block_non_contig[curb];
            if (!nc) {
                auto& be = block_ends[curb];
                if (be == 0) {
                    be = c + 1;
                } else if (be == c) {
                    ++be;
                } else {
                    nc = true;
                }
            }
        }
 
        sanisizer::resize(my_contigs, my_num_blocks);
 
        for (Block_ b = 0; b < my_num_blocks; ++b) {
            const auto length = my_block_sizes[b];
            if (block_non_contig[b]) {
                my_non_contig_offsets[b] = my_non_contig_total; 
                my_non_contig_total += length;
            } else if (length) {
                const auto start = block_ends[b] - length;
                my_contigs[b] = std::make_pair(start, length);
            }
        }
    }
 
public:
    const std::vector<Index_>& sizes() const {
        return my_block_sizes;
    }
 
    template<typename Input_>
    struct Buffers {
        std::vector<Index_> tmp_offsets;
        std::vector<Input_> tmp_buffer;
    };
 
    template<typename Input_>
    Buffers<Input_> create_buffers() const {
        return Buffers<Input_>();
    }
 
public:
    template<typename Input_, typename Distance_, class Matrix_ = knncolle::Matrix<Index_, Input_> >
    void build(
        const std::size_t num_dim,
        const Input_* const data,
        const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& builder,
        std::vector<std::shared_ptr<const knncolle::Prebuilt<Index_, Input_, Distance_> > >& output,
        Buffers<Input_>& work
    ) const {
        output.clear();
        sanisizer::resize(output, my_num_blocks);
 
        for (Block_ b = 0; b < my_num_blocks; ++b) {
            const auto& con = my_contigs[b];
            if (con.second) {
                const auto ptr = data + sanisizer::product_unsafe<std::size_t>(con.first, num_dim);
                output[b] = builder.build_shared(knncolle::SimpleMatrix(num_dim, con.second, ptr));
            }
        }
 
        if (my_non_contig_total) {
            work.tmp_buffer.resize(sanisizer::product<I<decltype(work.tmp_buffer.size())> >(my_non_contig_total, num_dim));
            work.tmp_offsets.clear();
            work.tmp_offsets.insert(work.tmp_offsets.end(), my_non_contig_offsets.begin(), my_non_contig_offsets.end());
 
            Index_ c = 0;
            while (c < my_num_cells) {
                const auto curb = my_block[c];
                const auto& con = my_contigs[curb];
                if (con.second) {
                    c += con.second; // skip past the contiguous stretch of observations.
                } else {
                    auto& curoff = work.tmp_offsets[curb];
                    std::copy_n(
                        data + sanisizer::product_unsafe<std::size_t>(c, num_dim),
                        num_dim,
                        work.tmp_buffer.data() + sanisizer::product_unsafe<std::size_t>(curoff, num_dim)
                    );
                    ++curoff;
                    ++c;
                }
            }
 
            for (Block_ b = 0; b < my_num_blocks; ++b) {
                if (my_contigs[b].second == 0) {
                    const auto length = my_block_sizes[b];
                    const auto ptr = work.tmp_buffer.data() + sanisizer::product_unsafe<std::size_t>(my_non_contig_offsets[b], num_dim);
                    output[b] = builder.build_shared(knncolle::SimpleMatrix(num_dim, length, ptr));
                }
            }
        }
    }
 
    template<typename Input_, typename Distance_, class Matrix_ = knncolle::Matrix<Index_, Input_> >
    std::vector<std::shared_ptr<const knncolle::Prebuilt<Index_, Input_, Distance_> > > build(
        const std::size_t num_dim,
        const Input_* const data,
        const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& builder
    ) const {
        std::vector<std::shared_ptr<const knncolle::Prebuilt<Index_, Input_, Distance_> > > prebuilts;
        auto bufs = create_buffers<Input_>();
        build(num_dim, data, builder, prebuilts, bufs);
        return prebuilts;
    }
};
 
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 block,
    const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& builder,
    const BlockedOptions& options
) {
    BlockedIndicesFactory<Index_, Block_> blocked_factory(num_cells, block);
    const auto prebuilts = blocked_factory.build(num_dim, data, builder);
    auto workspace = create_workspace<Distance_>(blocked_factory.sizes(), options);
    return compute_distance_blocked(prebuilts, workspace, options);
}
 
}
 
#endif