Ticket #11837: newton_basins_with_iterations.spyx

from sage.rings.complex_double cimport ComplexDoubleElement
cimport numpy as cnumpy
from sage.plot.primitive import GraphicPrimitive
from sage.misc.decorators import options
from sage.misc.misc import srange
from sage.symbolic.ring import var

cdef inline ComplexDoubleElement new_CDF_element(double w, double v):
    cdef ComplexDoubleElement z = <ComplexDoubleElement>PY_NEW(ComplexDoubleElement)
    z._complex.dat[0] = w
    z._complex.dat[1] = v
    return z



def complex_to_rgb(roots,z_values):
    """
    INPUT: 

    - ``z_values`` -- A grid of complex numbers, as a list of lists
    
    OUTPUT:

    An `N \\times M \\times 3` floating point Numpy array ``X``, where
    ``X[i,j]`` is an (r,g,b) tuple. 
    
    EXAMPLES::

        sage: from sage.plot.complex_plot import complex_to_rgb
        sage: complex_to_rgb([[0, 1, 1000]])
        array([[[ 0.        ,  0.        ,  0.        ],
                [ 0.77172568,  0.        ,  0.        ],
                [ 1.        ,  0.64421177,  0.64421177]]])
        sage: complex_to_rgb([[0, 1j, 1000j]])
        array([[[ 0.        ,  0.        ,  0.        ],
                [ 0.38586284,  0.77172568,  0.        ],
                [ 0.82210588,  1.        ,  0.64421177]]])
    """
    import numpy
    cdef unsigned int i, j, imax, jmax, n
    cdef ComplexDoubleElement zz
    from sage.rings.complex_double import CDF
    from sage.plot.colors import rainbow
    
    imax = len(z_values)
    jmax = len(z_values[0])
    cdef cnumpy.ndarray[cnumpy.float_t, ndim=3, mode='c'] rgb = numpy.empty(dtype=numpy.float, shape=(imax, jmax, 3))

    n = len(roots)
    R = rainbow(n)
    R = [Color(col).rgb() for col in R]
    
    sig_on()
    for i from 0 <= i < imax:
    
        row = z_values[i]
        for j from 0 <= j < jmax:
        
            zz = row[j][0]
            alpha = mag_to_lightness(row[j][1])

            try:
                color = R[roots.index(zz)]
                rgb[i, j, 0] = color[0]*alpha
                rgb[i, j, 1] = color[1]*alpha
                rgb[i, j, 2] = color[2]*alpha

            except:
                 print zz                
    sig_off()
    return rgb

cdef extern from "math.h":
    double atan(double)
    double log(double)
    double sqrt(double)
    double PI

cdef inline double mag_to_lightness(int r):
    """
    Tweak this to adjust how the magnitude affects the color.
    For instance, changing ``sqrt(r)`` to ``r`` will cause
    anything near a zero to be much darker and poles to be
    much lighter, while ``r**(.25)`` would cause the reverse
    effect.
    
    INPUT:

    - ``r`` - a non-negative real number

    OUTPUT:

    A value between `-1` (black) and `+1` (white), inclusive.

    EXAMPLES:

    This tests it implicitly::

        sage: from sage.plot.complex_plot import complex_to_rgb
        sage: complex_to_rgb([[0, 1, 10]])
        array([[[ 0.        ,  0.        ,  0.        ],
                [ 0.77172568,  0.        ,  0.        ],
                [ 1.        ,  0.22134776,  0.22134776]]])
    """
    return .5*atan(log(r+1)) * (4/PI)




def newt(func):
    vari = func.args()[0]
    fprime = diff(func,vari)
    return vari-func/fprime



class BasinPlot(GraphicPrimitive):
    def __init__(self, roots, z_values, xrange, yrange, options):
        """
        TESTS::

            sage: p = complex_plot(lambda z: z^2-1, (-2, 2), (-2, 2))
        """
        self.roots = roots
        self.xrange = xrange
        self.yrange = yrange
        self.z_values = z_values
        self.x_count = len(z_values)
        self.y_count = len(z_values[0])
        self.rgb_data = complex_to_rgb(roots,z_values)
        GraphicPrimitive.__init__(self, options)

    def get_minmax_data(self):
        """
        Returns a dictionary with the bounding box data.
        
        EXAMPLES::

            sage: p = complex_plot(lambda z: z, (-1, 2), (-3, 4))
            sage: sorted(p.get_minmax_data().items())
            [('xmax', 2.0), ('xmin', -1.0), ('ymax', 4.0), ('ymin', -3.0)]
        """
        from sage.plot.plot import minmax_data
        return minmax_data(self.xrange, self.yrange, dict=True)

    def _allowed_options(self):
        """
        TESTS::

            sage: isinstance(complex_plot(lambda z: z, (-1,1), (-1,1))[0]._allowed_options(), dict)
            True
        """
        return {'plot_points':'How many points to use for plotting precision',
                'interpolation':'What interpolation method to use'}

    def _repr_(self):
        """
        TESTS::

            sage: isinstance(complex_plot(lambda z: z, (-1,1), (-1,1))[0]._repr_(), str)
            True
        """
        return "BasinPlot defined by a %s x %s data grid"%(self.x_count, self.y_count)

    def _render_on_subplot(self, subplot):
        """
        TESTS::

            sage: complex_plot(lambda x: x^2, (-5, 5), (-5, 5))
        """
        options = self.options()
        x0,x1 = float(self.xrange[0]), float(self.xrange[1])
        y0,y1 = float(self.yrange[0]), float(self.yrange[1])
        subplot.imshow(self.rgb_data, origin='lower', extent=(x0,x1,y0,y1), interpolation=options['interpolation'])


cdef class RootFinder:
    cdef list roots
    cdef object newtf
    cdef ComplexDoubleElement cutoff
    def __init__(self, roots, newtf):
        self.newtf = newtf
        self.roots = roots
        self.cutoff = CDF(2)/3
    cpdef which_root(self,Varia):
        cdef int counter
        cdef ComplexDoubleElement root
        cdef ComplexDoubleElement varia = Varia
        for counter from 1<=counter<20:
            varia = self.newtf(varia)
            #print z
            for root in self.roots:
                if (<ComplexDoubleElement>varia._sub_(root)).abs()<self.cutoff:
                    return root,counter
        cdef list ls = []
        for root in self.roots:
            ls.append(varia._sub_(root).abs())
        return self.roots[ls.index(min(ls))],20


@options(plot_points=3, interpolation='nearest')
def basin_plot(list roots, xrange, yrange, **options):
    r"""    
    ``complex_plot`` takes a complex function of one variable, 
    `f(z)` and plots output of the function over the specified 
    ``xrange`` and ``yrange`` as demonstrated below. The magnitude of the 
    output is indicated by the brightness (with zero being black and
    infinity being white) while the argument is represented by the 
    hue (with red being positive real, and increasing through orange, 
    yellow, ... as the argument increases). 

    ``complex_plot(f, (xmin, xmax), (ymin, ymax), ...)``

    INPUT:

    - ``f`` -- a function of a single complex value `x + iy`

    - ``(xmin, xmax)`` -- 2-tuple, the range of ``x`` values

    - ``(ymin, ymax)`` -- 2-tuple, the range of ``y`` values

    The following inputs must all be passed in as named parameters:

    - ``plot_points`` -- integer (default: 100); number of points to plot
      in each direction of the grid

    - ``interpolation`` -- string (default: ``'catrom'``), the interpolation
      method to use: ``'bilinear'``, ``'bicubic'``, ``'spline16'``,
      ``'spline36'``, ``'quadric'``, ``'gaussian'``, ``'sinc'``,
      ``'bessel'``, ``'mitchell'``, ``'lanczos'``, ``'catrom'``,
      ``'hermite'``, ``'hanning'``, ``'hamming'``, ``'kaiser'``
        

    """
    from sage.plot.plot import Graphics
    from sage.plot.misc import setup_for_eval_on_grid
    from sage.ext.fast_callable import fast_callable
    from sage.rings.complex_double import CDF
    roots = [CDF(temp) for temp in roots]

    x = var('x')
    prefunc = prod([(x-root) for root in roots])
    f = fast_callable(prefunc, domain=CDF, expect_one_var=True)
    newtf = fast_callable(newt(prefunc), domain=CDF, expect_one_var=True)

    cdef ComplexDoubleElement cutoff = CDF(2./3)

    cdef RootFinder Finder = RootFinder(roots,newtf)
    f1 = Finder.which_root
    
    cdef double t, u
    ignore, ranges = setup_for_eval_on_grid([], [xrange, yrange], options['plot_points'])
    xrange,yrange=[r[:2] for r in ranges]
    sig_on()
    z_values = [[  f1(new_CDF_element(t, u)) for t in srange(*ranges[0], include_endpoint=True)]
                                            for u in srange(*ranges[1], include_endpoint=True)]
#    print z_values
    sig_off()
    g = Graphics()
    g._set_extra_kwds(Graphics._extract_kwds_for_show(options, ignore=['xmin', 'xmax']))
    g.add_primitive(BasinPlot(roots, z_values, xrange, yrange, options))
    return g