Phase 2.1: incidence DX (rolling/cumulative + examples)

PR_PHASE2_1.md

@@ -0,0 +1,29 @@
+Phase 2.1: incidence DX improvements (docs + rolling/cumulative)
+
+Key points
+- Add optional `cumulative=True` for wide output.
+- Add optional `rolling=<int>` for wide output with `rolling_kind` (sum/mean).
+- Expand README examples for incidence.
+
+Why this implementation
+- These options are common in epi workflows (cumulative incidence, smoothed curves) and improve usability without changing the core incidence semantics.
+- Implemented at the wide-output layer so behavior is consistent across strata columns.
+
+Multi-role debate (differences, not consensus)
+
+Role A — pragmatic developer
+- 👍 Likes: small API surface change, easy to test, minimal code.
+- ⚠️ Concern: rolling/cumulative after pivot means we assume the index ordering fully represents time (true for date and (epi_year, epi_week) sorted).
+
+Role B — architecture
+- 👍 Likes: keeps incidence core as counts; transformations are optional flags.
+- ⚠️ Concern: feature creep: might prefer a separate `transform_incidence()` pipeline later.
+
+Role C — developer user (DX)
+- 👍 Likes: fewer steps in notebooks; one-liners for rolling/cumulative.
+- ⚠️ Concern: wants more control (e.g., min_periods, centered rolling, cumulative per-calendar-year vs epi-year).
+
+Points of divergence to revisit later
+1) Whether these transforms belong inside `incidence()` long-term.
+2) Add `min_periods`, `center`, and `cumulative_by` options.
+3) Apply transforms to `output="long"` as well.

README.md

@@ -33,7 +33,9 @@ print(week_only)
 
 # Incidence (line list -> counts)
 # df = pandas.DataFrame({"onset_date": [...], "sex": [...]})
-# weekly = epydem.incidence(df, date_col="onset_date", freq="W-MMWR", by=["sex"])
+# weekly = epydem.incidence(df, date_col="onset_date", freq="W-MMWR", by=["sex"], fill_missing=True)
+# weekly_rolling2 = epydem.transform_incidence(weekly, rolling=2)
+# weekly_cum = epydem.transform_incidence(weekly, cumulative=True)
 ```
 
 ## Roadmap (high level)

epydem/__init__.py

@@ -3,6 +3,7 @@
 from .epiweek import calculate
 from .incidence import incidence
 from .time import epiweek, epiweek_number, mmwr_week, mmwr_week1_start, parse_ymd
+from .transform import transform_incidence
 
 __all__ = [
     "calculate",
@@ -12,4 +13,5 @@
     "mmwr_week1_start",
     "parse_ymd",
     "incidence",
+    "transform_incidence",
 ]

epydem/transform.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+import pandas as pd
+
+RollingKind = Literal["sum", "mean"]
+
+
+@dataclass(frozen=True)
+class TransformSpec:
+    rolling: int | None = None
+    rolling_kind: RollingKind = "sum"
+    min_periods: int = 1
+    center: bool = False
+    cumulative: bool = False
+
+
+def transform_incidence(
+    data: pd.DataFrame,
+    *,
+    rolling: int | None = None,
+    rolling_kind: RollingKind = "sum",
+    min_periods: int = 1,
+    center: bool = False,
+    cumulative: bool = False,
+    time_cols: Sequence[str] | None = None,
+) -> pd.DataFrame:
+    """Apply time-series transforms to incidence outputs.
+
+    This is intentionally separate from `incidence()` to keep `incidence()` as a
+    stable counts primitive.
+
+    Supported inputs:
+    - Wide incidence output: index is time (e.g. date or (epi_year, epi_week)).
+    - Long incidence output: requires `time_cols` (or inferred), and transforms
+      are applied per non-time column group.
+
+    Args:
+        data: Output from `epydem.incidence(...)`.
+        rolling: Rolling window size.
+        rolling_kind: "sum" or "mean".
+        min_periods: Passed to pandas rolling.
+        center: Passed to pandas rolling.
+        cumulative: If True, apply cumulative sum.
+        time_cols: For long-form data, the time columns. If None, we try to infer
+            ("date") or ("epi_year", "epi_week").
+
+    Returns:
+        Transformed DataFrame with the same shape/schema as input.
+    """
+
+    if rolling is None and not cumulative:
+        return data
+
+    # Wide: DataFrameIndex is time.
+    if time_cols is None and ("date" not in data.columns) and ("epi_year" not in data.columns):
+        wide = data.sort_index()
+        if rolling is not None:
+            r = wide.rolling(window=rolling, min_periods=min_periods, center=center)
+            wide = r.sum() if rolling_kind == "sum" else r.mean()
+        if cumulative:
+            wide = wide.cumsum()
+        return wide
+
+    # Long: operate on a value column per group.
+    df = data.copy()
+
+    if time_cols is None:
+        if "date" in df.columns:
+            time_cols = ["date"]
+        elif ("epi_year" in df.columns) and ("epi_week" in df.columns):
+            time_cols = ["epi_year", "epi_week"]
+        else:
+            raise ValueError("Cannot infer time_cols for long-form incidence")
+
+    value_cols = [c for c in df.columns if c not in set(time_cols)]
+    if len(value_cols) != 1:
+        raise ValueError(
+            "Long-form incidence must have exactly one value column (e.g., 'cases'). "
+            f"Got {value_cols}"
+        )
+    value_col = value_cols[0]
+
+    group_cols = [c for c in df.columns if c not in set(time_cols + [value_col])]
+
+    df = df.sort_values(list(time_cols))
+
+    def _apply(group: pd.DataFrame) -> pd.DataFrame:
+        s = group[value_col]
+        if rolling is not None:
+            r = s.rolling(window=rolling, min_periods=min_periods, center=center)
+            s = r.sum() if rolling_kind == "sum" else r.mean()
+        if cumulative:
+            s = s.cumsum()
+        group[value_col] = s
+        return group
+
+    if group_cols:
+        out = df.groupby(group_cols, dropna=False, sort=False, group_keys=False).apply(_apply)
+    else:
+        out = _apply(df)
+
+    return out

tests/test_incidence_phase2_1.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import pandas as pd
+
+import epydem
+
+
+def test_incidence_weekly_rolling_and_cumulative():
+    df = pd.DataFrame(
+        {
+            "onset_date": [
+                "2024-01-01",  # (2024,1)
+                "2024-01-08",  # (2024,2)
+                "2024-01-08",
+            ],
+            "province": ["A", "A", "A"],
+        }
+    )
+
+    base = epydem.incidence(
+        df,
+        date_col="onset_date",
+        freq="W-MMWR",
+        by=["province"],
+        fill_missing=True,
+    )
+    # week1=1, week2=2
+    assert base.loc[(2024, 1), "A"] == 1
+    assert base.loc[(2024, 2), "A"] == 2
+
+    roll2 = epydem.transform_incidence(
+        base,
+        rolling=2,
+        rolling_kind="sum",
+    )
+    assert roll2.loc[(2024, 1), "A"] == 1
+    assert roll2.loc[(2024, 2), "A"] == 3
+
+    cum = epydem.transform_incidence(
+        base,
+        cumulative=True,
+    )
+    assert cum.loc[(2024, 1), "A"] == 1
+    assert cum.loc[(2024, 2), "A"] == 3