f

telekom · Dec 8, 2023 · 3afdc8c · 3afdc8c
1 parent 35bc7bf
commit 3afdc8c
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 28 deletions.
diff --git a/mltb2/data.py b/mltb2/data.py
@@ -116,7 +116,7 @@ def load_colon() -> Tuple[pd.Series, pd.DataFrame]:
     return result
 
 
-def load_prostate_data() -> Tuple[pd.Series, pd.DataFrame]:
+def load_prostate() -> Tuple[pd.Series, pd.DataFrame]:
     """Load prostate data.
 
     The data is loaded and parsed from `prostate data
@@ -125,29 +125,35 @@ def load_prostate_data() -> Tuple[pd.Series, pd.DataFrame]:
     Returns:
         Tuple containing labels and data.
     """
-    # download data file
-    url = "https://web.stanford.edu/~hastie/CASI_files/DATA/prostmat.csv"
-    page = requests.get(url, timeout=10)
-    page_str = page.text
-
-    # check checksum of data file
-    page_hash = sha256(page_str.encode("utf-8")).hexdigest()
-    assert page_hash == "f1ccfd3c9a837c002ec5d6489ab139c231739c3611189be14d15ca5541b92036", page_hash
-
-    data_df = pd.read_csv(StringIO(page_str))
-    data_df = data_df.T
-
-    labels = []
-    for label in data_df.index:
-        if "control" in label:
-            labels.append(0)
-        elif "cancer" in label:
-            labels.append(1)
-        else:
-            assert False, "This must not happen!"
-
-    data_df = data_df.reset_index(drop=True)  # reset the index to default integer index
-    label_series = pd.Series(labels)
-    result = (label_series, data_df)
-
+    filename = "prostate.pkl.gz"
+    mltb2_data_home = get_and_create_mltb2_data_dir()
+    full_path = os.path.join(mltb2_data_home, filename)
+    if not os.path.exists(full_path):
+        # download data file
+        url = "https://web.stanford.edu/~hastie/CASI_files/DATA/prostmat.csv"
+        page = requests.get(url, timeout=10)
+        page_str = page.text
+
+        # check checksum of data file
+        page_hash = sha256(page_str.encode("utf-8")).hexdigest()
+        assert page_hash == "f1ccfd3c9a837c002ec5d6489ab139c231739c3611189be14d15ca5541b92036", page_hash
+
+        data_df = pd.read_csv(StringIO(page_str))
+        data_df = data_df.T
+
+        labels = []
+        for label in data_df.index:
+            if "control" in label:
+                labels.append(0)
+            elif "cancer" in label:
+                labels.append(1)
+            else:
+                assert False, "This must not happen!"
+
+        data_df = data_df.reset_index(drop=True)  # reset the index to default integer index
+        label_series = pd.Series(labels)
+        result = (label_series, data_df)
+        joblib.dump(result, full_path, compress=("gzip", 3))
+    else:
+        result = joblib.load(full_path)
     return result
diff --git a/tests/test_data.py b/tests/test_data.py
@@ -4,7 +4,7 @@
 
 import pandas as pd
 
-from mltb2.data import _load_colon_data, _load_colon_label, load_colon, load_prostate_data
+from mltb2.data import _load_colon_data, _load_colon_label, load_colon, load_prostate
 
 
 def test_load_colon_data():
@@ -33,7 +33,7 @@ def test_load_colon():
 
 
 def test_load_prostate_data():
-    result = load_prostate_data()
+    result = load_prostate()
     assert result is not None
     assert isinstance(result, tuple)
     assert len(result) == 2