ScaleByNeighbors.hpp

Go to the documentation of this file.

#ifndef SCRAN_NEIGHBOR_SCALING_HPP
#define SCRAN_NEIGHBOR_SCALING_HPP
 
#include "../utils/macros.hpp"
 
#include <vector>
#include <cmath>
#include <memory>
#include "knncolle/knncolle.hpp"
#include "tatami/stats/medians.hpp"
 
namespace scran {
 
class ScaleByNeighbors {
public:
    struct Defaults { 
        static constexpr int neighbors = 20;
 
        static constexpr bool approximate = false;
 
        static constexpr int num_threads = 1;
    };
 
    ScaleByNeighbors& set_neighbors(int n = Defaults::neighbors) {
        num_neighbors = n;
        return *this;
    }
 
    ScaleByNeighbors& set_approximate(int a = Defaults::approximate) {
        approximate = a;
        return *this;
    }
 
    ScaleByNeighbors& set_num_threads(int n = Defaults::num_threads) {
        nthreads = n;
        return *this;
    }
 
private:
    int num_neighbors = Defaults::neighbors;
    bool approximate = Defaults::approximate;
    int nthreads = Defaults::num_threads;
 
public:
    std::pair<double, double> compute_distance(int ndim, size_t ncells, const double* data) const {
        std::unique_ptr<knncolle::Base<> > ptr;
        if (!approximate) {
            ptr.reset(new knncolle::VpTreeEuclidean<>(ndim, ncells, data));
        } else {
            ptr.reset(new knncolle::AnnoyEuclidean<>(ndim, ncells, data));
        }
        return compute_distance(ptr.get());
    }
 
    template<class Search>
    std::pair<double, double> compute_distance(const Search* search) const {
        size_t nobs = search->nobs();
        std::vector<double> dist(nobs);
 
        tatami::parallelize([&](size_t, size_t start, size_t length) -> void {
            for (size_t i = start, end = start + length; i < end; ++i) {
                auto neighbors = search->find_nearest_neighbors(i, num_neighbors);
                dist[i] = neighbors.back().second;
            }
        }, nobs, nthreads);
 
        double med = tatami::stats::compute_median<double>(dist.data(), nobs);
        double rmsd = 0;
        for (auto d : dist) {
            rmsd += d * d;
        }
        rmsd = std::sqrt(rmsd);
        return std::make_pair(med, rmsd);
    }
 
public:
    static std::vector<double> compute_scale(const std::vector<std::pair<double, double> >& distances) {
        std::vector<double> output(distances.size());
 
        // Use the first entry with a non-zero RMSD as the reference.
        bool found_ref = false;
        size_t ref = 0;
        for (size_t e = 0; e < distances.size(); ++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 (size_t e = 0; e < distances.size(); ++e) {
                output[e] = (e == ref ? 1 : compute_scale(dref, distances[e]));
            }
        }
 
        return output;
    }
 
    static double compute_scale(const std::pair<double, double>& ref, const std::pair<double, double>& target) {
        if (target.first == 0 || ref.first == 0) {
            if (target.second == 0) {
                return std::numeric_limits<double>::infinity();
            } else if (ref.second == 0) {
                return 0;
            } else {
                return ref.second / target.second; 
            }
        } else {
            return ref.first / target.first;
        }
    }
 
    template<typename Embed>
    static void combine_scaled_embeddings(const std::vector<int>& ndims, size_t ncells, const std::vector<Embed>& embeddings, const std::vector<double>& scaling, double* output) {
        size_t nembed = ndims.size();
        if (embeddings.size() != nembed || scaling.size() != nembed) {
            throw std::runtime_error("'ndims', 'embeddings' and 'scale' should have the same length");
        }
 
        int ntotal = std::accumulate(ndims.begin(), ndims.end(), 0);
        size_t offset = 0;
 
        for (size_t e = 0; e < nembed; ++e) {
            size_t curdim = ndims[e];
            auto inptr = embeddings[e];
            auto outptr = output + offset;
            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 (size_t c = 0; c < ncells; ++c, inptr += curdim, outptr += ntotal) {
                    std::fill(outptr, outptr + curdim, 0);
                }
            } else {
                for (size_t c = 0; c < ncells; ++c, inptr += curdim, outptr += ntotal) {
                    for (size_t d = 0; d < curdim; ++d) {
                        outptr[d] = inptr[d] * s;
                    }
                }
            }
 
            offset += curdim;
        }
    }
};
 
}
 
#endif