Avoid overflow when adding numpy arrays

You can achieve this by creating a third array of dtype uint8, plus a bool array (which together are more memory efficient that one uint16 array).

np.putmask is useful for avoiding a temp array.

a = np.array([100, 200, 250], dtype=np.uint8)
b = np.array([50, 50, 50], dtype=np.uint8)
c = 255 - b  # a temp uint8 array here
np.putmask(a, c < a, c)  # a temp bool array here
a += b

However, as @moarningsun correctly points out, a bool array takes the the same amount of memory as a uint8 array, so this isn’t necessarily helpful. It is possible to solve this by avoiding having more than one temp array at any given time:

a = np.array([100, 200, 250], dtype=np.uint8)
b = np.array([50, 50, 50], dtype=np.uint8)
b = 255 - b  # old b is gone shortly after new array is created
np.putmask(a, b < a, b)  # a temp bool array here, then it's gone
a += 255 - b  # a temp array here, then it's gone

This approach trades memory consumption for CPU.


Another approach is to precalculate all possible results, which is O(1) extra memory (i.e. independent of the size of your arrays):

c = np.clip(np.arange(256) + np.arange(256)[..., np.newaxis], 0, 255).astype(np.uint8)
c
=> array([[  0,   1,   2, ..., 253, 254, 255],
          [  1,   2,   3, ..., 254, 255, 255],
          [  2,   3,   4, ..., 255, 255, 255],
          ..., 
          [253, 254, 255, ..., 255, 255, 255],
          [254, 255, 255, ..., 255, 255, 255],
          [255, 255, 255, ..., 255, 255, 255]], dtype=uint8)

c[a,b]
=> array([150, 250, 255], dtype=uint8)

This approach is the most memory-efficient if your arrays are very big. Again, it is expensive in processing time, because it replace the super-fast integer additions with the slower 2dim-array indexing.

EXPLANATION OF HOW IT WORKS

Construction of the c array above makes use of a numpy broadcasting trick. Adding an array of shape (N,) and array of shape (1,N) broadcast both to be (N,N)-like, thus the result is an NxN array of all possible sums. Then, we clip it. We get a 2dim array that satisfies: c[i,j]=min(i+j,255) for each i,j.

Then what’s left is using fancy indexing the grab the right values. Working with the input you provided, we access:

c[( [100, 200, 250] , [50, 50, 50] )]

The first index-array refers to the 1st dim, and the second to the 2nd dim. Thus the result is an array of the same shape as the index arrays ((N,)), consisting of the values [ c[100,50] , c[200,50] , c[250,50] ].

Leave a Comment