diff --git a/topobenchmarkx/io/load/us_county_demos.py b/topobenchmarkx/io/load/us_county_demos.py index ae026d2a..7c267393 100644 --- a/topobenchmarkx/io/load/us_county_demos.py +++ b/topobenchmarkx/io/load/us_county_demos.py @@ -5,6 +5,23 @@ def load_us_county_demos(path, year=2012, y_col="Election"): + r"""Load US County Demos dataset + + Parameters + ---------- + path: str + Path to the dataset. + year: int + Year to load the features. + y_col: str + Column to use as label. + + Returns + ------- + torch_geometric.data.Data + Data object of the graph for the US County Demos dataset. + """ + edges_df = pd.read_csv(f"{path}/county_graph.csv") stat = pd.read_csv(f"{path}/county_stats_{year}.csv", encoding="ISO-8859-1")