mumosa.hpp

Go to the documentation of this file.

#ifndef MUMOSA_HPP
#define MUMOSA_HPP
 
#include <vector>
#include <stdexcept>
#include <cmath>
#include <algorithm>
#include <limits>
#include <cstddef>
#include <type_traits>
 
#include "knncolle/knncolle.hpp"
#include "tatami_stats/tatami_stats.hpp"
#include "sanisizer/sanisizer.hpp"
 
namespace mumosa {
 
struct Options {
    int num_neighbors = 20;
 
    int num_threads = 1;
};
 
template<typename Input_>
std::remove_cv_t<std::remove_reference_t<Input_> > I(Input_ x) {
    return x;
}
template<typename Index_, typename Distance_>
std::pair<Distance_, Distance_> compute_distance(const Index_ num_cells, Distance_* const distances) {
    const Distance_ med = tatami_stats::medians::direct(distances, num_cells, /* skip_nan = */ false);
    Distance_ rmsd = 0;
    for (Index_ i = 0; i < num_cells; ++i) {
        const auto d = distances[i];
        rmsd += d * d;
    }
    rmsd = std::sqrt(rmsd);
    return std::make_pair(med, rmsd);
}
 
template<typename Index_, typename Input_, typename Distance_>
std::pair<Distance_, Distance_> compute_distance(const knncolle::Prebuilt<Index_, Input_, Distance_>& prebuilt, const Options& options) {
    const Index_ nobs = prebuilt.num_observations();
    const auto capped_k = knncolle::cap_k(options.num_neighbors, nobs);
    auto dist = sanisizer::create<std::vector<Distance_> >(nobs);
 
    knncolle::parallelize(options.num_threads, nobs, [&](const int, const Index_ start, const Index_ length) -> void {
        const auto searcher = prebuilt.initialize();
        std::vector<Distance_> distances;
        for (Index_ i = start, end = start + length; i < end; ++i) {
            searcher->search(i, capped_k, NULL, &distances);
            if (distances.size()) {
                dist[i] = distances.back();
            }
        }
    });
 
    return compute_distance(nobs, dist.data());
}
 
template<typename Index_, typename Input_, typename Distance_, class Matrix_ = knncolle::Matrix<Index_, Input_> >
std::pair<Distance_, Distance_> compute_distance(
    const std::size_t num_dim,
    const Index_ num_cells,
    const Input_* const data,
    const knncolle::Builder<Index_, Input_, Distance_, Matrix_>& builder,
    const Options& options)
{
    const auto prebuilt = builder.build_unique(knncolle::SimpleMatrix(num_dim, num_cells, data));
    return compute_distance(*prebuilt, options);
}
 
template<typename Distance_>
Distance_ compute_scale(const std::pair<Distance_, Distance_>& ref, const std::pair<Distance_, Distance_>& target) {
    if (target.first == 0 || ref.first == 0) {
        if (target.second == 0) {
            return std::numeric_limits<Distance_>::infinity();
        } else if (ref.second == 0) {
            return 0;
        } else {
            return ref.second / target.second; 
        }
    } else {
        return ref.first / target.first;
    }
}
 
template<typename Distance_>
std::vector<Distance_> compute_scale(const std::vector<std::pair<Distance_, Distance_> >& distances) {
    const auto ndist = distances.size();
    auto output = sanisizer::create<std::vector<Distance_> >(ndist);
 
    // Use the first entry with a non-zero RMSD as the reference.
    bool found_ref = false;
    decltype(I(ndist)) ref = 0;
    for (decltype(I(ndist)) e = 0; e < ndist; ++e) {
        if (distances[e].second) {
            found_ref = true;
            ref = e;
            break;
        }
    }
 
    // If all of them have a zero RMSD, then all scalings are zero, because it doesn't matter.
    if (found_ref) {
        const auto& dref = distances[ref];
        for (decltype(I(ndist)) e = 0; e < ndist; ++e) {
            output[e] = (e == ref ? static_cast<Distance_>(1) : compute_scale(dref, distances[e]));
        }
    }
 
    return output;
}
 
template<typename Index_, typename Input_, typename Scale_, typename Output_>
void combine_scaled_embeddings(const std::vector<std::size_t>& num_dims, const Index_ num_cells, const std::vector<Input_*>& embeddings, const std::vector<Scale_>& scaling, Output_* const output) {
    const auto nembed = num_dims.size();
    if (embeddings.size() != nembed || scaling.size() != nembed) {
        throw std::runtime_error("'num_dims', 'embeddings' and 'scale' should have the same length");
    }
 
    const std::size_t ntotal = std::accumulate(num_dims.begin(), num_dims.end(), static_cast<std::size_t>(0));
    std::size_t starting_dim = 0;
 
    for (decltype(I(nembed)) e = 0; e < nembed; ++e) {
        const auto curdim = num_dims[e];
        const auto inptr = embeddings[e];
        const auto s = scaling[e];
 
        if (std::isinf(s)) {
            // If the scaling factor is infinite, it implies that the current
            // embedding is all-zero, so we just fill with zeros, and move on.
            for (Index_ c = 0; c < num_cells; ++c) {
                const auto out_offset = sanisizer::nd_offset<std::size_t>(starting_dim, ntotal, c);
                std::fill_n(output + out_offset, curdim, 0);
            }
        } else {
            for (Index_ c = 0; c < num_cells; ++c) {
                for (decltype(I(curdim)) d = 0; d < curdim; ++d) {
                    const auto out_offset = sanisizer::nd_offset<std::size_t>(starting_dim + d, ntotal, c);
                    const auto in_offset = sanisizer::nd_offset<std::size_t>(d, curdim, c);
                    output[out_offset] = inptr[in_offset] * s;
                }
            }
        }
 
        starting_dim += curdim;
    }
}
 
}
 
#endif