在Python脚本中定义全局函数

Question

我是Python的新手.我正在编写一个脚本,它将使用Runge-Kutta方法在数值上集成一组常微分方程.由于Runge-Kutta方法是一种有用的数学算法,我把它放在自己的.py文件中,rk4.py.

def rk4(x,dt):
    k1=diff(x)*dt
    k2=diff(x+k1/2)*dt
    k3=diff(x+k2/2)*dt
    k4=diff(x+k3)*dt
    return x+(k1+2*k2+2*k3+k4)/6

该方法需要知道用户正在使用的方程组以执行算法,因此它调用一个函数diff(x),该函数将为rk4提供它需要工作的导数.由于方程式会随着使用而改变,我希望在运行特定问题的脚本中定义diff().在这种情况下,问题是汞的轨道,所以我写道mercury.py.(这不是它最终会看起来如何,但为了弄清楚我在做什么,我已经简化了它.)

from rk4 import rk4
import numpy as np

def diff(x):
    return x

def mercury(u0,phi0,dphi):
    x=np.array([u0,phi0])
    dt=2
    x=rk4(x,dt)
    return x

mercury(1,1,2)

当我运行mercury.py时,我收到一个错误:

  File "PATH/mercury.py", line 10, in mercury
    x=rk4(x,dt)
  File "PATH/rk4.py", line 2, in rk4
    k1=diff(x)*dt
NameError: global name 'diff' is not defined

我认为它diff()不是一个全局函数,当rk4运行时它对diff一无所知.显然rk4是一小段代码,我可以把它推到我当时正在使用的任何脚本中,但我认为Runge-Kutta积分器是一个基本的数学工具,就像在NumPy中定义的数组一样,所以它有意义的是使它成为一个被调用的函数,而不是在每个使用它的脚本中定义的函数(可能很多).但我也不能告诉rk4.py从特定的.py文件导入特定的差异,因为这首先破坏了我想要的rk4的普遍性.

有没有办法在像mercury.py这样的脚本中全局定义diff,这样当调用rk4时,它会知道diff？

Answer 1

接受函数作为参数:

def rk4(diff,  # accept an argument of the function to call
        x, dt)
    k1=diff(x)*dt
    k2=diff(x+k1/2)*dt
    k3=diff(x+k2/2)*dt
    k4=diff(x+k3)*dt
    return x+(k1+2*k2+2*k3+k4)/6

然后,当你调用时rk4,只需传入要执行的函数:

from rk4 import rk4
import numpy as np

def diff(x):
    return x

def mercury(u0,phi0,dphi):
    x=np.array([u0,phi0])
    dt=2
    x=rk4(diff,  # here we send the function to rk4
          x, dt)
    return x
mercury(1,1,2)

mercury接受diff作为参数也许是一个好主意,而不是从闭包(周围的代码)中获取它.然后你必须像往常一样传递它 - 你mercury在最后一行的调用会被读取mercury(diff, 1, 1, 2).

函数是Python中的"一等公民"(几乎所有东西,包括类和模块),在某种意义上它们可以用作参数,保存在列表中,分配给名称空间中的名称等等.