Ticket #11837: newton_basins_with_iterations_and_double.spyx

File newton_basins_with_iterations_and_double.spyx, 8.6 KB (added by gh-bryangingechen, 19 months ago)

A version of "newton_basins_with_iterations_and_double.spyx.txt​ (8.7 KB) - added by SimonKing?" minimally updated to work with SageMath 8.2

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