Poincare Section of a system of second order odes

It is the first time I am trying to write a Poincare section code at Python.

I borrowed the piece of code from here:

https://github.com/williamgilpin/rk4/blob/master/rk4_demo.py

and I have tried to run it for my system of second order coupled odes. The problem is that I do not see what I was expecting to. Actually, I need the Poincare section when x=0 and px>0.

I believe that my implementation is not the best out there. I would like to:

Improve the way that the initial conditions are chosen.
Apply the correct conditions (x=0 and px>0) in order to acquire the correct Poincare section.
Create one plot with all the collected poincare section data, not four separate ones.

I would appreciate any help.

This is the code:

from matplotlib.pyplot import *
from scipy import *
from numpy import *

# a simple Runge-Kutta integrator for multiple dependent variables and one independent variable

def rungekutta4(yprime, time, y0):
    # yprime is a list of functions, y0 is a list of initial values of y
    # time is a list of t-values at which solutions are computed
    #
    # Dependency: numpy

    N = len(time)

    y = array([thing*ones(N) for thing in y0]).T

    for ii in xrange(N-1):
        dt = time[ii+1] - time[ii]
        k1 = dt*yprime(y[ii], time[ii])
        k2 = dt*yprime(y[ii] + 0.5*k1, time[ii] + 0.5*dt)
        k3 = dt*yprime(y[ii] + 0.5*k2, time[ii] + 0.5*dt)
        k4 = dt*yprime(y[ii] + k3, time[ii+1])
        y[ii+1] = y[ii] + (k1 + 2.0*(k2 + k3) + k4)/6.0

    return y

# Miscellaneous functions
n= 1.0/3.0
kappa1 = 0.1
kappa2 = 0.1
kappa3 = 0.1
def total_energy(valpair):
    (x, y, px, py) = tuple(valpair)
    return .5*(px**2 + py**2) + (1.0/(1.0*(n+1)))*(kappa1*np.absolute(x)**(n+1)+kappa2*np.absolute(y-x)**(n+1)+kappa3*np.absolute(y)**(n+1))

def pqdot(valpair, tval):
    # input: [x, y, px, py], t
    # takes a pair of x and y values and returns \dot{p} according to the Hamiltonian
    (x, y, px, py) = tuple(valpair)
    return np.array([px, py, -kappa1*np.sign(x)*np.absolute(x)**n+kappa2*np.sign(y-x)*np.absolute(y-x)**n, kappa2*np.sign(y-x)*np.absolute(y-x)**n-kappa3*np.sign(y)*np.absolute(y)**n]).T

def findcrossings(data, data1):
    # returns indices in 1D data set where the data crossed zero. Useful for generating Poincare map at 0
    prb = list()
    for ii in xrange(len(data)-1):
        if (((data[ii] > 0) and (data[ii+1] < 0)) or ((data[ii] < 0) and (data[ii+1] > 0))) and data1[ii] > 0:
            prb.append(ii)
    return array(prb)

t = linspace(0, 1000.0, 100000)
print ("step size is " + str(t[1]-t[0]))

# Representative initial conditions for E=1
E = 1
x0=0
y0=0
init_cons = [[x0, y0, np.sqrt(2*E-(1.0*i/10.0)*(1.0*i/10.0)-2.0/(n+1)*(kappa1*np.absolute(x0)**(n+1)+kappa2*np.absolute(y0-x0)**(n+1)+kappa3*np.absolute(y0)**(n+1))), 1.0*i/10.0] for i in range(-10,11)]

outs = list()
for con in init_cons:
    outs.append( rungekutta4(pqdot, t, con) )


# plot the results
fig1 = figure(1)
for ii in xrange(4):
    subplot(2, 2, ii+1)
    plot(outs[ii][:,1],outs[ii][:,3])
    ylabel("py")
    xlabel("y")
    title("Full trajectory projected onto the plane")

fig1.suptitle('Full trajectories E = 1', fontsize=10)


# Plot Poincare sections at x=0 and px>0
fig2 = figure(2)
for ii in xrange(4):
    subplot(2, 2, ii+1)
    xcrossings = findcrossings(outs[ii][:,0], outs[ii][:,3])
    yints = [.5*(outs[ii][cross, 1] + outs[ii][cross+1, 1]) for cross in xcrossings]
    pyints = [.5*(outs[ii][cross, 3] + outs[ii][cross+1, 3]) for cross in xcrossings]
    plot(yints, pyints,'.')
    ylabel("py")
    xlabel("y")
    title("Poincare section x = 0")

fig2.suptitle('Poincare Sections E = 1', fontsize=10)

show()

You need to compute the derivatives of the Hamiltonian correctly. The derivative of |y-x|^n for x is

n*(x-y)*|x-y|^(n-2)=n*sign(x-y)*|x-y|^(n-1)

and the derivative for y is almost, but not exactly (as in your code), the same,

n*(y-x)*|x-y|^(n-2)=n*sign(y-x)*|x-y|^(n-1),

note the sign difference. With this correction you can take larger time steps, with correct linear interpolation probably even larger ones, to obtain the images

I changed the integration of the ODE to

t = linspace(0, 1000.0, 2000+1)
...
E_kin = E-total_energy([x0,y0,0,0])
init_cons = [[x0, y0, (2*E_kin-py**2)**0.5, py] for py in np.linspace(-10,10,8)]

outs = [ odeint(pqdot, con, t, atol=1e-9, rtol=1e-8) ) for con in init_cons[:8] ]

Obviously the number and parametrization of initial conditions may change.

The computation and display of the zero-crossings was changed to

def refine_crossing(a,b):
    tf = -a[0]/a[2]
    while abs(b[0])>1e-6:
        b = odeint(pqdot, a, [0,tf], atol=1e-8, rtol=1e-6)[-1];
        # Newton step using that b[0]=x(tf) and b[2]=x'(tf)
        tf -= b[0]/b[2]
    return [ b[1], b[3] ]

# Plot Poincare sections at x=0 and px>0
fig2 = figure(2)
for ii in xrange(8):
    #subplot(4, 2, ii+1)
    xcrossings = findcrossings(outs[ii][:,0], outs[ii][:,3])
    ycrossings = [ refine_crossing(outs[ii][cross], outs[ii][cross+1]) for cross in xcrossings]
    yints, pyints = array(ycrossings).T
    plot(yints, pyints,'.')
    ylabel("py")
    xlabel("y")
    title("Poincare section x = 0")

and evaluating the result of a longer integration interval