How to speed up python code with many variable assignments-CodePudding

I´m working on the following piece of code which creates a numpy array and uses many variable assignments to do that. The code was generated by Maple with optimization which results in lots of variable assignments. The code results from very complex mathematics and there is no further simplification possible with Maple.

    from math import sin, cos
    import numpy as np

    def __calc_tau(self, q, s_dot_v, omega_dot_ref, f_th, qm_ddot_ref):
        SC_param_est = self.SC_param_est
        SC_param = self.SC_param

        t1 = SC_param[3]   SC_param_est[0]   SC_param[10]   SC_param[9]   SC_param[8]
        t2 = cos(q[9])
        t3 = sin(q[9])
        t4 = SC_param[38] - SC_param[41]
        t5 = SC_param[35] - SC_param[38]   SC_param[40]
        t6 = t2 * SC_param_est[7]
        t7 = t3 * SC_param_est[9]
        t8 = t5 * SC_param[10]
        t9 = (-t6   t4 - t7) * SC_param_est[0]
        t10 = SC_param[32] * SC_param[9]
        ...
        t638 = ...
        MCQ = np.array([[t1, 0, 0, 0, t37, t54, t54, -t127, t61, -t128, 0, -t126, t125, 0, -t91, -t106, -t106,
                         -t45 * t108 * s_dot_v[6] - t46 * t91, -t114 * t45 - t16 * t46,
                         -(t123 * t45   t46 * t78) * SC_param_est[0], SC_param[7] * t124 * t1],
                        [0, t1, 0, -t31, 0, -t51, -t51, -t56, -t38, -t68, t126, 0, -t134, t129, 0, t131, t131,
                         t46 * t108 * s_dot_v[6] - t45 * t91, t114 * t46 - t16 * t45,
                         -SC_param_est[0] * (-t123 * t46   t45 * t78), -t135 * t133 * t1], t86, t34, t36, t35,
                        [t54, -t51, 0, -t300, -t58, t89, t89, t115, t121, -t69, -t106, t131, 0, t450, t43, t55,
                         t55, t170, t76, t239, -t47], t10, t5, t2], dtype='float64')

I measured the execution time of the parts of with timeit my program and this seems to be the bottleneck with an execution time of approx. 3ms (Intel i7-8565U @ 1.8 GHz).

Is there a proper way to speed this up?

CodePudding user response：

The many assignments are not the problem, but what is calculated in each case and then assigned to the variable.

From what I can see, there are a lot of mathematical functions like sin, cos and dot products. I think these functions may be your bottleneck. Maybe you could try to reduce the number of these operations to speed up your function.

CodePudding user response：

After some further research and testing I found a solution. Using the numba package the following code worked for me:

from numba import jit

@jit
def calc_MCQ_numba(SC_param, SC_param_est, q, s_dot_v):
    t1 = SC_param[3]   SC_param_est[0]   SC_param[10]   SC_param[9]   
            SC_param[8]
    t2 = np.cos(q[9])
    t3 = np.sin(q[9])
    t4 = SC_param[38] - SC_param[41]
    t5 = SC_param[35] - SC_param[38]   SC_param[40]
    t6 = t2 * SC_param_est[7]
    t7 = t3 * SC_param_est[9]
    ...
    return np.array([[t1, 0, 0, 0, t37, t54, t54, -t127, t61, -t128, 0, -t126, t125, 0, -t91, -t106, -t106,
                  -t45 * t108 * s_dot_v[6] - t46 * t91, -t114 * t45 - t16 * t46,
                  -(t123 * t45   t46 * t78) * SC_param_est[0], SC_param[7] * t124 * t1],
                 [0, t1, 0, -t31, 0, -t51, -t51, -t56, -t38, -t68, t126, 0, -t134, t129, 0, t131, t131,
                  t46 * t108 * s_dot_v[6] - t45 * t91, t114 * t46 - t16 * t45,
                  -SC_param_est[0] * (-t123 * t46   t45 * t78), -t135 * t133 * t1], t86, t34, t36, t35,
                 [t54, -t51, 0, -t300, -t58, t89, t89, t115, t121, -t69, -t106, t131, 0, t450, t43, t55,
                  t55, t170, t76, t239, -t47], t10, t5, t2], dtype='float64')

With this function the first time executing the code takes very long (approx. 25s on my machine) but each time the function gets called again is very fast. I managed to get an execution time of 10µs - which is stunning compared to my first result.