Ticket #11837: newton_basins_with_iterations.spyx

File newton_basins_with_iterations.spyx, 8.3 KB (added by kcrisman, 8 years ago)

Incorporates Simon's improvements and adds naive keeping track of iterations

Line 
1from sage.rings.complex_double cimport ComplexDoubleElement
2cimport numpy as cnumpy
3from sage.plot.primitive import GraphicPrimitive
4from sage.misc.decorators import options
5from sage.misc.misc import srange
6from sage.symbolic.ring import var
7
8cdef inline ComplexDoubleElement new_CDF_element(double w, double v):
9    cdef ComplexDoubleElement z = <ComplexDoubleElement>PY_NEW(ComplexDoubleElement)
10    z._complex.dat[0] = w
11    z._complex.dat[1] = v
12    return z
13
14
15
16def complex_to_rgb(roots,z_values):
17    """
18    INPUT:
19
20    - ``z_values`` -- A grid of complex numbers, as a list of lists
21   
22    OUTPUT:
23
24    An `N \\times M \\times 3` floating point Numpy array ``X``, where
25    ``X[i,j]`` is an (r,g,b) tuple.
26   
27    EXAMPLES::
28
29        sage: from sage.plot.complex_plot import complex_to_rgb
30        sage: complex_to_rgb([[0, 1, 1000]])
31        array([[[ 0.        ,  0.        ,  0.        ],
32                [ 0.77172568,  0.        ,  0.        ],
33                [ 1.        ,  0.64421177,  0.64421177]]])
34        sage: complex_to_rgb([[0, 1j, 1000j]])
35        array([[[ 0.        ,  0.        ,  0.        ],
36                [ 0.38586284,  0.77172568,  0.        ],
37                [ 0.82210588,  1.        ,  0.64421177]]])
38    """
39    import numpy
40    cdef unsigned int i, j, imax, jmax, n
41    cdef ComplexDoubleElement zz
42    from sage.rings.complex_double import CDF
43    from sage.plot.colors import rainbow
44   
45    imax = len(z_values)
46    jmax = len(z_values[0])
47    cdef cnumpy.ndarray[cnumpy.float_t, ndim=3, mode='c'] rgb = numpy.empty(dtype=numpy.float, shape=(imax, jmax, 3))
48
49    n = len(roots)
50    R = rainbow(n)
51    R = [Color(col).rgb() for col in R]
52   
53    sig_on()
54    for i from 0 <= i < imax:
55   
56        row = z_values[i]
57        for j from 0 <= j < jmax:
58       
59            zz = row[j][0]
60            alpha = mag_to_lightness(row[j][1])
61
62            try:
63                color = R[roots.index(zz)]
64                rgb[i, j, 0] = color[0]*alpha
65                rgb[i, j, 1] = color[1]*alpha
66                rgb[i, j, 2] = color[2]*alpha
67
68            except:
69                 print zz               
70    sig_off()
71    return rgb
72
73cdef extern from "math.h":
74    double atan(double)
75    double log(double)
76    double sqrt(double)
77    double PI
78
79cdef inline double mag_to_lightness(int r):
80    """
81    Tweak this to adjust how the magnitude affects the color.
82    For instance, changing ``sqrt(r)`` to ``r`` will cause
83    anything near a zero to be much darker and poles to be
84    much lighter, while ``r**(.25)`` would cause the reverse
85    effect.
86   
87    INPUT:
88
89    - ``r`` - a non-negative real number
90
91    OUTPUT:
92
93    A value between `-1` (black) and `+1` (white), inclusive.
94
95    EXAMPLES:
96
97    This tests it implicitly::
98
99        sage: from sage.plot.complex_plot import complex_to_rgb
100        sage: complex_to_rgb([[0, 1, 10]])
101        array([[[ 0.        ,  0.        ,  0.        ],
102                [ 0.77172568,  0.        ,  0.        ],
103                [ 1.        ,  0.22134776,  0.22134776]]])
104    """
105    return .5*atan(log(r+1)) * (4/PI)
106
107
108
109
110def newt(func):
111    vari = func.args()[0]
112    fprime = diff(func,vari)
113    return vari-func/fprime
114
115
116
117class BasinPlot(GraphicPrimitive):
118    def __init__(self, roots, z_values, xrange, yrange, options):
119        """
120        TESTS::
121
122            sage: p = complex_plot(lambda z: z^2-1, (-2, 2), (-2, 2))
123        """
124        self.roots = roots
125        self.xrange = xrange
126        self.yrange = yrange
127        self.z_values = z_values
128        self.x_count = len(z_values)
129        self.y_count = len(z_values[0])
130        self.rgb_data = complex_to_rgb(roots,z_values)
131        GraphicPrimitive.__init__(self, options)
132
133    def get_minmax_data(self):
134        """
135        Returns a dictionary with the bounding box data.
136       
137        EXAMPLES::
138
139            sage: p = complex_plot(lambda z: z, (-1, 2), (-3, 4))
140            sage: sorted(p.get_minmax_data().items())
141            [('xmax', 2.0), ('xmin', -1.0), ('ymax', 4.0), ('ymin', -3.0)]
142        """
143        from sage.plot.plot import minmax_data
144        return minmax_data(self.xrange, self.yrange, dict=True)
145
146    def _allowed_options(self):
147        """
148        TESTS::
149
150            sage: isinstance(complex_plot(lambda z: z, (-1,1), (-1,1))[0]._allowed_options(), dict)
151            True
152        """
153        return {'plot_points':'How many points to use for plotting precision',
154                'interpolation':'What interpolation method to use'}
155
156    def _repr_(self):
157        """
158        TESTS::
159
160            sage: isinstance(complex_plot(lambda z: z, (-1,1), (-1,1))[0]._repr_(), str)
161            True
162        """
163        return "BasinPlot defined by a %s x %s data grid"%(self.x_count, self.y_count)
164
165    def _render_on_subplot(self, subplot):
166        """
167        TESTS::
168
169            sage: complex_plot(lambda x: x^2, (-5, 5), (-5, 5))
170        """
171        options = self.options()
172        x0,x1 = float(self.xrange[0]), float(self.xrange[1])
173        y0,y1 = float(self.yrange[0]), float(self.yrange[1])
174        subplot.imshow(self.rgb_data, origin='lower', extent=(x0,x1,y0,y1), interpolation=options['interpolation'])
175
176
177cdef class RootFinder:
178    cdef list roots
179    cdef object newtf
180    cdef ComplexDoubleElement cutoff
181    def __init__(self, roots, newtf):
182        self.newtf = newtf
183        self.roots = roots
184        self.cutoff = CDF(2)/3
185    cpdef which_root(self,Varia):
186        cdef int counter
187        cdef ComplexDoubleElement root
188        cdef ComplexDoubleElement varia = Varia
189        for counter from 1<=counter<20:
190            varia = self.newtf(varia)
191            #print z
192            for root in self.roots:
193                if (<ComplexDoubleElement>varia._sub_(root)).abs()<self.cutoff:
194                    return root,counter
195        cdef list ls = []
196        for root in self.roots:
197            ls.append(varia._sub_(root).abs())
198        return self.roots[ls.index(min(ls))],20
199
200
201@options(plot_points=3, interpolation='nearest')
202def basin_plot(list roots, xrange, yrange, **options):
203    r"""   
204    ``complex_plot`` takes a complex function of one variable,
205    `f(z)` and plots output of the function over the specified
206    ``xrange`` and ``yrange`` as demonstrated below. The magnitude of the
207    output is indicated by the brightness (with zero being black and
208    infinity being white) while the argument is represented by the
209    hue (with red being positive real, and increasing through orange,
210    yellow, ... as the argument increases).
211
212    ``complex_plot(f, (xmin, xmax), (ymin, ymax), ...)``
213
214    INPUT:
215
216    - ``f`` -- a function of a single complex value `x + iy`
217
218    - ``(xmin, xmax)`` -- 2-tuple, the range of ``x`` values
219
220    - ``(ymin, ymax)`` -- 2-tuple, the range of ``y`` values
221
222    The following inputs must all be passed in as named parameters:
223
224    - ``plot_points`` -- integer (default: 100); number of points to plot
225      in each direction of the grid
226
227    - ``interpolation`` -- string (default: ``'catrom'``), the interpolation
228      method to use: ``'bilinear'``, ``'bicubic'``, ``'spline16'``,
229      ``'spline36'``, ``'quadric'``, ``'gaussian'``, ``'sinc'``,
230      ``'bessel'``, ``'mitchell'``, ``'lanczos'``, ``'catrom'``,
231      ``'hermite'``, ``'hanning'``, ``'hamming'``, ``'kaiser'``
232       
233
234    """
235    from sage.plot.plot import Graphics
236    from sage.plot.misc import setup_for_eval_on_grid
237    from sage.ext.fast_callable import fast_callable
238    from sage.rings.complex_double import CDF
239    roots = [CDF(temp) for temp in roots]
240
241    x = var('x')
242    prefunc = prod([(x-root) for root in roots])
243    f = fast_callable(prefunc, domain=CDF, expect_one_var=True)
244    newtf = fast_callable(newt(prefunc), domain=CDF, expect_one_var=True)
245
246    cdef ComplexDoubleElement cutoff = CDF(2./3)
247
248    cdef RootFinder Finder = RootFinder(roots,newtf)
249    f1 = Finder.which_root
250   
251    cdef double t, u
252    ignore, ranges = setup_for_eval_on_grid([], [xrange, yrange], options['plot_points'])
253    xrange,yrange=[r[:2] for r in ranges]
254    sig_on()
255    z_values = [[  f1(new_CDF_element(t, u)) for t in srange(*ranges[0], include_endpoint=True)]
256                                            for u in srange(*ranges[1], include_endpoint=True)]
257#    print z_values
258    sig_off()
259    g = Graphics()
260    g._set_extra_kwds(Graphics._extract_kwds_for_show(options, ignore=['xmin', 'xmax']))
261    g.add_primitive(BasinPlot(roots, z_values, xrange, yrange, options))
262    return g