Skip to content
Snippets Groups Projects
GeometryHelpers.cpp 5.18 KiB
Newer Older
  • Learn to ignore specific revisions
  • #include "GeometryHelpers.h"
    
    Chao Peng's avatar
    Chao Peng committed
    
    // some utility functions that can be shared
    
    namespace athena::geo {
    
    Chao Peng's avatar
    Chao Peng committed
    
    typedef ROOT::Math::XYPoint Point;
    
    
    // check if a 2d point is already in the container
    bool already_placed(const Point &p, const std::vector<Point> &vec, double xs = 1.0, double ys = 1.0, double tol = 1e-6)
    {
        for (auto &pt : vec) {
            if ((std::abs(pt.x() - p.x())/xs < tol) && std::abs(pt.y() - p.y())/ys < tol) {
                return true;
            }
        }
        return false;
    }
    
    
    Chao Peng's avatar
    Chao Peng committed
    // check if a square in a ring
    
    inline bool rec_in_ring(const Point &pt, double sx, double sy, double rmin, double rmax, double phmin, double phmax)
    
    Chao Peng's avatar
    Chao Peng committed
    {
        if (pt.r() > rmax || pt.r() < rmin) {
            return false;
        }
    
        // check four corners
        std::vector<Point> pts {
    
            Point(pt.x() - sx/2., pt.y() - sy/2.),
            Point(pt.x() - sx/2., pt.y() + sy/2.),
            Point(pt.x() + sx/2., pt.y() - sy/2.),
            Point(pt.x() + sx/2., pt.y() + sy/2.),
    
    Chao Peng's avatar
    Chao Peng committed
        };
        for (auto &p : pts) {
            if (p.r() > rmax || p.r() < rmin || p.phi() > phmax || p.phi() < phmin) {
                return false;
            }
        }
        return true;
    }
    
    // a helper function to recursively fill square in a ring
    
    void add_rectangle(Point p, std::vector<Point> &res, double sx, double sy,
                       double rmin, double rmax, double phmin, double phmax, int max_depth = 20, int depth = 0)
    
        // std::cout << depth << "/" << max_depth << std::endl;
        // exceeds the maximum depth in searching or already placed
        if ((depth > max_depth) || (already_placed(p, res, sx, sy))) {
    
        bool in_ring = rec_in_ring(p, sx, sy, rmin, rmax, phmin, phmax);
        if (in_ring) {
            res.emplace_back(p);
        }
    
        // continue search for a good placement or if no placement found yet
        if (in_ring || res.empty()) {
            // check adjacent squares
            add_rectangle(Point(p.x() + sx, p.y()), res, sx, sy, rmin, rmax, phmin, phmax, max_depth, depth + 1);
            add_rectangle(Point(p.x() - sx, p.y()), res, sx, sy, rmin, rmax, phmin, phmax, max_depth, depth + 1);
            add_rectangle(Point(p.x(), p.y() + sy), res, sx, sy, rmin, rmax, phmin, phmax, max_depth, depth + 1);
            add_rectangle(Point(p.x(), p.y() - sy), res, sx, sy, rmin, rmax, phmin, phmax, max_depth, depth + 1);
        }
    
    Chao Peng's avatar
    Chao Peng committed
    }
    
    // fill squares
    
    std::vector<Point> fillRectangles(Point ref, double sx, double sy, double rmin, double rmax, double phmin, double phmax)
    
        // convert (0, 2pi) to (-pi, pi)
        if (phmax > M_PI) {
            phmin -= M_PI;
            phmax -= M_PI;
        }
    
    Chao Peng's avatar
    Chao Peng committed
        // start with a seed square and find one in the ring
        // move to center
    
        ref = ref - Point(int(ref.x()/sx)*sx, int(ref.y()/sy)*sy);
    
        std::vector<Point> res;
        add_rectangle(ref, res, sx, sy, rmin, rmax, phmin, phmax, (int(rmax/sx) + 1)*(int(rmax/sy) + 1)*2);
        return res;
    }
    
    // check if a regular polygon is inside a ring
    bool poly_in_ring(const Point &p, int nsides, double lside, double rmin, double rmax, double phmin, double phmax)
    {
        // outer radius is contained
        if ((p.r() + lside <= rmax) && (p.r() - lside >= rmin)) {
            return true;
        }
    
        // inner radius is not contained
        double rin = std::cos(M_PI/nsides)*lside;
        if ((p.r() + rin > rmax) || (p.r() - rin < rmin)) {
            return false;
        }
    
        // in between, check every corner
        for (int i = 0; i < nsides; ++i) {
            double phi = (i + 0.5)*2.*M_PI/static_cast<double>(nsides);
            Point p2(p.x() + 2.*lside*std::sin(phi), p.y() + 2.*lside*std::cos(phi));
            if ((p2.r() > rmax) || (p2.r() < rmin)) {
                return false;
    
        }
        return true;
    }
    
    // recursively fill square (nside=4) or hexagon (nside=6) in a ring, other polygons won't work
    void add_poly(Point p, std::vector<Point> &res, int nsides, double lside,
                  double rmin, double rmax, double phmin, double phmax, int max_depth = 20, int depth = 0)
    {
        // std::cout << depth << "/" << max_depth << std::endl;
        // exceeds the maximum depth in searching or already placed
        if ((depth > max_depth) || (already_placed(p, res, lside, lside))) {
            return;
        }
    
        bool in_ring = poly_in_ring(p, nsides, lside, rmin, rmax, phmin, phmax);
        if (in_ring) {
            res.emplace_back(p);
        }
    
        // recursively add neigbors, continue if it was a good placement or no placement found yet
        if (in_ring || res.empty()) {
            double d = 2.*std::cos(M_PI/static_cast<double>(nsides))*lside;
            for (int i = 0; i < nsides; ++i) {
                double phi = i*2.*M_PI/static_cast<double>(nsides);
                add_poly(Point(p.x() + 2.*lside*std::sin(phi), p.y() + 2.*lside*std::cos(phi)), res, nsides, lside,
                         rmin, rmax, phmin, phmax, max_depth, depth + 1);
            }
        }
    }
    
    std::vector<Point> fillHexagons(Point ref, double lside, double rmin, double rmax, double phmin, double phmax)
    {
        // convert (0, 2pi) to (-pi, pi)
        if (phmax > M_PI) {
            phmin -= M_PI;
            phmax -= M_PI;
        }
        // start with a seed and find one in the ring
        // move to center
        ref = ref - Point(int(ref.x()/lside)*lside, int(ref.y()/lside)*lside);
    
    Chao Peng's avatar
    Chao Peng committed
    
        std::vector<Point> res;
    
        add_poly(ref, res, 6, lside, rmin, rmax, phmin, phmax, std::pow(int(rmax/lside) + 1, 2)*2);
    
    Chao Peng's avatar
    Chao Peng committed
        return res;
    }
    
    
    } // athena::geo