utils.py

import json
import logging
import os
import shutil

import torch
import numpy as np
from tqdm import tqdm
import CRPS.CRPS as pscore

import matplotlib
matplotlib.use('Agg')
#matplotlib.rcParams['savefig.dpi'] = 300 #Uncomment for higher plot resolutions
import matplotlib.pyplot as plt

import model.net as net
# import net as net
logger = logging.getLogger('DeepAR.Utils')


class Params:
    '''Class that loads hyperparameters from a json file.
    Example:
    params = Params(json_path)
    print(params.learning_rate)
    params.learning_rate = 0.5  # change the value of learning_rate in params
    '''

    def __init__(self, json_path):
        with open(json_path) as f:
            params = json.load(f)
            self.__dict__.update(params)

    def save(self, json_path):
        with open(json_path, 'w') as f:
            json.dump(self.__dict__, f, indent=4, ensure_ascii=False)

    def update(self, json_path):
        '''Loads parameters from json file'''
        with open(json_path) as f:
            params = json.load(f)
            self.__dict__.update(params)

    @property
    def dict(self):
        '''Gives dict-like access to Params instance by params.dict['learning_rate']'''
        return self.__dict__


class RunningAverage:
    '''A simple class that maintains the running average of a quantity
    Example:
    loss_avg = RunningAverage()
    loss_avg.update(2)
    loss_avg.update(4)
    loss_avg() = 3
    '''

    def __init__(self):
        self.steps = 0
        self.total = 0

    def update(self, val):
        self.total += val
        self.steps += 1

    def __call__(self):
        return self.total / float(self.steps)


def set_logger(log_path):
    '''Set the logger to log info in terminal and file `log_path`.
    In general, it is useful to have a logger so that every output to the terminal is saved
    in a permanent file. Here we save it to `model_dir/train.log`.
    Example:
    logging.info('Starting training...')
    Args:
        log_path: (string) where to log
    '''
    _logger = logging.getLogger('DeepAR')
    _logger.setLevel(logging.INFO)

    fmt = logging.Formatter('[%(asctime)s] %(name)s: %(message)s', '%H:%M:%S')

    class TqdmHandler(logging.StreamHandler):
        def __init__(self, formatter):
            logging.StreamHandler.__init__(self)
            self.setFormatter(formatter)

        def emit(self, record):
            msg = self.format(record)
            tqdm.write(msg)

    file_handler = logging.FileHandler(log_path)
    file_handler.setFormatter(fmt)
    _logger.addHandler(file_handler)
    _logger.addHandler(TqdmHandler(fmt))


def save_dict_to_json(d, json_path):
    '''Saves dict of floats in json file
    Args:
        d: (dict) of float-castable values (np.float, int, float, etc.)
        json_path: (string) path to json file
    '''
    with open(json_path, 'w') as f:
        # We need to convert the values to float for json (it doesn't accept np.array, np.float, )
        d = {k: float(v) for k, v in d.items()}
        json.dump(d, f, indent=4)


def save_checkpoint(state, is_best, epoch, checkpoint, ins_name=-1):
    '''Saves model and training parameters at checkpoint + 'last.pth.tar'. If is_best==True, also saves
    checkpoint + 'best.pth.tar'
    Args:
        state: (dict) contains model's state_dict, may contain other keys such as epoch, optimizer state_dict
        is_best: (bool) True if it is the best model seen till now
        checkpoint: (string) folder where parameters are to be saved
        ins_name: (int) instance index
    '''
    if ins_name == -1:
        filepath = os.path.join(checkpoint, f'epoch_{epoch}.pth.tar')
    else:
        filepath = os.path.join(checkpoint, f'epoch_{epoch}_ins_{ins_name}.pth.tar')
    if not os.path.exists(checkpoint):
        logger.info(f'Checkpoint Directory does not exist! Making directory {checkpoint}')
        os.mkdir(checkpoint)
    torch.save(state, filepath)
    logger.info(f'Checkpoint saved to {filepath}')
    if is_best:
        shutil.copyfile(filepath, os.path.join(checkpoint, 'best.pth.tar'))
        logger.info('Best checkpoint copied to best.pth.tar')


def load_checkpoint(checkpoint, model, optimizer=None):
    '''Loads model parameters (state_dict) from file_path. If optimizer is provided, loads state_dict of
    optimizer assuming it is present in checkpoint.
    Args:
        checkpoint: (string) filename which needs to be loaded
        model: (torch.nn.Module) model for which the parameters are loaded
        optimizer: (torch.optim) optional: resume optimizer from checkpoint
        gpu: which gpu to use
    '''
    if not os.path.exists(checkpoint):
        raise FileNotFoundError(f"File doesn't exist {checkpoint}")
    if torch.cuda.is_available():
        checkpoint = torch.load(checkpoint, map_location='cuda')
    else:
        checkpoint = torch.load(checkpoint, map_location='cpu')
    model.load_state_dict(checkpoint['state_dict'])

    if optimizer:
        optimizer.load_state_dict(checkpoint['optim_dict'])

    return checkpoint


def plot_all_epoch(variable, save_name, location='./figures/'):
    num_samples = variable.shape[0]
    x = np.arange(start=1, stop=num_samples + 1)
    f = plt.figure()
    plt.plot(x, variable[:num_samples])
    f.savefig(os.path.join(location, save_name + '_summary.png'))
    plt.close()


def init_metrics(sample=True):
    metrics = {
        'ND': np.zeros(2),  # numerator, denominator
        'RMSE': np.zeros(3),  # numerator, denominator, time step count
        'test_loss': np.zeros(2),
    }
    if sample:
        metrics['rou90'] = np.zeros(2)
        metrics['rou50'] = np.zeros(2)
    return metrics


def get_crps(sample_mu, sample_sigma, labels, predict_start):
    crps=[]
    for idx in range(sample_mu.shape[0]):
        pscore_l = []
        for t in range(sample_mu.shape[1]):
            mu = sample_mu[idx,t]
            sigma = sample_sigma[idx,t]
            torch.manual_seed(230)
            pred_sample = torch.normal(mu, sigma, size=(200,)) # sample num = 200
            # delete the negative sample
            pred_sample = pred_sample[pred_sample>0]
            if len(pred_sample) ==0:
                pred_sample=torch.tensor([0,0])
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            pred_sample=pred_sample.data.cpu().numpy()
            pscore_l.append(torch.tensor(pscore(pred_sample,labels[idx, predict_start+t].data.cpu().numpy()).compute()[0]).to(device))
        crps.append(torch.mean(torch.stack(pscore_l)).item())
    return np.array(crps)

def winkler_score(l_interval, u_interval, y_t, alpha):
    '''
    Calculates the Winkler score for an observation in 
    an prediction interval.
    
    Winkler scores penalise intervals that do not capture the 
    observation proportional to 2/alpha

    Params:
    -------
    l_interval: float
        Lower prediction interval
        
    u_interval: float
        Upper prediction interval
        
    y_t: float
        Observed ground truth value
        
    alpha: float
        The prediction interval alpha.  For an 80% pred intervals alpha=0.2
        
    Returns:
    -------
    float
    
    Example usage:
    --------------
    ```python
    >>> alpha = 0.2
    >>> interval = [744.54, 773.22]
    >>> y_t = 741.84
    >>> ws = winkler_score(interval, y_t, alpha)
    >>> print(round(ws, 2))
    
    56.68
    ```
    '''
    score = u_interval - l_interval
    if y_t < l_interval:
        score += ((2/alpha) * (l_interval - y_t))
    elif y_t > u_interval:
        score += ((2/alpha) * (y_t - u_interval))
        
    return score


def get_ws(sample_mu, sample_sigma, labels, predict_start):
    ws=[]
    for idx in range(sample_mu.shape[0]):
        ws_l = []
        for t in range(sample_mu.shape[1]):
            mu = sample_mu[idx,t]
            sigma = sample_sigma[idx,t]
            torch.manual_seed(230)
            pred_sample = torch.normal(mu, sigma, size=(200,)) # sample num = 200
            # delete the negative sample
            pred_sample = pred_sample[pred_sample>0]
            if len(pred_sample) ==0:
                pred_sample=torch.tensor([0,0])
            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
            ws_=winkler_score(np.percentile(np.array(pred_sample),0.1), np.percentile(np.array(pred_sample),0.9), labels[idx, predict_start+t], 0.2)
            ws_l.append(torch.tensor(ws_).to(device))
        ws.append(torch.mean(torch.stack(ws_l)).item())
    return np.array(ws)

def get_metrics(sample_mu, sample_sigma, labels, predict_start, samples=None, relative=False):
    metric = dict()
    metric['ND'] = net.accuracy_ND_(sample_mu, labels[:, predict_start:], relative=relative)
    metric['RMSE'] = net.accuracy_RMSE_(sample_mu, labels[:, predict_start:], relative=relative)
    metric['CRPS']= get_crps(sample_mu, sample_sigma, labels, predict_start)
    metric['WS']=get_ws(sample_mu, sample_sigma, labels, predict_start)
    if samples is not None:
        metric['rou90'] = net.accuracy_ROU_(0.9, samples, labels[:, predict_start:], relative=relative)
        metric['rou50'] = net.accuracy_ROU_(0.5, samples, labels[:, predict_start:], relative=relative)
    
    return metric

def update_metrics(raw_metrics, input_mu, input_sigma, sample_mu, labels, predict_start, lam, dataset,samples=None, relative=False):
    raw_metrics['ND'] = raw_metrics['ND'] + net.accuracy_ND(sample_mu, labels[:, predict_start:], relative=relative)
    raw_metrics['RMSE'] = raw_metrics['RMSE'] + net.accuracy_RMSE(sample_mu, labels[:, predict_start:], relative=relative)
    input_time_steps = input_mu.numel()
    ls = net.loss_fn2(input_mu, input_sigma, labels[:, :predict_start],lam,dataset) * input_time_steps
    ls = ls.data.cpu().numpy()
    raw_metrics['test_loss'] = raw_metrics['test_loss'] + [ls, input_time_steps]
    if samples is not None:
        raw_metrics['rou90'] = raw_metrics['rou90'] + net.accuracy_ROU(0.9, samples, labels[:, predict_start:], relative=relative)
        raw_metrics['rou50'] = raw_metrics['rou50'] + net.accuracy_ROU(0.5, samples, labels[:, predict_start:], relative=relative)
    return raw_metrics


def final_metrics(raw_metrics, sampling=False):
    summary_metric = {}
    summary_metric['ND'] = raw_metrics['ND'][0] / raw_metrics['ND'][1]
    summary_metric['RMSE'] = np.sqrt(raw_metrics['RMSE'][0] / raw_metrics['RMSE'][2]) / (
                raw_metrics['RMSE'][1] / raw_metrics['RMSE'][2])
    summary_metric['test_loss'] = (raw_metrics['test_loss'][0] / raw_metrics['test_loss'][1]).item()
    if sampling:
        summary_metric['rou90'] = raw_metrics['rou90'][0] / raw_metrics['rou90'][1]
        summary_metric['rou50'] = raw_metrics['rou50'][0] / raw_metrics['rou50'][1]
    return summary_metric