[FIXED] How do I efficiently add random vertical segments into a numpy array?

Issue

I am trying to simulate rain using NumPy, they say an image is more than a thousand words so here is a description longer than two thousand words:

I already wrote the code, but I think my implementation is inefficient, so I want to know if NumPy has any builtin functions that can speed up the process:

import numpy as np
from PIL import Image
from random import random, randbytes

def rain(width, strikes=360, color=True, lw=3):
    assert not width % 16
    height = int(width * 9 / 16)
    img = np.zeros((height, width, 3), dtype=np.uint8)
    half = height / 2
    for i in range(strikes):
        x = round(random() * width)
        y = round(height - random() * half)
        x1 = min(x + lw, width - 1)
        if color:
            rgb = list(randbytes(3))
        else:
            rgb = (178, 255, 255)
        img[0:y, x:x1] = rgb
    
    return img

img1 = Image.fromarray(rain(1920))
img1.show()
img1.save('D:/rain.jpg', format='jpeg', quality=80, optimize=True)

img2 = Image.fromarray(rain(1920, color=False))
img2.show()
img2.save('D:/rain_1.jpg', format='jpeg', quality=80, optimize=True)

Solution

I was able to improve by 2 to 4 times faster.

1. Since raindrops do not stop in the upper half of the image, the upper half can be stretched out from the lower half after all strikes end.

2. Since broadcasting tuples is relatively slow, use 32-bit format color instead.

def rain(width=1920, strikes=360, color=True, lw=3):
    assert not width % 16
    height = int(width * 9 / 16)
    img = np.zeros((height, width), dtype=np.uint32)
    half = height / 2
    upper_bottom = int(half) - 1
    alpha = 255 << 24

    # Paint background.
    # The upper half will always be overwritten and can be skipped.
    img[upper_bottom:] = alpha

    for i in range(strikes):
        x = round(random() * width)
        y = round(height - random() * half)
        x1 = min(x + lw, width - 1)
        if color:
            # Pack color into int. See below for details.
            rgb = int.from_bytes(randbytes(3), 'big') + alpha
        else:
            # This is how to pack color into int.
            r, b, g = 178, 255, 255
            rgb = r + (g << 8) + (b << 16) + alpha

        # Only the lower half needs to be painted in this loop.
        img[upper_bottom:y, x:x1] = rgb

    # The upper half can simply be stretched from the lower half.
    img[:upper_bottom] = img[upper_bottom]

    # Unpack uint32 to uint8 x4 without deep copying.
    img = img.view(np.uint8).reshape((height, width, 4))

    return img

Note:

• Endianness is ignored. May not work on some platforms.
• Performance is greatly degraded if the image width is very large.
• If you are going to convert img to PIL.Image, compare its performance too as it is also improved.

Because of the rain overlaps each other (which makes removing for-loop hard) and because the strikes are not so many (which makes the room for improvement small), I find it difficult to optimize further. Hopefully this is enough.

Answered By - ken