# Ticket #11837: newton_basins_with_iterations.spyx

File newton_basins_with_iterations.spyx, 8.3 KB (added by , 9 years ago) |
---|

Line | |
---|---|

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 |