🧪 EnglandCOVID and HundaryCP tests added

Author: nithinmanoj10

stgraph/dataset/dynamic/EnglandCovidDataLoader.py

@@ -51,6 +51,8 @@ def __init__(
             The number of time lags (default is 8)
         cutoff_time : int, optional
             The cutoff timestamp for the temporal dataset (default is None)
+        redownload : bool, optional (default is False)
+            Redownload the dataset online and save to cache
 
         Attributes
         ----------

stgraph/dataset/static/CoraDataLoader.py

@@ -52,6 +52,8 @@ def __init__(self, verbose=False, url=None, redownload=False) -> None:
             Flag to control whether to display verbose info (default is False)
         url : str, optional
             The URL from where the dataset is downloaded online (default is None)
+        redownload : bool, optional (default is False)
+            Redownload the dataset online and save to cache
 
         Attributes
         ----------

stgraph/dataset/temporal/HungaryCPDataLoader.py

@@ -4,7 +4,14 @@
 
 
 class HungaryCPDataLoader(STGraphTemporalDataset):
-    def __init__(self, verbose=False, url=None, lags=4, cutoff_time=None) -> None:
+    def __init__(
+        self,
+        verbose: bool = False,
+        url: str = None,
+        lags: int = 4,
+        cutoff_time: int = None,
+        redownload: bool = False,
+    ) -> None:
         r"""County level chicken pox cases in Hungary
 
         This dataset comprises information on weekly occurrences of chickenpox
@@ -48,6 +55,8 @@ def __init__(self, verbose=False, url=None, lags=4, cutoff_time=None) -> None:
             The number of time lags (default is 4)
         cutoff_time : int, optional
             The cutoff timestamp for the temporal dataset (default is None)
+        redownload : bool, optional (default is False)
+            Redownload the dataset online and save to cache
 
         Attributes
         ----------
@@ -69,10 +78,15 @@ def __init__(self, verbose=False, url=None, lags=4, cutoff_time=None) -> None:
 
         super().__init__()
 
-        assert lags > 0, "lags should be a positive integer"
-        assert type(lags) == int, "lags should be of type int"
-        assert cutoff_time > 0, "cutoff_time should be a positive integer"
-        assert type(cutoff_time) == int, "cutoff_time should be a positive integer"
+        if type(lags) != int:
+            raise TypeError("lags must be of type int")
+        if lags < 0:
+            raise ValueError("lags must be a positive integer")
+
+        if cutoff_time != None and type(cutoff_time) != int:
+            raise TypeError("cutoff_time must be of type int")
+        if cutoff_time != None and cutoff_time < 0:
+            raise ValueError("cutoff_time must be a positive integer")
 
         self.name = "Hungary_Chickenpox"
         self._verbose = verbose
@@ -84,6 +98,9 @@ def __init__(self, verbose=False, url=None, lags=4, cutoff_time=None) -> None:
         else:
             self._url = url
 
+        if redownload and self._has_dataset_cache():
+            self._delete_cached_dataset()
+
         if self._has_dataset_cache():
             self._load_dataset()
         else:
@@ -142,9 +159,11 @@ def _set_edge_weights(self):
     def _set_targets_and_features(self):
         r"""Calculates and sets the target and feature attributes"""
         stacked_target = np.array(self._dataset["FX"])
-        self._all_targets = np.array(
-            [stacked_target[i, :].T for i in range(stacked_target.shape[0])]
-        )
+
+        self._all_targets = [
+            stacked_target[i + self._lags, :].T
+            for i in range(self.gdata["total_timestamps"] - self._lags)
+        ]
 
     def get_edges(self):
         r"""Returns the edge list"""

tests/dataset/dynamic/test_EnglandCovidDataLoader.py

@@ -5,7 +5,6 @@
 
 
 def EnglandCovidDataCheck(eng_covid: EnglandCovidDataLoader):
-    # test for gdata
     assert eng_covid.gdata["total_timestamps"] == (
         61 if not eng_covid._cutoff_time else eng_covid._cutoff_time
     )
@@ -34,16 +33,32 @@ def EnglandCovidDataCheck(eng_covid: EnglandCovidDataLoader):
     for i in range(len(edge_list)):
         assert len(edge_list[i]) == len(edge_weights[i])
 
-    # test for features and targets
-    # TODO:
+    all_features = eng_covid.get_all_features()
+    all_targets = eng_covid.get_all_targets()
+
+    assert len(all_features) == eng_covid.gdata["total_timestamps"] - eng_covid._lags
+
+    assert all_features[0].shape == (
+        eng_covid.gdata["num_nodes"]["0"],
+        eng_covid._lags,
+    )
+
+    assert len(all_targets) == eng_covid.gdata["total_timestamps"] - eng_covid._lags
+
+    assert all_targets[0].shape == (eng_covid.gdata["num_nodes"]["0"],)
 
 
 def test_EnglandCovidDataLoader():
-    eng_covid = EnglandCovidDataLoader()
+    eng_covid = EnglandCovidDataLoader(verbose=True)
     eng_covid_1 = EnglandCovidDataLoader(cutoff_time=30)
     eng_covid_2 = EnglandCovidDataLoader(
         url="https://raw.githubusercontent.com/benedekrozemberczki/pytorch_geometric_temporal/master/dataset/england_covid.json"
     )
+    eng_covid_3 = EnglandCovidDataLoader(lags=12)
+    # eng_covid_4 = EnglandCovidDataLoader(redownload=True)
 
     EnglandCovidDataCheck(eng_covid)
     EnglandCovidDataCheck(eng_covid_1)
+    # EnglandCovidDataCheck(eng_covid_2)
+    EnglandCovidDataCheck(eng_covid_3)
+    # EnglandCovidDataCheck(eng_covid_4)

tests/dataset/static/test_CoraDataLoader.py

@@ -1,6 +1,3 @@
-import numpy as np
-import urllib.request
-
 from stgraph.dataset import CoraDataLoader
 
 

tests/dataset/temporal/test_HungaryCPDataLoader.py

@@ -0,0 +1,54 @@
+import pytest
+
+from stgraph.dataset import HungaryCPDataLoader
+
+
+def HungaryCPDataChecker(hungary: HungaryCPDataLoader):
+    assert hungary.gdata["total_timestamps"] == (
+        521 if not hungary._cutoff_time else hungary._cutoff_time
+    )
+
+    assert hungary.gdata["num_nodes"] == 20
+    assert hungary.gdata["num_edges"] == 102
+
+    edges = hungary.get_edges()
+    edge_weights = hungary.get_edge_weights()
+    all_targets = hungary.get_all_targets()
+
+    assert len(edges) == 102
+    assert len(edges[0]) == 2
+
+    assert len(edge_weights) == 102
+
+    assert len(all_targets) == hungary.gdata["total_timestamps"] - hungary._lags
+    assert all_targets[0].shape == (hungary.gdata["num_nodes"],)
+
+
+def test_HungaryCPDataLoader():
+    hungary_1 = HungaryCPDataLoader(verbose=True)
+    hungary_2 = HungaryCPDataLoader(lags=6)
+    hungary_3 = HungaryCPDataLoader(cutoff_time=100)
+    hungary_4 = HungaryCPDataLoader(
+        url="https://raw.githubusercontent.com/bfGraph/STGraph-Datasets/main/HungaryCP.json"
+    )
+
+    HungaryCPDataChecker(hungary_1)
+    HungaryCPDataChecker(hungary_2)
+    HungaryCPDataChecker(hungary_3)
+    # HungaryCPDataChecker(hungary_4)
+
+    with pytest.raises(TypeError) as exec:
+        HungaryCPDataLoader(lags="lags")
+    assert str(exec.value) == "lags must be of type int"
+
+    with pytest.raises(ValueError) as exec:
+        HungaryCPDataLoader(lags=-1)
+    assert str(exec.value) == "lags must be a positive integer"
+
+    with pytest.raises(TypeError) as exec:
+        HungaryCPDataLoader(cutoff_time="time")
+    assert str(exec.value) == "cutoff_time must be of type int"
+
+    with pytest.raises(ValueError) as exec:
+        HungaryCPDataLoader(cutoff_time=-1)
+    assert str(exec.value) == "cutoff_time must be a positive integer"