def flatten(cls, seq): """Return commutative, noncommutative and order arguments by combining related terms. ** Developer 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. """ 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 * ... 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 or coeff is S.ComplexInfinity: # we know for sure the result will be nan return [S.NaN], [], None coeff = S.ComplexInfinity 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) # We do want a combined exponent if it would not be an Add, such as # y 2y 3y # x * x -> x # We determine this if two exponents have the same term in 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... # in 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*Mul(*t))) c_powers = new_c_powers # and in num_exp new_num_exp = [] common_b = {} # b:e for b, e in num_exp: 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_num_exp.append((b,c*Mul(*t))) num_exp = new_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 = {} 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 = igcd(bi, bj) if g != 1: # 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: if obj.is_Mul: # sqrt(12) -> 2*sqrt(3) c, obj = obj.args # expecting only 2 args coeff *= c assert obj.is_Pow bi, ei = obj.args pnew.setdefault(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: if p.is_Mul: c, p = p.args coeff *= c assert p.is_Pow and p.base is S.NegativeOne neg1e = p.args[1] 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): new_c_part = [] coeff_sign = 1 for t in c_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_c_part.append(t) c_part = new_c_part new_nc_part = [] for t in nc_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_nc_part.append(t) nc_part = new_nc_part 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
def flatten(cls, seq): """Return commutative, noncommutative and order arguments by combining related terms. ** Developer 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. """ # 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 * ... 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 or coeff is S.ComplexInfinity: # we know for sure the result will be nan return [S.NaN], [], None coeff = S.ComplexInfinity 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) # We do want a combined exponent if it would not be an Add, such as # y 2y 3y # x * x -> x # We determine this if two exponents have the same term in 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... # in 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*Mul(*t))) c_powers = new_c_powers # and in num_exp new_num_exp = [] common_b = {} # b:e for b, e in num_exp: 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_num_exp.append((b,c*Mul(*t))) num_exp = new_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 not e is 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, 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 = {} 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 = igcd(bi, bj) if g != 1: # 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: if obj.is_Mul: # sqrt(12) -> 2*sqrt(3) c, obj = obj.args # expecting only 2 args coeff *= c assert obj.is_Pow bi, ei = obj.args pnew.setdefault(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: if p.is_Mul: c, p = p.args coeff *= c assert p.is_Pow and p.base is S.NegativeOne neg1e = p.args[1] 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): new_c_part = [] coeff_sign = 1 for t in c_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_c_part.append(t) c_part = new_c_part new_nc_part = [] for t in nc_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_nc_part.append(t) nc_part = new_nc_part 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=cmp_to_key(Basic.compare)) # 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
def flatten(cls, seq): # 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 * ... 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 or coeff is S.ComplexInfinity: # we know for sure the result will be nan return [S.NaN], [], None coeff = S.ComplexInfinity 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) # We do want a combined exponent if it would not be an Add, such as # y 2y 3y # x * x -> x # We determine this if two exponents have the same term in 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... # in 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 * Mul(*t))) c_powers = new_c_powers # and in num_exp new_num_exp = [] common_b = {} # b:e for b, e in num_exp: 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_num_exp.append((b, c * Mul(*t))) num_exp = new_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 not e is 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, 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 = {} 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 = igcd(bi, bj) if g != 1: # 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: if obj.is_Mul: # 12**(1/2) -> 2*sqrt(3) c, obj = obj.args # expecting only 2 args coeff *= c assert obj.is_Pow bi, ei = obj.args pnew.setdefault(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: if p.is_Mul: c, p = p.args coeff *= c assert p.is_Pow and p.base is S.NegativeOne neg1e = p.args[1] 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): new_c_part = [] coeff_sign = 1 for t in c_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_c_part.append(t) c_part = new_c_part new_nc_part = [] for t in nc_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_nc_part.append(t) nc_part = new_nc_part 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=cmp_to_key(Basic.compare)) # 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
def _rgcd(a, b): return Rational(Integer(igcd(a.p, b.p)), Integer(ilcm(a.q, b.q)))
def flatten(cls, seq): # 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 * ... 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 or coeff is S.ComplexInfinity: # we know for sure the result will be nan return [S.NaN], [], None coeff = S.ComplexInfinity 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) # We do want a combined exponent if it would not be an Add, such as # y 2y 3y # x * x -> x # We determine this if two exponents have the same term in 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... # in 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*Mul(*t))) c_powers = new_c_powers # and in num_exp new_num_exp = [] common_b = {} # b:e for b, e in num_exp: 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_num_exp.append((b,c*Mul(*t))) num_exp = new_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 not e is 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, 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 = {} 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 = igcd(bi, bj) if g != 1: # 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: if obj.is_Mul: # 12**(1/2) -> 2*sqrt(3) c, obj = obj.args # expecting only 2 args coeff *= c assert obj.is_Pow bi, ei = obj.args pnew.setdefault(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: if p.is_Mul: c, p = p.args coeff *= c assert p.is_Pow and p.base is S.NegativeOne neg1e = p.args[1] 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): new_c_part = [] coeff_sign = 1 for t in c_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_c_part.append(t) c_part = new_c_part new_nc_part = [] for t in nc_part: if t.is_positive: continue if t.is_negative: coeff_sign *= -1 continue new_nc_part.append(t) nc_part = new_nc_part 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=cmp_to_key(Basic.compare)) # 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