a Sage @interact for root finding with Newton's method

newton.sage

"""
My Newton's method Sage @interact. A reworked version of
http://wiki.sagemath.org/interact/calculus#Newton-Raphson_Root_Finding
by Neal Holtz, and which was in turn inspired by Marshall Hampton's
tangent line grapher.

This code is available as a published worksheet at
https://sagenb.kaist.ac.kr:8066/home/pub/53 and on the Sage cell server at
http://aleph.sagemath.org/?q=81e86e25-b27c-40c6-ac01-b7489b55e84e. (Note
that right now, you need to use

    step = ButtonBar(['Next', 'Prev', 'Zoom', 'Reset'])

to get the @interact to work on the cell server.)

You may modify and distribute this code under the terms of the
Creative Commons Attribution-ShareAlike 3.0 license:
http://creativecommons.org/licenses/by-sa/3.0/.

Dan Drake (http://mathsci.kaist.ac.kr/~drake)
"""

def newton(f, a, df=None):
    """
    return the next approximation after 'a' using Newton's method.
    """
    if not df:
        df = diff(f)
    return a - f(a)/df(a)

def newtonplot(f, xmin, xmax, points):
    """
    return the complete plot for the function f and the points given in
    `points`.
    """
    ret = plot(f, xmin, xmax)
    ret += line([(points[0], 0), (points[0], f(points[0]))], linestyle=':',
                color='red')
    ret += text('\n$x_0$', (points[0], 0), color='red', 
                vertical_alignment="bottom" if f(points[0]) < 0 else "top")
    for n, pair in enumerate(zip(points, points[1:])):
        curr, next = pair
        ret += newtonappend(f, curr, next, label='\n$x_{0}$'.format(n+1))
    return ret
    
def newtonappend(f, prev, curr, label=''):
    """
    return the new line segments going from prev to curr: we join the points

    (prev, f(prev)) to (curr, 0) to (curr, f(curr)).
    """
    p = line([(prev, f(prev)), (curr, 0)], color='red')
    p += line([(curr, 0), (curr, f(curr))], linestyle=':', color='red')
    if label:
        p += text(label, (curr, 0), color='red',
                  vertical_alignment="bottom" if f(curr) < 0 else "top" )
    return p

# Now the actual @interact stuff:

state = None
# we need user_state to distinguish between xmin/xmax changing when we
# zoom (don't reset) and changing from new input from the slider (yes,
# reset) -- as well as determining when f changes, since we munge f by
# running fast_float() on it.
user_state = None

@interact
def _(f = input_box(default=8*sin(x)*exp(-x)-1, label='f(x)'), 
      xmin = input_box(default=0), 
      xmax = input_box(default=4*pi), 
      x0 = input_box(default=3, label='x0'),
      step = ['Next', 'Prev', 'Zoom', 'Reset']):
    global state, user_state
    if (step == 'Reset' or 
       state is None or
       state['points'][0] != x0 or
       user_state['f'] != f or
       user_state['interval'] != (xmin, xmax)):
        state = {'f': fast_float(f, x), 'df': fast_float(diff(f), x),
                 'xmin': xmin, 'xmax': xmax, 'points': [x0]}
        
        p = newtonplot(state['f'], xmin, xmax, [x0])
        state['plot'] = p
        # find_minimum_on_interval doesn't work so well, see #2607
        state['ymin'] = p.ymin()
        state['ymax'] = p.ymax()
        user_state = {'f': f, 'interval': (xmin, xmax)}
    elif step == 'Next':
        next = newton(state['f'], state['points'][-1], df=state['df'])
        divbyzero = 10^5
        diverge = 1/10
        width = state['xmax'] - state['xmin']
        # are we *way* outside the interval?
        if not state['xmin'] - divbyzero * width <= next <= state['xmax'] + divbyzero * width:
            print 'Next point is supposedly {0}; derivative is probably zero at/near {1}!'.format(next, state['points'][-1])
            print 'Cowardly refusing to add that point.\n'
        # are we just a little bit outside the interval?
        elif not state['xmin'] - diverge * width <= next <= state['xmax'] + diverge * width:
            print 'Next point ({0}) is not within +/- 10% of given interval; possible divergence?'.format(next)
            print 'Cowardly refusing to add that point.\n'
        # okay, we're close enough:
        else:
            state['plot'] += newtonappend(state['f'], state['points'][-1], next,
                                          label='\n$x_{0}$'.format(len(state['points'])))
            state['points'].append(next)
    elif step == 'Prev':
        state['points'].pop()
        # sadly, "-=" does not work on graphics objects
        state['plot'] = newtonplot(state['f'], state['xmin'], state['xmax'], state['points'])
    elif step == 'Zoom':
        w = (state['xmax'] - state['xmin']) / 4
        state['xmax'] = state['points'][-1] + w
        state['xmin'] = state['points'][-1] - w
        p = plot(state['f'], (state['xmin'], state['xmax']))
        # new interval may not be subset of previous, so make sure min
        # isn't smaller and max isn't bigger
        state['ymin'] = max(state['ymin'], p.ymin())
        state['ymax'] = min(state['ymax'], p.ymax())
 
    curr = state['points'][-1]
    print 'Current estimate: {0}; f(that point): {1}'.format(curr, state['f'](curr))
    state['plot'].show(xmin=state['xmin'], xmax=state['xmax'], ymin=state['ymin'], ymax=state['ymax'])
    print 'n, x_n, f(x_n):'
    for n, x_ in enumerate(state['points']):
        print '{0}, {1}, {2}'.format(n, x_, state['f'](x_))