Оптимальный метод численного расчета 4D-интеграла на большой сетке ⇐ Python

[Anonymous]

Описание проблемы
Я пытаюсь численно вычислить 4D-интеграл:


где d^4 x = dt d^3 x и функция f(x)=f(t,x) следует разрешать в большой сетке. Чтобы уместить его в памяти (у меня ограничение в 1Тб ОЗУ), я дискретизирую интеграл по временной и пространственной сетке и разбиваю его на две части:


Затем я вычисляю сумму времени в цикле for, где для каждого времени шаг, я пересчитываю f(x)=f(t_n,x) и выполните 3D БПФ.
Заметки о производительностиВычисление f(x)=f(t,x) включает в себя несколько операций с массивами, для которых я использую numexpr и вычисляю БПФ с библиотекой pyfftw. Я не могу векторизовать во времени, потому что массив для БПФ тогда не поместится в памяти.
Я уже пробовал использовать более низкую точность (float32, complex64), чтобы ограничить потребление памяти и сократить время вычислений.
Для каждого временного шага выполняются следующие вычисления:

• Вычисление f(x)=f(t,x) требует некоторых операций с массивом (с numexpr), а затем обратного 3D-БПФ.< /li>
  Затем в теле Интеграл Я выполняю некоторые операции с массивом с помощью f(x)=f(t,x) и обрабатываю 3D-БПФ результата.

Я профилировал код с помощью Scalene и Memray, которые показывают, что большая часть времени тратится на БПФ и эти конкретные операции с массивами.
Вопрос
Буду очень признателен за любую подсказку о том, как я могу улучшить производительность этого расчета. Если бы вы могли указать на некоторые идеи/ресурсы для чтения, это тоже было бы здорово.

Минимальный рабочий пример показан ниже:import numpy as np
import numexpr as ne
import pyfftw

# Minimal working example
class Function:
def __init__(self, size=(20,20,1000), nthreads=1):
self.a = np.random.random(size)

self.tmp = np.empty(size, dtype='complex128')
self.tmp_fftw = pyfftw.FFTW(
self.tmp,
self.tmp,
axes=(0, 1, 2),
direction="FFTW_FORWARD",
flags=("FFTW_MEASURE",),
threads=nthreads,
)

def get_function(self, t):
ne.evaluate("exp(1j*t)*a", global_dict={"a": self.a}, out=self.tmp)
self.tmp_fftw.execute()
return self.tmp

class Integral:
def __init__(self, func, size=(20,20,1000), nthreads=1):
self.func = func

self.f = np.empty(size, dtype='complex128')
self.f_fftw = pyfftw.FFTW(
self.f,
self.f,
axes=(0, 1, 2),
direction="FFTW_FORWARD",
flags=("FFTW_MEASURE",),
threads=nthreads,
)

self.result = 0

def get_one_time_step(self, t):
self.f = self.func.get_function(t)
ne.evaluate("3*f + 2", global_dict={"f": self.f}, out=self.f)
self.f_fftw.execute()
self.result += self.f

def get_time_integral(self, t_grid):
for t in t_grid:
self.get_one_time_step(t)
return self.result

size = (30,30,1000)
func = Function(size)
integrator = Integral(func, size)

t_grid = np.linspace(-1, 1, 21)
result = integrator.get_time_integral(t_grid)


Подробнее здесь: https://stackoverflow.com/questions/793 ... large-grid