1
+ import pandas as pd
2
+ import requests
3
+ import io
4
+ import os
5
+ import igraph as ig
6
+
7
+ def load_got_network (local_path : str = './game_of_thrones/got-s1-edges.csv' , url : str = None ):
8
+ """
9
+ Loads the Game of Thrones network dataset from a local CSV file. If the file is not found,
10
+ it downloads the data from a public GitHub repository. The function can also export
11
+ the resulting graph to a GraphML file.
12
+
13
+ Args:
14
+ local_path (str, optional): The path to the local CSV file.
15
+ Defaults to './game_of_thrones/got-s1-edges.csv'.
16
+ url (str, optional): The URL of the raw CSV data. This is used as a fallback.
17
+ Defaults to "https://raw.githubusercontent.com/mathbeveridge/gameofthrones/master/data/got-s1-edges.csv".
18
+
19
+ Returns:
20
+ igraph.Graph: An igraph Graph object representing the network.
21
+ """
22
+ if url is None :
23
+ url = "https://raw.githubusercontent.com/mathbeveridge/gameofthrones/master/data/got-s1-edges.csv"
24
+
25
+ if os .path .exists (local_path ):
26
+ print (f"Loading data from local file: { local_path } " )
27
+ df_s1_edges = pd .read_csv (local_path )
28
+ else :
29
+ print (f"Local file not found at { local_path } . Attempting to download from GitHub..." )
30
+ try :
31
+ response = requests .get (url )
32
+ response .raise_for_status ()
33
+ season1_edges_data = io .StringIO (response .text )
34
+ df_s1_edges = pd .read_csv (season1_edges_data )
35
+ except requests .exceptions .RequestException as e :
36
+ print (f"Error fetching data from { url } : { e } " )
37
+ print ("Please ensure you have an active internet connection or download the file manually." )
38
+ return None
39
+
40
+ if 'Weight' in df_s1_edges .columns :
41
+ edges_for_tuplelist = df_s1_edges [['Source' , 'Target' , 'Weight' ]].values .tolist ()
42
+ g_s1 = ig .Graph .TupleList (edges_for_tuplelist , directed = False , weights = True )
43
+ print ("Graph created with edge weights." )
44
+ else :
45
+ g_s1 = ig .Graph .TupleList (df_s1_edges [['Source' , 'Target' ]].itertuples (index = False ), directed = False )
46
+ print ("Graph created without edge weights." )
47
+
48
+ return g_s1
49
+
50
+
51
+ def export_graph_to_graphml (graph : ig .Graph , export_path : str ):
52
+ """
53
+ Exports the given igraph Graph object to a GraphML file.
54
+
55
+ Args:
56
+ graph (igraph.Graph): The igraph Graph object to export.
57
+ export_path (str): The path where the GraphML file should be saved.
58
+ """
59
+ if export_path :
60
+ os .makedirs (os .path .dirname (export_path ), exist_ok = True )
61
+ graph .write_graphml (export_path )
62
+ print (f"Graph exported to { export_path } " )
63
+ else :
64
+ print ("No export path provided. Graph not exported." )
65
+
66
+ if __name__ == "__main__" :
67
+ local_path = './game_of_thrones/got-s1-edges.csv'
68
+ export_path = './game_of_thrones/GoT.graphml'
69
+ script_dir = os .path .dirname (os .path .abspath (__file__ ))
70
+
71
+ full_local_path = os .path .join (script_dir , local_path )
72
+ full_export_path = os .path .join (script_dir , export_path )
73
+
74
+ g_s1 = load_got_network (local_path = full_local_path )
75
+ if g_s1 :
76
+ export_graph_to_graphml (g_s1 , full_export_path )
77
+ else :
78
+ print ("Failed to load the Game of Thrones network." )
79
+
0 commit comments