Vectorized plot3d (#1522)

* Introduce a vectorized Plot3D in plot3d_vectorized.py. Since this
  generates Graphics3D containing GraphicsComplex with NumericArray,
  which no code will currently understand, for now it is enabled only
  if you set environment variable
  MATHICS_USE_VECTORIZED_PLOT=yes. (TBD how we will eventually
  deploy.)

* Introduce a "compiler" in plot_compile.py that uses expr.to_sympy()
  followed by sympy.lambdify to generate a Python function that calls
  numpy to evaluate. This is a huge win because most Python functions
  can do vectorized operations on entire arrays. For now it's specific
  to plotting; TBD whether it might find other uses. See
  test_plot_compiler to get a sense of what function is currently
  supported.

* Slightly refactors GraphicsGenerator for simplicity.

Here are some timings comparing classic and vectorized Plot3D. The
visual comparison was done using a prototype front-end that
understands GraphicsComplex and NumericArray.
            
Plot3D[Sin[x] Cos[y], {x,0,10}, {y,0,10}, PlotPoints->{p,p},
MaxRecursion->r]

                   p    r             eval_Plot3D
classic 50 2 (default) 10083 ms visible triangles on plot surface
classic 100 0 3232 ms less visible artifacts than above
    vectorized    200   n/a               5 ms      no visible artifacts

Plotting is now fast enough that useful interaction with Manipulate is
possible. I have a prototype of that and will submit a PR.

Other possible future work is to build out support for GraphicsComplex
and NumericArray in consumers of Graphics3D. It might be helpful to
have a "GraphicsParser" class that takes a Graphics3D expression and
turns it into Python data structures that are easy for consumers to
traverse.

Author: bdlucas1

mathics/builtin/drawing/plot.py

@@ -7,6 +7,7 @@
 """
 
 import numbers
+import os
 from abc import ABC
 from functools import lru_cache
 from math import cos, pi, sin
@@ -51,7 +52,18 @@
     get_plot_range,
     get_plot_range_option,
 )
-from mathics.eval.drawing.plot3d import eval_DensityPlot, eval_Plot3D
+
+# The vectorized plot function generates GraphicsComplex using NumericArray,
+# which no consumer will currently understand. So lets make it opt-in for now.
+# If it remains opt-in we'll probably want some combination of env variables,
+# Set option such as $UseVectorizedPlot, and maybe a non-standard Plot3D option.
+# For now an env variable is simplest.
+# TODO: work out exactly how to deploy.
+if os.getenv("MATHICS3_USE_VECTORIZED_PLOT", False):
+    from mathics.eval.drawing.plot3d_vectorized import eval_DensityPlot, eval_Plot3D
+else:
+    from mathics.eval.drawing.plot3d import eval_DensityPlot, eval_Plot3D
+
 from mathics.eval.nevaluator import eval_N
 
 # This tells documentation how to sort this module

mathics/core/systemsymbols.py

@@ -112,6 +112,7 @@
 SymbolGoldenRatio = Symbol("System`GoldenRatio")
 SymbolGraphics = Symbol("System`Graphics")
 SymbolGraphics3D = Symbol("System`Graphics3D")
+SymbolGraphicsComplex = Symbol("System`GraphicsComplex")
 SymbolGreater = Symbol("System`Greater")
 SymbolGreaterEqual = Symbol("System`GreaterEqual")
 SymbolGrid = Symbol("System`Grid")

mathics/eval/drawing/plot3d.py

@@ -22,6 +22,7 @@
     SymbolSlot,
 )
 from mathics.eval.drawing.plot import compile_quiet_function
+from mathics.timing import Timer
 
 from .util import GraphicsGenerator
 
@@ -388,6 +389,7 @@ def triangle(x1, y1, x2, y2, x3, y3, depth=0):
     return triangles, mesh_points, v_min, v_max
 
 
+@Timer("eval_Plot3D")
 def eval_Plot3D(
     plot_options,
     evaluation: Evaluation,

mathics/eval/drawing/plot3d_vectorized.py

@@ -0,0 +1,89 @@
+"""
+Vectorized evaluation routines for Plot3D and DensityPlot, which share a good bit of code.
+
+TODO: fill out eval_DensityPlot
+"""
+
+import math
+
+import numpy as np
+
+from mathics.core.evaluation import Evaluation
+from mathics.core.symbols import strip_context
+from mathics.timing import Timer
+
+from .plot_compile import plot_compile
+from .util import GraphicsGenerator
+
+
+@Timer("eval_Plot3D")
+def eval_Plot3D(
+    plot_options,
+    evaluation: Evaluation,
+):
+    graphics = GraphicsGenerator(dim=3)
+
+    for function in plot_options.functions:
+        # pull out plot options
+        _, xmin, xmax = plot_options.ranges[0]
+        _, ymin, ymax = plot_options.ranges[1]
+        nx, ny = plot_options.plotpoints
+        names = [strip_context(str(range[0])) for range in plot_options.ranges]
+
+        with Timer("compile"):
+            function = plot_compile(evaluation, function, names)
+
+        # compute (nx, ny) grids of xs and ys for corresponding vertexes
+        xs = np.linspace(xmin, xmax, nx)
+        ys = np.linspace(ymin, ymax, ny)
+        xs, ys = np.meshgrid(xs, ys)
+
+        # compute zs from xs and ys using compiled function
+        with Timer("compute zs"):
+            zs = function(**{str(names[0]): xs, str(names[1]): ys})
+
+        # sometimes expr gets compiled into something that returns a complex
+        # even though the imaginary part is 0
+        # TODO: check that imag is all 0?
+        # TODO: needed this for Hypergeometric - look into that
+        # assert np.all(np.isreal(zs)), "array contains complex values"
+        zs = np.real(zs)
+
+        with Timer("stack"):
+            # (nx*ny, 3) array of points, to be indexed by quads
+            xyzs = np.stack([xs, ys, zs]).transpose(1, 2, 0).reshape(-1, 3)
+
+            # (nx,ny) array of numbers from 0 to n-1 that are
+            # indexes into xyzs array for corresponding vertex
+            inxs = np.arange(math.prod(xs.shape)).reshape(xs.shape)
+
+            # shift inxs array four different ways and stack to form
+            # (4, nx-1, ny-1) array of quads represented as indexes into xyzs array
+            quads = np.stack(
+                [inxs[:-1, :-1], inxs[:-1, 1:], inxs[1:, 1:], inxs[1:, :-1]]
+            )
+
+            # transpose and flatten to ((nx-1)*(ny-1), 4) array, suitable for use in GraphicsComplex
+            quads = quads.T.reshape(-1, 4)
+
+            # ugh - indexes in Polygon are 1-based
+            quads += 1
+
+            # add a GraphicsComplex for this function
+            graphics.add_complex(xyzs, lines=None, polys=quads)
+
+    return graphics
+
+
+#
+#
+#
+
+
+def eval_DensityPlot(
+    plot_options,
+    evaluation: Evaluation,
+):
+    # TODO
+    # see plot3d.eval_DensityPlot for possible info on handling colors
+    pass

mathics/eval/drawing/plot_compile.py

@@ -0,0 +1,97 @@
+"""
+"Compile" an Expression by converting it to SymPy, then use sympy.lambdify
+to turn it into a Python function that calls NumPy functions to evaluate it.
+
+While possibly this provides an efficient function for point-wise evaluation,
+the main goal is to use NumPy to perform vectorized operations on arrays,
+which is a huge win for plotting.
+
+This will need testing and building out to make it robustly applicable
+to a wide array of expressions. So for now this is specific to plotting
+and is used only when enabled, so it lives with the plotting functions.
+Maybe eventually consider to move elsewhere if it seems to be more
+widely useful.
+"""
+
+import inspect
+
+import scipy
+import sympy
+
+from mathics.core.convert.sympy import SympyExpression, mathics_to_sympy
+from mathics.core.symbols import strip_context
+from mathics.core.util import print_expression_tree, print_sympy_tree
+
+
+# TODO: not in use yet
+# Add functions not found in scipy or numpy here.
+# Hopefully they are just thin adapater layers.
+# TODO: let's see how much is really needed here, and possibly consider moving
+# them to the Builtin somehow
+class AdditionalMappings:
+    def hyppfq(p, q, x):
+        if len(p) == 1 and len(q) == 1:
+            return scipy.special.hyp1f1(p[0], q[0], x)
+        else:
+            raise Exception(f"can't handle hyppfq({p}, {q}, x)")
+
+
+# mappings = [dict(AdditionalMappings.__dict__), "numpy"]
+
+
+class CompileError(Exception):
+    pass
+
+
+def plot_compile(evaluation, expr, names, debug=0):
+    """Compile the specified expression as a function of the given names"""
+
+    if debug >= 2:
+        print("=== compiling expr")
+        print_expression_tree(expr)
+
+    # Ask the expr Expression to generate a sympy expression and handle errors
+    try:
+        sympy_expr = expr.to_sympy()
+    except Exception as oops:
+        raise CompileError(f"{expr}.to_sympy() failed: {oops}")
+    if isinstance(sympy_expr, SympyExpression):
+        if debug:
+            # duplicates lookup logic in mathics.core.convert.sympy
+            lookup_name = expr.get_lookup_name()
+            builtin = mathics_to_sympy.get(lookup_name)
+            if builtin:
+                sympy_name = getattr(builtin, "sympy_name", None) if builtin else None
+                print(f"compile: Invalid sympy_expr {sympy_expr}")
+                print(f"compile: {builtin}.sympy_name is {repr(sympy_name)}")
+            else:
+                print(f"compile: {lookup_name} not registered with mathics_to_sympy")
+        raise CompileError(f"{expr.head}.to_sympy returns invalid sympy expr.")
+
+    # Strip symbols in sympy expression of context.
+    subs = {
+        sym: sympy.Symbol(strip_context(str(sym))) for sym in sympy_expr.free_symbols
+    }
+    sympy_expr = sympy_expr.subs(subs)
+
+    if debug >= 2:
+        print("=== equivalent sympy", type(sympy_expr))
+        print_sympy_tree(sympy_expr)
+
+    # Ask sympy to generate a function that will evaluate the expr.
+    # Use numpy and scipy to do the evaluation so that operations are vectorized.
+    # Augment the default numpy mappings with some additional ones not handled by default.
+    try:
+        symbols = sympy.symbols(names)
+        # compiled_function = sympy.lambdify(symbols, sympy_expr, mappings)
+        compiled_function = sympy.lambdify(
+            symbols, sympy_expr, modules=["numpy", "scipy"]
+        )
+    except Exception as oops:
+        raise CompileError(f"error compiling sympy expr {sympy_expr}: {oops}")
+
+    if debug >= 2:
+        print("=== compiled python")
+        print(inspect.getsource(compiled_function))
+
+    return compiled_function

mathics/eval/drawing/util.py

@@ -2,14 +2,15 @@
 Common utilities for plotting
 """
 
-import itertools
 
+from mathics.core.atoms import NumericArray
 from mathics.core.convert.expression import to_mathics_list
 from mathics.core.expression import Expression
 from mathics.core.list import ListExpression
 from mathics.core.systemsymbols import (
     SymbolGraphics,
     SymbolGraphics3D,
+    SymbolGraphicsComplex,
     SymbolLine,
     SymbolPolygon,
     SymbolRule,
@@ -25,69 +26,57 @@ class GraphicsGenerator:
     Support for generating Graphics and Graphics3D expressions
     """
 
-    # TODO: more precise types
-    # TODO: consider pre-zipping so only one for polys and one for lines
-    # TODO: consider whether we need to store these or can we just generate as we go?
-    poly_xyzs: list
-    poly_xyzs_colors: list
-    line_xyzs: list
-    line_xyzs_colors: list
+    # TODO: more precise typing?
+    graphics: list
 
     # 2 or 3
     dim: int
 
     def __init__(self, dim: int):
         self.dim = dim
-        self.poly_xyzs = []
-        self.poly_xyzs_colors = []
-        self.line_xyzs = []
-        self.line_xyzs_colors = []
+        self.graphics = []
+
+    # add polygons and lines, optionally with vertex colors
+    def add_thing(self, thing_symbol, things, colors):
+        arg = tuple(to_mathics_list(*thing) for thing in things)
+        arg = ListExpression(*arg) if len(arg) > 1 else arg[0]
+        if colors:
+            color_arg = tuple(to_mathics_list(*color) for color in colors)
+            color_arg = (
+                ListExpression(*color_arg) if len(color_arg) > 1 else color_arg[0]
+            )
+            color_rule = Expression(SymbolRule, SymbolVertexColors, color_arg)
+            self.graphics.append(Expression(thing_symbol, arg, color_rule))
+        else:
+            self.graphics.append(Expression(thing_symbol, arg))
 
-    # TODO: is this correct if some polys have colors and some don't?
     def add_polyxyzs(self, poly_xyzs, colors=None):
         """Add polygons specified by explicit xy[z] coordinates"""
-        self.poly_xyzs.append(poly_xyzs)
-        if colors:
-            self.poly_xyzs_colors.append(colors)
+        self.add_thing(SymbolPolygon, poly_xyzs, colors)
 
     def add_linexyzs(self, line_xyzs, colors=None):
         """Add lines specified by explicit xy[z] coordinates"""
-        self.line_xyzs.append(line_xyzs)
-        if colors:
-            self.line_xyzs_colors.append(colors)
+        self.add_thing(SymbolLine, line_xyzs, colors)
+
+    # TODO: color
+    def add_complex(self, xyzs, lines=None, polys=None):
+        complex = [NumericArray(xyzs)]
+        if polys is not None:
+            polys_expr = Expression(SymbolPolygon, NumericArray(polys))
+            complex.append(polys_expr)
+        if lines is not None:
+            polys_expr = Expression(SymbolLines, NumericArray(lines))
+            complex.append(lines_expr)
+        gc_expr = Expression(SymbolGraphicsComplex, *complex)
+        self.graphics.append(gc_expr)
 
     def generate(self, options):
         """
         Generates Graphics[3D] expression from supplied lines, polygons (etc.)
         """
-
-        # holds the elements of the final Graphics[3D] expr
-        graphics = []
-
-        # add polygons and lines, optionally with vertex colors
-        def add_thing(thing_symbol, thingss, colorss):
-            for things, colors in itertools.zip_longest(thingss, colorss):
-                arg = tuple(to_mathics_list(*thing) for thing in things)
-                arg = ListExpression(*arg) if len(arg) > 1 else arg[0]
-                if colors:
-                    color_arg = tuple(to_mathics_list(*color) for color in colors)
-                    color_arg = (
-                        ListExpression(*color_arg)
-                        if len(color_arg) > 1
-                        else color_arg[0]
-                    )
-                    color_rule = Expression(SymbolRule, SymbolVertexColors, color_arg)
-                    graphics.append(Expression(thing_symbol, arg, color_rule))
-                else:
-                    graphics.append(Expression(thing_symbol, arg))
-
-        add_thing(SymbolPolygon, self.poly_xyzs, self.poly_xyzs_colors)
-        add_thing(SymbolLine, self.line_xyzs, self.line_xyzs_colors)
-
-        # generate Graphics[3D] expression
         graphics_expr = Expression(
             SymbolGraphics3D if self.dim == 3 else SymbolGraphics,
-            ListExpression(*graphics),
+            ListExpression(*self.graphics),
             *options,
         )