Coverage for /usr/lib/python3/dist-packages/sympy/simplify/cse_main.py: 10%
372 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-14 15:55 +0200
1""" Tools for doing common subexpression elimination.
2"""
3from collections import defaultdict
5from sympy.core import Basic, Mul, Add, Pow, sympify
6from sympy.core.containers import Tuple, OrderedSet
7from sympy.core.exprtools import factor_terms
8from sympy.core.singleton import S
9from sympy.core.sorting import ordered
10from sympy.core.symbol import symbols, Symbol
11from sympy.matrices import (MatrixBase, Matrix, ImmutableMatrix,
12 SparseMatrix, ImmutableSparseMatrix)
13from sympy.matrices.expressions import (MatrixExpr, MatrixSymbol, MatMul,
14 MatAdd, MatPow, Inverse)
15from sympy.matrices.expressions.matexpr import MatrixElement
16from sympy.polys.rootoftools import RootOf
17from sympy.utilities.iterables import numbered_symbols, sift, \
18 topological_sort, iterable
20from . import cse_opts
22# (preprocessor, postprocessor) pairs which are commonly useful. They should
23# each take a SymPy expression and return a possibly transformed expression.
24# When used in the function ``cse()``, the target expressions will be transformed
25# by each of the preprocessor functions in order. After the common
26# subexpressions are eliminated, each resulting expression will have the
27# postprocessor functions transform them in *reverse* order in order to undo the
28# transformation if necessary. This allows the algorithm to operate on
29# a representation of the expressions that allows for more optimization
30# opportunities.
31# ``None`` can be used to specify no transformation for either the preprocessor or
32# postprocessor.
35basic_optimizations = [(cse_opts.sub_pre, cse_opts.sub_post),
36 (factor_terms, None)]
38# sometimes we want the output in a different format; non-trivial
39# transformations can be put here for users
40# ===============================================================
43def reps_toposort(r):
44 """Sort replacements ``r`` so (k1, v1) appears before (k2, v2)
45 if k2 is in v1's free symbols. This orders items in the
46 way that cse returns its results (hence, in order to use the
47 replacements in a substitution option it would make sense
48 to reverse the order).
50 Examples
51 ========
53 >>> from sympy.simplify.cse_main import reps_toposort
54 >>> from sympy.abc import x, y
55 >>> from sympy import Eq
56 >>> for l, r in reps_toposort([(x, y + 1), (y, 2)]):
57 ... print(Eq(l, r))
58 ...
59 Eq(y, 2)
60 Eq(x, y + 1)
62 """
63 r = sympify(r)
64 E = []
65 for c1, (k1, v1) in enumerate(r):
66 for c2, (k2, v2) in enumerate(r):
67 if k1 in v2.free_symbols:
68 E.append((c1, c2))
69 return [r[i] for i in topological_sort((range(len(r)), E))]
72def cse_separate(r, e):
73 """Move expressions that are in the form (symbol, expr) out of the
74 expressions and sort them into the replacements using the reps_toposort.
76 Examples
77 ========
79 >>> from sympy.simplify.cse_main import cse_separate
80 >>> from sympy.abc import x, y, z
81 >>> from sympy import cos, exp, cse, Eq, symbols
82 >>> x0, x1 = symbols('x:2')
83 >>> eq = (x + 1 + exp((x + 1)/(y + 1)) + cos(y + 1))
84 >>> cse([eq, Eq(x, z + 1), z - 2], postprocess=cse_separate) in [
85 ... [[(x0, y + 1), (x, z + 1), (x1, x + 1)],
86 ... [x1 + exp(x1/x0) + cos(x0), z - 2]],
87 ... [[(x1, y + 1), (x, z + 1), (x0, x + 1)],
88 ... [x0 + exp(x0/x1) + cos(x1), z - 2]]]
89 ...
90 True
91 """
92 d = sift(e, lambda w: w.is_Equality and w.lhs.is_Symbol)
93 r = r + [w.args for w in d[True]]
94 e = d[False]
95 return [reps_toposort(r), e]
98def cse_release_variables(r, e):
99 """
100 Return tuples giving ``(a, b)`` where ``a`` is a symbol and ``b`` is
101 either an expression or None. The value of None is used when a
102 symbol is no longer needed for subsequent expressions.
104 Use of such output can reduce the memory footprint of lambdified
105 expressions that contain large, repeated subexpressions.
107 Examples
108 ========
110 >>> from sympy import cse
111 >>> from sympy.simplify.cse_main import cse_release_variables
112 >>> from sympy.abc import x, y
113 >>> eqs = [(x + y - 1)**2, x, x + y, (x + y)/(2*x + 1) + (x + y - 1)**2, (2*x + 1)**(x + y)]
114 >>> defs, rvs = cse_release_variables(*cse(eqs))
115 >>> for i in defs:
116 ... print(i)
117 ...
118 (x0, x + y)
119 (x1, (x0 - 1)**2)
120 (x2, 2*x + 1)
121 (_3, x0/x2 + x1)
122 (_4, x2**x0)
123 (x2, None)
124 (_0, x1)
125 (x1, None)
126 (_2, x0)
127 (x0, None)
128 (_1, x)
129 >>> print(rvs)
130 (_0, _1, _2, _3, _4)
131 """
132 if not r:
133 return r, e
135 s, p = zip(*r)
136 esyms = symbols('_:%d' % len(e))
137 syms = list(esyms)
138 s = list(s)
139 in_use = set(s)
140 p = list(p)
141 # sort e so those with most sub-expressions appear first
142 e = [(e[i], syms[i]) for i in range(len(e))]
143 e, syms = zip(*sorted(e,
144 key=lambda x: -sum([p[s.index(i)].count_ops()
145 for i in x[0].free_symbols & in_use])))
146 syms = list(syms)
147 p += e
148 rv = []
149 i = len(p) - 1
150 while i >= 0:
151 _p = p.pop()
152 c = in_use & _p.free_symbols
153 if c: # sorting for canonical results
154 rv.extend([(s, None) for s in sorted(c, key=str)])
155 if i >= len(r):
156 rv.append((syms.pop(), _p))
157 else:
158 rv.append((s[i], _p))
159 in_use -= c
160 i -= 1
161 rv.reverse()
162 return rv, esyms
165# ====end of cse postprocess idioms===========================
168def preprocess_for_cse(expr, optimizations):
169 """ Preprocess an expression to optimize for common subexpression
170 elimination.
172 Parameters
173 ==========
175 expr : SymPy expression
176 The target expression to optimize.
177 optimizations : list of (callable, callable) pairs
178 The (preprocessor, postprocessor) pairs.
180 Returns
181 =======
183 expr : SymPy expression
184 The transformed expression.
185 """
186 for pre, post in optimizations:
187 if pre is not None:
188 expr = pre(expr)
189 return expr
192def postprocess_for_cse(expr, optimizations):
193 """Postprocess an expression after common subexpression elimination to
194 return the expression to canonical SymPy form.
196 Parameters
197 ==========
199 expr : SymPy expression
200 The target expression to transform.
201 optimizations : list of (callable, callable) pairs, optional
202 The (preprocessor, postprocessor) pairs. The postprocessors will be
203 applied in reversed order to undo the effects of the preprocessors
204 correctly.
206 Returns
207 =======
209 expr : SymPy expression
210 The transformed expression.
211 """
212 for pre, post in reversed(optimizations):
213 if post is not None:
214 expr = post(expr)
215 return expr
218class FuncArgTracker:
219 """
220 A class which manages a mapping from functions to arguments and an inverse
221 mapping from arguments to functions.
222 """
224 def __init__(self, funcs):
225 # To minimize the number of symbolic comparisons, all function arguments
226 # get assigned a value number.
227 self.value_numbers = {}
228 self.value_number_to_value = []
230 # Both of these maps use integer indices for arguments / functions.
231 self.arg_to_funcset = []
232 self.func_to_argset = []
234 for func_i, func in enumerate(funcs):
235 func_argset = OrderedSet()
237 for func_arg in func.args:
238 arg_number = self.get_or_add_value_number(func_arg)
239 func_argset.add(arg_number)
240 self.arg_to_funcset[arg_number].add(func_i)
242 self.func_to_argset.append(func_argset)
244 def get_args_in_value_order(self, argset):
245 """
246 Return the list of arguments in sorted order according to their value
247 numbers.
248 """
249 return [self.value_number_to_value[argn] for argn in sorted(argset)]
251 def get_or_add_value_number(self, value):
252 """
253 Return the value number for the given argument.
254 """
255 nvalues = len(self.value_numbers)
256 value_number = self.value_numbers.setdefault(value, nvalues)
257 if value_number == nvalues:
258 self.value_number_to_value.append(value)
259 self.arg_to_funcset.append(OrderedSet())
260 return value_number
262 def stop_arg_tracking(self, func_i):
263 """
264 Remove the function func_i from the argument to function mapping.
265 """
266 for arg in self.func_to_argset[func_i]:
267 self.arg_to_funcset[arg].remove(func_i)
270 def get_common_arg_candidates(self, argset, min_func_i=0):
271 """Return a dict whose keys are function numbers. The entries of the dict are
272 the number of arguments said function has in common with
273 ``argset``. Entries have at least 2 items in common. All keys have
274 value at least ``min_func_i``.
275 """
276 count_map = defaultdict(lambda: 0)
277 if not argset:
278 return count_map
280 funcsets = [self.arg_to_funcset[arg] for arg in argset]
281 # As an optimization below, we handle the largest funcset separately from
282 # the others.
283 largest_funcset = max(funcsets, key=len)
285 for funcset in funcsets:
286 if largest_funcset is funcset:
287 continue
288 for func_i in funcset:
289 if func_i >= min_func_i:
290 count_map[func_i] += 1
292 # We pick the smaller of the two containers (count_map, largest_funcset)
293 # to iterate over to reduce the number of iterations needed.
294 (smaller_funcs_container,
295 larger_funcs_container) = sorted(
296 [largest_funcset, count_map],
297 key=len)
299 for func_i in smaller_funcs_container:
300 # Not already in count_map? It can't possibly be in the output, so
301 # skip it.
302 if count_map[func_i] < 1:
303 continue
305 if func_i in larger_funcs_container:
306 count_map[func_i] += 1
308 return {k: v for k, v in count_map.items() if v >= 2}
310 def get_subset_candidates(self, argset, restrict_to_funcset=None):
311 """
312 Return a set of functions each of which whose argument list contains
313 ``argset``, optionally filtered only to contain functions in
314 ``restrict_to_funcset``.
315 """
316 iarg = iter(argset)
318 indices = OrderedSet(
319 fi for fi in self.arg_to_funcset[next(iarg)])
321 if restrict_to_funcset is not None:
322 indices &= restrict_to_funcset
324 for arg in iarg:
325 indices &= self.arg_to_funcset[arg]
327 return indices
329 def update_func_argset(self, func_i, new_argset):
330 """
331 Update a function with a new set of arguments.
332 """
333 new_args = OrderedSet(new_argset)
334 old_args = self.func_to_argset[func_i]
336 for deleted_arg in old_args - new_args:
337 self.arg_to_funcset[deleted_arg].remove(func_i)
338 for added_arg in new_args - old_args:
339 self.arg_to_funcset[added_arg].add(func_i)
341 self.func_to_argset[func_i].clear()
342 self.func_to_argset[func_i].update(new_args)
345class Unevaluated:
347 def __init__(self, func, args):
348 self.func = func
349 self.args = args
351 def __str__(self):
352 return "Uneval<{}>({})".format(
353 self.func, ", ".join(str(a) for a in self.args))
355 def as_unevaluated_basic(self):
356 return self.func(*self.args, evaluate=False)
358 @property
359 def free_symbols(self):
360 return set().union(*[a.free_symbols for a in self.args])
362 __repr__ = __str__
365def match_common_args(func_class, funcs, opt_subs):
366 """
367 Recognize and extract common subexpressions of function arguments within a
368 set of function calls. For instance, for the following function calls::
370 x + z + y
371 sin(x + y)
373 this will extract a common subexpression of `x + y`::
375 w = x + y
376 w + z
377 sin(w)
379 The function we work with is assumed to be associative and commutative.
381 Parameters
382 ==========
384 func_class: class
385 The function class (e.g. Add, Mul)
386 funcs: list of functions
387 A list of function calls.
388 opt_subs: dict
389 A dictionary of substitutions which this function may update.
390 """
392 # Sort to ensure that whole-function subexpressions come before the items
393 # that use them.
394 funcs = sorted(funcs, key=lambda f: len(f.args))
395 arg_tracker = FuncArgTracker(funcs)
397 changed = OrderedSet()
399 for i in range(len(funcs)):
400 common_arg_candidates_counts = arg_tracker.get_common_arg_candidates(
401 arg_tracker.func_to_argset[i], min_func_i=i + 1)
403 # Sort the candidates in order of match size.
404 # This makes us try combining smaller matches first.
405 common_arg_candidates = OrderedSet(sorted(
406 common_arg_candidates_counts.keys(),
407 key=lambda k: (common_arg_candidates_counts[k], k)))
409 while common_arg_candidates:
410 j = common_arg_candidates.pop(last=False)
412 com_args = arg_tracker.func_to_argset[i].intersection(
413 arg_tracker.func_to_argset[j])
415 if len(com_args) <= 1:
416 # This may happen if a set of common arguments was already
417 # combined in a previous iteration.
418 continue
420 # For all sets, replace the common symbols by the function
421 # over them, to allow recursive matches.
423 diff_i = arg_tracker.func_to_argset[i].difference(com_args)
424 if diff_i:
425 # com_func needs to be unevaluated to allow for recursive matches.
426 com_func = Unevaluated(
427 func_class, arg_tracker.get_args_in_value_order(com_args))
428 com_func_number = arg_tracker.get_or_add_value_number(com_func)
429 arg_tracker.update_func_argset(i, diff_i | OrderedSet([com_func_number]))
430 changed.add(i)
431 else:
432 # Treat the whole expression as a CSE.
433 #
434 # The reason this needs to be done is somewhat subtle. Within
435 # tree_cse(), to_eliminate only contains expressions that are
436 # seen more than once. The problem is unevaluated expressions
437 # do not compare equal to the evaluated equivalent. So
438 # tree_cse() won't mark funcs[i] as a CSE if we use an
439 # unevaluated version.
440 com_func_number = arg_tracker.get_or_add_value_number(funcs[i])
442 diff_j = arg_tracker.func_to_argset[j].difference(com_args)
443 arg_tracker.update_func_argset(j, diff_j | OrderedSet([com_func_number]))
444 changed.add(j)
446 for k in arg_tracker.get_subset_candidates(
447 com_args, common_arg_candidates):
448 diff_k = arg_tracker.func_to_argset[k].difference(com_args)
449 arg_tracker.update_func_argset(k, diff_k | OrderedSet([com_func_number]))
450 changed.add(k)
452 if i in changed:
453 opt_subs[funcs[i]] = Unevaluated(func_class,
454 arg_tracker.get_args_in_value_order(arg_tracker.func_to_argset[i]))
456 arg_tracker.stop_arg_tracking(i)
459def opt_cse(exprs, order='canonical'):
460 """Find optimization opportunities in Adds, Muls, Pows and negative
461 coefficient Muls.
463 Parameters
464 ==========
466 exprs : list of SymPy expressions
467 The expressions to optimize.
468 order : string, 'none' or 'canonical'
469 The order by which Mul and Add arguments are processed. For large
470 expressions where speed is a concern, use the setting order='none'.
472 Returns
473 =======
475 opt_subs : dictionary of expression substitutions
476 The expression substitutions which can be useful to optimize CSE.
478 Examples
479 ========
481 >>> from sympy.simplify.cse_main import opt_cse
482 >>> from sympy.abc import x
483 >>> opt_subs = opt_cse([x**-2])
484 >>> k, v = list(opt_subs.keys())[0], list(opt_subs.values())[0]
485 >>> print((k, v.as_unevaluated_basic()))
486 (x**(-2), 1/(x**2))
487 """
488 opt_subs = {}
490 adds = OrderedSet()
491 muls = OrderedSet()
493 seen_subexp = set()
494 collapsible_subexp = set()
496 def _find_opts(expr):
498 if not isinstance(expr, (Basic, Unevaluated)):
499 return
501 if expr.is_Atom or expr.is_Order:
502 return
504 if iterable(expr):
505 list(map(_find_opts, expr))
506 return
508 if expr in seen_subexp:
509 return expr
510 seen_subexp.add(expr)
512 list(map(_find_opts, expr.args))
514 if not isinstance(expr, MatrixExpr) and expr.could_extract_minus_sign():
515 # XXX -expr does not always work rigorously for some expressions
516 # containing UnevaluatedExpr.
517 # https://github.com/sympy/sympy/issues/24818
518 if isinstance(expr, Add):
519 neg_expr = Add(*(-i for i in expr.args))
520 else:
521 neg_expr = -expr
523 if not neg_expr.is_Atom:
524 opt_subs[expr] = Unevaluated(Mul, (S.NegativeOne, neg_expr))
525 seen_subexp.add(neg_expr)
526 expr = neg_expr
528 if isinstance(expr, (Mul, MatMul)):
529 if len(expr.args) == 1:
530 collapsible_subexp.add(expr)
531 else:
532 muls.add(expr)
534 elif isinstance(expr, (Add, MatAdd)):
535 if len(expr.args) == 1:
536 collapsible_subexp.add(expr)
537 else:
538 adds.add(expr)
540 elif isinstance(expr, Inverse):
541 # Do not want to treat `Inverse` as a `MatPow`
542 pass
544 elif isinstance(expr, (Pow, MatPow)):
545 base, exp = expr.base, expr.exp
546 if exp.could_extract_minus_sign():
547 opt_subs[expr] = Unevaluated(Pow, (Pow(base, -exp), -1))
549 for e in exprs:
550 if isinstance(e, (Basic, Unevaluated)):
551 _find_opts(e)
553 # Handle collapsing of multinary operations with single arguments
554 edges = [(s, s.args[0]) for s in collapsible_subexp
555 if s.args[0] in collapsible_subexp]
556 for e in reversed(topological_sort((collapsible_subexp, edges))):
557 opt_subs[e] = opt_subs.get(e.args[0], e.args[0])
559 # split muls into commutative
560 commutative_muls = OrderedSet()
561 for m in muls:
562 c, nc = m.args_cnc(cset=False)
563 if c:
564 c_mul = m.func(*c)
565 if nc:
566 if c_mul == 1:
567 new_obj = m.func(*nc)
568 else:
569 if isinstance(m, MatMul):
570 new_obj = m.func(c_mul, *nc, evaluate=False)
571 else:
572 new_obj = m.func(c_mul, m.func(*nc), evaluate=False)
573 opt_subs[m] = new_obj
574 if len(c) > 1:
575 commutative_muls.add(c_mul)
577 match_common_args(Add, adds, opt_subs)
578 match_common_args(Mul, commutative_muls, opt_subs)
580 return opt_subs
583def tree_cse(exprs, symbols, opt_subs=None, order='canonical', ignore=()):
584 """Perform raw CSE on expression tree, taking opt_subs into account.
586 Parameters
587 ==========
589 exprs : list of SymPy expressions
590 The expressions to reduce.
591 symbols : infinite iterator yielding unique Symbols
592 The symbols used to label the common subexpressions which are pulled
593 out.
594 opt_subs : dictionary of expression substitutions
595 The expressions to be substituted before any CSE action is performed.
596 order : string, 'none' or 'canonical'
597 The order by which Mul and Add arguments are processed. For large
598 expressions where speed is a concern, use the setting order='none'.
599 ignore : iterable of Symbols
600 Substitutions containing any Symbol from ``ignore`` will be ignored.
601 """
602 if opt_subs is None:
603 opt_subs = {}
605 ## Find repeated sub-expressions
607 to_eliminate = set()
609 seen_subexp = set()
610 excluded_symbols = set()
612 def _find_repeated(expr):
613 if not isinstance(expr, (Basic, Unevaluated)):
614 return
616 if isinstance(expr, RootOf):
617 return
619 if isinstance(expr, Basic) and (
620 expr.is_Atom or
621 expr.is_Order or
622 isinstance(expr, (MatrixSymbol, MatrixElement))):
623 if expr.is_Symbol:
624 excluded_symbols.add(expr)
625 return
627 if iterable(expr):
628 args = expr
630 else:
631 if expr in seen_subexp:
632 for ign in ignore:
633 if ign in expr.free_symbols:
634 break
635 else:
636 to_eliminate.add(expr)
637 return
639 seen_subexp.add(expr)
641 if expr in opt_subs:
642 expr = opt_subs[expr]
644 args = expr.args
646 list(map(_find_repeated, args))
648 for e in exprs:
649 if isinstance(e, Basic):
650 _find_repeated(e)
652 ## Rebuild tree
654 # Remove symbols from the generator that conflict with names in the expressions.
655 symbols = (symbol for symbol in symbols if symbol not in excluded_symbols)
657 replacements = []
659 subs = {}
661 def _rebuild(expr):
662 if not isinstance(expr, (Basic, Unevaluated)):
663 return expr
665 if not expr.args:
666 return expr
668 if iterable(expr):
669 new_args = [_rebuild(arg) for arg in expr.args]
670 return expr.func(*new_args)
672 if expr in subs:
673 return subs[expr]
675 orig_expr = expr
676 if expr in opt_subs:
677 expr = opt_subs[expr]
679 # If enabled, parse Muls and Adds arguments by order to ensure
680 # replacement order independent from hashes
681 if order != 'none':
682 if isinstance(expr, (Mul, MatMul)):
683 c, nc = expr.args_cnc()
684 if c == [1]:
685 args = nc
686 else:
687 args = list(ordered(c)) + nc
688 elif isinstance(expr, (Add, MatAdd)):
689 args = list(ordered(expr.args))
690 else:
691 args = expr.args
692 else:
693 args = expr.args
695 new_args = list(map(_rebuild, args))
696 if isinstance(expr, Unevaluated) or new_args != args:
697 new_expr = expr.func(*new_args)
698 else:
699 new_expr = expr
701 if orig_expr in to_eliminate:
702 try:
703 sym = next(symbols)
704 except StopIteration:
705 raise ValueError("Symbols iterator ran out of symbols.")
707 if isinstance(orig_expr, MatrixExpr):
708 sym = MatrixSymbol(sym.name, orig_expr.rows,
709 orig_expr.cols)
711 subs[orig_expr] = sym
712 replacements.append((sym, new_expr))
713 return sym
715 else:
716 return new_expr
718 reduced_exprs = []
719 for e in exprs:
720 if isinstance(e, Basic):
721 reduced_e = _rebuild(e)
722 else:
723 reduced_e = e
724 reduced_exprs.append(reduced_e)
725 return replacements, reduced_exprs
728def cse(exprs, symbols=None, optimizations=None, postprocess=None,
729 order='canonical', ignore=(), list=True):
730 """ Perform common subexpression elimination on an expression.
732 Parameters
733 ==========
735 exprs : list of SymPy expressions, or a single SymPy expression
736 The expressions to reduce.
737 symbols : infinite iterator yielding unique Symbols
738 The symbols used to label the common subexpressions which are pulled
739 out. The ``numbered_symbols`` generator is useful. The default is a
740 stream of symbols of the form "x0", "x1", etc. This must be an
741 infinite iterator.
742 optimizations : list of (callable, callable) pairs
743 The (preprocessor, postprocessor) pairs of external optimization
744 functions. Optionally 'basic' can be passed for a set of predefined
745 basic optimizations. Such 'basic' optimizations were used by default
746 in old implementation, however they can be really slow on larger
747 expressions. Now, no pre or post optimizations are made by default.
748 postprocess : a function which accepts the two return values of cse and
749 returns the desired form of output from cse, e.g. if you want the
750 replacements reversed the function might be the following lambda:
751 lambda r, e: return reversed(r), e
752 order : string, 'none' or 'canonical'
753 The order by which Mul and Add arguments are processed. If set to
754 'canonical', arguments will be canonically ordered. If set to 'none',
755 ordering will be faster but dependent on expressions hashes, thus
756 machine dependent and variable. For large expressions where speed is a
757 concern, use the setting order='none'.
758 ignore : iterable of Symbols
759 Substitutions containing any Symbol from ``ignore`` will be ignored.
760 list : bool, (default True)
761 Returns expression in list or else with same type as input (when False).
763 Returns
764 =======
766 replacements : list of (Symbol, expression) pairs
767 All of the common subexpressions that were replaced. Subexpressions
768 earlier in this list might show up in subexpressions later in this
769 list.
770 reduced_exprs : list of SymPy expressions
771 The reduced expressions with all of the replacements above.
773 Examples
774 ========
776 >>> from sympy import cse, SparseMatrix
777 >>> from sympy.abc import x, y, z, w
778 >>> cse(((w + x + y + z)*(w + y + z))/(w + x)**3)
779 ([(x0, y + z), (x1, w + x)], [(w + x0)*(x0 + x1)/x1**3])
782 List of expressions with recursive substitutions:
784 >>> m = SparseMatrix([x + y, x + y + z])
785 >>> cse([(x+y)**2, x + y + z, y + z, x + z + y, m])
786 ([(x0, x + y), (x1, x0 + z)], [x0**2, x1, y + z, x1, Matrix([
787 [x0],
788 [x1]])])
790 Note: the type and mutability of input matrices is retained.
792 >>> isinstance(_[1][-1], SparseMatrix)
793 True
795 The user may disallow substitutions containing certain symbols:
797 >>> cse([y**2*(x + 1), 3*y**2*(x + 1)], ignore=(y,))
798 ([(x0, x + 1)], [x0*y**2, 3*x0*y**2])
800 The default return value for the reduced expression(s) is a list, even if there is only
801 one expression. The `list` flag preserves the type of the input in the output:
803 >>> cse(x)
804 ([], [x])
805 >>> cse(x, list=False)
806 ([], x)
807 """
808 if not list:
809 return _cse_homogeneous(exprs,
810 symbols=symbols, optimizations=optimizations,
811 postprocess=postprocess, order=order, ignore=ignore)
813 if isinstance(exprs, (int, float)):
814 exprs = sympify(exprs)
816 # Handle the case if just one expression was passed.
817 if isinstance(exprs, (Basic, MatrixBase)):
818 exprs = [exprs]
820 copy = exprs
821 temp = []
822 for e in exprs:
823 if isinstance(e, (Matrix, ImmutableMatrix)):
824 temp.append(Tuple(*e.flat()))
825 elif isinstance(e, (SparseMatrix, ImmutableSparseMatrix)):
826 temp.append(Tuple(*e.todok().items()))
827 else:
828 temp.append(e)
829 exprs = temp
830 del temp
832 if optimizations is None:
833 optimizations = []
834 elif optimizations == 'basic':
835 optimizations = basic_optimizations
837 # Preprocess the expressions to give us better optimization opportunities.
838 reduced_exprs = [preprocess_for_cse(e, optimizations) for e in exprs]
840 if symbols is None:
841 symbols = numbered_symbols(cls=Symbol)
842 else:
843 # In case we get passed an iterable with an __iter__ method instead of
844 # an actual iterator.
845 symbols = iter(symbols)
847 # Find other optimization opportunities.
848 opt_subs = opt_cse(reduced_exprs, order)
850 # Main CSE algorithm.
851 replacements, reduced_exprs = tree_cse(reduced_exprs, symbols, opt_subs,
852 order, ignore)
854 # Postprocess the expressions to return the expressions to canonical form.
855 exprs = copy
856 for i, (sym, subtree) in enumerate(replacements):
857 subtree = postprocess_for_cse(subtree, optimizations)
858 replacements[i] = (sym, subtree)
859 reduced_exprs = [postprocess_for_cse(e, optimizations)
860 for e in reduced_exprs]
862 # Get the matrices back
863 for i, e in enumerate(exprs):
864 if isinstance(e, (Matrix, ImmutableMatrix)):
865 reduced_exprs[i] = Matrix(e.rows, e.cols, reduced_exprs[i])
866 if isinstance(e, ImmutableMatrix):
867 reduced_exprs[i] = reduced_exprs[i].as_immutable()
868 elif isinstance(e, (SparseMatrix, ImmutableSparseMatrix)):
869 m = SparseMatrix(e.rows, e.cols, {})
870 for k, v in reduced_exprs[i]:
871 m[k] = v
872 if isinstance(e, ImmutableSparseMatrix):
873 m = m.as_immutable()
874 reduced_exprs[i] = m
876 if postprocess is None:
877 return replacements, reduced_exprs
879 return postprocess(replacements, reduced_exprs)
882def _cse_homogeneous(exprs, **kwargs):
883 """
884 Same as ``cse`` but the ``reduced_exprs`` are returned
885 with the same type as ``exprs`` or a sympified version of the same.
887 Parameters
888 ==========
890 exprs : an Expr, iterable of Expr or dictionary with Expr values
891 the expressions in which repeated subexpressions will be identified
892 kwargs : additional arguments for the ``cse`` function
894 Returns
895 =======
897 replacements : list of (Symbol, expression) pairs
898 All of the common subexpressions that were replaced. Subexpressions
899 earlier in this list might show up in subexpressions later in this
900 list.
901 reduced_exprs : list of SymPy expressions
902 The reduced expressions with all of the replacements above.
904 Examples
905 ========
907 >>> from sympy.simplify.cse_main import cse
908 >>> from sympy import cos, Tuple, Matrix
909 >>> from sympy.abc import x
910 >>> output = lambda x: type(cse(x, list=False)[1])
911 >>> output(1)
912 <class 'sympy.core.numbers.One'>
913 >>> output('cos(x)')
914 <class 'str'>
915 >>> output(cos(x))
916 cos
917 >>> output(Tuple(1, x))
918 <class 'sympy.core.containers.Tuple'>
919 >>> output(Matrix([[1,0], [0,1]]))
920 <class 'sympy.matrices.dense.MutableDenseMatrix'>
921 >>> output([1, x])
922 <class 'list'>
923 >>> output((1, x))
924 <class 'tuple'>
925 >>> output({1, x})
926 <class 'set'>
927 """
928 if isinstance(exprs, str):
929 replacements, reduced_exprs = _cse_homogeneous(
930 sympify(exprs), **kwargs)
931 return replacements, repr(reduced_exprs)
932 if isinstance(exprs, (list, tuple, set)):
933 replacements, reduced_exprs = cse(exprs, **kwargs)
934 return replacements, type(exprs)(reduced_exprs)
935 if isinstance(exprs, dict):
936 keys = list(exprs.keys()) # In order to guarantee the order of the elements.
937 replacements, values = cse([exprs[k] for k in keys], **kwargs)
938 reduced_exprs = dict(zip(keys, values))
939 return replacements, reduced_exprs
941 try:
942 replacements, (reduced_exprs,) = cse(exprs, **kwargs)
943 except TypeError: # For example 'mpf' objects
944 return [], exprs
945 else:
946 return replacements, reduced_exprs