sparse_matrix.hpp

Go to the documentation of this file.

#ifndef TATAMI_PYTHON_SPARSE_MATRIX_HPP
#define TATAMI_PYTHON_SPARSE_MATRIX_HPP
 
#include "tatami/tatami.hpp"
#include "pybind11/pybind11.h"
 
#include "utils.hpp"
 
#include <algorithm>
#include <cstdint>
 
namespace tatami_python { 
 
template<typename Type_>
void dump_to_buffer(const pybind11::array& input, Type_* const buffer) {
    auto dtype = input.dtype();
    if (dtype.is(pybind11::dtype::of<double>())) {
        std::copy_n(static_cast<const double*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<float>())) {
        std::copy_n(static_cast<const float*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::int64_t>())) {
        std::copy_n(static_cast<const std::int64_t*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::int32_t>())) {
        std::copy_n(static_cast<const std::int32_t*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::int16_t>())) {
        std::copy_n(static_cast<const std::int16_t*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::int8_t>())) {
        std::copy_n(static_cast<const std::int8_t*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::uint64_t>())) {
        std::copy_n(static_cast<const std::uint64_t*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::uint32_t>())) {
        std::copy_n(static_cast<const std::uint32_t*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::uint16_t>())) {
        std::copy_n(static_cast<const std::uint16_t*>(input.request().ptr), input.size(), buffer);
 
    } else if (dtype.is(pybind11::dtype::of<std::uint8_t>())) {
        std::copy_n(static_cast<const std::uint8_t*>(input.request().ptr), input.size(), buffer);
 
    } else {
        throw std::runtime_error("unrecognized array type '" + std::string(dtype.kind(), 1) + std::to_string(dtype.itemsize()) + "' from 'extract_sparse_array()'");
    }
}
template<typename Value_, typename Index_, class Function_>
void parse_Sparse2darray(const pybind11::object& matrix, Value_* const vbuffer, Index_* const ibuffer, Function_ fun) {
    pybind11::object raw_svt = matrix.attr("contents");
    if (pybind11::isinstance<pybind11::none>(raw_svt)) {
        return;
    }
    auto svt = raw_svt.template cast<pybind11::list>();
 
    const auto shape = get_shape<Index_>(matrix);
    const auto NC = shape.second;
 
    for (I<decltype(NC)> c = 0; c < NC; ++c) {
        pybind11::object raw_inner(svt[c]);
        if (pybind11::isinstance<pybind11::none>(raw_inner)) {
            continue;
        }
 
        auto inner = raw_inner.template cast<pybind11::tuple>();
        if (inner.size() != 2) {
            auto ctype = get_class_name(matrix);
            throw std::runtime_error("each entry of '<" + ctype + ">.contents' should be a tuple of length 2 or None");
        }
 
        auto iinput = inner[0].template cast<pybind11::array>();
        if (ibuffer != NULL) {
            dump_to_buffer(iinput, ibuffer);
        }
        if (vbuffer != NULL) {
            auto vinput = inner[1].template cast<pybind11::array>();
            dump_to_buffer(vinput, vbuffer);
        }
 
        // cast is known to be safe as the length of these vectors cannot excced
        // the number of rows, the latter of which must fit in an Index_.
        fun(c, static_cast<Index_>(iinput.size()));
    }
}
 
template<typename CachedValue_, typename CachedIndex_, typename Index_>
void parse_sparse_matrix(
    const pybind11::object& matrix,
    bool row,
    std::vector<CachedValue_*>& value_ptrs, 
    CachedValue_* const vbuffer,
    std::vector<CachedIndex_*>& index_ptrs, 
    CachedIndex_* const ibuffer,
    Index_* const counts
) {
    const bool needs_value = !value_ptrs.empty();
    const bool needs_index = !index_ptrs.empty();
 
    parse_Sparse2darray(
        matrix,
        (needs_value ? vbuffer : NULL),
        (needs_index || row ? ibuffer : NULL),
        [&](const Index_ c, const Index_ nnz) -> void {
            // Note that non-empty value_ptrs and index_ptrs may be longer than the
            // number of rows/columns in the SVT matrix, due to the reuse of slabs.
            if (row) {
                if (needs_value) {
                    for (I<decltype(nnz)> i = 0; i < nnz; ++i) {
                        auto ix = ibuffer[i];
                        value_ptrs[ix][counts[ix]] = vbuffer[i];
                    }
                }
                if (needs_index) {
                    for (I<decltype(nnz)> i = 0; i < nnz; ++i) {
                        auto ix = ibuffer[i];
                        index_ptrs[ix][counts[ix]] = c;
                    }
                }
                for (I<decltype(nnz)> i = 0; i < nnz; ++i) {
                    ++(counts[ibuffer[i]]);
                }
 
            } else {
                if (needs_value) {
                    std::copy_n(vbuffer, nnz, value_ptrs[c]);
                }
                if (needs_index) {
                    std::copy_n(ibuffer, nnz, index_ptrs[c]);
                }
                counts[c] = nnz;
            }
        }
    );
}
}
 
#endif