ConfTimeTransformer.py

import os
import random
import yaml
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

###############################################################################
# 0) Environment Variables & Logging
###############################################################################
#  - Suppress most TF logs (errors only).
#  - Attempt to force deterministic ops (if supported by your TF version).
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"  # 0=ALL, 1=INFO, 2=WARNING, 3=ERROR
os.environ["TF_DETERMINISTIC_OPS"] = "1"

# If you want to disable GPU for maximum reproducibility (CPU only), uncomment:
# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

###############################################################################
# 1) TensorFlow & Other Imports
###############################################################################
import tensorflow as tf
from tensorflow.keras import layers, models
from sklearn.preprocessing import MinMaxScaler

# Try enabling TF deterministic ops in code (TF >= 2.8/2.9).
try:
    tf.config.experimental.enable_op_determinism()
except Exception as e:
    print("Note: Could not enable_op_determinism(). Some TF versions may not support it.", e)

###############################################################################
# 2) Set All Random Seeds
###############################################################################
# For reproducibility across runs (not guaranteed 100%, but as close as you can get).
random.seed(42)
np.random.seed(42)
tf.random.set_seed(42)

###############################################################################
# 3) Load Configuration
###############################################################################

def load_config(config_path: str = "config.yaml") -> dict:
    """Load YAML config file and return it as a dictionary."""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config

###############################################################################
# 4) Main Script
###############################################################################

def main():
    # Use dark background for all plots
    plt.style.use('dark_background')

    # -------------------------------------------------------------------------
    # 4.1) Load the config
    # -------------------------------------------------------------------------
    config = load_config("config.yaml")

    # Paths
    input_csv = config["paths"]["input_csv"]
    output_plots_dir = config["paths"].get("output_plots_dir", None)

    # Columns
    date_col = config["columns"]["date_col"]
    target_col = config["columns"]["target_col"]
    feature_cols = config["columns"]["feature_cols"]

    # Training params
    window_size = config["training"]["window_size"]
    train_split = config["training"]["train_split"]
    epochs = config["training"]["epochs"]
    batch_size = config["training"]["batch_size"]
    learning_rate = config["training"]["learning_rate"]

    # Transformer params
    d_model = config["transformer"]["d_model"]
    num_heads = config["transformer"]["num_heads"]
    ff_dim = config["transformer"]["ff_dim"]
    num_transformer_blocks = config["transformer"]["num_transformer_blocks"]
    dropout_rate = config["transformer"]["dropout_rate"]

    # Forecast params
    n_future = config["forecast"]["n_future"]

    # -------------------------------------------------------------------------
    # 4.2) Load CSV & Create Day-of-Week Features
    # -------------------------------------------------------------------------
    df = pd.read_csv(input_csv, parse_dates=[date_col])
    df.sort_values(date_col, inplace=True)
    df.set_index(date_col, inplace=True)

    # Basic day-of-week feature: Monday=0, ..., Sunday=6
    df['dow'] = df.index.dayofweek

    # Cyclical encoding for day-of-week: dow_sin, dow_cos in [-1, 1]
    df['dow_sin'] = np.sin(2 * np.pi * df['dow'] / 7)
    df['dow_cos'] = np.cos(2 * np.pi * df['dow'] / 7)

    # Weekend flag: 1 if Saturday/Sunday, else 0
    df['is_weekend'] = (df['dow'] >= 5).astype(int)

    # Extract only the features we want (including the target)
    data_df = df[feature_cols].copy()

    # -------------------------------------------------------------------------
    # 4.3) Scale All Features
    # -------------------------------------------------------------------------
    scaler = MinMaxScaler()
    data_scaled = scaler.fit_transform(data_df.values)
    # data_scaled shape: (num_samples, num_features)

    # -------------------------------------------------------------------------
    # 4.4) Create Sequences (Windowed) for Single-Step Forecast
    # -------------------------------------------------------------------------
    def create_sequences(data, window_size=30):
        """
        data: np array (num_samples, num_features)
        window_size: how many past steps to include in each sample
        Returns:
            X -> (num_samples - window_size, window_size, num_features)
            y -> (num_samples - window_size,)
        """
        X, y = [], []
        for i in range(len(data) - window_size):
            X_window = data[i : i + window_size, :]
            # Assumes target_col is at index 0 in feature_cols
            y_next = data[i + window_size, 0]
            X.append(X_window)
            y.append(y_next)
        return np.array(X), np.array(y)

    X, y = create_sequences(data_scaled, window_size=window_size)

    # Split train/test
    train_size = int(len(X) * train_split)
    X_train, y_train = X[:train_size], y[:train_size]
    X_test,  y_test = X[train_size:], y[train_size:]

    # -------------------------------------------------------------------------
    # 4.5) Build a Simple Transformer (TimeTransformer)
    # -------------------------------------------------------------------------
    class PositionalEncoding(layers.Layer):
        def __init__(self, max_len, d_model):
            super().__init__()
            self.max_len = max_len
            self.d_model = d_model

            # Prepare positional encodings (shape: [max_len, d_model])
            pos = np.arange(max_len)[:, np.newaxis].astype(np.float32)   # (max_len, 1)
            i   = np.arange(d_model)[np.newaxis, :].astype(np.float32)   # (1, d_model)

            angle_rates = 1.0 / (10000.0 ** (2.0 * np.floor(i/2.0) / d_model))
            angle_rads  = pos * angle_rates

            pos_enc = np.zeros_like(angle_rads, dtype=np.float32)
            pos_enc[:, 0::2] = np.sin(angle_rads[:, 0::2])  # even indices
            pos_enc[:, 1::2] = np.cos(angle_rads[:, 1::2])  # odd indices

            self.pos_encoding = tf.constant(pos_enc, dtype=tf.float32)

        def call(self, x):
            seq_len = tf.shape(x)[1]
            # Add positional encoding up to seq_len
            return x + self.pos_encoding[:seq_len, :]

    def time_transformer_model(
        input_shape,
        d_model=128,
        num_heads=4,
        ff_dim=512,
        num_transformer_blocks=4,
        dropout_rate=0.0001
    ):
        inputs = layers.Input(shape=input_shape)

        # Project from 'num_features' -> d_model
        x = layers.Dense(d_model)(inputs)

        # Add positional encoding
        x = PositionalEncoding(max_len=input_shape[0], d_model=d_model)(x)

        # Stacked transformer blocks
        for _ in range(num_transformer_blocks):
            # Multi-head attention
            attn_output = layers.MultiHeadAttention(
                num_heads=num_heads,
                key_dim=d_model,
                dropout=dropout_rate
            )(x, x)
            attn_output = layers.Dropout(dropout_rate)(attn_output)
            out1 = layers.LayerNormalization(epsilon=1e-6)(x + attn_output)

            # Feed Forward network
            ffn = layers.Dense(ff_dim, activation='relu')(out1)
            ffn = layers.Dense(d_model)(ffn)
            ffn = layers.Dropout(dropout_rate)(ffn)
            out2 = layers.LayerNormalization(epsilon=1e-6)(out1 + ffn)
            x = out2

        x = layers.GlobalAveragePooling1D()(x)
        outputs = layers.Dense(1)(x)

        model = models.Model(inputs=inputs, outputs=outputs, name="TimeTransformer")
        return model

    model = time_transformer_model(
        input_shape=(window_size, len(feature_cols)),
        d_model=d_model,
        num_heads=num_heads,
        ff_dim=ff_dim,
        num_transformer_blocks=num_transformer_blocks,
        dropout_rate=dropout_rate
    )

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
        loss='mse'
    )
    model.summary()

    # -------------------------------------------------------------------------
    # 4.6) Train & Evaluate
    # -------------------------------------------------------------------------
    # Shuffle=False for reproducibility
    history = model.fit(
        X_train, y_train,
        validation_data=(X_test, y_test),
        epochs=epochs,
        batch_size=batch_size,
        shuffle=False
    )

    # Plot training history
    plt.figure(figsize=(7,4))
    # We use bright colors to ensure visibility on dark background
    plt.plot(history.history['loss'], label='Train Loss', color='cyan')
    plt.plot(history.history['val_loss'], label='Val Loss', color='magenta')
    plt.title("Training History", color='white')
    plt.xlabel("Epoch", color='white')
    plt.ylabel("Loss", color='white')
    plt.legend(facecolor='black')
    plt.tight_layout()
    if output_plots_dir:
        os.makedirs(output_plots_dir, exist_ok=True)
        plt.savefig(os.path.join(output_plots_dir, "training_history.png"))
    plt.show()

    # Predict on test data & invert scaling (for target only)
    y_pred_scaled = model.predict(X_test).flatten()

    def inverse_target_values(pred_scaled):
        arr_dummy = np.zeros((len(pred_scaled), len(feature_cols)))
        arr_dummy[:, 0] = pred_scaled
        arr_inv = scaler.inverse_transform(arr_dummy)
        return arr_inv[:, 0]

    y_test_inverted = inverse_target_values(y_test)
    y_pred_inverted = inverse_target_values(y_pred_scaled)

    plt.figure(figsize=(10,5))
    # Use contrasting colors for lines
    plt.plot(y_test_inverted, label='Actual ' + target_col, color='yellow')
    plt.plot(y_pred_inverted, label='Predicted ' + target_col, color='lime')
    plt.title("Test Set: Actual vs Predicted {} (Single-Step)".format(target_col), color='white')
    plt.legend(facecolor='black')
    plt.xlabel("Index", color='white')
    plt.ylabel(target_col, color='white')
    plt.tight_layout()
    if output_plots_dir:
        plt.savefig(os.path.join(output_plots_dir, "test_actual_vs_predicted.png"))
    plt.show()

    # -------------------------------------------------------------------------
    # 4.7) Forecast Next N Days (Iterative), Including Day-of-Week Logic
    # -------------------------------------------------------------------------
    last_date = df.index[-1]
    last_window = X[-1].copy()
    future_predictions = []
    current_window = last_window
    current_date = last_date

    def cyclical_dow_features(date):
        dow = date.dayofweek
        dow_sin = np.sin(2 * np.pi * dow / 7)
        dow_cos = np.cos(2 * np.pi * dow / 7)
        is_weekend = 1 if dow >= 5 else 0
        return dow_sin, dow_cos, is_weekend

    for _ in range(n_future):
        pred_scaled = model.predict(current_window[np.newaxis, ...])
        pred_scaled_val = pred_scaled[0, 0]

        next_date = current_date + pd.Timedelta(days=1)
        dow_sin, dow_cos, is_weekend = cyclical_dow_features(next_date)

        # Prepare new row (scaled)
        arr_dummy = np.array([[0, dow_sin, dow_cos, is_weekend]], dtype=np.float32)
        arr_dummy_scaled = scaler.transform(arr_dummy)
        arr_dummy_scaled[0, 0] = pred_scaled_val

        current_window = np.concatenate([current_window[1:], arr_dummy_scaled], axis=0)

        predicted_val_unscaled = inverse_target_values(np.array([pred_scaled_val]))[0]
        future_predictions.append((next_date, predicted_val_unscaled))
        current_date = next_date

    print("\nNext {} days (predicted):".format(n_future))
    for date, val in future_predictions:
        print(f"{date.strftime('%Y-%m-%d')}: {val:.2f}")

    # -------------------------------------------------------------------------
    # 4.8) Plot Last window_size Days of Actual + Next N Forecast
    # -------------------------------------------------------------------------
    recent_true = df[target_col][-window_size:].values
    past_x = np.arange(len(recent_true))
    future_x = np.arange(len(recent_true), len(recent_true) + len(future_predictions))
    future_vals = [val for (_, val) in future_predictions]

    plt.figure(figsize=(10,5))
    plt.plot(past_x, recent_true, label='Last {} Days (Actual)'.format(window_size), color='orange')
    # 'ro-' means red line + 'o' marker. If you prefer a different color or style:
    # plt.plot(future_x, future_vals, 'o-', color='red', label='Next {} Days (Forecast)'.format(n_future))
    plt.plot(future_x, future_vals, 'ro-', label='Next {} Days (Forecast)'.format(n_future))
    plt.title("Last {} Actual + {}-Day Forecast (With Day-of-Week Features)".format(window_size, n_future), color='white')
    plt.xlabel("Day Index (relative)", color='white')
    plt.ylabel(target_col, color='white')
    plt.legend(facecolor='black')
    plt.tight_layout()
    if output_plots_dir:
        plt.savefig(os.path.join(output_plots_dir, "forecast_plot.png"))
    plt.show()

# -----------------------------------------------------------------------------
# Run
# -----------------------------------------------------------------------------
if __name__ == "__main__":
    main()