使用 python3 autograd 及 numpy 自定一個微分方程式

from autograd import elementwise_grad as grad
from autograd.extend import defvjp, primitive
import numpy as np

# 定義函式原型, 輸入資料型態必須是 np.array(.) 才能在 autograd 內運算
@primitive
def polynomial_f(x): # primitive 只能用 def, 不能用 lambda, 用來定義 forward function
    return x**3 # x³

# defvjp(polynomial_f, lambda o, x : (lambda g: 3*x**2))
#定義微分的 backward function, 需要將矩陣乘上微分方程
def grad_f(o, x): # 可以直接回傳 lambda, o 是 f(x) 的輸出, x 是 f(x) 輸入
    def derivative(_): # 綁定 x
        return 3*x**2
        # return np.full(x.shape, _) * 3*x**2 # 綁定 x, _
    return derivative
defvjp(polynomial_f, grad_f) # 將原型與微分方程綁定, 兜在一起

# 呼叫 autograd 的 grad function, 定義 n 階微分方程(order-n derivative): f'ⁿ(x) = dⁿf(x)/dx
df_n = lambda f, n = 1: df_n(grad(f), n - 1) if n > 1 else grad(f)
x = np.array([1, 2, 3, 4, 5])    # 測試 x 座標點, x 當成一組向量
print(f"f(x)   = x³  : {polynomial_f(x)},\tx = {x}")           # f(x) = x³
print(f"f'(x)  = 3x² : {df_n(polynomial_f)(x)   },\tx = {x}")  # f(x) 1 階微分f'(x)  = 3x²
print(f"f²'(x) = 6x  : {df_n(polynomial_f, 2)(x)},\tx = {x}")  # f(x) 2 階微分f"(x)  = 6x
print(f"f³'(x) = 6   : {df_n(polynomial_f, 3)(x)},\tx = {x}")  # f(x) 3 階微分f"'(x) = 6
print(f"f⁴'(x) = 0   : {df_n(polynomial_f, 4)(x)},\tx = {x}")  # f(x) 4 階微分f""(x) = 0