distribution.h

Go to the documentation of this file.

/* Author: Zachary Kingston */

#ifndef SE2EZ_CORE_KDE_
#define SE2EZ_CORE_KDE_

#include <ompl/datastructures/PDF.h>
#include <ompl/datastructures/Grid.h>
#include <ompl/datastructures/NearestNeighborsGNAT.h>

#include <se2ez/core/log.h>
#include <se2ez/core/math.h>
#include <se2ez/core/class_forward.h>

namespace se2ez
{
    SE2EZ_CLASS_FORWARD(Distribution)
    class Distribution
    {
    public:
        Distribution(unsigned int dimension) : dimension_(dimension)
        {
        }

        virtual ~Distribution() = default;

        virtual double pdf(const Eigen::Ref<const Eigen::VectorXd> &x) = 0;
        virtual void sample(Eigen::Ref<Eigen::VectorXd> sample) = 0;

        void addPoint(const Eigen::Ref<const Eigen::VectorXd> &x)
        {
            if (x.size() != dimension_)
                throw std::runtime_error(log::format("Invalid dimension %1% for data point (should be %2%)",
                                                     x.size(), dimension_));

            auto entry = new Eigen::VectorXd(dimension_);
            *entry = x;

            data_.emplace_back(entry);
            addedPoint(entry);
        }

        void addPoints(const tf::EigenVector<Eigen::VectorXd> &xs)
        {
            data_.reserve(data_.size() + xs.size());
            for (const auto &x : xs)
                addPoint(x);
        }

    protected:
        virtual void addedPoint(const Eigen::VectorXd * /*x*/)
        {
        }

        unsigned int dimension_;
        std::vector<const Eigen::VectorXd *> data_;
    };

    class Histogram : public Distribution
    {
    public:
        Histogram(unsigned int dimension,  //
                  const Eigen::Ref<const Eigen::VectorXd> &min, const Eigen::Ref<const Eigen::VectorXd> &max,
                  unsigned int cells)
          : Distribution(dimension)
          , min_(min)
          , max_(max)
          , width_((max_ - min_) / (double)cells)
          , cells_(cells)
          , histogram_(dimension_)
        {
            if (min.size() != dimension_)
                throw std::runtime_error(
                    log::format("Invalid dimension %1% for minimum (should be %2%)", min.size(), dimension_));

            if (max.size() != dimension_)
                throw std::runtime_error(
                    log::format("Invalid dimension %1% for maximum (should be %2%)", max.size(), dimension_));

            if (cells == 0)
                throw std::runtime_error(log::format("Number of cells must be positive"));

            for (unsigned int i = 0; i < dimension_; ++i)
                if (max_[i] < min_[i])
                    throw std::runtime_error(log::format("On index %1%: minimum %2% is greater than max %3%!",
                                                         i, min_[i], max_[i]));
        }

        ~Histogram()
        {
            Bins::CellArray cells;
            histogram_.getCells(cells);
            for (const auto &cell : cells)
                delete cell->data;
        }

        double pdf(const Eigen::Ref<const Eigen::VectorXd> &x) override
        {
            if (not data_.size())
                throw std::runtime_error(log::format("Cannot compute PDF without data."));

            if (x.size() != dimension_)
                throw std::runtime_error(log::format("Invalid dimension %1% for data point (should be %2%)",
                                                     x.size(), dimension_));
            Bin *bin = getCell(x);
            return bin->data->points.size() / data_.size();
        }

        void sample(Eigen::Ref<Eigen::VectorXd> sample) override
        {
            if (sample.size() != dimension_)
                throw std::runtime_error(log::format("Invalid dimension %1% for sample (should be %2%)",
                                                     sample.size(), dimension_));

            auto data = pdf_.sample(math::uniformReal(0, 1.));
            const Index &index = data->bin->coord;

            for (unsigned int i = 0; i < dimension_; ++i)
            {
                double l = min_[i] + index[i] * width_[i];
                double u = l + width_[i];

                sample[i] = math::uniformReal(l, u);
            }
        }

        void toString(std::ostream &out)
        {
            if (dimension_ > 2)
                throw std::runtime_error("Dimension to high to print!");

            if (dimension_ == 1)
            {
                unsigned int size = 30;
                for (unsigned int i = 0; i < cells_; ++i)
                {
                    Index index(dimension_);
                    index[0] = i;

                    Bin *bin = histogram_.getCell(index);
                    unsigned int v = (bin) ? bin->data->points.size() : 0;

                    unsigned int ch = math::remap(0, maxCell_, v, 0, size);
                    unsigned int j = 0;
                    for (; j < ch; ++j)
                        out << "@";
                    for (; j < size; ++j)
                        out << " ";

                    out << " : " << (double)v / data_.size();
                    out << std::endl;
                }
            }

            if (dimension_ == 2)
            {
                const char ramp[11] = " .:-=+*#%@";
                for (unsigned int i = 0; i < cells_; ++i)
                {
                    for (unsigned int j = 0; j < cells_; ++j)
                    {
                        Index index(dimension_);
                        index[0] = i;
                        index[1] = j;

                        Bin *bin = histogram_.getCell(index);
                        unsigned int v = (bin) ? bin->data->points.size() : 0;

                        unsigned int ch = math::remap(0, maxCell_, v, 0, 10);
                        out << ramp[ch];
                    }

                    out << std::endl;
                }
            }
        }

    protected:
        void addedPoint(const Eigen::VectorXd *x) override
        {
            Bin *bin = getCell(*x);
            bin->data->points.emplace_back(x);

            unsigned int points = bin->data->points.size();
            pdf_.update(bin->data->element, points);

            maxCell_ = (maxCell_ < points) ? points : maxCell_;
        }

    private:
        struct Data;

        using Bins = ompl::Grid<Data *>;
        using Bin = Bins::Cell;
        using Index = Bins::Coord;
        using PDF = ompl::PDF<Data *>;

        struct Data
        {
            Bin *bin;
            PDF::Element *element;
            std::vector<const Eigen::VectorXd *> points;
        };

        Index getIndex(const Eigen::Ref<const Eigen::VectorXd> &x)
        {
            Index index(dimension_);
            for (unsigned int i = 0; i < dimension_; ++i)
                index[i] = (x[i] - min_[i]) / width_[i];

            return index;
        }

        Bin *getCell(const Eigen::Ref<const Eigen::VectorXd> &x)
        {
            Index index = getIndex(x);
            Bin *bin = histogram_.getCell(index);

            if (not bin)
            {
                bin = histogram_.createCell(index);
                bin->data = new Data;
                bin->data->bin = bin;
                bin->data->element = pdf_.add(bin->data, 0.);

                histogram_.add(bin);
            }

            return bin;
        }

        const Eigen::VectorXd min_;
        const Eigen::VectorXd max_;
        const Eigen::VectorXd width_;
        const unsigned int cells_;

        unsigned int maxCell_{0};

        Bins histogram_;
        PDF pdf_;
    };  // namespace se2ez

    class KDE : public Distribution
    {
    public:
        KDE(unsigned int dimension)
          : Distribution(dimension)
          , min_(Eigen::VectorXd::Zero(dimension))
          , max_(Eigen::VectorXd::Zero(dimension))
          , mean_(Eigen::VectorXd::Zero(dimension))
          , sqMean_(Eigen::VectorXd::Zero(dimension))
          , H_(Eigen::VectorXd::Zero(dimension))
        {
            neighbors_.setDistanceFunction([](const auto &a, const auto &b) { return (*a - *b).norm(); });
        }

        double pdf(const Eigen::Ref<const Eigen::VectorXd> &x) override
        {
            if (not data_.size())
                throw std::runtime_error(log::format("Cannot compute PDF without data."));

            if (x.size() != dimension_)
                throw std::runtime_error(log::format("Invalid dimension %1% for data point (should be %2%)",
                                                     x.size(), dimension_));

            bandwidth();

            std::vector<const Eigen::VectorXd *> entries;
            entries.reserve(NEIGHBORS);

            const Eigen::VectorXd &v = x;
            neighbors_.nearestK(&v, NEIGHBORS, entries);

            // std::vector<const Eigen::VectorXd *> entries;
            // entries.reserve(NEIGHBORS);

            // const Eigen::VectorXd &v = x;
            // neighbors_.nearestR(&v, H_.maxCoeff() * 10, entries);

            double pdf = 0;
            for (const auto &entry : entries)
            {
                double a = 1.0;
                for (unsigned int i = 0; i < dimension_; ++i)
                    a *= kernel(x[i], (*entry)[i], H_[i]) / H_[i];

                pdf += a;
            }

            if (entries.size())
                return pdf / entries.size();
            return 0.;
        }

        void sample(Eigen::Ref<Eigen::VectorXd> sample) override
        {
            if (sample.size() != dimension_)
                throw std::runtime_error(log::format("Invalid dimension %1% for sample (should be %2%)",
                                                     sample.size(), dimension_));

            bandwidth();
            sample = *data_[math::uniformInteger(0, data_.size() - 1)];
            for (unsigned int i = 0; i < dimension_; ++i)
                sample[i] = generate(sample[i], H_[i]);
        }

        void toString(std::ostream &out, const Eigen::Ref<const Eigen::VectorXd> &min,
                      const Eigen::Ref<const Eigen::VectorXd> &max, unsigned int resolution)
        {
            if (dimension_ > 2)
                throw std::runtime_error("Dimension to high to print!");

            if (dimension_ == 1)
            {
                unsigned int size = 30;
                double width = (max[0] - min[0]) / resolution;
                double map[resolution];
                double m = 0;

                for (unsigned int i = 0; i < resolution; ++i)
                {
                    Eigen::VectorXd s(1);
                    s[0] = min[0] + i * width;

                    map[i] = pdf(s);
                    m = (m < map[i]) ? map[i] : m;
                }

                for (unsigned int i = 0; i < resolution; ++i)
                {
                    unsigned int ch = math::remap(0, m, map[i], 0, size);
                    unsigned int j = 0;
                    for (; j < ch; ++j)
                        out << "@";
                    for (; j < size; ++j)
                        out << " ";

                    out << " : " << map[i];

                    out << std::endl;
                }
            }

            if (dimension_ == 2)
            {
                const char ramp[11] = " .:-=+*#%@";

                Eigen::VectorXd width = (max - min) / resolution;
                double map[resolution][resolution];
                double m = 0;

                for (unsigned int i = 0; i < resolution; ++i)
                {
                    Eigen::VectorXd s(2);
                    s[0] = min[0] + i * width[0];
                    for (unsigned int j = 0; j < resolution; ++j)
                    {
                        s[1] = min[1] + j * width[1];

                        map[i][j] = pdf(s);
                        m = (m < map[i][j]) ? map[i][j] : m;
                    }
                }

                for (unsigned int i = 0; i < resolution; ++i)
                {
                    for (unsigned int j = 0; j < resolution; ++j)
                    {
                        unsigned int ch = math::remap(0, m, map[i][j], 0, 10);
                        ch = (ch == 10) ? 9 : ch;

                        out << ramp[ch];
                    }

                    out << std::endl;
                }
            }
        }

    protected:
        void addedPoint(const Eigen::VectorXd *x) override
        {
            neighbors_.add(x);

            mean_ += *x;
            sqMean_ += (*x).cwiseProduct(*x);

            dirty_ = true;
        }

        virtual double kernel(double x, double m, double s) = 0;
        virtual double generate(double m, double s) = 0;

    private:
        void bandwidth()
        {
            if (not dirty_)
                return;

            unsigned int num = data_.size();
            for (unsigned int i = 0; i < dimension_; ++i)
            {
                double mean = mean_[i] / num;
                double sqMean = sqMean_[i] / num;

                double sigma = std::sqrt(sqMean - mean * mean);
                H_[i] = std::pow(4. / (dimension_ + 2), 1. / (dimension_ + 4))  //
                        * std::pow(num, -1. / (dimension_ + 4))                 //
                        * sigma;
            }

            dirty_ = false;
        }

        Eigen::VectorXd min_;
        Eigen::VectorXd max_;

        Eigen::VectorXd mean_;
        Eigen::VectorXd sqMean_;

        bool dirty_{true};
        Eigen::VectorXd H_;
        ompl::NearestNeighborsGNAT<const Eigen::VectorXd *> neighbors_;

        const unsigned int NEIGHBORS = 1000;
    };

    class GaussianKDE : public KDE
    {
    public:
        GaussianKDE(unsigned int dimension) : KDE(dimension)
        {
        }

    protected:
        double kernel(double x, double m, double s) override
        {
            double z = (x - m) / s;
            // return std::exp(-0.5 * z * z) / (s * std::sqrt(math::pipi));
            // return std::exp(-1. * z * z / (2. * s * s));
            return std::exp(-1. * z * z / 2.);
        }

        double generate(double m, double s) override
        {
            return math::gaussianReal(m, s);
        }
    };
}  // namespace se2ez

#endif