added first dataset for testing

Author: Ashutosh Tiwari

datasets/dataset.py

@@ -1,24 +1,31 @@
 from typing import List, Tuple
+
+import pandas as pd
+import scipy
 from torch_geometric.data import download_url
 
 
 class Dataset(object):
     def __init__(self, root: str = '/tmp/', group_col: str = None, urls: List = None) -> None:
         self.root = root
         self.group_col = group_col
-        self.data = None
-        self.adj = None
         self.file_paths = []
         urls = urls or []
         for url in urls:
             self.file_paths.append(download_url(url, root))
-        self._set_data_df()
-        self._set_adj()
+        self.data = self._set_data_df()
+        self.adj = self._set_adj()
 
     @property
     def X(self):
         return self.data.loc[:, self.data.columns != self.group_col].values
 
     @property
     def y(self):
-        return self.data[:, self.group_col].values
+        return self.data[:, self.group_col].values
+
+    def _set_data_df(self) -> pd.DataFrame:
+        raise NotImplementedError
+
+    def _set_adj(self) -> scipy.sparse:
+        raise NotImplementedError

datasets/polbooks.py

@@ -1,5 +1,6 @@
 import pandas as pd
 import networkx as nx
+from torch_geometric.data import extract_zip
 
 
 from datasets import dataset
@@ -12,7 +13,16 @@ def __init__(self, root: str = "/tmp/"):
                          urls=["https://websites.umich.edu/~mejn/netdata/polbooks.zip"])
 
     def _set_data_df(self):
-        pass
+        extract_zip(self.file_paths[0], self.root)
+        self.graph = nx.read_gml(self.root + "polbooks.gml")
+        node_dict = dict(self.graph.nodes(data=True))
+        df = pd.DataFrame.from_dict(node_dict, orient='index').reset_index(drop=True). \
+            rename(columns={'value': 'political_leaning'})
+        df['political_leaning'] = df['political_leaning'].map(
+            {'n': 0, 'c': 1, 'l': 2}
+        )
+        return df
 
     def _set_adj(self):
+        return nx.adjacency_matrix(self.graph)
         pass