| """ |
| Plotting (requires matplotlib) |
| """ |
|
|
| from colorsys import hsv_to_rgb, hls_to_rgb |
| from .libmp import NoConvergence |
| from .libmp.backend import xrange |
|
|
| class VisualizationMethods(object): |
| plot_ignore = (ValueError, ArithmeticError, ZeroDivisionError, NoConvergence) |
|
|
| def plot(ctx, f, xlim=[-5,5], ylim=None, points=200, file=None, dpi=None, |
| singularities=[], axes=None): |
| r""" |
| Shows a simple 2D plot of a function `f(x)` or list of functions |
| `[f_0(x), f_1(x), \ldots, f_n(x)]` over a given interval |
| specified by *xlim*. Some examples:: |
| |
| plot(lambda x: exp(x)*li(x), [1, 4]) |
| plot([cos, sin], [-4, 4]) |
| plot([fresnels, fresnelc], [-4, 4]) |
| plot([sqrt, cbrt], [-4, 4]) |
| plot(lambda t: zeta(0.5+t*j), [-20, 20]) |
| plot([floor, ceil, abs, sign], [-5, 5]) |
| |
| Points where the function raises a numerical exception or |
| returns an infinite value are removed from the graph. |
| Singularities can also be excluded explicitly |
| as follows (useful for removing erroneous vertical lines):: |
| |
| plot(cot, ylim=[-5, 5]) # bad |
| plot(cot, ylim=[-5, 5], singularities=[-pi, 0, pi]) # good |
| |
| For parts where the function assumes complex values, the |
| real part is plotted with dashes and the imaginary part |
| is plotted with dots. |
| |
| .. note :: This function requires matplotlib (pylab). |
| """ |
| if file: |
| axes = None |
| fig = None |
| if not axes: |
| import pylab |
| fig = pylab.figure() |
| axes = fig.add_subplot(111) |
| if not isinstance(f, (tuple, list)): |
| f = [f] |
| a, b = xlim |
| colors = ['b', 'r', 'g', 'm', 'k'] |
| for n, func in enumerate(f): |
| x = ctx.arange(a, b, (b-a)/float(points)) |
| segments = [] |
| segment = [] |
| in_complex = False |
| for i in xrange(len(x)): |
| try: |
| if i != 0: |
| for sing in singularities: |
| if x[i-1] <= sing and x[i] >= sing: |
| raise ValueError |
| v = func(x[i]) |
| if ctx.isnan(v) or abs(v) > 1e300: |
| raise ValueError |
| if hasattr(v, "imag") and v.imag: |
| re = float(v.real) |
| im = float(v.imag) |
| if not in_complex: |
| in_complex = True |
| segments.append(segment) |
| segment = [] |
| segment.append((float(x[i]), re, im)) |
| else: |
| if in_complex: |
| in_complex = False |
| segments.append(segment) |
| segment = [] |
| if hasattr(v, "real"): |
| v = v.real |
| segment.append((float(x[i]), v)) |
| except ctx.plot_ignore: |
| if segment: |
| segments.append(segment) |
| segment = [] |
| if segment: |
| segments.append(segment) |
| for segment in segments: |
| x = [s[0] for s in segment] |
| y = [s[1] for s in segment] |
| if not x: |
| continue |
| c = colors[n % len(colors)] |
| if len(segment[0]) == 3: |
| z = [s[2] for s in segment] |
| axes.plot(x, y, '--'+c, linewidth=3) |
| axes.plot(x, z, ':'+c, linewidth=3) |
| else: |
| axes.plot(x, y, c, linewidth=3) |
| axes.set_xlim([float(_) for _ in xlim]) |
| if ylim: |
| axes.set_ylim([float(_) for _ in ylim]) |
| axes.set_xlabel('x') |
| axes.set_ylabel('f(x)') |
| axes.grid(True) |
| if fig: |
| if file: |
| pylab.savefig(file, dpi=dpi) |
| else: |
| pylab.show() |
|
|
| def default_color_function(ctx, z): |
| if ctx.isinf(z): |
| return (1.0, 1.0, 1.0) |
| if ctx.isnan(z): |
| return (0.5, 0.5, 0.5) |
| pi = 3.1415926535898 |
| a = (float(ctx.arg(z)) + ctx.pi) / (2*ctx.pi) |
| a = (a + 0.5) % 1.0 |
| b = 1.0 - float(1/(1.0+abs(z)**0.3)) |
| return hls_to_rgb(a, b, 0.8) |
|
|
| blue_orange_colors = [ |
| (-1.0, (0.0, 0.0, 0.0)), |
| (-0.95, (0.1, 0.2, 0.5)), |
| (-0.5, (0.0, 0.5, 1.0)), |
| (-0.05, (0.4, 0.8, 0.8)), |
| ( 0.0, (1.0, 1.0, 1.0)), |
| ( 0.05, (1.0, 0.9, 0.3)), |
| ( 0.5, (0.9, 0.5, 0.0)), |
| ( 0.95, (0.7, 0.1, 0.0)), |
| ( 1.0, (0.0, 0.0, 0.0)), |
| ( 2.0, (0.0, 0.0, 0.0)), |
| ] |
|
|
| def phase_color_function(ctx, z): |
| if ctx.isinf(z): |
| return (1.0, 1.0, 1.0) |
| if ctx.isnan(z): |
| return (0.5, 0.5, 0.5) |
| pi = 3.1415926535898 |
| w = float(ctx.arg(z)) / pi |
| w = max(min(w, 1.0), -1.0) |
| for i in range(1,len(blue_orange_colors)): |
| if blue_orange_colors[i][0] > w: |
| a, (ra, ga, ba) = blue_orange_colors[i-1] |
| b, (rb, gb, bb) = blue_orange_colors[i] |
| s = (w-a) / (b-a) |
| return ra+(rb-ra)*s, ga+(gb-ga)*s, ba+(bb-ba)*s |
|
|
| def cplot(ctx, f, re=[-5,5], im=[-5,5], points=2000, color=None, |
| verbose=False, file=None, dpi=None, axes=None): |
| """ |
| Plots the given complex-valued function *f* over a rectangular part |
| of the complex plane specified by the pairs of intervals *re* and *im*. |
| For example:: |
| |
| cplot(lambda z: z, [-2, 2], [-10, 10]) |
| cplot(exp) |
| cplot(zeta, [0, 1], [0, 50]) |
| |
| By default, the complex argument (phase) is shown as color (hue) and |
| the magnitude is show as brightness. You can also supply a |
| custom color function (*color*). This function should take a |
| complex number as input and return an RGB 3-tuple containing |
| floats in the range 0.0-1.0. |
| |
| Alternatively, you can select a builtin color function by passing |
| a string as *color*: |
| |
| * "default" - default color scheme |
| * "phase" - a color scheme that only renders the phase of the function, |
| with white for positive reals, black for negative reals, gold in the |
| upper half plane, and blue in the lower half plane. |
| |
| To obtain a sharp image, the number of points may need to be |
| increased to 100,000 or thereabout. Since evaluating the |
| function that many times is likely to be slow, the 'verbose' |
| option is useful to display progress. |
| |
| .. note :: This function requires matplotlib (pylab). |
| """ |
| if color is None or color == "default": |
| color = ctx.default_color_function |
| if color == "phase": |
| color = ctx.phase_color_function |
| import pylab |
| if file: |
| axes = None |
| fig = None |
| if not axes: |
| fig = pylab.figure() |
| axes = fig.add_subplot(111) |
| rea, reb = re |
| ima, imb = im |
| dre = reb - rea |
| dim = imb - ima |
| M = int(ctx.sqrt(points*dre/dim)+1) |
| N = int(ctx.sqrt(points*dim/dre)+1) |
| x = pylab.linspace(rea, reb, M) |
| y = pylab.linspace(ima, imb, N) |
| |
| |
| |
| |
| w = pylab.zeros((N, M, 3)) |
| for n in xrange(N): |
| for m in xrange(M): |
| z = ctx.mpc(x[m], y[n]) |
| try: |
| v = color(f(z)) |
| except ctx.plot_ignore: |
| v = (0.5, 0.5, 0.5) |
| w[n,m] = v |
| if verbose: |
| print(str(n) + ' of ' + str(N)) |
| rea, reb, ima, imb = [float(_) for _ in [rea, reb, ima, imb]] |
| axes.imshow(w, extent=(rea, reb, ima, imb), origin='lower') |
| axes.set_xlabel('Re(z)') |
| axes.set_ylabel('Im(z)') |
| if fig: |
| if file: |
| pylab.savefig(file, dpi=dpi) |
| else: |
| pylab.show() |
|
|
| def splot(ctx, f, u=[-5,5], v=[-5,5], points=100, keep_aspect=True, \ |
| wireframe=False, file=None, dpi=None, axes=None): |
| """ |
| Plots the surface defined by `f`. |
| |
| If `f` returns a single component, then this plots the surface |
| defined by `z = f(x,y)` over the rectangular domain with |
| `x = u` and `y = v`. |
| |
| If `f` returns three components, then this plots the parametric |
| surface `x, y, z = f(u,v)` over the pairs of intervals `u` and `v`. |
| |
| For example, to plot a simple function:: |
| |
| >>> from mpmath import * |
| >>> f = lambda x, y: sin(x+y)*cos(y) |
| >>> splot(f, [-pi,pi], [-pi,pi]) # doctest: +SKIP |
| |
| Plotting a donut:: |
| |
| >>> r, R = 1, 2.5 |
| >>> f = lambda u, v: [r*cos(u), (R+r*sin(u))*cos(v), (R+r*sin(u))*sin(v)] |
| >>> splot(f, [0, 2*pi], [0, 2*pi]) # doctest: +SKIP |
| |
| .. note :: This function requires matplotlib (pylab) 0.98.5.3 or higher. |
| """ |
| import pylab |
| import mpl_toolkits.mplot3d as mplot3d |
| if file: |
| axes = None |
| fig = None |
| if not axes: |
| fig = pylab.figure() |
| axes = mplot3d.axes3d.Axes3D(fig) |
| ua, ub = u |
| va, vb = v |
| du = ub - ua |
| dv = vb - va |
| if not isinstance(points, (list, tuple)): |
| points = [points, points] |
| M, N = points |
| u = pylab.linspace(ua, ub, M) |
| v = pylab.linspace(va, vb, N) |
| x, y, z = [pylab.zeros((M, N)) for i in xrange(3)] |
| xab, yab, zab = [[0, 0] for i in xrange(3)] |
| for n in xrange(N): |
| for m in xrange(M): |
| fdata = f(ctx.convert(u[m]), ctx.convert(v[n])) |
| try: |
| x[m,n], y[m,n], z[m,n] = fdata |
| except TypeError: |
| x[m,n], y[m,n], z[m,n] = u[m], v[n], fdata |
| for c, cab in [(x[m,n], xab), (y[m,n], yab), (z[m,n], zab)]: |
| if c < cab[0]: |
| cab[0] = c |
| if c > cab[1]: |
| cab[1] = c |
| if wireframe: |
| axes.plot_wireframe(x, y, z, rstride=4, cstride=4) |
| else: |
| axes.plot_surface(x, y, z, rstride=4, cstride=4) |
| axes.set_xlabel('x') |
| axes.set_ylabel('y') |
| axes.set_zlabel('z') |
| if keep_aspect: |
| dx, dy, dz = [cab[1] - cab[0] for cab in [xab, yab, zab]] |
| maxd = max(dx, dy, dz) |
| if dx < maxd: |
| delta = maxd - dx |
| axes.set_xlim3d(xab[0] - delta / 2.0, xab[1] + delta / 2.0) |
| if dy < maxd: |
| delta = maxd - dy |
| axes.set_ylim3d(yab[0] - delta / 2.0, yab[1] + delta / 2.0) |
| if dz < maxd: |
| delta = maxd - dz |
| axes.set_zlim3d(zab[0] - delta / 2.0, zab[1] + delta / 2.0) |
| if fig: |
| if file: |
| pylab.savefig(file, dpi=dpi) |
| else: |
| pylab.show() |
|
|
|
|
| VisualizationMethods.plot = plot |
| VisualizationMethods.default_color_function = default_color_function |
| VisualizationMethods.phase_color_function = phase_color_function |
| VisualizationMethods.cplot = cplot |
| VisualizationMethods.splot = splot |
|
|
