Traveler#

traveler is an experimental preview library for building and running ABM component models using JAX and Polars. It is designed to be flexible and efficient, allowing users to easily create custom models and pipelines for processing data.

Important

This documentation site is an early preview and no part of the traveler package is stable. What you see here represents an early preview of the intended functionality, which will evolve over time. The API can and will change without any deprecation warnings or other notice.

Why JAX?#

Activity-based models perform many of the same numerical operations for large numbers of households, people, tours, alternatives, and scenarios. JAX is a good fit for this work because it provides a NumPy-like array interface along with transformations that can compile and differentiate numerical code.

The main advantages are:

  • Speed. JAX can just-in-time compile array operations with jax.jit, combine operations that would otherwise require several passes over the data, and run the same model code on CPUs, GPUs, or other supported accelerators. Vectorized calculations also avoid slow Python loops over agents and alternatives.

  • Automatic differentiation. JAX can calculate derivatives of smooth model computations. This is useful for parameter estimation, calibration, sensitivity analysis, and optimization without requiring hand-written derivative code. It does not make discrete sampling differentiable, but the utility and probability calculations leading up to a choice can often be differentiated.

  • Composable numerical tools. Compilation, vectorization, differentiation, and parallelization are designed to work together. A well-structured model function can therefore be reused in several computational settings.

JAX also introduces tradeoffs:

  • Compilation has an up-front cost. Short calculations may not benefit, and changing array shapes, dtypes, or static arguments can trigger recompilation.

  • JAX works best with pure functions, immutable arrays, and regular array shapes. Side effects, data-dependent Python control flow, ragged data, and string-heavy processing generally require a different approach.

  • Moving data between JAX and non-JAX libraries can add time and memory overhead. Compiled code can also be less straightforward to inspect and debug than ordinary Python.

  • Accelerator memory is finite, and JAX commonly uses 32-bit numerical types by default. Large models must manage memory carefully and explicitly select higher precision where it matters.

Traveler therefore uses JAX for the numerical core of models and uses tools such as Polars for tabular data preparation. This division keeps the parts that benefit from compilation and automatic differentiation in JAX without forcing every data-processing task into its programming model. Care is taken to minimize the cost of conversions between JAX and non-JAX data structures, and zero-copy conversions are used where possible.

Licensing#

traveler is NOT an open-source package at this time. You may review the code and documentation here to evaluate the quality of these tools, but you must be granted a license to use it for any practical modeling.