Convert matrix into a dictionary whose keys are the row and column indices of the matrix and values correspond to the matrix values for given key indices.

matrix_to_dict.py

import numpy as np
from tqdm import trange
from collections import defaultdict
from typing import Dict, Tuple, DefaultDict


def get_matrix_idx_to_value_dict(
    matrix: np.ndarray,
    verbose: bool = True,
) -> DefaultDict[Tuple[int, int], int]:
    """
    Convert matrix into a dictionary whose:
    1) keys are the row and column indices of the matrix
    2) values correspond to the matrix values for given key indices
    Useful for converting big confusion matrices.
    
    :param np.ndarray matrix: matrix to convert to dict.
    :param bool verbose: verbose.
    :return: dictionary whose:
             1) keys are the row and column indices of the matrix
             2) values correspond to the matrix values for given key indices
    :rtype: DefaultDict[Tuple[int, int], int]
    """
    
    assert matrix.ndim == 2, "matrix should have only 2 dimensions."
    
    row_col_idx_to_value = defaultdict(int)
    
    n_rows, n_cols = matrix.shape
    cols = range(n_cols)
    
    if verbose:
        rows = trange(n_rows)
    else:
        rows = range(n_rows)
    
    for row_idx in rows:
        for col_idx in cols:
            value = matrix[row_idx][col_idx]
            if value != 0:  # remove sparsity
                row_col_idx_to_value[(row_idx, col_idx)] = value
                
    return row_col_idx_to_value


matrix = np.array([[1, 0], [0, 2]])
row_col_idx_to_value = get_matrix_idx_to_value_dict(matrix)  # defaultdict(int, {(0, 0): 1, (1, 1): 2})