Add optimization param to compute functions

Author: FrancescAlted

src/blosc2/lazyexpr.py

@@ -897,6 +897,7 @@ def fast_eval(  # noqa: C901
         The output array.
     """
     out = kwargs.pop("_output", None)
+    optimization = kwargs.pop("_optimization", "aggressive")
     dtype = kwargs.pop("dtype", None)
     where: dict | None = kwargs.pop("_where_args", None)
     if isinstance(out, blosc2.NDArray):
@@ -962,12 +963,12 @@ def fast_eval(  # noqa: C901
             expression(tuple(chunk_operands.values()), result, offset=offset)
         else:
             if where is None:
-                result = ne.evaluate(expression, chunk_operands)
+                result = ne.evaluate(expression, chunk_operands, optimization=optimization)
             else:
                 # Apply the where condition (in result)
                 if len(where) == 2:
                     new_expr = f"where({expression}, _where_x, _where_y)"
-                    result = ne.evaluate(new_expr, chunk_operands)
+                    result = ne.evaluate(new_expr, chunk_operands, optimization=optimization)
                 else:
                     # We do not support one or zero operands in the fast path yet
                     raise ValueError("Fast path: the where condition must be a tuple with two elements")
@@ -1068,6 +1069,7 @@ def slices_eval(  # noqa: C901
         The output array.
     """
     out: blosc2.NDArray | None = kwargs.pop("_output", None)
+    optimization = kwargs.pop("_optimization", "aggressive")
     chunks = kwargs.get("chunks")
     where: dict | None = kwargs.pop("_where_args", None)
     _indices = kwargs.pop("_indices", False)
@@ -1186,7 +1188,7 @@ def slices_eval(  # noqa: C901
             continue
 
         if where is None:
-            result = ne.evaluate(expression, chunk_operands)
+            result = ne.evaluate(expression, chunk_operands, optimization=optimization)
         else:
             # Apply the where condition (in result)
             if len(where) == 2:
@@ -1195,9 +1197,9 @@ def slices_eval(  # noqa: C901
                 # result = np.where(result, x, y)
                 # numexpr is a bit faster than np.where, and we can fuse operations in this case
                 new_expr = f"where({expression}, _where_x, _where_y)"
-                result = ne.evaluate(new_expr, chunk_operands)
+                result = ne.evaluate(new_expr, chunk_operands, optimization=optimization)
             elif len(where) == 1:
-                result = ne.evaluate(expression, chunk_operands)
+                result = ne.evaluate(expression, chunk_operands, optimization=optimization)
                 if _indices or _order:
                     # Return indices only makes sense when the where condition is a tuple with one element
                     # and result is a boolean array
@@ -1332,6 +1334,7 @@ def reduce_slices(  # noqa: C901
         The resulting output array.
     """
     out = kwargs.pop("_output", None)
+    optimization = kwargs.pop("_optimization", "aggressive")
     where: dict | None = kwargs.pop("_where_args", None)
     reduce_op = reduce_args.pop("op")
     axis = reduce_args["axis"]
@@ -1468,14 +1471,14 @@ def reduce_slices(  # noqa: C901
                 # We don't have an actual expression, so avoid a copy
                 result = chunk_operands["o0"]
             else:
-                result = ne.evaluate(expression, chunk_operands)
+                result = ne.evaluate(expression, chunk_operands, optimization=optimization)
         else:
             # Apply the where condition (in result)
             if len(where) == 2:
                 new_expr = f"where({expression}, _where_x, _where_y)"
-                result = ne.evaluate(new_expr, chunk_operands)
+                result = ne.evaluate(new_expr, chunk_operands, optimization=optimization)
             elif len(where) == 1:
-                result = ne.evaluate(expression, chunk_operands)
+                result = ne.evaluate(expression, chunk_operands, optimization=optimization)
                 x = chunk_operands["_where_x"]
                 result = x[result]
             else:
@@ -1579,6 +1582,8 @@ def chunked_eval(  # noqa: C901
             Default is False.
         _output: NDArray or np.ndarray, optional
             The output array to store the result.
+        _optimization: str, optional
+            The optimization level to use.  Default is 'aggressive'.
         _where_args: dict, optional
             Additional arguments for conditional evaluation.
     """
@@ -2399,6 +2404,8 @@ def compute(self, item=None, **kwargs) -> blosc2.NDArray:
         # When NumPy ufuncs are called, the user may add an `out` parameter to kwargs
         if "out" in kwargs:
             kwargs["_output"] = kwargs.pop("out")
+        if "optimization" in kwargs:
+            kwargs["_optimization"] = kwargs.pop("optimization")
         if hasattr(self, "_output"):
             kwargs["_output"] = self._output
         if hasattr(self, "_where_args"):
@@ -2538,7 +2545,7 @@ def save(self, urlpath=None, **kwargs):
         }
 
     @classmethod
-    def _new_expr(cls, expression, operands, guess, out=None, where=None):
+    def _new_expr(cls, expression, operands, guess, out=None, where=None, optimization="aggressive"):
         # Validate the expression
         validate_expr(expression)
         if guess:
@@ -2578,6 +2585,7 @@ def _new_expr(cls, expression, operands, guess, out=None, where=None):
             new_expr._output = out
         if where is not None:
             new_expr._where_args = where
+        new_expr._optimization = optimization
         return new_expr
 
 
@@ -2864,6 +2872,7 @@ def lazyexpr(
     where: tuple | list | None = None,
     local_dict: dict | None = None,
     global_dict: dict | None = None,
+    optimization: str = "aggressive",
 ) -> LazyExpr:
     """
     Get a LazyExpr from an expression.
@@ -2889,6 +2898,10 @@ def lazyexpr(
     global_dict: dict, optional
         The global dictionary to use when looking for operands in the expression.
         If not provided, the global dictionary of the caller will be used.
+    optimization: str, optional
+        The optimization level to use when evaluating the expression. Possible
+        values are "aggressive" and "moderate".  The default value is "aggressive".
+        This parameter has the same meaning as in `numexpr.evaluate()`.
 
     Returns
     -------
@@ -2925,13 +2938,16 @@ def lazyexpr(
             expression.operands.update(operands)
         if out is not None:
             expression._output = out
+        expression._optimization = optimization
         if where is not None:
             where_args = {"_where_x": where[0], "_where_y": where[1]}
             expression._where_args = where_args
         return expression
     elif isinstance(expression, blosc2.NDArray):
         operands = {"o0": expression}
-        return LazyExpr._new_expr("o0", operands, guess=False, out=out, where=where)
+        return LazyExpr._new_expr(
+            "o0", operands, guess=False, out=out, where=where, optimization=optimization
+        )
 
     if operands is None:
         # Try to get operands from variables in the stack
@@ -2948,7 +2964,9 @@ def lazyexpr(
             # _new_expr will take care of the constructor, but needs an empty dict in operands
             operands = {}
 
-    return LazyExpr._new_expr(expression, operands, guess=True, out=out, where=where)
+    return LazyExpr._new_expr(
+        expression, operands, guess=True, out=out, where=where, optimization=optimization
+    )
 
 
 def _open_lazyarray(array):
@@ -2985,7 +3003,7 @@ def _open_lazyarray(array):
 
 
 # Mimim numexpr's evaluate function
-def evaluate(ex, local_dict=None, global_dict=None):
+def evaluate(ex, local_dict=None, global_dict=None, optimization="aggressive"):
     """
     Evaluate a string expression using the Blosc2 compute engine.
 
@@ -3008,6 +3026,10 @@ def evaluate(ex, local_dict=None, global_dict=None):
     global_dict: dict, optional
         The global dictionary to use when looking for operands in the expression.
         If not provided, the global dictionary of the caller will be used.
+    optimization: str, optional
+        The optimization level to use when evaluating the expression. Possible
+        values are "aggressive" and "moderate".  The default value is "aggressive".
+        This parameter has the same meaning as in `numexpr.evaluate()`.
 
     Returns
     -------
@@ -3030,7 +3052,7 @@ def evaluate(ex, local_dict=None, global_dict=None):
     [ 5.515625  8.25     11.765625]
     [16.0625   21.140625 27.      ]]
     """
-    lexpr = lazyexpr(ex, local_dict=local_dict, global_dict=global_dict)
+    lexpr = lazyexpr(ex, local_dict=local_dict, global_dict=global_dict, optimization=optimization)
     return lexpr[:]