train_integrated.hpp

Go to the documentation of this file.

#ifndef SINGLEPP_TRAIN_INTEGRATED_HPP
#define SINGLEPP_TRAIN_INTEGRATED_HPP
 
#include "defs.hpp"
 
#include "tatami/tatami.hpp"
 
#include "build_reference.hpp"
#include "Markers.hpp"
#include "Intersection.hpp"
#include "utils.hpp"
 
#include <vector>
#include <algorithm>
#include <cstdint>
#include <memory>
#include <optional>
#include <cassert>
 
namespace singlepp {
 
template<typename Value_, typename Index_, typename Label_>
struct TrainIntegratedInput {
    std::shared_ptr<const tatami::Matrix<Value_, Index_> > ref;
    const Label_* labels;
    PerLabelMarkers<Index_> markers;
    Index_ test_nrow;
    std::optional<Intersection<Index_> > intersection;
};
 
template<typename Value_, typename Index_, typename Label_>
TrainIntegratedInput<Value_, Index_, Label_> prepare_integrated_input(
    std::shared_ptr<const tatami::Matrix<Value_, Index_> > ref,
    const Label_* labels, 
    PerLabelMarkers<Index_> markers
) {
    TrainIntegratedInput<Value_, Index_, Label_> output;
    output.ref = std::move(ref);
    output.labels = labels;
    output.markers = std::move(markers);
    output.test_nrow = output.ref->nrow(); // remember, test and ref are assumed to have the same features.
    return output;
}
 
template<typename Index_, typename Value_, typename Label_>
TrainIntegratedInput<Value_, Index_, Label_> prepare_integrated_input(
    Index_ test_nrow,
    Intersection<Index_> intersection,
    std::shared_ptr<const tatami::Matrix<Value_, Index_> > ref, 
    const Label_* labels, 
    PerLabelMarkers<Index_> markers
) {
    TrainIntegratedInput<Value_, Index_, Label_> output;
    output.ref = std::move(ref);
    output.labels = labels;
    output.markers = std::move(markers);
    output.test_nrow = test_nrow;
    output.intersection = std::move(intersection);
    return output;
}
 
template<typename Index_, typename Id_, typename Value_, typename Label_>
TrainIntegratedInput<Value_, Index_, Label_> prepare_integrated_input(
    Index_ test_nrow,
    const Id_* test_id, 
    std::shared_ptr<const tatami::Matrix<Value_, Index_> > ref, 
    const Id_* ref_id, 
    const Label_* labels,
    PerLabelMarkers<Index_> markers
) {
    TrainIntegratedInput<Value_, Index_, Label_> output;
    output.ref = std::move(ref);
    output.labels = labels;
    output.markers = std::move(markers);
    output.test_nrow = test_nrow;
    output.intersection = intersect_genes(test_nrow, test_id, output.ref->nrow(), ref_id);
    return output;
}
 
template<typename Index_>
struct IntegratedReference {
    struct DensePerLabel {
        Index_ num_samples;
        std::vector<Index_> markers; // indices to 'universe'
        RankedVector<Index_, Index_> all_ranked; // .second contains indices to 'universe'
    };
 
    struct SparsePerLabel {
        Index_ num_samples;
        std::vector<Index_> markers; // indices to 'universe'
        RankedVector<Index_, Index_> negative_ranked, positive_ranked; // .second contains indices to 'universe'
        std::vector<std::size_t> negative_indptrs, positive_indptrs; 
    };
 
    std::optional<std::vector<DensePerLabel> > dense;
    std::optional<std::vector<SparsePerLabel> > sparse;
};
template<typename Index_>
class TrainedIntegrated {
public:
    TrainedIntegrated(Index_ test_nrow, std::vector<Index_> universe, std::vector<IntegratedReference<Index_> > references) :
        my_test_nrow(test_nrow),
        my_universe(std::move(universe)),
        my_references(std::move(references))
    {
        assert(is_sorted_unique(my_universe.size(), my_universe.data()));
    }
 
    const auto& references() const {
        return my_references;
    }
private:
    Index_ my_test_nrow;
    std::vector<Index_> my_universe;
    std::vector<IntegratedReference<Index_> > my_references;
 
public:
    std::size_t num_references() const {
        return my_references.size();
    }
 
    Index_ test_nrow() const {
        return my_test_nrow;
    }
 
    const std::vector<Index_>& subset() const {
        return my_universe;
    }
 
    std::size_t num_labels(std::size_t r) const {
        const auto& ref = my_references[r];
        if (ref.dense.has_value()) {
            return ref.dense->size();
        } else {
            return ref.sparse->size();
        }
    }
 
    std::size_t num_profiles(std::size_t r) const {
        std::size_t num_prof = 0;
        const auto& ref = my_references[r];
        if (ref.dense.has_value()) {
            for (const auto& lab : *(ref.dense)) {
                num_prof += sanisizer::sum<std::size_t>(num_prof, lab.num_samples);
            }
        } else {
            for (const auto& lab : *(ref.sparse)) {
                num_prof += sanisizer::sum<std::size_t>(num_prof, lab.num_samples);
            }
        }
        return num_prof;
    }
};
 
struct TrainIntegratedOptions {
    int num_threads = 1;
};
 
template<bool ref_sparse_, typename Value_, typename Index_, typename Label_>
void train_integrated_per_reference_simple(
    const TrainIntegratedInput<Value_, Label_, Index_>& input,
    const std::vector<Index_>& universe,
    const std::vector<Index_>& remap_test_to_universe,
    const TrainIntegratedOptions& options,
    const std::vector<Index_>& positions,
    std::vector<std::vector<RankedVector<Index_, Index_> > >& out_ranked,
    typename std::conditional<ref_sparse_, std::vector<std::vector<RankedVector<Index_, Index_> > >&, bool>::type other_ranked 
) {
    const auto& ref = *(input.ref);
    const auto NC = ref.ncol();
    const auto num_universe = universe.size();
 
    tatami::parallelize([&](int, Index_ start, Index_ len) {
        auto vbuffer = sanisizer::create<std::vector<Value_> >(num_universe);
        auto ibuffer = [&](){
            if constexpr(ref_sparse_) {
                return sanisizer::create<std::vector<Index_> >(num_universe);
            } else {
                return false;
            }
        }();
 
        RankedVector<Value_, Index_> tmp_ranked;
        tmp_ranked.reserve(num_universe);
 
        // 'universe' technically refers to the row indices of the test matrix,
        // but in simple mode, the rows of the test and reference are the same, so we can use it directly here.
        tatami::VectorPtr<Index_> universe_ptr(tatami::VectorPtr<Index_>{}, &universe);
        auto ext = tatami::consecutive_extractor<ref_sparse_>(ref, false, start, len, std::move(universe_ptr)); 
 
        for (Index_ c = start, end = start + len; c < end; ++c) {
            tmp_ranked.clear();
 
            if constexpr(ref_sparse_) {
                auto info = ext->fetch(vbuffer.data(), ibuffer.data());
                for (I<decltype(info.number)> i = 0; i < info.number; ++i) {
                    const auto remapped = remap_test_to_universe[info.index[i]];
                    assert(sanisizer::is_less_than(remapped, num_universe));
                    tmp_ranked.emplace_back(info.value[i], remapped);
                }
            } else {
                auto ptr = ext->fetch(vbuffer.data());
                for (I<decltype(num_universe)> i = 0; i < num_universe; ++i) {
                    tmp_ranked.emplace_back(ptr[i], i); // a.k.a. remap_test_to_universe[universe[i]]. 
                }
            }
 
            std::sort(tmp_ranked.begin(), tmp_ranked.end());
 
            if constexpr(ref_sparse_) {
                const auto tStart = tmp_ranked.begin(), tEnd = tmp_ranked.end();
                auto zero_ranges = find_zero_ranges<Value_, Index_>(tStart, tEnd);
                simplify_ranks<Value_, Index_>(tStart, zero_ranges.first, out_ranked[input.labels[c]][positions[c]]);
                simplify_ranks<Value_, Index_>(zero_ranges.second, tEnd, other_ranked[input.labels[c]][positions[c]]);
            } else {
                simplify_ranks(tmp_ranked, out_ranked[input.labels[c]][positions[c]]);
            }
        }
    }, NC, options.num_threads);
}
 
template<bool ref_sparse_, typename Value_, typename Index_, typename Label_>
void train_integrated_per_reference_intersect(
    const TrainIntegratedInput<Value_, Label_, Index_>& input,
    const std::vector<Index_>& remap_test_to_universe,
    const Index_ test_nrow,
    const TrainIntegratedOptions& options,
    const std::vector<Index_>& positions,
    std::vector<std::vector<RankedVector<Index_, Index_> > >& out_ranked,
    typename std::conditional<ref_sparse_, std::vector<std::vector<RankedVector<Index_, Index_> > >&, bool>::type other_ranked 
) {
    const auto& ref = *(input.ref);
    const auto NC = ref.ncol();
 
    std::vector<Index_> ref_subset;
    sanisizer::reserve(ref_subset, input.intersection->size());
    auto remap_ref_subset_to_universe = sanisizer::create<std::vector<Index_> >(ref.nrow(), test_nrow); // all entries of remap_test_to_universe are less than test_nrow. 
    for (const auto& pair : *(input.intersection)) {
        const auto rdex = remap_test_to_universe[pair.first];
        if (rdex != test_nrow) {
            ref_subset.push_back(pair.second);
            remap_ref_subset_to_universe[pair.second] = rdex;
        }
    }
    std::sort(ref_subset.begin(), ref_subset.end());
 
    typename std::conditional<ref_sparse_, bool, std::vector<Index_> >::type remap_dense_to_universe;
    if constexpr(!ref_sparse_) {
        remap_dense_to_universe.reserve(ref_subset.size());
        for (auto r : ref_subset) {
            remap_dense_to_universe.push_back(remap_ref_subset_to_universe[r]);
        }
    }
 
    tatami::parallelize([&](int, Index_ start, Index_ len) {
        const auto ref_subset_size = ref_subset.size();
        auto vbuffer = sanisizer::create<std::vector<Value_> >(ref_subset_size);
        auto ibuffer = [&]() {
            if constexpr(ref_sparse_) {
                return sanisizer::create<std::vector<Index_> >(ref_subset_size);
            } else {
                return false;
            }
        }();
 
        RankedVector<Value_, Index_> tmp_ranked;
        tmp_ranked.reserve(ref_subset_size);
        tatami::VectorPtr<Index_> to_extract_ptr(tatami::VectorPtr<Index_>{}, &ref_subset);
        auto ext = tatami::consecutive_extractor<ref_sparse_>(ref, false, start, len, std::move(to_extract_ptr));
 
        for (Index_ c = start, end = start + len; c < end; ++c) {
            tmp_ranked.clear();
 
            if constexpr(ref_sparse_) {
                auto info = ext->fetch(vbuffer.data(), ibuffer.data());
                for (I<decltype(info.number)> i = 0; i < info.number; ++i) {
                    tmp_ranked.emplace_back(info.value[i], remap_ref_subset_to_universe[info.index[i]]);
                }
            } else {
                auto ptr = ext->fetch(vbuffer.data());
                for (I<decltype(ref_subset_size)> i = 0; i < ref_subset_size; ++i) {
                    tmp_ranked.emplace_back(ptr[i], remap_dense_to_universe[i]);
                }
            }
 
            std::sort(tmp_ranked.begin(), tmp_ranked.end());
 
            if constexpr(ref_sparse_) {
                const auto tStart = tmp_ranked.begin(), tEnd = tmp_ranked.end();
                auto zero_ranges = find_zero_ranges<Value_, Index_>(tStart, tEnd);
                simplify_ranks<Value_, Index_>(tStart, zero_ranges.first, out_ranked[input.labels[c]][positions[c]]);
                simplify_ranks<Value_, Index_>(zero_ranges.second, tEnd, other_ranked[input.labels[c]][positions[c]]);
            } else {
                simplify_ranks(tmp_ranked, out_ranked[input.labels[c]][positions[c]]);
            }
        }
    }, NC, options.num_threads);
}
template<typename Value_, typename Index_, typename Label_>
TrainedIntegrated<Index_> train_integrated(const std::vector<TrainIntegratedInput<Value_, Index_, Label_> >& inputs, const TrainIntegratedOptions& options) {
 
    // Checking that the number of genes in the test dataset are consistent.
    Index_ test_nrow = 0;
    if (inputs.size()) {
        test_nrow = inputs.front().test_nrow;
        for (const auto& in : inputs) {
            if (!sanisizer::is_equal(in.test_nrow, test_nrow)) {
                throw std::runtime_error("inconsistent number of rows in the test dataset across entries of 'inputs'");
            }
        }
    }
 
    // For references with intersections, we need to map the marker indices to the test rows, for comparability across references.
    const auto nrefs = inputs.size();
    std::vector<std::vector<Index_> > remap_intersection_to_test_index; 
    for (I<decltype(nrefs)> r  = 0; r < nrefs; ++r) {
        if (inputs[r].intersection.has_value()) {
            sanisizer::resize(remap_intersection_to_test_index, nrefs);
            break;
        }
    }
 
    // Identify the union of all marker genes as the universe, but excluding those markers that are not present in all intersections.
    // Specifically, 'universe' contains sorted and unique row indices for the test matrix, where 'remap_test_to_universe[universe[i]] == i'.
    // 'remap_test_to_universe[k]' defaults to the max number of rows in the test if 'k' is not present in all intersections. 
    std::vector<Index_> universe;
    auto remap_test_to_universe = sanisizer::create<std::vector<Index_> >(test_nrow, test_nrow); 
    {
        auto present = sanisizer::create<std::vector<char> >(test_nrow);
        auto count_refs = sanisizer::create<std::vector<I<decltype(nrefs)> > >(test_nrow);
        universe.reserve(test_nrow);
 
        for (I<decltype(nrefs)> r  = 0; r < nrefs; ++r) {
            const auto& markers = inputs[r].markers;
            const auto& inter = inputs[r].intersection;
 
            if (inter.has_value()) {
                auto& cur_test_remap = remap_intersection_to_test_index[r];
                sanisizer::resize(cur_test_remap, inputs[r].ref->nrow(), test_nrow); // again, default to the max number of rows if not present in the current intersection.
                for (const auto& pp : *inter) {
                    cur_test_remap[pp.second] = pp.first;
                    count_refs[pp.first] += 1;
                }
 
                for (const auto& labmrk : markers) {
                    for (const auto y : labmrk) {
                        const auto ty = cur_test_remap[y];
                        if (ty != test_nrow && !present[ty]) {
                            present[ty] = true;
                            universe.push_back(ty);
                        }
                    }
                }
 
            } else {
                for (const auto& labmrk : markers) {
                    for (const auto y : labmrk) {
                        if (!present[y]) {
                            present[y] = true;
                            universe.push_back(y);
                        }
                    }
                }
 
                for (auto& x : count_refs) {
                    x += 1;
                }
            }
        }
 
        std::sort(universe.begin(), universe.end());
        const auto num_universe = universe.size();
        I<decltype(num_universe)> keep = 0;
        for (I<decltype(num_universe)> u = 0; u < num_universe; ++u) {
            const auto marker = universe[u];
            if (count_refs[marker] == nrefs) {
                universe[keep] = marker;
                remap_test_to_universe[marker] = keep;
                ++keep;
            }
        }
        universe.resize(keep);
        universe.shrink_to_fit();
    }
 
    // Remapping the per-label markers for this reference to refer to our universe.
    auto references = sanisizer::create<std::vector<IntegratedReference<Index_> > >(nrefs);
    for (I<decltype(nrefs)> r = 0; r < nrefs; ++r) {
        const auto& curinput = inputs[r];
        const auto& currefmarkers = curinput.markers;
        const auto nlabels = currefmarkers.size();
        auto& currefout = references[r];
 
        const bool is_sparse = curinput.ref->is_sparse();
        if (is_sparse) {
            currefout.sparse.emplace(sanisizer::as_size_type<I<decltype(*(currefout.sparse))> >(nlabels));
        } else {
            currefout.dense.emplace(sanisizer::as_size_type<I<decltype(*(currefout.dense))> >(nlabels));
        }
 
        auto get_markers = [&](I<decltype(nlabels)> l) -> std::vector<Index_>& {
            if (is_sparse) {
                return (*(currefout.sparse))[l].markers; 
            } else {
                return (*(currefout.dense))[l].markers;
            }
        };
 
        if (curinput.intersection.has_value()) {
            auto& cur_test_remap = remap_intersection_to_test_index[r];
            for (I<decltype(nlabels)> l = 0; l < nlabels; ++l) {
                const auto& curlabmarkers = currefmarkers[l];
                auto& markers = get_markers(l);
                markers.reserve(curlabmarkers.size());
                for (const auto y : curlabmarkers) {
                    const auto ty = cur_test_remap[y];
                    if (ty != test_nrow) { // ignoring marker that can't be mapped via the current intersection.
                        const auto universe_index = remap_test_to_universe[ty];
                        if (universe_index != test_nrow) { // ignoring genes not present in all intersections.
                            markers.push_back(universe_index);
                        }
                    }
                }
            }
 
        } else {
            for (I<decltype(nlabels)> l = 0; l < nlabels; ++l) {
                const auto& curlabmarkers = currefmarkers[l];
                auto& markers = get_markers(l);
                markers.reserve(curlabmarkers.size());
                for (const auto y : curlabmarkers) {
                    const auto universe_index = remap_test_to_universe[y];
                    if (universe_index != test_nrow) { // ignoring genes not present in all intersections.
                        markers.push_back(universe_index);
                    }
                }
            }
        }
    }
 
    // Not needed after this so we try to free its allocation for recycling in the next malloc call.
    // Specifically, I want to give a chance for this memory to be re-used in train_integrated_per_reference_intersect.
    remap_intersection_to_test_index.clear();
 
    // Now, we create ranked vectors for each profile in the reference.
    for (I<decltype(nrefs)> r = 0; r < nrefs; ++r) {
        const auto& curinput = inputs[r];
        auto& currefout = references[r];
 
        const Index_ NC = curinput.ref->ncol();
        if (NC == 0) {
            throw std::runtime_error("reference dataset must have at least one column");
        }
        std::vector<Index_> positions;
        sanisizer::reserve(positions, NC);
 
        const auto nlabels = sanisizer::sum<std::size_t>(*std::max_element(curinput.labels, curinput.labels + NC), 1);
        auto samples_per_label = sanisizer::create<std::vector<Index_> >(nlabels);
        for (Index_ c = 0; c < NC; ++c) {
            auto& pos = samples_per_label[curinput.labels[c]];
            positions.push_back(pos);
            ++pos;
        }
 
        for (I<decltype(nlabels)> l = 0; l < nlabels; ++l) {
            if (samples_per_label[l] == 0) {
                throw std::runtime_error("no profiles available for label " + std::to_string(l) + " in reference " + std::to_string(r));
            }
        }
 
        if (!sanisizer::is_equal(curinput.markers.size(), nlabels)) {
            throw std::runtime_error("'markers' length should be equal to the number of unique labels");
        }
 
        if (curinput.ref->is_sparse()) {
            auto negative_ranked = sanisizer::create<std::vector<std::vector<RankedVector<Index_, Index_> > > >(nlabels);
            auto positive_ranked = sanisizer::create<std::vector<std::vector<RankedVector<Index_, Index_> > > >(nlabels);
            for (I<decltype(nlabels)> l = 0; l < nlabels; ++l) {
                const auto num_samples = samples_per_label[l];
                sanisizer::resize(negative_ranked[l], num_samples);
                sanisizer::resize(positive_ranked[l], num_samples);
            }
 
            if (curinput.intersection) {
                train_integrated_per_reference_intersect<true>(curinput, remap_test_to_universe, test_nrow, options, positions, negative_ranked, positive_ranked);
            } else {
                train_integrated_per_reference_simple<true, Value_>(curinput, universe, remap_test_to_universe, options, positions, negative_ranked, positive_ranked);
            }
 
            for (I<decltype(nlabels)> l = 0; l < nlabels; ++l) {
                auto& curlabout = (*(currefout.sparse))[l];
                const auto num_samples = samples_per_label[l];
                curlabout.num_samples = num_samples;
 
                I<decltype(curlabout.negative_ranked.size())> num_neg = 0;
                for (const auto& x : negative_ranked[l]) {
                    num_neg = sanisizer::sum<I<decltype(num_neg)> >(num_neg, x.size());
                }
 
                I<decltype(curlabout.positive_ranked.size())> num_pos = 0;
                for (const auto& x : positive_ranked[l]) {
                    num_pos = sanisizer::sum<I<decltype(num_pos)> >(num_pos, x.size());
                }
 
                curlabout.negative_ranked.reserve(num_neg);
                curlabout.negative_indptrs.reserve(sanisizer::sum<I<decltype(curlabout.negative_indptrs.size())> >(num_samples, 1));
                curlabout.negative_indptrs.push_back(0);
                for (const auto& x : negative_ranked[l]) {
                    curlabout.negative_ranked.insert(curlabout.negative_ranked.end(), x.begin(), x.end());
                    curlabout.negative_indptrs.push_back(curlabout.negative_ranked.size());
                }
 
                curlabout.positive_ranked.reserve(num_pos);
                curlabout.positive_indptrs.reserve(sanisizer::sum<I<decltype(curlabout.positive_indptrs.size())> >(num_samples, 1));
                curlabout.positive_indptrs.push_back(0);
                for (const auto& x : positive_ranked[l]) {
                    curlabout.positive_ranked.insert(curlabout.positive_ranked.end(), x.begin(), x.end());
                    curlabout.positive_indptrs.push_back(curlabout.positive_ranked.size());
                }
            }
 
        } else {
            auto out_ranked = sanisizer::create<std::vector<std::vector<RankedVector<Index_, Index_> > > >(nlabels);
            for (I<decltype(nlabels)> l = 0; l < nlabels; ++l) {
                const auto num_samples = samples_per_label[l];
                sanisizer::resize(out_ranked[l], num_samples);
            }
 
            if (curinput.intersection) {
                train_integrated_per_reference_intersect<false>(curinput, remap_test_to_universe, test_nrow, options, positions, out_ranked, true);
            } else {
                train_integrated_per_reference_simple<false, Value_>(curinput, universe, remap_test_to_universe, options, positions, out_ranked, true);
            }
 
            for (I<decltype(nlabels)> l = 0; l < nlabels; ++l) {
                auto& curlabout = (*(currefout.dense))[l];
                curlabout.num_samples = samples_per_label[l];
                curlabout.all_ranked.reserve(sanisizer::product<I<decltype(curlabout.all_ranked.size())> >(universe.size(), curlabout.num_samples));
                for (const auto& x : out_ranked[l]) {
                    curlabout.all_ranked.insert(curlabout.all_ranked.end(), x.begin(), x.end());
                }
            }
        }
    }
 
    return TrainedIntegrated<Index_>(test_nrow, std::move(universe), std::move(references));
}
 
}
 
#endif