Нейронный Фурье-оператор для решения уравнений Навье-Стокса

1. О проекте

Предлагается ML-модель, которая смогла бы прогнозировать приближенное поведение вязкой несжимаемой жидкости (газа) в форме завихрённости на единичный тор, описываемого дифференциальными уравнениями Навье-Стокса для случая движения жидкости (газа) на плоскости.

В качестве уравнений, описывающих движение жидкости, на некотором интервале времени используются уравнения следующего вида:

Где:

• $u \in C\left( \left\lbrack 0,\ T \right\rbrack \right)$ - поле скоростей,
• $w = \nabla u$ - завирхенность,
• $w_{0}\ \in L^{2}(\left( 0,\ 1 \right)^{2})$ - начальная завихренность,
• ν - коэффициент кинематической вязкости,
• $f\ \in L^{2}(\left( 0,\ 1 \right)^{2})$ - вектор воздействующей силы,
• $t$ - время,
• $[0, T]$ -- интервал времени, на котором рассматривается решение уравнения.

---

2. Набор данных

Для обучения построенной модели задана обучающая выборка объемом 1000 примеров решения предложенного уравнения на интервале времени в 20 секунд и размером расчетной сетки 64х64 ячейки. Необходимо построить генератор предобработки и подачи данных в модель, реализовать процесс ее обучения и вывода результатов моделирования.

Исходные данные, используемые в экспериментах можно найти здесь.

---

3. Описание работы

3.1. Проблематика

«Центральный институт авиационного моторостроения имени П.И. Баранова» (ЦИАМ им. Баранова) - головная научная организация российского авиадвигателестроения. Это единственная в стране научная организация, осуществляющая полный цикл исследований, необходимых при создании авиационных двигателей и газотурбинных установок на их основе, а также научно-техническое сопровождение изделий, находящихся в эксплуатации.

Одно из главных направлений деятельности организации является разработка программного обеспечения для авиационных силовых установок, а также разработка компьютерных моделей, описывающих процессы, происходящие внутри конструкции двигателя при его работе.

Отдел «Суперкомпьютерного моделирования» ЦИАМ предназначен для решения задач моделирования газодинамических процессов, происходящих в авиационных установках в интересах различных государственных заказчиков.

Одной из центральных задач, которая решается в отделе является задача расчета поведения газа в воздухозаборных устройствах в зависимости от их геометрических и физических показателей. Решение такой задачи основано на численном решении дифференциальных уравнений движения сплошной среды (жидкости, газа), учитывающее ее кинетическую вязкость и теплопроводность. Такие уравнения именуются уравнениями Навье-Стокса и по сей день имеют лишь численное решение для индивидуальных задач.

Уравнения Навье-Стокса являются общим случаем законов сохранения массы, импульса и энергии для вязкой жидкости (газа) и максимально точно описывают ее поведение.

Основные численные методы, которые используются для решения уравнений Навье-Стокса в частном виде:

• метод конечных разностей
• метод конечных объемов.

Оба этих метода предполагают дискретизацию пространства на заданное количество ячеек (сетку) и поиск сетчатых функций, описывающих давление, температуру и скорость среды в каждой ячейке пространства.

Однако у данных методов есть существенный недостаток - это необходимость поиска компромисса: с уменьшением размера сетки, разделяющей пространство, экспоненциально растет вычислительная сложность алгоритма расчета.

На сегодняшний день одной из самых актуальных задач отдела является поиск возможных методов оптимизации такого решения.

3.2. Нейросетевые подходы к решению уравнений в частных производных

Одним из самых современных и передовых подходов в решении задач оптимизации работы вычислительных алгоритмов является использование машинного обучения.

Идея такого подхода заключается в следующем: если у нас есть сложное дифференциальное уравнение (уравнение в частных производных) и есть случаи его решения, т.е. есть набор пар: начальное состояние объекта и его состояния через некоторый промежуток времени, то можно построить модель, которая бы аппроксимировал зависимость между начальным и конечным состоянием.

3.3. Нейронный Фурье-оператор

В 2020 году исследователями из Калифорнийского Технологического Института была выпущена статья, в которой они описывают новую методику обучения нейронных сетей для решения уравнений в частных производных, зависящих от времени. Согласно статье, предложенный метод на порядок точнее, чем традиционные методы и позволяет решать уравнения Навье-Стокса для любой жидкости или газа. При этом данный подход позволяет избавиться от зависимости от суперкомпьютеров и получать вполне неплохие результаты.

Подход основан на расчете весов нейронной сети в пространстве Фурье.

Как утверждает одна из авторов статьи - Анима Анандкумар, профессор Калифорнийского университета, это связано с тем, что движение воздуха на самом деле можно описать комбинацией волн различной частоты, а Фурье-преобразование отлично способно работать с такими частотами. Свою разработку ученые назвали Нейронным Фурье-оператором (Fourier Neural Operator).

Задачей данного проекта, является реализация и адаптация методики калифорнийских ученых для решения уравнений Навье-Стокса в случае несжимаемого вязкого газа. Потенциально данное решение можно использовать в качестве начального приближения уравнений Навье-Стокса, чтобы повысить скорость сходимости метода конечных объемов и получить высокую точность решения в расчетных задачах для авиационных силовых установок.

3.4. Математическая формализация

Пусть заданы множества состояний: $А$ - множество начальных состояний и $U$ - множество конечных состояний. Под состоянием мы будем понимать, например, поле скоростей или поле давлений, описывающие параметры вязкой несжимаемой среды в каждый момент времени в некоторой ограниченной области $D$. Таким образом, множество состояний А описывается как $A = A(D, \mathbb{R}^{d_{a}}) = (a_{j}^{d_{a}})$, а множество состояний U описывается как $U = U (D, \mathbb{R}^{d_{u}}) = (u_{j}^{d_{u}})$.

Пусть есть некоторое отображение $G$, такое, что $u_{j} = G( a_{j} )$.

Необходимо найти отображение $G^{*}: A \times \theta \rightarrow U$, где $\Theta$ - множество параметров (весов) или это можно записать как $G_{\theta}^{*}:A \rightarrow U,\ где\ \theta \in \theta$. При этом параметрическое отображение $G_{\theta}$ должно максимально полно аппроксимировать отображение $G$ т.е. $G \approx G_{\theta}^{*}$.

Зададимся некоторым функционалом ошибки $\mathcal{L}$, таким, что $\mathcal{L}:U \times U \rightarrow \mathbb{R}$ и поставим следующую задачу оптимизации:

$$\operatorname{}{\mathbb{E}\lbrack\mathcal{L}\left( G\left( a \right),\ G^{*}(a,\theta\ \right)\rbrack}$$

В общем случае решения дифференциальных уравнений в частных производных состояния $a_{j}$ и $u_{j}$ на самом деле являются векторными функциями. Мы будем работать с дискретными значениями этих функций. Разделим область $D$ на подобласти (ячейки) $D_{j} = { x_{1},\ x_{2},\ \ldots,\ x_{n}}$, где $n$ - частота дискретизации области $D$. И в каждой ячейке будем иметь значение состояний $a_{j}$ и $u_{j}$. Таким образом, мы имеем набор состояний $a\left( x \right)$ и $u\left( x \right)$, которые описывают поведение некоторой вязкой несжимаемой сплошной среды в течение заданного интервала времени T в каждой точке $x \in D$.

Результирующий набор обучающих пар представляет собой тензоры $\hat{A}$ - входной тензор и $\hat{U}$ -- тензор верных ответов, размерности которых $(S,\ S,\ T)\ $и $(\ S,\ S,T_{\text{out}})$ соответственно. Где $S$ - размерность сетки $T_{out}$ - размер интервала времени, на который мы хотим сделать прогноз поведения среды.

3.4.1. Архитектура сети

На входной слой нейронной сети поступает тензор $\hat{A}$ размерностью $(\ S,\ S,\ T)$, соответствующий функции $a\left( x \right)$. Сперва отобразим этот тензор в новое более высоко размерное пространство $\mathbb{R}^{v}$, где $d_{v}$ > T, с помощью линейного преобразования (полносвязного слоя) $P:\ \mathbb{R}^{u} \rightarrow \mathbb{R}^{v}$. Получим:

$$v_{0}\left( x \right) = P(a\left( x \right))$$

Сверточная часть нейронной сеть будет представлять собой итеративную процедуру применения следующей операции:

$$v_{l + 1}(x) {\sigma_{l}(Wv}{l}(x) + K\left( \varphi \right) \bullet v{l}(x))$$

Где $v_{l}$ - выход $l\ $-ого слоя сети, $W$ - тензор линейного преобразования (одномерная свертка Conv1D), $K\left( \varphi \right)$ -некоторое интегральное преобразование по области D с внутренними параметрами $\varphi$, $\sigma$ -- функция активации$\ \ l$-ого слоя.

$$u^{*}\left( x \right) = Q(v_{L}\left( x \right))$$

Остается вопрос, что взять в качестве интегрального преобразования $K\left( \varphi \right)$. Исследователи из калифорнийского университета предлагают использовать в качестве такого преобразования Фурье-преобразование:

$$K\left( \varphi \right)v_{l}(x) = F^{- 1}(R_{\varphi} \bullet \left( \text{Fv}\left( x \right) \right)\ \forall x\ \epsilon\ D$$

Где:

$(F^{- 1}f)(\omega) = \int_{D}^{}{f\left( x \right)e^{- 2\text{πxω}}\text{dx}}$- прямое преобразование Фурье

$(F^{- 1}f)(\omega) = \int_{D}^{}{f\left( x \right)e^{2\text{πxω}}\text{dω}}$ -- обратное преобразование Фурье

$R_{\varphi}$ -- обучаемое линейное преобразование в пространстве Фурье

$f:D \rightarrow \mathbb{R}^{v}$ -- некоторая вектор-функция

Главная особенность преобразования Фурье в том, что, представив предположительно периодическую функцию в пространстве Фурье в виде гармоник с различными частотами, мы можем выделить наиболее значимые моды, убрав ненужный шум в данных. Мы зададимся порогом моды $k_{\max}$ будем выделять наиболее высокие частоты из преобразованной функции и совершать над ними преобразование $R_{\varphi}$, остальные частоты при прохождении Фурье-слоя обращаются в 0.

Еще одной важной особенностью такого подхода является то, что, если дискретизация расчетной сетки равномерна $s_{1} \times \ s_{2}\ldots \times s_{d} = n\ $, то мы можем заменить интегральное преобразование Фурье на быстрое дискретное преобразование Фурье, практически не потеряв в точности разложения функции, но значительно увеличив производительность:

$$Ff_{l}(\omega) = \sum_{x_{1} = 0}^{s_{1} - 1}...{\sum_{x_{d} = 0}^{s_{d} - 1}{f_{l}(x_{1},..., x_{d})e^{- 2i\pi \sum_{j = 1}^{d}\frac{x_{j}\omega_{j}}{s_{j}}}}}$$

$$F^{- 1}f_{l}(x) = \sum_{\omega_{1} = 0}^{s_{1} - 1}{...\sum_{\omega_{d} = 0}^{s_{d} - 1}{f_{l}\left( \omega_{1},\ ...,\ \omega_{d} \right)e^{- 2i\pi\sum_{j = 1}^{d}\frac{x_{j}\omega_{j}}{s_{j}}}}}$$

Тогда линейное преобразование в пространстве Фурье будет иметь вид:

$$(R(Fv_{l})){ij} = \sum_k^{d{v}} R_{ijk} (Fv_{l}){ik}, k = 1,..., k{\max}, j = 1,..,d_{v}$$

В соответствии с приведенной математической формализацией результирующая архитектура нейронной сети имеет следующий вид (рисунок 1):

Рисунок 1. Архитектура модели. а) Полная схема архитектуры нейронного оператора б) Схема архитектуры сверточного Фурье-слоя

3.4.2. Алгоритм обучения с прогнозированием на 1 шаг

Пусть есть обучающая выборка объемом M пар тензоров: $\hat{A_{v}}$ - тензор размерностью $(S,\ S,\ T_{1})$, который описывает дискретные состояния газа на сетке размером (S, S) с коэффициентом вязкости $v$ на промежутке времени $(0, T)$, а также $\hat{U_{v}}$ - тензор размерностью $(S, S, T_{2})$, который описывает дискретные состояния газа на сетке размером $(S, S)$ с тем же коэффициентом вязкости $v$ на промежутке времени $(T, T + \delta T$).

Алгоритм обучения с предсказанием на один шаг $(\delta T = 1)$ строится следующим образом:

1. Для каждой пары тензоров $(\hat{A_{v}}, \hat{U_{v}})_{i}, i = 1,..., M$:

2. Проинициализировать общую ошибку $L$

3. Для каждого момента времени $t = 0,...,T$:

   • Пропустить тензор $\hat{A_{v}}$ через описанную нейронную сеть, получив на выходе тензор ${\hat{U_{v}}}_{t}^{*}\text{\ \ }$ размерностью $(S,\ S,\ 1)$ - предсказанное состояние сплошной среды на 1 шаг.
   • Вычислить функционал ошибки С между предсказанным состоянием ${\hat{U_{v}}}{t}^{*}\text{\ \ }$ и истинным ${\hat{U{v}}}{t}\text{\ \ }$: $l = \ \mathcal{L}\left( {\hat{U{v}}}{t}^{*},{\hat{U{v}}}_{t}\ \right)$
   • Прибавить значение функционала ошибки к общей ошибки $L\ : = \ L + l$
   • Обновить значение тензора $\hat{A_{v}}:= \text{concat}({\hat{A_{v}}}{t + 1},\ {\hat{U{v}}}_{t}^{*})$ - смещение на один шаг и замена последнего значения на предсказанное

4. Выполнить алгоритм обратного распространения ошибки нейронной сети

---

4. Реализация алгоритма

Реализация алгоритма обучения и тестирования нейронной сети для решения уравнений Навье-Стокса выполнена на языке программирования высокого уровня Python. В качестве фреймворка для глубокого обучения был выбран Tensorflow, который позволяет интерактивно следить за обучением и управлять им.

4.1. Зависимости

• numpy (>= 1.19.0)
• tensorflow (>= 2.8.0)
• scipy (>= 1.8.0)
• matplotlib (>=3.4.3)
• celluloid (>=0.2.0)
• h5py
• dotmap (>=1.3.30)

4.2. Описание реализации

Реализованный нейронный Фурье-оператор представляет собой последовательность 4ех Фурье-слоев с функциями активации ReLU и слоями BatchNormalization в узлах сложения результатов ветвей слоя. В качестве функций активации полносвязных слоев используется специализированная функция активации $tanhsoft$:

$$\sigma\left( x \right) = x\ tanh(\ln\left( 1 + e^{x} \right))$$

Функция потерь, используемая при обучении сети - LP-loss:

$$\mathcal{L =}\mathbb{E}\left\lbrack \frac{\left| \left| U^{*} - U \right| \right|}{\left| U \right|} \right\rbrack\ $$

Для обучения мы используем 1000 обучающих и 200 тестовых пар тензоров состояний $\hat{A_{v}}\ $ и $\hat{U_{v}}$. Выборка была составлена при решении уравнений Навье-Стокса (1) с коэффициентом вязкости $\nu = 10^{- 5}$. Размер расчетной сетки 64х64, интервал времени, на котором рассматриваются примеры - 20 секунд.

Алгоритм оптимизации - Adam с начальным шагом обучения $\gamma$, шаг обучения изменяется по убывающей экспоненте: $\gamma_{t + 1\ } = \gamma_{t}\exp\left( - kt \right),\ k$ - поправочный коэффициент.

Для удобства запуска обучения с различными параметрами введены конфигурационные файлы, они хранят в себе параметры настройки сети, которыми можно легко управлять при запуске новых экспериментов. Это позволяет не хранить исходный код и не производить запуск обучения у себя на компьютере, а обучаться удаленно или выводить обучение на виртуальную машину, выбирая необходимый файл конфигураций.

4.3. Установка и запуск обучения модели

git clone https://github.com/AndreyRysistov/NavierStokesFNO
pip install -r -q requirements.txt
python main.py -c configs/config.json 

---

5. Результаты

На рисунках 2 и 3 представлен график изменения функционала ошибки (Lp-loss) и значения метрики качества (MAE) в течение обучения на 70 эпохах.

а) График изменения функции ошибки б) График изменения метрики качества

На рисунке 3 представлена демонстрация результатов работы нейронной сети на тестовой выборке.

Рис 3. Демонстрация результатов работы нейронной сети на тестовой выборке. Слева – истинное состояние жидкости, справа – выход модели

---

6. Анализ результатов и выводы

Таким образом, в результате проделанной работы был реализован генератор предобработки и подачи данных в модель, процесс ее обучения и вывода результатов моделирования.

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

Такое решение в перспективе можно использовать в качестве начального приближения для решения уравнений Навье-Стокса на суперкомпьютерных вычислителях и повышать скорость сходимости более точных методов решения уравнений в тысячи раз.

Стоит отметить, что полученное решение обладает свойством инвариантности к размеру расчетной сетки. Благодаря тому, что преобразование Фурье инвариантно к размеру своего входа, то веса,полученные на расчетной сетке 64х64 ячейки можно перенести на любую другую расчетную сетку, например, 256х256.

---

7. Перспективы развития

Дальнейшее развитие работы предполагает производить обучение на сетках большего размера, чтобы приближаться к возможности решения реальных задач. Более того, предложенная модель способна выдавать первое приближение для уравнений, в которых не учитываются стенки твердого тела, в котором находится жидкость или газ, на которых может создаваться дополнительная завихренность.

Для решения таких уравнений нейросетевыми подходами требуется большая обучающая выборка, которая может быть получена только с помощью многочисленного решения уравнений Навье-Стокса с помощью классических методов, что занимает очень много времени.