using sympy lambdify for compilation

mmatera · mmatera · commit 26bcaf15e01c · 2025-12-01T00:39:45.000-03:00
diff --git a/mathics/builtin/compilation.py b/mathics/builtin/compilation.py
@@ -11,7 +11,7 @@
 from types import FunctionType
 
 from mathics.builtin.box.compilation import CompiledCodeBox
-from mathics.core.atoms import Integer, String
+from mathics.core.atoms import Complex, Integer, Real, String
 from mathics.core.attributes import A_HOLD_ALL, A_PROTECTED
 from mathics.core.builtin import Builtin
 from mathics.core.convert.expression import to_mathics_list
@@ -83,7 +83,6 @@ class Compile(Builtin):
 
     def eval(self, vars, expr, evaluation: Evaluation):
         "Compile[vars_, expr_]"
-
         if not vars.has_form("List", None):
             evaluation.message("Compile", "invars")
             return
@@ -167,7 +166,11 @@ def to_sympy(self, *args, **kwargs):
         raise NotImplementedError
 
     def __hash__(self):
-        return hash(("CompiledCode", ctypes.addressof(self.cfunc)))  # XXX hack
+        try:
+            return hash(("CompiledCode", ctypes.addressof(self.cfunc)))  # XXX hack
+        except TypeError:
+            return hash(("CompiledCode", self.cfunc,))  # XXX hack
+            
 
     def atom_to_boxes(self, f, evaluation: Evaluation):
         return CompiledCodeBox(String(self.__str__()), evaluation=evaluation)
@@ -191,27 +194,38 @@ class CompiledFunction(Builtin):
 
     """
 
-    messages = {"argerr": "Invalid argument `1` should be Integer, Real or boolean."}
+    messages = {"argerr": "Invalid argument `1` should be Integer, Real, Complex or boolean."}
     summary_text = "A CompiledFunction object."
 
     def eval(self, argnames, expr, code, args, evaluation: Evaluation):
         "CompiledFunction[argnames_, expr_, code_CompiledCode][args__]"
-
         argseq = args.get_sequence()
 
         if len(argseq) != len(code.args):
             return
 
         py_args = []
-        for arg in argseq:
-            if isinstance(arg, Integer):
-                py_args.append(arg.get_int_value())
+        args_spec = code.args or []
+        if len(args_spec)!= len(argseq):
+            evaluation.mesage("CompiledFunction","cfct", Integer(len(argseq)), Integer(len(args_spec)))
+            return
+        for arg, spec in zip(argseq, args_spec):
+            # TODO: check if the types are consistent.
+            # If not, show a message.
+            if isinstance(arg, (Integer, Real, Complex)):
+                val = arg.value
             elif arg.sameQ(SymbolTrue):
-                py_args.append(True)
+                val = True
             elif arg.sameQ(SymbolFalse):
-                py_args.append(False)
+                val = False
             else:
-                py_args.append(arg.round_to_float(evaluation))
+                val = arg.to_python()
+            try:
+                val = spec.type(val)
+            except TypeError:
+                # Fallback by replace values in expr?
+                return
+            py_args.append(val)                
         try:
             result = code.cfunc(*py_args)
         except (TypeError, ctypes.ArgumentError):
diff --git a/mathics/compile/base.py b/mathics/compile/base.py
@@ -10,3 +10,6 @@ class CompileArg:
     def __init__(self, name, type):
         self.name = name
         self.type = type
+
+    def __repr__(self):
+        return f"{self.name}:{self.type}"
diff --git a/mathics/core/convert/function.py b/mathics/core/convert/function.py
@@ -1,37 +1,39 @@
 # -*- coding: utf-8 -*-
+import numpy
+from typing import Callable, List, Optional, Tuple
 
-from typing import Callable, Optional, Tuple
 
 from mathics.core.evaluation import Evaluation
 from mathics.core.expression import Expression, from_python
 from mathics.core.symbols import Symbol, SymbolFalse, SymbolTrue
-from mathics.core.systemsymbols import SymbolBlank, SymbolInteger, SymbolReal
+from mathics.core.systemsymbols import SymbolBlank, SymbolInteger, SymbolReal, SymbolComplex, SymbolOr
 from mathics.eval.nevaluator import eval_N
+from mathics.core.convert.lambdify import lambdify_compile, CompileError as LambdifyCompileError
+
+
+PERMITTED_TYPES = {
+        Expression(SymbolBlank, SymbolInteger): int,
+        Expression(SymbolBlank, SymbolReal): float,
+        Expression(SymbolBlank, SymbolComplex): complex,        
+        Expression(SymbolOr, SymbolTrue, SymbolFalse): bool,
+        Expression(SymbolOr, SymbolFalse, SymbolTrue): bool,
+}
+
 
 try:
     from mathics.compile import CompileArg, CompileError, _compile
     from mathics.compile.types import bool_type, int_type, real_type
 
-    use_llvm = True
+    USE_LLVM = True
     # _Complex not implemented
-    permitted_types = {
-        Expression(SymbolBlank, SymbolInteger): int_type,
-        Expression(SymbolBlank, SymbolReal): real_type,
-        SymbolTrue: bool_type,
-        SymbolFalse: bool_type,
+    LLVM_TYPE_TRANSLATION = {
+        int: int_type,
+        float: real_type,
+        bool: bool_type,
     }
 except ImportError:
-    use_llvm = False
-    bool_type = bool
-    int_type = int
-    real_type = float
+    USE_LLVM = False
 
-    permitted_types = {
-        Expression(SymbolBlank, SymbolInteger): int,
-        Expression(SymbolBlank, SymbolReal): float,
-        SymbolTrue: bool,
-        SymbolFalse: bool,
-    }
 
 
 class CompileDuplicateArgName(Exception):
@@ -44,69 +46,66 @@ def __init__(self, var):
         self.var = var
 
 
-def expression_to_callable(
-    expr: Expression,
-    args: Optional[list] = None,
-    evaluation: Optional[Evaluation] = None,
-) -> Optional[Callable]:
+
+def expression_to_llvm(expr: Expression, args:Optional[list]=None, evaluation: Optional[Evaluation]=None):
     """
-    Return a Python callable from an expression. If llvm is available,
-    tries to produce llvm code. Otherwise, returns a Python function.
+    Convert an expression to LLVM code. None if it fails.
     expr: Expression
     args: a list of CompileArg elements
     evaluation: an Evaluation object used if the llvm compilation fails
     """
     try:
-        cfunc = _compile(expr, args) if (use_llvm and args is not None) else None
+        return _compile(expr, args) if (USE_LLVM and args is not None) else None
     except CompileError:
-        cfunc = None
-
-    if cfunc is None:
-        if evaluation is None:
-            raise CompileError
-        try:
-
-            def _pythonized_mathics_expr(*x):
-                inner_evaluation = Evaluation(definitions=evaluation.definitions)
-                x_mathics = (from_python(u) for u in x[: len(args)])
-                vars = dict(list(zip([a.name for a in args], x_mathics)))
-                pyexpr = expr.replace_vars(vars)
-                pyexpr = eval_N(pyexpr, inner_evaluation)
-                res = pyexpr.to_python()
-                return res
+        return None
 
-            # TODO: check if we can use numba to compile this...
-            cfunc = _pythonized_mathics_expr
-        except Exception:
-            cfunc = None
-    return cfunc
 
-
-def expression_to_callable_and_args(
-    expr: Expression,
-    vars: Optional[list] = None,
-    evaluation: Optional[Evaluation] = None,
-) -> Tuple[Optional[Callable], Optional[list]]:
+def expression_to_python_function(
+        expr: Expression,
+        args: Optional[list] = None,
+        evaluation: Optional[Evaluation] = None,
+) -> Optional[Callable]:
     """
-    Return a tuple of Python callable and a list of CompileArgs.
-    expr: A Mathics Expression object
-    vars: a list of Symbols or Mathics Lists of the form {Symbol, Type}
+    Return a Python function from an expression. 
+    expr: Expression
+    args: a list of CompileArg elements
+    evaluation: an Evaluation object used if the llvm compilation fails
+    """
+    try:
+        def _pythonized_mathics_expr(*x):
+            inner_evaluation = Evaluation(definitions=evaluation.definitions)
+            x_mathics = (from_python(u) for u in x[: len(args)])
+            vars = dict(list(zip([a.name for a in args], x_mathics)))
+            pyexpr = expr.replace_vars(vars)
+            pyexpr = eval_N(pyexpr, inner_evaluation)
+            res = pyexpr.to_python()
+            return res
+
+        # TODO: check if we can use numba to compile this...
+        return _pythonized_mathics_expr
+    except Exception:
+        return None
+
+
+def collect_args(vars)->Optional[List[CompileArg]]:
+    """
+    Convert a List expression into a list of CompileArg objects.
     """
     if vars is None:
-        args = None
+        return None
     else:
         args = []
         names = []
         for var in vars:
             if isinstance(var, Symbol):
                 symb = var
                 name = symb.get_name()
-                typ = real_type
+                typ = float
             elif var.has_form("List", 2):
                 symb, typ = var.elements
-                if isinstance(symb, Symbol) and typ in permitted_types:
+                if isinstance(symb, Symbol) and typ in PERMITTED_TYPES:
                     name = symb.get_name()
-                    typ = permitted_types[typ]
+                    typ = PERMITTED_TYPES[typ]
                 else:
                     raise CompileWrongArgType(var)
             else:
@@ -116,5 +115,46 @@ def expression_to_callable_and_args(
                 raise CompileDuplicateArgName(symb)
             names.append(name)
             args.append(CompileArg(name, typ))
+    return args
+    
+
+
+def expression_to_callable_and_args(
+    expr: Expression,
+    vars: Optional[list] = None,
+    evaluation: Optional[Evaluation] = None,
+    debug:int = 0,
+    vectorize=False    
+) -> Tuple[Callable, Optional[list]]:
+    """
+    Return a tuple of Python callable and a list of CompileArgs.
+    expr: A Mathics Expression object
+    vars: a list of Symbols or Mathics Lists of the form {Symbol, Type}
+    """
+    args = collect_args(vars)
+
+    # First, try to lambdify the expression:
+    try:
+        cfunc = lambdify_compile(evaluation, expr, [arg.name for arg in args], debug)
+        # lambdify_compile returns an already vectorized expression.
+        return cfunc, args
+    except LambdifyCompileError:
+        pass
+
+    # Then, try with llvm if available
+    if USE_LLVM:
+        try:            
+            llvm_args = None if args is None else [CompileArg(compile_arg.name, LLVM_TYPE_TRANSLATION[compile_arg.typ]) for compile_arg in args]
+            cfunc = expression_to_llvm(expr, llvm_args, evaluation)
+            if vectorize:
+                cfunc = numpy.vectorize(cfunc)
+            return cfunc, llvm_args
+        except KeyError:
+            pass
+        
+    # Last resource
+    cfunc = expression_to_python_function(expr, args, evaluation)
+    if vectorize:
+        cfunc = numpy.vectorize(cfunc)
+    return  cfunc, args
 
-    return expression_to_callable(expr, args, evaluation), args
diff --git a/mathics/core/convert/lambdify.py b/mathics/core/convert/lambdify.py
@@ -81,7 +81,7 @@ def lambdify_compile(evaluation, expr, names, debug=0):
     # 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)
+        symbols = sympy.symbols(tuple(strip_context(name) for name in names))
         # compiled_function = sympy.lambdify(symbols, sympy_expr, mappings)
         compiled_function = sympy.lambdify(
             symbols, sympy_expr, modules=["numpy", "scipy"]
