def flatten(cls, seq): """Return commutative, noncommutative and order arguments by combining related terms. Notes ===== * In an expression like ``a*b*c``, python process this through sympy as ``Mul(Mul(a, b), c)``. This can have undesirable consequences. - Sometimes terms are not combined as one would like: {c.f. http://code.google.com/p/sympy/issues/detail?id=1497} >>> from sympy import Mul, sqrt >>> from sympy.abc import x, y, z >>> 2*(x + 1) # this is the 2-arg Mul behavior 2*x + 2 >>> y*(x + 1)*2 2*y*(x + 1) >>> 2*(x + 1)*y # 2-arg result will be obtained first y*(2*x + 2) >>> Mul(2, x + 1, y) # all 3 args simultaneously processed 2*y*(x + 1) >>> 2*((x + 1)*y) # parentheses can control this behavior 2*y*(x + 1) Powers with compound bases may not find a single base to combine with unless all arguments are processed at once. Post-processing may be necessary in such cases. {c.f. http://code.google.com/p/sympy/issues/detail?id=2629} >>> a = sqrt(x*sqrt(y)) >>> a**3 (x*sqrt(y))**(3/2) >>> Mul(a,a,a) (x*sqrt(y))**(3/2) >>> a*a*a x*sqrt(y)*sqrt(x*sqrt(y)) >>> _.subs(a.base, z).subs(z, a.base) (x*sqrt(y))**(3/2) - If more than two terms are being multiplied then all the previous terms will be re-processed for each new argument. So if each of ``a``, ``b`` and ``c`` were :class:`Mul` expression, then ``a*b*c`` (or building up the product with ``*=``) will process all the arguments of ``a`` and ``b`` twice: once when ``a*b`` is computed and again when ``c`` is multiplied. Using ``Mul(a, b, c)`` will process all arguments once. * The results of Mul are cached according to arguments, so flatten will only be called once for ``Mul(a, b, c)``. If you can structure a calculation so the arguments are most likely to be repeats then this can save time in computing the answer. For example, say you had a Mul, M, that you wished to divide by ``d[i]`` and multiply by ``n[i]`` and you suspect there are many repeats in ``n``. It would be better to compute ``M*n[i]/d[i]`` rather than ``M/d[i]*n[i]`` since every time n[i] is a repeat, the product, ``M*n[i]`` will be returned without flattening -- the cached value will be returned. If you divide by the ``d[i]`` first (and those are more unique than the ``n[i]``) then that will create a new Mul, ``M/d[i]`` the args of which will be traversed again when it is multiplied by ``n[i]``. {c.f. http://code.google.com/p/sympy/issues/detail?id=2607} This consideration is moot if the cache is turned off. NB -- The validity of the above notes depends on the implementation details of Mul and flatten which may change at any time. Therefore, you should only consider them when your code is highly performance sensitive. Removal of 1 from the sequence is already handled by AssocOp.__new__. """ rv = None if len(seq) == 2: a, b = seq if b.is_Rational: a, b = b, a assert not a is S.One if a and a.is_Rational: r, b = b.as_coeff_Mul() a *= r if b.is_Mul: bargs, nc = b.args_cnc() rv = bargs, nc, None if a is not S.One: bargs.insert(0, a) elif b.is_Add and b.is_commutative: if a is S.One: rv = [b], [], None else: r, b = b.as_coeff_Add() bargs = [_keep_coeff(a, bi) for bi in Add.make_args(b)] bargs.sort(key=hash) ar = a*r if ar: bargs.insert(0, ar) bargs = [Add._from_args(bargs)] rv = bargs, [], None if rv: return rv # apply associativity, separate commutative part of seq c_part = [] # out: commutative factors nc_part = [] # out: non-commutative factors nc_seq = [] coeff = S.One # standalone term # e.g. 3 * ... iu = [] # ImaginaryUnits, I c_powers = [] # (base,exp) n # e.g. (x,n) for x num_exp = [] # (num-base, exp) y # e.g. (3, y) for ... * 3 * ... neg1e = 0 # exponent on -1 extracted from Number-based Pow pnum_rat = {} # (num-base, Rat-exp) 1/2 # e.g. (3, 1/2) for ... * 3 * ... order_symbols = None # --- PART 1 --- # # "collect powers and coeff": # # o coeff # o c_powers # o num_exp # o neg1e # o pnum_rat # # NOTE: this is optimized for all-objects-are-commutative case for o in seq: # O(x) if o.is_Order: o, order_symbols = o.as_expr_variables(order_symbols) # Mul([...]) if o.is_Mul: if o.is_commutative: seq.extend(o.args) # XXX zerocopy? else: # NCMul can have commutative parts as well for q in o.args: if q.is_commutative: seq.append(q) else: nc_seq.append(q) # append non-commutative marker, so we don't forget to # process scheduled non-commutative objects seq.append(NC_Marker) continue # 3 elif o.is_Number: if o is S.NaN or coeff is S.ComplexInfinity and o is S.Zero: # we know for sure the result will be nan return [S.NaN], [], None elif coeff.is_Number: # it could be zoo coeff *= o if coeff is S.NaN: # we know for sure the result will be nan return [S.NaN], [], None continue elif o is S.ComplexInfinity: if not coeff: # 0 * zoo = NaN return [S.NaN], [], None if coeff is S.ComplexInfinity: # zoo * zoo = zoo return [S.ComplexInfinity], [], None coeff = S.ComplexInfinity continue elif o is S.ImaginaryUnit: iu.append(o) continue elif o.is_commutative: # e # o = b b, e = o.as_base_exp() # y # 3 if o.is_Pow and b.is_Number: # get all the factors with numeric base so they can be # combined below, but don't combine negatives unless # the exponent is an integer if e.is_Rational: if e.is_Integer: coeff *= Pow(b, e) # it is an unevaluated power continue elif e.is_negative: # also a sign of an unevaluated power seq.append(Pow(b, e)) continue elif b.is_negative: neg1e += e b = -b if b is not S.One: pnum_rat.setdefault(b, []).append(e) continue elif b.is_positive or e.is_integer: num_exp.append((b, e)) continue c_powers.append((b, e)) # NON-COMMUTATIVE # TODO: Make non-commutative exponents not combine automatically else: if o is not NC_Marker: nc_seq.append(o) # process nc_seq (if any) while nc_seq: o = nc_seq.pop(0) if not nc_part: nc_part.append(o) continue # b c b+c # try to combine last terms: a * a -> a o1 = nc_part.pop() b1, e1 = o1.as_base_exp() b2, e2 = o.as_base_exp() new_exp = e1 + e2 # Only allow powers to combine if the new exponent is # not an Add. This allow things like a**2*b**3 == a**5 # if a.is_commutative == False, but prohibits # a**x*a**y and x**a*x**b from combining (x,y commute). if b1 == b2 and (not new_exp.is_Add): o12 = b1 ** new_exp # now o12 could be a commutative object if o12.is_commutative: seq.append(o12) continue else: nc_seq.insert(0, o12) else: nc_part.append(o1) nc_part.append(o) # handle the ImaginaryUnits if iu: if len(iu) == 1: c_powers.append((iu[0], S.One)) else: # a product of I's has one of 4 values; select that value # based on the length of iu: # len(iu) % 4 of (0, 1, 2, 3) has a corresponding value of # (1, I,-1,-I) niu = len(iu) % 4 if niu % 2: c_powers.append((S.ImaginaryUnit, S.One)) if niu in (2, 3): coeff = -coeff # We do want a combined exponent if it would not be an Add, such as # y 2y 3y # x * x -> x # We determine if two exponents have the same term by using # as_coeff_Mul. # # Unfortunately, this isn't smart enough to consider combining into # exponents that might already be adds, so things like: # z - y y # x * x will be left alone. This is because checking every possible # combination can slow things down. # gather exponents of common bases... def _gather(c_powers): new_c_powers = [] common_b = {} # b:e for b, e in c_powers: co = e.as_coeff_Mul() common_b.setdefault(b, {}).setdefault(co[1], []).append(co[0]) for b, d in common_b.items(): for di, li in d.items(): d[di] = Add(*li) for b, e in common_b.items(): for t, c in e.items(): new_c_powers.append((b, c*t)) return new_c_powers # in c_powers c_powers = _gather(c_powers) # and in num_exp num_exp = _gather(num_exp) # --- PART 2 --- # # o process collected powers (x**0 -> 1; x**1 -> x; otherwise Pow) # o combine collected powers (2**x * 3**x -> 6**x) # with numeric base # ................................ # now we have: # - coeff: # - c_powers: (b, e) # - num_exp: (2, e) # - pnum_rat: {(1/3, [1/3, 2/3, 1/4])} # 0 1 # x -> 1 x -> x for b, e in c_powers: if e is S.One: if b.is_Number: coeff *= b else: c_part.append(b) elif e is not S.Zero: c_part.append(Pow(b, e)) # x x x # 2 * 3 -> 6 inv_exp_dict = {} # exp:Mul(num-bases) x x # e.g. x:6 for ... * 2 * 3 * ... for b, e in num_exp: inv_exp_dict.setdefault(e, []).append(b) for e, b in inv_exp_dict.items(): inv_exp_dict[e] = Mul(*b) c_part.extend([Pow(b, e) for e, b in inv_exp_dict.iteritems() if e]) # b, e -> e' = sum(e), b # {(1/5, [1/3]), (1/2, [1/12, 1/4]} -> {(1/3, [1/5, 1/2])} comb_e = {} for b, e in pnum_rat.iteritems(): comb_e.setdefault(Add(*e), []).append(b) del pnum_rat # process them, reducing exponents to values less than 1 # and updating coeff if necessary else adding them to # num_rat for further processing num_rat = [] for e, b in comb_e.iteritems(): b = Mul(*b) if e.q == 1: coeff *= Pow(b, e) continue if e.p > e.q: e_i, ep = divmod(e.p, e.q) coeff *= Pow(b, e_i) e = Rational(ep, e.q) num_rat.append((b, e)) del comb_e # extract gcd of bases in num_rat # 2**(1/3)*6**(1/4) -> 2**(1/3+1/4)*3**(1/4) pnew = defaultdict(list) i = 0 # steps through num_rat which may grow while i < len(num_rat): bi, ei = num_rat[i] grow = [] for j in range(i + 1, len(num_rat)): bj, ej = num_rat[j] g = _rgcd(bi, bj) if g is not S.One: # 4**r1*6**r2 -> 2**(r1+r2) * 2**r1 * 3**r2 # this might have a gcd with something else e = ei + ej if e.q == 1: coeff *= Pow(g, e) else: if e.p > e.q: e_i, ep = divmod(e.p, e.q) # change e in place coeff *= Pow(g, e_i) e = Rational(ep, e.q) grow.append((g, e)) # update the jth item num_rat[j] = (bj/g, ej) # update bi that we are checking with bi = bi/g if bi is S.One: break if bi is not S.One: obj = Pow(bi, ei) if obj.is_Number: coeff *= obj else: # changes like sqrt(12) -> 2*sqrt(3) for obj in Mul.make_args(obj): if obj.is_Number: coeff *= obj else: assert obj.is_Pow bi, ei = obj.args pnew[ei].append(bi) num_rat.extend(grow) i += 1 # combine bases of the new powers for e, b in pnew.iteritems(): pnew[e] = Mul(*b) # see if there is a base with matching coefficient # that the -1 can be joined with if neg1e: p = Pow(S.NegativeOne, neg1e) if p.is_Number: coeff *= p else: c, p = p.as_coeff_Mul() coeff *= c if p.is_Pow and p.base is S.NegativeOne: neg1e = p.exp for e, b in pnew.iteritems(): if e == neg1e and b.is_positive: pnew[e] = -b break else: c_part.append(p) # add all the pnew powers c_part.extend([Pow(b, e) for e, b in pnew.iteritems()]) # oo, -oo if (coeff is S.Infinity) or (coeff is S.NegativeInfinity): def _handle_for_oo(c_part, coeff_sign): new_c_part = [] for t in c_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_c_part.append(t) return new_c_part, coeff_sign c_part, coeff_sign = _handle_for_oo(c_part, 1) nc_part, coeff_sign = _handle_for_oo(nc_part, coeff_sign) coeff *= coeff_sign # zoo if coeff is S.ComplexInfinity: # zoo might be # unbounded_real + bounded_im # bounded_real + unbounded_im # unbounded_real + unbounded_im # and non-zero real or imaginary will not change that status. c_part = [c for c in c_part if not (c.is_nonzero and c.is_real is not None)] nc_part = [c for c in nc_part if not (c.is_nonzero and c.is_real is not None)] # 0 elif coeff is S.Zero: # we know for sure the result will be 0 return [coeff], [], order_symbols # order commutative part canonically c_part.sort(key=cls._args_sortkey) # current code expects coeff to be always in slot-0 if coeff is not S.One: c_part.insert(0, coeff) # we are done if len(c_part) == 2 and c_part[0].is_Number and c_part[1].is_Add: # 2*(1+a) -> 2 + 2 * a coeff = c_part[0] c_part = [Add(*[coeff*f for f in c_part[1].args])] return c_part, nc_part, order_symbols