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

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)$ строится следующим образом:

Для каждой пары тензоров $(\hat{A_{v}}, \hat{U_{v}})_{i}, i = 1,..., M$:
Проинициализировать общую ошибку $L$
Для каждого момента времени $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. Реализация алгоритма

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

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

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. Перспективы развития

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

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

AndreyRysistov / NavierStokesFNO