PythonでScientific Programmingプチテク集

Integrator (Propagator) 作成について

Pythonでintegratorを作成する場合、scipyのsolve_ivp()関数などが便利。その際は、ODEの関数定義はnumpy可してあげるようにする。例えば、円制限三体問題の場合、

\begin{align}
    \ddot{x} -2\dot{y} 
        &= x -\dfrac{1-\mu}{r_{1}^3} (x+\mu) + \dfrac{\mu}{r_{2}^3} (1-\mu-x)
    \\
    \ddot{y} +2\dot{x} 
        &= y -\dfrac{1-\mu}{r_{1}^3}y - \dfrac{\mu}{r_{2}^3}y
    \\
    \ddot{z} 
        &= -\dfrac{1-\mu}{r_{1}^3}z - \dfrac{\mu}{r_{2}^3}z
\end{align}

となるので、これを6つのFirst Order ODEとして表すと

\begin{align}
    \dfrac{d x}{d t} &= \dot{x}
\\
    \dfrac{d y}{d t} &= \dot{y}
\\
    \dfrac{d z}{d t} &= \dot{z}
\\
    \dfrac{d \dot{x}}{d t} 
        &= \dfrac{d^2 x}{d t^2} = 2\dot{y} +x -\dfrac{1-\mu}{r_{1}^3} (x+\mu) + \dfrac{\mu}{r_{2}^3} (1-\mu-x)
    \\
    \dfrac{d \dot{y}}{d t} 
        &= \dfrac{d^2 y}{d t^2} =-2\dot{x} + y -\dfrac{1-\mu}{r_{1}^3}y - \dfrac{\mu}{r_{2}^3}y
    \\
    \dfrac{d \dot{z}}{d t} 
        &= \dfrac{d^2 z}{d t^2} =-\dfrac{1-\mu}{r_{1}^3}z - \dfrac{\mu}{r_{2}^3}z
\end{align}

関数にすると

import numpy as np
from numba import jit

# deifne RHS function in CR3BP
@jit(nopython=True)
def rhs_cr3bp(t,state,mu):
    """Equation of motion in CR3BP, formulated for scipy.integrate.solve=ivp(), compatible with njit

    Args:
        t (float): time
        state (np.array): 1D array of Cartesian state, length 6
        mu (float): CR3BP parameter
    Returns:
        (np.array): 1D array of derivative of Cartesian state    
    """

    # unpack positions
    x  = state[0]
    y  = state[1]
    z  = state[2]
    # unpack velocities
    vx = state[3]
    vy = state[4]
    vz = state[5]
    # compute radii to each primary
    r1 = np.sqrt( (x+mu)**2 + y**2 + z**2 )
    r2 = np.sqrt( (x-1+mu)**2 + y**2 + z**2 )
    # setup vector for dX/dt
    deriv = np.zeros((6,))
    # position derivatives
    deriv[0] = vx
    deriv[1] = vy
    deriv[2] = vz
    # velocity derivatives
    deriv[3] = 2*vy + x - ((1-mu)/r1**3)*(mu+x) + (mu/r2**3)*(1-mu-x)
    deriv[4] = -2*vx + y - ((1-mu)/r1**3)*y - (mu/r2**3)*y
    deriv[5] = -((1-mu)/r1**3)*z - (mu/r2**3)*z

    return deriv

となる。前期の通り、大概の場合ODEの表現はnumba化が比較的容易に可能と思われるのでおススメ。あとは、solve_ivp()関数を呼べばいいので

from scipy.integrate import solve_ivp

# integrate until final time tf
tf = 3.05
# initial state
state0 = np.array([ 9.83408400e-01, -9.42453366e-04, 1.27227988e-03, 
                    7.03724138e-01, -1.78296421e+00, 1.13566847e+00])
# cr3bp mass parameter mu
mu = 0.012150585609624
# call solve_ivp
sol = solve_ivp(fun=rhs_cr3bp, t_span=(0,tf), y0=state0, args=(mu,))

ここで、argsにはODEの状態行列 (state) と時間 (t) 以外に必要になるパラメータ（上記の例ではmuを入れる。tupleなので、たとえエントリーが一つでもカンマ（,）を忘れずに。