Visualize 1D numpy array as 2D array with matplotlib

Question

I have a 2D array of all the numbers 1 to 100 split by 10. And boolean values for each number being prime or not prime. I'm struggling to figure out how to visualize it like in the image below.

Here is my code to help understand what I have better.

I want to visualize it like this pic online. image

# excersize
is_prime = np.ones(100, dtype=bool) # array will be filled with Trues since 1 = True
# For each integer j starting from 2, cross out its higher multiples:
N_max = int(np.sqrt(len(is_prime) - 1))
for j in range(2, N_max + 1):
 is_prime[2*j::j] = False
# split an array up into multiple sub arrays
split_primes = np.split(is_prime, 10);
# create overlay for numbers
num_overlay = np.arange(100)
split_overlay = np.split(num_overlay, 10)
plt.plot(split_overlay)

• ImportanceOfBeingErnest: Check the docs at matplotlib.org/gallery/images_contours_and_fields/…
• Parker Shamblin: @ImportanceOfBeingErnest Thanks for finding this for me. I'm brand new and currently following along with the scipy lectures notes at scipy.org. It's a lot to take in but this looks exactly what I'm looking for.
• Parker Shamblin: @ImportanceOfBeingErnest I was curious if you could help lead me to where I can find out how to create a 2d array of all the numbers 1 - 100 right of initally instead of creating a 1d array first and then splitting it up into 10 chunks with np.split() ?
• ImportanceOfBeingErnest: np.arange(1, 101).reshape(10, 10) would give a 2D arrays with numbers 1 to 100 in 10 rows.

Answer 1

Creating 2D array of the numbers

Check out the documentation for numpy's reshape function. Here you can turn your array into a 2D array by doing:

data = is_prime.reshape(10,10)

we can also make an array of the first 100 integers to use for labeling in a similar fashion:

integers = np.arange(100).reshape(10,10)

Plotting the 2D array

When plotting in 2D you need to use one of the 2D functions that matplotlib provides: e.g. imshow, matshow, pcolormesh. You can either call these functions directly on your array, in which case they will use a colormap and each pixel's color will correspond to the value in associated spot in the array. Or you can explicitly make an RGB image which affords you a bit more control over the color of each box. For this case I think that that is a bit easier to do so the below solution uses that approach. However if you want to annotate heatmaps the matplolib documentation has a great resource for that here. For now we will create an array of RGB values (shape of 10 by 10 by 3) and change the colors of only the prime numbers using numpy's indexing abilities.

#create RGB array that we will fill in
rgb = np.ones((10,10,3)) #start with an array of white
rgb[data]=[1,1,0] # color the places where the data is prime to be white
plt.figure(figsize=(10,10))
plt.imshow(rgb)
# add number annotations
integers = np.arange(100).reshape(10,10)
#add annotations based on: https://stackoverflow.com/questions/20998083/show-the-values-in-the-grid-using-matplotlib
for (i, j), z in np.ndenumerate(integers):
 plt.text(j, i, '{:d}'.format(z), ha='center', va='center',color='k',fontsize=15)
# remove axis and tick labels
plt.axis('off')
plt.show()

Resulting in this image:array with primes highlighted in yellow

• Matt Hall: Nice! Just wanted to add that if the edges are important to the OP, they can use plt.pcolormesh, e.g. like plt.pcolormesh(rgb, edgecolor='black', lw=1); plt.axes().set_aspect('equal'). (AFAIK you don't get equal aspect ratio for free with pcolormesh.)