Dhaka Cholera: IF2 vs IFAD

Which algorithm reaches the higher likelihood for a comparable budget

import os
import sys
sys.path.append("..")
import report_utils as ru
from IPython.display import display, HTML, Markdown

display(HTML(f"<div style='margin-bottom: 20px;'>{ru.nav_bar('algorithms')}</div>"))

IF2 vs IFAD  |  Timing & Throughput  |  Likelihood Evaluation

Introduction

This report compares two pypomp fitting pipelines on the Dhaka cholera model:

1. IF2 — iterated filtering alone.
2. IFAD — iterated filtering followed by gradient-based training.

Both arms start from the same points drawn from the same box, and IF2 is given enough extra iterations that the two arms cost roughly the same wall clock. We would like to see which maximizes the likelihood more in the (roughly) same amount of time.

A more thorough version of this analysis is available here.

Parameter definitions

• gamma: Recovery rate from infection.
• m: Cholera mortality rate.
• rho: Reporting rate of cholera cases.
• epsilon: Rate of waning of immunity for inapparent infections.
• c: Fraction of infections that lead to severe infection.
• alpha: Non-linear transmission factor.
• delta: Cholera-induced mortality rate.
• beta_trend: Long-term secular trend in transmission.
• sigma: Environmental noise intensity.
• tau: Measurement noise (overdispersion parameter).
• bs1 - bs6: Spline coefficients modeling seasonal transmission rates.
• omegas1 - omegas6: Spline coefficients for seasonal environmental reservoir.
• S_0, I_0, Y_0, R1_0, R2_0, R3_0: Initial compartment fractions.

rho, c, alpha, delta and Y_0 are held fixed by both searches.

---

Benchmark Settings & Environment

The table below summarizes algorithmic parameters, software environments, and compute hardware recorded in latest.json for each configuration.

PLATFORMS = {
    "IF2": os.path.join("results", "if2"),
    "IFAD": os.path.join("results", "ifad"),
}

runs = ru.load_timing_data(PLATFORMS)
display(HTML(ru.build_settings_comparison_html(runs, is_panel=False)))

Setting / Parameter	IF2	IFAD
Algorithmic & Workload Settings
Run Level	4	4
Starting Searches ($N_{starts}$)	100	100
IF2 Iterations ($N_{iter}$)	1,000	175
Training Iterations ($N_{train}$)	0	175
IF2 Particles ($N_p$)	5,000	5,000
Pfilter Particles ($N_{p,eval}$)	5,000	5,000
Evaluation Replicates ($N_{reps}$)	36	36
Random Seed	631409	631409
Software & Environment
Pomp Framework	pypomp 1.0.0rc1	pypomp 1.0.0rc1
Backend / Engine	JAX 0.11.1	JAX 0.11.1
Quant Git Commit	630cc23	630cc23
Run Timestamp	2026-09-02 18:38:22	2026-09-02 18:41:53
Hardware & Compute
Compute Device	NVIDIA RTX PRO 6000 Blackwell Server Edition (1 GPU)	NVIDIA RTX PRO 6000 Blackwell Server Edition (1 GPU)
Slurm Partition	gpu-rtx6000	gpu-rtx6000
Slurm Job ID	59673746	59673747

missing = [(label, r["dir"]) for label, r in runs.items() if not r["available"]]
for label, path in missing:
    display(HTML(f"<div class='alert alert-warning'><strong>Missing results for {label}:</strong> Expected at <code>{path}</code>.</div>"))

import numpy as np
import pandas as pd
from plotnine import (
    ggplot, aes, geom_line, geom_density, facet_wrap, labs,
    coord_cartesian, theme
)

os.environ["JAX_PLATFORMS"] = "cpu"

ARMS = {"if2": "IF2", "ifad": "IFAD"}

meta = {arm: ru.load_json(os.path.join("results", arm, "latest.json")) for arm in ARMS}
final = {
    arm: pd.read_csv(os.path.join("results", arm, "results_final.csv"))
    for arm in ARMS
}
logliks = {
    arm: pd.read_csv(os.path.join("results", arm, "pfilter_logliks.csv"))
    for arm in ARMS
}
timings = {
    arm: pd.read_csv(os.path.join("results", arm, "timings.csv")) for arm in ARMS
}
traces = {
    arm: ru.load_traces(os.path.join("results", arm, "traces.csv.gz"), label)
    for arm, label in ARMS.items()
}

---

Runtimes

Validation expectation

IF2 alone should take a little longer than IFAD overall — the point of the comparison is that IFAD reaches a better likelihood in the same or less time. train should also cost no more than about three times what a mif iteration costs.

rows = []
for arm, label in ARMS.items():
    tm = timings[arm]
    cfg = meta[arm].get("run_config", {})
    per_stage = tm.groupby("method", sort=False)["time"].sum()
    n_mif = cfg.get("NFITR")
    n_train = cfg.get("NTRAIN") or 0
    rows.append({
        "Pipeline": label,
        "mif (s)": f"{per_stage.get('mif', np.nan):.1f}",
        "train (s)": f"{per_stage.get('train', 0.0):.1f}",
        "pfilter (s)": f"{per_stage.get('pfilter', np.nan):.1f}",
        "Total (s)": f"{tm['time'].sum():.1f}",
        "mif iters": n_mif,
        "train iters": n_train,
        "s / mif iter": f"{per_stage.get('mif', np.nan) / n_mif:.3f}" if n_mif else "—",
        "s / train iter": f"{per_stage.get('train', 0.0) / n_train:.3f}" if n_train else "—",
    })

display(HTML(pd.DataFrame(rows).to_html(classes="table table-striped table-hover", index=False)))

Pipeline	mif (s)	train (s)	pfilter (s)	Total (s)	mif iters	train iters	s / mif iter	s / train iter
IF2	536.0	0.0	19.8	555.8	1000	0	0.536	—
IFAD	99.0	413.5	19.7	532.2	175	175	0.566	2.363

---

Final log-likelihoods

Validation expectation

We expect a maximum IFAD log-likelihood around -3744.0 and a maximum IF2 log-likelihood around -3748.5.

summary = pd.DataFrame([
    {
        "Pipeline": label,
        "Starts": len(final[arm]),
        "Mean": final[arm]["logLik"].mean(),
        "Std": final[arm]["logLik"].std(),
        "Min": final[arm]["logLik"].min(),
        "Median": final[arm]["logLik"].median(),
        "Max": final[arm]["logLik"].max(),
    }
    for arm, label in ARMS.items()
]).round(2)

display(HTML(summary.to_html(classes="table table-striped table-hover", index=False)))

Pipeline	Starts	Mean	Std	Min	Median	Max
IF2	100	-3762.45	18.73	-3847.56	-3755.36	-3747.35
IFAD	100	-4174.36	2971.16	-24868.00	-3747.97	-3744.04

ll_df = pd.concat(
    [
        pd.DataFrame({"logLik": final[arm]["logLik"], "source": label})
        for arm, label in ARMS.items()
    ],
    ignore_index=True,
)
# A handful of starts stall far below the mode; clipping keeps the KDE readable.
ll_df = ll_df[ll_df["logLik"] >= -3800]

(
    ggplot(ll_df, aes(x="logLik", fill="source", color="source"))
    + geom_density(alpha=0.4)
    + labs(
        title="Final logLik density: IF2 vs IFAD",
        subtitle="Evaluated likelihood at each start's final parameter vector",
        x="Log-Likelihood", y="Density", fill="Pipeline", color="Pipeline",
    )
    + ru.scale_fill_premium()
    + ru.scale_color_premium()
    + ru.theme_premium
)

---

Final parameter estimates

Validation expectation

IFAD tends to yield parameter estimates with lower Monte Carlo variance than IF2 alone, so its density curves should be the narrower ones.

param_cols = [
    c for c in final["if2"].columns if c not in ("theta_idx", "logLik", "se")
]
params_df = pd.concat(
    [
        final[arm][param_cols].melt(var_name="quantity", value_name="param_value").assign(source=label)
        for arm, label in ARMS.items()
    ],
    ignore_index=True,
)

(
    ggplot(params_df, aes(x="param_value", fill="source", color="source"))
    + geom_density(alpha=0.4)
    + facet_wrap("quantity", scales="free", ncol=3)
    + labs(
        title="Final parameter estimates: IF2 vs IFAD",
        subtitle="Density across all starts",
        x="Estimated value", y="Density", fill="Pipeline", color="Pipeline",
    )
    + ru.scale_fill_premium()
    + ru.scale_color_premium()
    + ru.theme_premium
    + theme(figure_size=(10, 15), panel_spacing=0.02)
)

/home/aaronabk/research/quant/.venv/lib64/python3.12/site-packages/plotnine/layer.py:293: PlotnineWarning: stat_density : Removed 56 rows containing non-finite values.

---

Search traces

have_traces = all(t is not None for t in traces.values())
if have_traces:
    long = {arm: ru.to_long(traces[arm]) for arm in ARMS}

Validation expectation

The IF2 traces should not look terrible, but several parameters in this model are weakly identified, so tight convergence is not expected everywhere. Under IFAD, most parameters should tighten once the search switches from mif to train.

if have_traces:
    p = (
        ggplot(long["if2"], aes(x="iter", y="param_value", group="rep", color="factor(rep)"))
        + geom_line(show_legend=False, alpha=0.5)
        + facet_wrap("quantity", scales="free_y", ncol=3)
        + labs(
            title="IF2 parameter traces",
            subtitle="Individual start paths across iterations",
            x="Iteration", y="Parameter value",
        )
        + ru.theme_premium
        + theme(figure_size=(10, 15), panel_spacing=0.02)
    )
    display(p)
else:
    display(Markdown("_No traces on disk._"))

if have_traces:
    p = (
        ggplot(long["ifad"], aes(x="iter", y="param_value", group="rep", color="factor(rep)"))
        + geom_line(show_legend=False, alpha=0.5)
        + facet_wrap("quantity", scales="free_y", ncol=3)
        + labs(
            title="IFAD parameter traces",
            subtitle="Individual start paths across mif and train iterations",
            x="Iteration", y="Parameter value",
        )
        + ru.theme_premium
        + theme(figure_size=(10, 15), panel_spacing=0.02)
    )
    display(p)
else:
    display(Markdown("_No traces on disk._"))

if have_traces:
    ll_traces = pd.concat(
        [
            ru.thin(traces[arm])[["rep", "iter", "logLik", "source"]].dropna(subset=["logLik"])
            for arm in ARMS
        ],
        ignore_index=True,
    )
    p = (
        ggplot(ll_traces, aes(x="iter", y="logLik", group="rep", color="factor(rep)"))
        + geom_line(show_legend=False, alpha=0.6)
        + facet_wrap("source", ncol=1, scales="free_x")
        + coord_cartesian(ylim=(-4000, None))
        + labs(
            title="logLik traces",
            subtitle="Log-likelihood trajectory of each start",
            x="Iteration", y="Log-Likelihood",
        )
        + ru.theme_premium
        + theme(figure_size=(8, 8))
    )
    display(p)
else:
    display(Markdown("_No traces on disk._"))