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

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

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

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

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

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

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

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

Рис. 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. График обучения

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

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

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

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

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