Calculating the push or pull force of a molecule at a vertex

Question

The basic idea of the code is to calculate the push or pull force at a vertex, given the number of "push causing molecules" and "pull causing molecules" at a polygon vertex.

The code thus mainly performs arithmetic (including the function it calls, calculate_edge_force), and has a couple of if-statements. The function is called many times since it is required in order to compute the derivatives, which is later used by scipy.odeint for integration.

Here is my code, in plain Python (note that import numpy as np has been put in somewhere higher up in the file, and that the functions are inside a class, hence there are references to this as self.<something>):

def calculate_point_forces(self, PushMolecules, PushMolecules_iPlus1, PushMolecules_iMinus1, PullMolecules, PullMolecules_iPlus1, PullMolecules_iMinus1, point_index, num_nodes, polygon, polygon_compression, point_interpolygon_contact, push_vector, pull_vector, a1, a2):

    max_push_force = self.max_push_force
    max_pull_force = self.max_pull_force
    tangent_factor = self.tangent_factor
    rac_force_exponent = self.rac_force_exponent
    rho_force_exponent = self.rho_force_exponent

    avg_PushMolecules_plus1 = (PushMolecules + PushMolecules_iPlus1)*0.5
    avg_PullMolecules_plus1 = (PullMolecules + PullMolecules_iPlus1)*0.5

    avg_PushMolecules_minus1 = (PushMolecules + PushMolecules_iMinus1)*0.5
    avg_PullMolecules_minus1 = (PullMolecules + PullMolecules_iMinus1)*0.5

    if self.enable_compression_stiffening == True:
        if polygon_compression < 1:
            F_internal = ((1-polygon_compression)+self.cytosolic_baseline)*self.cytosolic_pressure
        else:
            F_internal = 0
    else:
        F_internal = 0

    if PushMolecules > PullMolecules:
        Fpush = int(not point_interpolygon_contact)*max_push_force*(1 - (1/(1 + (PushMolecules/a2)**rac_force_exponent)))*push_vector
        Fpull = 0
        Fe_pull = 0
    elif PushMolecules < PullMolecules:
        pull_magnitude = max_pull_force*(1 - (1/(1 + (PullMolecules/a1)**rho_force_exponent)))
        Fpull = (1-tangent_factor)*pull_magnitude*pull_vector
        Fe_pull = 0.5*tangent_factor*pull_magnitude
        Fpush = 0
    else:
        Fpush = 0
        Fpull = 0
        Fe_pull = 0

    Fe_pull_minus1 = 0
    Fe_pull_plus1 = 0
    if PushMolecules_iMinus1 < PullMolecules_iMinus1:
        pull_magnitude_iMinus1 = max_pull_force*(1 - (1/(1 + (PullMolecules_iMinus1/a1)**rho_force_exponent)))
        Fe_pull_minus1 = 0.5*tangent_factor*pull_magnitude_iMinus1
    if PushMolecules_iPlus1 < PullMolecules_iPlus1:
        pull_magnitude_iPlus1 = max_pull_force*(1 - (1/(1 + (PullMolecules_iPlus1/a1)**rho_force_exponent)))
        Fe_pull_plus1 = 0.5*tangent_factor*pull_magnitude_iPlus1

    Fe_pull  = np.max([Fe_pull_minus1, Fe_pull_plus1, Fe_pull])

    Fe_plus, Fe_plus_ = self.calculate_edge_force(
        avg_PushMolecules_plus1, avg_PullMolecules_plus1, point_index,
        (point_index+1) % num_nodes, polygon)
    Fe_minus, Fe_minus_ = self.calculate_edge_force(
        avg_PushMolecules_minus1, avg_PullMolecules_minus1, point_index,
        (point_index-1) % num_nodes, polygon)

    F = Fpush + Fpull + Fe_plus + Fe_minus + Fe_pull*Fe_plus_ + Fe_pull*Fe_minus_ + push_vector*F_internal

    return F

When I profile this, I get that the function itself takes about 3 seconds. Out of this, 1 second is spent in calls to the calculate_edge_force and 1 second is spent by np.max.

If I "naively" put in float and int wherever I can (I don't know how to type the NumPy arrays, e.g. push_vector or pull_vector, so I only do it for things like PushMolecules, max_push_force or point_index, which I know will be float or int), I actually get a very mild slowdown in speed (about 0.2 seconds)!

Clearly then, that's not the way to go about "Cythonizing" my code. What should I be doing instead?

Answer 1

First of all you probably should type your numpy arrays but be warned, whether you use the np.ndarray notation, or the memory view notation, only limited operations occur at the speed of c, that is indexing, 'shape', and a very few others. With a memory view you can only use these fast operations, with a ndarray type you can access all the python operations too but they still occur at the speed of python, not c.

In this case, your code doesn't actually use any array functions that would benefit from the arrays being properly typed so you'd gain nothing in the mere act of typing the arrays. The overhead of calling a python function is significant so you might benefit from not using the numpy array functions, if the numpy arrays are short and if the functionality is straightforward to reproduce as cython code using array indexing (fast) and if the function is called many times. The profiler should tell you this stuff.

I can see some low-hanging fruit in the above code, take this line, mentioned to be a major offender:

Fe_pull  = np.max([Fe_pull_minus1, Fe_pull_plus1, Fe_pull])

What happens here, is np.max is a call to a python function of a python module, the list will be constructed as a python list and the floats will be converted to python floats because a python list can't hold c floats. Then that code will all run at the speed of python. np.max will then unpack all that stuff and the result of np.max will be a python float which is converted to a c float. This will be slow!. In fact it may even be slower than python, because it involves creating extra temporary python objects.

What you would do is rewrite that to eliminate the python function call (and thus all the intermediate python objects). As it happens, if you use the python built-in 'max' on typed variables, cython is smart enough to eliminate the python call entirely:

Fe_pull = max(Fe_pull_minus1, Fe_pull_plus1, Fe_pull)

That will automatically become optimal inline c code for calculating the max of the 3 variables (Cython actually unrolls it as if/else statements)

The next low-hanging fruit, and another stated offender is this:

Fe_plus, Fe_plus_ = self.calculate_edge_force( ...

In some cases Cython optimizes away tuple-unpacking, but it doesn't in the case of function calls. So what is going to happen is cython will create an intermediate tuple (a python object), pack stuff into the tuple (Again, converting to python objects, as tuples can't hold c types), then unpack the tuple and convert the python objects to their c types. Again this should be turned into something that translates readily to c. Unfortunately in this case there is no faster better cleaner syntax, you would probably be best served by making a 'pair' struct and manually packing and unpacking it, something like:

#At the top level
cdef struct float_pair:
    float a
    float b

# In the class method definition
cdef floar_pair calculate_edge_force(...
    ...
    cdef float_pair result
    result.a = ...
    result.b = ...
    return result


# In your function   
    cdef float_pair result = self.calculate_edge_force( ...
    Fe_plus = result.a
    Fe_plus_ = result.b

Using a struct is by no means the only option here, but it's going to be faster than using a small array. You could also use the C++ 'pair' template which would be just as fast as a struct but it involves bringing in C++. For that matter you could also pass in the floats by pointers. None of these are clearly superior in either speed or elegance, but all will compile down to clean c code.

I imagine if you root out those big ghoulies of python calls you'll get a significant speed improvement. After that it's a matter of using cython -a on the .pyx file. The generated html file shows 'unoptimized' lines in yellow/orange, you can click on a line to see the generated c code. From there you can figure out how to eliminate the python calls - if it is worth it.

Answer 2

When you have code like

cdef Type x, y, z
z = some_python_function(x, y)

Cython is going t create a C type and then convert it to a Python type. When that is done, it converts the answer back. Both conversion steps can be slow, and could be causing the 0.2 second slowdown you mention.

Looking at your code, it's hard to tell what types things are. One thing to note is that if

avg_PushMolecules_plus1 = (PushMolecules + PushMolecules_iPlus1)*0.5

is operating on Numpy arrays, typing these will not give you any speed improvements. Instead, you should look to either

• Unwrap the array expressions and type them, and compile using Cython

• Unwrap the array expressions and use Numba's autojit

• Use something like numexpr to optimize the calculation

Even just rewriting this to

avg_PushMolecules_plus1 = PushMolecules + PushMolecules_iPlus1
acg_PushMolecules_plus1 *= 0.5

may have a small positive effect, by removing intermediates.

Now, here is the important part:

You are doing no heavy looping inside of the function. As it stands, there is nothing in the function to speed up.

One thing worth doing is to use line_profiler to get times per line. This will show you which expressions need most to be optimized.

You mention numpy.max as a potential candidate. One way to optimize this call would be two write it in pure Cython:

cdef double Fe_pull_minus1, Fe_pull_plus1, Fe_pull
...
cdef double maximum = Fe_pull_minus1
if Fe_pull_plus1 > maximum: maximum = Fe_pull_plus1
if Fe_pull > maximum: maximum = Fe_pull

although it sounds dubious that such a call could have significant time implications — certainly not more than a few microseconds!

---

It's really unobvious where the slow parts are, becaue

• I don't know what, if anything, is an array,

• I don't know how often this is getting called, so I don't know what kind of optimizations are appropriate and

• I don't have timings.

You can solve all of these points of confusion by using line_profiler.

Install it and put at the top of your file:

import line_profiler
profiler = line_profiler.LineProfiler()

then before you run anything, write:

profiler.enable()
profiler.add_function(calculate_point_forces)

and finally, on script end, write:

profiler.print_stats()

The output should then include line-based timings for the function.