mock_sparse_chunk.hpp

Go to the documentation of this file.

#ifndef TATAMI_CHUNKED_MOCK_SPARSE_CHUNK_HPP
#define TATAMI_CHUNKED_MOCK_SPARSE_CHUNK_HPP
 
#include <vector>
#include <cstdint>
 
namespace tatami_chunked {
 
namespace MockSparseChunk_internal {
 
template<typename InflatedValue_, typename InflatedIndex_>
struct Workspace {
    // Standard compressed sparse members:
    std::vector<InflatedValue_> values;
    std::vector<InflatedIndex_> indices;
    std::vector<size_t> indptrs;
 
    // Allocation to allow for O(1) mapping of requested indices to sparse indices.
    // This mimics what is done in the indexed sparse extractors in tatami proper.
    std::vector<uint8_t> remap;
};
 
template<class Blob_>
class Core {
public:
    Core() = default;
 
    Core(Blob_ chunk) : my_chunk(std::move(chunk)) {}
 
private:
    Blob_ my_chunk;
 
    typedef typename Blob_::value_type value_type;
 
    typedef typename Blob_::index_type index_type;
 
public:
    auto get_target_chunkdim(bool row) const {
        if (row) {
            return my_chunk.nrow();
        } else {
            return my_chunk.ncol();
        }
    }
 
    auto get_non_target_chunkdim(bool row) const {
        if (row) {
            return my_chunk.ncol();
        } else {
            return my_chunk.nrow();
        }
    }
 
    typedef Workspace<value_type, index_type> MyWorkspace;
 
private:
    template<typename Index_>
    static void refine_start_and_end(size_t& start, size_t& end, Index_ desired_start, Index_ desired_end, Index_ max_end, const std::vector<index_type>& indices) {
        if (desired_start) {
            auto it = indices.begin();
            // Using custom comparator to ensure that we cast to Index_ for signedness-safe comparisons.
            start = std::lower_bound(it + start, it + end, desired_start, [](Index_ a, Index_ b) -> bool { return a < b; }) - it;
        }
 
        if (desired_end != max_end) {
            if (desired_end == desired_start + 1) {
                if (start != end && static_cast<Index_>(indices[start]) == desired_start) {
                    end = start + 1;
                } else {
                    end = start;
                }
            } else {
                auto it = indices.begin();
                end = std::lower_bound(it + start, it + end, desired_end, [](Index_ a, Index_ b) -> bool { return a < b; }) - it;
            }
        }
    }
 
    // Building a present/absent mapping for the requested indices to allow O(1) look-up from the sparse matrix indices.
    template<typename Index_>
    static void configure_remap(MyWorkspace& work, const std::vector<Index_>& indices, size_t full) {
        work.remap.resize(full);
        for (auto i : indices) {
            work.remap[i] = 1;
        }
    }
 
    // Resetting just the affected indices so we can avoid a fill operation over the entire array.
    template<typename Index_>
    static void reset_remap(MyWorkspace& work, const std::vector<Index_>& indices) {
        for (auto i : indices) {
            work.remap[i] = 0;
        }
    }
 
private:
    template<bool is_block_, typename Index_>
    void fill_target(
        Index_ p, 
        Index_ non_target_start, 
        Index_ non_target_end, 
        Index_ non_target_chunkdim,
        MyWorkspace& work, 
        const std::vector<value_type*>& output_values, 
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        size_t start = work.indptrs[p], end = work.indptrs[p + 1];
        if (start >= end) {
            return;
        }
 
        refine_start_and_end(start, end, non_target_start, non_target_end, non_target_chunkdim, work.indices);
 
        auto& current_number = output_number[p];
        const bool needs_value = !output_values.empty();
        auto vptr = needs_value ? output_values[p] + current_number : NULL;
        const bool needs_index = !output_indices.empty();
        auto iptr = needs_index ? output_indices[p] + current_number : NULL;
 
        if constexpr(is_block_) {
            if (needs_value) {
                std::copy(work.values.begin() + start, work.values.begin() + end, vptr);
            }
            if (needs_index) {
                for (size_t i = start; i < end; ++i, ++iptr) {
                    *iptr = work.indices[i] + shift;
                }
            }
            current_number += end - start;
 
        } else {
            // Assumes that work.remap has been properly configured, see configure_remap().
            for (size_t i = start; i < end; ++i) {
                Index_ target = work.indices[i];
                if (work.remap[target]) {
                    if (needs_value) {
                        *vptr = work.values[i];
                        ++vptr;
                    }
                    if (needs_index) {
                        *iptr = target + shift;
                        ++iptr;
                    }
                    ++current_number;
                }
            }
        }
    }
 
    template<bool is_block_, typename Index_>
    void fill_secondary(
        Index_ s,
        Index_ target_start, 
        Index_ target_end, 
        Index_ target_chunkdim,
        MyWorkspace& work, 
        const std::vector<value_type*>& output_values, 
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        auto start = work.indptrs[s], end = work.indptrs[s + 1];
        if (start >= end) {
            return;
        }
 
        refine_start_and_end(start, end, target_start, target_end, target_chunkdim, work.indices);
 
        bool needs_value = !output_values.empty();
        bool needs_index = !output_indices.empty();
 
        if constexpr(is_block_) {
            for (size_t i = start; i < end; ++i) {
                auto p = work.indices[i];
                auto& num = output_number[p];
                if (needs_value) {
                    output_values[p][num] = work.values[i];
                }
                if (needs_index) {
                    output_indices[p][num] = s + shift;
                }
                ++num;
            }
 
        } else {
            // Assumes that work.remap has been properly configured, see configure_remap().
            for (size_t i = start; i < end; ++i) {
                Index_ target = work.indices[i];
                if (work.remap[target]) {
                    auto& num = output_number[target];
                    if (needs_value) {
                        output_values[target][num] = work.values[i];
                    }
                    if (needs_index) {
                        output_indices[target][num] = s + shift;
                    }
                    ++num;
                }
            }
        }
    }
 
public:
    template<typename Index_>
    void extract(
        bool row,
        Index_ target_start,
        Index_ target_length,
        Index_ non_target_start,
        Index_ non_target_length,
        MyWorkspace& work,
        const std::vector<value_type*>& output_values, 
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_chunk.inflate(work.values, work.indices, work.indptrs);
        Index_ target_end = target_start + target_length;
        Index_ non_target_end = non_target_start + non_target_length;
 
        if (my_chunk.is_csr() == row) {
            Index_ non_target_chunkdim = get_non_target_chunkdim(row);
            for (Index_ p = target_start; p < target_end; ++p) {
                fill_target<true>(p, non_target_start, non_target_end, non_target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
        } else {
            Index_ target_chunkdim = get_target_chunkdim(row);
            for (Index_ s = non_target_start; s < non_target_end; ++s) {
                fill_secondary<true>(s, target_start, target_end, target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
        }
    }
 
    template<typename Index_>
    void extract(
        bool row,
        Index_ target_start,
        Index_ target_length,
        const std::vector<Index_>& non_target_indices,
        MyWorkspace& work,
        const std::vector<value_type*>& output_values, 
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_chunk.inflate(work.values, work.indices, work.indptrs);
        Index_ target_end = target_start + target_length;
 
        if (my_chunk.is_csr() == row) {
            // non_target_indices is guaranteed to be non-empty, see contracts below.
            auto non_target_start = non_target_indices.front();
            auto non_target_end = non_target_indices.back() + 1; // need 1 past end.
            auto non_target_chunkdim = get_non_target_chunkdim(row);
 
            configure_remap(work, non_target_indices, non_target_chunkdim);
            for (Index_ p = target_start; p < target_end; ++p) {
                fill_target<false>(p, non_target_start, non_target_end, non_target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
            reset_remap(work, non_target_indices);
 
        } else {
            Index_ target_chunkdim = get_target_chunkdim(row);
            for (auto s : non_target_indices) {
                fill_secondary<true>(s, target_start, target_end, target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
        }
    }
 
public:
    template<typename Index_>
    void extract(
        bool row,
        const std::vector<Index_>& target_indices,
        Index_ non_target_start,
        Index_ non_target_length,
        MyWorkspace& work,
        const std::vector<value_type*>& output_values, 
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_chunk.inflate(work.values, work.indices, work.indptrs);
        Index_ non_target_end = non_target_start + non_target_length;
 
        if (my_chunk.is_csr() == row) {
            Index_ non_target_chunkdim = get_non_target_chunkdim(row);
            for (auto p : target_indices) {
                fill_target<true>(p, non_target_start, non_target_end, non_target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
 
        } else {
            // target_indices is guaranteed to be non-empty, see contracts below.
            auto target_start = target_indices.front();
            auto target_end = target_indices.back() + 1; // need 1 past end.
            auto target_chunkdim = get_target_chunkdim(row);
 
            configure_remap(work, target_indices, target_chunkdim);
            for (Index_ s = non_target_start; s < non_target_end; ++s) {
                fill_secondary<false>(s, target_start, target_end, target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
            reset_remap(work, target_indices);
        }
    }
 
    template<typename Index_>
    void extract(
        bool row,
        const std::vector<Index_>& target_indices,
        const std::vector<Index_>& non_target_indices,
        MyWorkspace& work,
        const std::vector<value_type*>& output_values, 
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_chunk.inflate(work.values, work.indices, work.indptrs);
 
        if (my_chunk.is_csr() == row) {
            // non_target_indices is guaranteed to be non-empty, see contracts below.
            auto non_target_start = non_target_indices.front();
            auto non_target_end = non_target_indices.back() + 1; // need 1 past end.
            auto non_target_chunkdim = get_non_target_chunkdim(row);
 
            configure_remap(work, non_target_indices, non_target_chunkdim);
            for (auto p : target_indices) {
                fill_target<false>(p, non_target_start, non_target_end, non_target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
            reset_remap(work, non_target_indices);
 
        } else {
            // target_indices is guaranteed to be non-empty, see contracts below.
            auto target_start = target_indices.front();
            auto target_end = target_indices.back() + 1; // need 1 past end.
            auto target_chunkdim = get_target_chunkdim(row);
 
            configure_remap(work, target_indices, target_chunkdim);
            for (auto s : non_target_indices) {
                fill_secondary<false>(s, target_start, target_end, target_chunkdim, work, output_values, output_indices, output_number, shift);
            }
            reset_remap(work, target_indices);
        }
    }
};
 
class MockBlob {
public:
    typedef int index_type;
    typedef double value_type;
 
    bool is_csr() const {
        return true;
    }
 
private:
    int my_nrow, my_ncol;
    std::vector<double> my_values;
    std::vector<int> my_indices;
    std::vector<size_t> my_pointers;
 
public:
    MockBlob() = default;
 
    MockBlob(int nrow, int ncol, std::vector<double> values, std::vector<int> indices, std::vector<size_t> pointers) : 
        my_nrow(nrow), my_ncol(ncol), my_values(std::move(values)), my_indices(std::move(indices)), my_pointers(std::move(pointers)) {}
 
    int nrow() const {
        return my_nrow;
    }
 
    int ncol() const {
        return my_ncol;
    }
 
    void inflate(std::vector<double>& vbuffer, std::vector<int>& ibuffer, std::vector<size_t>& pbuffer) const {
        vbuffer.resize(my_values.size());
        std::copy(my_values.begin(), my_values.end(), vbuffer.begin());
        ibuffer.resize(my_indices.size());
        std::copy(my_indices.begin(), my_indices.end(), ibuffer.begin());
        pbuffer.resize(my_pointers.size());
        std::copy(my_pointers.begin(), my_pointers.end(), pbuffer.begin());
    }
};
 
}
class MockSimpleSparseChunk {
public:
    typedef double value_type;
 
    struct Workspace {
        // Hiding this here to avoid giving the impression that we NEED to implement this.
        MockSparseChunk_internal::Workspace<value_type, int> work;
    };
 
    static constexpr bool use_subset = false;
 
public:
    // You can construct this however you like, I don't care.
    MockSimpleSparseChunk() = default;
    MockSimpleSparseChunk(int nr, int nc, std::vector<double> x, std::vector<int> i, std::vector<size_t> p) : 
        core(MockSparseChunk_internal::MockBlob(nr, nc, std::move(x), std::move(i), std::move(p))) {}
private:
    MockSparseChunk_internal::Core<MockSparseChunk_internal::MockBlob> core;
 
public:
    template<typename Index_>
    void extract(
        bool row,
        Index_ non_target_start, 
        Index_ non_target_length, 
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        core.template extract<Index_>(
            row, 
            0, 
            core.get_target_chunkdim(row), 
            non_target_start, 
            non_target_length, 
            work.work, 
            output_values, 
            output_indices, 
            output_number,
            shift
        );
    }
 
    template<typename Index_>
    void extract(
        bool row,
        const std::vector<Index_>& non_target_indices,
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        core.template extract<Index_>(
            row, 
            0, 
            core.get_target_chunkdim(row), 
            non_target_indices, 
            work.work, 
            output_values, 
            output_indices,
            output_number,
            shift
        );
    }
};
 
template<class Blob_>
class SimpleSparseChunkWrapper {
public:
    typedef typename Blob_::value_type value_type;
 
    typedef MockSparseChunk_internal::Workspace<value_type, typename Blob_::index_type> Workspace;
 
    static constexpr bool use_subset = false;
 
    SimpleSparseChunkWrapper() = default;
 
    SimpleSparseChunkWrapper(Blob_ core) : my_core(std::move(core)) {}
 
private:
    MockSparseChunk_internal::Core<Blob_> my_core;
 
public:
    template<typename Index_>
    void extract(
        bool row,
        Index_ non_target_start, 
        Index_ non_target_length, 
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_core.template extract<Index_>(
            row, 
            0, 
            my_core.get_target_chunkdim(row), 
            non_target_start, 
            non_target_length, 
            work, 
            output_values, 
            output_indices, 
            output_number,
            shift
        );
    }
 
    template<typename Index_>
    void extract(
        bool row,
        const std::vector<Index_>& non_target_indices,
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_core.template extract<Index_>(
            row, 
            0, 
            my_core.get_target_chunkdim(row), 
            non_target_indices, 
            work, 
            output_values, 
            output_indices,
            output_number,
            shift
        );
    }
};
 
class MockSubsettedSparseChunk {
public:
    typedef double value_type;
 
    struct Workspace {
        // Hiding this here to avoid giving the impression that we NEED to implement this.
        MockSparseChunk_internal::Workspace<value_type, int> work;
    };
 
    static constexpr bool use_subset = true;
 
    MockSubsettedSparseChunk() = default;
    MockSubsettedSparseChunk(int nrow, int ncol, std::vector<double> values, std::vector<int> indices, std::vector<size_t> pointers) : 
        my_core(MockSparseChunk_internal::MockBlob(nrow, ncol, std::move(values), std::move(indices), std::move(pointers))) {}
private:
    MockSparseChunk_internal::Core<MockSparseChunk_internal::MockBlob> my_core;
 
public:
    template<typename Index_>
    void extract(
        bool row,
        Index_ target_start, 
        Index_ target_length, 
        Index_ non_target_start, 
        Index_ non_target_length, 
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_core.extract(
            row, 
            target_start, 
            target_length, 
            non_target_start, 
            non_target_length, 
            work.work,
            output_values,
            output_indices,
            output_number,
            shift
        );
    }
 
    template<typename Index_>
    void extract(
        bool row,
        Index_ target_start, 
        Index_ target_length, 
        const std::vector<Index_>& non_target_indices, 
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_core.extract(
            row, 
            target_start, 
            target_length, 
            non_target_indices, 
            work.work, 
            output_values, 
            output_indices, 
            output_number,
            shift
        );
    }
 
public:
    template<typename Index_>
    void extract(
        bool row,
        const std::vector<Index_>& target_indices, 
        Index_ non_target_start, 
        Index_ non_target_length, 
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_core.extract(
            row, 
            target_indices, 
            non_target_start, 
            non_target_length, 
            work.work, 
            output_values,
            output_indices,
            output_number,
            shift
        );
    }
 
    template<typename Index_>
    void extract(
        bool row,
        const std::vector<Index_>& target_indices, 
        const std::vector<Index_>& non_target_indices, 
        Workspace& work, 
        const std::vector<value_type*>& output_values,
        const std::vector<Index_*>& output_indices,
        Index_* output_number,
        Index_ shift)
    const {
        my_core.extract(
            row, 
            target_indices, 
            non_target_indices, 
            work.work, 
            output_values, 
            output_indices, 
            output_number,
            shift 
        );
    }
};
 
}
 
#endif