FlowTaxi: Conditional Flow Matching for NYC Taxi Demand

August 2026 Python PyTorch FastAPI Streamlit DVC Docker

Overview

FlowTaxi trains a conditional continuous-time flow matching model (Rectified Flow) that transports 2D Gaussian noise $\mathbf{x}_0 \sim \mathcal{N}(0, \mathbf{I}_2)$ into Manhattan yellow taxi drop-off locations $\mathbf{x}_1$, conditioned on weekday clock hour $h \in {0, \dots, 23}$.

Users interact with a Streamlit web application to select a target weekday hour, trigger neural ODE trajectory generation via a FastAPI backend, and move a client-side time slider $t \in [0, 1]$ to watch particles flow into realistic NYC demand hotspots.

Key Engineering Highlights

Coordinate-Level Data Engineering: Built on historical 2010 NYC TLC yellow taxi records containing genuine GPS coordinates rather than coarse zone polygons.
Continuous-Time Flow Matching: Implements Rectified Flow Matching with straight vector field probability paths, enabling high-fidelity sample generation in significantly fewer ODE steps than standard diffusion models.
Strict Layer Separation: Pure PyTorch ML core (zero API/UI dependencies), FastAPI serving wrapper, and a Streamlit frontend communicating over HTTP.
FastAPI Lifespan Model Registry: Model weights, normalizer parameters, and metadata are loaded once into RAM at server startup, eliminating per-request disk I/O.
Client-Side Trajectory Caching: Backend generates the full trajectory sequence once. The Streamlit slider indexes into cached local frames client-side with zero repeated backend HTTP calls.
Production MLOps Pipeline: Utilizes DVC for versioning, W&B for experiment tracking, and GCP Cloud Run for fully dockerized deployment.

System Architecture

flowchart TD
    subgraph Data & Pipeline Layer ["Data Engineering & Versioning (DVC)"]
        TLC["NYC TLC Coordinate-Era Data (2010)"] --> Preprocess["Polars + GeoPandas Preprocessor"]
        Preprocess --> Normalizer["Train Normalizer (EPSG:32618)"]
        Preprocess --> Split["Temporal Split (70 / 15 / 15)"]
        Split --> DVC["DVC Versioned Artifact Store"]
    end

    subgraph Training Layer ["Pure PyTorch ML Core (src/flowtaxi/ml)"]
        DVC --> Train["Continuous Flow Matching Training"]
        Train --> WB["W&B Experiment Tracking"]
        Train --> Safetensors["Safetensors Model Weights & Metadata"]
    end

    subgraph Serving Layer ["FastAPI Backend (services/api)"]
        Safetensors --> Registry["Lifespan Model Registry (Load Once)"]
        Registry --> Lifespan["Verify SHA256 & Hold Model in Memory"]
        Lifespan --> SyncRoute["Sync Route: /api/v1/generate"]
        SyncRoute --> Solver["Neural ODE Solver (Euler / Midpoint / RK4)"]
    end

    subgraph User Experience Layer ["Streamlit Frontend (apps/streamlit_app)"]
        SyncRoute -- "HTTP JSON (Full Trajectory)" --> Streamlit["Streamlit UI (httpx)"]
        Streamlit --> ClientCache["Client-Side Trajectory Cache"]
        ClientCache --> Slider["Local t-Slider (0 <= t <= 1)"]
    end

Quantitative Evaluation

Evaluated against held-out test drop-offs ($N=50,000$) across 5 key weekday hours using $10,000$ generated samples per hour with the Midpoint solver (48 steps).

Model / Baseline	Sliced Wasserstein Distance (SWD) ↓	H3 Jensen-Shannon Divergence (Res 9) ↓	Top-20 Hotspot Recall ↑
Global Gaussian Baseline	0.1431	0.1807	24.0%
Spatial KDE Baseline	0.0616	0.1326	42.0%
FlowTaxi CFM (Ours)	0.0592	0.0589	62.0%

Serving Latency Benchmarks

Measured on GCP Cloud Run (1 vCPU, 2 GiB RAM, single Uvicorn worker, PyTorch CPU execution) for $800$ trajectory samples per request:

Quality Preset	ODE Solver	ODE Steps	Min Latency	p50 (Median) Latency	p95 Latency
Fast	Midpoint	32	7.50s	7.86s	8.40s
High Quality	RK4	64	24.51s	25.74s	27.95s

Local Development & Reproduction

To reproduce the environment and run the services locally using the uv package manager:

# Setup Environment
make setup

# Reproduce Data Preprocessing, Training & Evaluation
make dvc-repro

# Start FastAPI backend service (http://localhost:8080)
make api

# In a separate terminal, start Streamlit UI (http://localhost:8501)
make ui

Alternatively, run the containerized stack via Docker Compose:

make docker-up

Testing & Documentation

The project enforces a strict quality suite covering formatting (ruff), strict static typing (mypy), unit tests (pytest), API contracts, and pre-commit hooks.

Comprehensive documentation is available in the repository’s docs/ directory, including:

Model Card: Deep dive into the Flow Matching architecture, Spatial Fourier Features, loss function, and safety validations.
Data Card: Details on the NYC TLC January 2010 dataset and coordinate projection.
Architecture Specs: Breakdown of monorepo layer boundaries and frontend/backend state architecture.

```