Optimizing time in python

[Simplicity.]

Hi all,

This code solves the Euler Equation of Gas Dynamics which is a system of partial differential equations. The problem is that it takes very long to run than it would normally required. Things become worse after including the functions "S_p(a_p,a,b)" and "S_m(a_p,a,b)" which are called within other functions to add some viscosity to avoid oscillations in the solution, making the solution of high order in accuracy.To execute a single iteration would take more than a day and there are hundred of them. Without these functions, the program would run for about 6 hours and give results, which are first order accurate.

If anyone can help how to reformulate loops or do something else to optimize time, that will be greatly appreciated!

"""  """
from numpy import *
from scitools.std import *
from math import *

N = 100
M = 200; 
dx = 1.0/M    
cfl = 0.75
 
a_3 =   5.045; a_2 =1.68; a_1 = 1

dt = cfl*dx/sqrt(a_3)

f_rho = zeros(M); f_m = zeros(M); f_E = zeros(M);u = zeros(M);rho = zeros(M)

f_mss = zeros(M) ; f_Ess = zeros(M); p = zeros(M)

Gamma = 1.4 ; Eps = 10e-8
K  = linspace(0.0,1.0, M+1)

#Initial values

for i in range(M):			
		if K[i]>= 0 and K[i] <= 0.5:  u[i] = 0.0
		elif K[i] > 0.5 and K[i] <= 1 : u[i] = 0.0

def Initialise(u):
	
   	return u



for i in range(M):			
		if K[i]>= 0 and K[i] <= 0.5:  p[i] = 1.0
		elif K[i] > 0.5 and K[i] <= 1 : p[i] = 0.1	

def Initialization(p):
	
	return p


#Conserved Variables

for i in range(M):			
		if K[i]>=0 and K[i] <= 0.5:  rho[i] = 1.0
		elif K[i] > 0.5 and K[i] <= 1 : rho[i] = 0.125
		
def G(rho):

	return rho

m = G(rho)*Initialise(u)

def momentum(m):

	return m

E = Initialization(p)/(Gamma-1) + 0.5*(G(rho))*(momentum(m)/G(rho))**2

def Energy(E):

	return E


#Conserved fluxes

def flux_rho(f_rho):
   	
	f_rho = momentum(m)

	return f_rho

def flux_m(f_m):

	f_m = (momentum(m))**2/(G(rho)) + (Gamma -1)*(Energy(E)- (momentum(m))**2/(2*G(rho)))
	return f_m
 

def flux_E(f_E):

	
	f_E = (Energy(E) + (Gamma -1.0)*(Energy(E)- (momentum(m))**2/(2*G(rho))))*(momentum(m)/G(rho))
	return f_E


#v -equilibrium values

v_rho  = momentum(m)

def rho_v(v_rho):
	
	return v_rho

v_m = (momentum(m))**2/(G(rho)) + (Gamma -1)*(Energy(E)- (momentum(m))**2/(2*G(rho)))
def m_v(v_m):
	
	return v_m



v_E = (Energy(E) + (Gamma -1.0)*(Energy(E)- (momentum(m))**2/(2*G(rho))))*(momentum(m)/G(rho))
def E_v(v_E):
	
	return v_E

#Limiters to add some viscosity

def S_p(a_p,a,b):
	Sigma_p = zeros((M,), float32); Theta_p = zeros((M,), float32);	Theta_m =  zeros((M,), float32)
	for k in xrange (1,M-1):
		
		Theta_p[k]  = (a[k+1] + sqrt(a_p)*b[k+1] - a[k] - sqrt(a_p)*b[k])
		Theta_m[k] = (a[k] + sqrt(a_p)*b[k]- a[k-1] - sqrt(a_p)*b[k-1])
		
		#Minmod Slope limiter for positive theta
		if sign(Theta_p[k]) == sign(Theta_m[k]): x_p = sign(Theta_p[k])*min(abs(Theta_p[k]), abs(Theta_m[k]))
		else:
			x_p = 0
		Sigma_p[k] = (1/dx)*x_p	
		
		Sigma_p[0] = Sigma_p[1]
		Sigma_p[M-1] = Sigma_p[M-2]
	return Sigma_p

def S_m(a_p,a,b):

	
	Theta_p = zeros((M,), float32) ; Theta_m = zeros((M,), float32); Sigma_m = zeros((M,), float32)
	for k in xrange (1,M-1):
		

		Theta_p[k]  = (a[k+1] - sqrt(a_p)*b[k+1] - a[k] + sqrt(a_p)*b[k])
		Theta_m[k] = (a[k] - sqrt(a_p)*b[k] - a[k-1] + sqrt(a_p)*b[k-1])
		
		#Minmod Slope limiter for negative theta
		if sign(Theta_p[k]) == sign(Theta_m[k]): x_m = sign(Theta_p[k])*min(abs(Theta_p[k]), abs(Theta_m[k]))
		else:
			x_m = 0
		Sigma_m[k] = (1/dx)*x_m
		
		Sigma_m[0] = Sigma_m[1]
		Sigma_m[M-1] = Sigma_m[M-2]
	return Sigma_m


#Solve for v* of density 
def Evolve_rho_s(rho):
		
	rho_s =	Eps/(Eps -dt)*(rho_v(v_rho) - (dt/Eps)*flux_rho(f_rho))
	
	return rho_s



#Solve for u(1) of sensity

def Evolve_rho1(rho):
	
	
	rho1 = zeros((M,), float32)
	
	for k in range (1, M-1):	
		
		rho1[k] =  rho[k]-0.5*(dt/dx)*(Evolve_rho_s(rho)[k+1] - Evolve_rho_s(rho)[k-1]) + 0.5*(dt/dx)*sqrt(a_1)*(rho[k+1] - 2*rho[k] + rho[k-1])+ 0.25*(dt/dx)*dx*(S_m(a_1, Evolve_rho_s(rho),rho)[k+1] -(S_p(a_1, Evolve_rho_s(rho),rho)[k] + S_m(a_1,Evolve_rho_s(rho),rho)[k]) + S_p(a_1,Evolve_rho_s(rho),rho)[k-1])
		rho1[0] = rho1[1]
		rho1[M-1] = rho1[M-2]	
	return rho1


#Solve for v(1) of density

def Evolve_rho_v(rho):
		
	rho_v = zeros((M,), float32)
	for k in xrange (1,M-1):
	
		rho_v[k] = Evolve_rho_s(rho)[k]- 0.5*(dt/dx)*a_1*(rho[k+1] - rho[k-1]) + 0.5*(dt/dx)*(a_1/sqrt(a_1))*(Evolve_rho_s(rho)[k+1] - 2*Evolve_rho_s(rho)[k] + Evolve_rho_s(rho)[k-1])- 0.25*(dt/dx)*(a_1/sqrt(a_1))*dx*(S_m(a_1,Evolve_rho_s(rho),rho)[k+1] + (S_p(a_1,Evolve_rho_s(rho),rho)[k] - S_m(a_1,Evolve_rho_s(rho),rho)[k]) - S_p(a_1,Evolve_rho_s(rho),rho)[k-1])

		rho_v[0] = rho_v[1]
		rho_v[M-1] = rho_v[M-2]
	return rho_v

#Calculate m*

def Evolve_ms(m):
	ms =  Eps/(Eps -dt)*(m_v(v_m) - (dt/Eps)*(flux_m(f_m)))

	return ms

#u(1) for m 

#m1 = f(u**) for density

def Evolve_m1(m):
	m1 = zeros((M,), float32)	
	for k in xrange (1,M-1):
		m1[k] =  m[k]-0.5*(dt/dx)*(Evolve_ms(m)[k+1] - Evolve_ms(m)[k-1]) + 0.5*(dt/dx)*sqrt(a_2)*(m[k+1] - 2*m[k] + m[k-1]) + 0.25*(dt/dx)*dx*(S_m(a_2,Evolve_ms(m),m)[k+1] - (S_p(a_2,Evolve_ms(m),m)[k] + S_m(a_2,Evolve_ms(m),m)[k]) + S_p(a_2,Evolve_ms(m),m)[k-1])
		m1[0] = m1[1]
		m1[M-1] = m1[M-2]
	return m1

#calculate v(1) for m

def Evolve_mv(m):
	mv = zeros((M,), float32)
		
	for k in range (1,M-1):

		mv[k] = Evolve_ms(m)[k]- 0.5*(dt/dx)*a_2*(m[k+1] - m[k-1]) + 0.5*(dt/dx)*(a_2/sqrt(a_2))*(Evolve_ms(m)[k+1] - 2*Evolve_ms(m)[k] + Evolve_ms(m)[k-1]) - 0.25*(a_2/sqrt(a_2))*(dt/dx)*dx*(S_m(a_2,Evolve_ms(m),m)[k+1] + (S_p(a_2,Evolve_ms(m),m)[k] - S_m(a_2,Evolve_ms(m),m)[k]) - S_p(a_2,Evolve_ms(m),m)[k-1])
		mv[0] = mv[1]
		mv[M-1] = mv[M-2]
	return mv


#rho** = rho(1)

#Calc v** for rho
def Evolve_rho_ss(rho):
		
	rho_ss = Eps/(Eps + dt)*(Evolve_rho_v(rho) + (dt/Eps)*(Evolve_m1(m))) - 2*(dt/(Eps+dt))*( Evolve_rho_s(rho) - flux_rho(f_rho))

	return rho_ss


#u** = u(1)
# Solve for u(2) of density


def Evolve_rho2(rho):

	rho_v2 = zeros((M,), float32)
	for k in range (1,M-1):
			
		rho_v2[k] = Evolve_rho1(rho)[k]- 0.5*(dt/dx)*(Evolve_rho_ss(rho)[k+1] - Evolve_rho_ss(rho)[k-1]) + 0.5*(dt/dx)*sqrt(a_1)*(Evolve_rho1(rho)[k+1] - 2*Evolve_rho1(rho)[k] + Evolve_rho1(rho)[k-1]) + 0.25*(dt/dx)*dx*(S_m(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k+1] - (S_p(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k] + S_m(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k]) + S_p(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k-1])
	
		rho_v2[0] = rho_v2[1]
		rho_v2[M-1] = rho_v2[M-2]
	return rho_v2

# Solve for v(2) of density

def Evolve_rho_v2(rho):

	rho_v2 = zeros((M,), float32)
	for k in range (1,M-1):

		rho_v2[k] = Evolve_rho_ss(rho)[k]- 0.5*(dt/dx)*a_1*(Evolve_rho1(rho)[k+1] - Evolve_rho1(rho)[k-1]) + 0.5*(dt/dx)*(a_1/sqrt(a_1))*(Evolve_rho_ss(rho)[k+1] - 2*Evolve_rho_ss(rho)[k] + Evolve_rho_ss(rho)[k-1]) - 0.25*(dt/dx)*(a_1/sqrt(a_1))*dx*(S_m(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k+1] + (S_p(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k] - S_m(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k]) - S_p(a_1, Evolve_rho_ss(rho), Evolve_rho1(rho))[k-1])

		rho_v2[0] = rho_v2[1]
		rho_v2[M-1] = rho_v2[M-2]
	return rho_v2

#Find v* for E
def Evolve_Es(E):
	
	E_s =  Eps/(Eps -dt)*(E_v(v_E) - (dt/Eps)*(flux_E(f_E)))
		
	return E_s

#Solve u(1) for E

def Evolve_E1(E):
	
	E1 = zeros((M,), float32)

	for k in range (1,M-1):	

		E1[k] =  E[k]-0.5*(dt/dx)*(Evolve_Es(E)[k+1] - Evolve_Es(E)[k-1]) + 0.5*(dt/dx)*sqrt(a_3)*(E[k+1] - 2*E[k] + E[k-1]) + 0.25*(dt/dx)*dx*(S_m(a_3, Evolve_Es(E), E)[k+1] - (S_p(a_3, Evolve_Es(E), E)[k] + S_m(a_3, Evolve_Es(E), E)[k]) + S_p(a_3, Evolve_Es(E), E)[k-1])
		E1[0] = E1[1]
		E1[M-1] = E1[M-2]	
		
	return E1

#Calculate v(1) for E, EV1

def Evolve_Ev(E):

	Ev = zeros((M,), float32)
	
	for k in range (1,M-1):
	
		Ev[k] = Evolve_Es(E)[k] - 0.5*(dt/dx)*a_3*(E[k+1] - E[k-1]) + 0.5*(dt/dx)*(a_3/sqrt(a_3))*(Evolve_Es(E)[k+1] - 2*Evolve_Es(E)[k] + Evolve_Es(E)[k-1]) - 0.25*(a_3/sqrt(a_3))*(dt/dx)*dx*(S_m(a_3, Evolve_Es(E), E)[k+1] + (S_p(a_3, Evolve_Es(E), E)[k] - S_m(a_3, Evolve_Es(E), E)[k]) - S_p(a_3, Evolve_Es(E), E)[k-1])

		Ev[0] = Ev[1]
		Ev[M-1] = Ev[M-2]
	return Ev

#Lets calculate f(u**) for m

def flux_m_ss(f_mss):

	f_mss = (Evolve_m1(m))**2/(Evolve_rho1(rho)) + (Gamma -1)*(Evolve_E1(E) - (Evolve_m1(m))**2/(2*Evolve_rho1(rho)))
	
	return f_mss

#Solve for v** for m

def Evolve_mss(m):
	mss =  Eps/(Eps + dt)*(Evolve_mv(m) + (dt/Eps)*(flux_m_ss(f_mss))) - 2*(dt/(Eps + dt))*(Evolve_ms(m) - flux_m(f_m))

	return mss

#Now calculate u(2) for m
def Evolve_m2(m):

	m2 = zeros((M,), float32)
	
	for k in range (1,M-1):

		m2[k] =  Evolve_m1(m)[k]-0.5*(dt/dx)*(Evolve_mss(m)[k+1] - Evolve_mss(m)[k-1]) + 0.5*(dt/dx)*sqrt(a_2)*(Evolve_m1(m)[k+1] - 2*Evolve_m1(m)[k] + Evolve_m1(m)[k-1]) + 0.25*(dt/dx)*dx*(S_m(a_2, Evolve_mss(m), Evolve_m1(m))[k+1] - (S_p(a_2, Evolve_mss(m), Evolve_m1(m))[k] + S_m(a_2, Evolve_mss(m), Evolve_m1(m))[k]) + S_p(a_2, Evolve_mss(m), Evolve_m1(m))[k-1]) 
		m2[0] = m2[1]
		m2[M-1] = m2[M-2]
	return m2

#Calculate v(2) for m
def Evolve_mv2(m):
		
	mv2 = zeros((M,), float32)
		
	for k in range (1,M-1):

		mv2[k] = Evolve_mss(m)[k]- 0.5*(dt/dx)*a_2*(Evolve_m1(m)[k+1] - Evolve_m1(m)[k-1]) + 0.5*(dt/dx)*(a_2/sqrt(a_2))*(Evolve_mss(m)[k+1] - 2*Evolve_mss(m)[k] + Evolve_mss(m)[k-1]) - 0.25*(a_2/sqrt(a_2))*(dt/dx)*dx*(S_m(a_2, Evolve_mss(m), Evolve_m1(m))[k+1] + (S_p(a_2, Evolve_mss(m), Evolve_m1(m))[k] - S_m(a_2, Evolve_mss(m), Evolve_m1(m))[k]) - S_p(a_2, Evolve_mss(m), Evolve_m1(m))[k-1])
		mv2[0] = mv2[1]
		mv2[M-1] = mv2[M-2]
	return mv2

#Solve for f(u**) of E
def flux_E_ss(f_Ess):

	f_Ess = (Evolve_E1(E) + (Gamma -1.0)*(Evolve_E1(E)- (Evolve_m1(m))**2/(2*Evolve_rho1(rho))))*(Evolve_m1(m)/Evolve_rho1(rho))
	
	return f_Ess

#Now, calculate v** for E
def Evolve_Ess(E):
	
	E_ss =  Eps/(Eps + dt)*(Evolve_Ev(E) + (dt/Eps)*(flux_E_ss(f_Ess))) - 2*(dt/(Eps + dt))*(Evolve_Es(E) - flux_E(f_E))
			
	return E_ss

#Calculate u(2) for E
def Evolve_E2(E):

	E2 = zeros((M,), float32)

	for k in range (1,M-1):	

		E2[k] =  Evolve_E1(E)[k]-0.5*(dt/dx)*(Evolve_Ess(E)[k+1] - Evolve_Ess(E)[k-1]) + 0.5*(dt/dx)*sqrt(a_3)*(Evolve_E1(E)[k+1] - 2*Evolve_E1(E)[k] + Evolve_E1(E)[k-1]) + 0.25*(dt/dx)*dx*(S_m(a_3, Evolve_Ess(E), Evolve_E1(E))[k+1] - (S_p(a_3, Evolve_Ess(E), Evolve_E1(E))[k] + S_m(a_3, Evolve_Ess(E), Evolve_E1(E))[k]) + S_p(a_3, Evolve_Ess(E), Evolve_E1(E))[k-1])

		E2[0] = E2[1]
		E2[M-1] = E2[M-2]	
		
	return E2

#Calculate v(2) for E
def Evolve_Ev2(E):
	
	Ev2 = zeros((M,), float32)
	
	for k in range (1,M-1):

		Ev2[k] = Evolve_Ess(E)[k] - 0.5*(dt/dx)*a_3*(Evolve_E1(E)[k+1] - Evolve_E1(E)[k-1]) + 0.5*(dt/dx)*(a_3/sqrt(a_3))*(Evolve_Ess(E)[k+1] - 2*Evolve_Ess(E)[k] + Evolve_Ess(E)[k-1]) - 0.25*(a_3/sqrt(a_3))*(dt/dx)*dx*(S_m(a_3, Evolve_Ess(E), Evolve_E1(E))[k+1] + (S_p(a_3, Evolve_Ess(E), Evolve_E1(E))[k] - S_m(a_3, Evolve_Ess(E), Evolve_E1(E))[k]) - S_p(a_3, Evolve_Ess(E), Evolve_E1(E))[k-1])

		Ev2[0] = Ev2[1]
		Ev2[M-1] = Ev2[M-2]
	return Ev2


if __name__=="__main__":
	
	
	for n in xrange(N):
		
	
		plot(linspace(0.0,1.0,M),rho,'r0.2')


		rho = 0.5*( rho + Evolve_rho2(rho))
			
		v_rho = 0.5*(v_rho + Evolve_rho_v2(rho))

		m  = 0.5*( m  + Evolve_m2(m))
		v_m    = 0.5*(v_m + Evolve_mv2(m))

		E =   0.5*( E + Evolve_E2(E))
		v_E  =  0.5*( v_E + Evolve_Ev2(E))


xlabel("x")
ylabel("density")
title("Euler eqn of gas dynamics, second order in time and space")
hardcopy('modified_density1D2ndorder.png')
show()

[Gribouillis]

As a first hint, look at line 248 for example. Evolve_rho_ss(rho) is called 8 times in the same expression. I assume that all these calls involve arrays of real numbers. Obviously you should call the function only once temp = Evolve_rho_ss(rho) and then use the variable temp 8 times in the expression if this is needed. The call could perhaps even be done before the for loop.

The same problem arises in many of your functions for many calls. So The first thing to do is to remove all the calls which compute exactly the same quantity in each function and store the result of each single call in a local variable. This may speed up your program considerably.

Try to give descriptive names to these local variables (instead of 'temp').

[Simplicity.]

Thank you very much Gribouillis for your hint, this really solved the problem by increasing speed considerably!

With great appreciation,

Simplicity.