# Ticket #11837: newton_basins_with_iterations_and_double.spyx.txt

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

Line | |
---|---|

1 | %cython |

2 | from sage.rings.complex_double cimport ComplexDoubleElement |

3 | cimport numpy as cnumpy |

4 | from sage.plot.primitive import GraphicPrimitive |

5 | from sage.misc.decorators import options |

6 | from sage.misc.misc import srange |

7 | from sage.symbolic.ring import var |

8 | |

9 | cdef 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 | |

17 | def 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 | |

74 | cdef extern from "math.h": |

75 | double atan(double) |

76 | double log(double) |

77 | double sqrt(double) |

78 | double PI |

79 | |

80 | cdef 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 | |

111 | def newt(func): |

112 | vari = func.args()[0] |

113 | fprime = diff(func,vari) |

114 | return vari-func/fprime |

115 | |

116 | |

117 | |

118 | class 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 | |

178 | cdef 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') |

207 | def 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 |