# Improving Performance of Numba For-Loop and Matrix Multiplication for Time Integration Solver

I am writing some code that performs a simple time-history integration solver (Newmark method for structural dynamics) and I looking for some suggestions in how I can improve the performance of my code. Currently, I am calculating a parameter called displacements for many time steps (think on the order of 5,000,000 steps). From profiling the code without using numba it is apparent that the matrix multiplication seems to be slowing down the script in the for-loop. Some details about the input:

num_steps = number of time-steps to perform the calculation (5,000,000)

dofs = number of degrees of freedom in the system, this also defines the matrix sizes as (dofs, dofs) before this function is called

yF = vector of size (num_steps, 1)

k_hat_inv, a1, a2, a3 = are all square, symmetric matrices of size (dofs, dofs)

vel_const1, vel_const2, vel_const3, acc_const1, acc_const2, acc_const3 = are constant floats

disp_save, disp_steps = are all ints

For a simple example when dofs = 10 this code runs very quickly (~14s) but as dofs increases to 1280 the run time increases significantly (still running after 1.5 days) and I would like to find some ways to decrease this run time. I have tried running the algorithm using parallel = True and the prange function in Numba but this doesn't provide the correct solution. Are there any other ways that I can improve the performance of this code?

@jit(nopython=True, cache=True)
def newmark_looper(num_steps, dofs, yF, k_hat_inv, a1, a2, a3, vel_const1, vel_const2, vel_const3,
acc_const1, acc_const2, acc_const3, disp_save, disp_steps):

# Time-stepping solution as per Chopra, 2017
print('     Time History Started')

# Initialize some variables
cdisp = np.zeros(dofs-2)
cvel = np.zeros(dofs-2)
caccel = np.zeros(dofs-2)
disp_counter = 0
displacements = np.zeros((dofs-2, disp_steps))

for count in range(num_steps):

# Obtain the force for the next time step, i.e i+1

# Obtain the force for the next time step, i.e i+1

# Obtain P_i_p1_hat

# Calculate the displacement due to this force, u_i+1 = k_hat_inv * P_i_p1
u_i_p1 = k_hat_inv @ P_i_p1_hat

# Calculate the velocity due to this displacement, u_dot_i+1
udot_i_p1 = vel_const1*(u_i_p1-cdisp) + vel_const2*cvel + vel_const3*caccel

# Calculate the acceleration due to this displacement, u_2dot_i+1
u2dot_i_p1 = acc_const1*(u_i_p1-cdisp) - acc_const2*cvel - acc_const3*caccel

# Get current values for the next step
cdisp[:] = u_i_p1
cvel[:] = udot_i_p1
caccel[:] = u2dot_i_p1

# Write the displacement
if count % disp_save == 0:
disp_counter += 1
displacements[:, disp_counter] = u_i_p1

print('     Time History Analyzed')

return displacements