#!/net/python/bin/python import pylab as p from vector import * from math import log,sin,cos,sqrt,pi ################################################################## # Data Structure # # We will have a two dimensional array x_t with with x_t[t] referring # to the vector variable x of length n, where n is the degrees of # freedom in the system. x_t[t][i] is the value of the i-th variable # at t. For second order (eg, leapfrog integration), it will be assumed # that x[0]'=x[1], x[2]'=x[3], etc ################################################################### ################################################################### # Our library of vector functions # def Harmonic(x): """1-D Harmonic Oscillator Potential: U=-(r^2)/2 Takes an array x of length 2 and returns the function such that dx[i]/dt = f(x)[i]""" return vector([x[1],-x[0]]) def HarmonicEnergy(x): return (x[1]**2)/2 + sqrt(x[0]) def Toomre(x): """Toomre's Potential: U(r)=-(1+r^2)^-1/2. Takes an array x of length 4 and returns the function such that dx[i]/dt = f(x)[i]""" k=-(1+x[0]**2+x[2]**2)**-(3.0/2) return vector([x[1],x[0]*k,x[3],x[2]*k]) def ToomreEnergy(x): return (x[1]**2+x[3]**2)/2 - 1/sqrt(1+x[0]**2+x[2]**2) def Newton(x): """Newtonian Potential U(r)=1/r. Takes an array x of length 4 and returns the function such that dx[i]/dt = f(x)[i]""" k=-(x[0]**2+x[2]**2)**-(3.0/2) return vector([x[1],x[0]*k,x[3],x[2]*k]) def NewtonEnergy(x): return (x[1]**2+x[3]**2)/2 - 1/sqrt(x[0]**2+x[2]**2) def NonAxi(x,v0=1.0, rc=0.14, q=0.9): """Non-Axisymmetric potential U(x,y) = 0.5 v0^2*log(rc^2 + x^2 +(y/q)^2) Takes an array x of length 4 and returns the function such that dx[i]/dt = f(x)[i]""" k=-(v0**2)/(rc**2+x[0]**2+(x[2]/q)**2) return vector([x[1],k*x[0],x[3],k*x[2]/(q**2)]) def NonAxiEnergy(x, v0=1, rc=0.14, q=0.9): return (x[1]**2+x[3]**2)/2 - 0.5*v0**2*log(rc**2+x[0]**2+(x[2]/q)**2) def AngularMomentum(x): if len(x)==4: return x[0]*x[3]-x[2]*x[1] else: return 0 ################################################################### # Now for the knobs that make it all run... #function = one of the above functions #initial_conditions should have length matching input of function #Set potential to use: # 1 => Harmonic # 2 => Toomre # 3 => Newton # 4 => NonAxisymmetric potential = 3 initial_conditions = [1,0,0,0.5] #Set range and resolution of time values tstep=0.01 tmax=100 #Set method of numerical integration # 1 => Euler # 2 => Runge Kutta # 3 => LeapFrog method = 2 #Tell program which variables to graph # (range from 0 to DOF-1 xaxis,yaxis=0,1 ################################################################## #Define potential function and energy if potential==1: function=Harmonic energy=HarmonicEnergy elif potential==2: function=Toomre energy=ToomreEnergy elif potential==3: function=Newton energy=NewtonEnergy elif potential==4: function=NonAxi energy=NonAxiEnergy class NumInt: def __init__(self, DOF=len(initial_conditions)): self.x_t=vector([]) self.t=vector([]) self.DOF=DOF def EulerIntegrate(self, dt=1.0, t_final=7, x_0=initial_conditions,f=function): #check that Degrees of freedom are valid if len(x_0)!= self.DOF: print "Initial conditions must have same degrees of freedom as equations!" return False self.x_t=[vector(x_0)] self.t=vector([0]) t=0.0 x=vector(x_0) while t <= t_final: t+=dt new_x = x+dt*f(x) self.x_t.append(new_x) self.t.append(t) x = new_x def LeapfrogIntegrate(self,dt=1.0, t_final=7, x_0=initial_conditions,f=function): #check that Degrees of freedom are valid if len(x_0)!= self.DOF: print "Initial conditions must have same degrees of freedom as equations!" return False elif self.DOF%2==1: print "Leap frog invalid for odd Degrees of Freedom" self.x_t=[vector(x_0)] self.t=vector([0]) #first do a half-order Euler step for odd indices (x' variables) new_x=vector(x_0) for i in range(self.DOF/2): new_x[2*i+1] = x_0[2*i+1] + 0.5*dt*f(x_0)[2*i+1] x=new_x t=0.0 #now x is x0 for even indices, x1/2 for odd indices while t <= t_final: t+=dt new_x=x[:] #shallow copy - we don't want x to change for i in range(self.DOF/2): new_x[2*i+1] = x[2*i+1] + dt*f(x)[2*i+1] for i in range(self.DOF/2): new_x[2*i] = x[2*i] + dt*f(new_x)[2*i] self.x_t.append(new_x[:]) #change back half-stepped odd indices for i in range(self.DOF/2): j=2*i+1 self.x_t[-1][j]=0.5*(x[j]+new_x[j]) self.t.append(t) x = new_x #now undo the half-order Euler step for odd indices def RungeKuttaIntegrate(self, dt=1.0, t_final=7, x_0=initial_conditions,f=function): #check that Degrees of freedom are valid if len(x_0)!= self.DOF: print "Initial conditions must have same degrees of freedom as equations!" return False self.x_t=[vector(x_0)] self.t=vector([0]) t=0.0 x=vector(x_0) while t <= t_final: t+=dt k1 = dt*f(x) k2 = dt*f(x+k1/2) k3 = dt*f(x+k2/2) k4 = dt*f(x+k3) new_x = x + k1/6 + k2/3 + k3/3 + k4/6 self.x_t.append(new_x) self.t.append(t) x = new_x def grapht(self,val=[0]): """plot the variables of index in array (val) vs t""" for i in val: xi=[] for j in range(len(self.x_t)): xi.append(self.x_t[j][i]) p.plot(self.t,xi) p.ylabel('x[n]') p.xlabel('t') def graph(self,val0=0,val1=1,label=1): """plot a phase diagram of the variables of index in tuple (val)""" x0=[] for t in range(len(self.x_t)): x0.append(self.x_t[t][val0]) x1=[] for t in range(len(self.x_t)): x1.append(self.x_t[t][val1]) p.plot(x0,x1) if label: p.xlabel('x['+str(val0)+']') p.ylabel('x['+str(val1)+']') def graphE(self): """graph energies""" Energies=[] for x in self.x_t: Energies.append(energy(x)) p.plot(self.t, Energies) def graphL(self): """graph angular momenta""" Momenta=[] for x in self.x_t: Momenta.append(AngularMomentum(x)) p.plot(self.t, Momenta) #Now some functions used to display results in different formats: def graphphase(m=method, dt=tstep, tf=tmax, xa=xaxis, ya=yaxis, ic=initial_conditions): """ plots a phase portrait of xa,ya """ x=NumInt() if method==1: x.EulerIntegrate(dt,tf) elif method==2: x.RungeKuttaIntegrate(dt,tf) elif method==3: x.LeapfrogIntegrate(dt,tf) x.graph(xa,ya) string='x['+str(ya)+'] vs x['+str(xa)+'] (0 < t < '+str(x.t[-1])+ '):\nInitial Conditions: ' for i in range(x.DOF): string+='x['+str(i)+']='+str(ic[i])+' ' p.title(string) p.show() def graphinset(m=method, dt=tstep, tf=tmax, ic=initial_conditions): """ plots x-y alongside phase graphs of x-x' and y-y' """ x=NumInt() if method==1: x.EulerIntegrate(dt,tf) elif method==2: x.RungeKuttaIntegrate(dt,tf) elif method==3: x.LeapfrogIntegrate(dt,tf) p.axes([0.1,0.1,0.6,0.8]) x.graph(0,2,0) p.xlabel('x') p.ylabel('y') string='0 < t < '+str(x.t[-1])+ '\n' string+='Initial: (x,y)=('+str(ic[0])+','+str(ic[2])+'), ' string+="(x',y')= ("+str(ic[1])+','+str(ic[3])+')' p.title(string) a=p.axes([0.75,0.55,0.2,0.3],axisbg='30') x.graph(0,1,0) p.title('x phase') p.xlabel("x") p.ylabel("x'") p.setp(a, xticks=[], yticks=[]) b=p.axes([0.75,0.15,0.2,0.3],axisbg='30') x.graph(2,3,0) p.title('y phase') p.xlabel("y") p.ylabel("y'") p.setp(b, xticks=[], yticks=[]) p.show() def graphthree(dt=tstep, tf=tmax, ic=initial_conditions): """ for harmonic oscillator, plots results of all three integration schemes """ x=NumInt() y=NumInt() z=NumInt() x.EulerIntegrate(dt,tf) y.RungeKuttaIntegrate(dt,tf) z.LeapfrogIntegrate(dt,tf) x.graph(0,1) y.graph(0,1) z.graph(0,1) p.legend(['Euler','Runge Kutta','Leapfrog']) string='Phase Graph (0 < t < '+str(tf)+ ') [tstep=' +str(dt) + ']:\nInitially, ' string+='x='+str(ic[0])+", x'="+str(ic[1]) p.title(string) p.show() def harmonicphase(): """ plots an analytic phase portrait for a 1-D simple harmonic oscillator """ x=[] y=[] rt2=math.sqrt(2) for i in range(65): x.append(rt2*sin(i/10.0+pi/4)) y.append(rt2*cos(i/10.0+pi/4)) p.plot(x,y) p.title("Analytic Phase Graph (0