# Make an image where pixels are colored if they are prime

This code creates a black-and-white image of variable size where the n-th pixel is white if n is a prime number and black if it isn't. It works, but I wonder if it is well written or if it could be made faster.

Example:

Code:

import math

from PIL import Image

def is_prime(num: int) -> bool:
if not isinstance(num, int):
raise TypeError("num must be integer!")
if num < 2:
return False
if num == 2:
return True
for i in range(2, math.ceil(num ** 0.5) + 1):
if num % i == 0:
return False
return True

def bw(i: int) -> tuple[int, int, int]:
return (255 if is_prime(i) else 0,) * 3

def image(x_size: int, y_size: int = 0):
if y_size == 0:
y_size = x_size
image = Image.new("RGB", (x_size, y_size), "black")
for y in range(y_size):
print(y * x_size)
for x in range(x_size):
image.putpixel((x, y), bw(x + y * x_size + 1))
image.save("primes.png")

if __name__ == "__main__":
image(1000)

• Take a look at the questions under the tag sieve-of-eratosthenes. Apr 6, 2023 at 20:35
• That's a cool and beautiful visualization. Well done! Apr 7, 2023 at 9:07
• @LamarLatrell: The image code generates an image 1000 pixels wide, which is a multiple of 5. Hence all pixels in a column are congruent modulo 5, and those that are congruent to zero modulo 5 form columns of all black. (The same should be true of 2, every second column should be entirely black. Perhaps only black stripes of width 3, where a multiple of 5 falls between two multiples of 2 are wide enough to be visible.) Apr 7, 2023 at 21:49
• I can say for sure that the image storage is inefficient; why are you storing a series of 1-bit values in a 24bpp format? PIL.Image.new("1", (x_size, y_size), "black") for a 1bpp black/white format (often used for masks), s'il vous plaît! Your file will be a tiny fraction of the size. Or at the very least, PIL.Image.new("L", ...) (for 8bpp grayscale).
– FeRD
Apr 8, 2023 at 18:30
• I was occupied by writing my script so I failed to post an answer in time. Now I am late and there are already answers mentioning numpy and prime sieve... But they failed to mention that your image doesn't show the regularity of the prime gaps. If you use polar coordinates instead of Cartesian coordinates and put every number n at n radians and n distance from the origin, you will see patterns that your visualization failed to show. And if you create a new image using a different size, the relative location of the primes will change drastically, using polar coordinates will mitigate it. Apr 13, 2023 at 20:13

There are some things I would comment on, outside of what DeathIncarnate said in their answer.

The is_prime is fine, I would change two small details:

1. It is uncommon in Python to check the type of the input value (if not isinstance(num, int)). If the input is not an int, then all sorts of things are going wrong in the calling function. Is it really the responsibility of this function to avoid those? You could add an assert, which can be disabled if necessary. But mostly you should rely on the static type checking, which is what you have the type annotations for. This happens outside of executing the Python code, and therefore does not influence the execution speed.

2. The range range(2, math.ceil(num ** 0.5) + 1) is too wide. If num**0.5 is not an integer, then the square of its ceil will be larger than num, and therefore does not need to be checked. You should use math.floor instead.

But it generally contains too many numbers: if you can't divide by 2, you won't be able to divide by 4, 6, 8, ... This range could skip all even numbers. In fact, this range should skip all non-prime numbers, which is where my comment under your question hinted at.

The sieve of Erastothenes would be the standard way of producing a list of all primes up to a given number. It checks each new number to not be divisible by the primes found so far (in a manner of speaking, it actually doesn't use any divisions at all). In your case, you need a list of all primes up to and including x_size * y_size. Let's call this list primes (we'll come back to this later).

I agree with DeathIncarnate, you don't need an RGB image here, you just need a gray-scale image, which has a single channel. You can write a PNG of such a gray-scale image just fine. Because your code will be simpler, it is preferable.

To represent the image, I would use a NumPy array. Because NumPy arrays are easier to work with and more efficient to index than PIL.Image objects. Converting between the two is trivial, typically no data needs to be copied.

Earlier I referred primes, the list of primes. This list can be used to index into a NumPy array:

import numpy as np

# ...

img = np.zeros((y_size, x_size), dtype=np.uint8)
img.reshape(-1)[primes] = 255


That's it, that's all you need to do to create your image.

...well, in your code you start counting pixels from 1, not 0 as I did there. To subtract 1 from your primes, you can convert the list to a NumPy array first:

img.reshape(-1)[np.array(primes) - 1] = 255


Now, to save the image using Pillow, you can cast to a PIL.Image:

Image.fromarray(img).save('primes.png')


I wanted to try these ideas out. rdesparbes posted code with a better prime list generation similar to ideas I wrote above. This reduced the computation time on my machine from the 2.8s of OP's code to 0.63s. Using a NumPy array instead of a PIL.Image object, the time went down about 12%, to 0.55s. Implementing an actual sieve of Erastothenes (which doesn't do any divisions at all) reduced the time to 1/5th of that, 0.11s. This is about 25 times faster than OP's code.

This is the code I put together for testing. It's a bit messy, sorry for that, don't pay attention to the style, function names, or the functions themselves writing to file.

Note that about half of the 0.11s is converting the NumPy array to PIL.Image, and saving it to file.

import math
import time

import numpy as np
from PIL import Image

# --- OP ---

def is_prime(num: int) -> bool:
if not isinstance(num, int):
raise TypeError("num must be integer!")
if num < 2:
return False
if num == 2:
return True
for i in range(2, math.ceil(num ** 0.5) + 1):
if num % i == 0:
return False
return True

def bw(i: int) -> tuple[int, int, int]:
return (255 if is_prime(i) else 0,) * 3

def original_image(x_size: int, y_size: int = 0):
if y_size == 0:
y_size = x_size
image = Image.new("RGB", (x_size, y_size), "black")
for y in range(y_size):
for x in range(x_size):
image.putpixel((x, y), bw(x + y * x_size + 1))
image.save("original_image.png")

# --- RDESPARBES https://codereview.stackexchange.com/a/284350/151754 ---

def compute_primes(maximum: int) -> list[int]:
# updated by Cris to not test even numbers, and a few other tweaks
if maximum <= 2:
return []
primes = [2]

def is_prime(number_to_test_: int) -> bool:
max_threshold = int(number_to_test_ ** 0.5)
for prime in primes:
if prime > max_threshold:
break
if number_to_test_ % prime == 0:
return False
return True

for number_to_test in range(3, maximum, 2):
if is_prime(number_to_test):
primes.append(number_to_test)
return primes

def rdesparbes_image(x_size: int, y_size: int | None = None) -> None:
if y_size is None:
y_size = x_size
primes: list[int] = compute_primes(maximum=x_size * y_size + 1)  # fixed upper limit!
image = Image.new("RGB", (x_size, y_size), "black")
for prime in primes:
y, x = divmod(prime - 1, x_size)
image.putpixel((x, y), (255, 255, 255))
image.save("rdesparbes_image.png")

# --- CRIS #1: like above, but using NumPy array ---

def better_image(x_size: int, y_size: int | None = None) -> None:
if y_size is None:
y_size = x_size
primes = compute_primes(x_size * y_size + 1)
img = np.zeros((y_size, x_size), dtype=np.uint8)
img.reshape(-1)[np.array(primes) - 1] = 255
Image.fromarray(img).save('better_image.png')

# --- CRIS #2: sieve of Erastothenes ---

def sieve_image(x_size: int, y_size: int | None = None) -> None:
if y_size is None:
y_size = x_size
primes = np.ones(x_size * y_size + 1, dtype=np.bool_)
primes[0:2] = False
for ii in range(2, len(primes)):
if primes[ii]:
primes[ii**2::ii] = False
primes = np.reshape(primes[1:], (y_size, x_size)).astype(np.uint8)
primes *= 255
Image.fromarray(primes).save('sieve_image.png')

# --- MAIN ---

if __name__ == "__main__":
t = time.time()
original_image(1000)
print(f"{time.time() - t:0.3f}s")

t = time.time()
rdesparbes_image(1000)
print(f"{time.time() - t:0.3f}s")

t = time.time()
better_image(1000)
print(f"{time.time() - t:0.3f}s")

t = time.time()
sieve_image(1000)
print(f"{time.time() - t:0.3f}s")

• Would it be easier to work with a 1-D array and reshape it just before converting to an image? Apr 7, 2023 at 18:54
• @MarkRansom Yes, you could replace img = np.zeros((y_size, x_size), dtype=np.uint8); img.reshape(-1)[primes] = 255 with img = np.zeros((y_size * x_size), dtype=np.uint8); img[primes] = 255; img = np.reshape(img, (y_size, x_size)). Either way is quite simple, I think. :) Apr 7, 2023 at 19:33
• Ah I see, I missed that primes was already a 1-D array. Carry on. Apr 7, 2023 at 20:24
• Observation: A while back I did some speed testing on the Sieve and found that it was a bad idea as far along as I tested--main memory hits are far more expensive than divisions. It made sense in the old days when division was expensive and memory hits weren't nearly as bad. Apr 9, 2023 at 0:37
• @LorenPechtel Here we just do primes up to 1 million, it seems that the sieve performs very well in this range. See the update, I added some code and timing at the bottom of my post. Apr 17, 2023 at 23:00

Now, I will preface this by saying I've never used pillow before, so this is likely not the most effective means of doing this.

## Docstrings

You have used type annotations well in your functions which is a good start, but you may want to look into adding docstrings to describe their function.

## Arguments

Why do you default y_size to 0, as your not-entered? It is more idiomatic to use None (you may choose -1 if you want to avoid None) as a not supplied value rather than something which is potentially "valid". Consider if I was automating using this I could imagine a scenario where I requested a 0 element image (of course if I requested a -1 sized image that should probably error out)

Also, you don't allow the user to write to a file of their choosing. You should probably pass that as an argument through to the function to avoid users having to manually update their code or accidentally overwriting different sized images.

The main image variable inside your function has the same name as the function itself which is not a good thing and could lead to confusion. It is common practice to give functions names which are verbs such as generate_image, even this is a little generic and you might want generate_primes_image

## Image Type

You create a full RGB image while pillow offers a convenient bitwise image which seems to be what you want. Additionally, with the bitmap we would no longer need bw and simply is_prime can be our generator.

The pillow documentation also notes:

Image.putpixel(xy, value)[source]

Modifies the pixel at the given position. The color is given as a single numerical value for single-band images, and a tuple for multi-band images. In addition to this, RGB and RGBA tuples are accepted for P and PA images.

Note that this method is relatively slow. For more extensive changes, use paste() or the ImageDraw module instead.

So we probably might want to look at other means of generating it.

## Using generators

Looking into the documentation it seems we can generate our image more efficiently by running the loops as a generator:

def generate_primes_image(x_size: int, y_size: int = None, output_file: str = "primes.png"):
""" Generates a pixel representation of the primes as a 2d image saved to
output_file
"""
if y_size is None:
y_size = x_size

# Using 255*is_prime because this packs to length-8 bytes
# mode is greyscale ("L") to match
u = (is_prime(x)*255 for x in range(1, (x_size*y_size)+1))
image = Image.frombytes("L", (x_size, y_size), bytes(u))
image.save(output_file)


This should actually generate the data through C all at once which should be more efficient.

You could also look at the Sieve of Erastothenes (as suggested in the comments) or other prime generating methods and use the ImageDraw module to modify chunks of the image as you find the elements which are not prime.

• A good answer, except I don't see how y_size==0 could be a legal value. An image without pixels is not an image. Apr 6, 2023 at 21:07
• It's true that it only writes out an empty image, but I would still argue it's an image, just as an empty list is still a list. Certainly up to taste (and what's supported by pillow). I would also say that None (or a definitely illegal value such as -1) is more definitively "empty". Apr 6, 2023 at 21:24
• I agree, I would definitely choose None there. Apr 6, 2023 at 21:26
• I just tried this out, it seems that Pillow will happily create an Image object where one dimension is 0 (ridiculous!). But if you try to save that as a PNG, it will error out (no image file formats that I know of support images with zero pixels, this is not exclusive to PNG). Apr 6, 2023 at 21:43
• @CrisLuengo I imagine PIL allows 0-dimension images because it has tools like PIL.Image.Image.paste() — if you're assembling an image programmatically by pasting sections cropped from elsewhere (or generated outright), it's theoretically possible that you might occasionally end up taking/generating a no-sized section to insert — calling dest.paste(no_size_image) is a lot better than accidentally calling dest.paste(None) and taking home a traceback for your troubles.
– FeRD
Apr 8, 2023 at 19:00

# Style and design

if not isinstance(num, int):
raise TypeError("num must be integer!")


This piece of code is not really necessary: the type annotation makes it clear that you expect an int here. If you are worried about what the user of this function could provide as an argument, you can add a docstring, or use a static type checker like mypy

def image(x_size: int, y_size: int = 0):
if y_size == 0:
y_size = x_size

• This behavior could be described in a docstring, as it is not obvious from the signature of the function itself.
• Also, on a more general note, good job on adding type annotations. But it's a good practice to also indicate that a function does not return anything with -> None:.
• image is a confusing name for a function. Usually, functions start with a verb. This function's name conflicts with the name of the actual Image.Image instance. compute_image or generate_prime_image would probably be a better name for this function.

## Split responsibilities

image.save("primes.png")


This part makes the image function hard to reuse in another context, and should probably be moved elsewhere.

# Performance

This script is slow for two main reasons:

• The print statement, which makes expensive system calls
• Primes that have already been computed could be reused to compute new primes (see Sieve of Eratosthenes)

# Refactoring

Here is another version of the script that shows the same behavior, but is apparently faster because prime numbers are saved and reused to compute new ones:

from __future__ import annotations

import math

from PIL import Image

def compute_primes(maximum: int) -> List[int]:
"""Compute the list of the prime number

:param maximum: Upper bound of the list, not included
:return: The prime numbers, in ascending order
"""
if maximum <= 2:
return []
primes: List[int] = []

def is_prime(number_to_test_: int) -> bool:
max_threshold: int = math.ceil(number_to_test_ ** 0.5) + 1
for prime in primes:
if prime >= max_threshold:
break
if number_to_test_ % prime == 0:
return False
return True

for number_to_test in range(2, maximum):
if is_prime(number_to_test):
primes.append(number_to_test)
return primes

def generate_prime_image(x_size: int, y_size: int | None = None) -> Image.Image:
if y_size is None:
y_size = x_size
primes: list[int] = compute_primes(maximum=x_size * y_size)
image = Image.new("RGB", (x_size, y_size), "black")
for prime in primes:
y, x = divmod(prime - 1, x_size)
image.putpixel((x, y), (255, 255, 255))
return image

def main() -> None:
image = generate_prime_image(1000)
image.save("primes.png")

if __name__ == "__main__":
main()

• I would probably argue that you might want primes to be a set rather than a list for fast indexing although you're just looping over it and not reusing it so not necessary. I'm also a little confused why you're defining is_prime in the loop, is that an indentation issue from SE formatting? Apr 6, 2023 at 21:57
• We do not need to check if an element is present in the primes collection, so there is no need to use a set. lists are usually faster and more memory efficient (stackoverflow.com/a/17945009/17674349) is_prime is an inner function confined to the scope of the compute_primes function, as it is tightly coupled to it; it needs a read-only access to the primes list in a very specific state, which makes it hard to reuse if we defined it as a standard function. The prime function could be inlined, but I found it less readable Apr 7, 2023 at 9:32
• You can def the function once outside the loop and it would still have access to the variables in the function scope. Apr 7, 2023 at 10:08
• Oh yes indeed! I will modify my answer. (for those who are interested, I performed a benchmark with timeit, and measured a ~20% increase in speed when the function definition is outside the loop, which totally make sense...) Apr 7, 2023 at 10:35

I'm not a Python "expert", but, I hope you appreciate what I have to say.

Besides everything that was said, I do see 2 issues with your code:

# 1 - You have a magic number

You have the following lines:

if __name__ == "__main__":
image(1000)


You could use the argparse module to receive the size as an argument.

Here's my example on how it can be implemented:

if __name__ == "__main__":
import argparse

parser = argparse.ArgumentParser(
prog="Some Program name here",
description="Generates a black-and-white image of variable size where the n-th pixel is white if n is a prime number and black if it isn't."
)

args = parser.parse_args()

image(args.size)


# 2 - The image name and path are hardcoded

This makes it pretty inflexible to use.

You can change your image function to have the following:

def image(x_size: int, y_size: int = 0, fp = "primes.png"):
[...]
image.save(fp)


Then, you can modify the code I gave you before, to allow passing a file name:

if __name__ == "__main__":
import argparse

parser = argparse.ArgumentParser(
prog="Some Program name here",
description="Generates a black-and-white image of variable size where the n-th pixel is white if n is a prime number and black if it isn't."
)


This basically allows you to actually make something useful out of that if __name__ == "__main__":.
You can just run it like this: python3 file.py --size=2500 --file="huge.jpg".