Add existing vectorized Plot(2D) routines. Adjust tests and context

pymathics/vectorizedplot/__init__.py

@@ -27,16 +27,34 @@
 # The try block is needed because at installation time, dependencies are not
 # available. After the installation is successfull, we want to make this availabe.
 try:
+    from pymathics.vectorizedplot.plot_plot import (
+        LogPlot,
+        ParametricPlot,
+        Plot,
+        PolarPlot,
+    )
     from pymathics.vectorizedplot.plot_plot3d import (
+        ComplexPlot,
+        ComplexPlot3D,
         ContourPlot,
         ContourPlot3D,
+        DensityPlot,
         ParametricPlot3D,
+        Plot3D,
         SphericalPlot3D,
     )
 
     _BUILTINS_ = (
+        "ComplexPlot",
+        "ComplexPlot3D",
         "ContourPlot",
         "ContourPlot3D",
+        "DensityPlot",
+        "LogPlot",
+        "ParametricPlot",
+        "PolarPlot",
+        "Plot",
+        "Plot3D",
         "ParametricPlot3D",
         "SphericalPlot3D",
     )

pymathics/vectorizedplot/eval/plot_vectorized.py

@@ -0,0 +1,98 @@
+"""
+Vectorized evaluation routines for Plot and related subclasses of _Plot
+"""
+
+import numpy as np
+from mathics.builtin.graphics import Graphics
+from mathics.builtin.options import filter_from_iterable, options_to_rules
+from mathics.core.convert.lambdify import lambdify_compile
+from mathics.core.element import BaseElement
+from mathics.core.evaluation import Evaluation
+from mathics.core.expression import Expression
+from mathics.core.symbols import SymbolList, strip_context
+from mathics.timing import Timer
+
+from .colors import palette2, palette_color_directive
+from .util import GraphicsGenerator
+
+
+@Timer("eval_Plot_vectorized")
+def eval_Plot_vectorized(plot_options, options, evaluation: Evaluation):
+    # Note on naming: we use t to refer to the independent variable initially.
+    # For Plot etc. it will be x, but for ParametricPlot it is better called t,
+    # and for PolarPlot theta. After we call the apply_function supplied by
+    # the plotting class we will then have actual plot coordinate xs and ys.
+    tname, tmin, tmax = plot_options.ranges[0]
+    nt = plot_options.plot_points
+
+    # ParametricPlot passes a List of two functions, but lambdify_compile doesn't handle that
+    # TODO: we should be receiving this as a Python list not an expression List?
+    # TODO: can lambidfy_compile handle list with an appropriate to_sympy?
+    def compile_maybe_list(evaluation, function, names):
+        if isinstance(function, Expression) and function.head == SymbolList:
+            fs = [lambdify_compile(evaluation, f, names) for f in function.elements]
+
+            def compiled(vs):
+                return [f(vs) for f in fs]
+
+        else:
+            compiled = lambdify_compile(evaluation, function, names)
+        return compiled
+
+    # compile the functions
+    with Timer("compile"):
+        names = [strip_context(str(tname))]
+        compiled_functions = [
+            compile_maybe_list(evaluation, function, names)
+            for function in plot_options.functions
+        ]
+
+    # compute requested regularly spaced points over the requested range
+    ts = np.linspace(tmin, tmax, nt)
+
+    # 1-based indexes into point array to form a line
+    line = np.arange(nt) + 1
+
+    # compute the curves and accumulate in a GraphicsGenerator
+    graphics = GraphicsGenerator(dim=2)
+    for i, function in enumerate(compiled_functions):
+        # compute xs and ys from ts using the compiled function
+        # and the apply_function supplied by the plot class
+        with Timer("compute xs and ys"):
+            xs, ys = plot_options.apply_function(function, ts)
+
+        # If the result is not numerical, we assume that the plot have failed.
+        if isinstance(ys, BaseElement):
+            return None
+
+        # sometimes expr gets compiled into something that returns a complex
+        # even though the imaginary part is 0
+        # TODO: check that imag is all 0?
+        # assert np.all(np.isreal(zs)), "array contains complex values"
+        xs = np.real(xs)
+        ys = np.real(ys)
+
+        # take log if requested; downstream axes will adjust accordingly
+        if plot_options.use_log_scale:
+            ys = np.log10(ys)
+
+        # if it's a constant, make it a full array
+        if isinstance(xs, (float, int, complex)):
+            xs = np.full(ts.shape, xs)
+        if isinstance(ys, (float, int, complex)):
+            ys = np.full(ts.shape, ys)
+
+        # (nx, 2) array of points, to be indexed by lines
+        xys = np.stack([xs, ys]).T
+
+        # give it a color from the 2d graph default color palette
+        color = palette_color_directive(palette2, i)
+        graphics.add_directives(color)
+
+        # emit this line
+        graphics.add_complex(xys, lines=line, polys=None)
+
+    # copy options to output and generate the Graphics expr
+    options = options_to_rules(options, filter_from_iterable(Graphics.options))
+    graphics_expr = graphics.generate(options)
+    return graphics_expr

pymathics/vectorizedplot/plot.py

@@ -556,7 +556,6 @@ def to_list(expr):
             plot_range = [symbol_type("System`Automatic")] * dim
             plot_range[-1] = pr
         self.plot_range = plot_range
-
         # ColorFunction and ColorFunctionScaling options
         # This was pulled from construct_density_plot (now eval_DensityPlot).
         # TODO: What does pop=True do? is it right?

pymathics/vectorizedplot/plot_plot.py

@@ -5,7 +5,6 @@
 """
 
 from abc import ABC
-from functools import lru_cache
 from typing import Callable
 
 import numpy as np
@@ -17,8 +16,7 @@
 from mathics.core.symbols import SymbolTrue
 from mathics.core.systemsymbols import SymbolLogPlot, SymbolPlotRange, SymbolSequence
 
-from pymathics.vectorizedplot.eval.drawing.plot import eval_Plot
-from pymathics.vectorizedplot.eval.drawing.plot_vectorized import eval_Plot_vectorized
+from pymathics.vectorizedplot.eval.plot_vectorized import eval_Plot_vectorized
 
 from . import plot
 
@@ -48,7 +46,7 @@ class _Plot(Builtin, ABC):
             "appropriate list of constraints."
         ),
     }
-
+    context = "System`"
     options = Graphics.options.copy()
     options.update(
         {
@@ -84,8 +82,6 @@ def eval(self, functions, ranges, evaluation: Evaluation, options: dict):
         # because ndarray is unhashable, and in any case probably isn't useful
         # TODO: does caching results in the classic case have demonstrable performance benefit?
         apply_function = self.apply_function
-        if not plot.use_vectorized_plot:
-            apply_function = lru_cache(apply_function)
         plot_options.apply_function = apply_function
 
         # TODO: PlotOptions has already regularized .functions to be a list
@@ -98,7 +94,7 @@ def eval(self, functions, ranges, evaluation: Evaluation, options: dict):
         plot_options.use_log_scale = self.use_log_scale
         plot_options.expect_list = self.expect_list
         if plot_options.plot_points is None:
-            default_plot_points = 1000 if plot.use_vectorized_plot else 57
+            default_plot_points = 1000
             plot_options.plot_points = default_plot_points
 
         # pass through the expanded plot_range options
@@ -110,7 +106,7 @@ def eval(self, functions, ranges, evaluation: Evaluation, options: dict):
             options[str(SymbolLogPlot)] = SymbolTrue
 
         # this will be either the vectorized or the classic eval function
-        eval_function = eval_Plot_vectorized if plot.use_vectorized_plot else eval_Plot
+        eval_function = eval_Plot_vectorized
         with np.errstate(all="ignore"):  # suppress numpy warnings
             graphics = eval_function(plot_options, options, evaluation)
         return graphics
@@ -188,6 +184,7 @@ class Plot(_Plot):
      = -Graphics-
     """
 
+    context = "System`"
     summary_text = "plot curves of one or more functions"
 
     def apply_function(self, f: Callable, x_value):

pymathics/vectorizedplot/plot_plot3d.py

@@ -30,6 +30,7 @@ class _Plot3D(Builtin):
     """Common base class for Plot3D, DensityPlot, ComplexPlot, ComplexPlot3D"""
 
     attributes = A_HOLD_ALL | A_PROTECTED
+    context = "System`"
 
     # Check for correct number of args
     eval_error = Builtin.generic_argument_error
@@ -84,10 +85,6 @@ class _Plot3D(Builtin):
         # 'MaxRecursion': '2',  # FIXME causes bugs in svg output see #303
     }
 
-    def contribute(self, definitions, is_pymodule=False):
-        print("contribute with ", type(self))
-        super().contribute(definitions, is_pymodule)
-
     def eval(
         self,
         functions,

test/helper.py

@@ -2,34 +2,32 @@
 import time
 from typing import Optional
 
+from mathics.core.element import BaseElement
 from mathics.core.load_builtin import import_and_load_builtins
 from mathics.core.symbols import Symbol
 from mathics.session import MathicsSession
 
 import_and_load_builtins()
 
-# Set up a Mathics session with definitions.
+# Set up two Mathics session with definitions, one for the vectorized routines and
+# other for the standard.
 # For consistency set the character encoding ASCII which is
 # the lowest common denominator available on all systems.
-session = MathicsSession(character_encoding="ASCII")
 
+SESSIONS = {
+    # test.helper session is going to be set up with the library.
+    True: MathicsSession(character_encoding="ASCII"),
+    # Default non-vectorized
+    False: MathicsSession(character_encoding="ASCII"),
+}
 
-def reset_session(add_builtin=True, catch_interrupt=False):
-    global session
-    session.reset()
 
-
-def evaluate_value(str_expr: str):
-    expr = session.evaluate(str_expr)
+def expr_to_value(expr: BaseElement):
     if isinstance(expr, Symbol):
         return expr.name
     return expr.value
 
 
-def evaluate(str_expr: str):
-    return session.evaluate(str_expr)
-
-
 def check_evaluation(
     str_expr: str,
     str_expected: str,
@@ -39,6 +37,7 @@ def check_evaluation(
     to_string_expected: bool = True,
     to_python_expected: bool = False,
     expected_messages: Optional[tuple] = None,
+    use_vectorized: bool = True,
 ):
     """
     Helper function to test Mathics expression against
@@ -66,34 +65,41 @@ def check_evaluation(
     expected_messages ``Optional[tuple[str]]``: If a tuple of strings are passed into this parameter, messages and prints raised during
                     the evaluation of ``str_expr`` are compared with the elements of the list. If ``None``, this comparison
                     is ommited.
+
+    use_vectorized: bool
+          If True, use the session with `pymathics.vectorizedplot` loaded.
     """
+    current_session = SESSIONS[use_vectorized]
+
     if str_expr is None:
-        reset_session()
-        evaluate('LoadModule["pymathics.vectorizedplot"]')
+        current_session.reset()
+        current_session.evaluate('LoadModule["pymathics.vectorizedplot"]')
         return
 
     if to_string_expr:
         str_expr = f"ToString[{str_expr}]"
-        result = evaluate_value(str_expr)
+        result = expr_to_value(current_session.evaluate(str_expr))
     else:
-        result = evaluate(str_expr)
+        result = current_session.evaluate(str_expr)
 
-    outs = [out.text for out in session.evaluation.out]
+    outs = [out.text for out in current_session.evaluation.out]
 
     if to_string_expected:
         if hold_expected:
             expected = str_expected
         else:
             str_expected = f"ToString[{str_expected}]"
-            expected = evaluate_value(str_expected)
+            expected = expr_to_value(current_session.evaluate(str_expected))
     else:
         if hold_expected:
             if to_python_expected:
                 expected = str_expected
             else:
-                expected = evaluate(f"HoldForm[{str_expected}]").elements[0]
+                expected = current_session.evaluate(
+                    f"HoldForm[{str_expected}]"
+                ).elements[0]
         else:
-            expected = evaluate(str_expected)
+            expected = current_session.evaluate(str_expected)
             if to_python_expected:
                 expected = expected.to_python(string_quotes=False)
 

test/test_plot.py

@@ -3,7 +3,7 @@
 Unit tests from mathics.builtin.drawing.plot
 """
 
-from test.helper import check_evaluation, session
+from test.helper import SESSIONS, check_evaluation
 
 import pytest
 from mathics.core.expression import Expression
@@ -37,6 +37,7 @@ def test__listplot():
             hold_expected=True,
             failure_message=fail_msg,
             expected_messages=msgs,
+            use_vectorized=False,
         )
 
 
@@ -248,6 +249,7 @@ def test_plot(str_expr, msgs, str_expected, fail_msg):
         hold_expected=True,
         failure_message=fail_msg,
         expected_messages=msgs,
+        use_vectorized=False,
     )
 
 
@@ -302,6 +304,8 @@ def mark(parent_expr, marker):
 
 
 def eval_and_check_structure(str_expr, str_expected):
+    session = SESSIONS[False]
+    session.reset()
     expr = session.parse(str_expr)
     result = expr.evaluate(session.evaluation)
     expected = session.parse(str_expected)