589 lines
22 KiB
Python
589 lines
22 KiB
Python
import torch
|
|
from tqdm import tqdm
|
|
from src.losses.crps_metric import crps_from_samples
|
|
from src.trainers.trainer import Trainer
|
|
from src.trainers.autoregressive_trainer import AutoRegressiveTrainer
|
|
from src.data.preprocessing import DataProcessor
|
|
from src.utils.clearml import ClearMLHelper
|
|
from src.losses import PinballLoss, NonAutoRegressivePinballLoss, CRPSLoss
|
|
import plotly.graph_objects as go
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from scipy.interpolate import CubicSpline
|
|
import matplotlib.pyplot as plt
|
|
import seaborn as sns
|
|
import matplotlib.patches as mpatches
|
|
|
|
def sample_from_dist(quantiles, preds):
|
|
if isinstance(preds, torch.Tensor):
|
|
preds = preds.detach().cpu()
|
|
|
|
# if preds more than 2 dimensions, flatten to 2
|
|
if len(preds.shape) > 2:
|
|
preds = preds.reshape(-1, preds.shape[-1])
|
|
# target will be reshaped from (1024, 96, 15) to (1024*96, 15)
|
|
# our target (1024, 96) also needs to be reshaped to (1024*96, 1)
|
|
target = target.reshape(-1, 1)
|
|
|
|
# preds and target as numpy
|
|
preds = preds.numpy()
|
|
|
|
# random probabilities of (1000, 1)
|
|
import random
|
|
probs = np.array([random.random() for _ in range(1000)])
|
|
|
|
spline = CubicSpline(quantiles, preds, axis=1)
|
|
|
|
samples = spline(probs)
|
|
|
|
# get the diagonal
|
|
samples = np.diag(samples)
|
|
|
|
return samples
|
|
|
|
def auto_regressive(dataset, model, quantiles, idx_batch, sequence_length: int = 96):
|
|
device = next(model.parameters()).device
|
|
prev_features, targets = dataset.get_batch(idx_batch)
|
|
prev_features = prev_features.to(device)
|
|
targets = targets.to(device)
|
|
|
|
if len(list(prev_features.shape)) == 2:
|
|
initial_sequence = prev_features[:, :96]
|
|
else:
|
|
initial_sequence = prev_features[:, :, 0]
|
|
|
|
target_full = targets[:, 0].unsqueeze(1) # (batch_size, 1)
|
|
with torch.no_grad():
|
|
new_predictions_full = model(prev_features) # (batch_size, quantiles)
|
|
samples = (
|
|
torch.tensor(sample_from_dist(quantiles, new_predictions_full))
|
|
.unsqueeze(1)
|
|
.to(device)
|
|
) # (batch_size, 1)
|
|
predictions_samples = samples
|
|
predictions_full = new_predictions_full.unsqueeze(1)
|
|
|
|
for i in range(sequence_length - 1):
|
|
if len(list(prev_features.shape)) == 2:
|
|
new_features = torch.cat(
|
|
(prev_features[:, 1:96], samples), dim=1
|
|
) # (batch_size, 96)
|
|
|
|
new_features = new_features.float()
|
|
|
|
other_features, new_targets = dataset.get_batch_autoregressive(
|
|
np.array(idx_batch) + i + 1
|
|
) # (batch_size, new_features)
|
|
|
|
if other_features is not None:
|
|
prev_features = torch.cat(
|
|
(new_features.to(device), other_features.to(device)), dim=1
|
|
) # (batch_size, 96 + new_features)
|
|
else:
|
|
prev_features = new_features
|
|
|
|
else:
|
|
other_features, new_targets = dataset.get_batch_autoregressive(
|
|
np.array(idx_batch) + i + 1
|
|
) # (batch_size, 1, new_features)
|
|
|
|
# change the other_features nrv based on the samples
|
|
other_features[:, 0, 0] = samples.squeeze(-1)
|
|
# make sure on same device
|
|
other_features = other_features.to(device)
|
|
prev_features = prev_features.to(device)
|
|
prev_features = torch.cat(
|
|
(prev_features[:, 1:, :], other_features), dim=1
|
|
) # (batch_size, 96, new_features)
|
|
|
|
target_full = torch.cat(
|
|
(target_full, new_targets.to(device)), dim=1
|
|
) # (batch_size, sequence_length)
|
|
|
|
with torch.no_grad():
|
|
new_predictions_full = model(
|
|
prev_features
|
|
) # (batch_size, quantiles)
|
|
predictions_full = torch.cat(
|
|
(predictions_full, new_predictions_full.unsqueeze(1)), dim=1
|
|
) # (batch_size, sequence_length, quantiles)
|
|
|
|
samples = (
|
|
torch.tensor(sample_from_dist(quantiles, new_predictions_full))
|
|
.unsqueeze(-1)
|
|
.to(device)
|
|
) # (batch_size, 1)
|
|
predictions_samples = torch.cat((predictions_samples, samples), dim=1)
|
|
|
|
return (
|
|
initial_sequence,
|
|
predictions_full,
|
|
predictions_samples,
|
|
target_full.unsqueeze(-1),
|
|
)
|
|
|
|
class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
|
|
def __init__(
|
|
self,
|
|
model: torch.nn.Module,
|
|
input_dim: tuple,
|
|
optimizer: torch.optim.Optimizer,
|
|
data_processor: DataProcessor,
|
|
quantiles: list,
|
|
device: torch.device,
|
|
debug: bool = True,
|
|
):
|
|
|
|
self.quantiles = quantiles
|
|
|
|
criterion = PinballLoss(quantiles=quantiles)
|
|
super().__init__(
|
|
model=model,
|
|
input_dim=input_dim,
|
|
optimizer=optimizer,
|
|
criterion=criterion,
|
|
data_processor=data_processor,
|
|
device=device,
|
|
debug=debug,
|
|
)
|
|
|
|
def calculate_crps_from_samples(self, task, dataloader, epoch: int):
|
|
crps_from_samples_metric = []
|
|
|
|
with torch.no_grad():
|
|
total_samples = len(dataloader.dataset) - 96
|
|
for _, _, idx_batch in tqdm(dataloader):
|
|
idx_batch = [idx for idx in idx_batch if idx < total_samples]
|
|
|
|
if len(idx_batch) == 0:
|
|
continue
|
|
|
|
for idx in tqdm(idx_batch):
|
|
computed_idx_batch = [idx] * 100
|
|
_, _, samples, targets = self.auto_regressive(
|
|
dataloader.dataset, idx_batch=computed_idx_batch
|
|
)
|
|
samples = samples.unsqueeze(0)
|
|
targets = targets.squeeze(-1)
|
|
targets = targets[0].unsqueeze(0)
|
|
|
|
crps = crps_from_samples(samples, targets)
|
|
|
|
crps_from_samples_metric.append(crps[0].mean().item())
|
|
|
|
task.get_logger().report_scalar(
|
|
title="CRPS_from_samples", series="test", value=np.mean(crps_from_samples_metric), iteration=epoch
|
|
)
|
|
|
|
def log_final_metrics(self, task, dataloader, train: bool = True):
|
|
metrics = {metric.__class__.__name__: 0.0 for metric in self.metrics_to_track}
|
|
transformed_metrics = {
|
|
metric.__class__.__name__: 0.0 for metric in self.metrics_to_track
|
|
}
|
|
|
|
crps_from_samples_metric = []
|
|
|
|
with torch.no_grad():
|
|
total_samples = len(dataloader.dataset) - 96
|
|
batches = 0
|
|
for _, _, idx_batch in tqdm(dataloader):
|
|
idx_batch = [idx for idx in idx_batch if idx < total_samples]
|
|
|
|
if len(idx_batch) == 0:
|
|
continue
|
|
|
|
if train == False:
|
|
for idx in tqdm(idx_batch):
|
|
computed_idx_batch = [idx] * 100
|
|
_, outputs, samples, targets = self.auto_regressive(
|
|
dataloader.dataset, idx_batch=computed_idx_batch
|
|
)
|
|
samples = samples.unsqueeze(0)
|
|
targets = targets.squeeze(-1)
|
|
targets = targets[0].unsqueeze(0)
|
|
|
|
crps = crps_from_samples(samples, targets)
|
|
|
|
crps_from_samples_metric.append(crps[0].mean().item())
|
|
|
|
|
|
_, outputs, samples, targets = self.auto_regressive(
|
|
dataloader.dataset, idx_batch=idx_batch
|
|
)
|
|
|
|
samples = samples.to(self.device)
|
|
outputs = outputs.to(self.device)
|
|
targets = targets.to(self.device)
|
|
|
|
inversed_samples = self.data_processor.inverse_transform(samples)
|
|
inversed_targets = self.data_processor.inverse_transform(targets)
|
|
inversed_outputs = self.data_processor.inverse_transform(outputs)
|
|
|
|
inversed_samples = inversed_samples.to(self.device)
|
|
inversed_targets = inversed_targets.to(self.device)
|
|
inversed_outputs = inversed_outputs.to(self.device)
|
|
|
|
for metric in self.metrics_to_track:
|
|
if metric.__class__ != PinballLoss and metric.__class__ != CRPSLoss:
|
|
transformed_metrics[metric.__class__.__name__] += metric(
|
|
samples, targets.squeeze(-1)
|
|
)
|
|
metrics[metric.__class__.__name__] += metric(
|
|
inversed_samples, inversed_targets.squeeze(-1)
|
|
)
|
|
else:
|
|
transformed_metrics[metric.__class__.__name__] += metric(
|
|
outputs, targets
|
|
)
|
|
metrics[metric.__class__.__name__] += metric(
|
|
inversed_outputs, inversed_targets
|
|
)
|
|
batches += 1
|
|
|
|
for metric in self.metrics_to_track:
|
|
metrics[metric.__class__.__name__] /= batches
|
|
transformed_metrics[metric.__class__.__name__] /= batches
|
|
|
|
for metric_name, metric_value in metrics.items():
|
|
if PinballLoss.__name__ in metric_name:
|
|
continue
|
|
name = f"train_{metric_name}" if train else f"test_{metric_name}"
|
|
task.get_logger().report_single_value(name=name, value=metric_value)
|
|
|
|
for metric_name, metric_value in transformed_metrics.items():
|
|
name = (
|
|
f"train_transformed_{metric_name}"
|
|
if train
|
|
else f"test_transformed_{metric_name}"
|
|
)
|
|
task.get_logger().report_single_value(name=name, value=metric_value)
|
|
|
|
if train == False:
|
|
task.get_logger().report_single_value(
|
|
name="test_CRPS_from_samples_transformed", value=np.mean(crps_from_samples_metric)
|
|
)
|
|
|
|
# def get_plot_error(
|
|
# self,
|
|
# next_day,
|
|
# predictions,
|
|
# ):
|
|
# metric = PinballLoss(quantiles=self.quantiles)
|
|
# fig = go.Figure()
|
|
|
|
# next_day_np = next_day.view(-1).cpu().numpy()
|
|
# predictions_np = predictions.cpu().numpy()
|
|
|
|
# if True:
|
|
# next_day_np = self.data_processor.inverse_transform(next_day_np)
|
|
# predictions_np = self.data_processor.inverse_transform(predictions_np)
|
|
|
|
# # for each time step, calculate the error using the metric
|
|
# errors = []
|
|
# for i in range(96):
|
|
|
|
# target_tensor = torch.tensor(next_day_np[i]).unsqueeze(0)
|
|
# prediction_tensor = torch.tensor(predictions_np[i]).unsqueeze(0)
|
|
|
|
# errors.append(metric(prediction_tensor, target_tensor))
|
|
|
|
# # plot the error
|
|
# fig.add_trace(go.Scatter(x=np.arange(96), y=errors, name=metric.__class__.__name__))
|
|
# fig.update_layout(title=f"Error of {metric.__class__.__name__} for each time step")
|
|
|
|
# return fig
|
|
|
|
|
|
def get_plot(
|
|
self,
|
|
current_day,
|
|
next_day,
|
|
predictions,
|
|
show_legend: bool = True,
|
|
retransform: bool = True,
|
|
):
|
|
fig = go.Figure()
|
|
|
|
# Convert to numpy for plotting
|
|
current_day_np = current_day.view(-1).cpu().numpy()
|
|
next_day_np = next_day.view(-1).cpu().numpy()
|
|
predictions_np = predictions.cpu().numpy()
|
|
|
|
if retransform:
|
|
current_day_np = self.data_processor.inverse_transform(current_day_np)
|
|
next_day_np = self.data_processor.inverse_transform(next_day_np)
|
|
predictions_np = self.data_processor.inverse_transform(predictions_np)
|
|
|
|
ci_99_upper = np.quantile(predictions_np, 0.995, axis=0)
|
|
ci_99_lower = np.quantile(predictions_np, 0.005, axis=0)
|
|
|
|
ci_95_upper = np.quantile(predictions_np, 0.975, axis=0)
|
|
ci_95_lower = np.quantile(predictions_np, 0.025, axis=0)
|
|
|
|
ci_90_upper = np.quantile(predictions_np, 0.95, axis=0)
|
|
ci_90_lower = np.quantile(predictions_np, 0.05, axis=0)
|
|
|
|
ci_50_lower = np.quantile(predictions_np, 0.25, axis=0)
|
|
ci_50_upper = np.quantile(predictions_np, 0.75, axis=0)
|
|
|
|
# Add traces for current and next day
|
|
# fig.add_trace(go.Scatter(x=np.arange(96), y=current_day_np, name="Current Day"))
|
|
# fig.add_trace(go.Scatter(x=96 + np.arange(96), y=next_day_np, name="Next Day"))
|
|
|
|
# for i, q in enumerate(self.quantiles):
|
|
# fig.add_trace(
|
|
# go.Scatter(
|
|
# x=96 + np.arange(96),
|
|
# y=predictions_np[:, i],
|
|
# name=f"Prediction (Q={q})",
|
|
# line=dict(dash="dash"),
|
|
# )
|
|
# )
|
|
|
|
# # Update the layout
|
|
# fig.update_layout(
|
|
# title="Predictions and Quantiles of the Linear Model",
|
|
# showlegend=show_legend,
|
|
# )
|
|
|
|
sns.set_theme()
|
|
time_steps = np.arange(0, 96)
|
|
|
|
fig, ax = plt.subplots(figsize=(20, 10))
|
|
ax.plot(time_steps, predictions_np.mean(axis=0), label="Mean of NRV samples", linewidth=3)
|
|
# ax.fill_between(time_steps, ci_lower, ci_upper, color='b', alpha=0.2, label='Full Interval')
|
|
|
|
ax.fill_between(time_steps, ci_99_lower, ci_99_upper, color='b', alpha=0.2, label='99% Interval')
|
|
ax.fill_between(time_steps, ci_95_lower, ci_95_upper, color='b', alpha=0.2, label='95% Interval')
|
|
ax.fill_between(time_steps, ci_90_lower, ci_90_upper, color='b', alpha=0.2, label='90% Interval')
|
|
ax.fill_between(time_steps, ci_50_lower, ci_50_upper, color='b', alpha=0.2, label='50% Interval')
|
|
|
|
ax.plot(next_day_np, label="Real NRV", linewidth=3)
|
|
# full_interval_patch = mpatches.Patch(color='b', alpha=0.2, label='Full Interval')
|
|
ci_99_patch = mpatches.Patch(color='b', alpha=0.3, label='99% Interval')
|
|
ci_95_patch = mpatches.Patch(color='b', alpha=0.4, label='95% Interval')
|
|
ci_90_patch = mpatches.Patch(color='b', alpha=0.5, label='90% Interval')
|
|
ci_50_patch = mpatches.Patch(color='b', alpha=0.6, label='50% Interval')
|
|
|
|
|
|
ax.legend(handles=[ci_99_patch, ci_95_patch, ci_90_patch, ci_50_patch, ax.lines[0], ax.lines[1]])
|
|
return fig
|
|
|
|
def auto_regressive(self, dataset, idx_batch, sequence_length: int = 96):
|
|
return auto_regressive(dataset, self.model, self.quantiles, idx_batch, sequence_length)
|
|
|
|
def plot_quantile_percentages(
|
|
self, task, data_loader, train: bool = True, iteration: int = None, full_day: bool = False
|
|
):
|
|
quantiles = self.quantiles
|
|
total = 0
|
|
quantile_counter = {q: 0 for q in quantiles}
|
|
|
|
self.model.eval()
|
|
with torch.no_grad():
|
|
total_samples = len(data_loader.dataset) - 96
|
|
|
|
for inputs, targets, idx_batch in data_loader:
|
|
idx_batch = [idx for idx in idx_batch if idx < total_samples]
|
|
|
|
if full_day:
|
|
_, outputs, samples, targets = self.auto_regressive(
|
|
data_loader.dataset, idx_batch=idx_batch
|
|
)
|
|
# outputs: (batch, sequence_length, num_quantiles)
|
|
# targets: (batch, sequence_length, 1)
|
|
|
|
# reshape to (batch_size * sequence_length, num_quantiles)
|
|
outputs = outputs.reshape(-1, len(quantiles))
|
|
targets = targets.reshape(-1)
|
|
|
|
# to cpu
|
|
outputs = outputs.cpu().numpy()
|
|
targets = targets.cpu().numpy()
|
|
|
|
else:
|
|
inputs = inputs.to(self.device)
|
|
outputs = self.model(inputs).cpu().numpy() # (batch_size, num_quantiles)
|
|
targets = targets.squeeze(-1).cpu().numpy() # (batch_size, 1)
|
|
|
|
for i, q in enumerate(quantiles):
|
|
quantile_counter[q] += np.sum(
|
|
targets < outputs[:, i]
|
|
)
|
|
|
|
total += len(targets)
|
|
|
|
# to numpy array of length len(quantiles)
|
|
percentages = np.array(
|
|
[quantile_counter[q] / total for q in quantiles]
|
|
)
|
|
|
|
bar_width = 0.35
|
|
index = np.arange(len(quantiles))
|
|
|
|
# Plotting the bars
|
|
fig, ax = plt.subplots(figsize=(15, 10))
|
|
|
|
bar1 = ax.bar(
|
|
index, quantiles, bar_width, label="Ideal", color="brown"
|
|
)
|
|
bar2 = ax.bar(
|
|
index + bar_width, percentages, bar_width, label="NN model", color="blue"
|
|
)
|
|
|
|
# Adding the percentage values above the bars for bar2
|
|
for rect in bar2:
|
|
height = rect.get_height()
|
|
ax.text(
|
|
rect.get_x() + rect.get_width() / 2.0,
|
|
1.005 * height,
|
|
f"{height:.2}",
|
|
ha="center",
|
|
va="bottom",
|
|
) # Format the number as a percentage
|
|
|
|
series_name = "Training Set" if train else "Test Set"
|
|
full_day_str = "Full Day" if full_day else "Single Step"
|
|
|
|
# Adding labels and title
|
|
ax.set_xlabel("Quantile")
|
|
ax.set_ylabel("Fraction of data under quantile forecast")
|
|
ax.set_title(f"{series_name} {full_day_str} Quantile Performance Comparison")
|
|
ax.set_xticks(index + bar_width / 2)
|
|
ax.set_xticklabels(quantiles)
|
|
ax.legend()
|
|
|
|
task.get_logger().report_matplotlib_figure(
|
|
title="Quantile Performance Comparison",
|
|
series=f"{series_name} {full_day_str}",
|
|
report_image=True,
|
|
figure=plt,
|
|
iteration=iteration,
|
|
)
|
|
plt.close()
|
|
|
|
|
|
class NonAutoRegressiveQuantileRegression(Trainer):
|
|
def __init__(
|
|
self,
|
|
model: torch.nn.Module,
|
|
input_dim: tuple,
|
|
optimizer: torch.optim.Optimizer,
|
|
data_processor: DataProcessor,
|
|
quantiles: list,
|
|
device: torch.device,
|
|
debug: bool = True,
|
|
):
|
|
self.quantiles = quantiles
|
|
|
|
|
|
criterion = NonAutoRegressivePinballLoss(quantiles=quantiles)
|
|
super().__init__(
|
|
model=model,
|
|
input_dim=input_dim,
|
|
optimizer=optimizer,
|
|
criterion=criterion,
|
|
data_processor=data_processor,
|
|
device=device,
|
|
debug=debug,
|
|
)
|
|
|
|
def log_final_metrics(self, task, dataloader, train: bool = True):
|
|
metrics = {metric.__class__.__name__: 0.0 for metric in self.metrics_to_track}
|
|
transformed_metrics = {
|
|
metric.__class__.__name__: 0.0 for metric in self.metrics_to_track
|
|
}
|
|
|
|
with torch.no_grad():
|
|
for inputs, targets, _ in dataloader:
|
|
inputs, targets = inputs.to(self.device), targets.to(self.device)
|
|
|
|
outputs = self.model(inputs)
|
|
outputted_samples = [
|
|
sample_from_dist(self.quantiles, output.cpu().numpy())
|
|
for output in outputs
|
|
]
|
|
|
|
outputted_samples = torch.tensor(outputted_samples)
|
|
inversed_outputs_samples = self.data_processor.inverse_transform(
|
|
outputted_samples
|
|
)
|
|
|
|
outputs = outputs.reshape(inputs.shape[0], -1, len(self.quantiles))
|
|
inversed_outputs = self.data_processor.inverse_transform(outputs)
|
|
inversed_targets = self.data_processor.inverse_transform(targets)
|
|
|
|
inversed_outputs_samples = inversed_outputs_samples.to(self.device)
|
|
inversed_targets = inversed_targets.to(self.device)
|
|
outputted_samples = outputted_samples.to(self.device)
|
|
inversed_outputs = inversed_outputs.to(self.device)
|
|
|
|
for metric in self.metrics_to_track:
|
|
if metric.__class__ != PinballLoss and metric.__class__ != CRPSLoss:
|
|
transformed_metrics[metric.__class__.__name__] += metric(
|
|
outputted_samples, targets
|
|
)
|
|
metrics[metric.__class__.__name__] += metric(
|
|
inversed_outputs_samples, inversed_targets
|
|
)
|
|
else:
|
|
transformed_metrics[metric.__class__.__name__] += metric(
|
|
outputs, targets.unsqueeze(-1)
|
|
)
|
|
metrics[metric.__class__.__name__] += metric(
|
|
inversed_outputs, inversed_targets.unsqueeze(-1)
|
|
)
|
|
|
|
for metric in self.metrics_to_track:
|
|
metrics[metric.__class__.__name__] /= len(dataloader)
|
|
transformed_metrics[metric.__class__.__name__] /= len(dataloader)
|
|
|
|
for metric_name, metric_value in metrics.items():
|
|
if train:
|
|
metric_name = f"train_{metric_name}"
|
|
else:
|
|
metric_name = f"test_{metric_name}"
|
|
|
|
task.get_logger().report_single_value(
|
|
name=metric_name, value=metric_value
|
|
)
|
|
|
|
for metric_name, metric_value in transformed_metrics.items():
|
|
if train:
|
|
metric_name = f"train_transformed_{metric_name}"
|
|
else:
|
|
metric_name = f"test_transformed_{metric_name}"
|
|
|
|
task.get_logger().report_single_value(
|
|
name=metric_name, value=metric_value
|
|
)
|
|
|
|
def get_plot(self, current_day, next_day, predictions, show_legend: bool = True):
|
|
fig = go.Figure()
|
|
|
|
# Convert to numpy for plotting
|
|
current_day_np = current_day.view(-1).cpu().numpy()
|
|
next_day_np = next_day.view(-1).cpu().numpy()
|
|
|
|
# reshape predictions to (n, len(quantiles))$
|
|
predictions_np = predictions.cpu().numpy().reshape(-1, len(self.quantiles))
|
|
|
|
# Add traces for current and next day
|
|
fig.add_trace(go.Scatter(x=np.arange(96), y=current_day_np, name="Current Day"))
|
|
fig.add_trace(go.Scatter(x=96 + np.arange(96), y=next_day_np, name="Next Day"))
|
|
|
|
for i, q in enumerate(self.quantiles):
|
|
fig.add_trace(
|
|
go.Scatter(
|
|
x=96 + np.arange(96),
|
|
y=predictions_np[:, i],
|
|
name=f"Prediction (Q={q})",
|
|
line=dict(dash="dash"),
|
|
)
|
|
)
|
|
|
|
# Update the layout
|
|
fig.update_layout(title="Predictions and Quantiles", showlegend=show_legend)
|
|
|
|
return fig
|