Search code examples
pythonnumpypyopengl

fastest way to iterate in python


I've never had to concern myself with this problem so far but now I need to use some large number of vertices that need to be buffered by PyOpenGL and it seems like the python iteration is the bottleneck. Here is the situation. I have an array of 3D points vertices, and at each step I have to compute a 4D array of colors for each vertices. My approach so far is:

upper_border = len(self.vertices) / 3
#Only generate at first step, otherwise use old one and replace values
if self.color_array is None:
     self.color_array = numpy.empty(4 * upper_border)  

for i in range(upper_border):
     #Obtain a color between a start->end color
     diff_activity = (activity[i] - self.min) / abs_diff  
     clr_idx = i * 4
     self.color_array[clr_idx] = start_colors[0] + diff_activity * end_colors[0]
     self.color_array[clr_idx + 1] = start_colors[1] + diff_activity * end_colors[1]
     self.color_array[clr_idx + 2] = start_colors[2] + diff_activity * end_colors[2]
     self.color_array[clr_idx + 3] = 1

Now I don't think there's anything else I can do to eliminate the operations from each step of the loop, but I'm guessing there has to be a more optimal performance way to do that loop. I'm saying that because in javascript for example, the same calculus produces a 9FPS while in Python I'm only getting 2-3 FPS.

Regards, Bogdan


Solution

  • To make this code faster, you need to "vectorise" it: replace all explicit Python loops with implicit loops, using NumPy's broadcasting rules. I can try and give a vectorised version of your loop:

    if self.color_array is None:
         self.color_array = numpy.empty((len(activity), 4))
    diff_activity = (activity - self.min) / abs_diff
    self.color_array[:, :3] = (start_colors + 
                               diff_activity[:, numpy.newaxis] + 
                               end_colors)
    self.color_array[:, 3] = 1
    

    Note that I had to do a lot of guessing, since I'm not sure what all your variables are and what the code is supposed to do, so I can't guarantee this code runs. I turned color_array into a two-dimensional array, since this seems more appropriate. This probably requires changes in other parts of the code (or you need to flatten the array again).

    I assume that self.min and abs_diff are scalars and all other names reference NumPy arrays of the following shapes:

    activity.shape == (len(vertices) // 3,)
    start_colors.shape == (3,)
    end_colors.shape == (3,)
    

    It also looks as if vertices is a one-dimensional array and should be a two-dimensional array.