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

CategoricalType

Builds a shape- and category-aware annotation

No

Categorical

Names a shared category family such as TAZ

No

CategoricalType is the annotation helper behind TypeMaker.Categorical. It binds a symbolic array shape to a category definition and produces a typing.Annotated type whose underlying runtime type is CategoricalArray. Application code normally uses it through a TypeMaker rather than constructing a CategoricalType directly:

from traveler import TypeMaker
from traveler.trees 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 another annotation form for a named category family. It is used in typed stored-table views when the category values are defined elsewhere:

from traveler.core.stored import StoredTable
from traveler.core.typing import Categorical


class Tours(StoredTable):
    table_name = "tours"

    origin: Categorical.TAZ
    destination: Categorical.TAZ

Accessing Categorical.TAZ creates a lightweight categorical type whose mark is "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 schema information, not an alternative container.

Note

There is also a lower-level metaclass named traveler.core.categorical.CategoricalType that makes annotation forms such as CategoricalArray[...] possible. It serves the same general annotation layer and is not categorical data itself. Normal table and dataset schemas should prefer TypeMaker.Categorical, while stored-table views use Categorical.GROUP when they need to name a shared category family.

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.