Bayesian Congestion Model on AWS Lambda
Deploy a hierarchical PyMC congestion model as a serverless API on AWS Lambda.
Introduction
30 Second Summary
Every day, hundreds of buses crawl through Auckland traffic, running the same routes at the same times. Each one is quietly collecting data about how congested the roads actually are. If a bus is five minutes late at every stop along Great South Road, that road is probably jammed.
In this project, you will build a hierarchical Bayesian model in PyMC that treats buses as mobile congestion sensors. You'll fit the model on real-time delay data from Auckland Transport, then deploy it as a live prediction API on AWS Lambda.
What You'll Build
A Bayesian congestion model that uses Auckland bus delays as proxy measurements for road congestion, deployed as a serverless prediction API that returns full posterior distributions with credible intervals.
By the end of this project, you'll have:
- A data ingestion pipeline that polls the Auckland Transport GTFS real-time API and stores delay observations in Amazon S3.
- A hierarchical Bayesian model in PyMC with corridor-level partial pooling, Fourier time-of-day seasonality, and day-of-week effects.
- A containerized Lambda function behind API Gateway that returns posterior congestion estimates with 90% credible intervals.
- Secret Mission: Add weather covariates from MetService and build a congestion anomaly detector.
Are there any prerequisites?
This is an advanced project. You should be comfortable writing Python scripts and have some familiarity with Bayesian statistics (priors, posteriors, sampling). You also need an AWS account.
Before We Start
Before We Start...
Set Up Your Environment and AWS CLI
This project models Auckland's road congestion by treating buses as mobile sensors. Before you can collect real-time delay data and build the Bayesian model, your local environment and cloud access need to be ready.
This step gets three foundational pieces in place. You will install the AWS CLI for managing cloud resources, create a Python virtual environment with PyMC and its supporting libraries, and register for the Auckland Transport real-time data API.
In this step, get ready to:
- Install and configure AWS CLI v2 with your credentials.
- Create a Python virtual environment and install project dependencies.
- Register for the Auckland Transport Developer Portal and obtain your API subscription key.
Install and configure AWS CLI v2
The AWS CLI lets you manage AWS services directly from your terminal. You will use it later to push container images, create Lambda functions, and set up API Gateway endpoints.
- Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type Terminal and press Enter to open it.
- Check if the AWS CLI is already installed by running:
aws --version
✔️ I see a version number
Your AWS CLI is already installed. Skip ahead to the configuration section below.
ⓧ Command not found
No problem. Download and install the official macOS package by running these two commands:
curl "https://awscli.amazonaws.com/AWSCLIV2.pkg" -o "AWSCLIV2.pkg"
sudo installer -pkg AWSCLIV2.pkg -target /
What do these commands do?
The first command downloads the official AWS CLI v2 installer package. The second command runs that installer with admin privileges.
Your terminal will ask for your Mac login password after the sudo command. Type it in and press Enter. The characters will not appear as you type.
- Confirm the install succeeded by running:
aws --version
You should see output starting with aws-cli/2. followed by a version number.
Still seeing command not found?
Close your terminal window completely and open a new one. The installer updates your PATH, but existing terminal sessions do not pick up the change automatically.
If the problem persists, make sure the sudo installer command above finished without errors. Help me troubleshoot my AWS CLI installation.
Now you need to connect the CLI to your AWS account. This requires an access key from the IAM console.
Where do I find my AWS access keys?
Sign in to the AWS Management Console. Open the IAM service and navigate to Users. Select your user, then open the Security credentials tab.
Click Create access key and copy both the Access Key ID and Secret Access Key. Save them somewhere safe. You will not be able to see the secret key again after closing the dialog.
- Configure the CLI with your credentials by running:
aws configure
What does aws configure ask for?
This command prompts you for four values, one at a time:
- AWS Access Key ID - paste the access key you copied from the IAM console.
- AWS Secret Access Key - paste the secret key you copied.
- Default region name - enter ap-southeast-2. This is the Sydney region, the closest AWS region to Auckland.
- Default output format - enter json.
- Verify the configuration is working by running:
aws sts get-caller-identity
You should see a JSON response containing your AWS account ID, user ARN, and user ID.
Seeing an InvalidClientTokenId error?
This means your access key or secret key was entered incorrectly. Run aws configure again and carefully re-paste both values.
Make sure the access key in your IAM console is Active, not Inactive or Deleted. Help me fix my AWS CLI credentials.
Create your Python virtual environment
A virtual environment keeps this project's Python packages isolated from the rest of your system. PyMC and its dependencies are large, so isolation prevents version conflicts with other projects.
- Move to your Desktop by running this command:
cd ~/Desktop
- Create your project folder and move into it by running these commands:
mkdir auckland-congestion
cd auckland-congestion
- Confirm the folder exists by running:
ls ~/Desktop
You should see auckland-congestion in the output.
- Create a virtual environment by running:
python3 -m venv venv
- Activate the virtual environment:
macOS
source venv/bin/activate
Windows
venv\Scripts\activate
Your terminal prompt should now start with (venv). This confirms you are working inside the virtual environment.
Prompt does not show (venv)?
Make sure you ran the activation command from inside the auckland-congestion folder. Run pwd to check your current directory. If you are somewhere else, run cd ~/Desktop/auckland-congestion first, then try the activation command again. Help me activate my Python virtual environment.
Now you can install all the Python packages this project needs. PyMC is the core Bayesian modeling library. ArviZ provides diagnostics and visualization for the model output.
- Install all project dependencies by running:
pip install pymc==5.28.4 arviz pandas requests boto3 matplotlib
What are these packages?
- pymc is the hierarchical Bayesian modeling library you will use to build the congestion model.
- arviz provides trace plots, posterior predictive checks, and model summary tables.
- pandas handles the tabular delay data you will collect from Auckland Transport.
- requests makes HTTP calls to the Auckland Transport real-time API.
- boto3 is the AWS SDK for Python. You will use it to upload data to S3.
- matplotlib renders the data visualizations and model diagnostic plots.
This install takes a few minutes. PyMC pulls in several large numerical libraries including NumPy and PyTensor.
- Verify the install succeeded by running:
python3 -c "import pymc; print(pymc.__version__)"
You should see 5.28.4 printed in your terminal.
Seeing a ModuleNotFoundError?
Make sure your virtual environment is active. Your prompt should show (venv) at the beginning. If it does not, run the activation command from the platform tab above, then run the pip install command again.
If the install itself failed, scroll up in your terminal to find the error message. Help me fix my PyMC installation.
Register for the Auckland Transport API
Auckland Transport publishes live bus positions and delay data through a GTFS real-time API. You need a free account and a subscription key to access it.
What is GTFS?
GTFS stands for General Transit Feed Specification. It is a standardized format that transit agencies use to publish schedule and real-time vehicle data. The real-time extension provides live delay and position updates.
- Open dev-portal.at.govt.nz in your browser.
- Click Sign Up in the top right corner.
- Complete the registration form with your email address and a password.
- Verify your email address by clicking the confirmation link sent to your inbox.
- Sign in to the Developer Portal with your new account.
Now you need to subscribe to the specific API that provides real-time trip updates.
- Click APIs in the top navigation bar.
- Find and select the Realtime Compat API.
- Click Subscribe to add this API to your account.
After subscribing, the portal displays your subscription key. This is the value you will include in API requests.
- Navigate to your Profile page by clicking your username in the top right corner.
- Locate your Ocp-Apim-Subscription-Key under the Subscriptions section.
- Copy this key and save it somewhere accessible. You will need it in the next step when you build the data collector.
What are the API rate limits?
The Auckland Transport API allows 600 calls per minute and 35,000 calls per week. The real-time feed updates every 30 seconds, so polling once per minute is plenty for building up your dataset.
- Test the key by running this quick check in your terminal:
curl -s -o /dev/null -w "%{http_code}" -H "Ocp-Apim-Subscription-Key: your-subscription-key-here" https://api.at.govt.nz/realtime/legacy/tripupdates
- Replace your-subscription-key-here with the key you just copied from the portal.
You should see 200 printed in your terminal. This confirms the API accepted your subscription key.
Seeing 401 or 403 instead of 200?
A 401 means your subscription key was not recognized. Double-check that you copied the full key from the portal with no extra spaces.
A 403 means your subscription does not include this API. Go back to the portal and confirm you subscribed to the Realtime Compat API specifically. Help me fix my Auckland Transport API access.
Your environment is fully set up. You have the AWS CLI configured, a Python environment loaded with PyMC and its supporting libraries, and a live API key for Auckland Transport's bus data. Next up, you will build the data ingestion pipeline that collects real-time delay observations from every active bus trip.
Build the Data Ingestion Pipeline
Your environment is ready. Now you need data to feed your Bayesian model.
Auckland Transport's GTFS real-time API broadcasts live delay data for every active bus trip across the network. Each update tells you how many seconds late (or early) a bus is at each stop. Your job in this step is to capture those snapshots, map each observation to a road corridor, and store the results in S3 so your model has raw material to work with.
In this step, get ready to:
- Write a Python script that polls the AT Trip Update API and parses delay observations.
- Map stops to corridor segments using GTFS static data.
- Upload partitioned CSV files to an S3 bucket and accumulate 24-48 hours of data.
Poll the AT Trip Update API
The AT Trip Update API returns a JSON snapshot of every active bus trip on the Auckland network. Each snapshot contains nested stop_time_update records with delay values in seconds.
- Create a new file called ingest.py in your project folder.
- Add the imports, configuration, and API polling function by pasting this code into ingest.py:
import requests
import pandas as pd
import boto3
import os
import zipfile
from datetime import datetime, timezone
from io import BytesIO
# Auckland Transport API configuration
AT_API_URL = "https://api.at.govt.nz/realtime/legacy/tripupdates"
AT_API_KEY = "your-at-subscription-key-here"
# S3 configuration
S3_BUCKET = "[[BUCKET_NAME=\"auckland-congestion-yourname\"]]"
s3_client = boto3.client("s3")
def fetch_trip_updates():
"""Poll the AT Trip Update API for real-time delay data."""
headers = {"Ocp-Apim-Subscription-Key": AT_API_KEY}
response = requests.get(AT_API_URL, headers=headers)
response.raise_for_status()
return response.json()
What does this code do?
The imports pull in everything you need: requests for the HTTP call, pandas for tabular data, and boto3 for uploading to S3.
The fetch_trip_updates() function sends a GET request to AT's Trip Update endpoint with your subscription key in the Ocp-Apim-Subscription-Key header. The API returns a full snapshot of all active trips as JSON.
- Replace your-at-subscription-key-here with the Ocp-Apim-Subscription-Key you obtained from the AT Developer Portal in Step 1.
Now you need a function to extract the fields you care about from the nested JSON response.
- Add the parsing function below the fetch_trip_updates() function in ingest.py:
def parse_delay_observations(data):
"""Extract delay fields from GTFS real-time trip update response."""
observations = []
# Feed-level timestamp (seconds since epoch)
timestamp = data.get("header", {}).get("timestamp",
int(datetime.now(timezone.utc).timestamp()))
for entity in data.get("entity", []):
trip_update = entity.get("trip_update", {})
if not trip_update:
continue
# Extract the route ID from the trip descriptor
trip = trip_update.get("trip", {})
route_id = trip.get("route_id", "")
# Each stop_time_update has delay info for one stop
stop_time_updates = trip_update.get("stop_time_update", [])
if isinstance(stop_time_updates, dict):
stop_time_updates = [stop_time_updates]
for stu in stop_time_updates:
stop_id = stu.get("stop_id", "")
stop_sequence = stu.get("stop_sequence", 0)
# Delay comes from arrival or departure timing
arrival = stu.get("arrival", {})
departure = stu.get("departure", {})
delay = arrival.get("delay", departure.get("delay", 0))
observations.append({
"route_id": route_id,
"stop_id": stop_id,
"stop_sequence": stop_sequence,
"delay_seconds": delay,
"timestamp": timestamp,
})
return observations
What does this code do?
The AT API uses snake_case field names that match the GTFS Realtime spec: trip_update, stop_time_update, route_id, and so on.
The function loops through every entity in the feed. For each stop along each trip, it extracts the route_id, stop_id, stop_sequence, and the delay in seconds. Positive delay means late. Negative means early.
The isinstance(stop_time_updates, dict) check handles a known AT API quirk where a trip with a single stop update returns a dict instead of a list.
- Save ingest.py.
You will not see output yet. The next substep adds corridor mapping and the main execution block.
Seeing a syntax error?
Make sure the indentation is consistent. Python requires spaces (not tabs) for indentation. Check that both functions are at the top level of the file, not nested inside each other.
Help me debug a syntax error in my ingest.py file that polls the AT Trip Update API.
Map stops to corridor segments
Raw delay observations are tagged by stop_id and route_id, but your Bayesian model needs to reason about road corridors. Multiple bus routes share the same corridor, and that overlap is what gives partial pooling its power.
You will use the GTFS static data from Auckland Transport to map each stop to its corridor. The static data contains stops.txt, routes.txt, trips.txt, and stop_times.txt files that link routes to their stops.
- Download the GTFS static data zip by running this command:
curl -o gtfs.zip https://gtfs.at.govt.nz/gtfs.zip
- Extract the zip into a gtfs_static folder by running:
mkdir -p gtfs_static && unzip -o gtfs.zip -d gtfs_static
- Confirm the key files exist by running:
ls gtfs_static
You should see stops.txt, routes.txt, trips.txt, and stop_times.txt among other files in the output.
Don't see the expected files?
Make sure you ran the curl command from your project folder. If the download failed, check your internet connection and try the URL directly in your browser: https://gtfs.at.govt.nz/gtfs.zip.
Help me download and extract the Auckland Transport GTFS static data.
Now add the corridor mapping logic to ingest.py. This joins the static schedule data to assign each stop to a named corridor like "Great South Road SB" or "Dominion Road NB".
- Add the corridor mapping functions below parse_delay_observations() in ingest.py:
def load_corridor_mapping(gtfs_dir):
"""Build a mapping from (route_id, stop_id) to corridor segment
using GTFS static data."""
# Load the four key GTFS static files
stops = pd.read_csv(os.path.join(gtfs_dir, "stops.txt"))
routes = pd.read_csv(os.path.join(gtfs_dir, "routes.txt"))
stop_times = pd.read_csv(os.path.join(gtfs_dir, "stop_times.txt"))
trips = pd.read_csv(os.path.join(gtfs_dir, "trips.txt"))
# Join trips with routes to get route_short_name per trip
trips_routes = trips.merge(
routes[["route_id", "route_short_name"]], on="route_id"
)
# Join stop_times with trip/route info and stop names
merged = stop_times.merge(
trips_routes[["trip_id", "route_id", "route_short_name",
"direction_id"]],
on="trip_id",
)
merged = merged.merge(
stops[["stop_id", "stop_name", "stop_lat", "stop_lon"]],
on="stop_id",
)
# Corridor keywords matched against stop names
corridor_keywords = {
"Great South": "Great South Road",
"Gt South": "Great South Road",
"Dominion": "Dominion Road",
"Mt Eden": "Mt Eden Road",
"Symonds": "Symonds Street",
"Queen": "Queen Street",
"Manukau": "Manukau Road",
"Remuera": "Remuera Road",
"Sandringham": "Sandringham Road",
"New North": "New North Road",
}
def assign_corridor(row):
"""Match a stop's name to a corridor keyword."""
stop_name = str(row["stop_name"])
for keyword, corridor in corridor_keywords.items():
if keyword.lower() in stop_name.lower():
# direction_id 0 = outbound (NB), 1 = inbound (SB)
direction = "NB" if row.get("direction_id", 0) == 0 else "SB"
return f"{corridor} {direction}"
return "Other"
merged["corridor_id"] = merged.apply(assign_corridor, axis=1)
# Build the lookup dictionary: (route_id, stop_id) -> corridor_id
corridor_map = (
merged.groupby(["route_id", "stop_id"])["corridor_id"]
.first()
.to_dict()
)
return corridor_map
def add_corridors(observations, corridor_map):
"""Tag each observation with its corridor using the lookup."""
for obs in observations:
key = (obs["route_id"], obs["stop_id"])
obs["corridor_id"] = corridor_map.get(key, "Other")
return observations
What does this code do?
The load_corridor_mapping() function joins four GTFS static files (stops.txt, routes.txt, trips.txt, stop_times.txt) to connect every route-stop pair to a physical road corridor.
It matches stop names against corridor keywords like "Great South" or "Dominion". The direction_id from the trip data determines whether the bus is heading northbound (NB) or southbound (SB). Stops that do not match any keyword are labelled "Other" and filtered out later.
Why match by stop name?
Auckland bus stops are typically named after the road they sit on (e.g. "123 Great South Road"). This is a practical heuristic that works well for the major corridors you are modeling. A production system would use geospatial matching, but keyword matching keeps this project focused on the Bayesian modeling.
- Save ingest.py.
Getting a KeyError on a column name?
The GTFS static files use exact column names like route_id, stop_id, and direction_id. Check that your gtfs_static folder contains the correct files extracted from the AT zip. You can inspect column names with head -1 gtfs_static/trips.txt in your terminal.
Help me debug a KeyError when loading GTFS static data in my corridor mapping code.
Upload observations to S3 and collect data
Your model needs data from different times of day and days of the week to learn commute patterns. You will store each collection run as a partitioned CSV in S3 so the data accumulates over time.
- Create your S3 bucket by running this command:
aws s3 mb s3://[[BUCKET_NAME="auckland-congestion-yourname"]]
Why a unique bucket name?
S3 bucket names are globally unique across all AWS accounts. Include your name or a random string to avoid collisions. Make sure the bucket name in this command matches the S3_BUCKET value in your ingest.py.
- Verify the bucket was created by running:
aws s3 ls
You should see your bucket name in the output list.
Seeing an access denied error?
Make sure your AWS CLI is configured with credentials that have S3 permissions. Run aws sts get-caller-identity to verify your CLI is authenticated. If the bucket name is already taken, pick a different name and update S3_BUCKET in ingest.py to match.
Help me fix an S3 bucket creation error.
Now add the upload function and the main execution block that ties everything together.
- Add the upload function and main block at the bottom of ingest.py:
def upload_to_s3(observations, bucket):
"""Upload delay observations as a partitioned CSV to S3."""
if not observations:
print("No observations to upload.")
return
df = pd.DataFrame(observations)
# Add human-readable datetime columns for partitioning
df["datetime"] = pd.to_datetime(df["timestamp"], unit="s", utc=True)
df["date"] = df["datetime"].dt.strftime("%Y-%m-%d")
df["hour"] = df["datetime"].dt.hour
# Partition by date and hour for efficient querying later
date_str = df["date"].iloc[0]
hour_str = str(df["hour"].iloc[0]).zfill(2)
timestamp_str = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
key = f"delays/date={date_str}/hour={hour_str}/{timestamp_str}.csv"
csv_buffer = df.to_csv(index=False)
s3_client.put_object(Bucket=bucket, Key=key, Body=csv_buffer)
print(f"Uploaded {len(df)} observations to s3://{bucket}/{key}")
if __name__ == "__main__":
# Load corridor mapping from GTFS static data
gtfs_dir = "gtfs_static"
print("Loading corridor mapping from GTFS static data...")
corridor_map = load_corridor_mapping(gtfs_dir)
print(f"Loaded {len(corridor_map)} route-stop corridor mappings.")
# Fetch real-time trip updates from AT
print("Fetching trip updates from Auckland Transport API...")
data = fetch_trip_updates()
# Parse delay observations from the response
observations = parse_delay_observations(data)
print(f"Parsed {len(observations)} delay observations.")
# Tag each observation with its corridor
observations = add_corridors(observations, corridor_map)
# Keep only observations on known corridors
observations = [obs for obs in observations
if obs["corridor_id"] != "Other"]
print(f"Retained {len(observations)} observations with known corridors.")
# Upload the batch to S3
upload_to_s3(observations, S3_BUCKET)
print("Done!")
What does this code do?
The upload_to_s3() function converts the observations list into a pandas DataFrame, adds human-readable date and hour columns, and writes the CSV to an S3 key partitioned by date and hour (e.g. delays/date=2026-05-03/hour=08/20260503_083000.csv). This partition structure makes it easy to load specific time windows later.
The __main__ block orchestrates the full pipeline: load the corridor map, fetch live data, parse it, tag corridors, filter out unmatched stops, and upload to S3.
- Save ingest.py.
- Run the ingestion script by executing this command:
python ingest.py
You should see output similar to this:
Loading corridor mapping from GTFS static data...
Loaded 48230 route-stop corridor mappings.
Fetching trip updates from Auckland Transport API...
Parsed 12450 delay observations.
Retained 3120 observations with known corridors.
Uploaded 3120 observations to s3://auckland-congestion-yourname/delays/date=2026-05-03/hour=08/20260503_083000.csv
Done!
What should I see?
The exact numbers will vary depending on the time of day you run the script. During peak hours (7-9am, 4-6pm) you will see more active trips. Off-peak runs will have fewer observations. Any non-zero count of retained observations means the pipeline is working.
Seeing 0 retained observations?
If you see "Retained 0 observations", the corridor keyword matching did not find any stops. Open gtfs_static/stops.txt and check the stop_name column for the road names used. You may need to add keywords to the corridor_keywords dictionary that match the actual stop names in the data.
If you see "Parsed 0 delay observations", the API may have returned no active trips. Try running the script during Auckland business hours (NZST). The AT feed updates every 30 seconds, but there are fewer buses running late at night.
Help me troubleshoot why my ingestion script is returning 0 observations.
- Verify the data landed in S3 by running:
aws s3 ls s3://[[BUCKET_NAME="auckland-congestion-yourname"]]/delays/ --recursive
You should see at least one CSV file path in the output.
How much data do I need?
Your Bayesian model needs observations from different times of day to learn the commute cycle. Run the script several times across at least 24-48 hours. Aim for a mix of morning peak (7-9am), evening peak (4-6pm), off-peak, and at least one weekend day.
You can schedule runs manually or use a simple loop. The AT API allows up to 600 calls per minute and 35,000 calls per week, so running once every 5-10 minutes is safe.
- Run the script a few more times over the next day or two to build up data. You can continue to the next step once you have at least 5-10 successful runs across different times.
✔️ Awesome, I've got everything!
Double check you have saved ingest.py and that your S3 bucket has at least one CSV file in the delays/ prefix.
ⓧ I'd like to double check the full code
import requests
import pandas as pd
import boto3
import os
import zipfile
from datetime import datetime, timezone
from io import BytesIO
# Auckland Transport API configuration
AT_API_URL = "https://api.at.govt.nz/realtime/legacy/tripupdates"
AT_API_KEY = "your-at-subscription-key-here"
# S3 configuration
S3_BUCKET = "[[BUCKET_NAME=\"auckland-congestion-yourname\"]]"
s3_client = boto3.client("s3")
def fetch_trip_updates():
"""Poll the AT Trip Update API for real-time delay data."""
headers = {"Ocp-Apim-Subscription-Key": AT_API_KEY}
response = requests.get(AT_API_URL, headers=headers)
response.raise_for_status()
return response.json()
def parse_delay_observations(data):
"""Extract delay fields from GTFS real-time trip update response."""
observations = []
# Feed-level timestamp (seconds since epoch)
timestamp = data.get("header", {}).get("timestamp",
int(datetime.now(timezone.utc).timestamp()))
for entity in data.get("entity", []):
trip_update = entity.get("trip_update", {})
if not trip_update:
continue
# Extract the route ID from the trip descriptor
trip = trip_update.get("trip", {})
route_id = trip.get("route_id", "")
# Each stop_time_update has delay info for one stop
stop_time_updates = trip_update.get("stop_time_update", [])
if isinstance(stop_time_updates, dict):
stop_time_updates = [stop_time_updates]
for stu in stop_time_updates:
stop_id = stu.get("stop_id", "")
stop_sequence = stu.get("stop_sequence", 0)
# Delay comes from arrival or departure timing
arrival = stu.get("arrival", {})
departure = stu.get("departure", {})
delay = arrival.get("delay", departure.get("delay", 0))
observations.append({
"route_id": route_id,
"stop_id": stop_id,
"stop_sequence": stop_sequence,
"delay_seconds": delay,
"timestamp": timestamp,
})
return observations
def load_corridor_mapping(gtfs_dir):
"""Build a mapping from (route_id, stop_id) to corridor segment
using GTFS static data."""
# Load the four key GTFS static files
stops = pd.read_csv(os.path.join(gtfs_dir, "stops.txt"))
routes = pd.read_csv(os.path.join(gtfs_dir, "routes.txt"))
stop_times = pd.read_csv(os.path.join(gtfs_dir, "stop_times.txt"))
trips = pd.read_csv(os.path.join(gtfs_dir, "trips.txt"))
# Join trips with routes to get route_short_name per trip
trips_routes = trips.merge(
routes[["route_id", "route_short_name"]], on="route_id"
)
# Join stop_times with trip/route info and stop names
merged = stop_times.merge(
trips_routes[["trip_id", "route_id", "route_short_name",
"direction_id"]],
on="trip_id",
)
merged = merged.merge(
stops[["stop_id", "stop_name", "stop_lat", "stop_lon"]],
on="stop_id",
)
# Corridor keywords matched against stop names
corridor_keywords = {
"Great South": "Great South Road",
"Gt South": "Great South Road",
"Dominion": "Dominion Road",
"Mt Eden": "Mt Eden Road",
"Symonds": "Symonds Street",
"Queen": "Queen Street",
"Manukau": "Manukau Road",
"Remuera": "Remuera Road",
"Sandringham": "Sandringham Road",
"New North": "New North Road",
}
def assign_corridor(row):
"""Match a stop's name to a corridor keyword."""
stop_name = str(row["stop_name"])
for keyword, corridor in corridor_keywords.items():
if keyword.lower() in stop_name.lower():
# direction_id 0 = outbound (NB), 1 = inbound (SB)
direction = "NB" if row.get("direction_id", 0) == 0 else "SB"
return f"{corridor} {direction}"
return "Other"
merged["corridor_id"] = merged.apply(assign_corridor, axis=1)
# Build the lookup dictionary: (route_id, stop_id) -> corridor_id
corridor_map = (
merged.groupby(["route_id", "stop_id"])["corridor_id"]
.first()
.to_dict()
)
return corridor_map
def add_corridors(observations, corridor_map):
"""Tag each observation with its corridor using the lookup."""
for obs in observations:
key = (obs["route_id"], obs["stop_id"])
obs["corridor_id"] = corridor_map.get(key, "Other")
return observations
def upload_to_s3(observations, bucket):
"""Upload delay observations as a partitioned CSV to S3."""
if not observations:
print("No observations to upload.")
return
df = pd.DataFrame(observations)
# Add human-readable datetime columns for partitioning
df["datetime"] = pd.to_datetime(df["timestamp"], unit="s", utc=True)
df["date"] = df["datetime"].dt.strftime("%Y-%m-%d")
df["hour"] = df["datetime"].dt.hour
# Partition by date and hour for efficient querying later
date_str = df["date"].iloc[0]
hour_str = str(df["hour"].iloc[0]).zfill(2)
timestamp_str = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
key = f"delays/date={date_str}/hour={hour_str}/{timestamp_str}.csv"
csv_buffer = df.to_csv(index=False)
s3_client.put_object(Bucket=bucket, Key=key, Body=csv_buffer)
print(f"Uploaded {len(df)} observations to s3://{bucket}/{key}")
if __name__ == "__main__":
# Load corridor mapping from GTFS static data
gtfs_dir = "gtfs_static"
print("Loading corridor mapping from GTFS static data...")
corridor_map = load_corridor_mapping(gtfs_dir)
print(f"Loaded {len(corridor_map)} route-stop corridor mappings.")
# Fetch real-time trip updates from AT
print("Fetching trip updates from Auckland Transport API...")
data = fetch_trip_updates()
# Parse delay observations from the response
observations = parse_delay_observations(data)
print(f"Parsed {len(observations)} delay observations.")
# Tag each observation with its corridor
observations = add_corridors(observations, corridor_map)
# Keep only observations on known corridors
observations = [obs for obs in observations
if obs["corridor_id"] != "Other"]
print(f"Retained {len(observations)} observations with known corridors.")
# Upload the batch to S3
upload_to_s3(observations, S3_BUCKET)
print("Done!")
Your data pipeline is live. Every time you run ingest.py, another batch of corridor-level delay observations lands in S3. Next up, you will load this data into pandas and prepare it for the Bayesian model.
Explore and Prepare the Data
Your data pipeline is collecting real-time delay observations from Auckland buses into S3.
Before feeding this data into a Bayesian model, you need to understand its shape. How are delays distributed across corridors? Do peak hours show the congestion spikes you would expect? In this step, you will load your collected data, visualize these patterns, and transform everything into the format PyMC needs for hierarchical modeling.
In this step, get ready to:
- Load delay CSVs from S3 into a pandas DataFrame and engineer feature columns.
- Visualize delay distributions by corridor, time of day, and day type.
- Encode, standardize, and package the data for PyMC's hierarchical model.
Load delay data from S3 and create feature columns
Your ingestion script stored delay observations as partitioned CSV files in S3. You need to pull all of those files into a single DataFrame so you can analyze the full picture across corridors and time periods.
- Create a new file called prepare_data.py in the same project folder where ingest.py lives.
- Add the imports and S3 loading logic by copying this code into prepare_data.py:
import boto3
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pickle
from io import StringIO
# Connect to S3 and list all CSV files in the bucket
s3 = boto3.client("s3")
bucket_name = "[[BUCKET_NAME="auckland-congestion-data"]]"
response = s3.list_objects_v2(Bucket=bucket_name)
csv_keys = [obj["Key"] for obj in response["Contents"] if obj["Key"].endswith(".csv")]
# Download and concatenate all CSV files into one DataFrame
frames = []
for key in csv_keys:
obj = s3.get_object(Bucket=bucket_name, Key=key)
df_part = pd.read_csv(StringIO(obj["Body"].read().decode("utf-8")))
frames.append(df_part)
df = pd.concat(frames, ignore_index=True)
print(f"Loaded {len(df)} delay observations from {len(csv_keys)} files")
What does this code do?
- list_objects_v2 fetches the list of every object in your S3 bucket. The list comprehension filters down to only .csv files.
- The loop downloads each CSV with get_object, reads its body as text through StringIO, and parses it into a DataFrame.
- pd.concat stacks all the individual DataFrames into one unified table with a clean integer index.
- Save prepare_data.py.
- Run the script so far by running:
python prepare_data.py
- You should see a message like Loaded 5000 delay observations from 12 files. Your exact numbers depend on how long your collector ran.
Seeing NoSuchBucket or AccessDenied?
- Double check that the bucket name in prepare_data.py exactly matches the bucket you created in Step 2.
- Make sure your AWS CLI credentials are configured for the same account that owns the bucket. Run aws sts get-caller-identity to confirm.
- If you see an empty result, your collector may not have uploaded any files yet. Go back to Step 2 and run ingest.py a few more times.
Help me debug my S3 data loading in prepare_data.py
Now you need to extract the features the model will use. The raw data has timestamps, but the model needs numeric columns for the hour of day, whether it is a weekday or weekend, and corridor identity.
- Add the feature engineering code below the print statement at the bottom of prepare_data.py:
# Parse timestamp into datetime for feature extraction
df["timestamp"] = pd.to_datetime(df["timestamp"])
# Create hour_of_day feature (0-23)
df["hour_of_day"] = df["timestamp"].dt.hour
# Create binary weekday/weekend feature (1 = weekday, 0 = weekend)
df["day_of_week"] = (df["timestamp"].dt.dayofweek < 5).astype(int)
# Convert corridor_id to categorical type
df["corridor_id"] = df["corridor_id"].astype("category")
print(f"Corridors: {df['corridor_id'].nunique()}")
print(f"Date range: {df['timestamp'].min()} to {df['timestamp'].max()}")
print(df[["delay_seconds", "hour_of_day", "day_of_week"]].describe())
What do these features represent?
- hour_of_day captures the daily commute cycle. You expect delays to spike during morning (7-9am) and evening (4-6pm) peaks.
- day_of_week is a binary flag where 1 means weekday and 0 means weekend. Weekday congestion patterns differ significantly from weekends.
- corridor_id as a categorical type tells pandas to treat each corridor as a distinct group rather than a plain string.
- Save prepare_data.py and run it again:
python prepare_data.py
- Below the loading message, you should now see the corridor count, date range, and a statistical summary table showing the mean, standard deviation, and quartiles for delay_seconds, hour_of_day, and day_of_week.
Seeing a KeyError for one of the columns?
- The column names must match exactly what your ingest.py script wrote. Check your CSV headers by adding print(df.columns.tolist()) right after the pd.concat line.
- If you see column names with leading or trailing spaces, add df.columns = df.columns.str.strip() before the feature extraction code.
Help me fix my column name mismatch in prepare_data.py
Visualize delay distributions
Numbers alone do not tell the full story. Visualizing the data reveals whether your delay distributions match the expected patterns of urban congestion.
- Add the visualization code below the describe() print statement in prepare_data.py:
# Plot delay distributions by corridor, hour, and day type
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
# Box plot of delays per corridor
df.boxplot(column="delay_seconds", by="corridor_id", ax=axes[0], rot=90)
axes[0].set_title("Delay by Corridor")
axes[0].set_ylabel("Delay (seconds)")
# Bar chart of mean delay per hour
hourly = df.groupby("hour_of_day")["delay_seconds"].mean()
axes[1].bar(hourly.index, hourly.values)
axes[1].set_title("Mean Delay by Hour of Day")
axes[1].set_xlabel("Hour")
axes[1].set_ylabel("Mean Delay (seconds)")
# Box plot comparing weekday vs weekend
df.boxplot(column="delay_seconds", by="day_of_week", ax=axes[2])
axes[2].set_xticklabels(["Weekend", "Weekday"])
axes[2].set_title("Delay by Day Type")
axes[2].set_ylabel("Delay (seconds)")
plt.tight_layout()
plt.savefig("delay_distributions.png", dpi=150)
plt.show()
print("Saved visualization to delay_distributions.png")
What patterns should you look for?
The corridor box plot should show variation between corridors. Some corridors (like major arterials) will have higher median delays and wider spreads than quieter routes.
The hourly bar chart should show two humps: morning peak around hours 7-9 and evening peak around hours 16-18. This is the classic commute pattern the model's Fourier components will capture.
The day type plot should show weekdays with higher and more variable delays compared to weekends. If your data only covers weekdays, you may see little difference here.
- Save prepare_data.py and run it again:
python prepare_data.py
- A matplotlib window should appear showing three charts side by side. Close the window to let the script continue.
- Confirm the saved image exists by running:
ls delay_distributions.png
- You should see delay_distributions.png listed in the output.
No plot window appearing?
- If you are running on a remote server or headless environment, matplotlib may not have a display backend. Add matplotlib.use("Agg") before the import matplotlib.pyplot as plt line. The plot will still save to the PNG file.
- If the charts look empty or have very few data points, your collector may not have run long enough. You need at least a few hours of data across different times of day.
Help me get my matplotlib plots working in prepare_data.py
Prepare the data for PyMC
PyMC's hierarchical model needs corridors encoded as integer indices (not strings) and continuous features standardized for numerical stability during sampling.
- Add the encoding, standardization, and serialization code below the visualization print statement in prepare_data.py:
# Encode corridor_id as integer indices for PyMC
corridor_names = df["corridor_id"].cat.categories.tolist()
df["corridor_idx"] = df["corridor_id"].cat.codes.values
# Standardize hour_of_day for numerical stability
hour_mean = df["hour_of_day"].mean()
hour_std = df["hour_of_day"].std()
df["hour_scaled"] = (df["hour_of_day"] - hour_mean) / hour_std
# Create coordinate dictionaries for PyMC model dimensions
coords = {
"corridor": corridor_names,
"obs": np.arange(len(df)),
}
# Bundle everything the modeling step needs
model_data = {
"df": df,
"corridor_names": corridor_names,
"coords": coords,
"hour_mean": hour_mean,
"hour_std": hour_std,
"n_corridors": len(corridor_names),
}
# Save prepared data for the modeling step
with open("model_data.pkl", "wb") as f:
pickle.dump(model_data, f)
print(f"Prepared {len(df)} observations across {len(corridor_names)} corridors")
print(f"Saved model_data.pkl")
Why these transformations?
corridor_idx converts corridor names like "Great South Road SB" into integers (0, 1, 2...). PyMC uses these integers to index into the hierarchical parameter arrays.
Standardizing hour_of_day to zero mean and unit variance helps the NUTS sampler explore the posterior more efficiently. Raw hour values (0-23) can cause slow convergence.
The coords dictionary maps dimension names to their labels. When PyMC builds the model, these coordinates let you refer to corridors by name instead of raw index numbers.
Why save to a pickle file?
Serializing to model_data.pkl means you only run the data loading and preparation once. The modeling script in the next step loads this file directly, so you can iterate on the model without re-downloading from S3 every time.
- Save prepare_data.py and run it one final time:
python prepare_data.py
- After the plot window closes, you should see Prepared X observations across Y corridors and Saved model_data.pkl printed in the terminal.
- Confirm the pickle file exists by running:
ls model_data.pkl
- You should see model_data.pkl listed in the output.
Script errors out during the encoding step?
- If you see AttributeError: 'Series' object has no attribute 'cat', the astype("category") line from the feature engineering chunk did not run. Make sure it appears before the encoding code.
- If model_data.pkl is 0 bytes, check that the with open block uses "wb" mode (write binary).
Help me fix errors in the data preparation section of prepare_data.py
✔️ Awesome, I've got everything!
Great. Double check you have saved prepare_data.py and that both delay_distributions.png and model_data.pkl exist in your project folder.
ⓧ I'd like to double check the full code
Here is the complete prepare_data.py file. Compare it against yours and fix any differences.
import boto3
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pickle
from io import StringIO
# Connect to S3 and list all CSV files in the bucket
s3 = boto3.client("s3")
bucket_name = "[[BUCKET_NAME="auckland-congestion-data"]]"
response = s3.list_objects_v2(Bucket=bucket_name)
csv_keys = [obj["Key"] for obj in response["Contents"] if obj["Key"].endswith(".csv")]
# Download and concatenate all CSV files into one DataFrame
frames = []
for key in csv_keys:
obj = s3.get_object(Bucket=bucket_name, Key=key)
df_part = pd.read_csv(StringIO(obj["Body"].read().decode("utf-8")))
frames.append(df_part)
df = pd.concat(frames, ignore_index=True)
print(f"Loaded {len(df)} delay observations from {len(csv_keys)} files")
# Parse timestamp into datetime for feature extraction
df["timestamp"] = pd.to_datetime(df["timestamp"])
# Create hour_of_day feature (0-23)
df["hour_of_day"] = df["timestamp"].dt.hour
# Create binary weekday/weekend feature (1 = weekday, 0 = weekend)
df["day_of_week"] = (df["timestamp"].dt.dayofweek < 5).astype(int)
# Convert corridor_id to categorical type
df["corridor_id"] = df["corridor_id"].astype("category")
print(f"Corridors: {df['corridor_id'].nunique()}")
print(f"Date range: {df['timestamp'].min()} to {df['timestamp'].max()}")
print(df[["delay_seconds", "hour_of_day", "day_of_week"]].describe())
# Plot delay distributions by corridor, hour, and day type
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
# Box plot of delays per corridor
df.boxplot(column="delay_seconds", by="corridor_id", ax=axes[0], rot=90)
axes[0].set_title("Delay by Corridor")
axes[0].set_ylabel("Delay (seconds)")
# Bar chart of mean delay per hour
hourly = df.groupby("hour_of_day")["delay_seconds"].mean()
axes[1].bar(hourly.index, hourly.values)
axes[1].set_title("Mean Delay by Hour of Day")
axes[1].set_xlabel("Hour")
axes[1].set_ylabel("Mean Delay (seconds)")
# Box plot comparing weekday vs weekend
df.boxplot(column="delay_seconds", by="day_of_week", ax=axes[2])
axes[2].set_xticklabels(["Weekend", "Weekday"])
axes[2].set_title("Delay by Day Type")
axes[2].set_ylabel("Delay (seconds)")
plt.tight_layout()
plt.savefig("delay_distributions.png", dpi=150)
plt.show()
print("Saved visualization to delay_distributions.png")
# Encode corridor_id as integer indices for PyMC
corridor_names = df["corridor_id"].cat.categories.tolist()
df["corridor_idx"] = df["corridor_id"].cat.codes.values
# Standardize hour_of_day for numerical stability
hour_mean = df["hour_of_day"].mean()
hour_std = df["hour_of_day"].std()
df["hour_scaled"] = (df["hour_of_day"] - hour_mean) / hour_std
# Create coordinate dictionaries for PyMC model dimensions
coords = {
"corridor": corridor_names,
"obs": np.arange(len(df)),
}
# Bundle everything the modeling step needs
model_data = {
"df": df,
"corridor_names": corridor_names,
"coords": coords,
"hour_mean": hour_mean,
"hour_std": hour_std,
"n_corridors": len(corridor_names),
}
# Save prepared data for the modeling step
with open("model_data.pkl", "wb") as f:
pickle.dump(model_data, f)
print(f"Prepared {len(df)} observations across {len(corridor_names)} corridors")
print(f"Saved model_data.pkl")
Your data is loaded, visualized, and packaged. Next up, you will build the hierarchical Bayesian model in PyMC that turns these delay observations into corridor-level congestion estimates with full uncertainty quantification.
Build the Hierarchical Bayesian Model in PyMC
Your data is clean, encoded, and ready for modeling. Now comes the core of the project: building the Bayesian model that turns raw bus delays into corridor-level congestion estimates.
The key challenge is that some corridors have hundreds of observations while others have only a handful. A standard regression would either overfit the sparse corridors or ignore them entirely. Hierarchical partial pooling solves this by letting data-poor corridors borrow strength from the network-wide estimate.
In this step, get ready to:
- Define a hierarchical Bayesian model with corridor-level partial pooling in PyMC.
- Add Fourier time-of-day effects and a day-of-week covariate to capture commute patterns.
- Sample the posterior with NUTS and validate convergence through diagnostics and posterior predictive checks.
Define the hierarchical model structure
Each corridor gets its own congestion intercept, but all intercepts are drawn from a shared group-level distribution. Corridors with sparse data get pulled toward the network-wide mean. Corridors with lots of data keep their individual estimates.
What is partial pooling?
Imagine three approaches to estimating congestion per corridor. "No pooling" fits each corridor independently, which overfits corridors with few bus routes. "Complete pooling" ignores corridor identity and fits one global estimate, which underfits corridors with distinctive congestion patterns.
"Partial pooling" sits between these extremes. Each corridor has its own intercept, but all intercepts are drawn from a shared group distribution. The model learns both the individual corridor effects and the group-level spread. Corridors with thin data shrink toward the group mean automatically.
- Create a new file called model.py in the same folder as prepare_data.py.
- Start by adding the imports and data loading at the top of model.py:
import pickle
import numpy as np
import pymc as pm
import arviz as az
import matplotlib.pyplot as plt
# Load the prepared dataset from the previous step
with open("model_data.pkl", "rb") as f:
data = pickle.load(f)
corridor_idx = data["corridor_idx"]
hour_of_day = data["hour_of_day"]
day_of_week = data["day_of_week"]
delay_seconds = data["delay_seconds"]
coords = data["coords"]
n_corridors = len(coords["corridor"])
print(f"Loaded {len(delay_seconds)} observations across {n_corridors} corridors")
What does this code do?
This loads the prepared dataset you saved in the previous step. The corridor_idx array maps each observation to an integer corridor ID. The coords dictionary tells PyMC the names of each corridor dimension, which makes the output readable.
- Below the data loading code, add the model definition with group-level hyperpriors and corridor intercepts:
# Define the hierarchical Bayesian model
with pm.Model(coords=coords) as congestion_model:
# Register observed data as mutable containers for later prediction
corridor_idx_data = pm.MutableData("corridor_idx_data", corridor_idx)
hour_data = pm.MutableData("hour_data", hour_of_day)
day_data = pm.MutableData("day_data", day_of_week)
# Group-level hyperpriors: the shared distribution all corridors are drawn from
mu_corridor = pm.Normal("mu_corridor", mu=0, sigma=50)
sigma_corridor = pm.HalfNormal("sigma_corridor", sigma=30)
# Corridor-level intercepts with partial pooling
# Each corridor gets its own intercept, drawn from the group distribution
corridor_intercept = pm.Normal(
"corridor_intercept",
mu=mu_corridor,
sigma=sigma_corridor,
dims="corridor"
)
What does this code do?
The pm.MutableData containers wrap your input arrays so you can swap in new data at prediction time without redefining the model.
The two hyperpriors (mu_corridor and sigma_corridor) define the group-level distribution. Every individual corridor_intercept is drawn from this shared distribution. The dims="corridor" argument tells PyMC to create one intercept per corridor, labeled with the corridor names from your coordinates.
- Save model.py.
You will not see output yet. The model definition is incomplete until you add the time effects and likelihood in the next substep.
Add time-of-day and day-of-week effects
Congestion follows a daily cycle. Rush hour peaks at 7-9am and 4-6pm, and weekends look different from weekdays. A Fourier series captures this cyclical pattern using sine and cosine components at 24-hour and 12-hour periods.
Why a Fourier series instead of hour dummies?
You could create a separate parameter for each hour of the day, but that would add 24 parameters and assume adjacent hours are independent. A Fourier series captures the same daily cycle with just 4 parameters (two sine/cosine pairs). It enforces smooth transitions between hours, which matches how congestion actually builds and fades.
- In model.py, find the corridor_intercept block (ending with dims="corridor"). Add the following code below it, still inside the with pm.Model() block:
# Fourier series for time-of-day effects (24h and 12h periods)
beta_sin24 = pm.Normal("beta_sin24", mu=0, sigma=10)
beta_cos24 = pm.Normal("beta_cos24", mu=0, sigma=10)
beta_sin12 = pm.Normal("beta_sin12", mu=0, sigma=10)
beta_cos12 = pm.Normal("beta_cos12", mu=0, sigma=10)
# Combine the Fourier components into a single time effect
fourier_time = (
beta_sin24 * np.sin(2 * np.pi * hour_data / 24)
+ beta_cos24 * np.cos(2 * np.pi * hour_data / 24)
+ beta_sin12 * np.sin(2 * np.pi * hour_data / 12)
+ beta_cos12 * np.cos(2 * np.pi * hour_data / 12)
)
# Day-of-week effect (weekday vs weekend)
beta_day = pm.Normal("beta_day", mu=0, sigma=10)
# Expected delay combines corridor intercept, time, and day effects
mu_delay = pm.Deterministic(
"mu_delay",
corridor_intercept[corridor_idx_data] + fourier_time + beta_day * day_data
)
# Observation noise
sigma_obs = pm.HalfNormal("sigma_obs", sigma=30)
# Likelihood: observed delay given the expected delay
delay_obs = pm.Normal(
"delay_obs",
mu=mu_delay,
sigma=sigma_obs,
observed=delay_seconds
)
What does this code do?
The four beta_sin/beta_cos parameters capture the daily commute cycle. The 24-hour components model the overall day/night pattern. The 12-hour components add the double-peaked rush hour shape (morning and evening).
The mu_delay deterministic combines all three effects: the corridor-specific intercept, the Fourier time component, and the weekday/weekend shift. The likelihood (delay_obs) connects this expected delay to your actual observations through a Normal distribution with learned noise sigma_obs.
- Save model.py.
The model structure is now complete. You will not see output yet. The next substep runs the sampler.
Seeing an indentation error?
All the code from corridor_idx_data through delay_obs must be indented one level inside the with pm.Model() block. Check that every line uses four spaces of indentation, not a mix of tabs and spaces.
Help me fix an indentation error in my PyMC model definition
Sample the posterior and run diagnostics
With the model fully defined, you can now run the NUTS sampler to draw from the posterior distribution. After sampling, you need to verify that the chains actually converged.
- In model.py, find the delay_obs = pm.Normal( block. Add the following line below its closing parenthesis, still inside the with pm.Model() block:
# Sample the posterior using the NUTS sampler
idata = pm.sample(2000, tune=1000, cores=2, random_seed=42)
- Below the with pm.Model() block (back at the top indentation level with no indent), add the diagnostics code:
# Print convergence summary for all parameters
summary = az.summary(idata, var_names=[
"mu_corridor", "sigma_corridor", "corridor_intercept",
"beta_sin24", "beta_cos24", "beta_sin12", "beta_cos12",
"beta_day", "sigma_obs"
])
print(summary)
# Check that all R-hat values indicate convergence
rhat_values = summary["r_hat"]
print(f"\nMax R-hat: {rhat_values.max():.4f}")
if rhat_values.max() < 1.01:
print("Convergence confirmed: all R-hat values below 1.01")
else:
print("WARNING: some R-hat values exceed 1.01, investigate trace plots")
# Plot trace plots for key parameters
az.plot_trace(idata, var_names=[
"mu_corridor", "sigma_corridor", "beta_sin24", "beta_day"
])
plt.tight_layout()
plt.savefig("trace_plots.png", dpi=150)
plt.show()
What are R-hat and trace plots?
R-hat (also called the Gelman-Rubin statistic) compares within-chain variance to between-chain variance. If all chains converged to the same distribution, R-hat should be very close to 1.0. Values below 1.01 indicate good convergence.
Trace plots show the sampled values across iterations. Healthy traces look like "fuzzy caterpillars" that mix well. If you see chains stuck in different regions or trending upward/downward, that signals a convergence problem.
- Now add the posterior predictive checks below the diagnostics code. This generates simulated data from the posterior and compares it to your real observations:
# Run posterior predictive checks
with congestion_model:
idata = pm.sample_posterior_predictive(
idata, extend_inferencedata=True
)
# Visualize how well the model reproduces the observed data
az.plot_ppc(idata, num_pp_samples=100)
plt.savefig("ppc_plot.png", dpi=150)
plt.show()
print("Model fitting and diagnostics complete")
What is a posterior predictive check?
A posterior predictive check answers: "if the model were true, would the generated data look like the real data?" The pm.sample_posterior_predictive function simulates new delay observations using the fitted posterior. The extend_inferencedata=True flag adds these simulated samples directly to your existing idata object.
The az.plot_ppc visualization overlays 100 simulated datasets on top of the real data. If the model fits well, the simulated distributions should cover the observed distribution closely.
- Save model.py.
- Run the full script by executing:
python model.py
This will take several minutes
The NUTS sampler runs 2000 draws with 1000 tuning steps across 2 chains. Depending on your dataset size and machine, this can take 5-15 minutes. You will see a progress bar in your terminal showing the sampling speed and estimated time remaining.
When sampling finishes, you should see a summary table printed to the terminal with columns for mean, sd, r_hat, and ess_bulk for each parameter. Your max R-hat should be below 1.01. Two matplotlib windows will pop up: the trace plots and the PPC overlay.
R-hat values above 1.01?
High R-hat values mean the chains have not converged. Try increasing the tuning steps by changing tune=1000 to tune=2000 in the pm.sample() call. If you have very few observations for some corridors (fewer than 10), consider collecting more data before fitting.
Help me diagnose why my PyMC model has high R-hat values
Sampling is extremely slow or freezing?
If your dataset has more than 50,000 observations, sampling will be slow. Try reducing to a random subsample for initial fitting. You can also set cores=1 if parallel sampling causes issues on your machine.
Help me speed up PyMC sampling for my congestion model
✔️ Awesome, I've got everything!
Great. Double check you have saved model.py and confirmed all R-hat values are below 1.01.
ⓧ I'd like to double check the full code
Here is the complete model.py file. Compare it against yours to spot any differences.
import pickle
import numpy as np
import pymc as pm
import arviz as az
import matplotlib.pyplot as plt
# Load the prepared dataset from the previous step
with open("model_data.pkl", "rb") as f:
data = pickle.load(f)
corridor_idx = data["corridor_idx"]
hour_of_day = data["hour_of_day"]
day_of_week = data["day_of_week"]
delay_seconds = data["delay_seconds"]
coords = data["coords"]
n_corridors = len(coords["corridor"])
print(f"Loaded {len(delay_seconds)} observations across {n_corridors} corridors")
# Define the hierarchical Bayesian model
with pm.Model(coords=coords) as congestion_model:
# Register observed data as mutable containers for later prediction
corridor_idx_data = pm.MutableData("corridor_idx_data", corridor_idx)
hour_data = pm.MutableData("hour_data", hour_of_day)
day_data = pm.MutableData("day_data", day_of_week)
# Group-level hyperpriors: the shared distribution all corridors are drawn from
mu_corridor = pm.Normal("mu_corridor", mu=0, sigma=50)
sigma_corridor = pm.HalfNormal("sigma_corridor", sigma=30)
# Corridor-level intercepts with partial pooling
# Each corridor gets its own intercept, drawn from the group distribution
corridor_intercept = pm.Normal(
"corridor_intercept",
mu=mu_corridor,
sigma=sigma_corridor,
dims="corridor"
)
# Fourier series for time-of-day effects (24h and 12h periods)
beta_sin24 = pm.Normal("beta_sin24", mu=0, sigma=10)
beta_cos24 = pm.Normal("beta_cos24", mu=0, sigma=10)
beta_sin12 = pm.Normal("beta_sin12", mu=0, sigma=10)
beta_cos12 = pm.Normal("beta_cos12", mu=0, sigma=10)
# Combine the Fourier components into a single time effect
fourier_time = (
beta_sin24 * np.sin(2 * np.pi * hour_data / 24)
+ beta_cos24 * np.cos(2 * np.pi * hour_data / 24)
+ beta_sin12 * np.sin(2 * np.pi * hour_data / 12)
+ beta_cos12 * np.cos(2 * np.pi * hour_data / 12)
)
# Day-of-week effect (weekday vs weekend)
beta_day = pm.Normal("beta_day", mu=0, sigma=10)
# Expected delay combines corridor intercept, time, and day effects
mu_delay = pm.Deterministic(
"mu_delay",
corridor_intercept[corridor_idx_data] + fourier_time + beta_day * day_data
)
# Observation noise
sigma_obs = pm.HalfNormal("sigma_obs", sigma=30)
# Likelihood: observed delay given the expected delay
delay_obs = pm.Normal(
"delay_obs",
mu=mu_delay,
sigma=sigma_obs,
observed=delay_seconds
)
# Sample the posterior using the NUTS sampler
idata = pm.sample(2000, tune=1000, cores=2, random_seed=42)
# Print convergence summary for all parameters
summary = az.summary(idata, var_names=[
"mu_corridor", "sigma_corridor", "corridor_intercept",
"beta_sin24", "beta_cos24", "beta_sin12", "beta_cos12",
"beta_day", "sigma_obs"
])
print(summary)
# Check that all R-hat values indicate convergence
rhat_values = summary["r_hat"]
print(f"\nMax R-hat: {rhat_values.max():.4f}")
if rhat_values.max() < 1.01:
print("Convergence confirmed: all R-hat values below 1.01")
else:
print("WARNING: some R-hat values exceed 1.01, investigate trace plots")
# Plot trace plots for key parameters
az.plot_trace(idata, var_names=[
"mu_corridor", "sigma_corridor", "beta_sin24", "beta_day"
])
plt.tight_layout()
plt.savefig("trace_plots.png", dpi=150)
plt.show()
# Run posterior predictive checks
with congestion_model:
idata = pm.sample_posterior_predictive(
idata, extend_inferencedata=True
)
# Visualize how well the model reproduces the observed data
az.plot_ppc(idata, num_pp_samples=100)
plt.savefig("ppc_plot.png", dpi=150)
plt.show()
print("Model fitting and diagnostics complete")
Your hierarchical model is fitted and validated. The idata object holds thousands of posterior samples for every parameter in the model. Next up, you will dig into these posteriors to extract corridor-level congestion insights and save the fitted model for deployment.
Analyze the Posterior and Extract Insights
In the previous step you fitted your hierarchical Bayesian model and confirmed that the sampler converged. Now it is time to interrogate the posterior and understand what the model actually learned about Auckland congestion.
You will examine the group-level hyperparameters to see how much congestion varies across the network, compare individual corridor effects with a forest plot, and reconstruct the daily congestion curve from the Fourier components. At the end, you will save the entire posterior to a NetCDF file so the Lambda function can serve predictions without ever re-fitting the model.
In this step, get ready to:
- Examine group-level hyperparameters and compare corridor-level effects.
- Visualize the time-of-day congestion pattern from Fourier components.
- Save the fitted posterior to a NetCDF file for serverless deployment.
Examine group-level hyperparameters and corridor effects
The group-level hyperparameters mu_corridor and sigma_corridor control how the model pools information across corridors. A large sigma_corridor means corridors differ substantially from each other, while a small one means partial pooling is pulling them toward a common mean.
Before you can analyze the posterior in a separate script, you need to make the InferenceData object from model.py available on disk.
- In model.py, scroll to the very end of the file (after the posterior predictive check code).
- Find the last line of the script.
- Add these lines at the very end of model.py to pickle the inference data for the analysis script:
# Save inference data so the analysis script can load it
import pickle
with open("model_idata.pkl", "wb") as f:
pickle.dump(idata, f)
print("Saved idata to model_idata.pkl")
- Save model.py and re-run it by running this command:
python model.py
You should see all the previous model output, followed by Saved idata to model_idata.pkl at the end.
Model taking too long to re-run?
If you already have a successful model run from the previous step and do not want to wait for sampling again, you can temporarily reduce the number of draws (e.g. 500 instead of 1000) for this run. The purpose here is just to save the inference data to disk.
Help me speed up my model re-run so I can save idata to disk.
Now create the analysis script. This script loads the saved inference data and the prepared data (for corridor name mapping), then examines the posterior.
- Create a new file called analyze.py in your project folder.
- Add the imports, data loading, and hyperparameter summary by copying this code into analyze.py:
import pickle
import numpy as np
import matplotlib.pyplot as plt
import arviz as az
# Load prepared data for corridor name mapping
with open("model_data.pkl", "rb") as f:
model_data = pickle.load(f)
corridor_names = model_data["corridor_names"]
# Load the InferenceData saved by model.py
with open("model_idata.pkl", "rb") as f:
idata = pickle.load(f)
# --- Group-level hyperparameters ---
# Summarize the network-wide mean and spread of corridor intercepts
print("=== Group-Level Hyperparameters ===")
print(az.summary(idata, var_names=["mu_corridor", "sigma_corridor"]))
# --- Corridor-level forest plot ---
# Compare each corridor's congestion intercept side by side
az.plot_forest(
idata,
var_names=["corridor_offset"],
combined=True,
figsize=(9, 7)
)
plt.title("Corridor-Level Congestion Intercepts")
plt.tight_layout()
plt.savefig("forest_corridors.png", dpi=150)
plt.show()
What does this code do?
az.summary(idata, var_names=["mu_corridor", "sigma_corridor"]) prints a table showing the posterior mean, standard deviation, credible intervals, effective sample size, and R-hat for each hyperparameter. These tell you the network-wide average corridor congestion and how much corridors differ from each other.
az.plot_forest() creates a forest plot showing the posterior median and HDI (highest density interval) for each corridor's offset. The combined=True parameter merges all chains into a single distribution per corridor. Corridors with wide intervals have more uncertainty, often because they had fewer bus observations.
- Save analyze.py and run it by running this command:
python analyze.py
You should see a summary table printed in the terminal with columns for mean, sd, hdi_3%, hdi_97%, ess_bulk, and r_hat. A forest plot window should also appear showing all corridor intercepts with their credible intervals.
What should I look for in the summary?
Check the sigma_corridor row. If its posterior mean is large relative to mu_corridor, corridors genuinely differ in baseline congestion. If sigma_corridor is close to zero, the model has pooled all corridors toward a single estimate.
In the forest plot, corridors with narrow intervals have lots of bus data. Corridors with wide intervals are borrowing heavily from the group-level prior. This is partial pooling in action.
Seeing a KeyError for corridor_offset?
Make sure the variable name in var_names matches the name you used in model.py when defining the corridor offset. Check by printing list(idata.posterior.data_vars) to see all variable names in your posterior.
Help me fix a KeyError when plotting corridor offsets from my PyMC posterior.
Visualize the time-of-day congestion pattern
Your model captures the daily congestion cycle using a Fourier series with 24-hour and 12-hour components. By combining these posterior coefficients across a 24-hour axis, you can reconstruct the learned congestion curve and see whether the model picked up Auckland's known commute peaks at 7-9am and 4-6pm.
- In analyze.py, add the Fourier time-of-day visualization code at the bottom of the file:
# --- Time-of-day Fourier visualization ---
# Extract posterior means for each Fourier coefficient
posterior = idata.posterior
beta_sin24 = float(posterior["beta_sin24"].mean(dim=["chain", "draw"]))
beta_cos24 = float(posterior["beta_cos24"].mean(dim=["chain", "draw"]))
beta_sin12 = float(posterior["beta_sin12"].mean(dim=["chain", "draw"]))
beta_cos12 = float(posterior["beta_cos12"].mean(dim=["chain", "draw"]))
# Generate a 24-hour time axis
hours = np.linspace(0, 24, 200)
# Reconstruct the Fourier time-of-day effect
time_effect = (
beta_sin24 * np.sin(2 * np.pi * hours / 24)
+ beta_cos24 * np.cos(2 * np.pi * hours / 24)
+ beta_sin12 * np.sin(2 * np.pi * hours / 12)
+ beta_cos12 * np.cos(2 * np.pi * hours / 12)
)
# Plot the daily congestion pattern with known peak windows
plt.figure(figsize=(10, 5))
plt.plot(hours, time_effect, linewidth=2, color="steelblue")
plt.axvspan(7, 9, alpha=0.2, color="red", label="AM peak (7-9)")
plt.axvspan(16, 18, alpha=0.2, color="orange", label="PM peak (4-6)")
plt.xlabel("Hour of Day")
plt.ylabel("Congestion Effect (delay seconds)")
plt.title("Time-of-Day Congestion Pattern (Fourier Components)")
plt.legend()
plt.tight_layout()
plt.savefig("time_of_day_effect.png", dpi=150)
plt.show()
print("Saved time-of-day plot to time_of_day_effect.png")
What does this code do?
The four Fourier coefficients (beta_sin24, beta_cos24, beta_sin12, beta_cos12) combine sin and cos waves at 24-hour and 12-hour periods. The 24-hour pair captures the broad morning/evening cycle. The 12-hour pair captures the double-peaked commute pattern where congestion rises twice daily.
The red and orange shaded regions mark Auckland's known peak commute windows (7-9am and 4-6pm). If the model learned well, the curve should peak within or near these shaded regions.
- Save analyze.py and run it again by running this command:
python analyze.py
You should see the forest plot from before, followed by a new line plot showing the reconstructed daily congestion curve with the AM and PM peak windows shaded.
What should the curve look like?
A well-fitted model typically shows a curve that rises in the morning hours, peaks around 7-9am, dips around midday, rises again in the late afternoon (4-6pm), then drops to a minimum overnight. The exact shape depends on your data. If your collection period was dominated by one time of day, the curve may not fully resolve both peaks.
Curve looks flat or unexpected?
If the time-of-day curve is nearly flat, your data may not have enough variation across different hours. Check that your delay data covers both peak and off-peak periods by printing the distribution of hour_of_day in your original dataset. You may need to collect more data covering a wider range of times.
Help me understand why my Fourier time-of-day curve is flat.
Save the posterior for deployment
The Lambda function you will build in the next step needs the posterior samples to compute predictions. ArviZ's NetCDF format is the standard way to serialize InferenceData. It is portable, self-describing, and can be loaded with just az.from_netcdf().
- In analyze.py, add these lines at the very end of the file to save the posterior as NetCDF:
# --- Save posterior for Lambda deployment ---
# NetCDF is the standard format for ArviZ InferenceData
# This file contains all posterior samples needed for prediction
idata.to_netcdf("model_trace.nc")
print("Saved posterior samples to model_trace.nc")
Why NetCDF instead of pickle?
Pickle files are Python-version-specific and can break across environments. NetCDF is a cross-platform binary format designed for scientific array data. It preserves all the xarray metadata (dimensions, coordinates, attributes) and can be loaded with az.from_netcdf() on any machine with ArviZ installed.
The Lambda container only needs numpy, scipy, and netCDF4 to read this file. It does not need PyMC or PyTensor at inference time, which keeps the container small and fast.
- Save analyze.py and run it one more time by running this command:
python analyze.py
- Confirm the file was created by running:
ls -lh model_trace.nc
You should see model_trace.nc listed with a file size typically in the range of a few MB to tens of MB, depending on how many corridors and posterior draws you have.
File not found or zero bytes?
Make sure your script ran without errors by checking the terminal output. Look for the Saved posterior samples to model_trace.nc message. If it is missing, the script likely errored before reaching the save line. Check for tracebacks in the output above.
Also confirm you have the netCDF4 package installed. ArviZ uses it under the hood for .to_netcdf(). Install it if missing by running pip install netCDF4.
Help me fix issues saving my InferenceData to NetCDF.
✔️ Awesome, I've got everything!
Great. Double check you have saved analyze.py and that model_trace.nc exists in your project folder.
ⓧ I'd like to double check the full code
Here is the complete analyze.py file. Compare it against yours to make sure everything matches.
import pickle
import numpy as np
import matplotlib.pyplot as plt
import arviz as az
# Load prepared data for corridor name mapping
with open("model_data.pkl", "rb") as f:
model_data = pickle.load(f)
corridor_names = model_data["corridor_names"]
# Load the InferenceData saved by model.py
with open("model_idata.pkl", "rb") as f:
idata = pickle.load(f)
# --- Group-level hyperparameters ---
# Summarize the network-wide mean and spread of corridor intercepts
print("=== Group-Level Hyperparameters ===")
print(az.summary(idata, var_names=["mu_corridor", "sigma_corridor"]))
# --- Corridor-level forest plot ---
# Compare each corridor's congestion intercept side by side
az.plot_forest(
idata,
var_names=["corridor_offset"],
combined=True,
figsize=(9, 7)
)
plt.title("Corridor-Level Congestion Intercepts")
plt.tight_layout()
plt.savefig("forest_corridors.png", dpi=150)
plt.show()
# --- Time-of-day Fourier visualization ---
# Extract posterior means for each Fourier coefficient
posterior = idata.posterior
beta_sin24 = float(posterior["beta_sin24"].mean(dim=["chain", "draw"]))
beta_cos24 = float(posterior["beta_cos24"].mean(dim=["chain", "draw"]))
beta_sin12 = float(posterior["beta_sin12"].mean(dim=["chain", "draw"]))
beta_cos12 = float(posterior["beta_cos12"].mean(dim=["chain", "draw"]))
# Generate a 24-hour time axis
hours = np.linspace(0, 24, 200)
# Reconstruct the Fourier time-of-day effect
time_effect = (
beta_sin24 * np.sin(2 * np.pi * hours / 24)
+ beta_cos24 * np.cos(2 * np.pi * hours / 24)
+ beta_sin12 * np.sin(2 * np.pi * hours / 12)
+ beta_cos12 * np.cos(2 * np.pi * hours / 12)
)
# Plot the daily congestion pattern with known peak windows
plt.figure(figsize=(10, 5))
plt.plot(hours, time_effect, linewidth=2, color="steelblue")
plt.axvspan(7, 9, alpha=0.2, color="red", label="AM peak (7-9)")
plt.axvspan(16, 18, alpha=0.2, color="orange", label="PM peak (4-6)")
plt.xlabel("Hour of Day")
plt.ylabel("Congestion Effect (delay seconds)")
plt.title("Time-of-Day Congestion Pattern (Fourier Components)")
plt.legend()
plt.tight_layout()
plt.savefig("time_of_day_effect.png", dpi=150)
plt.show()
print("Saved time-of-day plot to time_of_day_effect.png")
# --- Save posterior for Lambda deployment ---
# NetCDF is the standard format for ArviZ InferenceData
# This file contains all posterior samples needed for prediction
idata.to_netcdf("model_trace.nc")
print("Saved posterior samples to model_trace.nc")
Your posterior is analyzed and saved. You now have a clear picture of which corridors are most congested, how congestion cycles through the day, and a portable model_trace.nc file ready for deployment. Next up, you will build a Lambda function that loads this file and serves live congestion predictions through an API.
Build the Prediction Lambda Function
Your posterior is fitted, diagnosed, and saved as model_trace.nc. Now you need a way for the outside world to query it.
Running PyMC inside a Lambda function would be far too slow and far too large. The trick is that you already have thousands of posterior draws saved in your NetCDF file. Prediction is just arithmetic on those draws: add the corridor intercept, the Fourier time components, and the day effect. That is fast, lightweight, and needs nothing more than NumPy.
In this step, get ready to:
- Write a Lambda handler that loads saved posterior samples and computes congestion predictions.
- Create a Dockerfile that packages the handler and model for AWS Lambda.
- Build the Docker container image locally.
Set up the Lambda handler's model loading
AWS Lambda keeps a function warm between invocations. Code at the global scope runs once during a cold start and is reused for every subsequent request. Loading the posterior samples at global scope means only the first invocation pays the I/O cost.
- Create a new file called lambda_handler.py in the same folder as your model_trace.nc file.
- Add the imports and global model loading code by copying the following into lambda_handler.py:
import json
import numpy as np
import arviz as az
# Load posterior samples once during cold start
# Lambda reuses this across warm invocations
idata = az.from_netcdf("model_trace.nc")
# Extract posterior sample arrays from the InferenceData
corridor_intercepts = idata.posterior["corridor_intercept"].values
beta_sin_24 = idata.posterior["beta_sin_24"].values
beta_cos_24 = idata.posterior["beta_cos_24"].values
beta_sin_12 = idata.posterior["beta_sin_12"].values
beta_cos_12 = idata.posterior["beta_cos_12"].values
beta_day = idata.posterior["beta_day"].values
# Build corridor name-to-index mapping from model coordinates
corridor_names = list(idata.posterior.coords["corridor"].values)
corridor_map = {name: idx for idx, name in enumerate(corridor_names)}
What does this code do?
az.from_netcdf("model_trace.nc") loads the full InferenceData object you saved in the previous step. This contains every posterior draw from the NUTS sampler.
Each idata.posterior["..."].values call extracts one parameter's draws as a NumPy array. The shape is (chains, draws) for scalar coefficients and (chains, draws, corridors) for the corridor intercepts.
The corridor map is built dynamically from the coordinates stored in the posterior. This means it automatically matches whatever corridor names you defined in your PyMC model.
Match your variable names
The strings "corridor_intercept", "beta_sin_24", etc. must match exactly the variable names you used in your pm.Model() block. If you named them differently, update them here to match.
You won't see output yet. The next chunk adds the handler function that uses these loaded samples.
Add the prediction handler function
The handler function is the entry point that Lambda calls for each request. It takes a corridor ID and time of day, then combines the saved posterior draws into a full predictive distribution.
- In lambda_handler.py, add the following function below the corridor_map line:
def handler(event, context):
"""Predict congestion for a given corridor and time of day."""
# Parse input from API Gateway
body = json.loads(event.get("body", "{}"))
corridor_id = body["corridor_id"]
hour_of_day = float(body["hour_of_day"])
is_weekday = body.get("is_weekday", True)
# Validate the corridor exists in the model
if corridor_id not in corridor_map:
return {
"statusCode": 400,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({
"error": f"Unknown corridor: {corridor_id}",
"valid_corridors": list(corridor_map.keys())
})
}
corridor_idx = corridor_map[corridor_id]
# Compute Fourier components for the requested hour
hour_rad = 2 * np.pi * hour_of_day / 24.0
sin_24 = np.sin(hour_rad)
cos_24 = np.cos(hour_rad)
sin_12 = np.sin(2 * hour_rad)
cos_12 = np.cos(2 * hour_rad)
# Combine saved posterior draws into a predictive distribution
# Each draw yields one sample of the predicted congestion
intercept_draws = corridor_intercepts[:, :, corridor_idx]
prediction_draws = (
intercept_draws
+ beta_sin_24 * sin_24
+ beta_cos_24 * cos_24
+ beta_sin_12 * sin_12
+ beta_cos_12 * cos_12
+ beta_day * float(is_weekday)
)
# Flatten across chains and draws
flat_draws = prediction_draws.flatten()
# Compute summary statistics with full uncertainty
result = {
"corridor_id": corridor_id,
"hour_of_day": hour_of_day,
"is_weekday": is_weekday,
"posterior_mean": float(np.mean(flat_draws)),
"posterior_median": float(np.median(flat_draws)),
"credible_interval_5th": float(np.percentile(flat_draws, 5)),
"credible_interval_95th": float(np.percentile(flat_draws, 95)),
"n_samples": len(flat_draws)
}
return {
"statusCode": 200,
"headers": {"Content-Type": "application/json"},
"body": json.dumps(result)
}
What does this code do?
- The function parses corridor_id, hour_of_day, and an optional is_weekday flag from the incoming JSON body.
- The Fourier computation converts the hour into sine and cosine waves at 24-hour and 12-hour periods. These capture the daily commute cycle and the midday lull.
- prediction_draws adds the corridor's intercept, time-of-day effects, and day effect element-wise across all posterior draws. Each element is one plausible congestion value.
- The 5th and 95th percentiles form the 90% credible interval. This tells the caller not just the best guess, but how uncertain the model is.
Why return a full distribution instead of a point estimate?
A single predicted delay number hides how confident the model is. A corridor with 50 observations might have a tight credible interval (say 15 to 20 seconds). A corridor with 3 observations will have a wide one (say 5 to 40 seconds). Both could have the same mean, but the uncertainty tells the caller very different stories.
- Save lambda_handler.py.
You won't be able to test this handler directly yet. The Dockerfile in the next substep packages everything into a container that Lambda can run.
Seeing a KeyError when you eventually test?
A KeyError on a posterior variable name (like "corridor_intercept") means the name in your handler does not match the name in your PyMC model. Open your model.py and check the exact strings you passed to pm.Normal() and other distributions.
help me debug a KeyError when loading posterior variables from my model_trace.nc file in lambda_handler.py
✔️ Awesome, I've got everything!
Great. Double check you have saved lambda_handler.py.
ⓧ I'd like to double check the full code
import json
import numpy as np
import arviz as az
# Load posterior samples once during cold start
# Lambda reuses this across warm invocations
idata = az.from_netcdf("model_trace.nc")
# Extract posterior sample arrays from the InferenceData
corridor_intercepts = idata.posterior["corridor_intercept"].values
beta_sin_24 = idata.posterior["beta_sin_24"].values
beta_cos_24 = idata.posterior["beta_cos_24"].values
beta_sin_12 = idata.posterior["beta_sin_12"].values
beta_cos_12 = idata.posterior["beta_cos_12"].values
beta_day = idata.posterior["beta_day"].values
# Build corridor name-to-index mapping from model coordinates
corridor_names = list(idata.posterior.coords["corridor"].values)
corridor_map = {name: idx for idx, name in enumerate(corridor_names)}
def handler(event, context):
"""Predict congestion for a given corridor and time of day."""
# Parse input from API Gateway
body = json.loads(event.get("body", "{}"))
corridor_id = body["corridor_id"]
hour_of_day = float(body["hour_of_day"])
is_weekday = body.get("is_weekday", True)
# Validate the corridor exists in the model
if corridor_id not in corridor_map:
return {
"statusCode": 400,
"headers": {"Content-Type": "application/json"},
"body": json.dumps({
"error": f"Unknown corridor: {corridor_id}",
"valid_corridors": list(corridor_map.keys())
})
}
corridor_idx = corridor_map[corridor_id]
# Compute Fourier components for the requested hour
hour_rad = 2 * np.pi * hour_of_day / 24.0
sin_24 = np.sin(hour_rad)
cos_24 = np.cos(hour_rad)
sin_12 = np.sin(2 * hour_rad)
cos_12 = np.cos(2 * hour_rad)
# Combine saved posterior draws into a predictive distribution
# Each draw yields one sample of the predicted congestion
intercept_draws = corridor_intercepts[:, :, corridor_idx]
prediction_draws = (
intercept_draws
+ beta_sin_24 * sin_24
+ beta_cos_24 * cos_24
+ beta_sin_12 * sin_12
+ beta_cos_12 * cos_12
+ beta_day * float(is_weekday)
)
# Flatten across chains and draws
flat_draws = prediction_draws.flatten()
# Compute summary statistics with full uncertainty
result = {
"corridor_id": corridor_id,
"hour_of_day": hour_of_day,
"is_weekday": is_weekday,
"posterior_mean": float(np.mean(flat_draws)),
"posterior_median": float(np.median(flat_draws)),
"credible_interval_5th": float(np.percentile(flat_draws, 5)),
"credible_interval_95th": float(np.percentile(flat_draws, 95)),
"n_samples": len(flat_draws)
}
return {
"statusCode": 200,
"headers": {"Content-Type": "application/json"},
"body": json.dumps(result)
}
Create the Dockerfile
Lambda needs a container image that includes your handler code, the saved posterior file, and the Python packages for prediction. The key insight is that you do not need PyMC or PyTensor in this container. All the heavy sampling already happened offline.
- Create a file called Dockerfile (no extension) in the same folder as lambda_handler.py and model_trace.nc.
- Add the following to your Dockerfile:
FROM public.ecr.aws/lambda/python:3.11
# Install prediction dependencies (no PyMC needed at inference time)
RUN pip install numpy pandas arviz netCDF4 scipy --no-cache-dir
# Copy the handler code and saved posterior samples
COPY lambda_handler.py ${LAMBDA_TASK_ROOT}/
COPY model_trace.nc ${LAMBDA_TASK_ROOT}/
# Set the Lambda handler entry point
CMD ["lambda_handler.handler"]
What does this Dockerfile do?
- FROM public.ecr.aws/lambda/python:3.11 uses the official AWS Lambda Python 3.11 base image. This image includes the Lambda runtime interface client so your function can communicate with the Lambda service.
- --no-cache-dir tells pip not to store downloaded package files. This keeps the container image smaller.
- ${LAMBDA_TASK_ROOT} is a built-in environment variable in the Lambda base image that points to the working directory where Lambda looks for your handler code.
- CMD ["lambda_handler.handler"] tells Lambda to call the handler function inside the lambda_handler module.
Why a container instead of a Lambda zip package?
Lambda zip deployments have a 50 MB limit for layers. NumPy, pandas, ArviZ, and SciPy together exceed that easily. Container images support up to 10 GB, which gives you plenty of room.
- Save the Dockerfile.
- Confirm all three files are in the same directory by running:
ls
You should see Dockerfile, lambda_handler.py, and model_trace.nc all listed in the output.
Missing model_trace.nc?
If model_trace.nc is not in the same folder as your Dockerfile, copy it over from wherever you saved it in the previous step. The Docker build context only sees files in the current directory.
help me find and copy my model_trace.nc file into the correct directory for the Docker build
Build the Docker image
With all three files in place, you can build the container image. Docker will pull the Lambda base image, install the Python packages, and copy your handler and model file into the image.
- Build the container image by running this command:
docker build -t auckland-congestion-lambda:latest .
This takes a few minutes the first time as Docker downloads the base image and installs packages. You should see each Dockerfile step complete with a numbered output like Step 1/5, Step 2/5, etc.
- Verify the image was built by running:
docker images auckland-congestion-lambda
You should see a row showing auckland-congestion-lambda with the tag latest and a size in the hundreds of MB range.
Docker build failing?
- If you see Cannot connect to the Docker daemon, make sure Docker Desktop is running. Open it from your system tray or start it through OS search.
- If pip install fails inside the build, check your internet connection. Docker needs to download packages during the build.
- If you see COPY failed: file not found, make sure you are running the docker build command from the same folder that contains lambda_handler.py and model_trace.nc.
help me debug my Docker build for the Lambda congestion container
Your prediction container is built and ready to go. Next up, you will push it to Amazon ECR, create the Lambda function, and wire it up to an API Gateway endpoint so anyone can query your congestion model over HTTPS.
Deploy to AWS Lambda and API Gateway
Your prediction handler and Docker container image are ready. The posterior computation logic is packaged and waiting to go live.
In this step, you will push the container to Amazon ECR, create an AWS Lambda function from that image, and wire up an API Gateway HTTP API in front of it. By the end, you will have a live HTTPS endpoint that returns Bayesian congestion estimates with credible intervals.
In this step, get ready to:
- Push your Docker image to Amazon ECR.
- Create an AWS Lambda function from the container image.
- Set up an API Gateway HTTP API and test the endpoint.
Push your Docker image to Amazon ECR
Lambda cannot pull a container image from your local machine. It needs the image stored in a container registry that AWS can access. Amazon ECR is AWS's managed container registry.
- Create a new ECR repository by running this command:
aws ecr create-repository --repository-name auckland-congestion-lambda --region [[AWS_REGION="your AWS region e.g. ap-southeast-2"]]
What does this command do?
This creates a private repository in ECR called auckland-congestion-lambda. It is the destination where your Docker image will live so Lambda can pull it.
Look for the repositoryUri in the output. It follows the pattern ACCOUNT_ID.dkr.ecr.REGION.amazonaws.com/auckland-congestion-lambda. You will need this URI in the next commands.
- Copy the repositoryUri value from the JSON output and save it somewhere you can access easily.
Seeing an AccessDeniedException?
Confirm your AWS CLI is configured with credentials that have ECR permissions. Run aws sts get-caller-identity to check which account and role you are using.
Help me fix my IAM permissions for ECR
- Authenticate Docker to your ECR registry by running this command:
aws ecr get-login-password --region [[AWS_REGION="your AWS region e.g. ap-southeast-2"]] | docker login --username AWS --password-stdin [[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]].dkr.ecr.[[AWS_REGION="your AWS region e.g. ap-southeast-2"]].amazonaws.com
What does this command do?
This retrieves a temporary authentication token from ECR and pipes it to docker login. The token is valid for 12 hours. After this, Docker can push images to your ECR registry.
You should see Login Succeeded in your terminal.
- Tag your local Docker image with the ECR repository URI by running this command:
docker tag auckland-congestion-lambda:latest [[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]].dkr.ecr.[[AWS_REGION="your AWS region e.g. ap-southeast-2"]].amazonaws.com/auckland-congestion-lambda:latest
- Push the tagged image to ECR by running this command:
docker push [[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]].dkr.ecr.[[AWS_REGION="your AWS region e.g. ap-southeast-2"]].amazonaws.com/auckland-congestion-lambda:latest
You should see layer upload progress followed by a latest: digest: sha256:... confirmation. This push may take a few minutes depending on image size and internet speed.
Push hanging or failing?
Make sure Docker Desktop is running. If you see no basic auth credentials, re-run the aws ecr get-login-password command above. The auth token may have expired.
Also confirm the tag matches your repository name exactly. A mismatch between the tag and the repository causes a repository does not exist error.
Help me troubleshoot my Docker push to ECR
Create the Lambda function
Lambda needs an IAM execution role that grants it permission to write logs. You will create this role first, then create the Lambda function from your ECR image.
- Create a trust policy file called trust-policy.json in your project folder with the following content:
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Principal": {
"Service": "lambda.amazonaws.com"
},
"Action": "sts:AssumeRole"
}
]
}
What is a trust policy?
A trust policy defines which AWS service is allowed to assume this role. This policy says the Lambda service can assume the role, which is required for any Lambda function to run.
- Create the IAM role by running this command:
aws iam create-role \
--role-name lambda-congestion-role \
--assume-role-policy-document file://trust-policy.json
- Attach the basic Lambda execution policy so the function can write to CloudWatch Logs:
aws iam attach-role-policy \
--role-name lambda-congestion-role \
--policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
Why AWSLambdaBasicExecutionRole?
This managed policy grants permission to write logs to CloudWatch. Without it, your Lambda function would run but you would have no way to see its log output for debugging.
Wait a few seconds
IAM role propagation takes a few seconds. Wait about 10 seconds before running the next command. If you create the Lambda function too quickly, you may see a The role defined for the function cannot be assumed by Lambda error.
- Create the Lambda function from your ECR image by running this command:
aws lambda create-function \
--function-name auckland-congestion-model \
--package-type Image \
--code ImageUri=[[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]].dkr.ecr.[[AWS_REGION="your AWS region e.g. ap-southeast-2"]].amazonaws.com/auckland-congestion-lambda:latest \
--role arn:aws:iam::[[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]]:role/lambda-congestion-role \
--memory-size 1024 \
--timeout 30
What do these flags mean?
- --package-type Image tells Lambda this function runs from a container image, not a zip file.
- --memory-size 1024 allocates 1024 MB of memory. The posterior computation loads the NetCDF file and performs arithmetic on numpy arrays, so it needs more than the default 128 MB.
- --timeout 30 sets the maximum execution time to 30 seconds. Prediction is fast since you are doing arithmetic on saved posterior samples, not re-sampling the model.
- Verify the function was created by running this command:
aws lambda get-function --function-name auckland-congestion-model --query 'Configuration.[FunctionName,State,MemorySize,Timeout]'
You should see the function name auckland-congestion-model, state Active (or Pending briefly while the container image loads), memory 1024, and timeout 30.
Function stuck in Pending?
Container-based Lambda functions take longer to initialize than zip-based ones. Wait 30-60 seconds and run the get-function command again. If the state is Failed, check that your ECR image URI is correct and that the role ARN matches your account ID.
Help me debug my Lambda function creation
Create an API Gateway HTTP API and test the endpoint
Your Lambda function exists but has no public URL. API Gateway creates an HTTPS endpoint that forwards incoming requests to your function.
- Create a new HTTP API by running this command:
aws apigatewayv2 create-api \
--name auckland-congestion-api \
--protocol-type HTTP
- Look for the ApiId in the JSON output. Copy this value. You will use it in every command that follows.
- Create a Lambda integration that connects the API to your function by running this command:
aws apigatewayv2 create-integration \
--api-id [[API_ID="your ApiId from the create-api output"]] \
--integration-type AWS_PROXY \
--integration-uri arn:aws:lambda:[[AWS_REGION="your AWS region e.g. ap-southeast-2"]]:[[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]]:function:auckland-congestion-model \
--payload-format-version 2.0
What does this integration do?
AWS_PROXY means API Gateway forwards the full HTTP request to your Lambda function and returns the function's response directly to the caller. payload-format-version 2.0 uses the newer, simpler event format for HTTP APIs.
- Copy the IntegrationId from the output.
- Create the POST /predict route by running this command:
aws apigatewayv2 create-route \
--api-id [[API_ID="your ApiId from the create-api output"]] \
--route-key "POST /predict" \
--target integrations/[[INTEGRATION_ID="your IntegrationId from the create-integration output"]]
- Create a default stage with auto-deploy so changes go live immediately:
aws apigatewayv2 create-stage \
--api-id [[API_ID="your ApiId from the create-api output"]] \
--stage-name '$default' \
--auto-deploy
- Grant API Gateway permission to invoke your Lambda function by running this command:
aws lambda add-permission \
--function-name auckland-congestion-model \
--statement-id apigateway-invoke \
--action lambda:InvokeFunction \
--principal apigateway.amazonaws.com \
--source-arn "arn:aws:execute-api:[[AWS_REGION="your AWS region e.g. ap-southeast-2"]]:[[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]]:[[API_ID="your ApiId from the create-api output"]]/*"
Why add a permission?
Lambda functions are locked down by default. Even though you created the integration, API Gateway cannot call your function unless you explicitly grant it permission with add-permission.
The source-arn scopes the permission to your specific API so no other API can invoke the function.
Your API endpoint URL follows this pattern:
https://[[API_ID="your ApiId from the create-api output"]].execute-api.[[AWS_REGION="your AWS region e.g. ap-southeast-2"]].amazonaws.com
- Test the endpoint by sending a POST request with a sample corridor and time:
curl -X POST https://[[API_ID="your ApiId from the create-api output"]].execute-api.[[AWS_REGION="your AWS region e.g. ap-southeast-2"]].amazonaws.com/predict \
-H "Content-Type: application/json" \
-d '{"corridor_id": "great_south_road_sb", "hour_of_day": 8}'
What should I see?
You should receive a JSON response containing four keys: posterior_mean, posterior_median, credible_interval_5th, and credible_interval_95th. These represent the full posterior congestion estimate for that corridor at 8:00 AM.
The credible interval gives you a 90% range. There is a 90% probability the true congestion level falls within this range, which is the Bayesian equivalent of a confidence interval with a more intuitive interpretation.
Getting an Internal Server Error?
Check your Lambda function logs in CloudWatch. The most common causes are a mismatched corridor_id that does not exist in your model, or the model_trace.nc file not being found inside the container.
If you see a 403 Forbidden error, the add-permission command may not have run successfully. Re-run it and try the curl request again.
Help me debug my Lambda API Gateway endpoint
Secret mission
Add Weather Covariates and Anomaly Detection
Your congestion model captures corridor-level patterns and time-of-day effects. But weather is one of the biggest drivers of traffic delays in Auckland. In this extension, you will add weather covariates from the MetService Point Forecast API and build a real-time anomaly detector that flags unusually congested corridors.
Before you go
Clean Up Your Resources
Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project uses several AWS services that have generous free tiers, but storage costs can accumulate over time.
Cost warning
Your Lambda function and API Gateway are covered by the AWS free tier (1M requests/month each). However, your S3 bucket and ECR repository incur small storage fees once the free tier expires.
If you are not planning to use the API again, choose Delete to avoid any ongoing charges.
Resources you used:
- S3 bucket containing delay observation CSV files.
- ECR repository storing your Lambda container image.
- Lambda function running your posterior prediction handler.
- API Gateway HTTP API exposing the POST endpoint.
- Local project files including your Python virtual environment, scripts, model trace file, and Docker image.
- Auckland Transport Developer Portal account with your GTFS API subscription.
Keep everything running
No action needed. Choose this if you are still actively building or want to keep the API as a portfolio piece.
- Your Lambda function and API Gateway cost nothing at low usage under the free tier.
- S3 storage for your delay data is pennies per month.
- ECR charges pennies for storing the container image (500 MB free for the first 12 months).
- You can continue running the data collector to accumulate more observations and refit the model over time.
Pause - I'll come back to this later
Stop collecting new data but leave the API running. This prevents new storage costs while keeping your endpoint live.
- Stop the data ingestion script. No new data means no additional S3 storage charges.
- Leave the Lambda function and API Gateway in place. They cost nothing when idle.
- Your existing model and API continue to serve predictions from the data you already collected.
- Deactivate your Python virtual environment by running deactivate in your terminal.
Delete - I don't want to use this again
Remove all project resources and start fresh. This eliminates all ongoing costs.
- Delete the Lambda function by running:
aws lambda delete-function --function-name your-function-name
- Delete the ECR repository and all its images by running:
aws ecr delete-repository --repository-name your-repo-name --force
- Empty and delete the S3 bucket by running:
aws s3 rb s3://your-bucket-name --force
- Delete the API Gateway HTTP API from the AWS console. Open the API Gateway service, select your API, and click Delete.
- Delete your local project folder from the Desktop.
- Remove local Docker images by running docker image prune -a to free up disk space.
- Your Auckland Transport Developer Portal account is free and does not need to be deleted.
Mission accomplished
That's a Wrap!
Nice work! You just built a full Bayesian congestion model for Auckland, trained it on real-time bus delay data, and deployed it as a live API.
You've learned how to:
- Collect real-time GTFS delay data from Auckland Transport buses and store it in Amazon S3, treating buses as mobile congestion sensors across the network.
- Build a hierarchical Bayesian model in PyMC with corridor-level partial pooling, Fourier time-of-day effects, and day-of-week seasonality, then validate it with convergence diagnostics and posterior predictive checks.
- Deploy the fitted posterior as a containerized AWS Lambda function behind API Gateway, returning congestion estimates with full credible intervals for any Auckland corridor and time of day.
- Secret Mission: Added weather covariates from MetService and built a posterior-based anomaly detector that flags unusual congestion events in real time.
Ready to quiz yourself?