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

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

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

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

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

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

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

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

Рис. 1. Схема грузика на пружинке, совершающего гармонические колебания

Рис. 1. Схема грузика на пружинке, совершающего гармонические колебания

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

\frac{\mathrm{d}x}{\mathrm{d}t} + \omega x = 0

Пускай грузик будет выведен из начального положение равновесия ($$x=0$$) в положение x=x_0 и отпущен без начальной скорости: \frac{\mathrm{d}x}{\mathrm{d}t}(t=0)= 0

.  Как следует из уравнения движения, \omega = \sqrt{\frac{k}{m}}.

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

x(t)=x_0cos(\omega t)

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

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

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

8a64cc25b28b210f1f8117b4931c82c1.png

Рис. 2. Схема полносвязной нейросети (PINN)

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

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

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

Собственно, вот сама 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 = 1e3*loss_{bc} + loss_{pde}. То есть Loss-функция состоит из двух частей: невязки от уравнений и отклонений от граничных условий. Более сложные PINN включают в себя еще больше частей, например, как отмечалось в введении, если есть экспериментальные данные, можно включить часть, связанную с отклонениями от этих экспериментальных данных. И важно правильно подобрать веса, чтобы нейросеть не стремилась минимизировать только какую-то одну часть. В своем примере я подобрала вес 1000 перед перед loss_{bc}.

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

Рис. 3. График обучения

Рис. 3. График обучения

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

Рис. 4. Результат для координаты и скорости, полученной с помощью автоматического дифференцирования.

Рис. 4. Результат для координаты и скорости, полученной с помощью автоматического дифференцирования.

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

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

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

© Habrahabr.ru