Python scipy.optimize 模块，newton() 实例源码

def main():
    # linear regression with gradient descent
    test_gd()
    # function optimization with gradient descent
    manual_gd(f_)
    # root finding methods
    # bisection method
    bisect(f, 0, 2)
    # secant method
    secant(f, 1, 2)
    # newton method
    newton(f, f_, 1)
    # scipy
    print("SB\tx=%f" % optimize.bisect(f, 0, 2))
    print("SS\tx=%f" % optimize.newton(f, 1))
    print("SN\tx=%f" % optimize.newton(f, 1, f_))

def find_loc(self, time, lat=None, lon=None):
        'Calculates position on sphere given time and one known coordinate'
        from scipy.optimize import newton
        def make_lon_func(time, lat, radius):
            'Closure for distance solver function'
            from sat_tools import Satellite
            satellite = Satellite()
            def lon_func(lon):
                'Function to solve for distance'
                nonlocal time, lat, radius
                return satellite.distance_to_ac(time, lat, lon, 3.5e4)-radius
            return lon_func
        if lat is None:
            raise NotImplementedError
        if lon is None:
            if not time > self.data[-1][0]:
                radius = list(filter(lambda x: x[0] == time, self.data))[0][1]
                offset = 0
            else:
                time = self.data[-1][0]
                radius = self.data[-1][1]
                offset = 1
            lon_func = make_lon_func(time, lat, radius)
            lon = newton(lon_func, 100.0, tol=1e-2) + offset
        return lat, round(lon, 2)

def test_intersection_point_function():
    for _ in xrange(1000):

        eps = np.random.rand() * 0.2

        x_old = - np.random.rand() * np.log(np.exp(0.2) - 1)
        y_old = function(x_old) - eps

        x_tangent = op.newton(LinearizeTwoTermPosynomials.tangent_point_func, x_old + 1, args=(x_old, eps))
        y_tangent = function(x_tangent)

        tangent_slope = (y_old - y_tangent) / (x_old - x_tangent)
        tangent_intercept = - (y_old - y_tangent) * x_tangent / (x_old - x_tangent) + y_tangent

        x_intersection = op.newton(LinearizeTwoTermPosynomials.intersection_point_func,
                                   x_tangent + 1, args=(tangent_slope, tangent_intercept, eps))
        assert (x_intersection > x_tangent)

        diff = function(x_intersection) - eps - tangent_slope * x_intersection - tangent_intercept

        assert (diff < 0.0001)
    return

def newton_rad_func(E_val,D,m,L0):

   E = E_val

   c2 = D*D*E/4.

   L = newton(newton_ang_func,L0,args=(c2,m),tol=1e-8,maxiter=500)

   slope = -(-L+m*(m+1.)+2.*D+c2)/(2.*(m+1.))

   z0 = [1+step*slope,slope]

   z = odeint(g,z0,x_rad,args=(c2,L,D,m))

   temp=pow(x_rad,2.0)-1.
   temp=pow(temp,m/2.)

   zz=temp*z[:,0]

   first_zz = np.array([1])
   zz=np.append(first_zz, zz)

   return z[:,1][-1]

def computeoptgam(self,cinfo,xmin,iStar,mingrad):
        ainv,ainv2,ainv3=cinfo
        a=float(np.matrix(xmin)*ainv*np.matrix(xmin).T)
        b=float(np.matrix(xmin)*ainv2*np.matrix(xmin).T)
        c=float(np.matrix(xmin)*ainv3*np.matrix(xmin).T)
        t=float(np.trace(ainv2))
        def computeGammaF3(a,b,c,t):
            def F(x=None,z=None):
                if x is None: return 0, cvxopt.matrix(0.2, (1,1))
                if x.size[0]!=1 or x[0]==1: return None
                f=cvxopt.matrix(0.0,(1,1))
                df=cvxopt.matrix(0.0,(1,1))
                f[0,0]=x**2*b**2/((1-x+a*x)**2*(1-x)**2)-2*x*c/((1-x+a*x)*(1-x)**2)+t/(1-x)**2
                df[0,0]=  (2*x*b**2)/((x - 1)**2*(a*x - x + 1)**2) - (2*c)/((x - 1)**2*(a*x - x + 1)) - (2*t)/(x - 1)**3 - (2*x**2*b**2)/((x - 1)**3*(a*x - x + 1)**2) + (4*x*c)/((x - 1)**3*(a*x - x + 1)) - (2*x**2*b**2*(a - 1))/((x - 1)**2*(a*x - x + 1)**3) + (2*x*c*(a - 1))/((x - 1)**2*(a*x - x + 1)**2)

                if z is None:return f,df
                H=cvxopt.matrix(0.0,(1,1))
                H[0,0]=z[0]*((6*t)/(x - 1)**4 + (8*c)/((x - 1)**3*(a*x - x + 1)) + (2*b**2)/((x - 1)**2*(a*x - x + 1)**2) - (8*x*b**2)/((x - 1)**3*(a*x - x + 1)**2) + (4*c*(a - 1))/((x - 1)**2*(a*x - x + 1)**2) + (6*x**2*b**2)/((x - 1)**4*(a*x - x+ 1)**2) - (12*x*c)/((x - 1)**4*(a*x - x + 1)) - (8*x*b**2*(a - 1))/((x - 1)**2*(a*x - x + 1)**3) - (4*x*c*(a - 1)**2)/((x - 1)**2*(a*x - x + 1)**3) + (8*x**2*b**2*(a - 1))/((x - 1)**3*(a*x - x + 1)**3) - (8*x*c*(a - 1))/((x - 1)**3*(a*x - x + 1)**2) + (6*x**2*b**2*(a - 1)**2)/((x - 1)**2*(a*x - x + 1)**4) )      
                return f,df,H
            G=cvxopt.matrix([[-1.0,1.0]]) 
            h=cvxopt.matrix([0.0,1.0]) 
            tol=1.e-1
            solvers.options['abstol']=tol
            solvers.options['reltol']=tol
            solvers.options['feastol']=tol
            solvers.options['show_progress'] = False
            return (solvers.cp(F, G=G, h=h)['x'])[0]   
        def GammaF3(a,b,c,t,x0,maxiter):
            def func(x):
                return (2*x*b**2)/((x - 1)**2*(a*x - x + 1)**2) - (2*c)/((x - 1)**2*(a*x - x + 1)) - (2*t)/(x - 1)**3 - (2*x**2*b**2)/((x - 1)**3*(a*x - x + 1)**2) + (4*x*c)/((x - 1)**3*(a*x - x + 1)) - (2*x**2*b**2*(a - 1))/((x - 1)**2*(a*x - x + 1)**3) + (2*x*c*(a - 1))/((x - 1)**2*(a*x - x + 1)**2)
            def fprime(x):
                return ((6*t)/(x - 1)**4 + (8*c)/((x - 1)**3*(a*x - x + 1)) + (2*b**2)/((x - 1)**2*(a*x - x + 1)**2) - (8*x*b**2)/((x - 1)**3*(a*x - x + 1)**2) + (4*c*(a - 1))/((x - 1)**2*(a*x - x + 1)**2) + (6*x**2*b**2)/((x - 1)**4*(a*x - x+ 1)**2) - (12*x*c)/((x - 1)**4*(a*x - x + 1)) - (8*x*b**2*(a - 1))/((x - 1)**2*(a*x - x + 1)**3) - (4*x*c*(a - 1)**2)/((x - 1)**2*(a*x - x + 1)**3) + (8*x**2*b**2*(a - 1))/((x - 1)**3*(a*x - x + 1)**3) - (8*x*c*(a - 1))/((x - 1)**3*(a*x - x + 1)**2) + (6*x**2*b**2*(a - 1)**2)/((x - 1)**2*(a*x - x + 1)**4) )

            return newton(func=func,x0=x0,fprime=fprime,tol=0.1,maxiter=maxiter)   
        Gamma=computeGammaF3(a,b,c,t)
        return Gamma

def getmass_zfourge(N, z):
    if isinstance(N, np.ndarray) or isinstance(N, list):
        mass = np.zeros([len(N)])
        for i, elem in enumerate(N):
            mass[i] = newton(getnum_zfourge, 10., args=(z,elem))
    else:
        mass = newton(getnum_zfourge, 10., args=(z,N))

    return mass

def addSSmNrm(self,solution):
        '''
        Finds steady state (normalized) market resources and adds it to the
        solution.  This is the level of market resources such that the expectation
        of market resources in the next period is unchanged.  This value doesn't
        necessarily exist.

        Parameters
        ----------
        solution : ConsumerSolution
            Solution to this period's problem, which must have attribute cFunc.
        Returns
        -------
        solution : ConsumerSolution
            Same solution that was passed, but now with the attribute mNrmSS.
        '''
        # Make a linear function of all combinations of c and m that yield mNext = mNow
        mZeroChangeFunc = lambda m : (1.0-self.PermGroFac/self.Rfree)*m + (self.PermGroFac/self.Rfree)*self.ExIncNext

        # Find the steady state level of market resources
        searchSSfunc = lambda m : solution.cFunc(m) - mZeroChangeFunc(m) # A zero of this is SS market resources
        m_init_guess = self.mNrmMinNow + self.ExIncNext # Minimum market resources plus next income is okay starting guess
        try:
            mNrmSS = newton(searchSSfunc,m_init_guess)
        except:
            mNrmSS = None

        # Add mNrmSS to the solution and return it
        solution.mNrmSS = mNrmSS
        return solution

def newton(f, f_, x0, e=1e-12):
    """
    Newton method for root finding
    """
    n = 0
    while True:
        n += 1
        x = x0 - (f(x0) / f_(x0))
        if abs(x - x0) < e:
            break
        x0 = x
    print("N\ti=%d\tx=%f\tf(x)=%f" % (n, x, f(x)))
    return x

def root_find(func, root_list, x, y, step_size, x_values):
    if len(root_list) == 0:
        root = sp.newton(func, y, maxiter = 2000)
        root_list.append(root)
        x_values.append(x)
    elif len(root_list) == 1:
        root = sp.newton(func, root_list[-1], maxiter = 2000)
        root_list.append(root)
        x_values.append(x)
    elif len(root_list) == 2:
        grad = ((root_list[-1] - root_list[-2])*
        np.abs(step_size/(x_values[-1] - x_values[-2])))
        root = sp.newton(func, root_list[-1] + grad, maxiter = 2000)
        if np.abs(root-root_list[-1]) > 0.1:
            raise Exception('largegrad at {}'.format(x_values[-1]))
        else:
            root_list.append(root)
            x_values.append(x)
    else:
        grad = ((root_list[-1] - root_list[-2]) + \
        (root_list[-1] - 2*root_list[-2] + root_list[-3]))* \
        np.abs(step_size/(x_values[-1] - x_values[-2]))
        root = sp.newton(func, root_list[-1] + grad, maxiter = 2000)
        if np.abs(root-root_list[-1]) > 0.1:
            raise Exception('largegrad at {}'.format(x_values[-1]))
        else:
            root_list.append(root)
            x_values.append(x)
    return x_values, root_list

# Attempts to trace a line of roots of a function.
# Arguments:
# func - The function. This can have any number of variables, 
# but the first one is always the one that the root finder will be used on.
# x - The x coordinate of the starting point.
# y - The y coordinate of the starting point.
# step_size - The step size to be used in calculating the next point.
# x_end_left, x_end_right - These define the limits of the interval in which
# line is to be traced.

def newton_method(self, initial_point, epsilon=1e-4):
        """ Newton-Raphson method
        :param epsilon: shock parameter for calculating numeric derivatives
        :param initial_point:
        :return: x(n + 1) = x(n) - f(x(n)) / (f'(x(n)), where f'(y) = (f(y + eps) - f(y - eps)) / (2 * eps)
        """

        def f_prime(x):
            f_x_plus = self.__objective_function(x + epsilon)
            f_x_minus = self.__objective_function(x - epsilon)
            f_p = 1.0 * (f_x_plus - f_x_minus) / (2 * epsilon)

            return f_p

        def f_second(x):
            f_x_plus = self.__objective_function(x + epsilon)
            f_x_minus = self.__objective_function(x - epsilon)
            f_x = self.__objective_function(x)
            f_s = 1.0 * (f_x_plus - 2 * f_x + f_x_minus) / (epsilon * epsilon)

            return f_s

        from scipy.optimize import newton
        optimum = newton(self.__objective_function, initial_point,
                         fprime=f_prime, fprime2=f_second,
                         tol=self.tolerance, maxiter=self.iteration_max)

        return optimum

def potential2radius(pot_func, pot, q, d=1, F=1.0, component=1, sma=1.0, loc='pole',
                     tol=1e-10, maxiter=50):
    """
    @param pot_func: the potential function to use
    @type pot_func: func
    @param pot: Roche potential value (unitless)
    @type pot: float
    @param q: mass ratio
    @type q: float
    @param d: separation (in units of semi-major axis)
    @type d: float
    @param F: synchronicity parameter
    @type F: float
    @param component: component in the system (1 or 2)
    @type component: integer
    @param sma: semi-major axis
    @type sma: value of the semi-major axis.
    @param loc: location of the radius ('pole' or 'eq')
    @type loc: str
    @return: potential value for a certain radius
    @rtype: float
    """


    if 'pol' in loc.lower():
        theta = 0
    elif 'eq' in loc.lower():
        theta = np.pi/2.
    else:
        ValueError,"don't understand loc=%s"%(loc)

    potential = lambda r, theta, phi, q, d, F, c: binary_potential(r, theta, phi, q, d, F, c)-pot

    try:
        r_pole = newton(potential, x0=1./pot, args=(theta, 0, q, d, F, component), tol=tol, maxiter=maxiter)
    except RuntimeError:
        raise ValueError("Failed to converge, potential {} is probably too low".format(pot))
    return r_pole*sma

def crossing(b, component, time, dynamics_method='keplerian', ltte=True, tol=1e-4, maxiter=1000):
    """
    tol in days
    """


    def projected_separation_sq(time, b, dynamics_method, cind1, cind2, ltte=True):
        """
        """
        #print "*** projected_separation_sq", time, dynamics_method, cind1, cind2, ltte


        times = np.array([time])

        if dynamics_method in ['nbody', 'rebound']:
            # TODO: make sure that this takes systemic velocity and corrects positions and velocities (including ltte effects if enabled)
            ts, xs, ys, zs, vxs, vys, vzs = dynamics.nbody.dynamics_from_bundle(b, times, compute=None, ltte=ltte)

        elif dynamics_method=='bs':
            ts, xs, ys, zs, vxs, vys, vzs = dynamics.nbody.dynamics_from_bundle_bs(b, times, compute, ltte=ltte)

        elif dynamics_method=='keplerian':
            # TODO: make sure that this takes systemic velocity and corrects positions and velocities (including ltte effects if enabled)
            ts, xs, ys, zs, vxs, vys, vzs = dynamics.keplerian.dynamics_from_bundle(b, times, compute=None, ltte=ltte, return_euler=False)

        else:
            raise NotImplementedError


        return (xs[cind2][0]-xs[cind1][0])**2 + (ys[cind2][0]-ys[cind1][0])**2


    # TODO: optimize this by allowing to pass cind1 and cind2 directly (and fallback to this if they aren't)
    starrefs = b.hierarchy.get_stars()
    cind1 = starrefs.index(component)
    cind2 = starrefs.index(b.hierarchy.get_sibling_of(component))

    # TODO: provide options for tol and maxiter (in the frontend computeoptionsp)?
    return newton(projected_separation_sq, x0=time, args=(b, dynamics_method, cind1, cind2, ltte), tol=tol, maxiter=maxiter)

def calc_smoothpar_logistic2(metapar):
    """Return the smoothing parameter corresponding to the given meta
    parameter when using :func:`~hydpy.cythons.smoothutils.smooth_logistic2`.

    Calculate the smoothing parameter value corresponding the meta parameter
    value 2.5:

    >>> from hydpy.auxs.smoothtools import calc_smoothpar_logistic2
    >>> smoothpar = calc_smoothpar_logistic2(2.5)

    Using this smoothing parameter value, the output of function
    :func:`~hydpy.cythons.smoothutils.smooth_logistic2` differs by
    1 % from the related `true` discontinuous step function for the
    input values -2.5 and 2.5 (which are located at a distance of 2.5
    from the position of the discontinuity):

    >>> from hydpy.cythons import smoothutils
    >>> from hydpy.core.objecttools import round_
    >>> round_(smoothutils.smooth_logistic2(-2.5, smoothpar))
    0.01
    >>> round_(smoothutils.smooth_logistic2(2.5, smoothpar))
    2.51

    For zero or negative meta parameter values, a zero smoothing parameter
    value is returned:

    >>> round_(calc_smoothpar_logistic2(0.0))
    0.0
    >>> round_(calc_smoothpar_logistic2(-1.0))
    0.0
    """
    if metapar <= 0.:
        return 0.
    else:
        return optimize.newton(_error_smoothpar_logistic2,
                               .3 * metapar**.84,
                               _smooth_logistic2_derivative,
                               args=(metapar,))

def inverse(fun, x):
  def err(xx):
    return fun(xx)-x;
  return optimize.newton(err, x)

def inverse(fun, x):
  def err(xx):
    return fun(xx)-x;
  return optimize.newton(err, x)

def iterate_two_term_posynomial_linearization_coeff(r, eps):
        """
        Finds the appropriate r, slope, and intercept for a given eps
        :param r: the number of PWL functions
        :param eps: error
        :return: the slope, intercept, and new x
        """
        if r < 2:
            raise Exception('The number of piece-wise sections should two or larger')

        a, b = [], []

        x_intersection = []
        x_tangent = []
        x_intersection.append(np.log(np.exp(eps) - 1))

        i = 1
        while i < r - 1:
            x_old = x_intersection[i - 1]
            try:
                tangent_point = op.newton(LinearizeTwoTermPosynomials.tangent_point_func, x_old + 1, args=(x_old, eps))
                slope = np.exp(tangent_point) / (1 + np.exp(tangent_point))
                intercept = -np.exp(tangent_point) * tangent_point / (1 + np.exp(tangent_point)) + np.log(
                    1 + np.exp(tangent_point))
                intersection_point = op.newton(LinearizeTwoTermPosynomials.intersection_point_func,
                                               tangent_point + 1, args=(slope, intercept, eps))
            except:
                return i, a, b, x_tangent, x_intersection

            x_tangent.append(tangent_point)
            a.append(slope)
            b.append(intercept)
            x_intersection.append(intersection_point)

            i += 1
        return r, a, b, x_tangent, x_intersection

def test_tangent_point_func():
    for _ in xrange(1000):

        eps = np.random.rand() * 0.2

        x_old = - np.random.rand() * np.log(np.exp(0.2) - 1)
        y_old = function(x_old) - eps

        x_tangent = op.newton(LinearizeTwoTermPosynomials.tangent_point_func, x_old + 1, args=(x_old, eps))
        y_tangent = function(x_tangent)

        assert (x_tangent > x_old)

        def tangent_line(x):
            return (y_old - y_tangent) * (x - x_tangent) / (x_old - x_tangent) + y_tangent

        assert (function(x_tangent) - tangent_line(x_tangent) < 0.0001)

        trial_points = list(np.arange(0, 20, 0.01))

        function_points = [function(i) for i in trial_points]
        tangent_points = [tangent_line(i) for i in trial_points]

        difference = [a - b for a, b in zip(function_points, tangent_points)]

        assert (all(i >= 0 for i in difference))

    return

def calc_rate(self):
        def fv_(r, self):
            return (-(self.pv + (1 - self.discount_factor) * self.pva)
                    / self.discount_factor)

        return newton(func=fv_, x0=.05, args=(self,),
                      maxiter=1000, tol=0.00001)

def _findroot(self, func, x0=None, xab=None, **kwds):
        """Find root of `func` by Newton's method if `x0` is given or Brent's
        method if `xab` is given. If neither is given, then
        ``xab=[self.x[0],self.x[-1]]`` and Brent's method is used.

        Parameters
        ----------
        func : callable, must accept a scalar and retun a scalar
        x0 : float
            start guess for Newton's secant method
        xab : sequence of length 2
            start bracket for Brent's method, root must lie in between
        **kwds : 
            passed to scipy root finder (newton() or brentq())

        Returns
        -------
        xx : scalar
            the root of func(x)
        """
        if x0 is not None:
            xx = optimize.newton(func, x0, **kwds)
        else:
            if xab is None:
                xab = [self.x[0], self.x[-1]]
            xx = optimize.brentq(func, xab[0], xab[1], **kwds)
        return xx

def solve_tank_for_V(self):
        '''Method which is called to solve for tank geometry when a certain
        volume is specified. Will be called by the __init__ method if V is set.

        Notes
        -----
        Raises an error if L and either of sideA_a or sideB_a are specified;
        these can only be set once D is known.
        Raises an error if more than one of D, L, or L_over_D are specified.
        Raises an error if the head ratios are not provided.

        Calculates initial guesses assuming no heads are present, and then uses
        fsolve to determine the correct dimentions for the tank.

        Tested, but bugs and limitations are expected here.
        '''
        if self.L and (self.sideA_a or self.sideB_a):
            raise Exception('Cannot specify head sizes when solving for V')
        if (self.D and self.L) or (self.D and self.L_over_D) or (self.L and self.L_over_D):
            raise Exception('Only one of D, L, or L_over_D can be specified\
            when solving for V')
        if ((self.sideA and not self.sideA_a_ratio) or (self.sideB and not self.sideB_a_ratio)):
            raise Exception('When heads are specified, head parameter ratios are required')

        if self.D:
            # Iterate until L is appropriate
            solve_L = lambda L: self._V_solver_error(self.V, self.D, L, self.horizontal, self.sideA, self.sideB, self.sideA_a, self.sideB_a, self.sideA_f, self.sideA_k, self.sideB_f, self.sideB_k, self.sideA_a_ratio, self.sideB_a_ratio)
            Lguess = self.V/(pi/4*self.D**2)
            self.L = float(newton(solve_L, Lguess))
        elif self.L:
            # Iterate until D is appropriate
            solve_D = lambda D: self._V_solver_error(self.V, D, self.L, self.horizontal, self.sideA, self.sideB, self.sideA_a, self.sideB_a, self.sideA_f, self.sideA_k, self.sideB_f, self.sideB_k, self.sideA_a_ratio, self.sideB_a_ratio)
            Dguess = (4*self.V/pi/self.L)**0.5
            self.D = float(newton(solve_D, Dguess))
        else:
            # Use L_over_D until L and D are appropriate
            Lguess = (4*self.V*self.L_over_D**2/pi)**(1/3.)
            solve_L_D = lambda L: self._V_solver_error(self.V, L/self.L_over_D, L, self.horizontal, self.sideA, self.sideB, self.sideA_a, self.sideB_a, self.sideA_f, self.sideA_k, self.sideB_f, self.sideB_k, self.sideA_a_ratio, self.sideB_a_ratio)
            self.L = float(newton(solve_L_D, Lguess))
            self.D = self.L/self.L_over_D

def Prandtl_von_Karman_Nikuradse_numeric(Re):
    def to_solve(f):
        # Good to 1E75, down to 1E-17
        return 1./f**0.5 + 2*log10(2.51/Re/f**0.5)
    return newton(to_solve, 0.000001)

def extract_pci_position(w):

    def theta(x):

        wval = w(fenics.Point(x))

        return wval[2]

    pci_pos = opt.newton(theta, 0.1)

    return pci_pos

def test_melt_toy_pcm__regression():

    w = melt_toy_pcm()

    """
    w = melt_toy_pcm(restart = True, restart_filepath = 'output/test_melt_toy_pcm/restart_t0.02.h5',
        start_time = 0.02,
        output_dir = 'output/test_melt_toy_pcm/restart_t0.02/')
    """


    # Verify against a reference solution.
    pci_y_position_to_check =  0.875

    reference_pci_x_position = 0.226

    def T_minus_T_f(x):

        wval = w.leaf_node()(fenics.Point(x, pci_y_position_to_check))

        return wval[3] - T_f


    pci_x_position = opt.newton(T_minus_T_f, 0.01)

    assert(abs(pci_x_position - reference_pci_x_position) < 1.e-2)

def verify_pci_position(true_pci_position, r, w):

    def theta(x):

        wval = w.leaf_node()(fenics.Point(x))

        return wval[2]

    pci_pos = opt.newton(theta, 0.1)

    assert(abs(pci_pos - true_pci_position) < r)

def L_calc(E_val,D,m,L0):
   E = E_val
   c2 = D*D*E/4.
   L = newton(newton_ang_func,L0,args=(c2,m),tol=1e-8,maxiter=500)
   return L

def pes(D_vec,m,L0,E_start):

   PES = np.zeros((D_vec.shape[0],2))
   i = 0
   for D in D_vec:    
     E_elec = newton(newton_rad_func,E_start,args=(D,m,L0),tol=1e-8,maxiter=200)
     L0 = L_calc(E_elec,D,m,L0)
     E_nuc = 2./D
     E_start = E_elec
     PES[i][0] = D
     PES[i][1] = E_elec+E_nuc
     i += 1

   return PES

def point_find(func, x_range, y_range, args=None):
    y_pts = []
    if args is None:
        args = [None]
    if x_range.shape == ():
        out = []
        for a in args:
            out.append(float(x_range))
        root_temp = []
        for y in y_range:
            try:
                root = sp.newton(func, y, args = tuple(out), tol = 1e-20)
                root_temp = t.check(root, root_temp, 10**(-6))
            except RuntimeError:
                pass
        y_pts.append(root_temp)
    else:
        for x in x_range:
            out = []
            for a in args:
                if a is None:
                    a = x
                out.append(a)
            root_temp = []
            for y in y_range:
                try:
                    root = sp.newton(func, y, args = tuple(out), tol = 1e-20)
                    root_temp = t.check(root, root_temp, 10**(-6))
                except RuntimeError:
                    pass
            y_pts.append(root_temp)
    y_array = t.list_to_array(y_pts)
    return y_array

# Attempts to find a root that is close to a root inputted into the finder.
# Arguments:
# func - The function. Must depend on a single variable.
# root_list - List of roots of the function found so far. Can be empty.
# x - The x coordinate of the starting point.
# y - The y coordinate of the starting point.
# step_size - The step size to be used in calculating the next point.
# x_values -

def BAWPremium( IsCall, Fwd, Strike, Vol, Texp, rd, rf ):
    '''
    Early exercise premium for american spot options based on Barone-Adesi, Whaley
    approximation formula. To compute the options prices, this needs to be added to
    the european options prices.
    '''
    if Texp <= 0. or Vol <=0:
        return 0.

    T   = Texp
    K   = Strike
    D   = exp( -rd * Texp)
    Dq  = exp( -rf * Texp )
    Phi = (IsCall and 1 or -1)

    k = (D==1.) and 2./Vol/Vol or 2.* rf/Vol/Vol/(1-D)
    # the expression in the middle is really the expression on the right in the limit D -> 1
    # note that lim_{D -> 1.} log(D)/(1-D) = -1.

    beta = 2.*(rd-rf)/Vol/Vol
    if Phi == 1:
        q2=(-(beta-1.)+sqrt((beta-1.)**2+4.*k))/2.
        def EarlyExerBdry( eeb ):
            x = D*BSFwd(True,eeb,K,Vol,T) + (1.-Dq*cnorm(fd1(eeb,K,Vol,T)))*eeb/q2 - eeb + K
            return x

        eeBdry = D*BSFwd(True,Fwd,K,Vol,T) + (1.-Dq*cnorm(fd1(Fwd,K,Vol,T)))* Fwd/q2 + K
        eeBdry = newton(EarlyExerBdry,eeBdry)
        if Fwd >= eeBdry:
            eePrem = -D*BSFwd(True,Fwd,K,Vol,T) + Fwd - K
        else:
            A2=(eeBdry/q2)*(1.-Dq*cnorm(fd1(eeBdry,K,Vol,T)))
            eePrem = A2 * pow(Fwd/eeBdry,q2)
    elif Phi == -1:
        q1=(-(beta-1.)-sqrt((beta-1.)**2+4.*k))/2.
        def EarlyExerBdry( eeb ):
            x = D*BSFwd(False,eeb,K,Vol,T) - (1.-Dq*cnorm(-fd1(eeb,K,Vol,T)))*eeb/q1 + eeb - K
            return x

        eeBdry = -D*BSFwd(False,Fwd,K,Vol,T) + (1.-Dq*cnorm(-fd1(Fwd,K,Vol,T)))*Fwd/q1 + K
        eeBdry = brentq(EarlyExerBdry,1e-12, K)
        if Fwd <= eeBdry:
            eePrem = -D*BSFwd(False,Fwd,K,Vol,T) + K - Fwd
        else:
            A1=-(eeBdry/q1)*(1.-Dq*cnorm(-fd1(eeBdry,K,Vol,T)))
            eePrem = A1 * pow(Fwd/eeBdry,q1)
    else:
        raise ValueError, 'option type can only be call or put'

    return eePrem

def test_valve_coefficients():
    Cv = Kv_to_Cv(2)
    assert_allclose(Cv, 2.3121984567073133)
    Kv = Cv_to_Kv(2.312)
    assert_allclose(Kv, 1.9998283393826013)
    K = Kv_to_K(2.312, .015)
    assert_allclose(K, 15.15337460039990)
    Kv = K_to_Kv(15.15337460039990, .015)
    assert_allclose(Kv, 2.312)

    # Two way conversions
    K = Cv_to_K(2.712, .015)
    assert_allclose(K, 14.719595348352552)
    assert_allclose(K, Kv_to_K(Cv_to_Kv(2.712), 0.015))

    Cv = K_to_Cv(14.719595348352552, .015)
    assert_allclose(Cv, 2.712)
    assert_allclose(Cv, Kv_to_Cv(K_to_Kv(14.719595348352552, 0.015)))

    # Code to generate the Kv Cv conversion factor
    # Round 1 trip; randomly assume Kv = 12, rho = 900; they can be anything
    # an tit still works
    dP = 1E5
    rho = 900.
    Kv = 12.
    Q = Kv/3600.
    D = .01
    V = Q/(pi/4*D**2)
    K = dP/(.5*rho*V*V)
    good_K = K

    def to_solve(x):
        from scipy.constants import gallon, minute, hour, psi
        conversion = gallon/minute*hour # from gpm to m^3/hr
        dP = 1*psi
        Cv = Kv*x*conversion
        Q = Cv/3600
        D = .01
        V = Q/(pi/4*D**2)
        K = dP/(.5*rho*V*V)
        return K - good_K

    from scipy.optimize import newton

    ans = newton(to_solve, 1.2)
    assert_allclose(ans, 1.1560992283536566)

def R_xi(E_vec,L0):

   traj = np.zeros([x_rad.shape[0]+1,E_vec.shape[0]])

   n=0
   for E in E_vec:

      c2 = D*D*E/4.

      L = newton(newton_ang_func,L0,args=(c2,m),tol=1e-8,maxiter=200)

      slope = -(-L+m*(m+1.)+2.*D+c2)/(2.*(m+1.))
      z0 = [1+step*slope,slope]

      z = odeint(g,z0,x_rad,args=(c2,L,D,m))

      temp=pow(x_rad,2.0)-1.
      temp=pow(temp,m/2.)

      zz=temp*z[:,0]

      first_zz = np.array([1])
      zz=np.append(first_zz, zz)

      traj[:,n] = zz

      n += 1

   xt = np.append(np.array([1]),x_rad)

   figure = plt.figure(figsize=(12, 11))
   plt.plot(xt,traj, linewidth=2.0, label = '')
   plt.ylabel('R($\\xi$)')#, fontdict=font)
   plt.xlabel('$\\xi$')#, fontdict=font)
   #plt.xlim(1.0,10.0)
   #plt.ylim(-1.0,1.0)
   plt.locator_params(axis='x', nbins=10)
   plt.locator_params(axis='y', nbins=10)
   plt.tick_params(axis='x', pad=15)
   #plt.legend(loc=1)

   plt.show()
   plt.close()