1 | from sage.rings.complex_double cimport ComplexDoubleElement |
---|
2 | cimport numpy as cnumpy |
---|
3 | from sage.plot.primitive import GraphicPrimitive |
---|
4 | from sage.misc.decorators import options |
---|
5 | from sage.misc.misc import srange |
---|
6 | from sage.symbolic.ring import var |
---|
7 | |
---|
8 | cdef 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 | |
---|
16 | def 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 | |
---|
73 | cdef extern from "math.h": |
---|
74 | double atan(double) |
---|
75 | double log(double) |
---|
76 | double sqrt(double) |
---|
77 | double PI |
---|
78 | |
---|
79 | cdef 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 | |
---|
110 | def newt(func): |
---|
111 | vari = func.args()[0] |
---|
112 | fprime = diff(func,vari) |
---|
113 | return vari-func/fprime |
---|
114 | |
---|
115 | |
---|
116 | |
---|
117 | class 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 | |
---|
177 | cdef 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') |
---|
202 | def 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 |
---|