from pylab import *
from scipy.integrate import odeint
from scipy.interpolate import interp1d
from numpy import *
def rhs(state,t):
global prevx
global prevy
global prevvx
global prevvy
global prevt
global poincaresection
x=state[0]
y=state[1]
vx=state[2]
vy=state[3]
#Our EOM's
A=(1-sigma)/sigma
vx_dot=-x/sigma+A*x/sqrt(pow(x,2)+pow(1-y,2))
vy_dot=A-y/sigma-A*(1-y)/sqrt(pow(x,2)+pow(1-y,2))
#The code below calculates the point at which our trajectory has crossed
# our poincare surface
#We only want to consider points where vy>0
if(vy>0.0):
#Our poincare surface is at y=0, so if y is now >0 and it was <0
# in the last timestep (prevy<0) then we know that sometime between
# the last timestep and now the solution had to cross the y=0 surface
if y>0.0 and (prevy<0.0):
# To figure out exactly when it crossed y=0 we are going to use the interp1d
# function built into python.
interpx=[prevx, x]
interpy=[prevy, y]
interpvx=[prevvx, vx]
interpvy=[prevvy, vy]
interpt=[prevt,t]
# We interpolate fy with y on the x axis and time on the y axis
# This is so we can set y=0 and find out what time that was at
# We interpolate all the other variables the opposite way, that is
# so we can find out what value they took at the time we got from the y
# interpolation.
fy=interp1d(interpy,interpt,kind='linear')
fx=interp1d(interpt,interpx,kind='linear')
fvx=interp1d(interpt,interpvx,kind='linear')
fvy=interp1d(interpt,interpvy,kind='linear')
#We figure out what time the solution crossed the y=0.0 surface
crossingtime=fy(0.0)
#Then figure out what all the other variables were at that time
poincarepointvx=fvx(crossingtime)
poincarepointx=fx(crossingtime)
poincarepointvy=fvy(crossingtime)
#We are interested in an x vs vx plot, so we add that point to our poincare section array
poincaresection=vstack((poincaresection,array([poincarepointx,poincarepointvx])))
prevx=x
prevy=y
prevvy=vy
prevvx=vx
prevt=t
return [vx,vy, vx_dot, vy_dot]
# Define the starttime, endtime and number of points
start=.0
end=17000.0
num=8000
# generate a list of times to output the solution at,
# "num" of them spaced equally going from "start", to "end"
t=linspace(start,end,num)
#Set sigma and the energy
sigma=0.25
E=1.0
# The poincare section is a collection of points, so lets define an empty array
# for that
poincaresection=array([0,0])
#My initial conditions, note I have to calculate what vy0 is from the energy (E)
x0=-0.75
vx0=1
y0=0.2
vy0=sqrt(2*E-2*y0-pow(vx0,2)-(1/sigma)*pow(sqrt(pow(x0,2)+pow(1-y0,2))-1+sigma,2))
print vy0
#The prev? variable is what x,y,vx,vy were on the previous timestep. We are going to use
# this to figure out when the solution has crossed our poincare surface
prevx=x0
prevy=y0
prevvx=vx0
prevvy=vy0
#Call our odeint routine
solution=odeint(rhs,[x0,y0,vx0,vy0],t)
figure()
plot(poincaresection[:,0],poincaresection[:,1],'o',markersize=.5)
xlabel('x')
ylabel('vx')
title('E='+str(E)+', sigma='+str(sigma))
savefig('poincare.pdf')