Skip to content

Commit 1c10a8a

Browse files
committed
add: game of thrones export script, csv and remove from functions
1 parent 01819fb commit 1c10a8a

File tree

5 files changed

+1358
-815
lines changed

5 files changed

+1358
-815
lines changed
Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
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

Comments
 (0)