Roessler flow and corresponding Poincare map for section x=0, using PyDSTool

gistfile1.py

# -*- coding: utf-8 -*-
"""
Created on Mon Mar 30 21:49:55 2015

@author: brian
"""

import PyDSTool as dst
from PyDSTool.Toolbox import phaseplane as pp
from matplotlib import pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import time

def build_sys(gentype, sys_name, pars, varspecs, ics, 
              xname, yname, zname, intdomain, 
              tdata=[0, 50], algparams=None, events=[], tdomain=None):
    # we must give a name
    DSargs = dst.args(name=sys_name)

    # parameters
    DSargs.pars = pars

    # rhs of the differential equation
    DSargs.varspecs = varspecs

    # initial conditions
    DSargs.ics = ics

    # set the domain of integration.
    DSargs.xdomain = intdomain
    
    if tdomain != None:
        DSargs.tdomain = [0, 1000]
    else:
        DSargs.tdata = tdata

    # to avoid typos / bugs, use built-in Symbolic differentation!
    f = [DSargs.varspecs[xname], DSargs.varspecs[yname], DSargs.varspecs[zname]]
    Df = dst.Diff(f, [xname, yname, zname])
    print "Jacobian: ", str(Df.renderForCode())
    DSargs.fnspecs = {'Jacobian': (['t', xname, yname, zname],
                                   str(Df.renderForCode()))}

    if algparams != None:
        DSargs.algparams = algparams

    # Make auxiliary functions to define event lines near saddle
    res = pp.make_distance_to_line_auxfn('Gamma_out_plus',
                                         'Gamma_out_plus_fn',
                                         (xname, yname), True)
    man_pars = res['pars']
    man_auxfns = res['auxfn']
    res = pp.make_distance_to_line_auxfn('Gamma_out_minus',
                                         'Gamma_out_minus_fn',
                                         (xname, yname), True)
    man_pars.extend(res['pars'])
    man_auxfns.update(res['auxfn'])

    # update param values with defaults (0)
    for p in man_pars:
        DSargs.pars[p] = 0

    if gentype in [dst.Generator.Vode_ODEsystem, dst.Generator.Euler_ODEsystem]:
        targetlang = 'python'
    else:
        targetlang = 'c'

    DSargs.fnspecs.update(man_auxfns)
    ev_plus = dst.Events.makeZeroCrossEvent(expr='Gamma_out_plus_fn(%s,%s)' % (xname, yname),
                                            dircode=0,
                                            argDict={'name': 'Gamma_out_plus',
                                                     'eventtol': 1e-5,
                                                     'eventdelay': 1e-3,
                                                     'starttime': 0,
                                                     'precise': False,
                                                     'active': False,
                                                     'term': True},
                                            targetlang=targetlang,
                                            varnames=[xname, yname],
                                            fnspecs=man_auxfns,
                                            parnames=man_pars
                                            )
    ev_minus = dst.Events.makeZeroCrossEvent(expr='Gamma_out_minus_fn(%s,%s)' % (xname, yname),
                                             dircode=0,
                                             argDict={'name': 'Gamma_out_minus',
                                                      'eventtol': 1e-5,
                                                      'eventdelay': 1e-3,
                                                      'starttime': 0,
                                                      'precise': False,
                                                      'active': False,
                                                      'term': True},
                                             targetlang=targetlang,
                                             varnames=[xname, yname],
                                             fnspecs=man_auxfns,
                                             parnames=man_pars
                                             )

    DSargs.events = [ev_plus, ev_minus] + events
    print "Generating ODE system..."
    return gentype(DSargs), DSargs.fnspecs['Jacobian']

def plot_3D_flow(gentype, sys_name, pars, varspecs, ics,
                 xname, yname, zname, intdomain,
                 tdata=[0, 100], algparams=None, tdomain=None):
                     
    fig = plt.figure()
    ax3d = fig.add_subplot(111, projection='3d')
    
    ode_sys, jac = build_sys(gentype, sys_name, pars, varspecs, ics, 
                             xname, yname, zname, intdomain, tdomain=None)
                             
    traj = ode_sys.compute("trajectory")
    pts = traj.sample()
    
    ax3d.plot(pts[xname], pts[yname], pts[zname], label=sysname)
    
    return ax3d

def poincare_map(gentype, sys_name, pars, varspecs, ics, 
                 xname, yname, zname, intdomain, section=None,
                 section_varnames=None, tdata=[0, 100], algparams=None, tdomain=None):
    
    st = time.time()
    print "Beginning to calculate Poincare map points..."
    
    if section == None:
        print "Nothing was done because no section was provided. Please provide section."
        pass
    elif section_varnames == None:
        print "Nothing was done because a section was provided, but section_varnames was not provided. Please provide section_varnames."
        pass
    
    if gentype in [dst.Generator.Vode_ODEsystem, dst.Generator.Euler_ODEsystem]:
        targetlang = 'python'
    else:
        targetlang = 'c'
                
    ev_args = {'name': 'poinc_map_event', 'eventtol': 1e-5, 'eventdelay': 1e-12, 'starttime': 0, 'active': True, 'term': False, 'precise': True}
                       
    section_crossing_event = dst.Events.makeZeroCrossEvent(section, 0, ev_args, varnames=section_varnames, targetlang=targetlang)
                             
    ode_sys, jac = build_sys(gentype, sys_name, pars, varspecs, ics, 
                             xname, yname, zname, intdomain, events=[section_crossing_event], tdomain=tdomain)
    traj = ode_sys.compute("poinc_map")
    evs = traj.getEvents('poinc_map_event')
    
    et = time.time()
    print "Finished. time = {} s".format(round(et-st, 3))
                             
    return np.array([evs[xname], evs[yname], evs[zname]])
    
def plot_surface(ax3d, surface_fn, xdomain=[-5, 5], ydomain=[-5, 5], numpoints=50):
    # incomplete function
    x = np.linspace(xdomain[0], xdomain[1], num=numpoints)
    y = np.linspace(ydomain[0], ydomain[1], num=numpoints)
    xx, yy = np.meshgrid(x, y)
    pass

gentype = dst.Generator.Vode_ODEsystem
pars = {'a': 0.36, 'b': 0.4, 'c': 4.5}
xname = 'x'
yname = 'y'
zname = 'z'
varspecs = {'x': '-y - z',
            'y': 'x + a*y',
            'z': 'b*x - c*z + x*z'}
sysname = "Roessler"
intdomain = {'x': [-1.7, 1.7], 'y': [-1.7, 1.7]}    
ics = {'x': 0, 'y': -4.4, 'z': -0.054}                    
ax3d = plot_3D_flow(gentype, sysname, pars, varspecs, ics, xname, yname, zname, intdomain)

gentype = dst.Generator.Vode_ODEsystem
algparams={'max_pts': 1e6}
poinc_map = poincare_map(gentype, sysname+'_poincmap', pars, varspecs, ics, xname, yname, zname, intdomain, section='x', section_varnames=[xname], tdomain=[0, 20000])

ax3d.plot(poinc_map[0], poinc_map[1], poinc_map[2], 'r.', markersize=3, label="Poincare map points")
ax3d.legend(loc='best')