Steps#

A Traveler simulation is an ordered series of operations on a Store. Each operation reads the current model state, computes something, and returns the state that the next operation should receive. Traveler calls one such operation a step.

A step might:

  • prepare or sample input records;

  • derive new columns from existing data;

  • simulate a choice and store the result;

  • update metadata about the model run; or

  • group several smaller steps into a reusable model phase.

The Step base class gives these operations a common interface. This lets a model pipeline focus on the sequence of work rather than on the implementation details of each operation.

The basic contract#

Every concrete Step has a name and implements one essential method:

def run_on(self, store):
    ...
    return updated_store

The input is the complete model store available at that point in the simulation. The return value is the store that should be passed to the next step. A step can therefore read several tables even when it writes only one of them—for example, a household annotation can use person and land-use data while adding fields to the household table.

A step can also carry:

  • name, a stable identifier used for display, logging, random-number derivation, or completion tracking;

  • param, a dictionary or Parameters object containing model coefficients or settings; and

  • on, the typed StoredTable class that the step primarily operates on.

Not every step needs parameters or a primary table. The only universal behavior is accepting a store and returning a store.

Running a sequence#

run_steps() executes its arguments in order. Internally, it keeps the return value from each step and supplies it to the next one:

store = store.run_steps(
    sample_households,
    annotate_persons,
    choose_auto_ownership,
)

Conceptually, the state flows through the pipeline like this:

        flowchart LR
    loaded["Loaded store"] -->|"sample households"| sampled["Sampled store"]
    sampled -->|"derive person fields"| annotated["Annotated store"]
    annotated -->|"simulate auto ownership"| choices["Store with choices"]
    choices --> outputs["Results and reports"]
    

Each step sees all results produced before it. Order therefore matters: a choice model cannot consume a derived field until the annotation step that creates it has run. run_steps() is deliberately a small, explicit sequencer—not a scheduler that discovers dependencies or reorders work.

Both Step instances and functions created by Traveler’s step decorators can appear in this sequence. A Step exposes itself through its __step__ property. A decorated function receives a __step__ attribute pointing to the step object built by the decorator. run_steps() uses that shared interface to handle both forms uniformly.

Updating the store#

A step mutates the logical simulation state, but it does not have to mutate the same Python object in place. In fact, Traveler’s built-in model steps generally use a functional update style:

  1. read from the input store;

  2. make a shallow store copy or construct a replacement store;

  3. replace the affected JaxTable or JaxDataset; and

  4. return the updated store.

For example, a small custom step can add a derived household field with assign_on():

from traveler.steps import Step


class IncomeThousands(Step):
    def __init__(self):
        super().__init__(name="income_thousands")

    def run_on(self, store):
        out = store.assign_on(
            "households",
            income_thousands=store["households"].income / 1_000,
        )
        out.completed_steps |= {self.name}
        return out

Running it produces a different store whose household table has the new column:

updated = store.run_steps(IncomeThousands())

updated["households"].income_thousands
"income_thousands" in store["households"]  # False

An application-specific step is allowed to update the input store directly, but it must still return a store. Code that invokes a step should never rely on object identity or ignore the return value:

# Correct: preserve the state returned by the step.
store = step.run_on(store)

# Also correct: run_steps handles the returned state automatically.
store = store.run_steps(step)

This contract allows different step implementations to choose the update strategy appropriate to their work while remaining composable.

Completion metadata#

A step can add its name to store.completed_steps. This metadata records which phases have run and can be used by validation. In particular, a field annotated with from_step is not required until its producing step is complete.

Completion tracking is not automatic in the Step base class; a concrete step decides when its work is complete and records the name as part of its returned state. This avoids marking a step complete before all of its updates have succeeded.

Example: computing derived fields#

compute_values turns a three-argument function into a step. The function receives:

  1. a detached ReferencedTable identifying the records being processed;

  2. the store, for access to other model data; and

  3. a parameter object.

It returns a dictionary of fields to add to its configured table:

from typing import Any

from traveler import Parameters
from traveler.core.stored import ReferencedTable
from traveler.steps import compute_values


class IncomeParameters(Parameters):
    high_income_threshold: float = 75_000.0


@compute_values(on=Households, param=IncomeParameters, use_jit=False)
def annotate_income(
    households_ref: ReferencedTable[Households],
    store: Store,
    param: IncomeParameters,
) -> dict[str, Any]:
    households = households_ref.with_store(store)
    return {
        "high_income": households.income >= param.high_income_threshold,
    }

The decorated name remains a function for direct inspection, while annotate_income.__step__ is the compute_values step. Either can be passed through the step workflow:

store = store.run_steps(annotate_income)

When run this way, compute_values resolves the complete chooser table, executes the calculation, adds the returned fields with assign_on(), and records the step name in completed_steps. With use_jit=True, compatible computations can use Traveler’s JAX and chunked execution path.

This pattern works well for preprocessing and annotation steps: the function describes the calculation, while the reusable step wrapper handles table selection and incorporation of the results into the store.

Example: simulating a choice#

simple_choice is a higher-level step for multinomial or nested logit models. Its utility function evaluates alternatives for a typed chooser table:

def mode_utilities(tours_ref, store, param):
    tours = tours_ref.with_store(store)
    return {
        "walk": param["walk_constant"] - param["time"] * tours.walk_time,
        "drive": param["drive_constant"] - param["time"] * tours.drive_time,
    }


choose_mode = simple_choice(
    name="choose_mode",
    utility_func=mode_utilities,
    altnames=("walk", "drive"),
    param={
        "walk_constant": 0.0,
        "drive_constant": 1.0,
        "time": 0.1,
    },
    on=Tours,
    result_name="mode",
)

choose_mode.run_on(store) computes utilities and probabilities, draws one alternative per tour using deterministic JAX random keys, and returns a store whose tours table contains the new mode field. String alternatives are represented as a categorical value, so their labels stay attached to the simulated codes.

The same model object also exposes methods for utilities, probabilities, logsums, analytic shares, and calibration. Those methods compute particular model quantities; run_on() is the operation that incorporates simulated choices into the store and participates in a step pipeline.

Grouping steps into a phase#

compound_step packages an ordered list of smaller steps as one reusable step. An initialization phase might derive land-use, person, and household fields in dependency order:

from traveler.steps import compound_step


initialize = compound_step(
    [
        annotate_land_use,
        annotate_persons,
        annotate_households,
        household_value_of_time,
    ],
    name="initialize",
)

store = store.run_steps(initialize)

The compound step does not merge the calculations or run them concurrently. It simply passes each component’s returned store to the next component and returns the final result. The group is useful for naming and reusing a stable phase of the model.

A complete simulation flow#

A small simulation can combine all of these forms while keeping the control flow visible:

store = load_inputs()
store.validate(missing_fields="ignore")

store = store.run_steps(
    SampleHouseholds(n=5_000),
    initialize,
    auto_ownership,
    work_location,
    choose_mode,
)

store.validate()
write_results(store)

The resulting pipeline has a straightforward interpretation:

  1. load and check the initial model state;

  2. reduce it to the desired simulation population;

  3. derive the fields required by later models;

  4. run choice models in dependency order;

  5. validate the completed state; and

  6. convert or write the resulting tables.

Because every component follows the same store-to-store contract, individual steps can be tested alone, grouped into phases, reused in other pipelines, or replaced without changing how the surrounding simulation is orchestrated.