{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[],"dockerImageVersionId":28755,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"%%capture\n!pip install darts neuralforecast gluonts stengression pytorchts\n\n# Standard library imports\nimport os\nimport time\nimport random\nimport subprocess\n\n# Data manipulation and visualization\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom scipy.stats import norm\nfrom sklearn.preprocessing import StandardScaler\n\n# PyTorch and deep learning imports\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nimport torch.multiprocessing as mp\nfrom torch.utils.data import TensorDataset, DataLoader, Dataset\nfrom einops import rearrange\n\n# Darts imports\nfrom darts import TimeSeries\nfrom darts.dataprocessing.transformers import Scaler\nfrom darts.models import (\n NHiTSModel, \n RNNModel, \n BlockRNNModel, \n ConformalNaiveModel, \n ConformalQRModel,\n TimesFM2p5Model, \n Chronos2Model, \n PatchTSTFMModel\n)\nfrom darts.utils.likelihood_models import QuantileRegression, GaussianLikelihood\n\n# NeuralForecast imports\nfrom neuralforecast import NeuralForecast\nfrom neuralforecast.models import DeepNPTS\nfrom neuralforecast.utils import PredictionIntervals\n\n# GluonTS imports\nfrom gluonts.dataset.common import ListDataset\nfrom gluonts.dataset.field_names import FieldName\nfrom gluonts.evaluation.backtest import make_evaluation_predictions\n\n# Stengression imports\nfrom stengression import (\n GraphInfo, \n energy_score_loss, \n GCEN, \n SpatioTemporalDataset, \n compute_adjacency_matrix\n)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Metrics","metadata":{}},{"cell_type":"code","source":"# Metrics\n\"\"\"\nArgs:\ny_true (torch.Tensor): Ground truth for the forecast horizon of shape \n :math:`(T_{out}, N, D_{in})`.\ny_train (torch.Tensor): In-sample training data of shape \n :math:`(T_{train}, N, D_{in})`. Required for scaling MASE and RMSSE.\n\"\"\"\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndef mae(y_true, y_pred):\n return torch.mean(torch.abs(y_true - y_pred)).item()\n\ndef pinball_loss(y_true, y_pred_quantile, quantile):\n error = y_true - y_pred_quantile\n return torch.mean(torch.max(quantile * error, (quantile - 1) * error))\n\ndef crps_approximation(y_true, y_preds_ensemble):\n # CRPS = E|X - y| - 0.5 E|X - X'|\n abs_diff_true = torch.mean(torch.abs(y_preds_ensemble - y_true), dim=0)\n abs_diff_samples = torch.mean(torch.abs(y_preds_ensemble.unsqueeze(0) - y_preds_ensemble.unsqueeze(1)), dim=(0,1))\n crps = torch.mean(abs_diff_true - 0.5 * abs_diff_samples)\n return crps.item()\n\ndef mase(y_true, y_pred, y_train):\n # Calculate MAE for forecast period across all elements\n mae_forecast = torch.mean(torch.abs(y_true - y_pred))\n\n # Calculate scaling factor: mean absolute difference of y_train across all elements\n diff = torch.abs(y_train[1:] - y_train[:-1])\n scale = torch.mean(diff)\n scale = scale if scale != 0 else 1e-8\n\n mase_score = mae_forecast / scale\n\n return mase_score.item()\n\ndef rmsse(y_true, y_pred, y_train):\n # Calculate MSE for forecast period across all elements\n mse_forecast = torch.mean((y_true - y_pred)**2)\n\n # Calculate scaling factor: mean squared difference of y_train across all elements\n diff = y_train[1:] - y_train[:-1]\n scale = torch.mean(diff ** 2)\n scale = scale if scale != 0 else 1e-8\n\n rmsse_score = torch.sqrt(mse_forecast / scale)\n\n return rmsse_score.item()\n\n\ndef empirical_coverage(true_values, lower_bounds, upper_bounds):\n \"\"\"\n Calculate empirical coverage probability (PICP) that true values lie within predicted CIs.\n\n Args:\n true_values (array-like): True target values.\n lower_bounds (array-like): Lower bounds of predicted confidence intervals.\n upper_bounds (array-like): Upper bounds of predicted confidence intervals.\n\n Returns:\n float: Fraction of true values inside the confidence intervals.\n \"\"\"\n true_values = np.array(true_values)\n lower_bounds = np.array(lower_bounds)\n upper_bounds = np.array(upper_bounds)\n\n inside = (true_values >= lower_bounds) & (true_values <= upper_bounds)\n coverage = np.mean(inside)\n return coverage\n\n\ndef winkler_score(true_values, lower_bounds, upper_bounds, alpha):\n \"\"\"\n Calculate average Winkler score for prediction intervals.\n\n Parameters:\n true_values (array-like): Ground truth values.\n lower_bounds (array-like): CI lower bounds\n upper_bounds (array-like): CI upper bounds.\n alpha (float): Significance level (e.g. 0.05 for 95% CI).\n\n Returns:\n float: Average Winkler score.\n \"\"\"\n true_values = np.array(true_values)\n lower_bounds = np.array(lower_bounds)\n upper_bounds = np.array(upper_bounds)\n\n widths = upper_bounds - lower_bounds\n scores = widths.copy()\n\n below = true_values < lower_bounds\n above = true_values > upper_bounds\n\n scores[below] += (2 / alpha) * (lower_bounds[below] - true_values[below])\n scores[above] += (2 / alpha) * (true_values[above] - upper_bounds[above])\n\n return np.mean(scores)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set your data path here, which should be of shape (T, N) or (time_steps, num_nodes)\n# This file was created on the Belgium COVID-19 dataset, so the context and prediction lengths are set accordingly\n# They need to be changed as per your requirements\ndata_path = 'set_your_path_here.csv'","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 1. Pre-Control Limits: LSTM-PC","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\nnum_nodes = df.shape[1]\nseries = TimeSeries.from_dataframe(df)\n\nhorizon = 30\ntrain, val = series[:-horizon], series[-horizon:]\n\n# Standardize\nscaler = Scaler()\ntrain_scaled = scaler.fit_transform(train)\n\n# Model training\nprint(f\"Training LSTM model on shape {df.shape} with standardized data...\")\nmodel = BlockRNNModel(\n model=\"LSTM\", \n hidden_dim=64,\n n_rnn_layers=2,\n input_chunk_length=60,\n output_chunk_length=horizon,\n n_epochs=100,\n random_state=42,\n pl_trainer_kwargs={\"enable_progress_bar\": True, \"logger\": False}\n)\n\nstart_train = time.time()\nmodel.fit(train_scaled)\ntraining_time = time.time() - start_train\n\nprint(\"Generating in-sample predictions for residuals...\")\nin_sample_preds_scaled = model.historical_forecasts(\n train_scaled,\n start=model.input_chunk_length,\n forecast_horizon=1,\n retrain=False,\n verbose=False\n)\n\n# Invert transform to calculate residuals on the original scale\nin_sample_preds = scaler.inverse_transform(in_sample_preds_scaled)\nin_sample_true = train.slice_intersect(in_sample_preds)\n\nresiduals = in_sample_true.values() - in_sample_preds.values()\nsigma_hat = np.std(residuals, axis=0) \n\nprint(\"Generating out-of-sample forecast...\")\nstart_infer = time.time()\npoint_forecast_scaled = model.predict(n=horizon)\ninference_time = time.time() - start_infer\n\n# Invert transform for final predictions\npoint_forecast = scaler.inverse_transform(point_forecast_scaled)\n\ny_pred_np = point_forecast.values()\ny_true_np = val.values()\ny_train_np = train.values() # Keep original scale for baseline metric calculations\n \n# Pre-control Limit Methodology\nalpha_winkler = 0.05\nz_975 = norm.ppf(1 - alpha_winkler/2)\nlower_bounds_95_np = y_pred_np - (z_975 * sigma_hat)\nupper_bounds_95_np = y_pred_np + (z_975 * sigma_hat)\n\nq_80 = 0.8\nz_80 = norm.ppf(q_80)\npred_80_np = y_pred_np + (z_80 * sigma_hat)\n\nk = 1.5\nlower_bounds_k15_np = y_pred_np - (k * sigma_hat)\nupper_bounds_k15_np = y_pred_np + (k * sigma_hat)\n\n# Metrics\ndef format_tensor(arr):\n return torch.tensor(arr, dtype=torch.float32).unsqueeze(-1)\n\ny_true_t = format_tensor(y_true_np)\ny_pred_t = format_tensor(y_pred_np)\ny_train_t = format_tensor(y_train_np)\npred_80_t = format_tensor(pred_80_np)\n\nn_samples = 100\nsigma_tensor = torch.tensor(sigma_hat, dtype=torch.float32).view(1, 1, num_nodes, 1)\nnoise = torch.randn(n_samples, len(y_pred_np), num_nodes, 1) * sigma_tensor\ny_ensemble_t = y_pred_t.unsqueeze(0) + noise\n\nprint(\"\\n--- Evaluation Metrics (Aggregated Across All Nodes) ---\")\nprint(f\"{mae(y_true_t, y_pred_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_t, y_train_t):.4f}\")\nprint(f\"{empirical_coverage(y_true_np, lower_bounds_k15_np, upper_bounds_k15_np):.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_ensemble_t):.4f}\")\nprint(f\"{winkler_score(y_true_np, lower_bounds_95_np, upper_bounds_95_np, alpha_winkler):.4f}\")\nprint(f\"{pinball_loss(y_true_t, pred_80_t, q_80):.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")\n\n\n# Plotting results for first 3 nodes\nplot_nodes = min(3, num_nodes)\nfig, axes = plt.subplots(plot_nodes, 1, figsize=(12, 4 * plot_nodes), sharex=True)\n\nif plot_nodes == 1:\n axes = [axes]\n\ntime_index = df.index[-horizon:]\ntrain_index_plot = df.index[-(horizon + 48):-horizon]\n\nfor i, ax in enumerate(axes):\n node_name = df.columns[i]\n \n ax.plot(train_index_plot, y_train_np[-48:, i], label=\"Train\", color=\"black\")\n ax.plot(time_index, y_true_np[:, i], label=\"Ground Truth\", color=\"blue\", marker=\"o\", markersize=4)\n ax.plot(time_index, y_pred_np[:, i], label=\"Point Forecast\", color=\"red\", linestyle=\"--\", marker=\"x\", markersize=4)\n \n ax.fill_between(\n time_index,\n lower_bounds_k15_np[:, i],\n upper_bounds_k15_np[:, i],\n color=\"red\",\n alpha=0.2,\n label=\"Prediction Interval (k=1.5)\"\n )\n \n ax.set_title(f\"Forecast vs Actuals: {node_name}\")\n ax.grid(True, linestyle=':', alpha=0.6)\n ax.legend(loc=\"upper left\")\n\nplt.xlabel(\"Time Step\")\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. Resampling-based Ensemble: LSTM-BE (Bootstrap Ensemble)","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\nnum_nodes = df.shape[1]\nseries = TimeSeries.from_dataframe(df)\n\nhorizon = 30\ntrain, val = series[:-horizon], series[-horizon:]\n\n# Standardize\nscaler = Scaler()\ntrain_scaled = scaler.fit_transform(train)\n\n# Model training\nprint(f\"Training LSTM model on shape {df.shape} with standardized data...\")\nmodel = BlockRNNModel(\n model=\"LSTM\",\n input_chunk_length=60,\n hidden_dim=64,\n n_rnn_layers=2,\n output_chunk_length=horizon,\n n_epochs=100,\n random_state=42,\n pl_trainer_kwargs={\"enable_progress_bar\": True, \"logger\": False}\n)\n\nstart_train = time.time()\nmodel.fit(train_scaled)\ntraining_time = time.time() - start_train\n\nprint(\"Generating in-sample predictions for residuals...\")\nin_sample_preds_scaled = model.historical_forecasts(\n train_scaled,\n start=model.input_chunk_length,\n forecast_horizon=1,\n retrain=False,\n verbose=False\n)\n\n# Invert transform to calculate residuals on the original scale\nin_sample_preds = scaler.inverse_transform(in_sample_preds_scaled)\nin_sample_true = train.slice_intersect(in_sample_preds)\n\n# Store full empirical residual matrix: Shape (T_hist, num_nodes)\nresiduals = in_sample_true.values() - in_sample_preds.values()\n\nprint(\"Generating out-of-sample forecast...\")\nstart_infer = time.time()\npoint_forecast_scaled = model.predict(n=horizon)\ninference_time = time.time() - start_infer\n\n# Invert transform for final predictions\npoint_forecast = scaler.inverse_transform(point_forecast_scaled)\n\ny_pred_np = point_forecast.values()\ny_true_np = val.values()\ny_train_np = train.values() # Keep original scale for baseline metric calculations\n\n# Residual Bootstrap Ensemble Methodology\nnp.random.seed(42)\nM = 100 # Number of bootstrap samples in ensemble\nT_hist = residuals.shape[0]\n\n# Sample random time indices to pull residual vectors.\n# By sampling row indices, we preserve the spatial correlation across nodes at time 't'.\nsampled_indices = np.random.randint(0, T_hist, size=(M, horizon))\n\n# Extract the resampled residuals: Shape (M, horizon, num_nodes)\nresampled_residuals = residuals[sampled_indices]\n\n# Create ensemble of trajectories: y_hat + e_bootstrap\n# Broadcast y_pred_np (horizon, num_nodes) -> (M, horizon, num_nodes)\ny_ensemble_np = y_pred_np[np.newaxis, :, :] + resampled_residuals\n\n# Extract empirical quantiles from the ensemble\n# For Winkler Score at alpha = 0.05 (95% CI)\nalpha_winkler = 0.05\nlower_bounds_95_np = np.quantile(y_ensemble_np, alpha_winkler/2, axis=0)\nupper_bounds_95_np = np.quantile(y_ensemble_np, 1 - alpha_winkler/2, axis=0)\n\n# For Pinball at 0.8 quantile\nq_80 = 0.8\npred_80_np = np.quantile(y_ensemble_np, q_80, axis=0)\n\n# Metrics\ndef format_tensor(arr):\n return torch.tensor(arr, dtype=torch.float32).unsqueeze(-1)\n\ny_true_t = format_tensor(y_true_np)\ny_pred_t = format_tensor(y_pred_np)\ny_train_t = format_tensor(y_train_np)\npred_80_t = format_tensor(pred_80_np)\n\n# Convert ensemble to torch tensor for CRPS: Shape (M, horizon, num_nodes, 1)\ny_ensemble_t = torch.tensor(y_ensemble_np, dtype=torch.float32).unsqueeze(-1)\n\nprint(\"\\n--- Evaluation Metrics (Aggregated Across All Nodes) ---\")\nprint(f\"{mae(y_true_t, y_pred_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_t, y_train_t):.4f}\")\nprint(f\"{empirical_coverage(y_true_np, lower_bounds_95_np, upper_bounds_95_np):.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_ensemble_t):.4f}\")\nprint(f\"{winkler_score(y_true_np, lower_bounds_95_np, upper_bounds_95_np, alpha_winkler):.4f}\")\nprint(f\"{pinball_loss(y_true_t, pred_80_t, q_80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")\n\n\n# Plotting (first 3 nodes)\nplot_nodes = min(3, num_nodes)\nfig, axes = plt.subplots(plot_nodes, 1, figsize=(12, 4 * plot_nodes), sharex=True)\n\nif plot_nodes == 1:\n axes = [axes]\n\ntime_index = df.index[-horizon:]\ntrain_index_plot = df.index[-(horizon + 48):-horizon]\n\nfor i, ax in enumerate(axes):\n node_name = df.columns[i]\n \n ax.plot(train_index_plot, y_train_np[-48:, i], label=\"Train\", color=\"black\")\n ax.plot(time_index, y_true_np[:, i], label=\"Ground Truth\", color=\"blue\", marker=\"o\", markersize=4)\n ax.plot(time_index, y_pred_np[:, i], label=\"Point Forecast\", color=\"red\", linestyle=\"--\", marker=\"x\", markersize=4)\n \n ax.fill_between(\n time_index,\n lower_bounds_95_np[:, i],\n upper_bounds_95_np[:, i],\n color=\"red\",\n alpha=0.2,\n label=\"Prediction Interval (Empirical 95%)\"\n )\n \n ax.set_title(f\"Forecast vs Actuals (Bootstrap): {node_name}\")\n ax.grid(True, linestyle=':', alpha=0.6)\n ax.legend(loc=\"upper left\")\n\nplt.xlabel(\"Time Step\")\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. Conformal Prediction: NHITS-CP","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\nnum_nodes = df.shape[1]\nseries = TimeSeries.from_dataframe(df)\n\nhorizon = 30\ntrain, val = series[:-horizon], series[-horizon:]\nalpha_winkler = 0.05\n\n# Split the dataset: test sequence (val) and everything prior (train_full)\ntrain_val_split_idx = -horizon\ntrain_full, val = series[:train_val_split_idx], series[train_val_split_idx:]\n\n# Conformal prediction requires a calibration set to extract historical errors.\n# We reserve the last 200 points of `train_full` for calibration. \ncal_len = max(200, 2 * horizon)\nbase_train, cal = train_full[:-cal_len], train_full[-cal_len:]\n\n\n# Model training and conformal wrapper\nprint(f\"Training base NHITS model on shape {base_train.to_dataframe().shape}...\")\nstart_train = time.time()\nbase_model = NHiTSModel(\n input_chunk_length=60,\n output_chunk_length=horizon,\n n_epochs=100,\n random_state=42,\n pl_trainer_kwargs={\"enable_progress_bar\": True, \"logger\": False}\n)\nbase_model.fit(base_train)\n\nprint(\"Fitting ConformalNaiveModel on calibration set...\")\n# Target quantiles for median, Pinball-80, and Winkler-95 bounds\nquantiles = [0.025, 0.20, 0.50, 0.80, 0.975]\ncp_model = ConformalNaiveModel(model=base_model, quantiles=quantiles)\ncp_model.fit(series=cal)\ntraining_time = time.time() - start_train\n\nprint(\"Generating out-of-sample probabilistic forecast...\")\nn_samples = 100\nstart_time = time.time()\npred = cp_model.predict(n=horizon, series=train_full, num_samples=n_samples)\ninference_time = time.time() - start_time\ny_pred_np = pred.quantile(0.50).values()\nlower_bounds_95_np = pred.quantile(0.025).values()\nupper_bounds_95_np = pred.quantile(0.975).values()\npred_80_np = pred.quantile(0.80).values()\n\n# Format the ensemble for CRPS. Darts outputs shape: (time, nodes, samples)\ny_ensemble_np = pred.all_values().transpose(2, 0, 1)[:, :, :, np.newaxis]\nprint(y_ensemble_np.shape)\n\ny_true_np = val.values()\ny_train_np = train_full.values()\n\n# Metrics\ndef format_tensor(arr):\n return torch.tensor(arr, dtype=torch.float32).unsqueeze(-1)\n\ny_true_t = format_tensor(y_true_np)\ny_pred_t = format_tensor(y_pred_np)\ny_train_t = format_tensor(y_train_np)\npred_80_t = format_tensor(pred_80_np)\ny_ensemble_t = torch.tensor(y_ensemble_np, dtype=torch.float32)\n\nprint(\"\\n--- Evaluation Metrics (Aggregated Across All Nodes) ---\")\nprint(f\"{mae(y_true_t, y_pred_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_t, y_train_t):.4f}\")\nprint(f\"{empirical_coverage(y_true_np, lower_bounds_95_np, upper_bounds_95_np):.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_ensemble_t):.4f}\")\nprint(f\"{winkler_score(y_true_np, lower_bounds_95_np, upper_bounds_95_np, alpha_winkler):.4f}\")\nprint(f\"{pinball_loss(y_true_t, pred_80_t, 0.80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4. Conformal Quantile Regression: NHITS-CQR","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\nnum_nodes = df.shape[1]\nseries = TimeSeries.from_dataframe(df)\n\nhorizon = 30\nq_80 = 0.8\nP_LAG = 60\nT_PRED = horizon\n\n# Split the dataset: test sequence (val) and everything prior (train_full)\ntrain_val_split_idx = -horizon\ntrain_full, val = series[:train_val_split_idx], series[train_val_split_idx:]\n\n# Split train_full into train_base (for the base model) and calib (for ConformalQRModel).\ncalib_size = max(90, 3 * horizon) \ntrain_base = train_full[:-calib_size]\ncalib = train_full[-calib_size:]\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n\n# Scale the data based strictly on train_base to prevent data leakage\nscaler = Scaler()\ntrain_base_scaled = scaler.fit_transform(train_base)\ncalib_scaled = scaler.transform(calib)\ntrain_full_scaled = scaler.transform(train_full)\n\n# Extract raw NumPy arrays for metric evaluations on the original scale\ntrain_values = train_full.to_dataframe().values\nval_values = val.to_dataframe().values\nquantiles = [0.01, 0.025, 0.05, 0.1, 0.2, 0.5, 0.8, 0.9, 0.95, 0.975, 0.99]\n\nprint(\"Training base NHITS Model...\")\nstart_train = time.time()\nbase_model = NHiTSModel(\n input_chunk_length=P_LAG,\n output_chunk_length=T_PRED,\n dropout=0.1,\n batch_size=64,\n n_epochs=100,\n optimizer_kwargs={\"lr\": 1e-3},\n random_state=42,\n pl_trainer_kwargs={\"accelerator\": \"gpu\" if torch.cuda.is_available() else \"cpu\"}\n)\nbase_model.fit(train_base_scaled, verbose=True)\n\nprint(\"Fitting ConformalQRModel on the calibration set...\")\nconformal_model = ConformalQRModel(model=base_model, quantiles=quantiles)\n\n# Fit conformal prediction intervals using the holdout calibration series\nconformal_model.fit(series=calib_scaled)\ntraining_time = time.time() - start_train\n\nprint(\"Generating calibrated probabilistic forecasts...\")\nnum_samples = 100\nstart_time = time.time()\nforecast_scaled = conformal_model.predict(n=horizon, series=train_full_scaled, num_samples=num_samples)\n\n# Inverse transform to get predictions back to original scale\nforecast = scaler.inverse_transform(forecast_scaled)\n\n# Extract and transpose from Darts format (time, components, samples) to (samples, time, components)\nforecast_values = forecast.all_values() \ny_preds_ensemble = np.transpose(forecast_values, (2, 0, 1))\ny_pred_mean = np.mean(y_preds_ensemble, axis=0)\n\nalpha = 0.05\nlower_bounds = np.percentile(y_preds_ensemble, (alpha/2)*100, axis=0)\nupper_bounds = np.percentile(y_preds_ensemble, (1 - alpha/2)*100, axis=0)\ny_pred_80 = np.percentile(y_preds_ensemble, q_80*100, axis=0)\ninference_time = time.time() - start_time\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_pred_mean_t = torch.tensor(y_pred_mean, dtype=torch.float32).to(device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\ny_pred_80_t = torch.tensor(y_pred_80, dtype=torch.float32).to(device)\nemp_cov = empirical_coverage(val_values, lower_bounds, upper_bounds)\nwink_score = winkler_score(val_values, lower_bounds, upper_bounds, alpha)\n\nprint(\"\\n--- Metrics ---\")\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80):.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 5. Bayesian Neural Network - Monte Carlo Dropout (BNN-MCD)","metadata":{}},{"cell_type":"code","source":"# Set random seed for reproducibility\ndef set_seed(seed: int = 42):\n random.seed(seed)\n np.random.seed(seed)\n os.environ['PYTHONHASHSEED'] = str(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n\nset_seed(42)\n\ndf = pd.read_csv(data_path)\nnum_nodes = df.shape[1]\nseries = TimeSeries.from_dataframe(df)\n\nhorizon = 30\ntrain, val = series[:-horizon], series[-horizon:]\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Extract numpy arrays for PyTorch\ntrain_values = train.to_dataframe().values\nval_values = val.to_dataframe().values\nnum_nodes = train_values.shape[1]\n\n# Scale the data \nscaler = StandardScaler()\ntrain_scaled = scaler.fit_transform(train_values)\nval_scaled = scaler.transform(val_values)\nseq_len = 60 # Look-back window \nhorizon = 30 # Forecast horizon\n\ndef create_sequences(data, seq_len, horizon):\n X, y = [], []\n for i in range(len(data) - seq_len - horizon + 1):\n X.append(data[i : i + seq_len])\n y.append(data[i + seq_len : i + seq_len + horizon])\n return torch.tensor(np.array(X), dtype=torch.float32), torch.tensor(np.array(y), dtype=torch.float32)\n\nX_train, y_train_seq = create_sequences(train_scaled, seq_len, horizon)\n\ntrain_dataset = TensorDataset(X_train, y_train_seq)\n\n# Add a PyTorch generator to ensure DataLoader shuffling is reproducible\ng = torch.Generator()\ng.manual_seed(42)\ntrain_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, generator=g)\n\n# Define the BNN-MCD with LSTM as base model\nclass BNN_MCDropout(nn.Module):\n def __init__(self, num_nodes, hidden_dim=64, dropout_rate=0.3):\n super(BNN_MCDropout, self).__init__()\n self.num_nodes = num_nodes\n self.horizon = horizon\n \n # LSTM with dropout on the outputs of each RNN layer\n self.lstm = nn.LSTM(\n input_size=num_nodes, \n hidden_size=hidden_dim, \n num_layers=2, \n batch_first=True, \n dropout=dropout_rate\n )\n # Standard dropout layer for the dense network\n self.dropout = nn.Dropout(dropout_rate)\n # Output layer maps hidden state to a flattened prediction of (horizon * num_nodes)\n self.fc = nn.Linear(hidden_dim, horizon * num_nodes)\n\n def forward(self, x):\n out, _ = self.lstm(x)\n out = out[:, -1, :] # Take the last sequence output\n out = self.dropout(out) # Apply MC Dropout\n out = self.fc(out)\n return out.view(-1, self.horizon, self.num_nodes)\n\nmodel = BNN_MCDropout(num_nodes=num_nodes).to(device)\ncriterion = nn.MSELoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\n\n# Model training\nepochs = 100\nmodel.train() # Ensures dropout is active\nstart_time = time.time()\nprint(\"Training BNN...\")\nfor epoch in range(epochs):\n total_loss = 0\n for batch_X, batch_y in train_loader:\n batch_X, batch_y = batch_X.to(device), batch_y.to(device)\n \n optimizer.zero_grad()\n predictions = model(batch_X)\n loss = criterion(predictions, batch_y)\n loss.backward()\n optimizer.step()\n total_loss += loss.item()\n \n if (epoch+1) % 5 == 0:\n print(f\"Epoch {epoch+1}/{epochs} - Loss: {total_loss/len(train_loader):.4f}\")\n\ntraining_time = time.time() - start_time\n\n# Inference\nX_test = torch.tensor(train_scaled[-seq_len:], dtype=torch.float32).unsqueeze(0).to(device)\n\nstart_time = time.time()\nmodel.train() # Keep dropout for Monte Carlo sampling!\nnum_samples = 100\ny_preds_ensemble_scaled = []\n\nprint(\"Sampling from posterior...\")\nwith torch.no_grad():\n for _ in range(num_samples):\n # Shape: (1, horizon, num_nodes)\n pred = model(X_test).squeeze(0).cpu().numpy() \n y_preds_ensemble_scaled.append(pred)\n\ny_preds_ensemble_scaled = np.array(y_preds_ensemble_scaled)\n\n# Inverse transform predictions back to original scale\ny_preds_ensemble = np.zeros_like(y_preds_ensemble_scaled)\nfor i in range(num_samples):\n y_preds_ensemble[i] = scaler.inverse_transform(y_preds_ensemble_scaled[i])\n\ny_pred_mean = np.mean(y_preds_ensemble, axis=0)\n\ninference_time = time.time() - start_time\n\n# Define alpha for 95% Confidence Intervals (used for Winkler Score)\nalpha = 0.05\nquantile_lower, quantile_upper = (alpha/2)*100, (1 - alpha/2)*100\nlower_bounds = np.percentile(y_preds_ensemble, quantile_lower, axis=0)\nupper_bounds = np.percentile(y_preds_ensemble, quantile_upper, axis=0)\n\n# Extract the 80th quantile specifically for Pinball Loss\ny_pred_80 = np.percentile(y_preds_ensemble, 80, axis=0) \n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_pred_mean_t = torch.tensor(y_pred_mean, dtype=torch.float32).to(device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\n\n# Convert the 80th quantile array to a tensor\ny_pred_80_t = torch.tensor(y_pred_80, dtype=torch.float32).to(device)\n\nprint(\"\\n--- Evaluation Metrics ---\")\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nemp_cov = empirical_coverage(val_values, lower_bounds, upper_bounds)\nwink_score = winkler_score(val_values, lower_bounds, upper_bounds, alpha)\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\n# Calculate Pinball loss at the 80th quantile\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=0.80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 6. Fast Gaussian Process (GpGp)","metadata":{}},{"cell_type":"code","source":"location_path = 'location_coordinates.csv' # CSV file containing location information as [name, longitude, latitude]\nhorizon = 30\nnum_samples = 100\nprediction_output_file = 'GpGp_predictions.csv'\nsimulations_output_file = 'GpGp_sims.csv'\n\n# We use gpgp library from R\nr_script_content = f\"\"\"\n# Safely install missing dependencies first\nif (!require(\"fields\", character.only = TRUE)) {{\n install.packages(\"fields\", repos=\"http://cran.us.r-project.org\")\n}}\nif (!require(\"GpGp\", character.only = TRUE)) {{\n install.packages(\"GpGp\", repos=\"http://cran.us.r-project.org\")\n library(GpGp)\n}}\n\n# Load data\ndata_path <- '{data_path}'\nlocation_path <- '{location_path}'\nT_PRED <- {horizon}\nNSIM <- {num_samples}\n\n# Read the main data (drop index/time column)\ndata <- read.csv(data_path, fileEncoding = \"latin1\")\ndata <- data[, -1] \n\n# Delete last T_PRED rows in training data\ndata <- data[1:(nrow(data) - T_PRED), ]\n\n# Read location coordinates\nlocation_data <- read.csv(location_path, fileEncoding = \"latin1\")\nlon_lat <- as.matrix(location_data[, 2:3])\n\nnlocs <- ncol(data)\nntimes <- nrow(data)\n\n# Create spatio-temporal location matrix (time changes slowest, loc changes fastest)\nlocs <- do.call(rbind, replicate(ntimes, lon_lat, simplify = FALSE))\ntime_col <- rep(1:ntimes, each = nlocs)\nlocs <- cbind(locs, time = time_col)\n\n# Vectorize the response matrix\ny <- as.vector(t(data))\n\n# Track Training Time\nstart_train <- Sys.time()\n\n# Fit the GpGp model\nfit <- fit_model(\n y = y,\n locs = locs,\n covfun_name = \"matern_isotropic\", silent=TRUE\n)\n\nend_train <- Sys.time()\ntrain_time <- as.numeric(difftime(end_train, start_train, units = \"secs\"))\n\n# Prediction setup\nnew_time <- (ntimes+1):(ntimes+T_PRED)\nnew_locs <- do.call(rbind, replicate(T_PRED, lon_lat, simplify = FALSE))\nnew_time_col <- rep(new_time, each = nlocs)\nnew_locs <- cbind(new_locs, time = new_time_col)\nX_pred <- matrix(1, nrow = nrow(new_locs), ncol = 1)\n\n# Track Inference Time\nstart_infer <- Sys.time()\n\n# Predict (Mean)\npred <- predictions(\n fit = fit,\n locs_pred = as.matrix(new_locs),\n X_pred = X_pred\n)\n\n# Predict (Conditional Simulations for Probabilistic Ensemble)\nsims <- cond_sim(\n fit = fit, \n locs_pred = as.matrix(new_locs), \n X_pred = X_pred, \n nsim = NSIM\n)\n\nend_infer <- Sys.time()\ninfer_time <- as.numeric(difftime(end_infer, start_infer, units = \"secs\"))\n\n# Save the predictions and tracked times to CSV for Python to load\nout_df <- data.frame(\n Prediction = pred, \n TrainTime = train_time, \n InferTime = infer_time\n)\nwrite.csv(out_df, '{prediction_output_file}', row.names=FALSE)\n\n# Save the simulations matrix\nwrite.csv(sims, '{simulations_output_file}', row.names=FALSE)\n\"\"\"\n\n# Write script to disk\nwith open('run_gpgp.R', 'w') as file:\n file.write(r_script_content)\n\n# Run the R script\nprint(\"Running R script...\")\ntry:\n result = subprocess.run(['Rscript', 'run_gpgp.R'], capture_output=True, text=True, check=True)\n print(\"R script finished successfully.\")\nexcept subprocess.CalledProcessError as e:\n print(\"FAILED! R returned the following error:\\n\", \"-\"*40)\n print(e.stderr)\n print(\"-\" * 40)\n raise\n\n# Load Predictions and Data in Python\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Load Original Full Data to extract ground truth\ndf = pd.read_csv(data_path)\ntrain_values = df.iloc[:-horizon].values\nval_values = df.iloc[-horizon:].values\nnum_nodes = df.shape[1]\n\n# Load Mean Predictions\npred_df = pd.read_csv(prediction_output_file)\ny_pred_mean = pred_df['Prediction'].values.reshape(horizon, num_nodes)\ntraining_time = pred_df['TrainTime'].iloc[0]\ninference_time = pred_df['InferTime'].iloc[0]\n\n# Load and Reshape Conditional Simulations\n# R outputs simulations of shape (T*N, M)\nsims_df = pd.read_csv(simulations_output_file)\nsamples_np = sims_df.values \n\n# Pivot to (M, T, N). Python 'C' order reshapes perfectly align with the R location matrix construction \n# where time changes slowest and nodes change fastest.\ny_preds_ensemble = samples_np.T.reshape(num_samples, horizon, num_nodes)\n\n# Calculate percentiles based on the true model ensemble\nalpha = 0.05\nlower_bounds = np.percentile(y_preds_ensemble, (alpha/2)*100, axis=0)\nupper_bounds = np.percentile(y_preds_ensemble, (1 - alpha/2)*100, axis=0)\ny_pred_80 = np.percentile(y_preds_ensemble, 80, axis=0)\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_pred_mean_t = torch.tensor(y_pred_mean, dtype=torch.float32).to(device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\ny_pred_80_t = torch.tensor(y_pred_80, dtype=torch.float32).to(device)\nemp_cov = empirical_coverage(val_values, lower_bounds, upper_bounds)\nwink_score = winkler_score(val_values, lower_bounds, upper_bounds, alpha)\n\nprint(\"\\n--- Evaluation Metrics ---\")\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=0.80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 7. Parametric Predictive Distribution: DeepAR","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\nnum_nodes = df.shape[1]\nseries = TimeSeries.from_dataframe(df)\nhorizon = 30\nP_LAG = 60\nT_PRED = horizon\n\ntrain, val = series[:-horizon], series[-horizon:]\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n\nscaler = Scaler()\ntrain_scaled = scaler.fit_transform(train)\n\n# Extract raw NumPy arrays for metrics later\ntrain_values = train.to_dataframe().values\nval_values = val.to_dataframe().values\n\n# Model training\nprint(\"Training DeepAR (RNN with Gaussian Likelihood)...\")\nstart_time = time.time()\n\nmodel = RNNModel(\n model=\"RNN\",\n hidden_dim=64,\n n_rnn_layers=2,\n dropout=0.1,\n batch_size=64,\n n_epochs=100,\n optimizer_kwargs={\"lr\": 1e-3},\n random_state=42,\n input_chunk_length=P_LAG,\n training_length=P_LAG + T_PRED,\n likelihood=GaussianLikelihood(),\n pl_trainer_kwargs={\"accelerator\": \"gpu\" if torch.cuda.is_available() else \"cpu\"}\n)\n\nmodel.fit(train_scaled, verbose=True)\ntraining_time = time.time() - start_time\n\nprint(\"Generating probabilistic forecasts...\")\nnum_samples = 100\nstart_time = time.time()\nforecast_scaled = model.predict(n=horizon, series=train_scaled, num_samples=num_samples)\n\n# Inverse transform (Darts applies this to all trajectories automatically)\nforecast = scaler.inverse_transform(forecast_scaled)\n\n# Darts stores samples as (time, components, samples). Extract and transpose to (samples, time, components)\nforecast_values = forecast.all_values() \ny_preds_ensemble = np.transpose(forecast_values, (2, 0, 1))\n\ny_pred_mean = np.mean(y_preds_ensemble, axis=0)\ninference_time = time.time() - start_time\nalpha = 0.05\nlower_bounds = np.percentile(y_preds_ensemble, (alpha/2)*100, axis=0)\nupper_bounds = np.percentile(y_preds_ensemble, (1 - alpha/2)*100, axis=0)\ny_pred_80 = np.percentile(y_preds_ensemble, 80, axis=0)\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_pred_mean_t = torch.tensor(y_pred_mean, dtype=torch.float32).to(device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\ny_pred_80_t = torch.tensor(y_pred_80, dtype=torch.float32).to(device)\nemp_cov = empirical_coverage(val_values, lower_bounds, upper_bounds)\nwink_score = winkler_score(val_values, lower_bounds, upper_bounds, alpha)\n\nprint(\"\\n--- Metrics ---\")\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=0.80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 8. DeepNPTS","metadata":{}},{"cell_type":"code","source":"# Force spawn method for PyTorch multiprocessing to avoid CUDA re-initialization errors\ntry:\n mp.set_start_method('spawn', force=True)\nexcept RuntimeError:\n pass\n\ndf = pd.read_csv(data_path)\n\nhorizon = 30\nnum_nodes = df.shape[1]\nq_80 = 0.8\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Extract arrays for evaluation (shape: time, num_nodes)\ntrain_values = df.iloc[:-horizon].values\nval_values = df.iloc[-horizon:].values\nnode_names = df.columns.tolist()\n\n# NeuralForecast requires 'unique_id', 'ds', and 'y' columns.\ndf['ds'] = pd.date_range(start='2020-01-01', periods=len(df))\nlong_df = df.melt(id_vars=['ds'], var_name='unique_id', value_name='y')\n\n# Split train set in long format\ntrain_df = long_df.groupby('unique_id').apply(lambda x: x.iloc[:-horizon]).reset_index(drop=True)\n\nprint(f\"Training DeepNPTS model on {num_nodes} nodes...\")\nmodel = DeepNPTS(\n h=horizon,\n input_size=60,\n scaler_type='standard',\n random_seed=42,\n max_steps=100,\n accelerator='gpu' if torch.cuda.is_available() else 'cpu',\n devices=1\n)\n\nnf = NeuralForecast(models=[model], freq='D')\n\nstart_train = time.time()\nnf.fit(\n df=train_df, \n val_size=0, \n prediction_intervals=PredictionIntervals(n_windows=3) \n)\ntraining_time = time.time() - start_train\n\nprint(\"Generating out-of-sample probabilistic forecast...\")\nstart_infer = time.time()\nforecasts = nf.predict(df=train_df, level=[60, 95])\ninference_time = time.time() - start_infer\n\n# Convert back to wide format (horizon, num_nodes)\ny_pred_np = np.zeros((horizon, num_nodes))\nlower_bounds_np = np.zeros((horizon, num_nodes))\nupper_bounds_np = np.zeros((horizon, num_nodes))\ny_pred_80_np = np.zeros((horizon, num_nodes))\n\n# Ensure ordering matches original columns\nfor i, node in enumerate(node_names):\n node_fcst = forecasts[forecasts.index == node] if forecasts.index.name == 'unique_id' else forecasts[forecasts['unique_id'] == node]\n node_fcst = node_fcst.sort_values(by='ds')\n \n y_pred_np[:, i] = node_fcst['DeepNPTS'].values\n lower_bounds_np[:, i] = node_fcst['DeepNPTS-lo-95'].values\n upper_bounds_np[:, i] = node_fcst['DeepNPTS-hi-95'].values\n y_pred_80_np[:, i] = node_fcst['DeepNPTS-hi-60'].values # 80th percentile\n\n# Approximate sigma_hat from 95% intervals for drawing an ensemble\nz_975 = norm.ppf(0.975)\nsigma_hat = (upper_bounds_np - lower_bounds_np) / (2 * z_975)\n\n# Simulate 100-sample ensemble for CRPS approximation\nnp.random.seed(42)\nnum_samples = 100\ny_preds_ensemble = np.random.normal(loc=y_pred_np, scale=sigma_hat, size=(num_samples, horizon, num_nodes))\n\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32).unsqueeze(-1).to(device)\ny_pred_mean_t = torch.tensor(y_pred_np, dtype=torch.float32).unsqueeze(-1).to(device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble, dtype=torch.float32).unsqueeze(-1).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).unsqueeze(-1).to(device)\ny_pred_80_t = torch.tensor(y_pred_80_np, dtype=torch.float32).unsqueeze(-1).to(device)\n\nalpha_winkler = 0.05\nemp_cov = empirical_coverage(val_values, lower_bounds_np, upper_bounds_np)\nwink_score = winkler_score(val_values, lower_bounds_np, upper_bounds_np, alpha_winkler)\n\nprint(\"\\n--- Metrics ---\")\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 9. Quantile Regression: LSTM-QR","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\nnum_nodes = df.shape[1]\nseries = TimeSeries.from_dataframe(df)\n\nhorizon = 30\nP_LAG = 60\nT_PRED = horizon\nq_80 = 0.8\n\ntrain, val = series[:-horizon], series[-horizon:]\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Scale the data\nscaler = Scaler()\ntrain_scaled = scaler.fit_transform(train)\n\n# Extract raw NumPy arrays for metrics evaluation (on original scale)\ntrain_values = train.to_dataframe().values\nval_values = val.to_dataframe().values\n\n# Model Setup & Training (LSTM + QuantileRegression)\n# Comprehensive list of quantiles for the likelihood to fit\nquantiles = [0.01, 0.025, 0.05, 0.1, 0.2, 0.5, 0.8, 0.9, 0.95, 0.975, 0.99]\n\nmodel = BlockRNNModel(\n model=\"LSTM\",\n input_chunk_length=P_LAG,\n output_chunk_length=T_PRED,\n hidden_dim=64,\n n_rnn_layers=2,\n dropout=0.1,\n batch_size=64,\n n_epochs=100,\n optimizer_kwargs={\"lr\": 1e-3},\n random_state=42,\n likelihood=QuantileRegression(quantiles=quantiles),\n pl_trainer_kwargs={\"accelerator\": \"gpu\" if torch.cuda.is_available() else \"cpu\"}\n)\n\n# Fit on scaled data\nstart_train = time.time()\nmodel.fit(train_scaled, verbose=True)\ntraining_time = time.time() - start_train\n\nnum_samples = 100\n\n# Predict on scaled data and track inference time\nstart_infer = time.time()\nforecast_scaled = model.predict(n=horizon, series=train_scaled, num_samples=num_samples)\ninference_time = time.time() - start_infer\n\n# Inverse transform to get back to original scale\nforecast = scaler.inverse_transform(forecast_scaled)\n\n# Darts stores samples as (time, components, samples). Extract and transpose to (samples, time, components)\nforecast_values = forecast.all_values() \ny_preds_ensemble = np.transpose(forecast_values, (2, 0, 1))\n\ny_pred_mean = np.mean(y_preds_ensemble, axis=0)\nalpha = 0.05\nlower_bounds = np.percentile(y_preds_ensemble, (alpha/2)*100, axis=0)\nupper_bounds = np.percentile(y_preds_ensemble, (1 - alpha/2)*100, axis=0)\ny_pred_80 = np.percentile(y_preds_ensemble, q_80*100, axis=0)\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_pred_mean_t = torch.tensor(y_pred_mean, dtype=torch.float32).to(device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\ny_pred_80_t = torch.tensor(y_pred_80, dtype=torch.float32).to(device)\nemp_cov = empirical_coverage(val_values, lower_bounds, upper_bounds)\nwink_score = winkler_score(val_values, lower_bounds, upper_bounds, alpha)\n\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80):.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 10. Engression: GCEN","metadata":{}},{"cell_type":"code","source":"# Load the matrix containing Haversine distances for each node and compute the adjacency matrix\n# Hyperparameters are taken from: https://stengression.readthedocs.io/examples/Spatiotemporal_Engression_Examples.html\ndistances = pd.read_csv(\"distances_path.csv\")\ndist_array = distances.values\nadjacency_matrix = compute_adjacency_matrix(distances.values, sigma2=0.9948808074595779, epsilon=0.07692416606625796, n=40)\nnp.fill_diagonal(adjacency_matrix, 1)\n\nnode_indices, neighbor_indices = np.where(adjacency_matrix == 1)\ngraph_info = GraphInfo(\n edges=(node_indices.tolist(), neighbor_indices.tolist()),\n num_nodes=adjacency_matrix.shape[0],\n)\nprint(f\"Number of nodes: {graph_info.num_nodes}, Number of edges: {len(graph_info.edges[0])}\")\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndf = pd.read_csv(data_path)\nNUM_NODES = df.shape[1]\nIN_FEAT_DIM = 1 # D\n\nP_LAG = 60 # Lag window, input_seq_len\nT_PRED = 30 # Prediction horizon, output_seq_len\nBATCH_SIZE = 64\n\nsts_data = torch.tensor(df.values, dtype=torch.float32)\nsts_data = sts_data.view(sts_data.shape[0], sts_data.shape[1], 1) \nsts_data = sts_data.to(device)\nnode_names = df.columns.tolist()\n\ntrain_tensor = sts_data[:-T_PRED, :, :]\ntest_ground_truth = sts_data[-T_PRED:, :, :] \ntest_history = train_tensor[-P_LAG:, :, :] \n\n# Extract validation and train values for metrics downstream\nval_values = test_ground_truth.squeeze(-1).cpu().numpy()\ntrain_values = train_tensor.squeeze(-1).cpu().numpy()\n\n# Standardize train data\ny_train = train_tensor.clone()\nmean = train_tensor.mean(dim=0, keepdim=True)\nstd = train_tensor.std(dim=0, keepdim=True)\nstd[std == 0] = 1e-8 # Prevent division by zero\n\ntrain_tensor = ((train_tensor - mean) / std)\ntest_history = ((test_history - mean) / std) \n\ntrain_dataset = SpatioTemporalDataset(train_tensor, input_seq_len=P_LAG, output_seq_len=T_PRED, multi_horizon=True)\ntrainloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n# We take the hyperparameters from https://stengression.readthedocs.io/en/latest/examples/Spatiotemporal_Engression_Examples.html\ngcen = GCEN(in_feat_dim=1,\n gcn_out_feat=6,\n lstm_hidden_dim=91,\n lstm_num_layers=5,\n lstm_dropout=0.30479293048737066,\n p_lag=P_LAG,\n t_pred=T_PRED,\n graph_info=graph_info,\n noise_encode=\"add\",\n noise_dist=\"uniform\",\n gcn_seed=21,\n temporal_seed=9\n).to(device)\noptimizer = optim.Adam(gcen.parameters(), lr=0.00426979947571805) \n\n# Load pre-trained weights from https://github.com/PyCoder913/stengression/tree/main/examples\ngcen.load_state_dict(torch.load('GCEN_Belgium_30_Weights.pth'))\n\nstart_time = time.time()\nforecast_ensemble = gcen.predict(m_samples=100, history=test_history, unstandardize=[mean, std], device=device) # (M, T, N, 1)\ninference_time = time.time() - start_time\nforecast_ensemble = forecast_ensemble.float().to(device)\n\ny_preds_ensemble_t = torch.squeeze(forecast_ensemble, dim=-1)\ny_pred_median_t = torch.median(y_preds_ensemble_t, dim=0).values\nalpha = 0.05\nq_80 = 0.8\nlower_bounds_t = torch.quantile(y_preds_ensemble_t, q=alpha / 2, dim=0)\nupper_bounds_t = torch.quantile(y_preds_ensemble_t, q=1 - alpha / 2, dim=0)\ny_pred_80_t = torch.quantile(y_preds_ensemble_t, q=q_80, dim=0)\nlower_bounds_np = lower_bounds_t.cpu().numpy()\nupper_bounds_np = upper_bounds_t.cpu().numpy()\n\nemp_cov = empirical_coverage(val_values, lower_bounds_np, upper_bounds_np)\nwink_score = winkler_score(val_values, lower_bounds_np, upper_bounds_np, alpha)\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\n\n# Metrics\nprint(f\"{mae(y_true_t, y_pred_median_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_median_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_median_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80):.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 11. K2VAE","metadata":{}},{"cell_type":"code","source":"# Import dependencies and define the model\n!git clone https://github.com/decisionintelligence/K2VAE.git\nsys.path.append('../K2VAE')\n!pip install lightning reformer_pytorch linear_attention_transformer\nfrom probts.model.nn.prob.k2VAE.k2vae import k2VAE\nimport torch\nfrom einops import rearrange\nfrom sympy.stats.rv import sampling_E\nfrom probts.model.nn.prob.k2VAE.k2vae import k2VAE\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nimport torch\nfrom torch import nn\nfrom typing import List\n\nfrom probts.utils import weighted_average\nfrom probts.data.data_utils.data_scaler import TemporalScaler\nfrom typing import Union\n\nclass Forecaster(nn.Module):\n def __init__(\n self,\n target_dim: int,\n context_length: Union[list,int],\n prediction_length: Union[list,int],\n freq: str ,\n use_lags: bool = False,\n use_feat_idx_emb: bool = False,\n use_time_feat: bool = False,\n lags_list: List[int] = None,\n feat_idx_emb_dim: int = 1,\n time_feat_dim: int = 1,\n use_scaling: bool = False,\n autoregressive: bool = False,\n no_training: bool = False,\n dataset: str = None,\n **kwargs\n ):\n super().__init__()\n \n self.context_length = context_length\n self.prediction_length = prediction_length\n \n if isinstance(self.context_length, list):\n self.max_context_length = max(self.context_length)\n else:\n self.max_context_length = self.context_length\n \n if isinstance(self.prediction_length, list):\n self.max_prediction_length = max(self.prediction_length)\n else:\n self.max_prediction_length = self.prediction_length\n \n self.target_dim = target_dim\n self.freq = freq\n self.use_lags = use_lags\n self.use_feat_idx_emb = use_feat_idx_emb\n self.use_time_feat = use_time_feat\n self.feat_idx_emb_dim = feat_idx_emb_dim\n self.time_feat_dim = time_feat_dim\n self.autoregressive = autoregressive\n self.no_training = no_training\n self.use_scaling = use_scaling\n self.dataset = dataset\n # Lag parameters\n self.lags_list = lags_list\n if self.use_scaling:\n self.scaler = TemporalScaler()\n else:\n self.scaler = None\n\n if self.lags_list is not None:\n self.lags_dim = len(self.lags_list) * target_dim\n \n if use_feat_idx_emb:\n self.feat_idx_emb = nn.Embedding(\n num_embeddings=self.target_dim, embedding_dim=self.feat_idx_emb_dim\n )\n else:\n self.feat_idx_emb = None\n \n self.input_size = self.get_input_size()\n \n\n @property\n def name(self):\n return self.__class__.__name__\n\n def get_input_size(self):\n input_size = self.target_dim if not self.use_lags else self.lags_dim\n if self.use_feat_idx_emb:\n input_size += self.use_feat_idx_emb * self.target_dim\n if self.use_time_feat:\n input_size += self.time_feat_dim\n return input_size\n\n def get_lags(self, sequence, lags_list, lags_length=1):\n \"\"\"\n Get several lags from the sequence of shape (B, L, C) to (B, L', C*N),\n where L' = lag_length and N = len(lag_list).\n \"\"\"\n assert max(lags_list) + lags_length <= sequence.shape[1]\n\n lagged_values = []\n for lag_index in lags_list:\n begin_index = -lag_index - lags_length\n end_index = -lag_index if lag_index > 0 else None\n lagged_value = sequence[:, begin_index:end_index, ...]\n if self.use_scaling:\n lagged_value = lagged_value / self.scaler.scale\n lagged_values.append(lagged_value)\n return torch.cat(lagged_values, dim=-1)\n\n def get_input_sequence(\n self,\n past_target_cdf,\n future_target_cdf,\n mode\n ):\n if mode == 'all':\n sequence = torch.cat((past_target_cdf, future_target_cdf), dim=1)\n seq_length = self.max_context_length + self.max_prediction_length\n elif mode == 'encode':\n sequence = past_target_cdf\n seq_length = self.max_context_length\n elif mode == 'decode':\n sequence = past_target_cdf\n seq_length = 1\n else:\n raise ValueError(f\"Unsupported input mode: {mode}\")\n \n if self.use_lags:\n input_seq = self.get_lags(sequence, self.lags_list, seq_length)\n else: \n input_seq = sequence[:, -seq_length:, ...]\n if self.use_scaling:\n input_seq = input_seq / self.scaler.scale\n return input_seq\n \n def get_input_feat_idx_emb(self, target_dimension_indicator, input_length):\n input_feat_idx_emb = self.feat_idx_emb(target_dimension_indicator) # [B K D]\n\n input_feat_idx_emb = (\n input_feat_idx_emb.unsqueeze(1)\n .expand(-1, input_length, -1, -1)\n .reshape(-1, input_length, self.target_dim * self.feat_idx_emb_dim)\n )\n return input_feat_idx_emb # [B L K*D]\n\n def get_input_time_feat(\n self,\n past_time_feat,\n future_time_feat,\n mode\n ):\n if mode == 'all':\n time_feat = torch.cat(\n (past_time_feat[:, -self.max_context_length:, ...], future_time_feat), dim=1)\n elif mode == 'encode':\n time_feat = past_time_feat[:, -self.max_context_length:, ...]\n elif mode == 'decode':\n time_feat = future_time_feat\n return time_feat\n\n def get_inputs(self, batch_data, mode):\n inputs_list = []\n\n input_seq = self.get_input_sequence(\n batch_data.past_target_cdf, batch_data.future_target_cdf, mode=mode)\n input_length = input_seq.shape[1] # [B L n_lags*K]\n inputs_list.append(input_seq)\n\n if self.use_feat_idx_emb:\n input_feat_idx_emb = self.get_input_feat_idx_emb(\n batch_data.target_dimension_indicator, input_length) # [B L K*D]\n inputs_list.append(input_feat_idx_emb)\n\n if self.use_time_feat:\n input_time_feat = self.get_input_time_feat(\n batch_data.past_time_feat, batch_data.future_time_feat, mode=mode) # [B L Dt]\n inputs_list.append(input_time_feat)\n return torch.cat(inputs_list, dim=-1).to(dtype=torch.float32)\n \n def get_scale(self, batch_data):\n self.scaler.fit(\n batch_data.past_target_cdf[:, -self.max_context_length:, ...],\n batch_data.past_observed_values[:, -self.max_context_length:, ...]\n )\n \n def get_weighted_loss(self, batch_data, loss):\n observed_values = batch_data.future_observed_values\n loss_weights, _ = observed_values.min(dim=-1, keepdim=True)\n loss = weighted_average(loss, weights=loss_weights, dim=1)\n return loss\n \n def loss(self, batch_data):\n raise NotImplementedError\n \n def forecast(self, batch_data=None, num_samples=None):\n raise NotImplementedError\n\n\nclass ConvertedParams:\n def __init__(self, params):\n for key, value in params.items():\n setattr(self, key, value)\n\n\nclass k2VAEModel(Forecaster):\n def __init__(\n self,\n d_model,\n d_ff,\n e_layers,\n dropout,\n activation,\n n_heads,\n factor,\n patch_len,\n multistep,\n dynamic_dim,\n hidden_layers,\n hidden_dim,\n weight_beta=0.001,\n sample_schedule=5,\n init_kalman='identity',\n init_koopman='both',\n **kwargs\n ):\n \"\"\"\n Initialize the model with parameters.\n \"\"\"\n super().__init__(**kwargs)\n # Initialize model parameters here\n config = ConvertedParams(kwargs)\n config.d_model = d_model\n config.d_ff = d_ff\n config.hidden_layers = hidden_layers\n config.dropout = dropout\n config.activation = activation\n config.e_layers = e_layers\n config.n_heads = n_heads\n config.factor = factor\n config.patch_len = patch_len\n config.multistep = multistep\n config.dynamic_dim = dynamic_dim\n config.hidden_dim = hidden_dim\n config.n_vars = self.input_size\n config.seq_len = self.context_length\n config.pred_len = self.prediction_length\n self.weight_beta = weight_beta\n\n config.sample_schedule = sample_schedule\n config.init_kalman = init_kalman\n config.init_koopman = init_koopman\n\n self.model = k2VAE(config)\n\n def forward(self, input):\n \"\"\"\n Forward pass for the model.\n\n Parameters:\n inputs [Tensor]: Input tensor for the model.\n\n Returns:\n Tensor: Output tensor.\n \"\"\"\n # Perform the forward pass of the model\n rec, prior_dist, post_dist = self.model(input)\n return rec, prior_dist, post_dist\n\n def kld_loss(self, dist):\n # Extract the mean and covariance matrix from the distribution\n mu_q = dist.loc # Shape: (B, P, hidden)\n cov_q = dist.covariance_matrix # Shape: (B, P, hidden, hidden)\n # Compute the log determinant of the covariance matrix\n log_det_q = torch.linalg.slogdet(cov_q)[1] # Shape: (B, P)\n\n # Compute the trace of the covariance matrix (sum of diagonal elements)\n trace_q = torch.einsum('btii->bt', cov_q) # Shape: (B, P)\n\n # Compute the squared norm of the mean for each time step\n mu_term = torch.sum(mu_q ** 2, dim=-1) # Shape: (B, P)\n\n # Get the latent space dimension (256 in this case)\n latent_dim = mu_q.size(-1)\n\n # Compute the KL divergence for each time step\n kld = 0.5 * (trace_q + mu_term - latent_dim - log_det_q) # Shape: (B, P)\n\n # Take the mean over the time steps, resulting in one KL value per batch sample\n kld = kld.mean(dim=-1) # Shape: (B,)\n\n # Take the mean over the batch to get the final average KL divergence\n kld = kld.mean() # Scalar\n\n return kld\n\n def loss(self, batch_data, threshold=1e2):\n \"\"\"\n Compute the loss for the given batch data.\n\n Parameters:\n batch_data [dict]: Dictionary containing input data and possibly target data.\n\n Returns:\n Tensor: Computed loss.\n \"\"\"\n # Extract inputs and targets from batch_data\n self.model.train()\n input = batch_data.past_target_cdf[:, -self.context_length:, :]\n target = batch_data.future_target_cdf\n # Forward pass\n rec, prior_dist, post_dist = self.forward(input)\n rec_loss = F.mse_loss(rec, input) + F.mse_loss(post_dist.loc, target)\n post_loss = -post_dist.log_prob(target).mean()\n kld_loss = self.kld_loss(prior_dist)\n weight_alpha = 1 if post_loss < threshold else 0\n # stabilize training process\n loss = rec_loss + weight_alpha * post_loss + self.weight_beta * kld_loss\n return loss\n\n def sample_from_distribution(self, input, num_samples):\n samples = self.model.sample(input, num_samples)\n return rearrange(samples, 'n b l c -> b n l c')\n\n def forecast(self, batch_data, num_samples=None):\n \"\"\"\n Generate forecasts for the given batch data.\n\n Parameters:\n batch_data [dict]: Dictionary containing input data.\n num_samples [int, optional]: Number of samples per distribution during evaluation. Defaults to None.\n\n Returns:\n Tensor: Forecasted outputs.\n \"\"\"\n # Perform the forward pass to get the outputs\n self.model.eval()\n input = batch_data.past_target_cdf[:, -self.context_length:, :]\n with torch.no_grad():\n if num_samples is not None:\n # If num_samples is specified, use it to sample from the distribution\n outputs = self.sample_from_distribution(input, num_samples)\n else:\n outputs = None\n return outputs # [batch_size, num_samples, prediction_length, var_num]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TimeSeriesBatch:\n def __init__(self, past_target_cdf, future_target_cdf):\n self.past_target_cdf = past_target_cdf\n self.future_target_cdf = future_target_cdf\n\nclass SlidingWindowDataset(Dataset):\n def __init__(self, data, context_length, prediction_length):\n self.data = torch.tensor(data, dtype=torch.float32)\n self.context_length = context_length\n self.prediction_length = prediction_length\n self.total_len = len(self.data) - context_length - prediction_length + 1\n\n def __len__(self):\n return max(0, self.total_len)\n\n def __getitem__(self, idx):\n past = self.data[idx : idx + self.context_length]\n future = self.data[idx + self.context_length : idx + self.context_length + self.prediction_length]\n return past, future\n\ndef collate_fn(batch):\n past_list = [item[0] for item in batch]\n future_list = [item[1] for item in batch]\n past_tensor = torch.stack(past_list, dim=0) # [B, context_len, n_vars]\n future_tensor = torch.stack(future_list, dim=0) # [B, pred_len, n_vars]\n return TimeSeriesBatch(past_tensor, future_tensor)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndf = pd.read_csv(data_path)\ndf = df.fillna(0.0)\n\nNUM_NODES = df.shape[1]\nP_LAG = 60\nT_PRED = 30\nBATCH_SIZE = 64\nq_80 = 0.8\nalpha = 0.05\n\ntrain_values = df.values[:-T_PRED]\nval_values = df.values[-T_PRED:]\n\n# Setup datasets and dataloaders\ntrain_dataset = SlidingWindowDataset(train_values, context_length=P_LAG, prediction_length=T_PRED)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_fn)\n\n# Instantiate k2VAEModel\nmodel = k2VAEModel(\n target_dim=11,\n freq=\"D\",\n d_model=64,\n d_ff=128,\n e_layers=2,\n dropout=0.1,\n activation='gelu',\n n_heads=4,\n factor=1,\n patch_len=10, \n multistep=T_PRED,\n dynamic_dim=64,\n hidden_layers=2,\n hidden_dim=128,\n weight_beta=0.001,\n sample_schedule=5,\n init_kalman='identity',\n init_koopman='both',\n input_size=NUM_NODES,\n context_length=P_LAG, \n prediction_length=T_PRED\n).to(device)\n\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\nEPOCHS = 100\n\nmodel.train()\nstart_time = time.time()\nfor epoch in range(EPOCHS):\n total_loss = 0.0\n for batch in train_loader:\n # Move batch data to device\n batch.past_target_cdf = batch.past_target_cdf.to(device)\n batch.future_target_cdf = batch.future_target_cdf.to(device)\n \n optimizer.zero_grad()\n loss = model.loss(batch)\n loss.backward()\n optimizer.step()\n total_loss += loss.item()\n \n print(f\"Epoch {epoch+1}/{EPOCHS} | Train Loss: {total_loss / max(1, len(train_loader)):.4f}\")\ntraining_time = time.time() - start_time\n\n# Metrics\nmodel.eval()\nstart_time = time.time()\nval_tensor = torch.tensor(val_values, dtype=torch.float32).to(device)\ncontext_tensor = torch.tensor(train_values[-P_LAG:], dtype=torch.float32).unsqueeze(0).to(device)\n\neval_batch = TimeSeriesBatch(\n past_target_cdf=context_tensor,\n future_target_cdf=val_tensor.unsqueeze(0)\n)\n\nNUM_SAMPLES = 100\nwith torch.no_grad():\n # Shape: [batch_size=1, num_samples, prediction_length, var_num]\n forecast_samples = model.forecast(eval_batch, num_samples=NUM_SAMPLES)\n\ninference_time = time.time()-start_time\n\ny_preds_ensemble_t = forecast_samples.squeeze(0).cpu() # [num_samples, prediction_length, var_num]\ny_true_t = val_tensor.cpu() # [prediction_length, var_num]\ny_train_t = torch.tensor(train_values, dtype=torch.float32) # [train_len, var_num]\n\ny_pred_mean_t = y_preds_ensemble_t.mean(dim=0)\ny_pred_80_t = torch.quantile(y_preds_ensemble_t, q_80, dim=0)\ny_lower_t = torch.quantile(y_preds_ensemble_t, alpha / 2, dim=0)\ny_upper_t = torch.quantile(y_preds_ensemble_t, 1 - (alpha / 2), dim=0)\n\nemp_cov = empirical_coverage(y_true_t.numpy(), y_lower_t.numpy(), y_upper_t.numpy())\nwink_score = winkler_score(y_true_t.numpy(), y_lower_t.numpy(), y_upper_t.numpy(), alpha)\n\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80):.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 12. Transformer-MAF","metadata":{}},{"cell_type":"code","source":"# Some patching and workarounds required because pytorch-ts library is no longer updated\n\nfile_path = '../pytorch-ts/pts/__init__.py'\nwith open(file_path, 'r') as f:\n lines = f.readlines()\n\n# Write it back without the broken pkg_resources import\nwith open(file_path, 'w') as f:\n for line in lines:\n if 'pkg_resources' in line:\n continue\n if 'get_distribution' in line:\n line = line.replace('get_distribution(__name__).version', '\"0.6.0\"')\n f.write(line)\n\nfile_path = '../pytorch-ts/pts/model/transformer_tempflow/transformer_tempflow_estimator.py'\n\nwith open(file_path, 'r') as f:\n content = f.read()\n\n# Remove the invalid 'freq=self.freq,' argument\npatched_content = content.replace('freq=self.freq,', '')\n\n# Write the fixed code back\nwith open(file_path, 'w') as f:\n f.write(patched_content)\n\nprint(\"Patches successful!\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from pts.model.transformer_tempflow import TransformerTempFlowEstimator\nfrom pts import Trainer\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\ndf = pd.read_csv(data_path)\ndf = df.fillna(0.0)\n\nNUM_NODES = df.shape[1]\nP_LAG = 60\nT_PRED = 30\nBATCH_SIZE = 64\nq_80 = 0.8\nalpha = 0.05\n\ntrain_values = df.values[:-T_PRED]\nval_values = df.values[-T_PRED:]\ntrain_target = train_values.T.astype(np.float32)\ntrain_target += np.random.normal(0, 1e-6, train_target.shape).astype(np.float32)\n\ntest_target = df.values.T.astype(np.float32)\n\nfreq = \"D\" \nstart_date = pd.Period(\"2020-01-01\", freq=freq)\n\n\ntrain_ds = ListDataset(\n [{FieldName.TARGET: train_target, FieldName.START: start_date}] * 500,\n freq=freq,\n one_dim_target=False\n)\n\n# test_ds does not need to be multiplied, as the evaluator extracts deterministically\ntest_ds = ListDataset(\n [{FieldName.TARGET: test_target, FieldName.START: start_date}],\n freq=freq,\n one_dim_target=False\n)\n\n\nEXPECTED_INPUT_SIZE = 46 # Dataset-specific, needs to be adjusted. This is specific to Belgium data\n\nestimator = TransformerTempFlowEstimator(\n d_model=120,\n num_heads=4,\n input_size=EXPECTED_INPUT_SIZE,\n target_dim=NUM_NODES,\n prediction_length=T_PRED,\n context_length=P_LAG,\n flow_type='MAF',\n dequantize=False, \n freq=freq,\n scaling=True, \n trainer=Trainer(\n device=device,\n epochs=100,\n learning_rate=1e-3,\n num_batches_per_epoch=50,\n batch_size=BATCH_SIZE,\n )\n)\n\nprint(\"Training Transformer-MAF...\")\nstart_train = time.time()\n\npredictor = estimator.train(\n train_ds, \n num_workers=0, \n prefetch_factor=None\n)\n\ntraining_time = time.time() - start_train\n\nprint(\"Generating probabilistic forecasts...\")\nstart_inf = time.time()\n\nforecast_it, _ = make_evaluation_predictions(\n dataset=test_ds,\n predictor=predictor,\n num_samples=100\n)\nforecasts = list(forecast_it)\ninference_time = time.time() - start_inf\n\ny_preds_ensemble_np = forecasts[0].samples\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble_np, dtype=torch.float32).to(device)\n\ny_pred_median_t = torch.median(y_preds_ensemble_t, dim=0).values\nlower_bounds_t = torch.quantile(y_preds_ensemble_t, q=alpha / 2, dim=0)\nupper_bounds_t = torch.quantile(y_preds_ensemble_t, q=1 - alpha / 2, dim=0)\ny_pred_80_t = torch.quantile(y_preds_ensemble_t, q=q_80, dim=0)\n\nlower_bounds_np = lower_bounds_t.cpu().numpy()\nupper_bounds_np = upper_bounds_t.cpu().numpy()\ny_pred_median_np = y_pred_median_t.cpu().numpy()\n\nemp_cov = empirical_coverage(val_values, lower_bounds_np, upper_bounds_np)\nwink_score = winkler_score(val_values, lower_bounds_np, upper_bounds_np, alpha)\n\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\n\nprint(\"\\n--- Metrics ---\")\nprint(f\"{mae(y_true_t, y_pred_median_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_median_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_median_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80):.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 13. TimeGrad","metadata":{}},{"cell_type":"code","source":"from pts.model.time_grad import TimeGradEstimator\nfrom pts import Trainer\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndf = pd.read_csv(data_path)\ndf = df.fillna(0.0)\n\nNUM_NODES = df.shape[1]\nP_LAG = 60\nT_PRED = 30\nBATCH_SIZE = 64\nq_80 = 0.8\nalpha = 0.05\n\ntrain_values = df.values[:-T_PRED]\nval_values = df.values[-T_PRED:]\n\ntrain_target = train_values.T.astype(np.float32)\ntrain_target += np.random.normal(0, 1e-6, train_target.shape).astype(np.float32)\n\ntest_target = df.values.T.astype(np.float32)\n\nfreq = \"D\" \nstart_date = pd.Period(\"2020-01-01\", freq=freq)\n\ntrain_ds = ListDataset(\n [{FieldName.TARGET: train_target, FieldName.START: start_date}] * 500,\n freq=freq,\n one_dim_target=False\n)\n\ntest_ds = ListDataset(\n [{FieldName.TARGET: test_target, FieldName.START: start_date}],\n freq=freq,\n one_dim_target=False\n)\n\nNUM_NODES = df.shape[1]\n\nestimator = TimeGradEstimator(\n target_dim=NUM_NODES,\n prediction_length=T_PRED,\n context_length=P_LAG,\n cell_type='GRU',\n input_size=46, # Dataset-specific\n freq=freq,\n loss_type='l2',\n scaling=True,\n #diff_steps=100,\n #beta_end=0.1,\n beta_schedule=\"linear\",\n trainer=Trainer(\n device=device,\n epochs=100,\n learning_rate=1e-3,\n #grad_clip_norm=1.0,\n num_batches_per_epoch=30,\n batch_size=BATCH_SIZE,\n )\n)\n\nprint(\"Training TimeGrad on Belgium Data...\")\nstart_train = time.time()\npredictor = estimator.train(train_ds, num_workers=0, prefetch_factor=None)\ntraining_time = time.time() - start_train\n\nprint(\"Generating probabilistic forecasts...\")\nstart_inf = time.time()\n\nforecast_it, _ = make_evaluation_predictions(\n dataset=test_ds,\n predictor=predictor,\n num_samples=100\n)\nforecasts = list(forecast_it)\ninference_time = time.time() - start_inf\n\ny_preds_ensemble_np = forecasts[0].samples\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble_np, dtype=torch.float32).to(device)\n\ny_pred_median_t = torch.median(y_preds_ensemble_t, dim=0).values\n\nlower_bounds_t = torch.quantile(y_preds_ensemble_t, q=alpha / 2, dim=0)\nupper_bounds_t = torch.quantile(y_preds_ensemble_t, q=1 - alpha / 2, dim=0)\ny_pred_80_t = torch.quantile(y_preds_ensemble_t, q=q_80, dim=0)\n\nlower_bounds_np = lower_bounds_t.cpu().numpy()\nupper_bounds_np = upper_bounds_t.cpu().numpy()\ny_pred_median_np = y_pred_median_t.cpu().numpy()\n\nemp_cov = empirical_coverage(val_values, lower_bounds_np, upper_bounds_np)\nwink_score = winkler_score(val_values, lower_bounds_np, upper_bounds_np, alpha)\n\ny_true_t = torch.tensor(val_values, dtype=torch.float32).to(device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32).to(device)\n\nprint(\"\\n--- Metrics ---\")\nprint(f\"{mae(y_true_t, y_pred_median_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_median_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_median_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80):.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Note: If the above cell throws an error, run this cell, restart session, then re-run\n!sed -i '/freq=self.freq/d' ../pytorch-ts/pts/model/time_grad/time_grad_estimator.py","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 14. TimesFM-2.5","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\n\nhorizon = 30\nnum_nodes = df.shape[1]\nq_80 = 0.8\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Extract arrays for evaluation (shape: time, num_nodes)\ntrain_values = df.iloc[:-horizon].values\nval_values = df.iloc[-horizon:].values\n\n# Convert to Darts TimeSeries\ntrain_ts = TimeSeries.from_values(train_values)\n\n# Quantiles should be subset of the ones with which TimesFM was trained in darts\nquantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]\nmodel = TimesFM2p5Model(input_chunk_length=60, output_chunk_length=30,\n likelihood=QuantileRegression(quantiles=quantiles))\n\nt0 = time.time()\nmodel.fit(train_ts)\ntraining_time = time.time() - t0\nnum_samples = 100\nt0 = time.time()\npred_ts = model.predict(n=horizon, series=train_ts, num_samples=num_samples)\ninference_time = time.time() - t0\npred_samples = pred_ts.all_values() \n\n# Rearrange to (samples, time, components) for easier quantile calculations\ny_preds_ensemble_np = np.transpose(pred_samples, (2, 0, 1))\n\n# Compute point predictions (mean) and quantiles\ny_pred_np = np.mean(y_preds_ensemble_np, axis=0) # Mean as point forecast\nlower_bounds_np = np.quantile(y_preds_ensemble_np, 0.025, axis=0) \nupper_bounds_np = np.quantile(y_preds_ensemble_np, 0.975, axis=0) \ny_pred_80_np = np.quantile(y_preds_ensemble_np, q_80, axis=0) # 80th percentile\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32, device=device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32, device=device)\ny_pred_mean_t = torch.tensor(y_pred_np, dtype=torch.float32, device=device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble_np, dtype=torch.float32, device=device)\ny_pred_80_t = torch.tensor(y_pred_80_np, dtype=torch.float32, device=device)\nalpha = 0.05 \nemp_cov = empirical_coverage(val_values, lower_bounds_np, upper_bounds_np)\nwink_score = winkler_score(val_values, lower_bounds_np, upper_bounds_np, alpha)\n\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 15. Chronos-2","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\n\nhorizon = 30\nnum_nodes = df.shape[1]\nq_80 = 0.8\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Extract arrays for evaluation (shape: time, num_nodes)\ntrain_values = df.iloc[:-horizon].values\nval_values = df.iloc[-horizon:].values\n\n# Convert to Darts TimeSeries\ntrain_ts = TimeSeries.from_values(train_values)\n\n# Quantiles should be subset of the ones with which TimesFM was trained in darts\nquantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]\nmodel = Chronos2Model(input_chunk_length=60, output_chunk_length=30,\n likelihood=QuantileRegression(quantiles=quantiles))\n\nt0 = time.time()\nmodel.fit(train_ts)\ntraining_time = time.time() - t0\nnum_samples = 100\nt0 = time.time()\npred_ts = model.predict(n=horizon, series=train_ts, num_samples=num_samples)\ninference_time = time.time() - t0\npred_samples = pred_ts.all_values() \n\n# Rearrange to (samples, time, components) for easier quantile calculations\ny_preds_ensemble_np = np.transpose(pred_samples, (2, 0, 1))\n\n# Compute point predictions (mean) and quantiles\ny_pred_np = np.mean(y_preds_ensemble_np, axis=0) # Mean as point forecast\nlower_bounds_np = np.quantile(y_preds_ensemble_np, 0.025, axis=0) \nupper_bounds_np = np.quantile(y_preds_ensemble_np, 0.975, axis=0) \ny_pred_80_np = np.quantile(y_preds_ensemble_np, q_80, axis=0) # 80th percentile\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32, device=device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32, device=device)\ny_pred_mean_t = torch.tensor(y_pred_np, dtype=torch.float32, device=device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble_np, dtype=torch.float32, device=device)\ny_pred_80_t = torch.tensor(y_pred_80_np, dtype=torch.float32, device=device)\nalpha = 0.05 \nemp_cov = empirical_coverage(val_values, lower_bounds_np, upper_bounds_np)\nwink_score = winkler_score(val_values, lower_bounds_np, upper_bounds_np, alpha)\n\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 16. PatchTSTFM","metadata":{}},{"cell_type":"code","source":"df = pd.read_csv(data_path)\n\nhorizon = 30\nnum_nodes = df.shape[1]\nq_80 = 0.8\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Extract arrays for evaluation (shape: time, num_nodes)\ntrain_values = df.iloc[:-horizon].values\nval_values = df.iloc[-horizon:].values\n\n# Convert to Darts TimeSeries\ntrain_ts = TimeSeries.from_values(train_values)\n\n# Quantiles should be subset of the ones with which TimesFM was trained in darts\nquantiles = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]\nmodel = PatchTSTFMModel(input_chunk_length=60, output_chunk_length=30,\n likelihood=QuantileRegression(quantiles=quantiles))\n\nt0 = time.time()\nmodel.fit(train_ts)\ntraining_time = time.time() - t0\nnum_samples = 100\nt0 = time.time()\npred_ts = model.predict(n=horizon, series=train_ts, num_samples=num_samples)\ninference_time = time.time() - t0\npred_samples = pred_ts.all_values() \n\n# Rearrange to (samples, time, components) for easier quantile calculations\ny_preds_ensemble_np = np.transpose(pred_samples, (2, 0, 1))\n\n# Compute point predictions (mean) and quantiles\ny_pred_np = np.mean(y_preds_ensemble_np, axis=0) # Mean as point forecast\nlower_bounds_np = np.quantile(y_preds_ensemble_np, 0.025, axis=0) \nupper_bounds_np = np.quantile(y_preds_ensemble_np, 0.975, axis=0) \ny_pred_80_np = np.quantile(y_preds_ensemble_np, q_80, axis=0) # 80th percentile\n\n# Metrics\ny_true_t = torch.tensor(val_values, dtype=torch.float32, device=device)\ny_train_t = torch.tensor(train_values, dtype=torch.float32, device=device)\ny_pred_mean_t = torch.tensor(y_pred_np, dtype=torch.float32, device=device)\ny_preds_ensemble_t = torch.tensor(y_preds_ensemble_np, dtype=torch.float32, device=device)\ny_pred_80_t = torch.tensor(y_pred_80_np, dtype=torch.float32, device=device)\nalpha = 0.05 \nemp_cov = empirical_coverage(val_values, lower_bounds_np, upper_bounds_np)\nwink_score = winkler_score(val_values, lower_bounds_np, upper_bounds_np, alpha)\n\nprint(f\"{mae(y_true_t, y_pred_mean_t):.4f}\")\nprint(f\"{mase(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{rmsse(y_true_t, y_pred_mean_t, y_train_t):.4f}\")\nprint(f\"{emp_cov:.4f}\")\nprint(f\"{crps_approximation(y_true_t, y_preds_ensemble_t):.4f}\")\nprint(f\"{wink_score:.4f}\")\nprint(f\"{pinball_loss(y_true_t, y_pred_80_t, quantile=q_80).item():.4f}\")\nprint(f\"{training_time:.4f}\")\nprint(f\"{inference_time:.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}