Merge pull request #21 from henryspatialanalysis/feature/jax

Rebuild POI change models in JAX

Author: njhenry

.claude/CLAUDE.md

@@ -39,7 +39,7 @@ Style: Black (format-on-save in VSCode). Lint: flake8 + pylint, configured in `p
 |---|---|
 | [src/openpois/io/](../src/openpois/io/) | I/O adapters: OSM history/snapshot, Overture, Foursquare, Census boundary |
 | [src/openpois/osm/](../src/openpois/osm/) | OSM-specific transforms: `format_observations`, `change_plots` |
-| [src/openpois/models/](../src/openpois/models/) | PyTorch empirical Bayes: `EventRate`, `ModelFitter`, model registry |
+| [src/openpois/models/](../src/openpois/models/) | JAX/BlackJAX empirical Bayes: `ModelFitter`, model registry |
 | [src/openpois/conflation/](../src/openpois/conflation/) | OSM×Overture matching: `taxonomy`, `match`, `merge` |
 | [scripts/](../scripts/) | End-to-end pipelines using config.yaml — not installed, reference only |
 | [site/](../site/) | Vue 3 + Vite frontend |
@@ -49,6 +49,7 @@ Style: Black (format-on-save in VSCode). Lint: flake8 + pylint, configured in `p
 - [docs/data-sources.md](docs/data-sources.md) — URLs, auth, schema quirks for every source
 - [docs/taxonomy-setup.md](docs/taxonomy-setup.md) — crosswalk CSVs, build_taxonomy.py, frontend sync
 - [docs/data-versioning.md](docs/data-versioning.md) — `versions:` block, path resolution, external references
+- [docs/turnover-model-methodology.md](docs/turnover-model-methodology.md) — statistical derivation of the POI turnover model with ZIE extension
 
 ## Running to-do
 

.claude/docs/turnover-model-methodology.md

@@ -9,20 +9,21 @@ Rerun just the modeling step (step 3 of the full pipeline) for a fixed source-da
 
 ## Prerequisites
 
-- `osm_observations_{tag_key}.csv` already generated by `scripts/osm_data/format_tabular.py` at the pinned `versions.osm_data`.
+- `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
 
 1. **Pin `versions.osm_data`** in `config.yaml`. Do *not* change it — that's the whole point.
 
 2. **Bump `versions.model_output`** using the convention:
-   - `{date}_by_{group_key}` for grouped models (e.g., `20260416_by_leisure`, `20260416_by_amenity`)
+   - `{date}_by_shared_label` for the unified random-effects model (one intercept per taxonomy category)
+   - `{date}_by_{group_key}` for ad-hoc groupings on other columns
    - `{date}_constant` for the single-rate baseline
 
 3. **Edit `osm_turnover_model`** in `config.yaml`:
-   - `model_type`: one of `constant`, `random_by_type`, `pseudo_varying` (registry at [src/openpois/models/osm_models.py](../../../src/openpois/models/osm_models.py))
-   - `group_key`: column to group by (e.g., `leisure_last_value`, `shop`, `amenity`). Null for constant.
+   - `model_type`: one of `constant`, `random_by_type` (registry at [src/openpois/models/osm_models.py](../../../src/openpois/models/osm_models.py))
+   - `group_key`: column to group by. Default `shared_label` (unified taxonomy). Null for constant. Other observation columns (raw OSM keys like `shop`, `amenity`, ...) are still accepted.
    - `group_values`: restrict to specific values, or null for all
    - `min_value_count`: drop groups below this count
 
@@ -35,7 +36,7 @@ Rerun just the modeling step (step 3 of the full pipeline) for a fixed source-da
    ```bash
    python scripts/osm_snapshot/apply_model.py
    ```
-   `apply_model.py` picks up every `{stub}_by_*` directory at the stub date and falls back to `{stub}_constant` for unmatched groups.
+   `apply_model.py` loads a single `{stub}_by_shared_label` random-effects model (if present) and falls back to `{stub}_constant` for rows with no matching taxonomy label.
 
 ## Comparing variants
 
@@ -46,5 +47,5 @@ Rerun just the modeling step (step 3 of the full pipeline) for a fixed source-da
 ## Pitfalls
 
 - Forgetting to change `versions.model_output` overwrites the previous variant's outputs.
-- `group_key` must exist as a column in the observations CSV; run `format_tabular.py` first if adding a new tag key to `osm_data.tag_key`.
+- `group_key` must exist as a column in the observations CSV. `shared_label` is populated by `format_tabular.py` from the conflation taxonomy crosswalk; if you change the crosswalk, rerun `format_tabular.py`.
 - `min_value_count` filters groups silently — check `fitted_params.csv` row count vs. expected group count.

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

@@ -21,23 +21,23 @@ 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_{tag_key}.csv`
+2. **Format tabular observations** → `osm_observations.parquet`
    ```bash
    python scripts/osm_data/format_tabular.py
    ```
-   Uses `osm_data.tag_key` (e.g., `name`) to produce one observation row per version with change/deletion flags.
+   Uses `osm_data.tag_key` (e.g., `name`) to flag change/deletion per POI version, then assigns shared taxonomy labels from the conflation crosswalk and explodes rows per label. One row = (POI version, shared_label). Rows with no matching taxonomy category are dropped.
 
 3. **Pick a modeling config and fit λ** — see [skills/iterate-model-types](../iterate-model-types/SKILL.md) for choosing `model_type` / `group_key`.
    ```bash
    python scripts/models/osm_turnover.py
    ```
-   Writes `fitted_params.csv`, `param_draws.csv`, `predictions.csv`, `fitted_model.pt` to `{date}_by_{group_key}` (or `{date}_constant`) under `directories.model_output.path`.
+   Writes `fitted_params.csv`, `param_draws.csv`, `predictions.csv` to `{date}_by_shared_label` (the unified random-effects model) or `{date}_constant` (single-rate baseline) under `directories.model_output.path`.
 
 4. **Apply predictions to the OSM snapshot** → `osm_snapshot_rated.parquet`
    ```bash
    python scripts/osm_snapshot/apply_model.py
    ```
-   Reads the `osm_data.apply_model.model_stub` date, loads all `{stub}_by_*` dirs (plus a `{stub}_constant` fallback), and rates every POI in `osm_snapshot.parquet`.
+   Reads the `osm_data.apply_model.model_stub` date, loads the `{stub}_by_shared_label` random-effects model (if present), falls back to `{stub}_constant` for rows with no matching taxonomy label, and rates every POI in `osm_snapshot.parquet`.
 
 ## Verification
 
@@ -51,5 +51,5 @@ Hand off to [skills/verify-pipeline-run](../verify-pipeline-run/SKILL.md) — in
 
 - Entry: [src/openpois/io/osm_history_pbf.py](../../../src/openpois/io/osm_history_pbf.py) (`download_osm_history`)
 - Entry: [src/openpois/osm/format_observations.py](../../../src/openpois/osm/format_observations.py)
-- Entry: [src/openpois/models/](../../../src/openpois/models/) — `ModelFitter`, `EventRate`, `pytorch_setup`
+- Entry: [src/openpois/models/](../../../src/openpois/models/) — `ModelFitter` (JAX/BlackJAX), model classes
 - Registry: [src/openpois/models/osm_models.py](../../../src/openpois/models/osm_models.py) — `MODEL_REGISTRY`, `get_model_class`

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

@@ -32,7 +32,6 @@ Flag >5% drops. Known regression patterns:
   fitted_params.csv     # λ and σ per group
   param_draws.csv       # uncertainty bounds
   predictions.csv       # predictions per POI
-  fitted_model.pt       # torch state
 ```
 
 Checks:

.claude/skills/verify-pipeline-run/SKILL.md

@@ -70,7 +70,7 @@ The core workflow models how long POI tags remain stable over time using histori
 ```bash
 python exploratory/osm_data/download.py       # Download OSM history for a bounding box
 python exploratory/osm_data/format_tabular.py # Format into observation records
-python exploratory/models/pytorch_simple.py   # Fit Poisson change-rate model
+python scripts/models/osm_turnover.py         # Fit Poisson change-rate model (JAX)
 ```
 
 ---

README.md

@@ -1,12 +1,12 @@
 # Versioned directories (used with config.get_dir_path())
 versions:
   osm_data: "20260416"
-  model_output: "20260416_by_leisure"
+  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,17 +126,29 @@ osm_data:
     - last_obs_timestamp
     - last_tag_timestamp
   apply_model:
-    model_stub: '20260416'
+    model_stub: '20260422'
 
-# Settings for exploratory/models/pytorch_simple.py
+# Settings for scripts/models/osm_turnover.py (JAX turnover model)
 osm_turnover_model:
-  model_type: random_by_type
-  tag_key: name
+  # Overridable at the CLI via --model-type {constant,random_by_type}.
+  default_model_type: constant
   var_prior: [-1.0, 5.0]
-  group_key: leisure_last_value
+  # 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
   min_value_count: 5
-  n_draws: 250
+  # NUTS warmup (window adaptation) and retained-sample counts. Warmup should
+  # generally be >= n_samples for hierarchical models.
+  n_warmup: 500
+  n_samples: 500
+  # Number of independent chains (vmapped in parallel). n_chains > 1 enables
+  # R-hat and bulk ESS diagnostics at roughly linear wall-time cost on CPU.
+  n_chains: 4
   save_full_model: true
 
 # Directory definitions (used with config.get_dir_path())
@@ -161,14 +173,17 @@ directories:
       # Legacy Overpass-based pipeline (still used by scripts/osm_data/download.py)
       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.parquet
   model_output:
     versioned: true
     path: ~/data/openpois/osm_turnover_model
     files:
       fitted_params: fitted_params.csv
       param_draws: param_draws.csv
       predictions: predictions.csv
-      fitted_model: fitted_model.pt
+      diagnostics: diagnostics.csv
+      inference_data: inference_data.nc
   snapshot_foursquare:
     versioned: true
     path: ~/data/openpois/snapshots/foursquare
@@ -247,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:

config.yaml

@@ -136,38 +136,48 @@ AWS credentials via environment variables or ``~/.aws/credentials``.
 models
 ------
 
-openpois.models.event_rate
-~~~~~~~~~~~~~~~~~~~~~~~~~~
+openpois.models.jax_core
+~~~~~~~~~~~~~~~~~~~~~~~~
 
-Representation of a Poisson event rate (λ) used by the change-rate model.
-Wraps a constant or time-varying λ tensor and computes the probability that
-at least one change event occurs within a given time interval via numerical
-or closed-form integration.
+JAX/BlackJAX helpers: a PRNG factory, a jitted Markov-chain scan, a NUTS
+sampler with window adaptation, and a vmap-based predictive-draw utility.
 
-.. automodule:: openpois.models.event_rate
+.. automodule:: openpois.models.jax_core
    :members:
    :undoc-members:
    :show-inheritance:
 
 openpois.models.model_fitter
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-L-BFGS optimizer wrapper for POI change-rate models. Fits model parameters
-using PyTorch and the ``torchmin`` optimizer, generates posterior parameter
-draws for uncertainty quantification, and produces prediction tables of
-change probability versus time.
+BlackJAX NUTS fitter for POI change-rate models. Takes an ``event_rate_fun``
+plus starting parameters as a pytree, runs window-adapted NUTS to draw from
+the posterior, and produces posterior summaries and predictive distributions
+of change probability versus time.
 
 .. automodule:: openpois.models.model_fitter
    :members:
    :undoc-members:
    :show-inheritance:
 
+openpois.models.osm_models
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+JAX model classes for OSM turnover. ``ConstantModel`` and
+``RandomByTypeModel`` package their own data, priors, and event-rate
+functions to hand to ``ModelFitter``. Selectable via ``get_model_class``.
+
+.. automodule:: openpois.models.osm_models
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 openpois.models.setup
 ~~~~~~~~~~~~~~~~~~~~~
 
-Environment setup utilities for PyTorch model runs. Selects GPU or CPU
-device, configures ``torch_continuum`` optimisation level, and prepares
-filtered and grouped observation data for model fitting.
+Data-preparation helpers. ``prepare_data_for_model`` filters and groups
+observation records and computes the ``tag_years`` elapsed column used as
+the per-observation interval length in fitting.
 
 .. automodule:: openpois.models.setup
    :members:

docs/api.rst

@@ -28,7 +28,9 @@
 # Mock heavy imports so Sphinx can parse docstrings without installing the
 # full conda environment in CI.
 autodoc_mock_imports = [
-    "torch",
+    "jax",
+    "jaxlib",
+    "blackjax",
     "numpy",
     "pandas",
     "geopandas",

docs/conf.py

@@ -92,8 +92,9 @@ See :mod:`openpois.io.osm_history`.
    python scripts/osm_data/format_tabular.py
 
 Converts raw version histories into one-row-per-observation records, each
-flagged for whether the configured tag changed. Output:
-``osm_observations_{tag_key}.csv``.
+flagged for whether the configured ``osm_data.tag_key`` changed, then
+assigns a shared taxonomy label and explodes rows for POIs mapping to
+multiple labels. Output: ``osm_observations.csv``.
 
 See :mod:`openpois.osm.format_observations`.
 
@@ -103,13 +104,13 @@ See :mod:`openpois.osm.format_observations`.
 
    python scripts/models/osm_turnover.py
 
-Fits an empirical Bayes PyTorch model (constant or random-effects by type)
-estimating the Poisson change rate λ per group. Outputs ``fitted_params.csv``
-and ``predictions.csv`` (and optionally ``param_draws.csv`` /
-``fitted_model.pt``).
+Fits an empirical Bayes JAX model (constant or random-effects by type)
+estimating the Poisson change rate λ per group via BlackJAX NUTS. Outputs
+``fitted_params.csv`` and ``predictions.csv`` (and optionally
+``param_draws.csv``).
 
-See :mod:`openpois.models.model_fitter`, :mod:`openpois.models.setup`, and
-:mod:`openpois.models.event_rate`.
+See :mod:`openpois.models.model_fitter`, :mod:`openpois.models.osm_models`,
+and :mod:`openpois.models.setup`.
 
 **Step 4 — Visualise stability curves** *(optional)*