Эффективное стохастическое численное интегрирование по многим траекториям
Я реализую численный метод решения стохастических дифференциальных уравнений с использованием метода Эйлера-Маруямы.
То, что у меня есть, работает, но неэффективно. Причина в том, что из-за стохастической природы проблемы у меня много траекторий. Прямо сейчас я решаю их одну за другой. У меня такое чувство, что я должен иметь возможность их распараллелить, поскольку они независимы.
Рабочий код выглядит так
%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, который я определил.
@jared О, ты прав: я получаю сообщение об ошибке. Когда я пытаюсь использовать u[:,0,:]=u0, я получаю сообщение «не удалось передать входной массив из формы (3,) в форму (3,50)». Это имеет смысл, поскольку LHS на самом деле не знает, как разместить u0. Затем я предполагаю, что другая часть кода, цикл, также выдаст мне аналогичные ошибки из-за матричных произведений, которые я определяю в A и B.
А что, если бы ты это сделал u[:,0,:] = u0[:,None]? Это должно транслироваться.
@jared, это сработало, я не знал, что ты можешь это сделать. Однако теперь ошибка распространяется на то, как я определил A и B. Например, моя следующая ошибка теперь связана с: np.reshape(B*w,len(u0)); return _wrapfunc(a, 'reshape', newshape, order=order); невозможно преобразовать массив размером 150 в форму (3,)
При векторизации кода вам необходимо учитывать размерность (форму) всех терминов и убедиться, что вы имеете дело со всеми из них. Вы просите нас решить все проблемы с формой за вас?
Я рекомендую начать с нового файла (где вы напишете векторизованную версию) и запустить код в отладчике, чтобы пройти построчно и проверить форму каждого термина, выяснить, какой формы он должен быть, и правильно изменить форму или добавить размеры. (как я уже сказал с помощью u0), пока ваш код не станет согласованным.
Вы используете np.complex64, но это выглядит подозрительно неправильно. Знаете ли вы, что np.complex64 на самом деле представляет собой комплексные числа простой точности, а не двойной точности? Это предназначено? Если да (что необычно для такого кода), почему бы вам не использовать повсюду простые точности?
@JérômeRichard Я не знал об этом... Я думал, что это комплексный эквивалент np.float64 с точки зрения точности. Каким же тогда должен быть правильный тип данных для комплексных чисел двойной точности?
np.complex128 подходит для двойной точности. Я должен сказать, что соглашение об именах не очень интуитивно понятно...
Ответы 1
Вот проблемы с производительностью, которые я обнаружил в предоставленном коде:
A.dot(u)здесь явно излишен, поскольку матрицаAпредставляет собой матрицу 3x3 только с 4 ненулевыми значениями. Лучше вычислить выражение вручную, особенно с помощью Numba.- создание множества небольших массивов, таких как 2D-массивы 3x3 или даже 1D-массивы всего с 3 элементами, очень дорого, даже в Numba (поскольку массивы распределяются динамически и подсчитываются ссылки). Вам следует предварительно выделить массив или даже просто использовать 3 переменные (гораздо более эффективно). В результате в Numba возврат новых массивов часто менее эффективен, чем запись в существующий предварительно выделенный массив, переданный в параметре (по вышеуказанной причине).
- выделения обычно не масштабируются, поэтому всегда следует уменьшать количество выделений перед выполнением вычислений с несколькими потоками.
- итерация по последней оси обычно выполняется быстрее (в смежных массивах), поскольку элементы последней оси хранятся в памяти непрерывно, в то время как элементы другой оси не являются таковыми. В результате замена первой и последней осей при вычислении
uдолжна происходить быстрее. Его форма будет(Mmax,Nmax+1,len(u0)). Для получения дополнительной информации прочтите: AoS против SoA. - Вы можете предварительно вычислить некоторые константы, особенно когда значения умножаются на 0.
Кроме того, np.reshape(B*w,len(u0)) сбивает с толку, поскольку B имеет размер 3, поэтому u0 должно быть также размера 3, а reshape кажется бесполезным. Обратите внимание, что np.complex64 означает простую точность (как указано в комментариях).
Вот полученный код Numba:
import numba as nb # new
@nb.njit('(complex128[:], complex128[::1])')
def A(u, res):
x=u[0]
y=u[1]
z=u[2]
omega = 0.5 * np.sqrt((1+8*kappa*z*z))
res[0] = -2 * x + omega * y
res[1] = -omega * x
res[2] = -kappa * z
return res
@nb.njit('(complex128[:], float64, complex128[::1])')
def B(u, w, res):
g = np.sqrt(kappa*nth)
res[0] = 0
res[1] = g * w
res[2] = g * w
return res
# No signature is provided so the first call will be much slower
# But providing a signature here is complicated since A and B are functions
@nb.njit
def SDE(A,B):
u = np.zeros((len(u0),Nmax+1,Mmax), dtype=np.complex128)
sqrt_dt = np.sqrt(dt)
for m in range(Mmax):
u[:,0,m] = u0
tmp1 = np.empty(3, dtype=np.complex128)
tmp2 = np.empty(3, dtype=np.complex128)
for n in range(0,Nmax):
A(u[:,n,m],tmp1)
B(u[:,n,m],w[n,m],tmp2)
for i in range(3):
u[i,n+1,m] = u[i,n,m] + dt * tmp1[i] + tmp2[i] * sqrt_dt
return u
На моей машине (с процессором i5-9600KF) это примерно в 500 раз быстрее.
Я думаю, что после оптимизации кода нет необходимости использовать несколько потоков, поскольку вычисления, наконец, стали довольно быстрыми. Если этого недостаточно, вы можете добавить флаг parallel=True и заменить for m in range(Mmax) на for m in nb.prange(Mmax). Однако это не будет хорошо масштабироваться из-за ложного совместного использования, вызванного плохой структурой памяти. Как указано в приведенном выше списке, вам следует поменять местами оси 1 и 3, чтобы решить эту проблему.
В конечном итоге окончательный код после распараллеливания и с лучшим расположением памяти должен выглядеть так:
# A and B are the same as before
# No signature is provided so the first call will be much slower
# But providing a signature here is complicated since A and B are functions
@nb.njit(parallel=True)
def SDE(A,B):
u = np.zeros((Mmax,Nmax+1,len(u0)), dtype=np.complex128)
sqrt_dt = np.sqrt(dt)
for m in nb.prange(Mmax):
u[m,0,:] = u0
tmp1 = np.empty(3, dtype=np.complex128)
tmp2 = np.empty(3, dtype=np.complex128)
for n in range(0,Nmax):
A(u[m,n,:],tmp1)
B(u[m,n,:],w[m,n],tmp2)
for i in range(3):
u[m,n+1,i] = u[m,n,i] + dt * tmp1[i] + tmp2[i] * sqrt_dt
return u
# [...] same code
w = np.random.randn(Nmax+1,Mmax).T.copy()
u0 = np.array([1., 0., np.sqrt(nth)/2])
u_t=SDE(A,B)
u_mean=np.mean(u_t,axis=0)
Этот код примерно в 2000 раз быстрее.
Примечание. w — это глобальная переменная, поэтому Numba считает ее константой времени компиляции (т. е. она никогда не должна меняться в течение срока службы приложения). Если это не так, вам следует передать ее в качестве параметра.
Кстати, обратите внимание, что Julia может быть лучшим языком для таких вычислений, поскольку стандартная реализация — это JIT-компилятор, и мы можем легко избежать создания новых временных массивов (хотя создание нового массива все еще довольно дорого даже в Julia).
какой отличный ответ, и я пойду подробно все проверять. Однако отправная точка для меня несколько проблематична: приведенный мной пример представлял собой игрушечную модель для моего реального кода, в которой матрицы A и B намного больше (10x10) и (10x1), поэтому писать его вручную не всегда является лучшим решением. Наиболее эффективным. Что мне нужно изменить в вашем ответе, чтобы сохранить A и B в качестве матриц, которые я определил?
Если A на практике является матрицей 10x10, то вам следует предварительно выделить ее один раз (как я сделал с tmp1 и tmp2), а затем написать цикл для самостоятельного умножения матриц (поскольку я не думаю, что np.dot будет эффективным, а Numba кажется пока не поддерживать np.dot для комплексных чисел). Если ваша реальная матрица имеет >60% нулей, все равно рекомендуется выполнить операцию вручную, хотя это обременительно, если вы действительно заботитесь о производительности.
В B вы можете просто заменить ручные вычисления циклом. Это будет быстро, если 1. u0.size известно во время компиляции и 2. не создается временный массив. Если (1) неверно, то скорость будет медленнее, но не намного больше. Если (2) неверно, то накладные расходы на распределение наверняка будут огромными. Обратите внимание, что такие детали имеют большое значение, поскольку узкое место может отличаться от u0.size==10.
Спасибо за подробный ответ. Я попробовал то, что вы сказали о матрице А, и в некоторой степени это работает. Под этим я имею в виду, что Numba не позволяет мне «вызывать» переменные, то есть z=u[2]. Итак, я могу создать матрицу A и выполнить умножение матрицы с помощью Numba, пока z=cte, но если z=u[2], я получаю некоторую ошибку, которую невозможно описать. (Ошибки Numba обычно для меня совершенно непонятны). Тем не менее, я попробую то, что вы предложили, — выполнить операцию с матрицей вручную, поскольку меня волнует производительность. Спасибо!
Да, ошибки Numba не очень удобны для пользователя. Мне сложно понять, что происходит с предоставленной информацией. Советую вам проверить, чтобы тип константы и u[2] был одинаковым. Вы можете вручную указать значения, например, с помощью z = np.complex128(value).
Я только что адаптировал код для своего реального случая, и у меня возникла неожиданная проблема: ядро умирает в моем файле блокнота Jupyter. Размер даже не такой уж и большой: Nmax:200, Mmax=2 и len(u0)=10. Раньше я моделировал гораздо более крупные системы с помощью Numba и не сталкивался с этой проблемой.
Это может произойти в двух случаях: доступ к памяти за пределами границ и состояние гонки. Наиболее распространенным является первый. Вы можете добавить флаг boundscheck=True и отключить параллелизм с помощью parallel=False. Numba сообщит вам, если произошел неверный доступ к памяти (но за счет более медленного выполнения). Это обычная проблема с кодами нижнего уровня.
На самом деле я узнал, что проблема была в w = np.random.randn(Nmax+1,Mmax).T.copy(). Если я удалю .copy(), ядро не умрет, и я действительно смогу увеличить масштаб проблемы. Совершенно неожиданно, и я не знаю, почему
Для меня это не имеет смысла, если вы не используете w[n,m] вместо w[m,n] (что в любом случае должно быть обнаружено Numba с помощью boundscheck=True). Также убедитесь, что функция перекомпилирована, поскольку предполагается, что w не изменяется, поскольку она является глобальной. (Конечно, безопаснее передать его в параметре, ИМХО).
И последнее: меня немного смущают временные массивы tmp. В вашем коде вы вызываете это: A(u[:,n,m],tmp1) и B(u[:,n,m],w[n,m],tmp2), не присваивая его какой-либо переменной. Затем вы вызываете u[i,n+1,m] = u[i,n,m] + dt * tmp1[i] + tmp2[i] * sqrt_dt. Правильно ли сказать, что значения A(u[:,n,m],tmp1) были неявно сохранены в tmp1[i]?
Да, функции A и B сохраняют свои элементы в массиве, указанном в последнем аргументе (tmp1 и tmp2 здесь). Это позволит избежать создания дорогостоящего нового временного массива. В Python аргументы не могут быть изменены (поэтому ссылка на массив является неизменяемой), но элементы массива, передаваемые в параметре, могут быть изменены (поэтому элементы массива изменяемы - и неявной копии массива нет). Эквивалентный пример функции, выполняющей это в Numpy: np.multiply(arr1, arr2, out=resArr), которая сохраняет результат в resArr без создания временного массива (out= не является обязательным).
Давайте продолжим обсуждение в чате.
«Однако наивно это не сработает». А что, не работает? Вы получаете ошибку?