Python sympy.Add() Examples

def _print_Sum(self, expr):
        if len(expr.limits) == 1:
            tex = r"\sum_{%s=%s}^{%s} " % \
                tuple([ self._print(i) for i in expr.limits[0] ])
        else:
            def _format_ineq(l):
                return r"%s \leq %s \leq %s" % \
                    tuple([self._print(s) for s in (l[1], l[0], l[2])])

            tex = r"\sum_{\substack{%s}} " % \
                str.join('\\\\', [ _format_ineq(l) for l in expr.limits ])

        if isinstance(expr.function, Add):
            tex += r"\left(%s\right)" % self._print(expr.function)
        else:
            tex += self._print(expr.function)

        return tex 

def _eval_expand_tensorproduct(self, **hints):
        """Distribute TensorProducts across addition."""
        args = self.args
        add_args = []
        stop = False
        for i in range(len(args)):
            if isinstance(args[i], Add):
                for aa in args[i].args:
                    tp = TensorProduct(*args[:i] + (aa,) + args[i + 1:])
                    if isinstance(tp, TensorProduct):
                        tp = tp._eval_expand_tensorproduct()
                    add_args.append(tp)
                break

        if add_args:
            return Add(*add_args)
        else:
            return self 

def _normal_ordered_form_terms(expr, independent=False, recursive_limit=10,
                               _recursive_depth=0):
    """
    Helper function for normal_ordered_form: loop through each term in an
    addition expression and call _normal_ordered_form_factor to perform the
    factor to an normally ordered expression.
    """

    new_terms = []
    for term in expr.args:
        if isinstance(term, Mul):
            new_term = _normal_ordered_form_factor(
                term, recursive_limit=recursive_limit,
                _recursive_depth=_recursive_depth, independent=independent)
            new_terms.append(new_term)
        elif isinstance(term, Expectation):
            term = Expectation(normal_ordered_form(term.args[0]), term.is_normal_order)
            new_terms.append(term)
        else:
            new_terms.append(term)

    return Add(*new_terms) 

def _eval_expand_expectation(self, **hints):
        A = self.args[0]
        if isinstance(A, Add):
            # <A + B> = <A> + <B>
            return Add(*(Expectation(a, self.is_normal_order).expand(expectation=True) for a in A.args))

        if isinstance(A, Mul):
            # <c A> = c<A> where c is a commutative term
            A = A.expand()
            cA, ncA = A.args_cnc()
            return Mul(Mul(*cA), Expectation(Mul._from_args(ncA), self.is_normal_order).expand())
        
        if isinstance(A, Integral):
            # <∫adx> ->  ∫<a>dx
            func, lims = A.function, A.limits
            new_args = [Expectation(func, self.is_normal_order).expand()]
            for lim in lims:
                new_args.append(lim)
            return Integral(*new_args)
        
        return self 

def _replace_op_func(e, variable):

    if isinstance(e, Operator):
        return OperatorFunction(e, variable)
    
    if e.is_Number:
        return e
    
    if isinstance(e, Pow):
        return Pow(_replace_op_func(e.base, variable), e.exp)
    
    new_args = [_replace_op_func(arg, variable) for arg in e.args]
    
    if isinstance(e, Add):
        return Add(*new_args)
    
    elif isinstance(e, Mul):
        return Mul(*new_args)
    
    else:
        return e 

def test_differentiable():
    a = Function(name="a", grid=Grid((10, 10)))
    e = Function(name="e", grid=Grid((10, 10)))

    assert isinstance(1.2 * a.dx, Mul)
    assert isinstance(e + a, Add)
    assert isinstance(e * a, Mul)
    assert isinstance(a * a, Pow)
    assert isinstance(1 / (a * a), Pow)

    addition = a + 1.2 * a.dx
    assert isinstance(addition, Add)
    assert all(isinstance(a, Differentiable) for a in addition.args)

    addition2 = a + e * a.dx
    assert isinstance(addition2, Add)
    assert all(isinstance(a, Differentiable) for a in addition2.args) 

def test_diffify():
    a = Function(name="a", grid=Grid((10, 10)))
    e = Function(name="e", grid=Grid((10, 10)))

    assert isinstance(diffify(sympy.Mul(*[1.2, a.dx])), Mul)
    assert isinstance(diffify(sympy.Add(*[a, e])), Add)
    assert isinstance(diffify(sympy.Mul(*[e, a])), Mul)
    assert isinstance(diffify(sympy.Mul(*[a, a])), Pow)
    assert isinstance(diffify(sympy.Pow(*[a*a, -1])), Pow)

    addition = diffify(sympy.Add(*[a, sympy.Mul(*[1.2, a.dx])]))
    assert isinstance(addition, Add)
    assert all(isinstance(a, Differentiable) for a in addition.args)

    addition2 = diffify(sympy.Add(*[a, sympy.Mul(*[e, a.dx])]))
    assert isinstance(addition2, Add)
    assert all(isinstance(a, Differentiable) for a in addition2.args) 

def _sympystr(self, print_obj):
        if len(self.terms) == 0:
            return ZERO_STR

        self = self._with_sorted_terms()
        s = ''
        for coef, pdop in self.terms:
            coef_str = print_obj.doprint(coef)
            pd_str = print_obj.doprint(pdop)

            if coef == S(1):
                s += pd_str
            elif coef == S(-1):
                s += '-' + pd_str
            else:
                if isinstance(coef, Add):
                    s += '(' + coef_str + ')*' + pd_str
                else:
                    s += coef_str + '*' + pd_str
            s += ' + '

        s = s.replace('+ -', '- ')
        s = s[:-3]
        return s 

def indexify(self, indices=None, lshift=False, subs=None):
        """Create a types.Indexed from the current object."""
        if indices is not None:
            return Indexed(self.indexed, *indices)

        # Substitution for each index (spacing only used in own dimension)
        subs = subs or {}
        subs = [{**{d.spacing: 1, -d.spacing: -1}, **subs} for d in self.dimensions]

        # Add halo shift
        shift = self._size_nodomain.left if lshift else tuple([0]*len(self.dimensions))
        # Indices after substitutions
        indices = [sympy.sympify((a - o + f).xreplace(s)) for a, o, f, s in
                   zip(self.args, self.origin, shift, subs)]
        indices = [i.xreplace({k: sympy.Integer(k) for k in i.atoms(sympy.Float)})
                   for i in indices]
        return self.indexed[indices] 

def _print_Product(self, expr):
        if len(expr.limits) == 1:
            tex = r"\prod_{%s=%s}^{%s} " % \
                tuple([ self._print(i) for i in expr.limits[0] ])
        else:
            def _format_ineq(l):
                return r"%s \leq %s \leq %s" % \
                    tuple([self._print(s) for s in (l[1], l[0], l[2])])

            tex = r"\prod_{\substack{%s}} " % \
                str.join('\\\\', [ _format_ineq(l) for l in expr.limits ])

        if isinstance(expr.function, Add):
            tex += r"\left(%s\right)" % self._print(expr.function)
        else:
            tex += self._print(expr.function)

        return tex 

def sympy_to_grim(expr, **kwargs):
    import sympy
    assert isinstance(expr, sympy.Expr)
    if expr.is_Integer:
        return Expr(int(expr))
    if expr.is_Symbol:
        return Expr(symbol_name=expr.name)
    if expr is sympy.pi:
        return Pi
    if expr is sympy.E:
        return ConstE
    if expr is sympy.I:
        return ConstI
    if expr.is_Add:
        args = [sympy_to_grim(x, **kwargs) for x in expr.args]
        return Add(*args)
    if expr.is_Mul:
        args = [sympy_to_grim(x, **kwargs) for x in expr.args]
        return Mul(*args)
    if expr.is_Pow:
        args = [sympy_to_grim(x, **kwargs) for x in expr.args]
        b, e = args
        return Pow(b, e)
    raise NotImplementedError("converting %s to Grim", type(expr)) 

def __new__(cls, *args, **kwargs):
        obj = cls.__base__.__new__(cls, *args, **kwargs)

        # Unfortunately SymPy may build new sympy.core objects (e.g., sympy.Add),
        # so here we have to rebuild them as devito.core objects
        if kwargs.get('evaluate', True):
            obj = diffify(obj)

        return obj 

def convert_postfix(postfix):
    if hasattr(postfix, 'exp'):
        exp_nested = postfix.exp()
    else:
        exp_nested = postfix.exp_nofunc()

    exp = convert_exp(exp_nested)
    for op in postfix.postfix_op():
        if op.BANG():
            if isinstance(exp, list):
                raise Exception("Cannot apply postfix to derivative")
            exp = sympy.factorial(exp, evaluate=False)
        elif op.eval_at():
            ev = op.eval_at()
            at_b = None
            at_a = None
            if ev.eval_at_sup():
                at_b = do_subs(exp, ev.eval_at_sup()) 
            if ev.eval_at_sub():
                at_a = do_subs(exp, ev.eval_at_sub())
            if at_b != None and at_a != None:
                exp = sympy.Add(at_b, -1 * at_a, evaluate=False)
            elif at_b != None:
                exp = at_b
            elif at_a != None:
                exp = at_a
            
    return exp 

def convert_add(add):
    if add.ADD():
       lh = convert_add(add.additive(0))
       rh = convert_add(add.additive(1))
       return sympy.Add(lh, rh, evaluate=False)
    elif add.SUB():
        lh = convert_add(add.additive(0))
        rh = convert_add(add.additive(1))
        return sympy.Add(lh, -1 * rh, evaluate=False)
    else:
        return convert_mp(add.mp()) 

def test_parameter_add():
    """
    Makes sure the __add__ method of Parameters behaves as expected.
    """
    a = Parameter(value=1.0, min=0.5, max=1.5)
    b = Parameter(value=1.0, min=0.0)
    new = a + b
    assert isinstance(new, sympy.Add) 

def collect_const(expr):
    """
    *Much* faster alternative to sympy.collect_const. *Potentially* slightly less
    powerful with complex expressions, but it is equally effective for the
    expressions we have to deal with.
    """
    # Running example: `a*3. + b*2. + c*3.`

    # -> {a: 3., b: 2., c: 3.}
    mapper = expr.as_coefficients_dict()

    # -> {3.: [a, c], 2: [b]}
    inverse_mapper = defaultdict(list)
    for k, v in mapper.items():
        if v >= 0:
            inverse_mapper[v].append(k)
        else:
            inverse_mapper[-v].append(-k)

    terms = []
    for k, v in inverse_mapper.items():
        if len(v) == 1 and not v[0].is_Add:
            # Special case: avoid e.g. (-2)*a
            mul = Mul(k, *v)
        elif all(i.is_Mul and len(i.args) == 2 and i.args[0] == -1 for i in v):
            # Other special case: [-a, -b, -c ...]
            add = Add(*[i.args[1] for i in v], evaluate=False)
            mul = Mul(-k, add, evaluate=False)
        else:
            # Back to the running example
            # -> (a + c)
            add = Add(*v)
            # -> 3.*(a + c)
            mul = Mul(k, add, evaluate=False)

        terms.append(mul)

    return Add(*terms) 

def normal_order(expr, recursive_limit=10, _recursive_depth=0):
    """Normal order an expression with bosonic or fermionic operators. Note
    that this normal order is not equivalent to the original expression, but
    the creation and annihilation operators in each term in expr is reordered
    so that the expression becomes normal ordered.

    Parameters
    ==========

    expr : expression
        The expression to normal order.

    recursive_limit : int (default 10)
        The number of allowed recursive applications of the function.

    Examples
    ========

    >>> from sympsi import Dagger
    >>> from sympsi.boson import BosonOp
    >>> from sympsi.operatorordering import normal_order
    >>> a = BosonOp("a")
    >>> normal_order(a * Dagger(a))
    Dagger(a)*a
    """
    if _recursive_depth > recursive_limit:
        warnings.warn("Too many recursions, aborting")
        return expr

    if isinstance(expr, Add):
        return _normal_order_terms(expr, recursive_limit=recursive_limit,
                                   _recursive_depth=_recursive_depth)
    elif isinstance(expr, Mul):
        return _normal_order_factor(expr, recursive_limit=recursive_limit,
                                    _recursive_depth=_recursive_depth)
    else:
        return expr 

def doit(self, **hints):
        """Expand the density operator into an outer product format.

        Examples
        =========

        >>> from sympsi.state import Ket
        >>> from sympsi.density import Density
        >>> from sympsi.operator import Operator
        >>> A = Operator('A')
        >>> d = Density([Ket(0), 0.5], [Ket(1),0.5])
        >>> d.doit()
        0.5*|0><0| + 0.5*|1><1|

        """

        terms = []
        for (state, prob) in self.args:
            state = state.expand()  # needed to break up (a+b)*c
            if (isinstance(state, Add)):
                for arg in product(state.args, repeat=2):
                    terms.append(prob *
                                 self._generate_outer_prod(arg[0], arg[1]))
            else:
                terms.append(prob *
                             self._generate_outer_prod(state, state))

        return Add(*terms) 

def _is_supported_formula(formula: sympy.Basic) -> bool:
    if isinstance(formula, (sympy.Symbol, sympy.Integer, sympy.Float,
                            sympy.Rational, sympy.NumberSymbol)):
        return True
    if isinstance(formula, (sympy.Add, sympy.Mul)):
        return all(_is_supported_formula(f) for f in formula.args)
    return False 

def _implicit_conversion(obj):
    if isinstance(obj, (int, float)):
        return Constant(obj)
    elif isinstance(obj, Expr):
        return obj
    elif isinstance(obj, str):
        return Symbol(unique_keys=(obj,))

    if sympy is not None:
        if isinstance(obj, sympy.Mul):
            if len(obj.args) != 2:
                raise NotImplementedError("Did you use evaluate=False?")
            return _MulExpr([_implicit_conversion(obj.args[0]), _implicit_conversion(obj.args[1])])
        elif isinstance(obj, sympy.Add):
            if len(obj.args) != 2:
                raise NotImplementedError("Did you use evaluate=False?")
            return _AddExpr([_implicit_conversion(obj.args[0]), _implicit_conversion(obj.args[1])])
        elif isinstance(obj, sympy.Pow):
            return _PowExpr(_implicit_conversion(obj.base), _implicit_conversion(obj.exp))
        elif isinstance(obj, sympy.Float):
            return Constant(float(obj))
        elif isinstance(obj, sympy.Symbol):
            return Symbol(unique_keys=(obj.name,))

    raise NotImplementedError(
        "Don't know how to convert %s (of type %s)" % (obj, type(obj))) 

def _doit(obj, args=None):
        cls = diffify._cls(obj)
        args = args or obj.args

        if cls is obj.__class__:
            # Try to just update the args if possible (Add, Mul)
            try:
                return obj._new_rawargs(*args, is_commutative=obj.is_commutative)
            # Or just return the object (Float, Symbol, Function, ...)
            except AttributeError:
                return obj

        # Create object directly from args, avoid any rebuild
        return cls(*args, evaluate=False) 

def __rsub__(self, other):
        return Add(other, -self) 

def biharmonic(self, weight=1):
        """
        Generates a symbolic expression for the weighted biharmonic operator w.r.t.
        all spatial Dimensions Laplace(weight * Laplace (self))
        """
        space_dims = [d for d in self.dimensions if d.is_Space]
        derivs = tuple('d%s2' % d.name for d in space_dims)
        return Add(*[getattr(self.laplace * weight, d) for d in derivs]) 

def symbolic_shape(self):
        """
        The symbolic shape of the object. This includes the domain, halo, and
        padding regions. While halo and padding are known quantities (integers),
        the domain size is given as a symbol.
        """
        halo = [sympy.Add(*i, evaluate=False) for i in self._size_halo]
        padding = [sympy.Add(*i, evaluate=False) for i in self._size_padding]
        domain = [i.symbolic_size for i in self.dimensions]
        ret = tuple(sympy.Add(i, j, k, evaluate=False)
                    for i, j, k in zip(domain, halo, padding))
        return DimensionTuple(*ret, getters=self.dimensions) 

def laplace(self):
        """
        Generates a symbolic expression for the Laplacian, the second
        derivative w.r.t all spatial Dimensions.
        """
        space_dims = [d for d in self.dimensions if d.is_Space]
        derivs = tuple('d%s2' % d.name for d in space_dims)
        return Add(*[getattr(self, d) for d in derivs]) 

def _surd_coefficients(sympy_exp):
  """Extracts coefficients a, b, where sympy_exp = a + b * sqrt(base)."""
  sympy_exp = sympy.simplify(sympy.expand(sympy_exp))

  def extract_b(b_sqrt_base):
    """Returns b from expression of form b * sqrt(base)."""
    if isinstance(b_sqrt_base, sympy.Pow):
      # Just form sqrt(base)
      return 1
    else:
      assert isinstance(b_sqrt_base, sympy.Mul)
      assert len(b_sqrt_base.args) == 2
      assert b_sqrt_base.args[0].is_rational
      assert isinstance(b_sqrt_base.args[1], sympy.Pow)  # should be sqrt.
      return b_sqrt_base.args[0]

  if sympy_exp.is_rational:
    # Form: a.
    return sympy_exp, 0
  elif isinstance(sympy_exp, sympy.Add):
    # Form: a + b * sqrt(base)
    assert len(sympy_exp.args) == 2
    assert sympy_exp.args[0].is_rational
    a = sympy_exp.args[0]
    b = extract_b(sympy_exp.args[1])
    return a, b
  else:
    # Form: b * sqrt(base).
    return 0, extract_b(sympy_exp) 

def __add__(self, other):
        return Add(self, other) 

def __radd__(self, other):
        return Add(other, self) 

def __sub__(self, other):
        return Add(self, -other) 

def __isub__(self, other):
        return Add(self, -other)