def psi(time, L=0, M=1000, aperiodicity=0,kind=0):
    """Return the probability amplitude the M-site lattice at the
    given time.  The initial condition is amplitude 1 at the central site,
    zero at all other sites.
    Parameters
    ----------
    time : float
        End time of the integration.
    L : float
        Nonlinearity parameter.
    M : int
        Number of lattice sites. (Default 101.)
    aperiodicity : float
        Aperiodicity parameter, defined in terms of the hoppings as (b/a - 1).
        (Default 0, which corresponds to a periodic chain.)
    Returns
    -------
    y : float
        Amplitude at the central site, $|\psi_{M/2}|$.
    """
    integrator = complex_ode(dnls_rhs(M, L, aperiodicity,kind))

    central_site_index = (M - 1)/2
    ic = np.zeros(shape=(M,))
    ic[central_site_index] = 1
    integrator.set_initial_value(ic)
    y= integrator.integrate(time)
    integrator = complex_ode(dnls_rhs(M, L, aperiodicity,kind))
    return y

    def k_int(mani,k_radii,k_powers,p,m,e):

        # time interval
        tspan = [0,2*np.pi]
        # initial condition
        ynot = np.array([0,0])

        out = np.zeros((len(k_powers)),dtype=np.complex)

        for j in range(len(k_powers)): # 1:length(k_powers)

            pre_k_ode = lambda t,y: k_ode(t,y,mani,k_radii[j],k_powers[j],p,e)
            # solve ode
            integrator = complex_ode(pre_k_ode).set_integrator('dopri5',
                            atol=m['k_int_options']['AbsTol'],
                            rtol=m['k_int_options']['RelTol'])
            integrator.set_initial_value(ynot,tspan[0])
            integrator.integrate(tspan[-1])
            Y = integrator.y
            Y = np.array([Y.T]).T

            # gather output
            out[j] = Y[0,-1]+Y[1,-1]*1j

        return out

def _evolve_cont(i,H,T,atol=1E-9,rtol=1E-9):
	"""
	This function evolves the ith local basis state under the Hamiltonian H up to period T. 
	This is used to construct the stroboscpoic evolution operator
	"""
	
	nsteps=_np.iinfo(_np.int32).max # huge number to make sure solver is successful.
	psi0=_np.zeros((H.Ns,),dtype=_np.complex128) 
	psi0[i]=1.0

	solver=complex_ode(H._hamiltonian__SO)
	solver.set_integrator('dop853', atol=atol,rtol=rtol,nsteps=nsteps) 
	solver.set_initial_value(psi0,t=0.0)
	t_list = [0,T]
	nsteps = 1
	while True:
		for t in t_list[1:]:
			solver.integrate(t)
			if solver.successful():
				if t == T:
					return solver.y
				continue
			else:
				break

		nsteps *= 10
		t_list = _np.linspace(0,T,num=nsteps+1,endpoint=True)

def psi(t,M=10,steps=1, L=0,  aperiodicity=0,kind=0):
    """
    Parameters
    ----------
    time : float
        End time of the integration.
    L : float
        Nonlinearity parameter.
    M : int
        Number of lattice sites. (Default 101.)
    aperiodicity : float
        Aperiodicity parameter, defined in terms of the hoppings as (b/a - 1).
        (Default 0, which corresponds to a periodic chain.)
    kind: 0(periodic),1 (fib),2 (TM),3 (RS)
    Returns
    -------
    y : float
    wavefn at time t
    """
    r = integrate.complex_ode(dnls_rhs(M, L, aperiodicity,kind))
    central_site_index = (M - 1)/2
    ic = np.zeros(shape=(M,))
    ic[central_site_index] = 1.0
    r.set_initial_value(ic)
    dt = float(t)/steps
    wavefn = np.ones((steps,M), dtype=np.complex)
    t_arr=np.ones((steps), dtype=np.float)
    for idx in xrange(steps): #steps are needed here
      if r.successful(): #essentially it's an saying **if True:**
		wavefn[idx,:] = r.y[:]
		t_arr[idx]=r.t
      else:
		raise ValueError("Integration failed at t = {}".format(r.t)) #i don't understand this line
      r.integrate(r.t + dt)#note that if dt is too big, then you will get an error. So, steps should at least be same as 
    return wavefn,t_arr     #time units

def psi(t,M=10,steps=1, L=0,  aperiodicity=0,kind=0):
    r = integrate.complex_ode(dnls_rhs(M, L, aperiodicity,kind))
    central_site_index = (M - 1)/2
    ic = np.zeros(shape=(M,))
    ic[central_site_index] = 1.0
    r.set_initial_value(ic)
    dt = t/steps
    wavefn = np.ones((steps,M), dtype=np.complex)
    t_arr=np.ones((steps), dtype=np.float)
    for idx in xrange(steps):
        if r.successful(): #essentially, **if True:**
                if abs(r.y[0])< 1e-8 and abs(r.y[-1])<1e-8: #to make sure the lattice is big enough so that wavefn doesn't touch the boundary.
                    wavefn[idx,:] = r.y[:]
                    t_arr[idx]=r.t
                else:  		
                    #raise ValueError("wavefunction touching the boundary at t = {}".format(r.t))
                    print ("wfn touching the boundary at t = {} for p={}, U={}".format(r.t,aperiodicity,L))
                    wave_fn_truncated=wavefn[0:idx,:]
                    t_arr_truncated=t_arr[0:idx]
                    return wave_fn_truncated, t_arr_truncated
        else:
            raise ValueError("Integration failed at t = {}".format(r.t)) 
        r.integrate(r.t + dt) 
        #print r.t
    return wavefn,t_arr     

    def _run_solout_break_test(self, integrator):
        # Check correct usage of stopping via solout
        ts = []
        ys = []
        t0 = 0.0
        tend = 20.0
        y0 = [0.0]

        def solout(t, y):
            ts.append(t)
            ys.append(y.copy())
            if t > tend/2.0:
                return -1

        def rhs(t, y):
            return [1.0/(t - 10.0 - 1j)]

        ig = complex_ode(rhs).set_integrator(integrator)
        ig.set_solout(solout)
        ig.set_initial_value(y0, t0)
        ret = ig.integrate(tend)
        assert_array_equal(ys[0], y0)
        assert_array_equal(ys[-1], ret)
        assert_equal(ts[0], t0)
        assert_(ts[-1] > tend/2.0)
        assert_(ts[-1] < tend)

def manifold_compound(x, z, lamda, s, p, m, A, k, pmMU):
    """
    manifold_compound(x,z,lambda,s,p,m,A,k,pmMU)

    Returns the vector representing the manifold evaluated at x(2).

    Input "x" is the interval the manifold is computed on, "z" is the
    initializing vector for the ode solver, "lambda" is the point on the
    complex plane where the Evans function is computed, s,p,m are structures
    explained in the STABLAB documentation, "A" is the function handle to the
    desired Evans matrix, "k" is the dimension of the manifold sought, and
    "pmMU" is 1 or -1 depending on if respectively the growth or decay
    manifold is sought.
    """

    eigenVals,eigenVects = np.linalg.eig(A(x[0],lamda,s,p))

    ind = np.argmax(np.real(pmMU*eigenVals))
    MU = eigenVals[ind]

    # Solve the ODE
    def ode_f(x,y):
        return capa(x,y,lamda,s,p,A,m['n'],k,MU)
    if 'options' in m:
        try:
            integrator = complex_ode(ode_f).set_integrator('dopri5',
                         atol=m['options']['AbsTol'],
                         rtol=m['options']['RelTol'],
                         nsteps=m['options']['nsteps'])
        except KeyError:
            integrator = complex_ode(ode_f).set_integrator('dopri5',atol=1e-6,
                                                            rtol=1e-5,
                                                            nsteps=10000)
    else:
        integrator = complex_ode(ode_f).set_integrator('dopri5',atol=1e-6,
                                                                rtol=1e-5,
                                                                nsteps=10000)
    x0 = x[0]
    z0 = z.T[0]
    integrator.set_initial_value(z0,x0)
    integrator.integrate(x[-1])
    Z = integrator.y

    out = Z
    return out

def manifold_polar(x,y,lamda,A,s,p,m,k,mu):
    """
     Returns "Omega", the orthogonal basis for the manifold evaluated at x[-1]
     and "gamma" the radial equation evaluated at x[-1].

     Input "x" is the interval on which the manifold is solved, "y" is the
     initializing vector, "lambda" is the point in the complex plane where the
     Evans function is evaluated, "A" is a function handle to the Evans
     matrix, s, p,and m are structures explained in the STABLAB documentation,
     and k is the dimension of the manifold sought.
    """

    def ode_f(x,y):
        return m['method'](x,y,lamda,A,s,p,m['n'],k,mu,m['damping'])

    t0 = x[0]
    y0 = y.reshape(m['n']*k,1,order='F')
    y0 = np.concatenate((y0,np.array([[0.0]],dtype=np.complex)),axis=0)
    y0 = y0.T[0]

    #initiate integrator object
    if 'options' in m:
        try:
            integrator = complex_ode(ode_f).set_integrator('dopri5',
                         atol=m['options']['AbsTol'],
                         rtol=m['options']['RelTol'],
                         nsteps=m['options']['nsteps'])
        except KeyError:
            integrator = complex_ode(ode_f).set_integrator('dopri5',atol=1e-6,
                                                            rtol=1e-5,
                                                            nsteps=10000)
    else:
        integrator = complex_ode(ode_f).set_integrator('dopri5',atol=1e-6,
                                                                rtol=1e-5,
                                                                nsteps=10000)

    integrator.set_initial_value(y0,t0) # set initial time and initial value
    integrator.integrate(x[-1])
    Y = integrator.y
    Y = np.array([Y.T]).T

    omega = Y[0:k*m['n'],-1].reshape(m['n'],k,order = 'F')
    gamma = np.exp(Y[m['n']*k,-1])

    return omega, gamma

def _testcase_one_mode(h_sys, rho0, g, w, L, timesteps):
    """Integration of the single environment-mode case for debugging. The exact
    reduced dynamics is described by the reduced density operator p00 as well
    as three auxiliary states (p01, p10, p11). Their equations of motion read


    :param h_sys: @todo
    :param rho0: @todo
    :param g: @todo
    :param w: @todo
    :param L: @todo
    :param timesteps: @todo
    :returns: @todo

    """
    from numpy import dot
    from scipy.integrate import complex_ode

    dim = h_sys.shape[0]
    adj = lambda A: np.conj(np.transpose(A))

    prop = sp.lil_matrix((dim**2 * 4, dim**2 * 4), dtype=complex)
    i = [[(slice(m*dim**2, (m+1)*dim**2), slice(n*dim**2, (n+1)*dim**2))
          for n in range(4)] for m in range(4)]

    prop[i[0][0]] += -1.j * commutator(h_sys)
    prop[i[0][2]] += -multiply_raveled(adj(L), 'l') + multiply_raveled(adj(L), 'r')
    prop[i[0][1]] += multiply_raveled(L, 'l') - multiply_raveled(L, 'r')

    prop[i[1][1]] += -1.j * commutator(h_sys) - np.conj(w) * np.identity(dim**2)
    prop[i[1][0]] += np.conj(g) * multiply_raveled(adj(L), 'r')
    prop[i[1][3]] += -multiply_raveled(adj(L), 'l') - multiply_raveled(adj(L), 'r')

    prop[i[2][2]] += -1.j * commutator(h_sys) - w * np.identity(dim**2)
    prop[i[2][0]] += g * multiply_raveled(L, 'l')
    prop[i[2][3]] += -multiply_raveled(L, 'l') - multiply_raveled(L, 'r')

    prop[i[3][3]] += -1.j * commutator(h_sys) - (w + np.conj(w)) * np.identity(dim**2)
    prop[i[3][1]] += g * multiply_raveled(L, 'l')
    prop[i[3][2]] += np.conj(g) * multiply_raveled(adj(L), 'r')

    print('CORRECT')
    print(prop)

    y0 = np.zeros((4, dim, dim), dtype=complex)
    y0[0] = rho0
    rho = np.empty((len(timesteps), dim, dim), dtype=complex)
    rho[0] = rho0

    r = complex_ode(lambda t, y: prop.dot(y))\
        .set_integrator('vode', atol=1e-10, rtol=1e-10, nsteps=100) \
        .set_initial_value(y0.ravel())
    for i, t in enumerate(timesteps[1:]):
        r.integrate(t)
        rho[i + 1] = r.y.reshape((4, dim, dim))[0]

    return timesteps, rho

    def solve_ODE(self, H=None):
        """Iteratively solve the ODE dy/dt = f(t,y) on a discretized time-grid.

            Returns:
            --------
                    t:  (N,)  ndarray
                        Time array.
                phi_a:  (N,2) ndarray
                        Overlap <phi_a|psi>.
                phi_b:  (N,2) ndarray
                        Overlap <phi_b|psi>.
        """

        if H is None:
            H = self.H

        # set initial conditions
        self.get_c_eigensystem()        # calculate eigensystem for all times
        if self.init_state_method == 'gain':
            self._find_gain_state()
        elif self.init_state_method == 'energy':
            self._find_lower_energy_state()
        self.eVec0 = self._get_init_state()

        # create ode object to solve Schroedinger equation (SE)
        ode_kwargs = {'rtol': 1e-9,
                      'atol': 1e-9}
        SE = complex_ode(lambda t, phi: -1j*H(t).dot(phi))
        SE.set_integrator('dopri5', **ode_kwargs)
        SE.set_initial_value(self.eVec0, t=0.0)

        # iterate SE
        for n, tn in enumerate(self.t):
            if SE.successful():
                self.Psi[n,:] = SE.y
                SE.integrate(SE.t + self.dt)
            else:
                raise Exception("ODE convergence error!")

        if self.calc_adiabatic_state:
            self._get_adiabatic_state()

        # replace projection of states by dot product via Einstein sum
        projection = np.einsum('ijk,ij -> ik',
                               self.eVecs_l, self.Psi)
        # use alternative means to obtain coefficients:
        #  (c1, c2) = X^-1^T psi
        # from scipy.linalg import inv
        # projection = [np.einsum('jk,j -> k', inv(self.eVecs_r[n,:]).T, self.Psi[n,:])
        #                for n, _ in enumerate(self.t)]
        # projection = np.asarray(projection)

        self.phi_a, self.phi_b = [projection[:,n] for n in (0,1)]

        return self.t, self.phi_a, self.phi_b

 def _start_integrator(self, ham, small_step):
     """ Initialize a stepping integrator. """
     self.sparse_ham = issparse(ham)
     evo_eq = calc_evo_eq(self.isdop, self.sparse_ham)
     self.stepper = complex_ode(evo_eq(ham))
     int_mthd, step_fct = ('dopri5', 150) if small_step else ('dop853', 50)
     first_step = norm(ham, 'f') / step_fct
     self.stepper.set_integrator(int_mthd, nsteps=0, first_step=first_step)
     self.stepper.set_initial_value(self.p0.A.reshape(-1), self.t0)
     self.update_to = self._update_to_integrate
     self.solved = False

def sol(kx, ky):
	def ham_k(t):
		return ham(kx, ky, t)
	def f(y_vec, t):
		y_1, y_2 = y_vec
		fun = -1j * np.dot( ham_k(t), np.array([ [y_1],[y_2] ]) )
		return [fun[0,0], fun[1,0]]

	y_result = complex_ode(f, y0, t_output)
	
	return np.array([y_result.real, y_result.imag])

def f(df, u, h, theta):
    dt = h / 1e1
    df_args = functools.partial(df, theta = theta)
    r = si.complex_ode(df_args)
    sol = []
    for v in u:
        x0 = [0., v * 1j]
        r.set_initial_value(x0, 0)
        while r.successful() and r.t < h:
            r.integrate(r.t + dt)
        sol.append(r.y)
    return np.array(sol).T

    def _integrator(self, f, **kwargs):
        from scipy.integrate import complex_ode

        defaults = {
            'nsteps': 1e9,
            'with_jacobian': False,
            'method': 'bdf',
        }
        defaults.update(kwargs)

        r = complex_ode(f).set_integrator('vode', atol=self.error_abs, rtol=self.error_rel, **defaults)
        return r

def central_amplitude(time, L, M=101, aperiodicity=0):


integrator = complex_ode(dnls_rhs(M, L, aperiodicity))

    central_site_index = (M - 1)/2
    ic = np.zeros(shape=(M,))
    ic[central_site_index] = 1
    integrator.set_initial_value(ic)

    y = integrator.integrate(time)
    return np.abs(y[central_site_index])

def compare_performance(func, y0, t0, tf, y_ref, problem_name, tol_boundary=(0,6), is_complex=False, nsteps=10e5, solout=(lambda t: t)):
    print 'RUNNING COMPARISON TEST FOR ' + problem_name
    tol = [1.e-3,1.e-5,1.e-7,1.e-9,1.e-11,1.e-13]
    a, b = tol_boundary
    tol = tol[a:b]

    extrap = {}
    dopri5 = {}
    dop853 = {}

    for method in [extrap, dopri5, dop853]:
        for diagnostic in ['runtime','fe_seq','fe_tot','yerr','nstp']:
            method[diagnostic] = np.zeros(len(tol))

    def func2(t,y):
        return func(y,t)

    for i in range(len(tol)):
        print 'Tolerance: ', tol[i]

        for method, name in [(extrap,'ParEx'), (dopri5,'DOPRI5'), (dop853,'DOP853')]:
            print 'running ' + name
            start_time = time.time()
            if name == 'ParEx':
                y, infodict = parex.solve(func, [t0, tf], y0, solver=parex.Solvers.EXPLICIT_MIDPOINT, atol=tol[i], rtol=tol[i], max_steps=nsteps, adaptive=True, diagnostics=True)
                y[-1] = solout(y[-1])
                method['yerr'][i] = relative_error(y[-1], y_ref)
            else: # scipy solvers DOPRI5 and DOP853
                if is_complex:
                    r = complex_ode(func2, jac=None).set_integrator(name.lower(), atol=tol[i], rtol=tol[i], verbosity=10, nsteps=nsteps)
                else:
                    r = ode(func2, jac=None).set_integrator(name.lower(), atol=tol[i], rtol=tol[i], verbosity=10, nsteps=nsteps)
                r.set_initial_value(y0, t0)
                r.integrate(r.t+(tf-t0))
                assert r.t == tf, "Integration did not converge. Try increasing the max number of steps"
                y = solout(r.y)
                method['yerr'][i] = relative_error(y, y_ref)

            method['runtime'][i] = time.time() - start_time
            method['fe_seq'][i], method['fe_tot'][i], method['nstp'][i] = infodict['fe_seq'], infodict['nfe'], infodict['nst']
            print 'Runtime: ', method['runtime'][i], ' s   Error: ', method['yerr'][i], '   fe_seq: ', method['fe_seq'][i], '   fe_tot: ', method['fe_tot'][i], '   nstp: ', method['nstp'][i]
            print ''
        
        print ''

    for method, name in [(extrap,'ParEx'), (dopri5,'DOPRI5'), (dop853,'DOP853')]:
        print "Final data: " + name
        print method['runtime'], method['fe_seq'], method['fe_tot'], method['yerr'], method['nstp']
    print ''
    print ''

    return (extrap, dopri5, dop853)

    def plot_this(ynot,tspan,ode,rp):
        sol = stablab.Struct()
        sol.t = []
        sol.y = []
        def updateSol(tVal,yVals):
            sol.t.append(tVal)
            sol.y.append(yVals)
            return None

        integrator = complex_ode(ode)
        integrator.set_integrator('dopri5',atol=1e-10,rtol=1e-10)
        integrator.set_solout(updateSol)
        integrator.set_initial_value(ynot,tspan[0])
        integrator.integrate(tspan[-1])
        sol.t, sol.y = np.array(sol.t), np.array(sol.y)
        return sol

    def _do_problem(self, problem, integrator, method='adams'):

        # ode has callback arguments in different order than odeint
        f = lambda t, z: problem.f(z, t)
        jac = None
        if hasattr(problem, 'jac'):
            jac = lambda t, z: problem.jac(z, t)
        ig = complex_ode(f, jac)
        ig.set_integrator(integrator,
                          atol=problem.atol/10,
                          rtol=problem.rtol/10,
                          method=method)
        ig.set_initial_value(problem.z0, t=0.0)
        z = ig.integrate(problem.stop_t)

        assert_(ig.successful(), (problem, method))
        assert_(problem.verify(array([z]), problem.stop_t), (problem, method))

	def __init__(self,myhamgen,param):
		self.hamgen = myhamgen
		self.param=param

		#Set up the solver
		self.norm=0
		#self.r=ode(self.func).set_integrator('zvode', method='bdf',rtol=1e-6)
		#self.r=complex_ode(self.func).set_integrator('dopri5',rtol=1e-6)
		self.r=complex_ode(self.func).set_integrator('dopri5',rtol=1e-6,nsteps=100000)


		#Generate basis vectors
		self.b=[]
		for x in range(0,param.numneu):
			self.b.append(np.zeros(param.numneu))
			self.b[x][x]=1.0

		self.splines=Splines.Spline()

def psi_old(t,M=10,steps=1, L=0,  aperiodicity=0,kind=0):
    """
    Parameters
    ----------
    time : float
        End time of the integration.
    L : float
        Nonlinearity parameter.
    M : int
        Number of lattice sites. (Default 101.)
    aperiodicity : float
        Aperiodicity parameter, defined in terms of the hoppings as (b/a - 1).
        (Default 0, which corresponds to a periodic chain.)
    kind: 0(periodic),fib,tm,rs
    Returns
    -------
    y : float
    wavefn at time t
    """
    r = integrate.complex_ode(dnls_rhs(M, L, aperiodicity,kind))
    central_site_index = (M - 1)/2
    ic = np.zeros(shape=(M,))
    ic[central_site_index] = 1.0
    r.set_initial_value(ic)
    dt = float(t)/steps
    wavefn = np.zeros((steps,M), dtype=np.complex)
    t_arr=np.zeros((steps), dtype=np.float)
    for idx in xrange(steps):
      if r.successful(): #essentially, **if True:**
		if abs(r.y[0])< 1e-10 and abs(r.y[-1])<1e-10: #to make sure the lattice is big enough so that wavefn doesn't touch the boundary.
		    wavefn[idx,:] = r.y[:]
		    t_arr[idx]=r.t
		    #print "y"
		else:
		    print ("wfn touching the boundary at t = {} for p={}".format(r.t,aperiodicity))
		    wave_fn_truncated=wavefn[0:idx,:]
		    t_arr_truncated=t_arr[0:idx]
		    return wave_fn_truncated, t_arr_truncated
      else:
		raise ValueError("Integration failed at t = {}".format(r.t)) 
      r.integrate(r.t + dt)#note that if dt is too big, then you will get an error. So, steps should at least be same as 
    return wavefn,t_arr     #time units