bayesianbandits.ArmColumnFeaturizer#

class bayesianbandits.ArmColumnFeaturizer(column_name: str = 'arm_token')#

Bases: ArmFeaturizer[T]

Appends arm tokens as a column to context features.

This featurizer simply adds the arm token as an additional column to each context’s features. This allows users to then apply standard sklearn transformers (OneHotEncoder, PolynomialFeatures, etc.) to create the desired feature representation.

The output can be used with sklearn’s ColumnTransformer to apply different transformations to context features vs. the arm column.

Parameters:: column_name (str, default "arm_token") – Name of the column to add containing the arm tokens.

Examples

>>> import numpy as np
>>> from sklearn.compose import ColumnTransformer
>>> from sklearn.preprocessing import OneHotEncoder, PolynomialFeatures
>>>
>>> # Create featurizer
>>> featurizer = ArmColumnFeaturizer()
>>>
>>> # Context features
>>> X = np.array([[1.0, 2.0], [3.0, 4.0]])  # 2 contexts, 2 features
>>>
>>> # Example 1: Discrete arm tokens with one-hot encoding
>>> X_with_arms = featurizer.transform(X, action_tokens=[0, 1, 2])
>>> X_with_arms.shape
(6, 3)
>>>
>>> # Apply sklearn transformers
>>> transformer = ColumnTransformer([
...     ('context', 'passthrough', [0, 1]),
...     ('arm', OneHotEncoder(sparse_output=False), [2])
... ])
>>> X_final = transformer.fit_transform(X_with_arms)
>>> # Result has 2 context features + 3 one-hot encoded arm features
>>> X_final.shape
(6, 5)

>>> # Example 2: Continuous arm values with polynomial features
>>> X_with_arms = featurizer.transform(X, action_tokens=[0.1, 0.5, 0.9])
>>> transformer = ColumnTransformer([
...     ('context', 'passthrough', [0, 1]),
...     ('arm', PolynomialFeatures(degree=2, include_bias=False), [2])
... ])
>>> X_final = transformer.fit_transform(X_with_arms)
>>> # Result has 2 context features + 2 polynomial features (x, x^2)
>>> X_final.shape
(6, 4)

>>> # Example 3: With pandas DataFrame (preserves DataFrame structure)
>>> import pandas as pd
>>> df = pd.DataFrame({'feature1': [1.0, 3.0], 'feature2': [2.0, 4.0]})
>>> df_with_arms = featurizer.transform(df, action_tokens=['A', 'B', 'C'])
>>> # Returns DataFrame with 'arm_token' column added
>>> list(df_with_arms.columns)
['feature1', 'feature2', 'arm_token']

>>> # Example 4: Using custom column name
>>> featurizer = ArmColumnFeaturizer(column_name='product_id')
>>> df_with_arms = featurizer.transform(df, action_tokens=['A', 'B', 'C'])
>>> list(df_with_arms.columns)
['feature1', 'feature2', 'product_id']

__init__(column_name: str = 'arm_token')#

Initialize the ArmColumnFeaturizer.

Parameters:: column_name (str, default "arm_token") – Name of the column to add containing the arm tokens.

transform(X: Sized, *, action_tokens: Sequence[T]) → Any#

Append arm tokens as a column to context features.

Parameters:

X (Sized) – Input context features. Can be a pandas DataFrame, numpy array, or other array-like structure of shape (n_contexts, n_features).
action_tokens (sequence of T) – Action tokens for each arm. Can be any type (int, float, str, etc.).

Returns:

X_transformed – If X is a pandas DataFrame, returns a DataFrame with arm tokens appended as a new column. Otherwise, returns an ndarray of shape (n_contexts * n_arms, n_features + 1).

Return type:

DataFrame or ndarray

bayesianbandits.ArmColumnFeaturizer#

This Page