CODEWITHSUNDEEP
pythonsympypretty-print
Michael M
Michael M

Reputation: 121

How to pretty-print Sympy Functions without arguments?

I am working with partial differential equations using a method of multiple scales and am trying to use Sympy to check my analysis. However, using multiple scales leads to many functional arguments which in turn leads to difficult to read output. I would like to suppress the arguments in the printing of the functions to make the output easier to read. For example,

would be replaced by

or

The closest answer I was able to find by searching online comes from the answer on SymPy print only function name, but my understanding is that this approach is only for outputting the input to LaTeX, rather than pretty-printing the output in a Jupyter notebook. Perhaps I am misunderstanding that answer, though.

Thank you!

Upvotes: 2

Views: 819

Answers (2)

Sam
Sam

Reputation: 46

I've been fighting sympy with the same thing; too many arguments just makes expressions hard to read. Although the 'name' attribute is pretty good at identifying the sympy Functions that we made, a better way to just process the sy.Functions you created is to use the AppliedUndef option when getting the 'atoms' of the expression.

I've ended up making this helper function. The problem becomes it is too effective in a way. Let's say you replace g(t) with just g, and there is a Derivative(g(t), t) in the expression. That just became 0. I put in a basic check for this, but clearly use this carefully. I fear there might be some other edge case I haven't thought of, but so far for the few cases I've worked with, this does exactly what I think you are asking.

import sympy as sy
from sympy.core.function import AppliedUndef
from typing import List
def cleanOutUnwantedArguments(exp : sy.Expr, argsToClean : List[sy.Symbol] = None) -> sy.Expr:
    """
    For sympy Functions you have made yourself from the Function type, 
    create an expression that doesn't have extra arguments.  For example:

    myExp = sy.Function('g')(x,y,t)*sy.cos(x)
    cleanedExpression = cleanOutUnwantedArguments(myExp, [x,t])
    
    ... will return:
    sy.Function('g')(y)*sy.cos(x)

    This is mainly used for pretty printing of expressions.  
    An empty or null argsToClean will remove all arguments

    Args:
        exp (sy.Expr): The expression to clean
        argsToClean (List[sy.Symbol]): The symbol arguments to clean from exp.  
        A null or empty list will clean all arguments (free_symbols)

    Returns:
        sy.Expr: An expression with the desired arguments cleaned 
    """

    if argsToClean == None:
        argsToClean = []
    for arg in exp.atoms(AppliedUndef) :
       
        symbolsToLeaveInFinalTerm = []
        for freeSymbol in arg.free_symbols :
            if len(argsToClean) == 0 :
                continue
            if freeSymbol not in argsToClean: 
                symbolsToLeaveInFinalTerm.append(freeSymbol)
                
        if len(symbolsToLeaveInFinalTerm) == 0 :
            rewrittenArg = sy.Symbol(arg.name)
        else:
            rewrittenArg = sy.Function(arg.name)(*symbolsToLeaveInFinalTerm)
        
        maybeExp = exp.subs(arg, rewrittenArg)    
        if maybeExp.simplify() != 0 : # if we (for example) replace a derivative(g(t), t) with just symbol g, that derivative becomes 0, not what we want
            exp = maybeExp 
    return exp

# adapted from a real unit test
x = sy.Symbol('x')
y = sy.Symbol('y')
t = sy.Symbol('t')
z = sy.Function('z')(t)
testExpression = sy.Function("g")(x,y,t)*sy.cos(x)*sy.Derivative(z, t)
cleanedExpression = cleanOutUnwantedArguments(testExpression, [x,t])

print(cleanedExpression)
expectedExpression = sy.Function("g")(y)*sy.cos(x)*sy.Derivative(z, t)
if(str(cleanedExpression) != "g(y)*cos(x)*Derivative(z(t), t)") :
    raise Exception("oh no")
if(cleanedExpression != expectedExpression) :
    raise Exception("oh no")
print("Success")

Upvotes: 0

Yang
Yang

Reputation: 1

I found a simple way to improve the format when viewing it in Jupyter Notebook, but frankly, this is NOT a complete solution, as you will see below. When you are going to display an expression containing the lengthy SymPy function itself or its derivatives, substitute the function with a SymPy symbol that is literally the same as the function's name.

In your case, the default output could be obtained by

from sympy import *
init_printing()
s0,s1,s2,t0,t1 = symbols('sigma_0,sigma_1,sigma_2,tau_0,tau_1')
Pi0 = Function('Pi_0')(s0,s1,s2,t0,t1)
diff(Pi0,t1)

It may be improved by replacing the last line with

diff(Pi0,t1).subs(Pi0,Symbol('Pi_0'))

As you may notice, although the expression becomes short, at the same time all the $\partial$ symbols are replaced with $d$ for the total derivative. I think it's fine just to check the results in Jupyter Notebook with this mathematically unstrict format, but some extra effort is needed when editing it in Latex (for example, use the 'find and replace' functionality of the editor).

Hope that helps a bit. And looking forward to a better answer to completely solve this issue.

Upvotes: 0

Related Questions