refactor: move db helpers and fix version plots

Author: oesteban

src/db.py

@@ -0,0 +1,58 @@
+# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
+"""Database helpers used by the CLI and plotting utilities."""
+
+from __future__ import annotations
+
+import datetime
+from typing import Tuple
+
+import pandas as pd
+from pymongo import MongoClient
+
+
+def load_event(event_name: str, unique: bool = True) -> pd.DataFrame:
+    """Load one event collection from MongoDB."""
+    db = MongoClient().fmriprep_stats
+    data = pd.DataFrame(list(db[event_name].find()))
+    if len(data) == 0:
+        raise RuntimeError(f"No records of event '{event_name}'")
+
+    data["dateCreated"] = pd.to_datetime(data["dateCreated"])
+    data["date_minus_time"] = data["dateCreated"].apply(
+        lambda df: datetime.datetime(year=df.year, month=df.month, day=df.day)
+    )
+    if unique:
+        data = data.drop_duplicates(subset=["run_uuid"])
+    return data
+
+
+def massage_versions(
+    started: pd.DataFrame, success: pd.DataFrame
+) -> Tuple[pd.DataFrame, pd.DataFrame]:
+    """Normalize version strings as done in the analysis notebook."""
+    started = started.copy()
+    success = success.copy()
+
+    started = started.fillna(value={"environment_version": "older"})
+    success = success.fillna(value={"environment_version": "older"})
+
+    started.loc[started.environment_version == "v0.0.1", "environment_version"] = "older"
+    success.loc[success.environment_version == "v0.0.1", "environment_version"] = "older"
+
+    started.loc[started.environment_version.str.startswith("20.0"), "environment_version"] = "older"
+    success.loc[success.environment_version.str.startswith("20.0"), "environment_version"] = "older"
+    started.loc[started.environment_version.str.startswith("20.1"), "environment_version"] = "older"
+    success.loc[success.environment_version.str.startswith("20.1"), "environment_version"] = "older"
+
+    versions = sorted(
+        {
+            ".".join(v.split(".")[:2])
+            for v in started.environment_version.unique()
+            if "." in str(v)
+        }
+    )
+    for ver in versions:
+        started.loc[started.environment_version.str.startswith(ver), "environment_version"] = ver
+        success.loc[success.environment_version.str.startswith(ver), "environment_version"] = ver
+
+    return started, success

src/run.py

@@ -28,7 +28,8 @@
 
 import click
 from api import parallel_fetch, ISSUES, DEFAULT_MAX_ERRORS
-from viz import load_event, plot_performance, plot_version_stream
+from db import load_event, massage_versions
+from viz import plot_performance, plot_version_stream
 
 DEFAULT_DAYS_WINDOW = 90
 DEFAULT_CHUNK_DAYS = 1
@@ -163,7 +164,9 @@ def plot(output_dir, drop_cutoff):
     unique_success = load_event("success")
 
     plot_performance(unique_started, unique_success, drop_cutoff=drop_cutoff, out_file=out_perf)
-    plot_version_stream(unique_started, unique_success, drop_cutoff=drop_cutoff, out_file=out_ver)
+
+    started_v, success_v = massage_versions(unique_started, unique_success)
+    plot_version_stream(started_v, success_v, drop_cutoff=drop_cutoff, out_file=out_ver)
     click.echo(f"Saved plots to {output_dir}")
 
 

src/viz.py

@@ -7,7 +7,6 @@
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
-from pymongo import MongoClient
 from scipy.interpolate import RBFInterpolator
 
 
@@ -27,22 +26,6 @@ def _parse(vstr):
 _vparse = np.vectorize(_parse)
 
 
-def load_event(event_name: str, unique: bool = True) -> pd.DataFrame:
-    """Load one event collection from MongoDB."""
-    db = MongoClient().fmriprep_stats
-    data = pd.DataFrame(list(db[event_name].find()))
-    if len(data) == 0:
-        raise RuntimeError(f"No records of event '{event_name}'")
-
-    data["dateCreated"] = pd.to_datetime(data["dateCreated"])
-    data["date_minus_time"] = data["dateCreated"].apply(
-        lambda df: datetime.datetime(year=df.year, month=df.month, day=df.day)
-    )
-    if unique:
-        data = data.drop_duplicates(subset=["run_uuid"])
-    return data
-
-
 # -----------------------------------------------------------------------------
 # Plotting functions
 # -----------------------------------------------------------------------------