GeoGrid.h

Go to the documentation of this file.

/*
 *  This is free software: you can redistribute it and/or modify it
 *  under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  any later version.
 *  The software is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *  You should have received a copy of the GNU General Public License
 *  along with the software. If not, see <http://www.gnu.org/licenses/>.
 *
 *  Copyright 2018, 2019, Jan Broeckhove and Bistromatics group.
 */

#pragma once

#include "contact/ContactType.h"
#include "geopop/geo/GeoGridKdTree.h"

#include <set>
#include <unordered_map>
#include <vector>

namespace stride {
class ContactPool;
class Population;
} // namespace stride

namespace geopop {

class Location;

template <typename Policy, typename... F>
class GeoAggregator;

class GeoGrid
{
public:
        explicit GeoGrid(stride::Population* population);

        GeoGrid(const GeoGrid&) = delete;

        GeoGrid operator=(const GeoGrid&) = delete;

        void AddLocation(std::shared_ptr<Location> location);

        void Finalize();

        std::shared_ptr<Location> GetById(unsigned int id) const { return m_locations[m_id_to_index.at(id)]; }

        stride::Population* GetPopulation() const { return m_population; }

        std::vector<const Location*> LocationsInRadius(const Location& start, double radius) const;

        std::vector<stride::ContactPool*> GetNearbyPools(stride::ContactType::Id id, const Location& start,
                                                         double startRadius = 10.0) const;

        std::set<const Location*> LocationsInBox(double long1, double lat1, double long2, double lat2) const;

        std::set<const Location*> LocationsInBox(Location* loc1, Location* loc2) const;

        std::vector<Location*> TopK(size_t k) const;

public:
        template <typename Policy, typename F>
        GeoAggregator<Policy, F> BuildAggregator(F functor, typename Policy::Args&& args) const;

        template <typename Policy>
        GeoAggregator<Policy> BuildAggregator(typename Policy::Args&& args) const;

public:
        using iterator       = std::vector<std::shared_ptr<Location>>::iterator;
        using const_iterator = std::vector<std::shared_ptr<Location>>::const_iterator;

        iterator begin() { return m_locations.begin(); }

        iterator end() { return m_locations.end(); }

        const_iterator cbegin() const { return m_locations.cbegin(); }

        const_iterator cend() const { return m_locations.cend(); }

        std::shared_ptr<Location>& operator[](size_t index) { return m_locations[index]; }

        const std::shared_ptr<Location>& operator[](size_t index) const { return m_locations[index]; }

        size_t size() const { return m_locations.size(); }

private:
        void CheckFinalized(const std::string& functionName) const;

private:
        std::vector<std::shared_ptr<Location>> m_locations;

        std::unordered_map<unsigned int, unsigned int> m_id_to_index;

        stride::Population* m_population;

        bool m_finalized;

        GeoGridKdTree m_tree;
};

} // namespace geopop