Successful run of JAX model by shared_label with a zero-inflation term.

Author: njhenry

.claude/skills/iterate-model-types/SKILL.md

@@ -9,7 +9,7 @@ Rerun just the modeling step (step 3 of the full pipeline) for a fixed source-da
 
 ## Prerequisites
 
-- `osm_observations.csv` already generated by `scripts/osm_data/format_tabular.py` at the pinned `versions.osm_data`. Each row is one (POI version, shared_label) pair — POIs mapping to multiple taxonomy categories are exploded into multiple rows.
+- `osm_observations.parquet` already generated by `scripts/osm_data/format_tabular.py` at the pinned `versions.osm_data`. Each row is one (POI version, shared_label) pair — POIs mapping to multiple taxonomy categories are exploded into multiple rows.
 - You know which model variant you want next.
 
 ## Steps

.claude/skills/model-history-pipeline/SKILL.md

@@ -21,7 +21,7 @@ End-to-end: Geofabrik full-history PBFs → observations table → fitted λ →
    ```
    Runs `osmium tags-filter --omit-referenced` then `osmium time-filter`, then pyosmium streams results. Controlled by `download.osm.*` in config.yaml.
 
-2. **Format tabular observations** → `osm_observations.csv`
+2. **Format tabular observations** → `osm_observations.parquet`
    ```bash
    python scripts/osm_data/format_tabular.py
    ```

config.yaml

@@ -1,12 +1,12 @@
 # Versioned directories (used with config.get_dir_path())
 versions:
   osm_data: "20260416"
-  model_output: "20260421_by_shared_label"
+  model_output: "20260422_by_shared_label"
   snapshot_osm: "20260417"
   snapshot_overture: "20260417"
   snapshot_foursquare: "20260416"
-  aws: "20260417"
-  conflation: "20260417"
+  aws: "20260422"
+  conflation: "20260422"
 
 # Settings for downloading data
 download:
@@ -126,14 +126,18 @@ osm_data:
     - last_obs_timestamp
     - last_tag_timestamp
   apply_model:
-    model_stub: '20260421'
+    model_stub: '20260422'
 
 # Settings for scripts/models/osm_turnover.py (JAX turnover model)
 osm_turnover_model:
   # Overridable at the CLI via --model-type {constant,random_by_type}.
   default_model_type: constant
   var_prior: [-1.0, 5.0]
-  # Column in osm_observations.csv for grouping random effects.
+  # Tight hyperprior on log_tau (random-effect scale for per-group logit_delta
+  # in RandomByTypeModel). Tau median ≈ exp(-2) ≈ 0.135 on the logit scale —
+  # shrinks per-group δ toward the global intercept.
+  logit_delta_var_prior: [-2.0, 0.5]
+  # Column in osm_observations.parquet for grouping random effects.
   # "shared_label" = shared taxonomy category
   group_key: shared_label
   group_values: null
@@ -170,7 +174,7 @@ directories:
       osm_elements: osm_elements.csv
       osm_failed_elements: osm_failed_elements.csv
       # Modelling-ready observations (one row per POI version × shared_label)
-      osm_observations: osm_observations.csv
+      osm_observations: osm_observations.parquet
   model_output:
     versioned: true
     path: ~/data/openpois/osm_turnover_model
@@ -258,7 +262,7 @@ upload:
   s3_prefix_osm: "snapshots/osm"
   s3_prefix_conflation: "snapshots/conflated"
   latest_url_osm: "https://openpois-public.s3.us-west-2.amazonaws.com/snapshots/osm/20260417/osm_snapshot_partitioned/"
-  latest_url_conflation: "https://openpois-public.s3.us-west-2.amazonaws.com/snapshots/conflated/20260417/conflated_partitioned/"
+  latest_url_conflation: "https://openpois-public.s3.us-west-2.amazonaws.com/snapshots/conflated/20260422/conflated_partitioned/"
   geohash_precision_partition: 4  # ~39 km x 20 km cells; ~1,000–3,000 cells over CONUS
   geohash_precision_sort: 6       # ~0.6 km x 1.2 km; fine-grained sort within each partition
   pmtiles:

scripts/exploratory/bench_random_by_type.py

@@ -24,7 +24,7 @@
     small  — n = 10 000, K = 20
     medium — n = 1 000 000, K = 91
     large  — n = 4 200 000, K = 91  (slow; ~matches production scale)
-    real   — reads real osm_observations.csv via config.yaml
+    real   — reads real osm_observations.parquet via config.yaml
 """
 
 from __future__ import annotations
@@ -117,7 +117,7 @@ def _load_real_observations() -> pd.DataFrame:
     group_key = cfg.get(
         "osm_turnover_model", "group_key", fail_if_none = False
     )
-    df = pd.read_csv(path)
+    df = pd.read_parquet(path)
     prepared = prepare_data_for_model(
         data = df,
         group_key = group_key,
@@ -199,6 +199,7 @@ def _run_one(
         param_likelihood = model.param_likelihood,
         derive_draws = model.derive_draws,
         log_likelihood_fun = model.log_likelihood_fun,
+        log_1md_fun = getattr(model, "log_1md_fun", None),
         verbose = False,
     )
 

scripts/models/osm_turnover.py

@@ -1,7 +1,7 @@
 """
 Fit an empirical Bayes JAX model for OSM POI tag change rates.
 
-Reads ``osm_observations.csv`` (produced by ``osm_data/format_tabular.py``,
+Reads ``osm_observations.parquet`` (produced by ``osm_data/format_tabular.py``,
 one row per (POI version, shared_label)) and fits a Poisson change-rate
 model using BlackJAX NUTS. The model estimates a per-group change rate λ
 (events per year). Predictions give the probability that a tag changes
@@ -24,6 +24,10 @@
     osm_turnover_model.default_model_type   — "constant" or "random_by_type"
                                               (overridable via --model-type)
     osm_turnover_model.var_prior            — (loc, scale) hyperprior on log_sigma
+    osm_turnover_model.logit_delta_prior    — (loc, scale) prior on logit_delta_0
+                                              intercept
+    osm_turnover_model.logit_delta_var_prior — (loc, scale) tight hyperprior on
+                                              log_tau (per-group δ scale)
     osm_turnover_model.n_warmup             — NUTS warmup steps (adaptation)
     osm_turnover_model.n_samples            — posterior draws retained
     osm_turnover_model.n_chains             — number of NUTS chains (vmapped)
@@ -120,13 +124,15 @@ def flatten_param_draws(
     config.write_self("model_output")
 
     # Data preparation ------------------------------------------------------>
-    observations_df = pd.read_csv(OBSERVATIONS_PATH)
+    observations_df = pd.read_parquet(OBSERVATIONS_PATH)
+    # t1_col defaults to "last_obs_timestamp" in prepare_data_for_model, so
+    # tag_years is the inter-observation interval the per-row Bernoulli-on-
+    # Poisson likelihood requires (methodology §1.2).
     obs_sub = prepare_data_for_model(
         data = observations_df,
         group_key = GROUP_KEY,
         group_values = GROUP_VALUES,
         min_value_count = MIN_VALUE_COUNT,
-        t1_col = "last_tag_timestamp",
         t2_col = "obs_timestamp",
     )
 
@@ -135,15 +141,26 @@ def flatten_param_draws(
         "osm_turnover_model", "default_model_type"
     )
     print(f"Model type: {model_type}")
+    metadata = {
+        "dt_col": "tag_years",
+        "group": GROUP_KEY,
+        "var_prior": tuple(
+            config.get("osm_turnover_model", "var_prior")
+        ),
+    }
+    logit_delta_prior = config.get(
+        "osm_turnover_model", "logit_delta_prior", fail_if_none = False
+    )
+    if logit_delta_prior is not None:
+        metadata["logit_delta_prior"] = tuple(logit_delta_prior)
+    logit_delta_var_prior = config.get(
+        "osm_turnover_model", "logit_delta_var_prior", fail_if_none = False
+    )
+    if logit_delta_var_prior is not None:
+        metadata["logit_delta_var_prior"] = tuple(logit_delta_var_prior)
     model = get_model_class(model_type)(
         dataset = obs_sub,
-        metadata = {
-            "dt_col": "tag_years",
-            "group": GROUP_KEY,
-            "var_prior": tuple(
-                config.get("osm_turnover_model", "var_prior")
-            ),
-        },
+        metadata = metadata,
     )
 
     fitter = ModelFitter(
@@ -157,6 +174,7 @@ def flatten_param_draws(
         param_likelihood = model.param_likelihood,
         derive_draws = model.derive_draws,
         log_likelihood_fun = model.log_likelihood_fun,
+        log_1md_fun = getattr(model, "log_1md_fun", None),
         verbose = True,
     )
     fitter.fit()
@@ -172,10 +190,23 @@ def flatten_param_draws(
         )
 
     # Predictions ----------------------------------------------------------->
+    # Emit both regimes (methodology §4.2 Step G): the conditional formula
+    # populates p_mean/p_lower/p_upper (δ-independent, right for rating
+    # already-observed POIs); the fresh formula populates p_fresh_* (uses δ,
+    # right for rating a hypothetical freshly tagged POI).
     predict_times = jnp.arange(101) / 10.0
     predict_data = model.build_predict_data(predict_times)
+    conditional = fitter.predict(data = predict_data, mode = "conditional")
+    fresh = (
+        fitter.predict(data = predict_data, mode = "fresh")
+        .rename(columns = {
+            "p_mean": "p_fresh_mean",
+            "p_lower": "p_fresh_lower",
+            "p_upper": "p_fresh_upper",
+        })
+    )
     predictions = (
-        fitter.predict(data = predict_data)
+        pd.concat([conditional, fresh], axis = 1)
         .assign(t1 = 0.0, units = "years")
     )
     predictions["t2"] = np.asarray(predict_data["dt"])

scripts/osm_data/data_viz.py

@@ -1,7 +1,7 @@
 """
 Plot OSM tag stability curves from observation data.
 
-Reads ``osm_observations.csv`` (one row per (POI version, shared_label)) and
+Reads ``osm_observations.parquet`` (one row per (POI version, shared_label)) and
 computes Kaplan-Meier-style survival estimates showing what fraction of
 observations remain unchanged over time. Saves two types of PNG figures:
 
@@ -12,7 +12,7 @@
        count, shown as separate facets.
 
 Config keys used (config.yaml):
-    directories.osm_data           — directory containing input CSV and viz/ output
+    directories.osm_data           — directory containing input parquet and viz/ output
     osm_data.tag_key               — tag key whose changes define observation
                                      events (used only in plot titles)
     osm_data.timestamp_cols        — columns to parse as timestamps (rows with nulls dropped)
@@ -112,9 +112,9 @@ def get_preds_df(version: str) -> pd.DataFrame | None:
             return None
         return pd.read_csv(preds_fp).assign(
             year = pd.col('t2'),
-            conf_mean = (1.0 - pd.col('p_mean')),
-            conf_lower = (1.0 - pd.col('p_upper')),
-            conf_upper = (1.0 - pd.col('p_lower')),
+            conf_mean = (1.0 - pd.col('p_fresh_mean')),
+            conf_lower = (1.0 - pd.col('p_fresh_upper')),
+            conf_upper = (1.0 - pd.col('p_fresh_lower')),
         )
 
     return {
@@ -135,7 +135,7 @@ def get_preds_df(version: str) -> pd.DataFrame | None:
     #   observation timestamp will be missing for these rows
     timestamp_cols = config.get("osm_data", "timestamp_cols")
     observations_df = (
-        pd.read_csv(OBSERVATIONS_PATH)
+        pd.read_parquet(OBSERVATIONS_PATH)
         .dropna(subset = timestamp_cols)
     )
     for timestamp_col in timestamp_cols:

scripts/osm_data/download_history.py

@@ -111,3 +111,21 @@
         chunk_size = CHUNK_SIZE,
         verbose = VERBOSE,
     )
+
+    # -------------------------------------------------------------------------
+    # Clean up intermediates
+    # -------------------------------------------------------------------------
+    finals_ok = all(
+        p.exists() and p.stat().st_size > 0
+        for p in (OUTPUT_VERSIONS, OUTPUT_CHANGES)
+    )
+    if finals_ok:
+        intermediates = (
+            RAW_PBF, FILTERED_PBF, TIME_FILTERED_PBF,
+            RAW_PR_PBF, FILTERED_PR_PBF, TIME_FILTERED_PR_PBF,
+            US_VERSIONS, US_CHANGES, PR_VERSIONS, PR_CHANGES,
+        )
+        for p in intermediates:
+            if p.exists():
+                print(f"Removing intermediate {p} ...")
+                p.unlink()

scripts/osm_data/format_tabular.py

@@ -32,7 +32,7 @@
     osm_changes.parquet.
 
 Output file (in osm_data directory):
-    osm_observations.csv — one row per (POI version, shared_label). Columns:
+    osm_observations.parquet — one row per (POI version, shared_label). Columns:
         id, version, tag_key, last_tag_timestamp, obs_timestamp, changed,
         shared_label, plus every filter_keys column for reference.
 """
@@ -86,7 +86,7 @@
     print(f"DuckDB wrote {n_written:,} raw observations to {OUT_PATH}")
 
     print("Loading raw observations for shared-label assignment ...")
-    obs_df = pd.read_csv(OUT_PATH, dtype = {k: str for k in OSM_KEYS})
+    obs_df = pd.read_parquet(OUT_PATH)
 
     print("Assigning shared taxonomy labels (multi-label, exploded) ...")
     labels_per_row, _ = assign_osm_shared_label(
@@ -109,5 +109,5 @@
     print("Top shared labels by row count:")
     print(obs_df["shared_label"].value_counts().head(15).to_string())
 
-    obs_df.to_csv(OUT_PATH, index = False)
+    obs_df.to_parquet(OUT_PATH, index = False)
     print(f"Saved {len(obs_df):,} observations to {OUT_PATH}")

scripts/osm_snapshot/download.py

@@ -87,3 +87,12 @@
         chunk_dir = CHUNK_DIR,
         verbose = VERBOSE,
     )
+
+    # -------------------------------------------------------------------------
+    # Clean up intermediates
+    # -------------------------------------------------------------------------
+    if OUTPUT_PATH.exists() and OUTPUT_PATH.stat().st_size > 0:
+        for p in (RAW_PBF, FILTERED_PBF, RAW_PR_PBF, FILTERED_PR_PBF):
+            if p.exists():
+                print(f"Removing intermediate {p} ...")
+                p.unlink()

scripts/overture/download.py

@@ -30,6 +30,8 @@
         Columns: overture_id, overture_name, taxonomy_l0, taxonomy_l1,
         taxonomy_l2, brand_name, confidence, geometry, source
 """
+import shutil
+
 import pyarrow.parquet as pq
 from config_versioned import Config
 from openpois.io.boundary import get_us_pr_boundary
@@ -91,3 +93,12 @@
     )
     n_rows = pq.read_metadata(output_path).num_rows
     print(f"Saved {n_rows:,} Overture POIs to {output_path}")
+
+    # -------------------------------------------------------------------------
+    # Clean up intermediates
+    # -------------------------------------------------------------------------
+    if OUTPUT_PATH.exists() and OUTPUT_PATH.stat().st_size > 0:
+        parts_dir = SAVE_DIR / ".parts"
+        if parts_dir.is_dir():
+            print(f"Removing intermediate {parts_dir} ...")
+            shutil.rmtree(parts_dir)