smolloy
smolloy

Reputation: 368

Speed up numerical integration in Cython

As part of a physics calculation, I load a list of electrical field values from a csv file, and pass them to a Cython function for integration. This is not a simple summation of the values in the file, but requires a recalculation of a certain parameter at each step of the integration. By reorganising the code, and moving the slowest parts to Cython, I have achieved a significant boost in speed, however I suspect there are further optimisations that could be done.

Here is the Cython function that I am trying to speed up:

import numpy as np
cimport numpy as np
DTYPE = np.float
ctypedef np.float_t DTYPE_t
cdef float m = 938272046
cdef float c = 299792458
cimport cython

@cython.boundscheck(False)
def ELTintegration(float xelmax, float deltaT, float phi, float omega, 
                np.ndarray[DTYPE_t] zExtract, np.ndarray[DTYPE_t] EzExtract, float KEin):
    cdef float integ = 0 
    cdef float xelmax_deltaT = xelmax * deltaT
    cdef unsigned int i, loopLen=len(EzExtract)
    cdef float v
    cdef float z, Ez, phaseTerm, diff
    cdef float newE, gamma, beta
    for i in range(len(EzExtract)):
        z = zExtract[i]
        Ez = EzExtract[i]
        newE = KEin + integ
        gamma = newE/m + 1 
        beta = (1 - 1/gamma**2)**0.5
        v = beta * c 
        phaseTerm = np.cos(omega*z/v + phi)
        diff = phaseTerm * Ez * xelmax_deltaT
        integ = integ + diff

    return integ

Each of the ndarrays is one dimensional, and contains ~4000 floats.

As you can see, I have declared all variables I use, switched off bounds checking, made the array indexing efficient, and pulled all necessary calculations out of the loop.

Have I reached the limit of what I can reasonably do to speed this up? Is there anything I have missed in terms of achieving a more efficient function? Is it realistic to think I could speed this up by another factor of 2?

Thanks for reading.

Upvotes: 1

Views: 1490

Answers (1)

Veedrac
Veedrac

Reputation: 60137

Tim Pietzcker is right that it's better suited for Code Review but it's also almost on-topic here, so I'll go for it.

If you run

cython -a myfile.pyx

You will see yellow. Often, but not always, yellow means slow. The most important yellow is with the np.cos call. This will be extremely slow because you're going through Python and doing lots of type creation to get it done.

Using

from libc.math cimport cos

will let you use C's cos, which will be much faster.

There are also some checks for division by zero and the sign of division. These are avoidable by turning the compiler directive cdivision to True.

Note that the modern way of using Numpy arrays is with the memoryview syntax (eg. DTYPE_t[:] for a 1D memoryview of type DTYPE_t).

Also note that you might be as good off or even better off using numba and autojit on the plain Python code. numba does crazy things.

Upvotes: 5

Related Questions