refactor: cache

Author: tristan-f-r

cache/.gitignore

@@ -0,0 +1 @@
+artifacts

cache/README.md

@@ -1,3 +1,3 @@
 # cache
 
-Handles artifact fetching and cache. 
+Handles artifact fetching and cache. This folder has a `Snakefile` which only contains a single function used for producing fetching rules.

cache/Snakefile

@@ -0,0 +1,34 @@
+from cache import link
+from cache.util import uncompress
+import urllib.parse
+from dataclasses import dataclass
+from typing import Union
+from pathlib import Path
+
+@dataclass
+class FetchConfig:
+    directive: list[str]
+    uncompress: bool = False
+
+def produce_fetch_rules(input_dict: dict[str, Union[FetchConfig, list[str]]]):
+    """
+    Produces fetch rules based on a dictionary mapping
+    output files to their directory.py-based directive.
+    """
+    # Map inputs to be wrapped with FetchConfig if list[str]
+    input_dict = {k: FetchConfig(v) if isinstance(v, list) else v for k, v in input_dict.items()}
+
+    directives = [urllib.parse.quote_plus("/".join(directive.directive)) for directive in input_dict.values()]
+    assert len(directives) == len(set(directives)), "Directives aren't unique!"
+
+    for output_file, config in input_dict.items():
+        # Since placeholders are evaluated when the job is actually ran,
+        # we pass data using params and output.
+        rule:
+            name: f"fetch_{urllib.parse.quote_plus("/".join(config.directive))}_to_{urllib.parse.quote_plus(output_file)}"
+            output: file=output_file
+            params:
+                config=config
+            run:
+                Path(output.file).parent.mkdir(exist_ok=True)
+                link(Path(output.file), params.config.directive, uncompress=params.config.uncompress)

cache/__init__.py

@@ -2,17 +2,53 @@
 This is how spras-benchmarking handles artifact caching. `cache` should be used specifically inside `Snakefile`
 """
 
+from cache.util import uncompress as uncompress_file
 from cache.directory import get_cache_item
 from pathlib import Path
 import os
 from urllib.parse import quote_plus
+import pickle
 
 __all__ = ["link"]
 
 dir_path = Path(os.path.dirname(os.path.realpath(__file__)))
+artifacts_dir = dir_path / "artifacts"
 
+def get_artifact_name(directive: list[str]) -> str:
+    return quote_plus("/".join(directive))
 
-def link(output: str, directive: list[str]):
+def has_expired(directive: list[str]) -> bool:
+    """
+    Check if the artifact metadata associated with a directive has expired.
+    Avoids re-downloading the artifact if nothing has changed.
+    """
+    artifact_name = get_artifact_name(directive)
+    cache_item = get_cache_item(directive)
+
+    metadata_dir = artifacts_dir / 'metadata'
+    metadata_dir.mkdir(exist_ok=True)
+    metadata_file = (artifacts_dir / 'metadata' / artifact_name).with_suffix((artifacts_dir / artifact_name).suffix + '.metadata')
+
+    # metadata never existed: we need to retrieve the new file
+    if not metadata_file.exists():
+        with open(metadata_file, 'wb') as f:
+            pickle.dump(cache_item, f)
+        return True
+
+    old_cache_item = None
+    with open(metadata_file, 'rb') as f:
+        old_cache_item = pickle.load(f)
+
+    # metadata expired: re-retrieve the item
+    if old_cache_item != cache_item:
+        with open(metadata_file, 'wb') as f:
+            pickle.dump(cache_item, f)
+        return True
+    
+    # metadata hasn't changed and already existed: this hasn't expired
+    return False
+
+def link(output: str, directive: list[str], uncompress=False):
     """
     Links output files from cache.directory directives.
     For example,
@@ -25,10 +61,22 @@ def link(output: str, directive: list[str]):
     (lying somewhere in the cache folder) with the desired `output`.
     """
 
-    artifacts_dir = dir_path / "artifacts"
     artifacts_dir.mkdir(exist_ok=True)
 
-    artifact_name = quote_plus("/".join(directive))
+    artifact_name = get_artifact_name(directive)
+
+    Path(output).unlink(missing_ok=True)
+
+    # Re-download if the directive has expired.
+    cache_item = get_cache_item(directive)
+    if has_expired(directive):
+        (artifacts_dir / artifact_name).unlink(missing_ok=True)
+        cache_item.download(artifacts_dir / artifact_name)
 
-    get_cache_item(directive).download(artifacts_dir / artifact_name)
-    (artifacts_dir / artifact_name).symlink_to(output)
+    if uncompress:
+        uncompressed_artifact_path = Path(str(artifacts_dir / artifact_name) + '.uncompressed')
+        uncompressed_artifact_path.unlink(missing_ok=True)
+        uncompress_file(artifacts_dir / artifact_name, uncompressed_artifact_path)
+        Path(output).symlink_to(uncompressed_artifact_path)
+    else:
+        Path(output).symlink_to(artifacts_dir / artifact_name)

cache/directory.py

@@ -66,7 +66,7 @@ def download(self, output: str | PathLike):
                 cached="https://drive.google.com/uc?id=1fvjdIbgzbgJrdJxWRRRwwS1zuegf6DOj",
                 online="http://stringdb-downloads.org/download/protein.links.v12.0/9606.protein.links.v12.0.txt.gz",
             ),
-            "9606.protein.alliases.txt.gz": CacheItem(
+            "9606.protein.aliases.txt.gz": CacheItem(
                 name="STRING 9606 protein aliases",
                 cached="https://drive.google.com/uc?id=1IWrQeTVCcw1A-jDk-4YiReWLnwP0S9bY",
                 online="https://stringdb-downloads.org/download/protein.aliases.v12.0/9606.protein.aliases.v12.0.txt.gz",

databases/README.md

@@ -1,4 +1,4 @@
-from cache import link
+include: "../../cache/Snakefile"
 
 rule all:
     input:
@@ -7,29 +7,21 @@ rule all:
         "prize_files/alopecia_areata_prizes.txt",
         "prize_files/diabetes_mellitus_prizes.txt"
 
-rule fetch:
-    output:
-        a="raw/human_disease_knowledge_filtered.tsv",
-        b="raw/human_disease_textmining_filtered.tsv",
-        c="raw/HumanDO.tsv",
-        d="raw/tiga_gene-trait_stats.tsv",
-        e="raw/ensg-ensp.tsv"
-        f="raw/9606.protein.links.v12.0.txt",
-        g="raw/9606.protein.aliases.v12.0.txt"
-    run:
-        link(output.a, ["DISEASES", "human_disease_textmining_filtered.tsv"])
-        link(output.b, ["DISEASES", "human_disease_textmining_filtered.tsv"])
-        link(output.c, ["DISEASES", "HumanDO.tsv"])
-        link(output.d, ["DISEASES", "tiga_gene-tra)it_stats.tsv"])
-        link(output.e, ["BioMart", "ensg-ensp.tsv"])
-        link(output.f, ["STRING", "9606", "9606.protein.links.v12.0.txt"])
-        link(output.g, ["STRING", "9606", "9606.protein.aliases.v12.0.txt"])
+produce_fetch_rules({
+    "raw/human_disease_textmining_filtered.tsv": ["DISEASES", "human_disease_textmining_filtered.tsv"],
+    "raw/human_disease_knowledge_filtered.tsv": ["DISEASES", "human_disease_knowledge_filtered.tsv"],
+    "raw/HumanDO.tsv": ["DISEASES", "HumanDO.tsv"],
+    "raw/tiga_gene-trait_stats.tsv": ["DISEASES", "tiga_gene-trait_stats.tsv"],
+    "raw/ensg-ensp.tsv": ["BioMart", "ensg-ensp.tsv"],
+    "raw/9606.protein.links.txt": FetchConfig(["STRING", "9606", "9606.protein.links.txt.gz"], uncompress=True),
+    "raw/9606.protein.aliases.txt": FetchConfig(["STRING", "9606", "9606.protein.aliases.txt.gz"], uncompress=True),
+})
 
 rule inputs:
     input:
         "raw/HumanDO.tsv",
         "raw/tiga_gene-trait_stats.tsv",
-        "raw/9606.protein.aliases.v12.0.txt"
+        "raw/9606.protein.aliases.txt"
     output:
         "data/inputs.csv"
     shell:
@@ -39,7 +31,8 @@ rule gold_standard:
     input:
         "raw/human_disease_knowledge_filtered.tsv",
         "raw/human_disease_textmining_filtered.tsv",
-        "raw/9606.protein.aliases.v12.0.txt"
+        "raw/9606.protein.aliases.txt",
+        "raw/ensg-ensp.tsv"
     output:
         "data/gold_standard.csv"
     shell:
@@ -49,7 +42,7 @@ rule files:
     input:
         "data/inputs.csv",
         "data/gold_standard.csv",
-        "raw/9606.protein.links.v12.0.txt"
+        "raw/9606.protein.links.txt"
     output:
         # These are the two we use for the SPRAS run for now
         "GS_files/Alopecia_areata_GS.txt",

databases/__init__.py

@@ -40,9 +40,9 @@ def main():
         df = df[["str_id"]]
         df.to_csv(diseases_path / "GS_files" / f"{disease.replace(' ', '_')}_GS.txt", sep="\t", index=False, header=None)
 
-    # See /databases/stringdb.py for information on how this was grabbed.
+    # See /cache/directory.py for information on how this was grabbed.
     # 9606 is the organism code for homo sapiens and the required background interactome of DISEASES.
-    string = pd.read_csv(diseases_path / ".." / ".." / "databases" / "string" / "9606.protein.links.v12.0.txt", sep=" ", skiprows=[0], header=None)
+    string = pd.read_csv(diseases_path / "raw" / "9606.protein.links.txt", sep=" ", skiprows=[0], header=None)
 
     # Threshold anything above a confidence score of 900 to trim down the background interactome
     string = string[string.iloc[:, 2] > 900]