Dhaka Cholera: IF2 vs IFAD

Which algorithm reaches the higher likelihood for a comparable budget

Published

October 16, 2026

Show Code

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 | Scaling | Likelihood Evaluation

Introduction

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

IF2: iterated filtering alone.
IFAD: iterated filtering followed by gradient-based training.

Both pipelines start from the same points drawn from the same box, and IF2 is given enough extra iterations that the two cost roughly the same wall-clock time. We would like to see which maximizes the likelihood more in roughly the 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.

Show Code

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"

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

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()
}

Show Code

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>"))

Visual review

Runtimes: IF2 alone should take a little longer than IFAD overall, and a train iteration should cost no more than about three times a mif iteration.
Final log-likelihoods: the maximum IFAD log-likelihood should be around -3744.0, and the maximum IF2 log-likelihood around -3748.5.
Final parameter estimates: the IFAD densities should be narrower than the IF2 densities.
Search traces: under IFAD, most parameters should tighten once the search switches from mif to train.

Runtimes

Show Code

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

Show Code

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

Starts with a final log-likelihood below -3800 are left out of the plot.

Show Code

ll_df = pd.concat(
    [
        pd.DataFrame({"logLik": final[arm]["logLik"], "source": label})
        for arm, label in ARMS.items()
    ],
    ignore_index=True,
)
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

Show Code

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/lib/python3.14/site-packages/plotnine/layer.py:293: PlotnineWarning: stat_density : Removed 56 rows containing non-finite values.

Search traces

Show Code

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}

Show Code

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._"))

Show Code

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._"))

Show Code

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._"))

Appendix: settings and environment

Show Code

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