Add support for Mesh

Author: bdlucas1

mathics/eval/drawing/plot3d_vectorized.py

@@ -10,7 +10,7 @@
 
 from mathics.core.evaluation import Evaluation
 from mathics.core.symbols import strip_context
-from mathics.core.systemsymbols import SymbolRGBColor, SymbolNone
+from mathics.core.systemsymbols import SymbolRGBColor, SymbolNone, SymbolEdgeForm
 from mathics.timing import Timer
 
 from .plot_compile import plot_compile
@@ -27,13 +27,14 @@ def eval_Plot3D(
     # 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]
 
-    # 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)
+    # Mesh option
+    nmesh = None
+    if isinstance(plot_options.mesh, int):
+        nmesh = plot_options.mesh
+    elif plot_options.mesh is not SymbolNone:
+        nmesh = 20
 
     # https://davidmathlogic.com/colorblind
     palette = [
@@ -48,53 +49,81 @@ def eval_Plot3D(
         #(0, 0, 0), # black
     ]
 
-    for i, function in enumerate(plot_options.functions):
-        with Timer("compile"):
-            function = plot_compile(evaluation, function, names)
+    # compile the functions
+    with Timer("compile"):
+        compiled_functions = [plot_compile(evaluation, function, names) for function in plot_options.functions]
 
-        # compute zs from xs and ys using compiled function
-        with Timer("compute zs"):
-            zs = function(**{str(names[0]): xs, str(names[1]): ys})
+    def compute_over_grid(nx, ny):
 
-        # 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)
+        # 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)
 
-        # if it's a constant, make it a full array
-        if isinstance(zs, (float, int, complex)):
-            zs = np.full(xs.shape, zs)
+        # (nx,ny) array of numbers from 0 to n-1 that are
+        # indexes into xyzs array for corresponding vertex
+        # +1 because these will be used as WL indexes, which are 1-based
+        inxs = np.arange(math.prod(xs.shape)).reshape(xs.shape) + 1
 
-        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)
+        for function in compiled_functions:
 
-            # 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]]
-            )
+            # compute zs from xs and ys using compiled function
+            with Timer("compute zs"):
+                zs = function(**{str(names[0]): xs, str(names[1]): ys})
 
-            # transpose and flatten to ((nx-1)*(ny-1), 4) array, suitable for use in GraphicsComplex
-            quads = quads.T.reshape(-1, 4)
+            # 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)
 
-            # ugh - indexes in Polygon are 1-based
-            quads += 1
+            # if it's a constant, make it a full array
+            if isinstance(zs, (float, int, complex)):
+                zs = np.full(xs.shape, zs)
 
-            # choose a color
-            rgb = palette[i%len(palette)]
-            rgb = [c/255.0 for c in rgb]
-            graphics.add_color(SymbolRGBColor, rgb)
-
-            # add a GraphicsComplex for this function
-            graphics.add_complex(xyzs, lines=None, polys=quads)
+            # (nx*ny, 3) array of points, to be indexed by quads
+            xyzs = np.stack([xs, ys, zs]).transpose(1, 2, 0).reshape(-1, 3)
 
+            yield xyzs, inxs
+
+    # generate the quads and emit a GraphicsComplex containing them
+    for i, (xyzs, inxs) in enumerate(compute_over_grid(*plot_options.plotpoints)):
+
+        # 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)
+
+        # choose a color
+        rgb = palette[i%len(palette)]
+        rgb = [c/255.0 for c in rgb]
+        #graphics.add_color(SymbolRGBColor, rgb)
+        graphics.add_directives([SymbolRGBColor, *rgb])
+
+        # add a GraphicsComplex for this function
+        graphics.add_complex(xyzs, lines=None, polys=quads)
+
+    # if requested by the Mesh attribute create a mesh of lines covering the surfaces
+    if nmesh:
+
+        # meshes are black for now
+        graphics.add_directives([SymbolRGBColor, 0,0,0])
+
+        with Timer("Mesh"):
+            nmesh = 20 # TODO: use supplied option
+            nx, ny = plot_options.plotpoints
+            # Do nmesh lines in each direction.
+            # Each mesh line has high res (nx or ny) so it follows
+            # the contours of the surface.
+            for xyzs, inxs in compute_over_grid(nx, nmesh):
+                graphics.add_complex(xyzs, lines=inxs, polys=None)
+            for xyzs, inxs in compute_over_grid(nmesh, ny):
+                graphics.add_complex(xyzs, lines=inxs.T, polys=None)
 
     return graphics
 

mathics/eval/drawing/util.py

@@ -4,9 +4,10 @@
 
 
 from mathics.core.atoms import NumericArray
-from mathics.core.convert.expression import to_mathics_list
+from mathics.core.convert.expression import to_mathics_list, to_expression
 from mathics.core.expression import Expression
 from mathics.core.list import ListExpression
+from mathics.core.symbols import Symbol
 from mathics.core.systemsymbols import (
     SymbolGraphics,
     SymbolGraphics3D,
@@ -59,17 +60,28 @@ def add_linexyzs(self, line_xyzs, colors=None):
         self.add_thing(SymbolLine, line_xyzs, colors)
 
     def add_color(self, symbol, components):
-        from mathics.core.convert.expression import to_expression
         expr = to_expression(symbol, *components)
         self.graphics.append(expr)
 
+    def add_directives(self, *ds):
+        def cvt(d):
+            if isinstance(d, list) and len(d) > 0 and isinstance(d[0], Symbol):
+                expr = to_expression(d[0], *(cvt(dd) for dd in d[1:]))
+                return expr
+            else:
+                return d
+        for d in ds:
+            expr = cvt(d)
+            self.graphics.append(expr)
+        
+
     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))
+            lines_expr = Expression(SymbolLine, NumericArray(lines))
             complex.append(lines_expr)
         gc_expr = Expression(SymbolGraphicsComplex, *complex)
         self.graphics.append(gc_expr)