t6harsh/tb-vulnerability-hotspot-predictor

TB Vulnerability Hotspot Predictor - India

Model Description

An industry-grade hybrid ML pipeline for predicting tuberculosis (TB) vulnerability hotspots across India at district level. Combines gradient-boosted decision trees with transformer-based contextual embeddings for high-resolution vulnerability mapping.

Architecture

Input (153 features) → Quantile Transform
    ├── FT-Transformer (d=64, L=2, h=4) → [CLS] embedding (64-dim)
    └── Quantile-transformed features
         ↓
    Concatenated features (153 + 64 = 217)
         ↓
    ├── XGBoost (500 trees, depth=5)
    └── LightGBM (500 trees, depth=5)
         ↓
    Ensemble (0.5 × XGB + 0.5 × LGB)
         ↓
    Probabilistic Calibration (Platt + Isotonic + Temperature)
         ↓
    Calibrated Risk Probabilities

Performance

Classification (Hotspot Detection)

Regression (TB Notification Rate)

Ablation Study Results

Spatial Encoding Impact (AUC-ROC)

Feature Group Importance (Single Group AUC)

Calibration Methods

Visualizations

The model repository includes 10 publication-quality visualizations:

1. 1.Vulnerability Map - District-level TB vulnerability across India
2. 2.Hotspot/Coldspot Map - Getis-Ord Gi* classification
3. 3.Feature Importance - Top 20 predictive features
4. 4.ROC Curve - Classification performance
5. 5.Calibration Curve - Reliability diagram
6. 6.Ablation Heatmap - Spatial encoding & feature group comparisons
7. 7.Confusion Matrix - Classification outcomes
8. 8.Regional TB Rates - Distribution by Indian region
9. 9.Correlation Matrix - Key TB risk factor correlations
10. 10.State Vulnerability - Distribution across top 15 states

Key Findings

1. 1.Spatial features are critical: +1.3% AUC from spatial lag, Fourier encoding, graph proximity, and geographic clustering
2. 2.Environmental features (PM2.5, temperature, rainfall, altitude) are the strongest individual predictors
3. 3.Platt scaling achieves best calibration (ECE: 0.029) while preserving AUC
4. 4.Transformer embeddings capture cross-feature interactions that trees miss
5. 5.Multi-scale clustering (K=5,10,20,50 + DBSCAN) outperforms any single granularity

Technical Stack

• —Models: XGBoost, LightGBM, FT-Transformer (PyTorch)
• —Spatial: Getis-Ord Gi*, Spatial Lag, Fourier PE, K-hop Graph Aggregation
• —Calibration: Platt Scaling, Isotonic Regression, Temperature Scaling
• —Data: 729 Indian districts × 153 features

Usage

# Clone and run the full pipeline
git clone https://huggingface.co/t6harsh/tb-vulnerability-hotspot-predictor
cd tb-vulnerability-hotspot-predictor/src
python main_pipeline.py

Files

├── src/
│   ├── main_pipeline.py          # Main orchestrator
│   ├── data_generation.py        # Synthetic data based on real distributions
│   ├── spatial_features.py       # All spatial feature engineering
│   ├── ft_transformer.py         # FT-Transformer implementation
│   ├── hybrid_model.py           # Hybrid GBDT + Transformer model
│   ├── calibration.py            # Platt/Isotonic/Temperature calibration
│   └── ablation_study.py         # Comprehensive ablation framework
├── vulnerability_map.png
├── hotspot_coldspot_map.png
├── feature_importance.png
├── roc_curve.png
├── calibration_curve.png
├── ablation_heatmap.png
├── confusion_matrix.png
├── regional_tb_rates.png
├── correlation_matrix.png
├── state_vulnerability.png
├── metrics.json
├── ablation_results.json
└── REPORT.md

License

Apache 2.0

Acknowledgments

Based on published distributions from NFHS-5, Census 2011, WHO GTB, NIKSHAY, and methodology from:

• —Gorishniy et al., "Revisiting Deep Learning Models for Tabular Data" (FT-Transformer)
• —Guo et al., "On Calibration of Modern Neural Networks" (Calibration)
• —Spatial analysis informed by PySAL methodology