09.4 運算式樹與自訂函數#

本節目標#

理解 SymPy 運算式本質上是樹狀結構
用 srepr() 觀看運算式的完整樹狀表示
用 func 與 args 遍歷樹、用 preorder_traversal 走訪所有節點
學會用 Function 類定義自訂函數，並實作數值近似
認識 rewrite：把運算式改寫成另一種形式
動手實驗：寫一個函數統計運算式中出現的運算子種類

運算式是樹#

SymPy 的任何運算式都是一個物件樹。 \(x^2 + 2x + 1\) 不是「一段文字」，而是：

        Add
       / | \
     x^2 2x  1

x**2 本身又是 Pow(x, 2)，2*x 是 Mul(2, x)。整體是：

from sympy import symbols, srepr

x = symbols('x')
expr = x**2 + 2*x + 1
print(srepr(expr))
# Add(Pow(Symbol('x'), Integer(2)), Mul(Integer(2), Symbol('x')), Integer(1))

srepr 完整揭露了樹的結構：最外層是 Add，其子節點依序是 Pow、Mul 與 Integer。每一節點都有兩個基本屬性：

func：這個節點的「運算類別」，如 Add、Mul、sin
args：子節點組成的 tuple

用 func(*args) 可以精確地重建同一個運算式：

print(expr.func)          # <class 'sympy.core.add.Add'>
print(expr.args)          # (x**2, 2*x, 1)
print(expr.func(*expr.args) == expr)   # True

這個性質是遍歷與改寫運算式的基礎——任何運算式都能拆成「頭」加「子運算式們」，也能原樣重建。

遍歷樹：func 與 args、preorder_traversal#

因為每個節點都有 args，可以遞迴走完整棵樹（複習 4.1「函數即物件」：類別也是物件，可以傳遞）：

from sympy import symbols, Pow, Symbol, Integer

x, y = symbols('x y')
expr = x**2*y + sin(x*y)

def show_tree(e, depth=0):
    print("  " * depth + str(e.func.__name__) + f": {e}")
    for a in e.args:
        show_tree(a, depth + 1)

show_tree(expr)
# Mul: x**2*y*sin(x*y)
#   Pow: x**2
#     Symbol: x
#     Integer: 2
#   Symbol: y
#   sin: sin(x*y)
#     Mul: x*y
#       Symbol: x
#       Symbol: y

也可以用內建的 preorder_traversal（前序走訪：先訪節點、再訪子節點）：

from sympy import preorder_traversal

for node in preorder_traversal(expr):
    print(node.func.__name__, "->", node)

常見應用：找出運算式中所有的自由符號、或所有特定型態的子運算式：

syms = {a for a in preorder_traversal(expr) if isinstance(a, Symbol)}
print(syms)   # {x, y}

自訂函數：Function 類與數值近似#

SymPy 允許你定義自己的符號函數：繼承 Function，並在類別方法中定義行為。eval 定義「遇到特定輸入時自動化簡」的規則：

from sympy import Function, sin, cos, sqrt, symbols

class unit_step(Function):
    """階梯函數：負數為 0，非負數為 1"""
    @classmethod
    def eval(cls, x):
        if x.is_Number:
            if x < 0:
                return 0
            else:
                return 1

x = symbols('x')
print(unit_step(-3))        # 0
print(unit_step(2.5))       # 1
print(unit_step(x))         # unit_step(x)：符號輸入保持原樣

要讓自訂函數能被 evalf 求數值，定義 _eval_evalf 方法：

class gauss(Function):
    """高斯函數 exp(-x^2)"""
    @classmethod
    def eval(cls, x):
        if x.is_Number:
            return None   # 不自動展開，保持符號形式

    def _eval_evalf(self, prec):
        from sympy import exp, Float
        return exp(-self.args[0]**2).evalf(prec)

print(gauss(1))         # gauss(1)
print(gauss(1).evalf()) # 0.367879441171442
print(gauss(2).evalf(20))
# 0.01831563888873418029372

自訂函數可以參與微分、展開等符號運算；若沒有定義對應規則，SymPy 會保留 Derivative 形式不動。

替代與改寫：subs 與 rewrite#

subs 把符號或子運算式換掉（複習前面章節的用法）：

from sympy import sin, cos, exp

print((x**2 + 1).subs(x, 3))          # 10
print(sin(x)**2 + cos(x)**2)          # sin(x)**2 + cos(x)**2（不自動化簡）
print(sin(x).rewrite(cos))            # 轉成 cos 的形式：1 - 2*sin(x/2)**2...

rewrite(target) 是更強大的「換形式」工具，指定目標函數或模組：

print(sin(x).rewrite(exp))    # -I*(exp(2*I*x) - 1)/(2*exp(I*x))
print(cos(x).rewrite(sin))    # sin(x + pi/2)
print(exp(x).rewrite(sin))    # 以 sin/cos 表示
print((1/sin(x)).rewrite(cos)) # 用餘切表示

rewrite 常用於：把雙曲函數換成指數、把階乘換成伽瑪函數、把三角函數統一成一種再做化簡：

from sympy import tan, cosh, sinh, simplify

e = tanh(x)
print(e.rewrite(exp))            # (exp(2*x) - 1)/(exp(2*x) + 1)
print(simplify(sin(x).rewrite(cos) - cos(x - x)))   # 0

動手實驗：統計運算式中的運算子種類#

結合本章所學：走訪運算式樹，統計出現了哪些「運算子類別」（Add、Mul、Pow……）與各自次數：

from sympy import symbols, sin, cos, exp, preorder_traversal, Rational

x, y = symbols('x y')

def count_operators(expr):
    """統計運算式樹中每種 func（運算子）出現的次數"""
    counts = {}
    for node in preorder_traversal(expr):
        name = type(node).__name__
        counts[name] = counts.get(name, 0) + 1
    return counts

exprs = [
    x**2 + 2*x + 1,
    sin(x)*cos(x) + exp(-x),
    (x + y)**3 / (x - y) + Rational(1, 2),
]
for e in exprs:
    print(f"{e}  ->  {count_operators(e)}")

輸出範例：

x**2 + 2*x + 1  ->  {'Add': 1, 'Mul': 1, 'Pow': 1, 'Symbol': 1, 'Integer': 2}
sin(x)*cos(x) + exp(-x)  ->  {'Add': 1, 'Mul': 1, 'sin': 1, 'cos': 1,
                              'exp': 1, 'Symbol': 1, 'Integer': 1}

注意細節：Integer(2) 這類「原子」節點沒有子運算式（args 是空 tuple），所以遞迴會自然停在葉節點。若只想統計真正的「運算子」（有子節點的節點），加一個過濾：

def count_operators_only(expr):
    counts = {}
    for node in preorder_traversal(expr):
        if node.args:   # 有子節點才算運算子
            name = type(node).__name__
            counts[name] = counts.get(name, 0) + 1
    return counts

print(count_operators_only(sin(x)**2 + x*y))
# {'Add': 1, 'Mul': 1, 'Pow': 1, 'sin': 1}

小結與練習#

本節重點

SymPy 運算式是樹：func 是節點類別、args 是子節點，func(*args) 可重建
srepr() 顯示完整樹狀結構；preorder_traversal 前序走訪所有節點
自訂函數：繼承 Function，用 eval 定義化簡規則、_eval_evalf 定義數值近似
rewrite(target) 把運算式改寫成其他形式（exp、sin、cos…）

練習

對 \(e = \sin(x)\cos(y) + x^2 y\) 用 srepr 列印，並手畫出它的樹狀圖。
寫一個遞迴函數 depth(expr)，回傳運算式樹的深度（葉節點深度為 0）。
定義自訂函數 sgn2(x)（符號函數），讓 sgn2(-2) 回傳 \(-1\) 、sgn2(0) 回傳 \(0\) 。
實驗：修改 count_operators_only，把結果按出現次數由大到小排序輸出；再對 \(\sin^2 x + \cos^2 x\) 展開（expand）前後分別統計，比較節點總數的變化。