Эффективное стохастическое численное интегрирование по многим траекториям ⇐ Python

[Anonymous]

Я реализую численный метод решения стохастических дифференциальных уравнений с использованием метода Эйлера-Маруямы.
То, что у меня есть, работает, но неэффективно. Причина в том, что из-за стохастической природы проблемы у меня много траекторий. Прямо сейчас я решаю их одну за другой. У меня такое ощущение, что я смогу их распараллелить, поскольку они независимы.
Рабочий код выглядит так

Код: Выделить всё

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import time
from numba import jit, njit
import os

def A(u):
x=u[0]
y=u[1]
z=u[2]

omega=1/2*np.sqrt((1+8*kappa*z*z))

A=np.array([[-2,omega,0],
[-omega,0,0],
[0,0,-kappa]])

du=A.dot(u)
return du

def B(u,w):
x=u[0]
y=u[1]
z=u[2]

g=np.sqrt(kappa*nth)

B=np.array([[0],
[1],
[1]])*g

return np.reshape(B*w,len(u0))

def SDE(A,B):
u = np.zeros((len(u0),Nmax+1,Mmax),dtype=np.complex64)
for m in range(Mmax):
u[:,0,m]=u0
for n in range(0,Nmax):
u[:,n+1,m] = u[:,n,m]+dt*A(u[:,n,m])+B(u[:,n,m],w[n,m])*np.sqrt(dt)

return u

#Parameters
kappa=0.05
nth=1.
gamma=1

Mmax=100 #number of trajectories

Tmax=10. ##max value for time
dt=0.05
Nmax=int(Tmax/dt) ##number of steps

t_list=np.arange(0,Tmax+dt/2,dt)

w = np.random.randn(Nmax+1,Mmax)

u0 = np.array([1., 0., np.sqrt(nth)/2])

u_t=SDE(A,B)

u_mean=np.mean(u_t,axis=2)

Этот код представляет собой упрощение моего реального кода, в котором у меня гораздо большее измерение системы и гораздо больше траекторий.

Обратите внимание, что это неэффективно. , потому что по мере увеличения Mmax мне приходится их перебирать.
В идеале я бы хотел, чтобы мой решатель выглядел примерно так:

Код: Выделить всё

def SDE(A,B):
u = np.zeros((len(u0),Nmax+1,Mmax),dtype=np.complex64)
u[:,0,:]=u0
for n in range(0,Nmax):
u[:,n+1,:] = u[:,n,:]+dt*A(u[:,n,:])+B(u[:,n,:],w[n,:])*np.sqrt(dt)

return u

т. е. иметь возможность пренебречь циклом над m и просто делать это параллельно. Однако по наивности это не сработает.
Еще один идеальный способ сделать его более эффективным — использовать Numba. Однако после многих попыток мне не удалось реализовать njit с помощью решателя SDE, который я определил.

Подробнее здесь: https://stackoverflow.com/questions/787 ... ajectories