Traffic Speed Prediction with Metaheuristic Optimizers

Predicting average traffic speed (km/h) with a small MLP, and comparing how well population-based metaheuristics — Particle Swarm Optimization (PSO), a hybrid memetic PSO, and a Genetic Algorithm (GA) — train its weights compared to the gradient-based Adam baseline.

All code was extracted unchanged from the original Colab notebook (notebooks/optimization-full-project.ipynb) and organized into an importable Python package plus runnable training scripts.

Project overview

The task

Given hourly traffic-density records (location, timestamp, vehicle count), predict the average speed for that road segment and hour. The data follows the Istanbul Metropolitan Municipality traffic density format (DATE_TIME, LATITUDE, LONGITUDE, NUMBER_OF_VEHICLES, AVERAGE_SPEED); the notebook used traffic_density_202501.csv (January 2025).

The model

TrafficSpeedMLP — a compact fully-connected network:

13 input features → Linear(13, 32) → ReLU → Linear(32, 16) → ReLU → Linear(16, 1)

Its small parameter count (~1k weights) is deliberate: it keeps the search space tractable for the population-based optimizers, which treat the flattened weight vector as an individual/particle in R^D.

The features (13 dimensions)

Built in a fully vectorized pass over the DataFrame (build_feature_tensors):

Feature	Dims	Notes
Latitude, longitude	2	MinMax-scaled to [-1, 1] (fit on train split only)
sin(hour), cos(hour)	2	Cyclic encoding of the hour of day
Day of week	7	One-hot
Public-holiday flag	1	From an optional holiday set
Vehicle count	1	MinMax-scaled to [-1, 1]

The target AVERAGE_SPEED is standardized (StandardScaler); its training-set standard deviation σ is kept to rescale the Huber loss threshold.

The loss

WeightedHuberLoss = per-sample Huber loss (δ = 5 km/h, rescaled by the target σ) × sample weights (peak-hour flag and traffic-volume bonus, clamped ≥ 1)

• λ‖θ‖² L2 regularization (λ = 1e-5). Peak hours are 06–10 and 16–20.

The data split

Chronological (no leakage): train ≤ Jan 20, validation Jan 21–25, test > Jan 25, with a 70/15/15 percentage fallback if any split comes out empty. All scalers are fit on the training split only.

The optimizers

Optimizer	Module	How it searches
PSO	optimizers/pso.py	Classic global-best PSO over the flattened weight vector: fixed inertia ω = 0.6, cognitive/social coefficients c1 = c2 = 1.7, He-initialized swarm, early stopping on validation loss. Train fitness is estimated on a ~5000-sample subset per particle.
Hybrid PSO	optimizers/pso_hybrid.py	Memetic PSO: linearly decreasing inertia (0.9 → 0.4), velocity and position clipping (box constraint ±3), plus Lamarckian "micro" Adam steps every 5 iterations on a shared mini-batch — gradient refinement written back into the particle.
GA (loop)	optimizers/ga.py	Generational GA on the weight vector: k-tournament selection, simulated binary crossover (SBX, η = 10), per-gene Gaussian mutation, elitism (5%), early stopping when the top-5% mean fitness stops improving. Fitness is validation loss.
GA (vectorized)	optimizers/ga_vectorized.py	Same algorithm re-written with tensor ops: the whole population is a (P, D) tensor and fitness for all individuals is computed in one pass per batch via torch.func.functional_call + vmap. Requires PyTorch ≥ 2.1.
Adam (from scratch)	optimizers/adam.py	A faithful re-implementation of Adam (bias-corrected first/second moments, optional weight decay). Defined in the notebook for reference; the Adam training run itself uses torch.optim.Adam, as the notebook did.

Every optimizer tracks the best weights by validation loss and restores them at the end of training. Final comparison is on the held-out test set: weighted Huber loss, plus MAE / RMSE / R² computed on inverse-scaled (km/h) predictions.

Repository structure

.
├── notebooks/
│   └── optimization-full-project.ipynb   # Original notebook (untouched)
├── src/traffic_opt/                      # The extracted package
│   ├── model.py                          # TrafficSpeedMLP
│   ├── loss.py                           # WeightedHuberLoss
│   ├── param_utils.py                    # flatten/unflatten model weights <-> 1D tensor
│   ├── evaluation.py                     # evaluate_model_on_test_set (loss, MAE, RMSE, R²)
│   ├── optimizers/
│   │   ├── adam.py                       # From-scratch Adam
│   │   ├── pso.py                        # PSOOptimizer
│   │   ├── pso_hybrid.py                 # PSOOptimizerHybrid (memetic)
│   │   ├── ga.py                         # GAOptimizer (loop-based)
│   │   └── ga_vectorized.py              # Vectorized GA (exported as GAOptimizerVectorized)
│   └── data/
│       ├── features.py                   # build_feature_tensors (13-dim feature builder)
│       ├── dataset.py                    # TrafficDataset, preprocess_data (filter + split + scalers)
│       └── loaders.py                    # create_dataloaders, INPUT_DIM, BATCH_SIZE
├── scripts/
│   ├── train_pso.py                      # Train + test-evaluate plain PSO
│   ├── train_pso_hybrid.py               # Train + test-evaluate hybrid PSO
│   ├── train_ga.py                       # Train + test-evaluate GA
│   └── train_adam.py                     # Train + test-evaluate Adam baseline
├── data/                                 # Put the CSV here (git-ignored); see data/README.md
├── pyproject.toml
├── requirements.txt
└── README.md

Notebook → module mapping

Notebook section	Extracted to
Imports	Distributed per module
Implementation of the Base of the Model	src/traffic_opt/model.py
Implementation of Loss Function	src/traffic_opt/loss.py
Utility Functions	src/traffic_opt/param_utils.py
Adam Optimizer	src/traffic_opt/optimizers/adam.py
PSO Optimizer	src/traffic_opt/optimizers/pso.py
PSO Optimizer (Hybrid)	src/traffic_opt/optimizers/pso_hybrid.py
GA Optimizier	src/traffic_opt/optimizers/ga.py
GA Optimizer (Vectorized)	src/traffic_opt/optimizers/ga_vectorized.py
build_feature_tensors cell	src/traffic_opt/data/features.py
Set Up the Dataset (definitions)	src/traffic_opt/data/dataset.py
Set Up the Dataset (loader creation)	src/traffic_opt/data/loaders.py
Evaluate the PSO Optimizer	scripts/train_pso.py
Hybrid PSO Training	scripts/train_pso_hybrid.py
GA Training	scripts/train_ga.py
Final Evaluation on Test Set (function)	src/traffic_opt/evaluation.py
Adam training + evaluation cells	scripts/train_adam.py
google.colab drive mount	Dropped (Colab-only); the CSV path is now a --csv argument

Getting started

1. Install

python -m venv .venv && source .venv/bin/activate
pip install -e .          # installs the traffic_opt package + dependencies

(or pip install -r requirements.txt and run with PYTHONPATH=src.)

2. Get the data

Drop the traffic density CSV into data/traffic_density_202501.csv (see data/README.md for the expected columns).

3. Train

python scripts/train_adam.py                 # gradient baseline
python scripts/train_pso.py                  # plain PSO
python scripts/train_pso_hybrid.py           # memetic PSO + Adam micro-steps
python scripts/train_ga.py                   # genetic algorithm

Every script accepts --csv <path> and --batch-size <n> (default 256), prints training progress, and finishes with test-set metrics (weighted Huber loss, MAE, RMSE, R² in km/h).

Library usage:

from traffic_opt import TrafficSpeedMLP, WeightedHuberLoss
from traffic_opt.optimizers import PSOOptimizer, GAOptimizerVectorized
from traffic_opt.data import create_dataloaders, INPUT_DIM

Faithfulness notes

The code is a verbatim extraction of the notebook — including its quirks, which were intentionally not fixed:

• Two GAOptimizer classes. The notebook defines the class twice (loop-based and vectorized). The GA training cell consumes history['val_loss'], which matches the loop-based version's return type, so traffic_opt.optimizers.GAOptimizer is the loop-based one; the vectorized variant is exported as GAOptimizerVectorized (its train() returns a plain list of per-generation losses).
• Loop-based GA is CUDA-only as written. Its _initialize_population moves the template model to 'cuda' (hardcoded, despite the adjacent "Init on CPU" comment), so scripts/train_ga.py requires a GPU — exactly as in the notebook, which ran on a Colab T4. Use GAOptimizerVectorized on CPU-only machines.
• The custom AdamOptimizer is never used in the pipeline. The notebook's Adam run uses torch.optim.Adam; the from-scratch implementation is kept as reference.
• Loss weighting is a placeholder. In WeightedHuberLoss, the peak-hour multiplier and volume bonus factor are both 1.0 (marked "refine as needed" in the notebook). The vectorized GA's inlined loss uses a peak-hour multiplier of 2.0 — a pre-existing inconsistency between the two notebook cells, preserved as-is.
• PSO train-fitness subsampling. PSOOptimizer estimates each particle's train fitness on roughly the first 5000 samples of the (shuffled) train loader.

The only adaptations made during modularization (no logic changed):

1. The Colab Google Drive mount and hardcoded CSV path were replaced by the --csv CLI argument (default data/traffic_density_202501.csv), and the BATCH_SIZE = 256 constant became the --batch-size argument (same default).
2. In scripts/train_pso_hybrid.py, the parameter-count print references model_pso_hybrid; the notebook cell referenced model_pso from an earlier cell, which doesn't exist in a standalone script.
3. scripts/train_pso_hybrid.py ends with a test-set evaluation for parity with the other scripts (the notebook only test-evaluated PSO, GA and Adam).
4. Unused imports from the notebook's import cell (torch.nn.functional, datetime, scipy.stats.kstest) were not carried into the modules.

Requirements

• Python ≥ 3.9
• PyTorch ≥ 2.1 (torch.func is needed by the vectorized GA)
• numpy, pandas, scikit-learn, scipy
• A CUDA GPU for scripts/train_ga.py (see faithfulness notes); everything else runs on CPU or GPU.