Cohort Tracker

docs/usage/usage.md

@@ -253,6 +253,16 @@ In contrast to a preprocessing function, a tool usually adds an easily interpret
     tools.causal_inference
 ```
 
+### Cohort Tracking
+
+```{eval-rst}
+.. autosummary::
+    :toctree: tools
+    :nosignatures:
+
+    tools.CohortTracker
+```
+
 ## Plotting
 
 The plotting module `ehrapy.pl.\*` largely parallels the `tl.\*` and a few of the `pp.\*` functions.

ehrapy/tools/__init__.py

@@ -1,6 +1,7 @@
 from ehrapy.tools._sa import anova_glm, cox_ph, glm, kmf, ols, test_kmf_logrank, test_nested_f_statistic
 from ehrapy.tools._scanpy_tl_api import *  # noqa: F403
 from ehrapy.tools.causal._dowhy import causal_inference
+from ehrapy.tools.cohort_tracking._cohort_tracker import CohortTracker
 from ehrapy.tools.feature_ranking._rank_features_groups import filter_rank_features_groups, rank_features_groups
 
 try:  # pragma: no cover

tests/conftest.py

@@ -1,8 +1,39 @@
+import os
 from pathlib import Path
 
 import pytest
+from matplotlib import pyplot as plt
+from matplotlib.figure import Figure
+from matplotlib.testing.compare import compare_images
 
 
 @pytest.fixture
 def root_dir():
     return Path(__file__).resolve().parent
+
+
+# simplified from https://github.com/scverse/scanpy/blob/main/scanpy/tests/conftest.py
+@pytest.fixture
+def check_same_image(tmp_path):
+    def check_same_image(
+        fig: Figure,
+        base_path: Path | os.PathLike,
+        *,
+        tol: float,
+    ) -> None:
+        expected = Path(base_path).parent / (Path(base_path).name + "_expected.png")
+        if not Path(expected).is_file():
+            raise OSError(f"No expected output found at {expected}.")
+        actual = tmp_path / "actual.png"
+
+        fig.tight_layout()
+        fig.savefig(actual, dpi=80)
+
+        result = compare_images(expected, actual, tol=tol, in_decorator=True)
+
+        if result is None:
+            return None
+
+        raise AssertionError(result)
+
+    return check_same_image

tests/tools/cohort_tracking/test_cohort_tracking.py

@@ -0,0 +1,145 @@
+from pathlib import Path
+
+import pytest
+
+import ehrapy as ep
+from ehrapy.io._read import read_csv
+
+CURRENT_DIR = Path(__file__).parent
+_TEST_DATA_PATH = f"{CURRENT_DIR.parent}/test_data_features_ranking"
+_TEST_IMAGE_PATH = f"{CURRENT_DIR.parent}/_images"
+
+
+@pytest.fixture
+def adata_mini():
+    return read_csv(f"{_TEST_DATA_PATH}/dataset1.csv", columns_obs_only=["glucose", "weight", "disease", "station"])
+
+
+@pytest.mark.parametrize("columns", [None, ["glucose", "weight", "disease", "station"]])
+def test_CohortTracker_init_vanilla(columns, adata_mini):
+    ct = ep.tl.CohortTracker(adata_mini, columns)
+    assert ct._tracked_steps == 0
+    assert ct.tracked_steps == 0
+    assert ct._tracked_text == []
+    assert ct._tracked_operations == []
+
+
+def test_CohortTracker_type_detection(adata_mini):
+    ct = ep.tl.CohortTracker(adata_mini, ["glucose", "weight", "disease", "station"])
+    assert set(ct.categorical) == {"disease", "station"}
+
+
+def test_CohortTracker_init_set_columns(adata_mini):
+    # limit columns
+    ep.tl.CohortTracker(adata_mini, columns=["glucose", "disease"])
+
+    # invalid column
+    with pytest.raises(ValueError):
+        ep.tl.CohortTracker(
+            adata_mini,
+            columns=["glucose", "disease", "non_existing_column"],
+        )
+
+    # force categoricalization
+    ep.tl.CohortTracker(adata_mini, columns=["glucose", "disease"], categorical=["glucose", "disease"])
+
+    # invalid category
+    with pytest.raises(ValueError):
+        ep.tl.CohortTracker(
+            adata_mini,
+            columns=["glucose", "disease"],
+            categorical=["station"],
+        )
+
+
+def test_CohortTracker_call(adata_mini):
+    ct = ep.tl.CohortTracker(adata_mini)
+
+    ct(adata_mini)
+    assert ct.tracked_steps == 1
+    assert ct._tracked_text == ["Cohort 0\n (n=12)"]
+
+    ct(adata_mini)
+    assert ct.tracked_steps == 2
+    assert ct._tracked_text == ["Cohort 0\n (n=12)", "Cohort 1\n (n=12)"]
+
+
+def test_CohortTracker_reset(adata_mini):
+    ct = ep.tl.CohortTracker(adata_mini)
+
+    ct(adata_mini)
+    ct(adata_mini)
+
+    ct.reset()
+    assert ct.tracked_steps == 0
+    assert ct._tracked_text == []
+    assert ct._tracked_operations == []
+
+
+def test_CohortTracker_plot_cohort_change_test_sensitivity(adata_mini, check_same_image):
+    ct = ep.tl.CohortTracker(adata_mini)
+
+    # check that e.g. different color triggers error
+    ct(adata_mini, label="First step", operations_done="Some operations")
+    fig1, _ = ct.plot_cohort_change(show=False, color_palette="husl")
+
+    with pytest.raises(AssertionError):
+        check_same_image(
+            fig=fig1,
+            base_path=f"{_TEST_IMAGE_PATH}/cohorttracker_adata_mini_step1",
+            tol=1e-1,
+        )
+
+
+def test_CohortTracker_plot_cohort_change(adata_mini, check_same_image):
+    ct = ep.tl.CohortTracker(adata_mini)
+
+    ct(adata_mini, label="First step", operations_done="Some operations")
+    fig1, _ = ct.plot_cohort_change(show=False)
+
+    check_same_image(
+        fig=fig1,
+        base_path=f"{_TEST_IMAGE_PATH}/cohorttracker_adata_mini_step1",
+        tol=1e-1,
+    )
+
+    ct(adata_mini, label="Second step", operations_done="Some other operations")
+    fig2, _ = ct.plot_cohort_change(show=False)
+
+    check_same_image(
+        fig=fig2,
+        base_path=f"{_TEST_IMAGE_PATH}/cohorttracker_adata_mini_step2",
+        tol=1e-1,
+    )
+
+
+def test_CohortTracker_flowchart_sensitivity(adata_mini, check_same_image):
+    ct = ep.tl.CohortTracker(adata_mini)
+
+    ct(adata_mini, label="Base Cohort")
+    ct(adata_mini, operations_done="Some processing")
+
+    # check that e.g. different arrow size triggers error
+    fig, _ = ct.plot_flowchart(show=False, arrow_size=0.5)
+
+    with pytest.raises(AssertionError):
+        check_same_image(
+            fig=fig,
+            base_path=f"{_TEST_IMAGE_PATH}/cohorttracker_adata_mini_flowchart",
+            tol=1e-1,
+        )
+
+
+def test_CohortTracker_flowchart(adata_mini, check_same_image):
+    ct = ep.tl.CohortTracker(adata_mini)
+
+    ct(adata_mini, label="Base Cohort")
+    ct(adata_mini, operations_done="Some processing")
+
+    fig, _ = ct.plot_flowchart(show=False)
+
+    check_same_image(
+        fig=fig,
+        base_path=f"{_TEST_IMAGE_PATH}/cohorttracker_adata_mini_flowchart",
+        tol=1e-1,
+    )

tests/tools/ehrapy_data/dataset1.csv

@@ -0,0 +1,13 @@
+idx,sys_bp_entry,dia_bp_entry,glucose,weight,disease,station
+1,138,78,80,77,A,ICU
+2,139,79,90,76,A,ICU
+3,140,80,120,60,A,MICU
+4,141,81,130,90,A,MICU
+5,148,77,80,110,B,ICU
+6,149,78,135,78,B,ICU
+7,150,79,125,56,B,MICU
+8,151,80,95,76,B,MICU
+9,158,55,70,67,C,ICU
+10,159,56,85,82,C,ICU
+11,160,57,125,59,C,MICU
+12,161,58,125,81,C,MICU