Initial WIP

Author: bdlucas1

SYMBOLS_MANIFEST.txt

@@ -244,6 +244,8 @@ System`Complement
 System`Complex
 System`ComplexExpand
 System`ComplexInfinity
+System`ComplexPlot
+System`ComplexPlot3D
 System`Complexes
 System`CompositeQ
 System`Composition

mathics/builtin/drawing/plot.py

@@ -26,6 +26,7 @@
 from mathics.core.evaluation import Evaluation
 from mathics.core.expression import Expression
 from mathics.core.list import ListExpression
+from mathics.eval.nevaluator import eval_N
 from mathics.core.symbols import Symbol, SymbolList
 from mathics.core.systemsymbols import (
     SymbolAll,
@@ -38,6 +39,7 @@
     SymbolLog10,
     SymbolNone,
     SymbolRGBColor,
+    SymbolSequence,
     SymbolStyle,
 )
 from mathics.eval.drawing.charts import draw_bar_chart, eval_chart
@@ -62,11 +64,9 @@
 # TODO: work out exactly how to deploy.
 use_vectorized_plot = os.getenv("MATHICS3_USE_VECTORIZED_PLOT", False)
 if use_vectorized_plot:
-    from mathics.eval.drawing.plot3d_vectorized import eval_DensityPlot, eval_Plot3D
+    from mathics.eval.drawing.plot3d_vectorized import eval_DensityPlot, eval_Plot3D, eval_ComplexPlot, eval_ComplexPlot3D
 else:
-    from mathics.eval.drawing.plot3d import eval_DensityPlot, eval_Plot3D
-
-from mathics.eval.nevaluator import eval_N
+    from mathics.eval.drawing.plot3d import eval_DensityPlot, eval_Plot3D, eval_ComplexPlot, eval_ComplexPlot3D
 
 # This tells documentation how to sort this module
 # Here we are also hiding "drawing" since this erroneously appears at the top level.
@@ -137,6 +137,7 @@ def eval(self, points, evaluation: Evaluation, options: dict):
             )
 
         points = points.evaluate(evaluation)
+        print("xxx points", points)
         if not isinstance(points, ListExpression):
             evaluation.message(class_name, "lpn", points)
             return
@@ -460,17 +461,20 @@ def __init__(self, expr, range_exprs, options, evaluation):
         # plot ranges
         self.ranges = []
         for range_expr in range_exprs:
+            print("xxx range_expr", range_expr)
             if not range_expr.has_form("List", 3):
                 self.error(expr, "invrange", range_expr)
             if not isinstance(range_expr.elements[0], Symbol):
                 self.error(expr, "invrange", range_expr)
             range = [range_expr.elements[0]]
             for limit_expr in range_expr.elements[1:3]:
-                limit = limit_expr.round_to_float(evaluation)
-                if limit is None:
+                limit = eval_N(limit_expr, evaluation).to_python()
+                if not isinstance(limit, (int,float,complex)):
                     self.error(expr, "plln", limit_expr, range_expr)
                 range.append(limit)
-            if range[2] <= range[1]:
+            if isinstance(limit, (int,float)) and range[2] <= range[1]:
+                self.error(expr, "invrange", range_expr)
+            if isinstance(limit, complex) and (range[2].real <= range[1].real or range[2].imag <= range[1].imag):
                 self.error(expr, "invrange", range_expr)
             self.ranges.append(range)
 
@@ -539,7 +543,9 @@ def check_plotpoints(steps):
 
 
 class _Plot3D(Builtin):
-    """Common base class for Plot3D and DensityPlot"""
+    """Common base class for Plot3D, DensityPlot, ComplexPlot, ComplexPlot3D"""
+
+    attributes = A_HOLD_ALL | A_PROTECTED
 
     # Check for correct number of args
     eval_error = Builtin.generic_argument_error
@@ -565,29 +571,61 @@ class _Plot3D(Builtin):
         ),
     }
 
+    # Plot3D, ComplexPlot3D
+    options3d = Graphics3D.options | {
+        "Axes": "True",
+        "AspectRatio": "1",
+        "Mesh": "Full",
+        "PlotPoints": "None",
+        "BoxRatios": "{1, 1, 0.4}",
+        "MaxRecursion": "2",
+    }
+
+    # DensityPlot, ComplexPlot
+    options2d = Graphics.options | {
+        "Axes": "False",
+        "AspectRatio": "1",
+        "Mesh": "None",
+        "Frame": "True",
+        "ColorFunction": "Automatic",
+        "ColorFunctionScaling": "True",
+        "PlotPoints": "None",
+        "MaxRecursion": "0",
+        # 'MaxRecursion': '2',  # FIXME causes bugs in svg output see #303
+    }
+
     def eval(
         self,
         functions,
-        xrange,
-        yrange,
+        ranges,
         evaluation: Evaluation,
         options: dict,
     ):
-        """%(name)s[functions_, xrange_, yrange_, OptionsPattern[%(name)s]]"""
+        """%(name)s[functions_, ranges__, OptionsPattern[%(name)s]]"""
 
         # TODO: test error for too many, too few, no args
 
         # parse options, bailing out if anything is wrong
         try:
-            plot_options = PlotOptions(self, [xrange, yrange], options, evaluation)
+            ranges = ranges.elements if ranges.head is SymbolSequence else [ranges]
+            plot_options = PlotOptions(self, ranges, options, evaluation)
         except ValueError:
             return None
 
-        # ask the subclass to get one or more functions as appropriate
-        plot_options.functions = self.get_functions_param(functions)
+        # TODO: consult many_functions variable set by subclass and error
+        # if many_functions is False but multiple are supplied
+        if functions.has_form("List", None):
+            plot_options.functions = functions.elements
+        else:
+            plot_options.functions = [functions]
+
+        # subclass must set eval_function and graphics_class
+        graphics = self.eval_function(plot_options, evaluation)
+        if not graphics:
+            return
+        graphics_expr = graphics.generate(options_to_rules(options, self.graphics_class.options))
+        return graphics_expr
 
-        # delegate to subclass, which will call the appropriate eval_* function
-        return self.do_eval(plot_options, evaluation, options)
 
 
 class BarChart(_Chart):
@@ -712,6 +750,27 @@ class ColorDataFunction(Builtin):
     summary_text = "color scheme object"
 
 
+class ComplexPlot3D(_Plot3D):
+
+    summary_text = "plots one or more complex functions as a surface"
+    expected_args = 2
+    options = _Plot3D.options3d
+
+    many_functions = True
+    eval_function = staticmethod(eval_ComplexPlot3D)
+    graphics_class = Graphics3D
+
+class ComplexPlot(_Plot3D):
+    
+    summary_text = "plots a complex function showing amplitude and phase using colors"
+    expected_args = 2
+    options = _Plot3D.options2d
+
+    many_functions = False
+    eval_function = staticmethod(eval_ComplexPlot)
+    graphics_class = Graphics
+
+
 class DensityPlot(_Plot3D):
     """
     <url>:WMA link: https://reference.wolfram.com/language/ref/DensityPlot.html</url>
@@ -736,35 +795,13 @@ class DensityPlot(_Plot3D):
      = -Graphics-
     """
 
-    attributes = A_HOLD_ALL | A_PROTECTED
-
-    options = Graphics.options.copy()
-    options.update(
-        {
-            "Axes": "False",
-            "AspectRatio": "1",
-            "Mesh": "None",
-            "Frame": "True",
-            "ColorFunction": "Automatic",
-            "ColorFunctionScaling": "True",
-            "PlotPoints": "None",
-            "MaxRecursion": "0",
-            # 'MaxRecursion': '2',  # FIXME causes bugs in svg output see #303
-        }
-    )
     summary_text = "density plot for a function"
+    expected_args = 3
+    options = _Plot3D.options2d
 
-    # TODO: error if more than one function here
-    def get_functions_param(self, functions):
-        """can only have one function"""
-        return [functions]
-
-    # called by superclass
-    def do_eval(self, plot_options, evaluation, options):
-        """called by superclass to call appropriate eval_* function"""
-        graphics = eval_DensityPlot(plot_options, evaluation)
-        graphics_expr = graphics.generate(options_to_rules(options, Graphics.options))
-        return graphics_expr
+    many_functions = False
+    eval_function = staticmethod(eval_DensityPlot)
+    graphics_class = Graphics
 
 
 class DiscretePlot(_Plot):
@@ -1878,30 +1915,10 @@ class Plot3D(_Plot3D):
      = -Graphics3D-
     """
 
-    attributes = A_HOLD_ALL | A_PROTECTED
-
-    options = Graphics.options.copy()
-    options.update(
-        {
-            "Axes": "True",
-            "AspectRatio": "1",
-            "Mesh": "Full",
-            "PlotPoints": "None",
-            "BoxRatios": "{1, 1, 0.4}",
-            "MaxRecursion": "2",
-        }
-    )
     summary_text = "plots 3D surfaces of one or more functions"
+    expected_args = 3
+    options = _Plot3D.options3d
 
-    def get_functions_param(self, functions):
-        """May have a function or a list of functions"""
-        if functions.has_form("List", None):
-            return functions.elements
-        else:
-            return [functions]
-
-    def do_eval(self, plot_options, evaluation, options):
-        """called by superclass to call appropriate eval_* function"""
-        graphics = eval_Plot3D(plot_options, evaluation)
-        graphics_expr = graphics.generate(options_to_rules(options, Graphics3D.options))
-        return graphics_expr
+    many_functions = True
+    eval_function = staticmethod(eval_Plot3D)
+    graphics_class = Graphics3D

mathics/eval/drawing/plot3d.py

@@ -492,3 +492,17 @@ def eval_color(x, y, v):
         graphics.add_linexyzs([mesh_points[xi]])
 
     return graphics
+
+
+def eval_ComplexPlot3D(
+    plot_options,
+    evaluation: Evaluation,
+):
+    return None
+
+
+def eval_ComplexPlot(
+    plot_options,
+    evaluation: Evaluation,
+):
+    return None

mathics/eval/drawing/plot3d_vectorized.py

@@ -17,12 +17,18 @@
 from .util import GraphicsGenerator
 
 
-def make_plot(plot_options, evaluation: Evaluation, dim: int):
+def make_plot(plot_options, evaluation: Evaluation, dim: int, is_complex: bool):
     graphics = GraphicsGenerator(dim)
 
     # pull out plot options
-    _, xmin, xmax = plot_options.ranges[0]
-    _, ymin, ymax = plot_options.ranges[1]
+    if not is_complex:
+        _, xmin, xmax = plot_options.ranges[0]
+        _, ymin, ymax = plot_options.ranges[1]
+    else:
+        # will generate xs and ys as for real, then combine to form complex cs
+        _, cmin, cmax = plot_options.ranges[0]
+        xmin, xmax = cmin.real, cmax.real
+        ymin, ymax = cmin.imag, cmax.imag
     names = [strip_context(str(range[0])) for range in plot_options.ranges]
 
     # Mesh option
@@ -73,7 +79,11 @@ def compute_over_grid(nx, ny):
         for function in compiled_functions:
             # compute zs from xs and ys using compiled function
             with Timer("compute zs"):
-                zs = function(**{str(names[0]): xs, str(names[1]): ys})
+                if not is_complex:
+                    zs = function(**{str(names[0]): xs, str(names[1]): ys})
+                else:
+                    cs = xs + ys * 1j # TODO: fast enough?
+                    zs = function(**{str(names[0]): cs})
 
             # sometimes expr gets compiled into something that returns a complex
             # even though the imaginary part is 0
@@ -155,12 +165,28 @@ def eval_Plot3D(
     plot_options,
     evaluation: Evaluation,
 ):
-    return make_plot(plot_options, evaluation, dim=3)
+    return make_plot(plot_options, evaluation, dim=3, is_complex=False)
 
 
 @Timer("eval_DensityPlot")
 def eval_DensityPlot(
     plot_options,
     evaluation: Evaluation,
 ):
-    return make_plot(plot_options, evaluation, dim=2)
+    return make_plot(plot_options, evaluation, dim=2, is_complex=False)
+
+
+@Timer("eval_ComplexPlot3D")
+def eval_ComplexPlot3D(
+    plot_options,
+    evaluation: Evaluation,
+):
+    return None
+
+
+@Timer("eval_ComplexPlot")
+def eval_ComplexPlot(
+    plot_options,
+    evaluation: Evaluation,
+):
+    return make_plot(plot_options, evaluation, dim=2, is_complex=True)