Spaces:
Sleeping
Sleeping
File size: 3,619 Bytes
230416a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import xgboost as xgb
import joblib
import os
import warnings
warnings.filterwarnings('ignore')
class ImpactForecaster:
def __init__(self):
self.model = xgb.XGBRegressor(
n_estimators=150,
learning_rate=0.1,
max_depth=6,
random_state=42
)
self.encoders = {}
self.features = ['latitude', 'longitude', 'event_type_encoded',
'event_cause_encoded', 'priority_encoded',
'start_hour', 'day_of_week', 'is_weekend']
self.is_trained = False
def train(self, df):
"""Trains the XGBoost model to predict impact duration."""
print("Training ImpactForecaster...")
cat_cols = ['event_type', 'event_cause', 'priority']
df_encoded = df.copy()
for col in cat_cols:
le = LabelEncoder()
df_encoded[col] = df_encoded[col].astype(str)
le.fit(df_encoded[col].tolist() + ['unknown'])
self.encoders[col] = le
df_encoded[col + '_encoded'] = le.transform(df_encoded[col])
X = df_encoded[self.features]
y = np.log1p(df_encoded['impact_duration_mins'])
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
self.model.fit(
X_train, y_train,
eval_set=[(X_test, y_test)],
verbose=False
)
self.is_trained = True
# Save model
os.makedirs('models', exist_ok=True)
joblib.dump(self.model, 'models/forecaster_model.pkl')
joblib.dump(self.encoders, 'models/forecaster_encoders.pkl')
print("Model training complete and saved.")
def predict(self, event_data):
"""Predicts impact duration and calculates congestion radius."""
if not self.is_trained:
self._load_model()
df = pd.DataFrame([event_data])
for col, le in self.encoders.items():
if col in df.columns:
val = str(df[col].iloc[0])
if val not in le.classes_:
val = 'unknown'
df[col + '_encoded'] = le.transform([val])
else:
df[col + '_encoded'] = le.transform(['unknown'])
for feat in self.features:
if feat not in df.columns:
df[feat] = 0
X = df[self.features]
log_pred = self.model.predict(X)[0]
duration_mins = np.expm1(log_pred)
# Congestion radius scales with duration and priority
base_radius = 500
priority_val = str(event_data.get('priority', '')).lower()
if priority_val == 'high':
base_radius = 800
elif priority_val == 'low':
base_radius = 300
congestion_radius = min(base_radius + (duration_mins * 6), 4000)
return {
'predicted_duration_mins': float(duration_mins),
'congestion_radius_meters': float(congestion_radius)
}
def _load_model(self):
try:
self.model = joblib.load('models/forecaster_model.pkl')
self.encoders = joblib.load('models/forecaster_encoders.pkl')
self.is_trained = True
except FileNotFoundError:
raise RuntimeError("Model not trained yet. Call train() first.")
|