|
| 1 | +# ------------------------------------------------------------- |
| 2 | +# Copyright (c) Henry Spatial Analysis. All rights reserved. |
| 3 | +# Licensed under the MIT License. See LICENSE in project root for information. |
| 4 | +# ------------------------------------------------------------- |
| 5 | + |
| 6 | +""" |
| 7 | +Fit a declining-hazard turnover model (``weibull`` / ``lomax`` / |
| 8 | +``constant_breakpoint``) and a ``constant`` baseline on the same data, then run |
| 9 | +the overdispersion diagnostics on the full population and compare the tag-age |
| 10 | +φ̂ gradient + WAIC. |
| 11 | +
|
| 12 | +The point: the production random-effects fit and the two-rate breakpoint both |
| 13 | +leave a steep φ̂-by-tag-age gradient (5.3 → 8.2 → ~106/135), because past the |
| 14 | +breakpoint they revert to a single constant rate. A continuous declining hazard |
| 15 | +— Weibull (duration dependence) or Lomax (gamma-frailty, the NB analogue) — |
| 16 | +tests whether modelling the decline as a continuum flattens that gradient. |
| 17 | +
|
| 18 | +Global parameters (3 each) are fit on a random subsample (ample for population- |
| 19 | +level params, keeps the dense NUTS fast); diagnostics + WAIC run on / sample |
| 20 | +from the full modelled population. Memory-safe on a 16 GB host (the dispersion |
| 21 | +loop streams one draw at a time; WAIC uses a bounded eval subset). |
| 22 | +
|
| 23 | + python -u scripts/exploratory/fit_hazard_model_and_diagnose.py \ |
| 24 | + --model-type lomax \ |
| 25 | + --observations ~/data/openpois/osm_data/20260521/osm_observations.parquet \ |
| 26 | + 2>&1 | tee ~/data/openpois/logs/hazard_lomax.log |
| 27 | +""" |
| 28 | +from __future__ import annotations |
| 29 | + |
| 30 | +import argparse |
| 31 | +import time |
| 32 | +from pathlib import Path |
| 33 | + |
| 34 | +import jax |
| 35 | +import jax.numpy as jnp |
| 36 | +import numpy as np |
| 37 | +import pandas as pd |
| 38 | +from jax.scipy.special import logsumexp |
| 39 | + |
| 40 | +from openpois.models import dispersion |
| 41 | +from openpois.models.model_fitter import ModelFitter |
| 42 | +from openpois.models.osm_models import get_model_class |
| 43 | +from openpois.models.setup import prepare_data_for_model |
| 44 | + |
| 45 | +CELL_COLS = ["shared_label", "msa_code", "urban_rural"] |
| 46 | +LOAD_COLS = [ |
| 47 | + "id", "shared_label", "msa_code", "urban_rural", "changed", |
| 48 | + "last_obs_timestamp", "obs_timestamp", "last_tag_timestamp", |
| 49 | +] |
| 50 | +NATURAL = ["lambda", "shape", "theta", "t_breakpoint", "lambda_1", "lambda_2", |
| 51 | + "delta"] |
| 52 | + |
| 53 | + |
| 54 | +def _log(msg: str) -> None: |
| 55 | + print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush = True) |
| 56 | + |
| 57 | + |
| 58 | +def _metadata(model_type: str, args) -> dict: |
| 59 | + md = {"dt_col": "tag_years"} |
| 60 | + if model_type == "weibull": |
| 61 | + md["weibull_log_shape_prior"] = [args.shape_loc, args.shape_scale] |
| 62 | + elif model_type == "lomax": |
| 63 | + md["lomax_log_theta_prior"] = [args.theta_loc, args.theta_scale] |
| 64 | + elif model_type == "constant_breakpoint": |
| 65 | + md["t_breakpoint_prior"] = [float(np.log(args.tb_median)), args.tb_scale] |
| 66 | + return md |
| 67 | + |
| 68 | + |
| 69 | +def _fit(model_type: str, df_fit: pd.DataFrame, metadata: dict, args): |
| 70 | + model = get_model_class(model_type)(dataset = df_fit, metadata = metadata) |
| 71 | + fitter = ModelFitter( |
| 72 | + event_rate_fun = model.event_rate_fun, |
| 73 | + starting_params = model.starting_params, |
| 74 | + data = model.data, target = model.target, |
| 75 | + num_warmup = args.n_warmup, num_samples = args.n_samples, |
| 76 | + num_chains = args.n_chains, |
| 77 | + param_likelihood = model.param_likelihood, |
| 78 | + derive_draws = model.derive_draws, |
| 79 | + log_likelihood_fun = model.log_likelihood_fun, |
| 80 | + rng_key = jax.random.PRNGKey(args.seed), |
| 81 | + verbose = True, |
| 82 | + ) |
| 83 | + fitter.fit() |
| 84 | + return model, fitter |
| 85 | + |
| 86 | + |
| 87 | +def _waic(model_type, fitter, df_eval, metadata, n_use): |
| 88 | + """In-sample WAIC on a bounded eval subset (deviance scale, lower better).""" |
| 89 | + model = get_model_class(model_type)(dataset = df_eval, metadata = metadata) |
| 90 | + n_avail = int(jax.tree_util.tree_leaves(fitter.param_draws)[0].shape[0]) |
| 91 | + idx = np.linspace(0, n_avail - 1, min(n_use, n_avail)).round().astype(int) |
| 92 | + draws = {k: v[idx] for k, v in fitter.param_draws.items()} |
| 93 | + target = model.target |
| 94 | + |
| 95 | + @jax.jit |
| 96 | + def ll_matrix(d): |
| 97 | + return jax.vmap( |
| 98 | + lambda p: model.pointwise_log_likelihood(p, model.data, target) |
| 99 | + )(d) |
| 100 | + |
| 101 | + ll = ll_matrix(draws) # (S, M) |
| 102 | + s = ll.shape[0] |
| 103 | + lppd = logsumexp(ll, axis = 0) - jnp.log(s) |
| 104 | + p_waic = jnp.var(ll, axis = 0, ddof = 1) |
| 105 | + return float(-2.0 * jnp.sum(lppd - p_waic)), int(len(target)) |
| 106 | + |
| 107 | + |
| 108 | +def main() -> None: |
| 109 | + ap = argparse.ArgumentParser() |
| 110 | + ap.add_argument("--model-type", required = True, |
| 111 | + choices = ["weibull", "lomax", "constant_breakpoint"]) |
| 112 | + ap.add_argument("--observations", required = True) |
| 113 | + ap.add_argument("--fit-sample", type = int, default = 600_000) |
| 114 | + ap.add_argument("--eval-sample", type = int, default = 500_000) |
| 115 | + ap.add_argument("--n-ppc-draws", type = int, default = 150) |
| 116 | + ap.add_argument("--n-warmup", type = int, default = 300) |
| 117 | + ap.add_argument("--n-samples", type = int, default = 300) |
| 118 | + ap.add_argument("--n-chains", type = int, default = 2) |
| 119 | + ap.add_argument("--shape-loc", type = float, default = 0.0) |
| 120 | + ap.add_argument("--shape-scale", type = float, default = 1.0) |
| 121 | + ap.add_argument("--theta-loc", type = float, default = 0.0) |
| 122 | + ap.add_argument("--theta-scale", type = float, default = 1.5) |
| 123 | + ap.add_argument("--tb-median", type = float, default = 4.0) |
| 124 | + ap.add_argument("--tb-scale", type = float, default = 0.5) |
| 125 | + ap.add_argument("--amenity", default = None, |
| 126 | + help = "Restrict to a single shared_label (e.g. 'School'); " |
| 127 | + "fits on the full filtered set without subsampling.") |
| 128 | + ap.add_argument("--seed", type = int, default = 0) |
| 129 | + args = ap.parse_args() |
| 130 | + |
| 131 | + _log(f"Loading observations: {args.observations}") |
| 132 | + df = pd.read_parquet(Path(args.observations).expanduser(), columns = LOAD_COLS) |
| 133 | + df = prepare_data_for_model(df) |
| 134 | + df = df[df[CELL_COLS].notna().all(axis = 1)].reset_index(drop = True) |
| 135 | + keep = CELL_COLS + ["id", "changed", "tag_years", "is_first_interval", |
| 136 | + "age_start", "age_end"] |
| 137 | + df = df[keep].copy() |
| 138 | + _log(f" modelled rows: {len(df):,}; change rate {df['changed'].mean():.4f}") |
| 139 | + |
| 140 | + if args.amenity: |
| 141 | + df = df[df["shared_label"] == args.amenity].reset_index(drop = True) |
| 142 | + _log(f" filtered to amenity '{args.amenity}': {len(df):,} rows; " |
| 143 | + f"change rate {df['changed'].mean():.4f}") |
| 144 | + # Single small amenity: fit + evaluate on the full set, no subsampling. |
| 145 | + df_fit = df_eval = df |
| 146 | + else: |
| 147 | + df_fit = df.sample(n = min(args.fit_sample, len(df)), |
| 148 | + random_state = args.seed).reset_index(drop = True) |
| 149 | + df_eval = df.sample(n = min(args.eval_sample, len(df)), |
| 150 | + random_state = args.seed + 1).reset_index(drop = True) |
| 151 | + |
| 152 | + md = _metadata(args.model_type, args) |
| 153 | + _log(f"Fitting '{args.model_type}' on {len(df_fit):,} rows; metadata={md}") |
| 154 | + model_full_meta = md |
| 155 | + model, fitter = _fit(args.model_type, df_fit, md, args) |
| 156 | + tab = fitter.get_parameter_table() |
| 157 | + _log(f"Fitted '{args.model_type}' parameters:") |
| 158 | + print(tab[tab["parameter"].isin(NATURAL)].to_string(index = False), |
| 159 | + flush = True) |
| 160 | + |
| 161 | + # Constant baseline on the same fit subsample, for WAIC reference. |
| 162 | + _log("Fitting 'constant' baseline on the same rows...") |
| 163 | + const_model, const_fitter = _fit("constant", df_fit, {"dt_col": "tag_years"}, |
| 164 | + args) |
| 165 | + |
| 166 | + waic_m, n_eval = _waic(args.model_type, fitter, df_eval, model_full_meta, |
| 167 | + args.n_ppc_draws) |
| 168 | + waic_c, _ = _waic("constant", const_fitter, df_eval, {"dt_col": "tag_years"}, |
| 169 | + args.n_ppc_draws) |
| 170 | + _log(f"=== WAIC on {n_eval:,} eval rows (deviance scale, lower=better) ===") |
| 171 | + print(f" {args.model_type:>20}: {waic_m:,.1f}", flush = True) |
| 172 | + print(f" {'constant':>20}: {waic_c:,.1f}", flush = True) |
| 173 | + print(f" Δ (model − constant): {waic_m - waic_c:,.1f} " |
| 174 | + f"({'better' if waic_m < waic_c else 'worse'})", flush = True) |
| 175 | + |
| 176 | + # Dispersion diagnostics on the full population for the fitted model. |
| 177 | + _log(f"Running dispersion diagnostics on full data ('{args.model_type}')...") |
| 178 | + model_full = get_model_class(args.model_type)(dataset = df, metadata = md) |
| 179 | + report = dispersion.dispersion_report( |
| 180 | + model_full, fitter, n_ppc_draws = args.n_ppc_draws, seed = args.seed, |
| 181 | + ) |
| 182 | + _log("=== SUMMARY ===") |
| 183 | + print(report["summary"].to_string(index = False), flush = True) |
| 184 | + _log("=== φ̂ by tag-age bin (target: a flatter gradient than 5.3→8.2→~106) ===") |
| 185 | + sub = report["subgroup"] |
| 186 | + print(sub[sub["grouping"] == "age_bin"].to_string(index = False), flush = True) |
| 187 | + |
| 188 | + # Same-handicap (no-covariate) constant baseline age gradient, so the |
| 189 | + # comparison isolates the hazard SHAPE rather than the missing covariates. |
| 190 | + _log("Running dispersion diagnostics on full data ('constant' baseline)...") |
| 191 | + const_full = get_model_class("constant")( |
| 192 | + dataset = df, metadata = {"dt_col": "tag_years"}) |
| 193 | + report_c = dispersion.dispersion_report( |
| 194 | + const_full, const_fitter, n_ppc_draws = args.n_ppc_draws, seed = args.seed, |
| 195 | + ) |
| 196 | + _log("=== CONSTANT baseline — φ̂ by tag-age bin (same rows) ===") |
| 197 | + subc = report_c["subgroup"] |
| 198 | + print(subc[subc["grouping"] == "age_bin"].to_string(index = False), |
| 199 | + flush = True) |
| 200 | + phi_m = report["summary"].set_index("statistic").loc[ |
| 201 | + "covariate_cell", "phi_hat"] |
| 202 | + phi_c = report_c["summary"].set_index("statistic").loc[ |
| 203 | + "covariate_cell", "phi_hat"] |
| 204 | + _log(f"covariate_cell φ̂ — {args.model_type}: {phi_m:.2f} vs " |
| 205 | + f"constant: {phi_c:.2f}") |
| 206 | + |
| 207 | + out = Path(args.observations).expanduser().parent |
| 208 | + report["summary"].to_csv( |
| 209 | + out / f"hazard_{args.model_type}_dispersion_summary.csv", index = False) |
| 210 | + report["subgroup"].to_csv( |
| 211 | + out / f"hazard_{args.model_type}_dispersion_subgroup.csv", index = False) |
| 212 | + _log(f"Wrote hazard_{args.model_type}_dispersion_*.csv to {out}") |
| 213 | + |
| 214 | + |
| 215 | +if __name__ == "__main__": |
| 216 | + main() |
0 commit comments