Ticket #11837: newton_basins_with_iterations_and_double.spyx.txt

File newton_basins_with_iterations_and_double.spyx.txt, 8.7 KB (added by SimonKing, 8 years ago)

Naive keeping track of iterations, using a simplified distance test

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