Categorical Data#

Categorical data represents values drawn from a finite set of labels. A travel mode, for example, might be one of "walk", "bike", or "car"; a zone field might contain one of the zone identifiers defined by a land-use table.

Traveler separates a categorical variable into two parts:

  • categories define the possible values and their order;

  • codes are integer positions into that category list.

The runtime representation is CategoricalArray:

from traveler import CategoricalArray


modes = ["walk", "bike", "car"]
mode = CategoricalArray.from_values(
    categories=modes,
    values=["car", "walk", "car"],
)

mode.categories  # ["walk", "bike", "car"]
mode.codes       # Array([2, 0, 2], dtype=uint8)
mode.ordered     # False

The labels carry meaning for people and for validation. The codes are the compact numerical representation used by model calculations. Codes should not be interpreted independently of their categories: code 2 means "car" only for this particular category list.

Why JAX needs a special representation#

JAX has no native categorical dtype. Its arrays contain numerical, boolean, and a few other JAX-supported scalar types, but they cannot directly contain arbitrary Python strings together with category metadata. Converting a column of labels with jnp.asarray() therefore cannot produce an array suitable for JAX transformations.

CategoricalArray bridges that gap. It is registered as a JAX pytree with:

  • codes, a JAX integer array, as its dynamic array data;

  • categories, ordered, and mark as the metadata describing those codes.

Numerical work can operate on the codes under jax.jit while the category definition remains attached to the value:

import jax


@jax.jit
def is_car(value):
    return value.codes == 2


is_car(mode)  # Array([ True, False,  True], dtype=bool)

Traveler also has special handlers for categorical values when flattening pytrees, selecting rows, assigning fields, and converting tables. A generic array conversion would either reject string labels or discard the metadata needed to decode and validate the integer codes.

Note

Categories are static pytree metadata from JAX’s perspective. Passing a categorical array with a different category definition to compiled code may therefore result in a different JAX compilation. The codes can vary normally between calls.

Runtime data and annotation helpers#

Traveler uses three related concepts on the data and schema sides:

Name

Role

Contains data?

CategoricalArray

Runtime categorical value

Yes: JAX codes plus category metadata

TypeMaker.Categorical(...)

Declares shape, categories, and ordering

No

Categorical

Names a shared category family such as TAZ

No

TypeMaker binds a symbolic array shape to a category definition through its Categorical(...) factory. The result is a typing.Annotated type whose underlying runtime type is CategoricalArray:

from traveler import TypeMaker
from traveler.tables import JaxTable


trip = TypeMaker("trip")
modes = ["walk", "bike", "car"]


class Trips(JaxTable):
    mode: trip.Categorical(categories=modes)

Conceptually, the annotation records all of the following:

runtime type: CategoricalArray
categories:   ["walk", "bike", "car"]
ordered:      False
shape:        "trip"

Categorical provides the corresponding annotation form for a named category family. It is used when the category values are defined elsewhere, such as by another table in the same store:

from traveler import Categorical
from traveler.store import StoredTable


class Tours(StoredTable):
    table_name = "tours"

    origin: Categorical.TAZ
    destination: Categorical.TAZ

Accessing Categorical.TAZ produces a categorical annotation carrying the group name "TAZ". Both fields therefore declare that they use the same category family. The actual columns returned by tours.origin and tours.destination are still CategoricalArray instances; Categorical is only an annotation namespace, not another data container.

Creating and loading categorical data#

When source data contains labels, from_values() is the clearest constructor. It creates the integer codes and rejects values that are not present in the provided category list:

mode = CategoricalArray.from_values(
    categories=["walk", "bike", "car"],
    values=["car", "walk", "car"],
    ordered=False,
    mark="MODE",
)

If the source already contains codes and the code-to-category mapping is known, construct the array directly:

import jax.numpy as jnp


mode = CategoricalArray(
    categories=["walk", "bike", "car"],
    codes=jnp.array([2, 0, 2]),
    mark="MODE",
)

The direct constructor checks that codes are non-negative and below the number of categories. For small category sets it stores them in a compact integer dtype. Prefer from_values() when beginning with labels because it establishes the mapping instead of assuming the input codes already follow it.

Traveler preserves categorical information at its supported table-conversion boundaries:

  • a Pandas categorical Series becomes a CategoricalArray when used to construct a JaxDataset or JaxTable;

  • a Polars Enum column becomes a CategoricalArray in JaxTable.from_polars();

  • converting a JaxTable back to Pandas or Polars reconstructs categorical columns from their codes and categories.

For example, a Polars enum carries the complete category vocabulary into a typed JAX table:

import polars as pl


frame = pl.DataFrame(
    {
        "mode": pl.Series(
            ["car", "walk", "car"],
            dtype=pl.Enum(modes),
        )
    }
)

trips = Trips.from_polars(frame)
trips.mode.categories  # ["walk", "bike", "car"]
trips.mode.codes       # Array([2, 0, 2], dtype=uint8)

CSV and ordinary Parquet columns do not necessarily carry a complete, application-defined category vocabulary. A loader should therefore cast a prepared Polars column to Enum or explicitly call from_values() before placing the data in a JAX-backed object. Nullable source values must also be resolved during preparation because CategoricalArray does not use a negative or null code.

Validation#

Construction establishes a valid code range. Dataset validation goes further by comparing the runtime value with its schema annotation:

trips = Trips({"mode": mode})
trips.validate(coerce=False)

For a TypeMaker.Categorical annotation, validation checks that:

  • the field is actually a CategoricalArray;

  • its codes are JAX integer data;

  • its code-array shape matches the declared symbolic dimensions;

  • its categories match in both values and order; and

  • its ordered flag matches the annotation.

Categorical dimensions participate in the same named-dimension consistency checks as numerical arrays. If mode and another field both use the "trip" dimension, their lengths must agree.

Some category sets are defined by model data rather than a literal list in a class definition. Prefixing a category name with @ declares a shared category group:

trip = TypeMaker("trip")


class Trips(JaxTable):
    origin: trip.Categorical(categories="@TAZ")

    def __category_validation__(self):
        return {"TAZ": [101, 102, 103]}


origin = CategoricalArray.from_values(
    categories=[101, 102, 103],
    values=[103, 101, 102],
    mark="TAZ",
)

trips = Trips({"origin": origin})
trips.validate(coerce=False)

Here @TAZ means “validate against the category group named TAZ.” Validation checks both the group’s category values and the array’s mark. A dataset can provide groups through __category_validation__(), and a Store collects compatible groups from its contents so related tables can use the same vocabulary. For example, a LandUseMAZ table might define a MAZ group, which will serve as the category labels for any other table in the same store that uses Categorical.MAZ.

import numpy as np

microzone = TypeMaker("MAZ")
trip = TypeMaker("trip")


class LandUseMAZ(StoredTable):

    MAZ: np.int32
    """Unique identifier for the MAZ."""

    def __category_validation__(self) -> dict[str, Any]:
        return {
            "MAZ": self.MAZ,
        }


class Trip(StoredTable):

    origin: Categorical.MAZ
    """Origin MAZ for the trip."""

    destination: Categorical.MAZ
    """Destination MAZ for the trip."""

When that store already defines a group, CategoricalArray.from_values_predefined() and CategoricalArray.from_codes_predefined() can construct an array from the shared definition.

Using categorical values in model code#

Use codes for JAX indexing, comparisons, and other numerical operations. Use the category metadata or convert to Pandas or Polars when labels are needed for display or data preparation:

mode.codes                    # numerical model representation
mode.to_pandas().tolist()     # ["car", "walk", "car"]

For nominal data, the numerical ordering of codes is only an encoding detail. For ordinal data, set ordered=True and keep the category list in its intended order. In either case, retaining the CategoricalArray wrapper ensures that the codes, their meaning, and the validation rules continue to travel together.