dcg and ndcg implementation in python

Question

I have implemented the Discounted Cumulative Gain (DCG) and Normalized Discounted Cumulative Gain (NDCG) in python. I am not sure if the code is correct or did I forget some important criteria for DCG and NDCG. Here is my code so far:

import numpy as np

def get_dcg_score(predictions: np.ndarray, test_interaction_matrix: np.ndarray, topK = 10) -> float:
"""
predictions - np.ndarray - predictions of the recommendation algorithm for each user.
test_interaction_matrix - np.ndarray - test interaction matrix for each user.

returns - float - mean dcg score over all user.
"""    
score = None

# TODO: YOUR IMPLEMENTATION.

score = []

for idx, (pred,test) in enumerate(zip(predictions,test_interaction_matrix)):
    print(idx,pred,test)
    for i, (j,jj) in enumerate(zip(pred[:topK], test[:topK])):

        if i == 0 and jj == 1:
            sc = jj
            score.append(sc)
        if i != 0 and jj == 1:
            sc = jj / np.log2(j+2)
            score.append(sc)
        if (i != 0 and jj == 0) or ( i == 0 and jj == 0):
            continue

               
score = sum(score)/len(predictions)

return score

I evaluate this on the two arrays.

predictions = np.array([[0, 1, 2, 3], [3, 2, 1, 0]])
test_interaction_matrix = np.array([[1, 0, 0, 0], [0, 0, 0, 1]])

dcg_score = get_dcg_score(predictions, test_interaction_matrix, topK=4)
print(dcg_score)

assert np.isclose(dcg_score, 1), "1 expected"

Now for the NDCG I need to implement Ideal Discounted Cumulative Gain (IDCG) first and then divide DCG by IDCG. Here what I have for NDCG.

def get_ndcg_score(predictions: np.ndarray, test_interaction_matrix: np.ndarray, topK = 10) -> float:
"""
predictions - np.ndarray - predictions of the recommendation algorithm for each user.
test_interaction_matrix - np.ndarray - test interaction matrix for each user.
topK - int - topK recommendations should be evaluated.

returns - average ndcg score over all users.
"""    
score = None

# TODO: YOUR IMPLEMENTATION.

score_idcg = []
for i, (vp, vt) in enumerate(zip(predictions,test_interaction_matrix)):
    element_sorted = sorted(vp,reverse=True)
    for j, (ele_p, ele_vt) in enumerate(zip(element_sorted, vt)):

        if j == 0 and ele_vt == 1:
            scr = ele_vt
            score_idcg.append(scr)
        if j != 0 and ele_vt == 1:
            scr = ele_vt / np.log2(j+2)
            score_idcg.append(scr)
        if (j != 0 and ele_vt == 0) or (j == 0 and ele_vt == 0):
            continue

print(score_idcg)                
score_idcg = sum(score_idcg)/len(predictions)
print(score_idcg) 

score_dcg = get_dcg_score(predictions, test_interaction_matrix, topK = 4)

score_ndcg = score_dcg / score_idcg

    
return score_ndcg

Again I test it on these two arrays:

predictions = np.array([[0, 1, 2, 3], [3, 2, 1, 0], [1, 2, 3, 0], [-1, -1, -1, -1]])
test_interaction_matrix = np.array([[1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 0, 0], [0, 0, 0, 0]])

ndcg_score = get_ndcg_score(predictions, test_interaction_matrix, topK=4)

assert np.isclose(ndcg_score, 1), "ndcg score is not correct."

Could somebody please look at my code and find why I don't get the right result for ndcg test? I just can't figure it out. Please also look at dcg implementation as well if it is faulty. Sorry for the horrible code. Write me if you need more info. Any suggestion is appreciated.