Определение глобальной функции в Python script

Я новичок в Python. Я пишу script, который будет численно интегрировать набор обыкновенных дифференциальных уравнений с использованием метода Рунге-Кутты. Поскольку метод Рунге-Кутты является полезным математическим алгоритмом, я поставил его в свой собственный .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() определялся в script, который будет запускать конкретную проблему. В этом случае проблема - это орбита ртути, поэтому я написал 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 - небольшая часть кода, и я мог бы просто вставить его в любой script, который я использую в то время, но я думаю, что интегратор Runge-Kutta - это фундаментальный математический инструмент, как и массив, определенный в NumPy, и поэтому имеет смысл сделать его функцией, которая называется скорее той, которая определена в каждом script, который ее использует (чего может быть много). Но я также не могу сказать rk4.py, чтобы импортировать конкретный diff из определенного .py файла, потому что это разрушает общность rk4, которую я хочу в первую очередь.

Есть ли способ определить diff в глобальном масштабе в script, таком как mercury.py, так что, когда вызывается rk4, он будет знать о diff?

Ответ 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 (как и все, включая классы и модули), в том смысле, что их можно использовать в качестве аргументов, вести в списках, назначаться именам в пространствах имен и т.д. и т.д..

Ответ 2

diff уже является глобальным в модуле mercury.py. Но чтобы использовать его в rk4.py, вам нужно будет импортировать его следующим образом:

from mercury import diff

Это прямой ответ на ваш вопрос.

Однако передача функции diff в rk4, как предложено @poorsod, намного более элегантна, а также позволяет избежать круговой зависимости между mercury.py и rk4.py, поэтому я предлагаю вам сделать это таким образом.