utils.py

import numpy as np
import math
import random
import os
import json
import pickle
from scipy.sparse import csr_matrix
from tqdm import tqdm
import multiprocessing

import torch
import torch.nn.functional as F

def set_seed(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True

def check_path(path):
    if not os.path.exists(path):
        os.makedirs(path)
        print(f'{path} created')

def neg_sample(item_set, item_size):  # random sample an item id that is not in the user's interact history
    item = random.randint(1, item_size - 1)
    while item in item_set:
        item = random.randint(1, item_size - 1)
    return item

class EarlyStopping:
    """Early stops the training if validation loss doesn't improve after a given patience."""
    def __init__(self, checkpoint_path, patience=7, verbose=False, delta=0):
        """
        Args:
            patience (int): How long to wait after last time validation loss improved.
                            Default: 7
            verbose (bool): If True, prints a message for each validation loss improvement.
                            Default: False
            delta (float): Minimum change in the monitored quantity to qualify as an improvement.
                            Default: 0
        """
        self.checkpoint_path = checkpoint_path
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.delta = delta

    def compare(self, score):
        for i in range(len(score)):
            if score[i] > self.best_score[i]+self.delta:
                return False
        return True

    def __call__(self, score, model):

        if self.best_score is None:
            self.best_score = score
            self.score_min = np.array([0]*len(score))
            self.save_checkpoint(score, model)
        elif self.compare(score):
            self.counter += 1
            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(score, model)
            self.counter = 0

    def save_checkpoint(self, score, model):
        '''Saves model when validation loss decrease.'''
        if self.verbose:
            print(f'Validation score increased.  Saving model ...')
        torch.save(model.state_dict(), self.checkpoint_path)
        self.score_min = score





def generate_rating_matrix_train(user_seq, num_users, num_items):
    row = []
    col = []
    data = []
    for user_id, item_list in enumerate(user_seq):
        for item in item_list[:-3]:
            row.append(user_id)
            col.append(item)
            data.append(1)
    row = np.array(row)
    col = np.array(col)
    data = np.array(data)
    rating_matrix = csr_matrix((data, (row, col)), shape=(num_users, num_items))
    return rating_matrix

# def generate_rating_matrix_valid(user_seq, num_users, num_items):
#     # three lists are used to construct sparse matrix
#     row = []
#     col = []
#     data = []
#     for user_id, item_list in enumerate(user_seq):
#         for item in item_list[:-2]: #
#             row.append(user_id)
#             col.append(item)
#             data.append(1)

#     row = np.array(row)
#     col = np.array(col)
#     data = np.array(data)
#     rating_matrix = csr_matrix((data, (row, col)), shape=(num_users, num_items))

#     return rating_matrix

def generate_rating_matrix_valid(user_dict, num_users, num_items):
    row = []
    col = []
    data = []

    # Iterate over each user's sequences
    for user_id in user_dict:
        for (item_list, _) in user_dict[user_id]:
            # Exclude the last 2 items for validation
            for item in item_list[:-2]:
                row.append(user_id)
                col.append(item)
                data.append(1)

    # Convert lists to numpy arrays
    row = np.array(row)
    col = np.array(col)
    data = np.array(data)

    # Create the rating matrix in CSR format
    rating_matrix = csr_matrix((data, (row, col)), shape=(num_users, num_items))

    return rating_matrix

# def generate_rating_matrix_test(user_seq, num_users, num_items):
#     # three lists are used to construct sparse matrix
#     row = []
#     col = []
#     data = []
#     for user_id, item_list in enumerate(user_seq):
#         for item in item_list[:-1]: #
#             row.append(user_id)
#             col.append(item)
#             data.append(1)

#     row = np.array(row)
#     col = np.array(col)
#     data = np.array(data)
#     rating_matrix = csr_matrix((data, (row, col)), shape=(num_users, num_items))

#     return rating_matrix

def generate_rating_matrix_test(user_dict, num_users, num_items):
    """
    Generate the test rating matrix from user sequences.
    :param user_seq: List of item sequences for each user.
    :param user_dict: Dictionary to store sequences for each user.
    :param num_users: Number of users.
    :param num_items: Number of items.
    :return: Test rating matrix in CSR format.
    """
    row = []
    col = []
    data = []

    # Iterate over each user's sequences
    for user_id in user_dict:
        for item_list, _ in user_dict[user_id]:
            # Exclude the last 1 item for test
            for item in item_list[:-1]:
                row.append(user_id)
                col.append(item)
                data.append(1)

    # Convert lists to numpy arrays
    row = np.array(row)
    col = np.array(col)
    data = np.array(data)

    # Create the rating matrix in CSR format
    rating_matrix = csr_matrix((data, (row, col)), shape=(num_users, num_items))

    return rating_matrix

# def get_user_seqs(data_file):
#     """
#     load txt data file
#     :param data_file: path of dataset, every line: [user_id item1 item2 item3 ...]
#     :return:
#     """
#     lines = open(data_file).readlines()
#     user_seq = []
#     item_set = set()
#     for line in lines:
#         user, items = line.strip().split(' ', 1)
#         items = items.split(' ')
#         items = [int(item) for item in items]
#         user_seq.append(items)
#         item_set = item_set | set(items)
#     max_item = max(item_set)

#     num_users = len(lines)
#     num_items = max_item + 2

#     valid_rating_matrix = generate_rating_matrix_valid(user_seq, num_users, num_items)
#     test_rating_matrix = generate_rating_matrix_test(user_seq, num_users, num_items)
#     return user_seq, max_item, valid_rating_matrix, test_rating_matrix, num_users

def get_user_seqs(data_file):
    """
    Load txt data file with time information.
    Each line format: [user_id item1,time1 item2,time2 item3,time3 ...]
    Each user may have multiple lines.
    :param data_file: Path of dataset.
    :return: user_seq, max_item, time_seq, valid_rating_matrix, test_rating_matrix, num_users
    """
    user_seq = []  # List of item sequences
    time_seq = []  # List of time sequences
    id_seq = []
    item_set = set()  # Set of all items
    user_dict = {}  # Dictionary to store sequences for each user

    lines = open(data_file).readlines()
    for line in lines:
        user, items = line.strip().split(' ', 1)
        user = int(user)
        items = items.split(' ')

        # Parse item and time
        item_list = []
        time_list = []
        for item_time in items:
            item, time = item_time.split(',')
            item_list.append(int(item))
            time_list.append(int(time))

        # Add to user_dict
        if user not in user_dict:
            user_dict[user] = []
        user_dict[user].append((item_list, time_list))

    # Extract sequences for each user
    for user in user_dict:
        for item_list, time_list in user_dict[user]:
            user_seq.append(item_list)
            time_seq.append(time_list)
            item_set.update(item_list)
            id_seq.append(int(user))

    max_item = max(item_set)
    num_users = len(user_dict)

    # Generate rating matrices
    valid_rating_matrix = generate_rating_matrix_valid(user_dict, num_users, max_item + 2)
    test_rating_matrix = generate_rating_matrix_test(user_dict, num_users, max_item + 2)

    return user_seq, max_item, time_seq, valid_rating_matrix, test_rating_matrix, num_users, id_seq




def get_metric(pred_list, topk=10):
    NDCG = 0.0
    HIT = 0.0
    MRR = 0.0
    # [batch] the answer's rank
    for rank in pred_list:
        MRR += 1.0 / (rank + 1.0)
        if rank < topk:
            NDCG += 1.0 / np.log2(rank + 2.0)
            HIT += 1.0
    return HIT /len(pred_list), NDCG /len(pred_list), MRR /len(pred_list)



def recall_at_k(actual, predicted, topk):
    sum_recall = 0.0
    num_users = len(predicted)
    true_users = 0
    recall_dict = {}
    for i in range(num_users):
        act_set = set(actual[i])
        pred_set = set(predicted[i][:topk])
        if len(act_set) != 0:
            one_user_recall = len(act_set & pred_set) / float(len(act_set))
            recall_dict[i] = one_user_recall
            sum_recall += one_user_recall
            true_users += 1
    return sum_recall / true_users, recall_dict

def cal_mrr(actual, predicted):
    sum_mrr = 0.
    true_users = 0
    num_users = len(predicted)
    mrr_dict = {}
    for i in range(num_users):
        r = []
        act_set = set(actual[i])
        pred_list = predicted[i]
        for item in pred_list:
            if item in act_set:
                r.append(1)
            else:
                r.append(0)
        r = np.array(r)
        if np.sum(r) > 0:
            one_user_mrr = np.reciprocal(np.where(r==1)[0]+1, dtype=np.float64)[0]
            sum_mrr += one_user_mrr
            true_users += 1
            mrr_dict[i] = one_user_mrr
        else:
            mrr_dict[i] = 0.
    return sum_mrr / len(predicted), mrr_dict




def ndcg_k(actual, predicted, topk):
    res = 0
    ndcg_dict = {}
    for user_id in range(len(actual)):
        k = min(topk, len(actual[user_id]))
        idcg = idcg_k(k)
        dcg_k = sum([int(predicted[user_id][j] in
                         set(actual[user_id])) / math.log(j+2, 2) for j in range(topk)])
        res += dcg_k / idcg
        ndcg_dict[user_id] = dcg_k / idcg
    return res / float(len(actual)), ndcg_dict


# Calculates the ideal discounted cumulative gain at k
def idcg_k(k):
    res = sum([1.0/math.log(i+2, 2) for i in range(k)])
    if not res:
        return 1.0
    else:
        return res

def dcg_at_k(r, k, method=1):
    """Score is discounted cumulative gain (dcg)
    Relevance is positive real values.  Can use binary
    as the previous methods.
    Returns:
        Discounted cumulative gain
    """
    r = np.asfarray(r)[:k]
    if r.size:
        if method == 0:
            return r[0] + np.sum(r[1:] / np.log2(np.arange(2, r.size + 1)))
        elif method == 1:
            return np.sum(r / np.log2(np.arange(2, r.size + 2)))
        else:
            raise ValueError('method must be 0 or 1.')
    return 0.


def ndcg_at_k(r, k, method=1):
    """Score is normalized discounted cumulative gain (ndcg)
    Relevance is positive real values.  Can use binary
    as the previous methods.
    Returns:
        Normalized discounted cumulative gain
    """
    dcg_max = dcg_at_k(sorted(r, reverse=True), k, method)
    if not dcg_max:
        return 0.
    return dcg_at_k(r, k, method) / dcg_max