How to not double count sympy numbers

Question

I need for a project to calculate some sort of simplicity of some expressions and among other things I need to assign to each number appearing in the expression a complexity (something like log2(x), such that the complexity is measured in bits, ie log2(4)=2 bits). Anyway, for a give sympy expression I have this piece of code:

is_atomic_number = lambda expr: expr.is_Atom and expr.is_number
eq_numbers = [subexpression for subexpression in preorder_traversal(expr) if is_atomic_number(subexpression)] 

It mainly works, but if I have something of the form 2*(x+y), that returns as the numbers appearing in the expression (2,2) instead of just 2, which is what I want. Is there a way to make sure that when something gets factored out, I count it only once? Thank you!

Answer 1

The atoms method might be what you want:

>>> (2*(x+y)).atoms(Number)
{2}

The replace method has the option too keep track of the mapping used, but it doesn't handle identical items that were replaced by different values

>>> (2*x + x**2).replace(lambda x:x.is_Number, lambda x: Dummy(), map=True)
(Dummy_1*x + x**Dummy_2, {2: Dummy_2})

So the following might work:

>>> def flatargs(expr):
...     if expr.is_Atom: yield expr
...     else:
...         for a in expr.args:
...             for i in flatargs(a): yield i
>>> list(flatargs(2*x+x**2))
[x, 2, 2, x]
>>> from sympy.utilities.iterables import multiset
>>> multiset([i for i in flatargs(2*x+x**2+y/2) if i.is_Number)
{1/2: 1, 2: 2}

If you want the numerator and denominator separately from flatargs you would have to handle the Rational atom as a special case and deal with the cases when numerator or denomitor is 1. But hopefully what is here will get you on your way.