Batch algebraic operations in Python

Question

I am using a for in Python to populate an array. Currently I iterate over all the elements in the array, and for each index I retrieve some relevant information from another array, and then I perform an operation.

This is the current code

for idx, vertex_color in enumerate(self.vertex_colors):
    coefficients = self.cubic_coefficients[idx*3:(idx*3)+3]

    # Applies cubic regression to obtain regressed mean amplitude
    c[idx] = coefficients[0] * current_beta**2 + coefficients[1] * current_beta + coefficients[2]

But my vector self.vertex_colors is very large, therefore this for loop is my bottle neck. Is there anyway to parallelize or perform with a single command these kind of operations? In a way I want to do what bsxfun does in MATLAB, but for an arbitrary function.

Answer 1

For your specific case you can do without the loop in this way if coefficients is a numpy.array.

import numpy

coefficients = numpy.array([1, 2, 3, 1, 2, 3])
current_beta = 1
c = (coefficients[0::3] * current_beta**2 
     + coefficients[1::3] * current_beta 
     + coefficients[2::3])

This will work for any length of coefficients. The syntax used for the slicing here basically starts at the point and then takes every 3rd element.