ewars_sandbox.py

import streamlit as st
import pandas as pd
import numpy as np
from itertools import product
import plotly.express as px

# Machine Learning Models
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.metrics import mean_absolute_error, accuracy_score, precision_score

# Advanced Forecasting Models
from statsmodels.tsa.statespace.sarimax import SARIMAX
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense

# --- Enhanced Custom CSS for Stunning Visuals ---
st.markdown(
    """
    <style>
    body {
        background: linear-gradient(135deg, #e0eafc, #cfdef3);
    }
    .main {
        background-color: transparent;
        font-family: 'Roboto', sans-serif;
    }
    .sidebar .sidebar-content {
        background: linear-gradient(180deg, #2e7bcf, #1a4e8a);
        color: #ffffff;
    }
    h1, h2, h3 {
        color: #1a4e8a;
        font-family: 'Roboto', sans-serif;
    }
    .stButton>button {
        background-color: #1a4e8a;
        color: white;
        border-radius: 5px;
        border: none;
        padding: 10px 20px;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# --- Transform Weather Features Function ---
def transform_weather_features(X_train, X_test, weather_cols, method, n_components=3):
    if method == "PCA":
        from sklearn.decomposition import PCA
        pca = PCA(n_components=n_components)
        train_weather = X_train[weather_cols]
        test_weather = X_test[weather_cols]
        X_train_trans = pca.fit_transform(train_weather)
        X_test_trans = pca.transform(test_weather)
        pca_cols = [f'PCA_{i+1}' for i in range(n_components)]
        X_train_pca = pd.DataFrame(X_train_trans, columns=pca_cols, index=X_train.index)
        X_test_pca = pd.DataFrame(X_test_trans, columns=pca_cols, index=X_test.index)
        X_train = X_train.drop(columns=weather_cols)
        X_test = X_test.drop(columns=weather_cols)
        X_train = pd.concat([X_train, X_train_pca], axis=1)
        X_test = pd.concat([X_test, X_test_pca], axis=1)
        return X_train, X_test
    elif method == "Autoencoder":
        # Build a simple autoencoder using TensorFlow Keras for the weather features.
        input_dim = len(weather_cols)
        encoding_dim = n_components
        input_layer = tf.keras.layers.Input(shape=(input_dim,))
        encoded = tf.keras.layers.Dense(16, activation='relu')(input_layer)
        encoded = tf.keras.layers.Dense(encoding_dim, activation='relu')(encoded)
        decoded = tf.keras.layers.Dense(16, activation='relu')(encoded)
        decoded = tf.keras.layers.Dense(input_dim, activation='linear')(decoded)
        autoencoder = tf.keras.models.Model(inputs=input_layer, outputs=decoded)
        encoder = tf.keras.models.Model(inputs=input_layer, outputs=encoded)
        autoencoder.compile(optimizer='adam', loss='mse')
        autoencoder.fit(X_train[weather_cols], X_train[weather_cols], epochs=50, batch_size=16, verbose=0)
        X_train_encoded = encoder.predict(X_train[weather_cols])
        X_test_encoded = encoder.predict(X_test[weather_cols])
        encoded_cols = [f'AE_{i+1}' for i in range(encoding_dim)]
        X_train_ae = pd.DataFrame(X_train_encoded, columns=encoded_cols, index=X_train.index)
        X_test_ae = pd.DataFrame(X_test_encoded, columns=encoded_cols, index=X_test.index)
        X_train = X_train.drop(columns=weather_cols)
        X_test = X_test.drop(columns=weather_cols)
        X_train = pd.concat([X_train, X_train_ae], axis=1)
        X_test = pd.concat([X_test, X_test_ae], axis=1)
        return X_train, X_test
    else:
        return X_train, X_test

# --- Data Generation Functions ---

@st.cache_data
def generate_synthetic_data():
    # Parameters
    num_locations = 10
    dates = pd.date_range('2015-01-01', '2020-12-31')
    num_days = len(dates)
    diseases = ['malaria', 'dengue', 'cholera', 'river_blindness']
    
    # --- Create dummy location data with extra attributes ---
    # First half: Bangladesh, Second half: Sierra Leone
    countries = ['Bangladesh'] * (num_locations // 2) + ['Sierra Leone'] * (num_locations - num_locations // 2)
    populations = [np.random.randint(1000000, 5000000) if c == 'Bangladesh' else np.random.randint(200000, 1000000) for c in countries]
    economics = [np.random.randint(1500, 3000) if c == 'Bangladesh' else np.random.randint(500, 1500) for c in countries]
    lulc = np.random.choice(["urban", "rural", "suburban"], size=num_locations)
    vector_density = np.random.uniform(0, 1, size=num_locations)
    # Generate lat/lon based on country:
    lats = [np.random.uniform(20.5, 26) if c == "Bangladesh" else np.random.uniform(7.5, 10) for c in countries]
    lons = [np.random.uniform(88, 92) if c == "Bangladesh" else np.random.uniform(-13, -10) for c in countries]
    
    locations = pd.DataFrame({
        'location_id': range(num_locations),
        'lat': lats,
        'lon': lons,
        'country': countries,
        'population': populations,
        'economics': economics,
        'lulc': lulc,
        'vector_density': vector_density
    })

    # Weather parameters
    weather_params = {
        "10u": {"mean": 0, "std": 5},
        "10v": {"mean": 0, "std": 5},
        "2d": {"mean": 20, "std": 5},
        "2t": {"mean": 25, "std": 5},
        "msl": {"mean": 1013, "std": 10},
        "sp": {"mean": 1013, "std": 10},
        "tp": {"mean": 0.01, "std": 0.02, "clip": [0, None]},
        "tcc": {"mean": 0.5, "std": 0.2, "clip": [0, 1]},
        "lc": {"mean": 0, "std": 0},
        "ld": {"mean": 0, "std": 0},
        "lmlt": {"mean": 25, "std": 5},
        "cp": {"mean": 0.005, "std": 0.01, "clip": [0, None]},
        "crr": {"mean": 0.001, "std": 0.002, "clip": [0, None]},
        "tcrw": {"mean": 0.1, "std": 0.1, "clip": [0, None]},
        "swvl1": {"mean": 0.3, "std": 0.1, "clip": [0, 1]},
        "cvh": {"mean": 0.5, "std": 0.2, "clip": [0, 1]},
        "cvl": {"mean": 0.5, "std": 0.2, "clip": [0, 1]},
        "skt": {"mean": 25, "std": 5},
        "e": {"mean": -0.001, "std": 0.001},
        "ro": {"mean": 0.001, "std": 0.001, "clip": [0, None]},
        "sro": {"mean": 0.001, "std": 0.001, "clip": [0, None]},
        "tcwv": {"mean": 30, "std": 10}
    }

    weather_df = pd.DataFrame({
        'date': np.repeat(dates, num_locations),
        'location_id': np.tile(range(num_locations), num_days)
    })
    for var, params in weather_params.items():
        mean = params['mean']
        std = params['std']
        data = np.random.normal(mean, std, num_days * num_locations)
        if 'clip' in params:
            clip_min, clip_max = params['clip']
            data = np.clip(data, clip_min, clip_max)
        weather_df[var] = data

    weather_df = weather_df.sort_values(['location_id', 'date'])
    weather_df['tp_lag7'] = weather_df.groupby('location_id')['tp'].transform(
        lambda x: x.rolling(7, min_periods=1).sum()
    )

    # Generate disease data and merge weather and location info BEFORE computing cases
    disease_df = pd.DataFrame(list(product(dates, range(num_locations), diseases)),
                              columns=['date', 'location_id', 'disease'])
    disease_df = pd.merge(disease_df, weather_df, on=['date', 'location_id'])
    disease_df = pd.merge(disease_df, locations, on='location_id', how='left')

    def compute_cases(row):
        if row['disease'] == 'malaria':
            lam = max(0, 0.5 * row['2t'] + 10 * row['tp'] + 5 + 10 * row['vector_density'])
            return np.random.poisson(lam)
        elif row['disease'] == 'dengue':
            lam = max(0, 0.3 * row['2t'] + 15 * row['tp'] + 3 + 8 * row['vector_density'])
            return np.random.poisson(lam)
        elif row['disease'] == 'cholera':
            lam = 0.01 + 0.1 * row['tp_lag7']
            return np.random.poisson(lam)
        elif row['disease'] == 'river_blindness':
            lam = 0.01 + 0.05 * row['swvl1']
            return np.random.poisson(lam)
        return 0

    disease_df['cases'] = disease_df.apply(compute_cases, axis=1)
    return locations, weather_df, disease_df

# --- Data Preparation Functions ---

def prepare_forecasting_data(disease_df, disease, location_id, lags=14):
    df = disease_df[(disease_df['disease'] == disease) & (disease_df['location_id'] == location_id)].copy()
    df = df.sort_values('date')
    for lag in range(1, lags + 1):
        df[f'cases_lag{lag}'] = df['cases'].shift(lag)
    exclude_cols = ['date', 'location_id', 'disease', 'cases', 'lat', 'lon', 'country',
                    'population', 'economics', 'lulc', 'vector_density'] + [f'cases_lag{i}' for i in range(1, lags + 1)]
    weather_vars = [col for col in df.columns if col not in exclude_cols]
    df['target'] = df['cases'].shift(-1)
    df = df.dropna()
    features = [f'cases_lag{i}' for i in range(1, lags + 1)] + weather_vars
    X = df[features]
    y = df['target']
    train_size = int(len(df) * 0.8)
    X_train, X_test = X.iloc[:train_size], X.iloc[train_size:]
    y_train, y_test = y.iloc[:train_size], y.iloc[train_size:]
    return X_train, X_test, y_train, y_test, df, features

def prepare_hotspot_data(disease_df, disease, date, lags=14):
    df = disease_df[disease_df['disease'] == disease].copy()
    past_df = df[df['date'] <= date].sort_values(['location_id', 'date'])
    for lag in range(1, lags + 1):
        past_df[f'cases_lag{lag}'] = past_df.groupby('location_id')['cases'].shift(lag)
    latest = past_df[past_df['date'] == date].dropna()
    exclude_cols = ['date', 'location_id', 'disease', 'cases', 'lat', 'lon', 'country',
                    'population', 'economics', 'lulc', 'vector_density'] + [f'cases_lag{i}' for i in range(1, lags + 1)]
    weather_vars = [col for col in latest.columns if col not in exclude_cols]
    X = latest[[f'cases_lag{i}' for i in range(1, lags + 1)] + weather_vars]
    next_day = df[df['date'] == date + pd.Timedelta(days=1)][['location_id', 'cases']].rename(columns={'cases': 'target'})
    latest = latest.merge(next_day, on='location_id', how='left')
    y = latest['target']
    return X, y, latest

def prepare_event_data(disease_df, disease, location_id, window=14, threshold=5):
    df = disease_df[(disease_df['disease'] == disease) & (disease_df['location_id'] == location_id)].copy()
    df = df.sort_values('date')
    df['future_cases'] = df['cases'].shift(-window).rolling(window, min_periods=1).sum()
    df['target'] = (df['future_cases'] > threshold).astype(int)
    for lag in range(1, window + 1):
        df[f'cases_lag{lag}'] = df['cases'].shift(lag)
    exclude_cols = ['date', 'location_id', 'disease', 'cases', 'future_cases', 'target', 'lat', 'lon', 'country',
                    'population', 'economics', 'lulc', 'vector_density'] + [f'cases_lag{i}' for i in range(1, window + 1)]
    weather_vars = [col for col in df.columns if col not in exclude_cols]
    df = df.dropna()
    features = [f'cases_lag{i}' for i in range(1, window + 1)] + weather_vars
    X = df[features]
    y = df['target']
    train_size = int(len(df) * 0.8)
    X_train, X_test = X.iloc[:train_size], X.iloc[train_size:]
    y_train, y_test = y.iloc[:train_size], y.iloc[train_size:]
    return X_train, X_test, y_train, y_test, df, features

# --- Sidebar Feature Engineering Option ---
feat_eng_method = st.sidebar.selectbox("Select Feature Engineering Method", ["None", "PCA", "Autoencoder"])
if feat_eng_method != "None":
    n_components = st.sidebar.slider("Number of Components", 1, 10, 3)

# --- Main App ---

st.title("Multidisease EWARS Modeling Sandbox")

# Sidebar Country Selection
selected_country = st.sidebar.selectbox("Select Country", ["Bangladesh", "Sierra Leone"])

# Load data
locations, weather_df, disease_df = generate_synthetic_data()

# Filter data based on selected country
disease_df = disease_df[disease_df['country'] == selected_country]
country_locations = locations[locations['country'] == selected_country]

# Sidebar: Display location info for the selected country
with st.sidebar.expander("Location Info"):
    st.dataframe(country_locations.reset_index(drop=True))

# Sidebar navigation for tasks
task = st.sidebar.selectbox("Select Task", ["Forecasting", "Hotspot Prediction", "Event-based Outbreak Risk Prediction"])

if task == "Forecasting":
    st.header("Disease Case Forecasting")
    disease = st.selectbox("Select Disease", ["malaria", "dengue"])
    loc_options = country_locations['location_id'].tolist()
    location_id = st.selectbox("Select Location ID", loc_options)
    model_choice = st.selectbox("Select Model", ["Linear Regression", "Random Forest", "SARIMAX", "LSTM"])
    lags = st.slider("Number of Lag Days", 1, 30, 14)
    
    if st.button("Run Forecasting"):
        X_train, X_test, y_train, y_test, df_model, features = prepare_forecasting_data(disease_df, disease, location_id, lags)
        # Identify weather columns (exclude lag features)
        lag_cols = [f'cases_lag{i}' for i in range(1, lags + 1)]
        weather_cols = [col for col in features if col not in lag_cols]
        if feat_eng_method != "None":
            X_train, X_test = transform_weather_features(X_train, X_test, weather_cols, feat_eng_method, n_components)
        
        if model_choice == "Linear Regression":
            model = LinearRegression()
            model.fit(X_train, y_train)
            y_pred = model.predict(X_test)
        elif model_choice == "Random Forest":
            model = RandomForestRegressor(n_estimators=100, random_state=42)
            model.fit(X_train, y_train)
            y_pred = model.predict(X_test)
        elif model_choice == "SARIMAX":
            train_size = int(len(df_model) * 0.8)
            train_series = df_model['cases'].iloc[:train_size]
            sarimax_model = SARIMAX(train_series, order=(1, 0, 0), enforce_stationarity=False, enforce_invertibility=False)
            sarimax_result = sarimax_model.fit(disp=False)
            y_pred = sarimax_result.forecast(steps=len(y_test)).values
        elif model_choice == "LSTM":
            lag_cols = [f'cases_lag{i}' for i in range(1, lags + 1)]
            X_train_lstm = X_train[lag_cols].values.reshape((X_train.shape[0], lags, 1))
            X_test_lstm = X_test[lag_cols].values.reshape((X_test.shape[0], lags, 1))
            lstm_model = Sequential([
                LSTM(50, activation='tanh', input_shape=(lags, 1)),
                Dense(1)
            ])
            lstm_model.compile(optimizer='adam', loss='mse')
            lstm_model.fit(X_train_lstm, y_train.values, epochs=20, batch_size=16, verbose=0)
            y_pred = lstm_model.predict(X_test_lstm).flatten()
        
        mae = mean_absolute_error(y_test, y_pred)
        st.write(f"Mean Absolute Error: {mae:.2f}")
        test_df = df_model.iloc[-len(y_test):].copy()
        test_df['predicted'] = y_pred
        fig = px.line(test_df, x='date', y=['cases', 'predicted'],
                      title=f"Forecasting for {disease} at Location {location_id} using {model_choice}")
        st.plotly_chart(fig, use_container_width=True)

elif task == "Hotspot Prediction":
    st.header("Hotspot Prediction")
    disease = st.selectbox("Select Disease", ["malaria", "dengue", "cholera", "river_blindness"])
    date = st.date_input("Select Prediction Date", value=pd.to_datetime("2020-12-01"))
    model_choice = st.selectbox("Select Model", ["Linear Regression", "Random Forest"])
    threshold = st.slider("Hotspot Threshold (Cases)", 1, 20, 5)
    
    if st.button("Run Hotspot Prediction"):
        date = pd.to_datetime(date)
        X, y, latest = prepare_hotspot_data(disease_df, disease, date)
        if model_choice == "Linear Regression":
            model = LinearRegression()
        else:
            model = RandomForestRegressor(n_estimators=100, random_state=42)
        model.fit(X, y)
        predictions = model.predict(X)
        latest['predicted_cases'] = predictions
        latest['is_hotspot'] = (latest['predicted_cases'] > threshold).astype(int)
        
        hotspots = latest[latest['is_hotspot'] == 1][['location_id', 'lat', 'lon', 'predicted_cases']]
        if hotspots.empty:
            st.warning("No hotspots predicted with current settings. Try lowering the threshold.")
        else:
            st.write("Predicted Hotspots:")
            st.dataframe(hotspots)
            fig = px.scatter_mapbox(
                hotspots,
                lat="lat", lon="lon", size="predicted_cases",
                color="predicted_cases",
                mapbox_style="open-street-map",
                title=f"Hotspots for {disease} on {date.date()}"
            )
            st.plotly_chart(fig, use_container_width=True)

else:  # Event-based Outbreak Risk Prediction
    st.header("Event-based Outbreak Risk Prediction")
    disease = st.selectbox("Select Disease", ["cholera", "river_blindness"])
    loc_options = country_locations['location_id'].tolist()
    location_id = st.selectbox("Select Location ID", loc_options)
    model_choice = st.selectbox("Select Model", ["Logistic Regression", "Random Forest"])
    window = st.slider("Lookback Window (Days)", 7, 30, 14)
    threshold = st.slider("Outbreak Threshold (Cases in Next 7 Days)", 1, 10, 5)
    
    if st.button("Run Outbreak Prediction"):
        X_train, X_test, y_train, y_test, df_model, features = prepare_event_data(disease_df, disease, location_id, window, threshold)
        # Optionally, feature engineering could be applied here as well.
        if len(np.unique(y_train)) < 2:
            st.warning("Insufficient outbreak events in training data. Please adjust the threshold or lookback window.")
        else:
            if model_choice == "Logistic Regression":
                model = LogisticRegression(max_iter=1000)
            else:
                model = RandomForestClassifier(n_estimators=100, random_state=42)
            model.fit(X_train, y_train)
            y_pred = model.predict(X_test)
            y_pred_proba = model.predict_proba(X_test)[:, 1]
            accuracy = accuracy_score(y_test, y_pred)
            precision = precision_score(y_test, y_pred)
            st.write(f"Accuracy: {accuracy:.2f}, Precision: {precision:.2f}")
            test_df = df_model.iloc[-len(y_test):].copy()
            test_df['predicted_proba'] = y_pred_proba
            fig = px.line(test_df, x='date', y='predicted_proba', 
                          title=f"Outbreak Risk for {disease} at Location {location_id}")
            st.plotly_chart(fig, use_container_width=True)

st.sidebar.write("This sandbox now confines location generation within Bangladesh and Sierra Leone and includes feature engineering options (PCA & Autoencoder) for weather variables.")