San Francisco Crime Classification with Gators#

This notebook demonstrates how to use the gators library for feature engineering in a multiclass classification problem. We’ll predict crime categories from the San Francisco Crime dataset using various gators transformers.

Key Features Demonstrated:#

  • Data preprocessing and cleaning

  • DateTime feature engineering (cyclic, ordinal, business time, holidays)

  • String feature engineering (pattern detection)

  • Spatial feature engineering (coordinate rotation)

  • Group-based imputation

  • Feature encoding and interaction

1. Import Libraries#

Import the necessary libraries including gators transformers for feature engineering.

[1]:

import polars as pl
import pandas as pd
from IPython.display import display

from gators.pipeline import Pipeline
from gators.encoders import OneHotEncoder, CountEncoder
from gators.data_cleaning import DropColumns, CastColumns
from gators.feature_generation import PlanRotationFeatures
from gators.imputers import GroupByImputer
from gators.feature_generation_dt import (
    CyclicFeatures,
    OrdinalFeatures,
    BusinessTimeFeatures,
    TimeBinFeatures,
    HolidayFeatures,
)
from gators.feature_generation_str import (
    Contains,
    InteractionFeatures,
)

from xgboost import XGBClassifier

2. Load and Preprocess Data#

Load the San Francisco Crime dataset and perform initial preprocessing:

  • Remove redundant columns

  • Handle outlier coordinate values (replace invalid coordinates with null)

  • Create a distance-to-center feature using San Francisco’s geographic center

[2]:

# Load data
X_train = pl.read_parquet('../../../../../Documents/kaggle/sf/train.parquet')
X_test = pl.read_parquet('../../../../../Documents/kaggle/sf/test.parquet')

# Drop unnecessary columns
X_train = X_train.drop(["Descript", "Resolution", "DayOfWeek"])
X_test = X_test.drop(["DayOfWeek"])

# Extract target variable
target = 'Category'
y_train = X_train[target]
X_train = X_train.drop(target)

# Replace invalid coordinate values with null
X_train = X_train.with_columns([
    pl.when(pl.col('X') == -120.5).then(None).otherwise(pl.col('X')).alias('X'),
    pl.when(pl.col('Y') == 90.0).then(None).otherwise(pl.col('Y')).alias('Y')
])

X_test = X_test.with_columns([
    pl.when(pl.col('X') == -120.5).then(None).otherwise(pl.col('X')).alias('X'),
    pl.when(pl.col('Y') == 90.0).then(None).otherwise(pl.col('Y')).alias('Y')
])

# Create distance to San Francisco center feature
sf_center_x, sf_center_y = -122.4194, 37.7749
X_train = X_train.with_columns([
    (((pl.col('X') - sf_center_x)**2 + (pl.col('Y') - sf_center_y)**2)**0.5).alias('distance_to_center')
])
X_test = X_test.with_columns([
    (((pl.col('X') - sf_center_x)**2 + (pl.col('Y') - sf_center_y)**2)**0.5).alias('distance_to_center')
])

# Parse Dates into a native Datetime column up front (in Python, not inside the
# exported pipeline) - ONNX has no generic string-to-datetime parse operator, so
# this must happen before pipeline.fit()/transform() rather than via CastColumns.
X_train = X_train.with_columns(pl.col('Dates').str.to_datetime())
X_test = X_test.with_columns(pl.col('Dates').str.to_datetime())

print(f"Training data shape: {X_train.shape}")
print(f"Test data shape: {X_test.shape}")

Training data shape: (878049, 6)
Test data shape: (884262, 7)

3. Prepare Target Variable#

Convert crime category labels to numeric format for model training.

[3]:

# Define all crime categories
crime_categories = [
    "ARSON", "ASSAULT", "BAD CHECKS", "BRIBERY", "BURGLARY",
    "DISORDERLY CONDUCT", "DRIVING UNDER THE INFLUENCE", "DRUG/NARCOTIC",
    "DRUNKENNESS", "EMBEZZLEMENT", "EXTORTION", "FAMILY OFFENSES",
    "FORGERY/COUNTERFEITING", "FRAUD", "GAMBLING", "KIDNAPPING",
    "LARCENY/THEFT", "LIQUOR LAWS", "LOITERING", "MISSING PERSON",
    "NON-CRIMINAL", "OTHER OFFENSES", "PORNOGRAPHY/OBSCENE MAT",
    "PROSTITUTION", "RECOVERED VEHICLE", "ROBBERY", "RUNAWAY",
    "SECONDARY CODES", "SEX OFFENSES FORCIBLE", "SEX OFFENSES NON FORCIBLE",
    "STOLEN PROPERTY", "SUICIDE", "SUSPICIOUS OCC", "TREA", "TRESPASS",
    "VANDALISM", "VEHICLE THEFT", "WARRANTS", "WEAPON LAWS"
]

# Convert target to numeric labels
mapping = {cat: str(i) for i, cat in enumerate(crime_categories)}
y_train = y_train.replace(mapping).cast(pl.Int64)

print(f"Number of classes: {len(crime_categories)}")
print(f"Target shape: {y_train.shape}")

Number of classes: 39
Target shape: (878049,)

4. Build Feature Engineering Pipeline#

The gators Pipeline orchestrates multiple feature transformers:

Data Cleaning & Casting:#

  • GroupByImputer: Fill missing X/Y coordinates using mean values grouped by police district

  • CastColumns: Convert date strings to datetime type

Datetime Features:#

  • TimeBinFeatures: Create time-of-day bins (morning, afternoon, evening, night)

  • DatetimeCyclicFeatures: Generate cyclic features for temporal patterns (sin/cos transformations)

  • BusinessTimeFeatures: Extract business-related time features (weekday vs weekend)

  • HolidayFeatures: Identify whether crimes occurred on holidays

  • DatetimeOrdinalFeatures: Create ordinal features (month, day, hour, etc.)

String Features:#

  • Contains: Detect patterns in address strings (e.g., contains ‘Block’, ‘AV’, ‘ST’)

  • InteractionFeatures: Create interactions between categorical features (e.g., day_of_week × part_of_day)

Spatial Features:#

  • PlanRotationFeatures: Rotate X/Y coordinates at multiple angles to capture spatial patterns

Encoding:#

  • OneHotEncoder: Encode categorical variables as binary features

  • CountEncoder: Encode categories by their frequency

[4]:

# Define datetime components to extract
ordinal_components = [
    "month", "week", "day_of_week", "day_of_month",
    "day_of_year", "hour", "minute", "weekend"
]

cyclic_components = [
    "month", "week", "day_of_week", "day_of_month",
    "day_of_year", "hour", "minute"
]

# Build the pipeline with gators transformers
steps = [
    # Impute missing coordinates using group averages
    ("group_imputer", GroupByImputer(
        group_by_column='PdDistrict',
        strategy='mean',
        subset=['X', 'Y']
    )),

    # Extract time-of-day bins
    ("dt_timebin", TimeBinFeatures(
        subset=['Dates'],
        bin_types=['part_of_day', 'season', 'time_of_month', 'time_of_year', 'rush_hour', 'day']
    )),

    # Create cyclic datetime features (captures cyclical nature of time)
    ("dt_cyclic", CyclicFeatures(
        subset=['Dates'],
        angles=[180 * i / 4 for i in range(8)],
        components=cyclic_components
    )),

    # Extract business time features
    ("dt_business", BusinessTimeFeatures(subset=['Dates'])),

    # Add holiday indicator
    ("dt_holiday", HolidayFeatures(subset=['Dates'], features=['is_holiday'])),

    # Extract ordinal datetime components
    ("dt_ordinal", OrdinalFeatures(subset=['Dates'], components=ordinal_components)),

    # Extract patterns from address strings
    ("contains", Contains(contains_dict={'Address': ['/', 'Block', 'AV', 'ST']})),

    # Drop original columns after feature extraction
    ("drop", DropColumns(subset=["Address", "Dates"])),

    # Create interaction features
    # Dates__day (string weekday name, from TimeBinFeatures) rather than the numeric
    # Dates__day_of_week (OrdinalFeatures) - ONNX can't string-concatenate a numeric
    # column and exactly match Polars' float-to-string formatting.
    ("interaction_time", InteractionFeatures(subset=['Dates__day', 'Dates__part_of_day'])),
    ("interaction_district", InteractionFeatures(subset=['PdDistrict', 'Dates__part_of_day'])),

    # Rotate coordinates to capture spatial patterns
    ("plan_rotation", PlanRotationFeatures(
        columns=[['X', 'Y']],
        angles=[180 * i / 4 for i in range(8)]
    )),

    # Encode categorical variables
    ("onehot_encoder", OneHotEncoder()),
    ("count_encoder", CountEncoder()),
]

# Create and fit the pipeline
pipeline = Pipeline(steps=steps, verbose=True)
X_train_transformed = pipeline.fit_transform(X_train, y_train)
X_test_transformed = pipeline.transform(X_test)

print(f"\nOriginal features: {X_train.shape[1]}")
print(f"Engineered features: {X_train_transformed.shape[1]}")

[Pipeline] fit+transform   1/13 · group_imputer  |  in: rows=878049  cols=6  nulls=201  →  out: rows=878049  cols=6  nulls=67  (0.017s)
[Pipeline] fit+transform   2/13 · dt_timebin  |  in: rows=878049  cols=6  nulls=67  →  out: rows=878049  cols=12  nulls=67  (0.050s)
[Pipeline] fit+transform   3/13 · dt_cyclic  |  in: rows=878049  cols=12  nulls=67  →  out: rows=878049  cols=68  nulls=67  (0.088s)
[Pipeline] fit+transform   4/13 · dt_business  |  in: rows=878049  cols=68  nulls=67  →  out: rows=878049  cols=71  nulls=67  (0.037s)
[Pipeline] fit+transform   5/13 · dt_holiday  |  in: rows=878049  cols=71  nulls=67  →  out: rows=878049  cols=72  nulls=67  (0.104s)
[Pipeline] fit+transform   6/13 · dt_ordinal  |  in: rows=878049  cols=72  nulls=67  →  out: rows=878049  cols=80  nulls=67  (0.014s)
[Pipeline] fit+transform   7/13 · contains  |  in: rows=878049  cols=80  nulls=67  →  out: rows=878049  cols=84  nulls=67  (0.022s)
[Pipeline] fit+transform   8/13 · drop  |  in: rows=878049  cols=84  nulls=67  →  out: rows=878049  cols=82  nulls=67  (0.001s)
[Pipeline] fit+transform   9/13 · interaction_time  |  in: rows=878049  cols=82  nulls=67  →  out: rows=878049  cols=83  nulls=67  (0.025s)
[Pipeline] fit+transform   10/13 · interaction_district  |  in: rows=878049  cols=83  nulls=67  →  out: rows=878049  cols=84  nulls=67  (0.030s)
[Pipeline] fit+transform   11/13 · plan_rotation  |  in: rows=878049  cols=84  nulls=67  →  out: rows=878049  cols=100  nulls=67  (0.005s)
[Pipeline] fit+transform   12/13 · onehot_encoder  |  in: rows=878049  cols=100  nulls=67  →  out: rows=878049  cols=196  nulls=67  (0.139s)
[Pipeline] fit+transform   13/13 · count_encoder  |  in: rows=878049  cols=196  nulls=67  →  out: rows=878049  cols=196  nulls=67  (0.000s)
[Pipeline] transform   1/13 · group_imputer  |  in: rows=884262  cols=7  nulls=228  →  out: rows=884262  cols=7  nulls=76  (0.013s)
[Pipeline] transform   2/13 · dt_timebin  |  in: rows=884262  cols=7  nulls=76  →  out: rows=884262  cols=13  nulls=76  (0.047s)
[Pipeline] transform   3/13 · dt_cyclic  |  in: rows=884262  cols=13  nulls=76  →  out: rows=884262  cols=69  nulls=76  (0.086s)
[Pipeline] transform   4/13 · dt_business  |  in: rows=884262  cols=69  nulls=76  →  out: rows=884262  cols=72  nulls=76  (0.034s)
[Pipeline] transform   5/13 · dt_holiday  |  in: rows=884262  cols=72  nulls=76  →  out: rows=884262  cols=73  nulls=76  (0.005s)
[Pipeline] transform   6/13 · dt_ordinal  |  in: rows=884262  cols=73  nulls=76  →  out: rows=884262  cols=81  nulls=76  (0.014s)
[Pipeline] transform   7/13 · contains  |  in: rows=884262  cols=81  nulls=76  →  out: rows=884262  cols=85  nulls=76  (0.013s)
[Pipeline] transform   8/13 · drop  |  in: rows=884262  cols=85  nulls=76  →  out: rows=884262  cols=83  nulls=76  (0.001s)
[Pipeline] transform   9/13 · interaction_time  |  in: rows=884262  cols=83  nulls=76  →  out: rows=884262  cols=84  nulls=76  (0.024s)
[Pipeline] transform   10/13 · interaction_district  |  in: rows=884262  cols=84  nulls=76  →  out: rows=884262  cols=85  nulls=76  (0.030s)
[Pipeline] transform   11/13 · plan_rotation  |  in: rows=884262  cols=85  nulls=76  →  out: rows=884262  cols=101  nulls=76  (0.005s)
[Pipeline] transform   12/13 · onehot_encoder  |  in: rows=884262  cols=101  nulls=76  →  out: rows=884262  cols=197  nulls=76  (0.047s)
[Pipeline] transform   13/13 · count_encoder  |  in: rows=884262  cols=197  nulls=76  →  out: rows=884262  cols=197  nulls=76  (0.000s)

Original features: 6
Engineered features: 196

5. Train Model and Generate Predictions#

Train an XGBoost classifier with the engineered features and generate predictions.

[5]:

[c for c in X_test_transformed.columns if c not in X_train_transformed.columns]

[5]:

['Id']

[6]:

# Define model parameters
xgb_params = {
    'max_depth': 6,
    'learning_rate': 0.1,
    'n_estimators': 100,
    'min_child_weight': 5,
    'subsample': 0.85,
    'colsample_bytree': 0.85,
    'max_delta_step': 2,
    'tree_method': 'hist',
    'eval_metric': 'mlogloss',
    'random_state': 0,
    'n_jobs': -1
}

# Train the model
estimator = XGBClassifier(**xgb_params)
estimator.fit(X_train_transformed, y_train)

print("Model training completed!")

# Generate predictions
submission = pd.DataFrame(
    estimator.predict_proba(X_test_transformed.drop("Id")),
    index=X_test["Id"].to_pandas(),
    columns=crime_categories
)
submission.to_csv("submission.csv.zip", compression="zip")
print("Submission file created successfully!")

Model training completed!
Submission file created successfully!

6. Feature Importance Analysis#

Analyze which features contribute most to the model’s predictions.

[7]:

# Extract feature importances
feat_imp = pl.DataFrame({
    "feature": estimator.feature_names_in_,
    "importance": estimator.feature_importances_
})
feat_imp = feat_imp.sort("importance", descending=True)

# Display top 20 most important features
print("Top 20 Most Important Features:")
display(feat_imp.head(20))

Top 20 Most Important Features:
shape: (20, 2)
featureimportance
strf32
"Address__contains_Block"0.055322
"Address__contains_/"0.043314
"Dates__minute__sin270"0.041261
"Dates__minute"0.037772
"Dates__minute__sin90"0.034877
……
"PdDistrict__Dates__part_of_day…0.010086
"PdDistrict__Dates__part_of_day…0.009924
"XY_y45"0.009858
"XY_x90"0.00976
"PdDistrict__SOUTHERN"0.009686

7. ONNX Export#

Convert the fitted pipeline and XGBoost classifier to ONNX for portable, runtime-independent inference.

Function

Purpose

check_pipeline_onnx_compatibility

Audit which steps have native ONNX converters

pipeline_to_onnx

Export preprocessing pipeline to a self-contained ONNX graph

pipeline_to_scoring_onnx

Chain preprocessing graph with an ML model ONNX graph

create_session

Create an ORT_ENABLE_ALL-optimised InferenceSession

run_session

Run session on a Polars DataFrame (supports batching)

 
[8]:

from gators.onnx_converters import (
    check_pipeline_onnx_compatibility,
    pipeline_to_onnx,
    pipeline_to_scoring_onnx,
    create_session,
    run_session,
)
import onnx
import os

ONNX_DIR = "onnx_models"
os.makedirs(ONNX_DIR, exist_ok=True)

# ── 7.1  Audit ONNX compatibility ───────────────────────────────────
compat = check_pipeline_onnx_compatibility(pipeline)
supported     = sorted(name for name, ok in compat.items() if ok)
not_supported = sorted(name for name, ok in compat.items() if not ok)
print(f"Native ONNX converters  ({len(supported)}): {supported}")
print(f"Identity pass-throughs  ({len(not_supported)}): {not_supported}")

# ── 7.2  Export preprocessing pipeline ───────────────────────────────────
preprocessing_onnx = pipeline_to_onnx(pipeline, errors="coerce")
onnx.save(preprocessing_onnx, f"{ONNX_DIR}/sf_crime_preprocessing.onnx")

print(f"\nPreprocessing graph saved.")
print(f"  Nodes  : {len(preprocessing_onnx.graph.node)}")

print(f"  Inputs : {len(preprocessing_onnx.graph.input)}")
print(f"  Outputs: {len(preprocessing_onnx.graph.output)}")

Native ONNX converters  (13): ['contains', 'count_encoder', 'drop', 'dt_business', 'dt_cyclic', 'dt_holiday', 'dt_ordinal', 'dt_timebin', 'group_imputer', 'interaction_district', 'interaction_time', 'onehot_encoder', 'plan_rotation']
Identity pass-throughs  (0): []

Preprocessing graph saved.
  Nodes  : 1520
  Inputs : 6
  Outputs: 196

[9]:

from onnxmltools import convert_xgboost
from onnxmltools.convert.common.data_types import FloatTensorType

feature_columns = X_train_transformed.columns
n_features = len(feature_columns)

# ── 7.3  Export XGBoost classifier to ONNX ───────────────────────────────────
# estimator was fit() on a named Polars DataFrame, so its booster stores real column
# names instead of the "f0","f1",... pattern onnxmltools expects when converting a
# raw Booster - clear them so the dump uses generic feature indices.
booster = estimator.get_booster()
booster.feature_names = None
xgb_onnx = convert_xgboost(
    booster,
    initial_types=[("input", FloatTensorType([None, n_features]))],
)
onnx.save(xgb_onnx, f"{ONNX_DIR}/sf_crime_xgb.onnx")
print(f"XGBoost model exported  ({n_features} input features)")

# ── 7.4  Build end-to-end scoring model: raw data → crime category probabilities ──
scoring_onnx = pipeline_to_scoring_onnx(
    pipeline,
    xgb_onnx,
    feature_columns=feature_columns,
    errors="coerce",
)
onnx.save(scoring_onnx, f"{ONNX_DIR}/sf_crime_scoring.onnx")
print(f"\nEnd-to-end scoring model saved.")
print(f"  Inputs : {len(scoring_onnx.graph.input)}")
print(f"  Outputs: {[o.name for o in scoring_onnx.graph.output]}")

XGBoost model exported  (196 input features)

End-to-end scoring model saved.
  Inputs : 6
  Outputs: ['label', 'probabilities']

[10]:

import numpy as np

# ── 7.5  Create ORT sessions and run inference ────────────────────────────────
pre_session     = create_session(preprocessing_onnx)
scoring_session = create_session(scoring_onnx, optimized_model_path=f"{ONNX_DIR}/sf_crime_scoring_opt.onnx")

# Raw test features (no Id) → engineered features
X_test_raw = X_test.select(pipeline._input_columns)
X_test_onnx = run_session(pre_session, X_test_raw)
print(f"ONNX preprocessing output: {X_test_onnx.shape}")

# End-to-end: raw data → crime category probabilities
onnx_preds = run_session(scoring_session, X_test_raw)
onnx_probs = onnx_preds["probabilities"].to_numpy()

# ── 7.6  Compare xgb_onnx vs scoring_onnx ────────────────────────────────────
# xgb_onnx expects a [N, F] float32 matrix — feed the same features used to fit()
xgb_session     = create_session(xgb_onnx)
xgb_input_name  = xgb_session.get_inputs()[0].name
X_xgb           = X_train_transformed.to_pandas().values.astype("float32")
xgb_train_probs = np.asarray(xgb_session.run(None, {xgb_input_name: X_xgb})[1])
native_probs    = estimator.predict_proba(X_train_transformed)
max_diff        = float(abs(xgb_train_probs - native_probs).max())

print(f"\nMax |xgb_onnx - native|: {max_diff:.2e}")
assert max_diff < 1e-4, f"Models diverge: {max_diff:.2e}"
print("ONNX predictions match native predictions ✓")

2026-08-24 10:27:54.745 Python[55744:332825] 2026-08-24 10:27:54.745931 [W:onnxruntime:, graph.cc:5337 CleanUnusedInitializersAndNodeArgs] Removing initializer '__hf_big'. It is not used by any node and should be removed from the model.
2026-08-24 10:27:54.746 Python[55744:332825] 2026-08-24 10:27:54.745993 [W:onnxruntime:, graph.cc:5337 CleanUnusedInitializersAndNodeArgs] Removing initializer '__hf_neg1'. It is not used by any node and should be removed from the model.
2026-08-24 10:27:54.823 Python[55744:332825] 2026-08-24 10:27:54.823361 [W:onnxruntime:, graph.cc:5337 CleanUnusedInitializersAndNodeArgs] Removing initializer '__hf_big'. It is not used by any node and should be removed from the model.
2026-08-24 10:27:54.823 Python[55744:332825] 2026-08-24 10:27:54.823434 [W:onnxruntime:, graph.cc:5337 CleanUnusedInitializersAndNodeArgs] Removing initializer '__hf_neg1'. It is not used by any node and should be removed from the model.

ONNX preprocessing output: (884262, 196)

Max |xgb_onnx - native|: 1.28e-06
ONNX predictions match native predictions ✓

Summary#

This notebook demonstrates the power of the gators library for feature engineering in machine learning pipelines:

Key Takeaways:#

  1. Modular Pipeline: Easily compose complex feature engineering workflows using gators transformers

  2. Datetime Engineering: Rich set of temporal features (cyclic, ordinal, business time, holidays)

  3. Spatial Features: Coordinate rotation and distance-based features for geographic data

  4. String Processing: Pattern extraction from text fields

  5. Smart Imputation: Group-based imputation leveraging categorical relationships

  6. Seamless Integration: Works with Polars DataFrames and scikit-learn/XGBoost models

The gators library significantly reduced feature engineering complexity while creating a comprehensive feature set for multiclass classification.