PINN или нейросети, знающие физику

Что такое PINN и какова область их применения?

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

Можно сказать, что PINN — это замена численному моделированию и тогда, возникает вопрос: “А нужны ли нейросети там, где хорошо работают численные методы”. Но не все так просто. Представьте, что Вы или Ваш коллега провел эксперимент, например, измерил скорости частиц в потоке жидкости, или получил точечные измерения температуры. Если Вы действительно имеете опыт проведения эксперимента, наверняка знаете, что экспериментальные данные далеко не идеальны и могут принести много головной боли при обработке. А теперь представим, что все же Вы эту обработку провели, получили датасет из эксперимента, и теперь хотите эти данные использовать в уравнениях, чтобы получить другие параметры течения. Например, измерили скорость, а из уравнений гидродинамики, хотите получить давление. Или другими словами, провести ассимиляцию данных, говоря на современном околонаучным языком. Численное моделирование в таком случае можно дать сбой, потому что даже тщательно отфильтрованные данные могут быть шумными (особенно если от них требуется брать производные, а если еще и вторые, то совсем, шляпа). Или их может быть мало (например, температуру измеряли термопарой в нескольких точках). В данном случае, вроде, эксперимент есть и потенциально восстанавливать одни величины по другим можно, решая уравнения. И тут на помощь могут прийти PINN. Потому что они работают иначе, чем численные методы. Они не используют схем переноса, а параметры нейросети минимизируются в выбранных точках.

Кроме того, в PINN не используются традиционные численные производные, а используется такой инструмент — автоматическое дифференцирование. Ведь, напомню, PINN существуют, чтобы минимизировать невязки от дифференциальных уравнений. Постараюсь объяснить “наивно”, что такое “автоматические” производные. PINN — полносвязная нейросеть, то есть несколько слоев с определенным количеством нейронов на каждом слое. При обратном проходе по входным параметрам (напомню, что здесь входные данные — это пространственные координаты, время) берутся производные. А теперь задумаемся: ведь это не что иное, как обычные производные, только, можно сказать, “производная” в точке, в пределе. Обычная численная производная — это разности, а здесь — это просто вложенные в нейросеть весовые коэффициенты. Мы же можем использовать эти производные, как необходимые нам. И вторые производные, и третьи… Да, можем. И вот мощное преимущество PINN. Вспомним, что экспериментальные данные — это часто шумные, разреженные данные. А нейросеть, при правильном обучении, может выучить эти данные и предоставит нам автоматические производные в нужным нам точках.

В этом состоит идея использования нейросетей для ассимиляции данных.

Простой пример с гармоническим осциллятором

В данном обзоре я хочу на простом примере показать, как обучить нейросеть решать уравнение для обычного гармонического осциллятора (грузика массы m на пружине жесткостью k) и продемонстрировать, что такое “автоматические” производные, указать на некоторые особенности PINN. Возьмем самый обычный грузик на пружине. Такой пример изучается еще в школе, и хорошо известен. На рис. 1 изображена знакомая схема.

Уравнение, которым будем физически “информировать” нашу PINN, то есть уравнение движения гармонического осциллятора:

Пускай грузик будет выведен из начального положение равновесия () в положение и отпущен без начальной скорости:

.  Как следует из уравнения движения, .

. Решая данное уравнение, получаем закон движения грузика:

Это и хотелось бы получить после обучения нейросети. 

Архитектура нейросети

Теперь самое интересное — как создать PINN и обучить ее? На самом деле, ничего сложного здесь нет. Очень важны такие гиперпараметры, как количество слоев и количество нейронов на каждом слое (то есть глубина нейросети), так как это влияет на время обучения и естественно, на результат. Однозначных ответов, почему надо выбрать конкретное количество нейронов и слоев в разных статьях я не нашла, поэтому тут есть некоторый элемент подбора. Хочу отметить, что лучше начинать с маленького количества слоев (скажем, 5) и маленького количества нейронов (например, по 20 на каждом слое). Конечно, это зависит от количества выходов и входов нейросети. В данном примере я буду показывать результат с нейросетью из 4 скрытых слоев и 20 нейронов на каждом слое. Между слоями используется функция активации гиперболический тангенс, на выходе последнего слоя также есть гиперболический тангенс, чтобы выходы нейросети находились в пределах от -1 до 1. Предварительно я выбрала , чтобы удовлетворять этому условию, но всегда можно сделать надлежащую нормировку, поэтому здесь не так важно, какое начальное условие выбрать. выбрано для простоты.

Далее я приведу части кода. Использовалась библиотека pytorch.

В качестве оптимизатора был выбран L-BFGS (о методе можно прочесть здесь). Чаще всего в статьях встречается либо этот оптимизатор, либо использование двух оптимизаторов: предварительное обучение с помощью Adam, затем L-BFGS. Так как задача достаточно простая, я решила не перегружать обучение и использовала только последний оптимизатор.

Вот параметры задачи: возьмем грузик с собственной частотой колебаний () и рассмотрим промежуток, то есть наш грузик должен будет совершить два колебания. Тут я намеренно опускаю все размерности, будет считать, что все величины предварительно обезразмерены.

Собственно, вот сама PINN:

class simpleModel(nn.Module):   def __init__(self,                hidden_size=20):     super().__init__()     self.layers_stack = nn.Sequential(         nn.Linear(1, hidden_size),         nn.Tanh(),         nn.Linear(hidden_size, hidden_size), #1         nn.Tanh(),         nn.Linear(hidden_size, hidden_size), #2         nn.Tanh(),         nn.Linear(hidden_size, hidden_size), #3         nn.Tanh(),         nn.Linear(hidden_size, hidden_size), #4         nn.Tanh(),         nn.Linear(hidden_size, 1),         nn.Tanh(),     )    def forward(self, x):     return self.layers_stack(x)

А вот сам процесс обучения:

steps=1000 pbar = tqdm(range(steps), desc='Training Progress') t = (torch.linspace(0, 1, 100).unsqueeze(1)).to(device) t.requires_grad = True  metric_data = nn.MSELoss() writer = SummaryWriter() optimizer = torch.optim.LBFGS(model.parameters(), lr=0.1)  def train():          for step in pbar:         def closure():             optimizer.zero_grad()             loss = pdeBC(t)             loss.backward()             return loss          optimizer.step(closure)         if step % 2 == 0:             current_loss = closure().item()             pbar.set_description("Step: %d | Loss: %.6f" %                                  (step, current_loss))             writer.add_scalar('Loss/train', current_loss, step)  train() writer.close()

Пожалуй, самое главное здесь, это функция pdeBC():

def pdeBC(t):     out = model(t).to(device)     f1 = pde(out, t)      inlet_mask = (t[:, 0] == 0)     t0 = t[inlet_mask]     x0 = model(t0).to(device)     dx0dt = torch.autograd.grad(x0, t0, torch.ones_like(t0), create_graph=True, \                         retain_graph=True)[0]      loss_bc = metric_data(x0, x0_true) + \                 metric_data(dx0dt, dx0dt_true.to(device))     loss_pde = metric_data(f1, torch.zeros_like(f1))      loss = 1e3*loss_bc + loss_pde      metric_x = metric_data(out, x0_true * torch.sin(omega*t + torch.pi / 2))     metric_x0 = metric_data(x0, x0_true)     metric_dx0dt = metric_data(dx0dt, dx0dt_true.to(device))      acc_metrics = {'metric_x': metric_x,                 'metric_x0': metric_x0,                 'metric_dx0dt': metric_dx0dt,                 }      metrics = {'loss': loss,                 'loss_bc': loss_bc,                 'loss_pde': loss_pde,                 }     for k, v in metrics.items():         run[k].append(v)     for k, v in acc_metrics.items():         run[k].append(v)      return loss 

Вместе с самим уравнением:

def pde(out, t, nu=2):     omega = 2 * torch.pi * nu     dxdt = torch.autograd.grad(out, t, torch.ones_like(t), create_graph=True, \                             retain_graph=True)[0]     d2xdt2 = torch.autograd.grad(dxdt, t, torch.ones_like(t), create_graph=True, \                             retain_graph=True)[0]     f = d2xdt2 + (omega ** 2) * out     return f

Обращаю внимание, что torch.autograd.grad — это и есть автоматическая производная. И, как можно заметить, . То есть Loss-функция состоит из двух частей: невязки от уравнений и отклонений от граничных условий. Более сложные PINN включают в себя еще больше частей, например, как отмечалось в введении, если есть экспериментальные данные, можно включить часть, связанную с отклонениями от этих экспериментальных данных. И важно правильно подобрать веса, чтобы нейросеть не стремилась минимизировать только какую-то одну часть. В своем примере я подобрала вес 1000 перед перед .

Привожу график обучения:

А вот результат предсказанных нейросетью значений и аналитических для самой координаты и полученной с помощью torch.autograd.grad скорости :

Получилось достаточно неплохо (а именно, L2-норма отклонения предсказанного результата от аналитического решения составляет менее 2%).

Заключение. Мы разобрали простой пример решения задачи гармонического осциллятора с помощью PINN. Хотелось показать функционал PINN и инструмент автоматического дифференцирования. Существуют намного более сложные примеры с включением экспериментальных данных, где геометрия трехмерная и задачи нестационарные, но об этом в следующий раз 🙂

P.S. Код с сохраненными весами модели доступен в репозитории.