utils_lib.py

import matplotlib.pyplot as plt  # requires: pip install matplotlib
import numpy as np
import pdb 
import os

def plot_pred(org_data, median, _dir='None', gt=None, low=None, high=None, forecast_index=None, title=None):
    # print("Visualizing prediction data...")
    if forecast_index is None:
        forecast_index = range(len(org_data), len(org_data) + len(median))
    plt.figure(figsize=(8, 4))
    plt.plot(org_data, color="royalblue", label="historical data")
    plt.plot(forecast_index, median, color="tomato", label="median forecast")
    if low is not None and high is not None:
        plt.fill_between(forecast_index, low, high, color="tomato", alpha=0.3, label="80% prediction interval")
    if gt is not None:
        plt.plot(range(len(org_data), len(org_data) + len(gt)), gt, color="green", label="gt")
    plt.legend()
    plt.grid()
    directory = f"./{_dir}"
    if not os.path.exists(directory):
        os.makedirs(directory, exist_ok=True)
    if title is not None:
        plt.savefig(f'{directory}/{title}.jpg')
    plt.close()