ArogyaJal Early Warning System for Waterborne Disease Outbreaks
Model Description
The ArogyaJal Early Warning System v2.0 is a production-grade machine learning model designed to predict waterborne disease outbreaks. Developed by the ML Engineering Team, this system provides early alerts for potential outbreaks, enabling timely interventions and safeguarding public health. It is built upon an ensemble of LightGBM and XGBoost classifiers, trained on a large-scale dataset from 50 villages over a 730-day period.
Problem Formulation
The model addresses a binary classification task: predicting whether the 7-day forward sum of reported cases of waterborne diseases will be greater than or equal to 3 (indicating an outbreak). This formulation allows for clear, actionable alerts.
Dataset
The model was trained and validated on a comprehensive dataset with the following key characteristics:
Total Observations: 36,200 samples
Number of Villages: 50
Date Range: January 1, 2023, to December 24, 2024
Class Imbalance: 72.0% negative (no outbreak), 28.0% positive (outbreak), reflecting a realistic 2.57:1 imbalance ratio.
Validation: Panel-Safe Stratified time-based validation was used to ensure robustness and prevent temporal leakage.
Feature Engineering
The model utilizes 44 strictly causal features, engineered through a highly optimized, vectorized pipeline. These features capture critical information from IoT sensor data and reported cases, including:
Missingness indicators
Imputed IoT features
Temporal lags (t-1, t-3, t-7, t-14) for IoT parameters and reported cases
Rolling statistics (7-day and 14-day mean/std) for turbidity and reported cases
3-day rate of change (derivatives) for pH, turbidity, TDS, and conductivity
Cyclical seasonal features (sin/cos of day of year)
Water quality composite score
Interaction features (turb_x_cases_lag1, ph_deviation, tds_x_turb)
Cumulative burden (7-day and 14-day rolling sum of reported cases)
Model Architecture
The final production model is a weighted ensemble of two Optuna-tuned gradient boosting models:
LightGBM Binary Classifier
XGBoost Binary Classifier
The ensemble combines their predictions with weights: 0.90 * P_LightGBM + 0.10 * P_XGBoost.
Performance Metrics
The model demonstrates excellent performance on a holdout test set (7,250 samples):
Metric
Value
Target
Status
PR-AUC
0.9437
≥ 0.60
✓ EXCEEDED
ROC-AUC
0.9528
≥ 0.85
✓ EXCEEDED
Recall (Sensitivity)
0.8712 (87.1%)
≥ 80.0%
✓ EXCEEDED
Precision
0.9060 (90.6%)
≥ 50.0%
✓ EXCEEDED
F1-Score
0.8883
≥ 0.60
✓ EXCEEDED
False Positive Rate
4.5%
≤ 30.0%
✓ EXCEEDED
Optimal Threshold
0.3774
-
-
Confusion Matrix (Holdout Test Set):
True Negatives (TN): 4,633
False Positives (FP): 217
False Negatives (FN): 309
True Positives (TP): 2,091
Explainability (SHAP)
SHAP analysis identified the most influential features for outbreak prediction:
turbidity_roll_std7: Short-term volatility in water turbidity.
turbidity_roll_mean7: Baseline water cloudiness.
turbidity_roc3: 3-day rate of change in turbidity.
reported_cases_roll_std7: Volatility of active outbreaks.
This indicates that rapid increases and high volatility in water turbidity, coupled with seasonal patterns, are strong predictors of outbreaks.
Usage
To use the ArogyaJal Early Warning System for inference, you will need the following artifacts:
model_lgbm_v2.pkl
model_xgb_v2.pkl
ensemble_config.json
arogyajal_inference.py (the inference script)
Example Inference Code
python
1import pandas as pd
2from arogyajal_inference import OutbreakWarningSystem
34# Assuming model files and config are in the same directory5warning_system = OutbreakWarningSystem(6 lgb_path="model_lgbm_v2.pkl",7 xgb_path="model_xgb_v2.pkl",8 config_path="ensemble_config.json"9)1011# Prepare new data for prediction (example DataFrame structure)12df_new = pd.DataFrame({13'timestamp': pd.to_datetime(['2024-12-25','2024-12-26']),14'village_id':['VIL_001','VIL_001'],15'ph':[7.2,7.3],16'turbidity':[0.8,1.5],17'tds':[32000,32500],18'conductivity':[460,470],19'reported_cases':[0,1]20})2122results = warning_system.predict(df_new)23alerts = warning_system.get_alerts(results)2425print("Outbreak Alerts:")26print(alerts)
Installation
To run the inference script, you will need the following Python packages:
