docs: Equations and plots (#339)

- Tweak plots, including color scale
- bug: Preprocess filtering now removes dates after end of last season
- Improve equation notation
- Add tool for rendering equation

Author: swo

Makefile

@@ -1,4 +1,3 @@
-CONFIG = scripts/config_vignette.yaml
 RAW_DATA = data/raw.parquet
 
 RUN_ID = $(shell python scripts/get_run_id.py --config=$(CONFIG))

README.md

@@ -1,6 +1,6 @@
 # Vaccination coverage forecasting
 
-_Formerly known as Immunization Uptake Projections, or `vcf`._
+_Formerly known as Immunization Uptake Projections, or `iup`._
 
 ## Getting started
 
@@ -22,9 +22,8 @@ The vignette uses monthly estimates of season flu vaccination coverage, from the
 
 ### Running the vignette
 
-1. Run the pipeline with `make`. (You can run steps in parallel with, e.g., `make -j4`.)
-   - By default, `make` will use `scripts/config_vignette.yaml` for its configuration.
-   - You can use different configs by running `make CONFIG=/path/to/config.yaml`.
+1. Run the pipeline with `mise run vignette`.
+   - Inspect `mise.toml` for how to use `make` to run with different numbers of threads and different configs.
 2. Inspect `output/vignette/`:
    - `config.yaml`: a copy of the input config
    - `data.parquet`: the preprocessed, observed data

docs/lpl.md

@@ -4,50 +4,44 @@ This model proposes a latent true coverage curve, which is subject to observatio
 
 ## Terminology and notation
 
-- _feature_: A categorical feature of the data (e.g., season, geography) that partially determines coverage
-- _level_: A value taken on by a feature (e.g., New Jersey, or 2018/2019)
-- _group_: A unique combination of features (e.g., New Jersey in 2018/2019)
-- $t$: time since the start of the season, measured in $\text{year}^{-1}$
-- $n_{gt}$: number of people in group $g$, surveyed at time $t$. Drawn from the `sample_size` column of the NIS data; called `N_tot` in the codebase.
-- $x_{gt}$: number of people in group $g$, surveyed at time $t$, who are vaccinated. Approximated as $\mathrm{round}(\hat{v}_{gt}, n_{gt})$, where $\hat{v}_{gt}$ is the `estimate` column.
-- $v_g(t)$: latent true coverage among group $g$ at time $t$
-- $z_{gj}$: integer index indicating the level of the $j$-th feature for group $g$.
-
-For example, let the features be season and geography, in that order. Let group 5 be associated with the fourth season and the third geography. Then $z_{51} = 4$ and $z_{52} = 3$.
-
-## Model overview
-
-For each group $g$, the latent coverage $v_g(t)$ is assumed to be a sum of a logistic curve and a line with intercept at $t=0$. The shape parameter $K$ and midpoint $\tau$ of the logistic curve are assumed to be common to all groups. The height $A_g$ of the logistic curve is a grand mean $\mu_A$ plus effects $\delta_{A,j,z_{gj}}$ for each feature $j$. For example, the $A_g$ for Alaska in 2018/2019 will be the grand mean $\mu_A$, plus the Alaska effect, plus the 2018/2019 effect. In fact, the model uses a third interaction term season-geography, so there is also an Alaska-in-2018/2019 term.
-
-The slopes $M_g$ follow a similar pattern.
-
-The observations $x_{gt}$ are beta binomial-distributed around the mean $v_g(t) \cdot n_{gt}$, with variance modified by an extra parameter $D$.
+- $C_{sgm}$: reported, estimated coverage (proportion of people vaccinated) in season $s$, geography (i.e., state) $g$, and month $m$ (`estimate` in the data)
+- $N_{sgm}$: number of people surveyed (`sample_size` in the data)
+- $v_{sg}(t)$: latent true coverage at time $t$ (measured in years)
 
 ## Model equations
 
 ```math
 \begin{align*}
-x_{gt} &\sim \mathrm{BetaBinom}\big(v_g(t) \cdot D, [1-v_g(t)] \cdot D, n_{gt}\big) \\
-v_g(t) &= \frac{A_g}{1 + \exp\{- K \cdot (t - \tau)\}} + M_g t \\
-A_g &= \mu_A + \sum_j \delta_{Aj z_{gj}} \\
-M_g &= \mu_M + \sum_j \delta_{Mj z_{gj}} \\
-\mu_A &\sim \text{Beta}(100.0, 180.0) \\
-\mu_M &\sim \text{Gamma}(\text{shape} = 1.0, \text{rate} = 10.0) \\
-\delta_{Ajk} &\sim \mathcal{N}(0, \sigma_{Aj}) \\
-\delta_{Mjk} &\sim \mathcal{N}(0, \sigma_{Mj}) \\
-\sigma_{Aj} &\sim \text{Exp}(40.0) \\
-\sigma_{Mj} &\sim \text{Exp}(40.0) \\
+X_{sgm} &= \mathrm{round}(C_{sgm} \cdot N_{sgm}) \\
+v_{sg}(t) &= \frac{A_{sg}}{1 + \exp\{- K \cdot (t - \tau)\}} + B_{sg} t \\
+A_{sg} &= \beta^{(A)} + \beta_s^{(AS)} + \beta_g^{(AG)} + \beta_{sg}^{(ASG)} \\
+B_{sg} &= \beta^{(B)} + \beta_s^{(BS)} + \beta_g^{(BG)} + \beta_{sg}^{(BSG)} \\
+\\
+X_{sgm} &\sim \mathrm{BetaBinom}\big( N_{sgm}, v_{sg}(t_m) \cdot D, [1-v_{sg}(t_m)] \cdot D \big) \\
 K &\sim \text{Gamma}(\text{shape} = 25.0, \text{rate} = 1.0) \\
 \tau &\sim \text{Beta}(100.0, 225.0) \\
 D &\sim \text{Gamma}(\text{shape} = 350.0, \text{rate} = 1.0) \\
+\beta^{(A)} &\sim \text{Beta}(100.0, 180.0) \\
+\beta_s^{(AS)} &\sim \mathcal{N}\left(0, \sigma^{(AS)}\right) \\
+\beta_g^{(AG)} &\sim \mathcal{N}\left(0, \sigma^{(AG)}\right) \\
+\beta_{sg}^{(ASG)} &\sim \mathcal{N}\left(0, \sigma^{(ASG)}\right) \\
+\beta^{(B)} &\sim \text{Gamma}(\text{shape} = 1.0, \text{rate} = 10.0) \\
+\beta_s^{(BS)} &\sim \mathcal{N}\left(0, \sigma^{(BS)}\right) \\
+\beta_g^{(BG)} &\sim \mathcal{N}\left(0, \sigma^{(BG)}\right) \\
+\beta_{sg}^{(BSG)} &\sim \mathcal{N}\left(0, \sigma^{(BSG)}\right) \\
+\sigma^{(\bullet)} &\sim \text{Exp}(40.0) \\
 \end{align*}
 ```
 
-Note that:
+Notes:
+
+- The latent coverage $v_{sg}(t)$ is assumed to be a sum of a logistic curve and a line with intercept at $t=0$
+- The shape parameter $K$ and midpoint $\tau$ of the logistic curve are assumed to be common to all groups
+- The height $A_{sg}$ of the logistic curve is a grand mean $\beta^{(A)}$ plus effects for the season, state, and season-state interaction. The slopes $M_g$ follow a similar pattern.
 
 ```math
 \begin{align*}
-\mathbb{E}[x_{gt}] &= v_g(t) \cdot n_{gt} \\
-\mathrm{Var}[x_{gt}] &= v_g(t) \cdot [1-v_g(t)] \cdot \frac{n_{gt} (n_{gt} + D)}{D+1}
+\mathbb{E}[X_{sgk}] &= v_{sg}(t_k) \cdot N_{sgk} \\
+\mathrm{Var}[X_{sgk}] &= v_{sg}(t_k) \cdot [1-v_{sg}(t_k)] \cdot \frac{N_{sgk} (N_{sgk} + D)}{D+1}
 \end{align*}
 ```

docs/render_eq.py

@@ -0,0 +1,71 @@
+"""
+Render equation block as an image
+"""
+
+import subprocess
+import tempfile
+from pathlib import Path
+
+import pdf2image
+
+
+def get_eqn(path="docs/lpl.md") -> str:
+    """Collect the first back-ticked math block from `path`"""
+    lines = []
+    in_eq = False
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if not in_eq and line == "```math":
+                in_eq = True
+            elif in_eq and line == "```":
+                in_eq = False
+                break
+            elif in_eq:
+                lines.append(line)
+            else:
+                pass
+
+    return "\n".join(lines)
+
+
+def latex_to_png(latex_str: str, out_path: str, dpi: int = 400, verbose: bool = False):
+    """Render `latex_str` using pdflatex"""
+    content = (
+        r"""
+\documentclass[preview,border=2pt]{standalone}
+\usepackage{amsmath,amssymb}
+\begin{document}
+%s
+\end{document}
+"""
+        % latex_str
+    )
+
+    with tempfile.TemporaryDirectory() as td:
+        path = Path(td) / "eq.tex"
+        path.write_text(content)
+        result = subprocess.run(
+            [
+                "pdflatex",
+                "-interaction=nonstopmode",
+                "-output-directory",
+                td,
+                str(path),
+            ],
+            capture_output=True,
+        )
+
+        if verbose:
+            print(result.stdout.decode("utf-8"))
+            print(result.stderr.decode("utf-8"))
+
+        pages = pdf2image.convert_from_path(str(Path(td) / "eq.pdf"), dpi=dpi)
+
+    pages[0].save(out_path, "PNG")
+
+
+if __name__ == "__main__":
+    eqn = get_eqn()
+    print(eqn)
+    latex_to_png(eqn, "tmp_eqn.png", verbose=False)

mise.toml

@@ -0,0 +1,3 @@
+[tasks]
+vignette = "make -j4 CONFIG=scripts/config_vignette.yaml"
+full = "make -j4 CONFIG=scripts/config_full.yaml"

scripts/eval.py

@@ -34,10 +34,10 @@
     assert len(forecast_season) == 1, "Can only score forecasts from one season"
     forecast_season = forecast_season[0]
 
-    eos_abs_diff = vcf.eos_abs_diff(
-        obs=data.filter(pl.col("season") == pl.lit(forecast_season)),
-        pred=pred.filter(pl.col("season") == pl.lit(forecast_season)),
-        features=["season", "geography"],
-    )
+    obs = data.filter(pl.col("season") == pl.lit(forecast_season))
+    pred = pred.filter(pl.col("season") == pl.lit(forecast_season))
+    features = ["season", "geography"]
 
-    eos_abs_diff.write_parquet(args.output)
+    scores = vcf.eos_abs_diff(obs=obs, pred=pred, features=features)
+
+    scores.write_parquet(args.output)

scripts/plot_data.py

@@ -8,6 +8,7 @@
 import yaml
 
 AXIS_PERCENT = alt.Axis(format=".0%")
+SCALE = alt.Scale(domain=(0, 0.70))
 LINE_OPACITY = 0.25
 
 MEDIAN_POINT_KWARGS = {
@@ -18,14 +19,10 @@
 }
 
 MEDIAN_ENCODINGS = [
-    alt.Color(
-        "type",
-        scale=alt.Scale(domain=["datum", "median"], range=["black", "red"]),
-        legend=None,
-    ),
     alt.Shape(
         "type",
         scale=alt.Scale(domain=["datum", "median"], range=["circle", "stroke"]),
+        legend=None,
     ),
     alt.Size("type", scale=alt.Scale(domain=["datum", "median"], range=[20, 200])),
 ]
@@ -187,20 +184,26 @@ def hightlight_state(df, month, order_n, value="score_value", state_var="geograp
     eos = data.filter((pl.col("time_end") == pl.col("time_end").max()).over("season"))
 
     # for each season, show eos spread over states
-    alt.Chart(add_medians(eos, "season")).mark_point().encode(
+    alt.Chart(eos).mark_boxplot(color="black").encode(
         alt.X("season", title=None),
-        alt.Y("estimate", title="End of season coverage", axis=AXIS_PERCENT),
-        *MEDIAN_ENCODINGS,
+        alt.Y(
+            "estimate", title="End of season coverage", axis=AXIS_PERCENT, scale=SCALE
+        ),
     ).save(out_dir / "coverage_by_season.svg")
 
     # for each state, show eos spread over seasons
-    alt.Chart(add_medians(eos, "geography")).mark_point().encode(
+    alt.Chart(add_medians(eos, "geography")).mark_point(color="black").encode(
         alt.X(
             "geography",
             title=None,
             sort=alt.EncodingSortField("estimate", "median", "descending"),
         ),
-        alt.Y("estimate", title="End of season coverage", axis=AXIS_PERCENT),
+        alt.Y(
+            "estimate",
+            title="End of season coverage",
+            axis=AXIS_PERCENT,
+            scale=SCALE,
+        ),
         *MEDIAN_ENCODINGS,
     ).save(out_dir / "coverage_by_state.svg")
 

scripts/plot_preds.py

@@ -9,27 +9,32 @@
 import vcf
 
 LINE_OPACITY = 0.4
+MODEL_COLOR_SCALE = alt.Scale(
+    domain=["LPLModel", "RFModel"], range=["#a6611a", "#018571"]
+)
 
 
 def plot_data_cone(
     chart_data: pl.DataFrame,
     facet_kwargs: dict,
-    out_dir: str,
+    out_dir: Path,
     filename: str,
     properties_kwargs: dict | None = None,
     config_legend_kwargs: dict | None = None,
     config_axis_kwargs: dict | None = None,
 ):
     base = alt.Chart(chart_data).encode(
-        alt.X("time_end", title=None, axis=alt.Axis(format="%b"))
+        alt.X("time_end", title=None, axis=alt.Axis(format="%b", grid=False))
     )
     fc_cone = base.mark_area(opacity=0.25).encode(
-        alt.Y("pred_lci", title="", axis=AXIS_PERCENT),
+        alt.Y(
+            "pred_lci", title="", axis=AXIS_PERCENT, scale=alt.Scale(domain=(0.0, 0.70))
+        ),
         alt.Y2("pred_uci"),
-        alt.Color("model"),
+        alt.Color("model", scale=MODEL_COLOR_SCALE),
     )
     fc_points = base.mark_line(opacity=0.75).encode(
-        alt.Y("pred_estimate"), alt.Color("model")
+        alt.Y("pred_estimate"), alt.Color("model", scale=MODEL_COLOR_SCALE)
     )
     data_points = base.mark_point(color="black").encode(alt.Y("obs_estimate"))
     data_error = base.mark_rule(color="black").encode(
@@ -135,7 +140,7 @@ def plot_data_cone(
     plot_data_cone(
         chart_data=chart_data,
         facet_kwargs={
-            "column": "forecast_date",
+            "column": alt.Column("forecast_date", header=alt.Header(format="%b %Y")),
             "row": "geography",
         },
         out_dir=out_dir,
@@ -174,7 +179,9 @@ def plot_data_cone(
                 sort=["South Dakota", "North Dakota"],
             ),
             "column": alt.Column(
-                "forecast_date", header=alt.Header(labelFontSize=20), title=""
+                "forecast_date",
+                header=alt.Header(format="%b %Y", labelFontSize=20),
+                title="",
             ),
         },
         config_axis_kwargs={"labelFontSize": 20, "titleFontSize": 24},
@@ -196,7 +203,9 @@ def plot_data_cone(
                 sort=["Wyoming", "Vermont"],
             ),
             "column": alt.Column(
-                "forecast_date", header=alt.Header(labelFontSize=20), title=""
+                "forecast_date",
+                header=alt.Header(format="%b %Y", labelFontSize=20),
+                title="",
             ),
         },
         config_axis_kwargs={"labelFontSize": 20, "titleFontSize": 24},

scripts/plot_scores.py

@@ -4,6 +4,7 @@
 import altair as alt
 import polars as pl
 import yaml
+from plot_preds import MODEL_COLOR_SCALE
 
 LINE_OPACITY = 0.4
 
@@ -26,34 +27,17 @@
     out_dir = out_flag.parent
     out_dir.mkdir(parents=True, exist_ok=True)
 
-    # sis = score in season
-    data = scores.filter(pl.col("score_fun") == pl.lit("eos_abs_diff"))
-
-    base = alt.Chart(data).encode(
-        alt.X("forecast_date", type="temporal", axis=alt.Axis(format="%b"))
-    )
-    line_chart = base.mark_line(point=True, opacity=LINE_OPACITY).encode(
-        alt.Y("score_value", title="Score (abs. end-of-season diff.)"),
-        alt.Detail("geography"),
-        alt.Color("model"),
-    )
-
-    # Filter for the final forecast date and use the gather_n function to
-    # make plots with ticks
-
-    line_chart.save(out_dir / "scores.svg")
-
     ## boxplot of scores across states by forecast date ##
     alt.Chart(scores).mark_boxplot(ticks=True).encode(
-        x=alt.X(
+        alt.X(
             "model",
             title=None,
             axis=alt.Axis(labels=False, ticks=False),
             scale=alt.Scale(padding=1),
         ),
-        y=alt.Y("score_value", title="End-of-season Abs Diff"),
-        color="model",
-        column=alt.Column(
+        alt.Y("score_value", title="End-of-season abs. diff. (p.p.)"),
+        alt.Color("model", scale=MODEL_COLOR_SCALE),
+        alt.Column(
             "forecast_date",
             title="",
             header=alt.Header(orient="bottom", labelFontSize=20, format="%b"),
@@ -62,6 +46,6 @@
         stroke=None
     ).configure_axis(labelFontSize=20, titleFontSize=24).configure_legend(
         labelFontSize=20, title=None
-    ).save(out_dir / "scores_boxplot.svg")
+    ).save(out_dir / "scores.svg")
 
     out_flag.touch()

scripts/preprocess.py

@@ -46,6 +46,9 @@ def preprocess(
 
     """
 
+    start_date = pl.date(start_year, season_start_month, season_start_day)
+    end_date = pl.date(end_year, season_end_month, season_end_day)
+
     def geo_filter(df: pl.DataFrame) -> pl.DataFrame:
         """Optionally retain rows belonging to selected geographies."""
         if geographies is None:
@@ -70,8 +73,8 @@ def geo_filter(df: pl.DataFrame) -> pl.DataFrame:
             )
         )
         .filter(
-            # drop dates before or after the outermost season
-            pl.col(date_col).dt.year().is_between(start_year, end_year),
+            # keep dates only within the desired data range
+            pl.col(date_col).is_between(start_date, end_date),
             # drop out-of-season dates between seasons
             pl.col("season").is_null().not_(),
         )