|
| 1 | +""" |
| 2 | +Vectorized evaluation routines for Plot and related subclasses of _Plot |
| 3 | +""" |
| 4 | + |
| 5 | +import numpy as np |
| 6 | +from mathics.builtin.graphics import Graphics |
| 7 | +from mathics.builtin.options import filter_from_iterable, options_to_rules |
| 8 | +from mathics.core.convert.lambdify import lambdify_compile |
| 9 | +from mathics.core.element import BaseElement |
| 10 | +from mathics.core.evaluation import Evaluation |
| 11 | +from mathics.core.expression import Expression |
| 12 | +from mathics.core.symbols import SymbolList, strip_context |
| 13 | +from mathics.timing import Timer |
| 14 | + |
| 15 | +from .colors import palette2, palette_color_directive |
| 16 | +from .util import GraphicsGenerator |
| 17 | + |
| 18 | + |
| 19 | +@Timer("eval_Plot_vectorized") |
| 20 | +def eval_Plot_vectorized(plot_options, options, evaluation: Evaluation): |
| 21 | + # Note on naming: we use t to refer to the independent variable initially. |
| 22 | + # For Plot etc. it will be x, but for ParametricPlot it is better called t, |
| 23 | + # and for PolarPlot theta. After we call the apply_function supplied by |
| 24 | + # the plotting class we will then have actual plot coordinate xs and ys. |
| 25 | + tname, tmin, tmax = plot_options.ranges[0] |
| 26 | + nt = plot_options.plot_points |
| 27 | + |
| 28 | + # ParametricPlot passes a List of two functions, but lambdify_compile doesn't handle that |
| 29 | + # TODO: we should be receiving this as a Python list not an expression List? |
| 30 | + # TODO: can lambidfy_compile handle list with an appropriate to_sympy? |
| 31 | + def compile_maybe_list(evaluation, function, names): |
| 32 | + if isinstance(function, Expression) and function.head == SymbolList: |
| 33 | + fs = [lambdify_compile(evaluation, f, names) for f in function.elements] |
| 34 | + |
| 35 | + def compiled(vs): |
| 36 | + return [f(vs) for f in fs] |
| 37 | + |
| 38 | + else: |
| 39 | + compiled = lambdify_compile(evaluation, function, names) |
| 40 | + return compiled |
| 41 | + |
| 42 | + # compile the functions |
| 43 | + with Timer("compile"): |
| 44 | + names = [strip_context(str(tname))] |
| 45 | + compiled_functions = [ |
| 46 | + compile_maybe_list(evaluation, function, names) |
| 47 | + for function in plot_options.functions |
| 48 | + ] |
| 49 | + |
| 50 | + # compute requested regularly spaced points over the requested range |
| 51 | + ts = np.linspace(tmin, tmax, nt) |
| 52 | + |
| 53 | + # 1-based indexes into point array to form a line |
| 54 | + line = np.arange(nt) + 1 |
| 55 | + |
| 56 | + # compute the curves and accumulate in a GraphicsGenerator |
| 57 | + graphics = GraphicsGenerator(dim=2) |
| 58 | + for i, function in enumerate(compiled_functions): |
| 59 | + # compute xs and ys from ts using the compiled function |
| 60 | + # and the apply_function supplied by the plot class |
| 61 | + with Timer("compute xs and ys"): |
| 62 | + xs, ys = plot_options.apply_function(function, ts) |
| 63 | + |
| 64 | + # If the result is not numerical, we assume that the plot have failed. |
| 65 | + if isinstance(ys, BaseElement): |
| 66 | + return None |
| 67 | + |
| 68 | + # sometimes expr gets compiled into something that returns a complex |
| 69 | + # even though the imaginary part is 0 |
| 70 | + # TODO: check that imag is all 0? |
| 71 | + # assert np.all(np.isreal(zs)), "array contains complex values" |
| 72 | + xs = np.real(xs) |
| 73 | + ys = np.real(ys) |
| 74 | + |
| 75 | + # take log if requested; downstream axes will adjust accordingly |
| 76 | + if plot_options.use_log_scale: |
| 77 | + ys = np.log10(ys) |
| 78 | + |
| 79 | + # if it's a constant, make it a full array |
| 80 | + if isinstance(xs, (float, int, complex)): |
| 81 | + xs = np.full(ts.shape, xs) |
| 82 | + if isinstance(ys, (float, int, complex)): |
| 83 | + ys = np.full(ts.shape, ys) |
| 84 | + |
| 85 | + # (nx, 2) array of points, to be indexed by lines |
| 86 | + xys = np.stack([xs, ys]).T |
| 87 | + |
| 88 | + # give it a color from the 2d graph default color palette |
| 89 | + color = palette_color_directive(palette2, i) |
| 90 | + graphics.add_directives(color) |
| 91 | + |
| 92 | + # emit this line |
| 93 | + graphics.add_complex(xys, lines=line, polys=None) |
| 94 | + |
| 95 | + # copy options to output and generate the Graphics expr |
| 96 | + options = options_to_rules(options, filter_from_iterable(Graphics.options)) |
| 97 | + graphics_expr = graphics.generate(options) |
| 98 | + return graphics_expr |
0 commit comments