Work Location Choice

A work-location model evaluates possible destination zones for each person. The simplest specification can use commute time alone, while a richer destination utility can include a tour-mode logsum that summarizes the accessibility of all available travel modes. That richer calculation is informative but expensive when repeated for every person-zone pair.

This walkthrough prepares the MTC mini data and compares three implementations: a time-only model, a full model that computes mode-choice logsums for every zone, and a sample-of-alternatives model that computes those logsums only for sampled destinations.

Set up Traveler and the example data

We enable informational logging so that the chunked model steps report their progress. Calling max_num_threads() without an argument reports the current thread limit without changing it.

import logging

import jax.numpy as jnp

import mtc
import traveler as tv

tv.log_to_stdout(logging.INFO)
tv.max_num_threads()

[00:00.99] INFO: Logging to stdout set to level 20

The mini loader returns related household, person, tour, land-use, and skim tables. We then select 14,250 households and their associated people and tours. SampleHouseholds returns a new store and rebuilds the relationship indexes, so the sampled tables remain aligned.

from traveler.steps import SampleHouseholds

store = mtc.mini()
store = store.run_steps(SampleHouseholds(n=14_250))

/Users/jpn/Git/traveler/src/traveler/steps/_sample_households.py:262: UserWarning: n (14250) is greater than the number of households in the store, returning all households.
  return self.run_via_polars(store)

The store summary confirms the dimensions that drive the computation: people are the choosers, zones are the alternatives, and the skim matrices provide travel attributes between origins and destinations.

store.info()

<mtc.tables.store.Store>
  households = <JaxTable with shape (5000,)>
  persons = <JaxTable with shape (8212,)>
  skims = <Skims with 21 keys>
  land_use = <JaxTable with shape (25,)>
  time_periods = {'EA': 0, 'AM': 1, 'MD': 2, 'PM': 3, 'EV': 4}
  tours = <JaxTable with shape (20000,)>

Derive the model inputs

Work-location utility depends on fields that are not present in the raw input tables. The MTC initialize compound step runs the land-use, person, household, and value-of-time annotations in dependency order and returns a store containing their results.

from mtc.models import initialize

store = store.run_steps(initialize)
store.completed_steps

[00:02.10] INFO: compute_values step annotate_landuse completed in 0:00:00.06

[00:02.42] INFO: compute_values step annotate_persons completed in 0:00:00.32

[00:02.58] INFO: compute_values step annotate_households completed in 0:00:00.16

[00:02.64] INFO: compute_values step household_value_of_time completed in 0:00:00.06

frozenset({'annotate_households',
           'annotate_landuse',
           'annotate_persons',
           'household_value_of_time'})

The initialized store has the same table relationships as before, now augmented with the derived fields needed by the destination and lower-level mode-choice utilities.

store.info()

<mtc.tables.store.Store>
  households = <JaxTable with shape (5000,)>
  persons = <JaxTable with shape (8212,)>
  skims = <Skims with 21 keys>
  land_use = <JaxTable with shape (25,)>
  time_periods = {'EA': 0, 'AM': 1, 'MD': 2, 'PM': 3, 'EV': 4}
  tours = <JaxTable with shape (20000,)>

Configure the three specifications

All three models use the same round-trip time coefficient. The two logsum specifications also receive the tour-mode coefficients and a mu parameter that scales the mode-choice logsum in destination utility. Sharing these values isolates the computational strategy as the important difference among the runs.

from mtc.models.tour_mode import tour_mode
from mtc.models.work_location import (
    work_location,
    work_location_no_logsums,
    work_location_sampled,
)

parameters = {"time": -0.5, "mu": 0.59, "tour_mode": tour_mode.param}
work_location_no_logsums.param = parameters.copy()
work_location.param = parameters.copy()
work_location_sampled.param = parameters.copy()

Run the time-only baseline

The baseline evaluates all destination zones but omits mode-choice logsums. It is useful both as a fast model and as a reference for seeing how the accessibility term changes choices.

no_logsums_store = store.run_steps(work_location_no_logsums)

Run full logsum enumeration

The full specification evaluates the lower-level tour-mode model for every person and every destination. This is the direct formulation of the model and provides the comparison point for the sampled approximation.

full_logsum_store = store.run_steps(work_location)

Run the sampled alternative model

The sampled model first uses the inexpensive time-only utility to draw 30 destination slots per person. It computes tour-mode logsums only for those slots, applies the sampling-probability correction, and then simulates the final choice. This preserves the richer utility while avoiding a logsum calculation for every zone.

sampled_logsum_store = store.run_steps(work_location_sampled)

Inspect and compare the choices

Each step writes its selected TAZ to the returned persons table as work_location_taz. The categorical column displays zone labels while retaining compact integer codes for model calculations.

sampled_logsum_store.persons.to_polars().select(
    "person_id", "home_zone_id", "work_location_taz"
).head(10)

shape: (10, 3)

person_id	home_zone_id	work_location_taz
i32	enum	enum
25671	"TAZ-5"	"TAZ-5"
25675	"TAZ-5"	"TAZ-5"
25678	"TAZ-6"	"TAZ-6"
25683	"TAZ-6"	"TAZ-6"
25684	"TAZ-6"	"TAZ-6"
25691	"TAZ-6"	"TAZ-6"
25692	"TAZ-6"	"TAZ-6"
25704	"TAZ-6"	"TAZ-6"
25725	"TAZ-6"	"TAZ-6"
25734	"TAZ-6"	"TAZ-6"

The alternatives are simulated rather than selected by a deterministic maximum, so exact row-by-row agreement is not the goal. Still, comparing choices from the same people and stable random streams gives a quick diagnostic of how much the logsum term and destination sampling alter outcomes in this run.

{
    "time only vs. full logsums": jnp.mean(
        no_logsums_store.persons["work_location_taz"].codes
        == full_logsum_store.persons["work_location_taz"].codes
    ),
    "full logsums vs. sampled logsums": jnp.mean(
        full_logsum_store.persons["work_location_taz"].codes
        == sampled_logsum_store.persons["work_location_taz"].codes
    ),
}

{'time only vs. full logsums': Array(0.30017048, dtype=float32),
 'full logsums vs. sampled logsums': Array(0.99318075, dtype=float32)}

The three stores make the modeling tradeoff explicit. Omitting logsums is cheapest but changes the behavioral specification; full enumeration is the clearest reference implementation; sampling keeps the logsum specification while reducing its most expensive dimension. A production comparison should examine aggregate destination distributions and runtime over repeated runs, not rely only on the illustrative row-level agreement above.