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 orParametersobject containing model coefficients or settings; andon, the typedStoredTableclass 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:
read from the input store;
make a shallow store copy or construct a replacement store;
replace the affected
JaxTableorJaxDataset; andreturn 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:
a detached
ReferencedTableidentifying the records being processed;the store, for access to other model data; and
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:
load and check the initial model state;
reduce it to the desired simulation population;
derive the fields required by later models;
run choice models in dependency order;
validate the completed state; and
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.