random vertex colors using Numpy

The bit of research presented in a previous article was not completely academical: speed improvements using Numpy are very welcome when for example generating random vertex colors for very large meshes.
With the results of the measurements in the back of my mind and spurred by a discussion on BlenderArtists I created a new version of the random vertex colors addon. (available on GitHub), which reduced the time to generate random vertex colors for approximately 1 million faces from 3.2 seconds to 0.8 seconds on my machine (or even 0.4s for a 32 bit variant), i.e. up to 8 times faster which is not bad.

Some explanation

Assigning random vertex colors to faces means that you have to assign colors to all the loops of a polygon the same random color (for the difference between vertices and loops, check the BMesh design document. The same applies to the regular Mesh objects we use in the code below). This can be done in a straight forward manner once you have located the vertex color layer:

mesh = context.scene.objects.active.data
vertex_colors = mesh.vertex_colors.active.data
polygons = mesh.polygons
for poly in polygons:
 color = [random(), random(), random()]
 for loop_index in range(poly.loop_start, poly.loop_start + poly.loop_total):
  vertex_colors[loop_index].color = color

Straight forward as this may be, a loop inside a loop is time consuming and so is generating lots of random numbers, especially because Python does not do this in parallel even if you have a processor with multiple cores. Fortunately for us Blender comes bundled with Numpy, which is a library that can manipulate huge arrays of numbers in a very efficient manner. This allows for a much more efficient approach (although as shown previously a significant speed increase is only noticeable for large meshes).

startloop = np.empty(npolygons, dtype=np.int)
numloops = np.empty(npolygons, dtype=np.int)
polygon_indices = np.empty(npolygons, dtype=np.int)

polygons.foreach_get('index', polygon_indices)
polygons.foreach_get('loop_start', startloop)
polygons.foreach_get('loop_total', numloops)

colors = np.random.random_sample((npolygons,3))

We can even reduce storage (and get an additional speedup) if we change the types to 32 bit variants. There will be no loss of accuracy as these are the sizes used by Blender internally. (Would you do a lot of additional calculations this might be different of course). The change would only alter the declarations of the arrays:

startloop = np.empty(npolygons, dtype=np.int32)
numloops = np.empty(npolygons, dtype=np.int32)
polygon_indices = np.empty(npolygons, dtype=np.int32)

polygons.foreach_get('index', polygon_indices)
polygons.foreach_get('loop_start', startloop)
polygons.foreach_get('loop_total', numloops)

colors = np.random.random_sample((npolygons,3)).astype(np.float32)

As shown above we start out by creating Numpy array that will hold the startloop indices and the number of of loops in each polygon as well as an array that will hold the polygon indices itself. This last one isn't strictly needed for assigning random values because we don't care which random color we assign to which polygon but for other scenarios it might make more sense so we keep it here. We get all these indices from the Mesh object using the fast foreach_get method. We then use Numpy buil-in random_sample function the generate the random colors (3 random floats between 0 and 1) for all polygons.

loopcolors = np.empty((nloops,3)) # or loopcolors = np.empty((nloops,3), dtype=np.float32)
loopcolors[startloop] = colors[polygon_indices]
numloops -= 1
nz = np.flatnonzero(numloops)
while len(nz):
 startloop[nz] += 1
 loopcolors[startloop[nz]] = colors[polygon_indices[nz]]
 numloops[nz] -= 1
 nz = np.flatnonzero(numloops)

The real work is done in the code above: we first create an empty array to hold the colors for all individual loops. Then we assign all the loops with index startloop with the colors that corresponds to the color of the polygon it belongs to. Note that startloop and polygon_indices are arrays with the same length, i.e. the number of polygons. Now we also have an array numloops which holds for each polygon the number of loops. We decrement this number of loops by one for all polygons in one go (line 3) and create an array of indices of those elements in numloops that are still greater than zero (line 4). If we are still left with one or more indices (line 5) we increment the index of the startloop for all those nonzero indices (line 6).
Line 7 then assigns again in one go a polygon color to a loop at a certain index for all polygons where there are still a non zero number of loops. And finally we again reduce our numloop counters. Note that this all works because the loop indices of all loops of a polygon are consecutive.

loopcolors = loopcolors.flatten()
vertex_colors.foreach_set("color", loopcolors)

The final lines flatten the array of loop colors to a 1-D array and write back the colors to the Mesh object with the fast foreach_set method.

Discussion

Now even though we end up with about four times as much code, it is much faster while still quite readable, as long as you get your head around using index arrays. Basically our single while loops executes the inner loop of assigning colors to loops for all polygons in parallel. The penalty is for this speed increase is memory consumption: We use 5 large numpy arrays here.