Representing array elements on a square in Python (as image)

June 28, 2022

I have an array A with shape (3,3). Is there a way to represent the array elements on a square of size 3×3? In general, I would like to represent nxn arrays on a square of size nxn? The expected output is attached.

import numpy as np

A=np.array([[10,20,30],[40,50,60],[70,80,90]])

The expected output is

Output

>Solution :

You could use seaborn.heatmap that has a nice API:

import numpy as np
import seaborn as sns
from matplotlib.colors import ListedColormap

A = np.array([[10, 20, 30], [40, 50, 60], [70, 80, 90]])

ax = sns.heatmap(A,
            annot=True, square=True, cbar=False,
            xticklabels=False, yticklabels=False,
            cmap=ListedColormap(['white']),
            linecolor='k', lw=2,
            annot_kws={'size': 30}
           )

ax.figure.savefig('img.png')

output: