RK4 用 numba 加速

Answer 1

解决方案

在 mycode.py 中

import numpy as np
from scipy import integrate
from typing import Union, List
import numba

def AdjMtoAdjL(adjM: np.ndarray) -> list:
    return [np.argwhere(adjM[:,i] > 0).flatten() for i in range(len(adjM))]
def AdjMtoEdgL(adjM: np.ndarray) -> np.ndarray:
    return np.argwhere(adjM > 0)

@numba.jit(nopython=True)
# def swing(t, y, model_param, model):
def swing(t, y, phi, m, gamma, P, K, model):
    if model == "swing":
        T, O = y
        T = np.array([T])
        O = np.array([O])
    else:
        T = y

    # Get Interaction
    Interaction = K*np.sin(T-phi)
    """
    \dot{\theta} &= \omega \\
    \dot{\omega} &= \frac{1}{m}(P-\gamma\omega+\Sigma K\sin(\theta-\phi))
    """
    if model == "swing":
        dT = O
        dO = 1/m*(P - gamma*O - Interaction)
        dydt = np.concatenate((dT, dO))#, dtype=np.float64)
    else:
        dydt = P + Interaction
    return dydt

@numba.jit(nopython=True)
def RK4(func, t_end, X0, dt, phi, m, gamma, P, K, model):
    t = np.arange(0,t_end, dt, dtype=np.float64)
    X  = np.zeros((t.shape[0], X0.shape[0]))
    X[0] = X0
    hdt = dt*.5
    for i in range(t.shape[0]-1):
        t1 = t[i]
        x1 = X[i]
        k1 = func(t[i], X[i], phi, m, gamma, P, K, model)
        
        t2 = t[i] + hdt
        x2 = X[i] + hdt * k1
        k2 = func(t2, x2, phi, m, gamma, P, K, model)
        
        t3 = t[i] + hdt
        x3 = X[i] + hdt * k2
        k3 = func(t3, x3, phi, m, gamma, P, K, model)
        
        t4 = t[i] + dt
        x4 = X[i] + dt * k3
        k4 = func(t4, x4, phi, m, gamma, P, K, model)
        X[i+1] = X[i] + dt / 6. * (k1 + 2. * k2 + 2. * k3 + k4)
    return X

主代码.ipynb

import networkx as nx
import os
import multiprocessing as mp
from multiprocessing import Pool
import time
import numpy as np
import swing

def multiprocess(Ngrid=101, t_end=30., omega_lim=30, dt=.001, n_cpu=19):
    start = int(time.time())
    
    T_range = np.linspace(0, 2*np.pi, Ngrid)
    O_range = np.linspace(-omega_lim, omega_lim, Ngrid)
    
    paramss = []
    for theta in T_range:
        for omega in O_range:
            y0 = np.hstack((
                theta,  # Theta
                omega,  # Omega
            ))
            params = {}
            params['sparam'] = Swing_Parameters
            params['t_end'] = t_end
            params['init'] = y0
            params['dt'] = dt
            paramss.append(params)
            del([[params]])

    p = Pool(processes=n_cpu)
    result = p.map(solve_func, paramss)
    
    end = int(time.time())
    print("***run time(sec) : ", end-start)
    print("Number of Core : " + str(n_cpu))
    return result

def solve_func(params):
    Swing_Parameters = params['sparam']
    t_end = params['t_end']
    y0 = params['init']
    dt = params['dt']
    
    # model = swing.SwingSingle(**Swing_Parameters)
    t_eval = np.arange(0,t_end, dt)
    # solution = integrate.solve_ivp(model, [0,t_end], y0, dense_output=False, 
                       # t_eval=t_eval, vectorized=True, method="LSODA")
    phi = Swing_Parameters["phi"]
    m = Swing_Parameters["m"]
    gamma = Swing_Parameters["gamma"]
    P = Swing_Parameters["P"]
    K = Swing_Parameters["K"]
    _model = Swing_Parameters["model"]

    solution = swing.RK4(swing.swing, t_end, y0, dt, phi, m, gamma, P, K, _model)
    return solution
    
Ngrid = 301
t_end = 24.
omega_lim = 30
dt = .05

Ngrid = 301
t_end = 24.
omega_lim = 30
dt = .05

Swing_Parameters = {
    "phi": np.pi,
    "m": 1.,
    "gamma": 0.3,
    "P": 2.,
    "K": 8.,
    "model": "swing"
}

model = swing.SwingSingle(**Swing_Parameters)

res = multiprocess(Ngrid=Ngrid, t_end=t_end, omega_lim=omega_lim, dt=dt, n_cpu=19)