Generate the Helium network graph for a city.

helium_network_graphs.py

from get_data import *
import math
import networkx as nx
import h3
import json
from torch_geometric.utils import from_networkx
from torch_geometric.data import Data
import pickle


class CityGraph:
    def __init__(self, city_id):
        self.city_id = city_id
        self.g = None
        self.hotspot_list = list_hotspots_in_city(city_id)
        self.positions = None

    def generate_graph(self, calculate_rewards=False):
        """Uses list of hotspots and witness data to create nodes and adjacency with networkx"""
        g = nx.Graph()
        positions = {}

        # create nodes for all hotspots in list
        for hotspot in self.hotspot_list:
            if calculate_rewards:
                hotspot['rewards'] = get_hotspot_rewards(hotspot['address'], 5)
            else:
                hotspot['rewards'] = 0
            coords = (hotspot['lng'], hotspot['lat'])
            g.add_node(hotspot['address'], pos=coords, lat=hotspot['lat'], lng=hotspot['lng'],
                       location_hex=hotspot['location_hex'], name=hotspot['name'], rewards=hotspot['rewards'],
                       gain=hotspot['gain'], elevation=hotspot['elevation'])
            positions[hotspot['address']] = coords

        # create edges representing hotspots that witness each other
        i = 0
        for hotspot in self.hotspot_list:
            witness_list = list_witnesses_for_hotspot(hotspot['address'])
            hotspot['num_witnesses'] = len(witness_list)
            g.nodes[hotspot['address']]['num_witnesses'] = hotspot['num_witnesses']
            for witness in witness_list:

                # make sure that all witnesses are also in this subset
                if not any(witness['address'] == h['address'] for h in self.hotspot_list):
                    continue

                # calculate distance between witness and challengee
                dist = math.sqrt((hotspot['lat'] - witness['lat']) ** 2 + (hotspot['lng'] - witness['lng']) ** 2)
                g.add_edge(hotspot['address'], witness['address'], weight=dist)

            i += 1
            if i % 10 == 0:
                print(f"{str(i)} out of {str(len(self.hotspot_list))} hotspots complete...")

        self.g = g
        self.positions = positions
        return g

    def pagerank(self):
        if not self.g:
            raise NameError('You must generate a graph first with CityGraph.generate_graph()')
        return nx.pagerank(self.g)

    def betweenness_centrality(self):
        if not self.g:
            raise NameError('You must generate a graph first with CityGraph.generate_graph()')
        return nx.betweenness_centrality(self.g)