transform_to_graphml.py

import os.path

import networkx as nx
import pandas as pd

if os.path.exists('data/air_traffic.graphml'):
    print('data/air_traffic.graphml file already exists.')
    exit(0)

df_airports = pd.read_csv('data/airports.csv')
df_flights = pd.read_csv('data/anac.csv')

# Create graph
G = nx.Graph()

# Add nodes
for index, row in df_airports.iterrows(): # noqa: B007
    G.add_node(row['code'],
               name=row['name'],
               country=row['country'],
               latitude=row['lat_geo_point'],
               longitude=row['lon_geo_point']
               )

# Add edges
df_edges = df_flights[[
    'origin_airport_abbreviation',
    'destination_airport_abbreviation',
]].dropna()
df_edges = df_edges.groupby(df_edges.columns.tolist(), as_index=False).size()
for index, row in df_edges.iterrows(): # noqa: B007
    if row['origin_airport_abbreviation'] == row['destination_airport_abbreviation']:
        continue
    G.add_edge(
        row['origin_airport_abbreviation'],
        row['destination_airport_abbreviation'],
        flight_count=row['size']
    )

# Export to graphml
nx.write_graphml(G, 'data/air_traffic.graphml')