{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Код для построения и обучения модели PINN для решения двумерного уравнения теплопроводности"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В данном коде к практическому руководству представлена инструкция по созданию PINN — физически-информированной нейронной сети — для моделирования процесса теплопереноса в двумерном теле, её доработки синусоидальной функцией активации, обучению с использованием физических законов и двух оптимизаторов и экспорту в формат [ONNX](https://onnx.ai).\n",
    "\n",
    "Код практического руководства подготовлен компанией [ООО «ДЖЭТ ЛАБ»](https://get-lab.ru) для демонстрации интеграции полученной модели на [платформу REPEAT](https://app.repeatlab.ru) для выполнения инженерных расчётов с применением моделей машинного обучения.\n",
    "\n",
    "Данное практическое руководство является частью [документации пользователя платформы REPEAT](https://app.repeatlab.ru/docs/).\n",
    "\n",
    "Для более полного ознакомления с поставленной задачей рекомендуется ознакомиться с уроком на [сайте](https://app.repeatlab.ru/docs/howto/ml_approximation/heat_equation/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Загрузка необходимых библиотек"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install torch numpy onnx netron matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Импорт библиотек\n",
    "\n",
    "Для работы потребуются:\n",
    "- torch - библиотека для работы с тензорами и нейронными сетями.\n",
    "- torch.nn - модуль для создания нейронных сетей.\n",
    "- numpy - библиотека для работы с массивами и математическими операциями.\n",
    "- onnx библиотека для работы с моделями в формате ONNX.\n",
    "- json - модуль для работы с JSON-данными.\n",
    "- netron - инструмент для визуализации моделей нейронных сетей.\n",
    "- matplotlib - библиотека для построения графиков и визуализации данных.\n",
    "- IPython.display - модуль для отображения изображений и других медиа в Jupyter-ноутбуках."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import onnx\n",
    "import json\n",
    "\n",
    "import netron\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import IFrame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Определение аналитического решения\n",
    "\n",
    "Для проверки обученной модели получем функцию аналитического решения. Она позволит как посчитать ошибку по $ L2 $ норме в конкретный момент времени, так и провести сравнительный анализ эволюции решения полученного с помощью PINN и аналитическим решением."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def analytical_solution(x: np.ndarray, \n",
    "                        y: np.ndarray, \n",
    "                        t: np.ndarray,\n",
    "                        bounds: dict,\n",
    "                        alpha: float) -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Аналитическое решение\n",
    "\n",
    "    Параметры:\n",
    "        x: np.ndarray\n",
    "            Вектор координат x.\n",
    "        y: np.ndarray\n",
    "            Вектор координат y.\n",
    "        t: np.ndarray\n",
    "            Вектор координат t.\n",
    "        bounds: dict\n",
    "            Границы расчетной области.\n",
    "        alpha: float\n",
    "            Коэффициент температурапроводности материала.\n",
    "\n",
    "    Возвращает:\n",
    "        np.ndarray\n",
    "            Вектор значений температур, являющихся аналитическим решением.\n",
    "    \"\"\"\n",
    "    u = (np.sin(np.pi * x/bounds['x']) * np.sin(np.pi * y/bounds['y'])\n",
    "         * np.exp(-2 * np.pi**2 * alpha * t/(bounds['x']*bounds['y'])))\n",
    "    \n",
    "    return u"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Определение линейного слоя с синусоидальной функцией активации\n",
    "\n",
    "Пользуясь выводами из [статьи](https://arxiv.org/abs/2006.09661) и используя за основу представленный авторами [код](https://github.com/vsitzmann/siren) реализуем класс линейного слоя с синусоидальной функцией активации, как доработку линейного слоя.\n",
    "\n",
    "`SinLayer` - класс слоя, представляющий собой слой нейронной сети с синусоидальной функцией активации.\n",
    "- `__init__` - конструктор класса, который принимает параметры слоя.\n",
    "- `init_weights()` - метод, который задает начальные значения весов.\n",
    "- `forward` - метод,  который выполняет прямой проход через слой.\n",
    "\n",
    "Этот слой используется в архитектуре нейросетей, таких как SIREN (Sinusoidal Representation Networks), где синусоидальная активация позволяет эффективно представлять сложные функции."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SinLayer(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    Слой нейронной сети с синусоидальной функцией активации.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    in_features: int\n",
    "        Количество входных признаков.\n",
    "    out_features: int\n",
    "        Количество выходных признаков.\n",
    "    bias: bool\n",
    "        Использовать ли смещение в линейном слое.\n",
    "    is_first: bool\n",
    "        Является ли слой первым в сети.\n",
    "    omega_0: float\n",
    "        Частота для функции синуса.\n",
    "\n",
    "    Methods\n",
    "    -------\n",
    "    init_weights\n",
    "        Инициализирует веса слоя с учетом его положения в сети.\n",
    "    forward: torch.tensor -> torch.tensor\n",
    "        Прямой проход через слой, реализующий линейный слой и функцию активации.\n",
    "    \"\"\"\n",
    "    def __init__(self, \n",
    "                 input: float, \n",
    "                 output: float, \n",
    "                 bias: bool=True,\n",
    "                 is_first: bool=False, \n",
    "                 omega_0: float=30):\n",
    "        super().__init__()\n",
    "        self.omega_0 = omega_0\n",
    "        self.is_first = is_first\n",
    "        self.in_features = input\n",
    "        self.linear = torch.nn.Linear(input, output, bias=bias)\n",
    "        self.init_weights()\n",
    "    \n",
    "    def init_weights(self):\n",
    "        \"\"\"\n",
    "        Инициализация весов линейного слоя.\n",
    "        \"\"\"\n",
    "        with torch.no_grad():\n",
    "            if self.is_first:\n",
    "                self.linear.weight.uniform_(-1 / self.in_features, \n",
    "                                            1 / self.in_features)\n",
    "            else:\n",
    "                self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) \n",
    "                                            / self.omega_0, \n",
    "                                            np.sqrt(6 / self.in_features) \n",
    "                                            / self.omega_0)\n",
    "        \n",
    "    def forward(self, input: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Прямой проход через слой.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        input: torch.Tensor\n",
    "            Входной тензор.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Результат применения синусоидальной функции активации.\n",
    "        \"\"\"\n",
    "        return torch.sin(self.omega_0 * self.linear(input))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Определение модели PINN\n",
    "\n",
    "`SinPINN` - класс модели, представляющий полносвязную физически-информированную нейронную сеть.\n",
    "- `__init__` - конструктор класса, который инициализирует слои сети. Сеть состоит из нескольких линейных слоев `SinLayer`.\n",
    "- `forward` - метод, который задаёт порядок прохождения входных данных через сеть. Этот метод принимает на вход тензоры `x`, `y` и `t` и возвращает предсказание модели (температуру)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SinPINN(nn.Module):\n",
    "    \"\"\"\n",
    "    Физически-информированная нейронная сеть (PINN) с функцией активации синус для моделирования теплопереноса.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    layers: nn.ModuleList\n",
    "        Список линейных слоёв нейронной сети.\n",
    "\n",
    "    Methods\n",
    "    -------\n",
    "    forward: (torch.tensor, torch.tensor, torch.tensor) -> torch.tensor\n",
    "        Прямой проход через сеть, возвращающий предсказанные значения температуры для входных данных x, y, t.\n",
    "    \"\"\"\n",
    "    def __init__(self, \n",
    "                 layers: list):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList()        \n",
    "        for i in range(len(layers) - 1):\n",
    "            if i == 0:\n",
    "                self.layers.append(SinLayer(layers[i], \n",
    "                                             layers[i+1], is_first=True))\n",
    "            elif i == len(layers) - 2:\n",
    "                self.layers.append(nn.Linear(layers[i], \n",
    "                                             layers[i+1],))\n",
    "            else:\n",
    "                self.layers.append(SinLayer(layers[i], \n",
    "                                             layers[i+1], is_first=False))\n",
    "    \n",
    "    def forward(self, \n",
    "                x: torch.Tensor, \n",
    "                y: torch.Tensor, \n",
    "                t: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Прямой проход через сеть.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x: torch.Tensor\n",
    "            Тензор координат x.\n",
    "        y: torch.Tensor\n",
    "            Тензор координат y.\n",
    "        t: torch.Tensor\n",
    "            Тензор время t.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        torch.Tensor\n",
    "            Выход сети.\n",
    "        \"\"\"\n",
    "        inputs = torch.cat([x.reshape(-1, 1), \n",
    "                            y.reshape(-1, 1), \n",
    "                            t.reshape(-1, 1)], dim=1)\n",
    "        for layer in self.layers:\n",
    "            inputs = layer(inputs)\n",
    "        return inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Определение функции потерь\n",
    "\n",
    "Функция потерь включает три компонента:\n",
    "- **Потери на начальных условиях** – сравнение предсказанного значения температуры в момент времени $ t=0 $.\n",
    "- **Потери на граничных условиях** – сравнение предсказанного значения температуры в точках $ (x, y) \\in \\partial \\Omega $, где $ \\Omega = [0,L] \\times [0,L] $.\n",
    "- **Потери для уравнения** – невязка, посчитанная по дифференциальному уравнению $ u_t - \\alpha (u_{xx} + u_{yy})  = 0 $.\n",
    "\n",
    "Для начала рассмотрим компонент потерь для уравнения и напишем функцию для его расчета."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pde_residual(model: SinPINN, \n",
    "                 x: torch.Tensor, \n",
    "                 y: torch.Tensor, \n",
    "                 t: torch.Tensor, \n",
    "                 alpha: float):\n",
    "    \"\"\"\n",
    "    Вычисляет остаток по уравнению теплопроводности для модели PINN.\n",
    "    Вид уравнения:\n",
    "        u_t - alpha * (u_xx + u_yy) = 0\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model: torch.nn.Module\n",
    "        Модель, вычисляющая u(x, y, t).\n",
    "    x: torch.Tensor\n",
    "        Тензор координаты x.\n",
    "    y: torch.Tensor\n",
    "        Тензор координаты y.\n",
    "    t: torch.Tensor\n",
    "        Тензор времени t.\n",
    "    alpha: float\n",
    "        Коэффициент теплопроводности.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "        torch.Tensor\n",
    "            Остаток уравнения f = u_t - alpha*(u_xx + u_yy).\n",
    "    \"\"\"\n",
    "    u = model(x, y, t)\n",
    "\n",
    "    u_x, u_y, u_t = torch.autograd.grad(u, [x, y, t], \n",
    "                                        grad_outputs=torch.ones_like(u), \n",
    "                                        create_graph=True)\n",
    "    \n",
    "    u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), \n",
    "                               create_graph=True)[0]\n",
    "    u_yy = torch.autograd.grad(u_y, y, grad_outputs=torch.ones_like(u_y), \n",
    "                               create_graph=True)[0]\n",
    "    \n",
    "    f = u_t - alpha * (u_xx + u_yy)\n",
    "    return f"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Подготовка данных"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Для проверки выполнения уравнений задачи на предсказаниях модели, необходимо создать множество точке, на которых будет считаться функция потерь. Поскольку задача формулируется тремя уравнениями (уравнение теплопереноса, начальное и граничное условие) создадим множества точек коллокации, для каждого из уравнений соответственно."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_pde_points(n: int, \n",
    "                        bounds: dict, \n",
    "                        device: torch.device):\n",
    "    \"\"\"\n",
    "    Генерирует точки коллокации для уравнения PDE.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    n: int\n",
    "        Число точек коллокации.\n",
    "    bounds: dict\n",
    "        Границы области.\n",
    "    device: torch.device\n",
    "        Устройство (CPU или GPU).\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Tuple[torch.Tensor, torch.Tensor, torch.Tensor]\n",
    "        Тензоры x, y, t с requires_grad=True.\n",
    "    \"\"\"\n",
    "    x = (bounds['x'] * torch.rand(n, 1, device=device)).requires_grad_()\n",
    "    y = (bounds['y'] * torch.rand(n, 1, device=device)).requires_grad_()\n",
    "    t = (bounds['t'] * torch.rand(n, 1, device=device)).requires_grad_()\n",
    "    return x, y, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_initial_condition(n: int, \n",
    "                               bounds: dict, \n",
    "                               device: torch.device):\n",
    "    \"\"\"\n",
    "    Генерирует точки для начального условия.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    n: int\n",
    "        Число точек начального условия.\n",
    "    bounds: dict\n",
    "        Границы расчетной области.\n",
    "    device: torch.device\n",
    "        Устройство (CPU или GPU).\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "        Тензоры x, y, t, u, \n",
    "        где t — нулевой тензор, а u = sin(pi * x/L) * sin(pi * y/L).\n",
    "    \"\"\"\n",
    "    x = bounds['x'] * torch.rand(n, 1, device=device)\n",
    "    y = bounds['y'] * torch.rand(n, 1, device=device)\n",
    "    t = torch.zeros_like(x, device=device)\n",
    "    u = torch.sin(np.pi * x/bounds['x']) * torch.sin(np.pi * y/bounds['y'])\n",
    "    return x, y, t, u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_boundary_conditions(n: int, \n",
    "                                 bounds: dict, \n",
    "                                 device: torch.device):\n",
    "    \"\"\"\n",
    "    Генерирует точки граничных условий.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_b: int \n",
    "        Число точек для граничных условий.\n",
    "    bounds: dict\n",
    "        Границы области, например, {'x': 1.0, 'y': 1.0, 't': 1.0}.\n",
    "    device: torch.device\n",
    "        Устройство (CPU или GPU).\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:\n",
    "        x_b, y_b, t_b, u_b, где u_b = 0.\n",
    "    \"\"\"\n",
    "    x = bounds['x'] * torch.rand(n, 1, device=device)\n",
    "    y = bounds['y'] * torch.rand(n, 1, device=device)\n",
    "    t = bounds['t'] * torch.rand(n, 1, device=device)\n",
    "\n",
    "    border_selector = torch.randint(0, 4, (n, 1), device=device)\n",
    "    x[border_selector == 0] = 0.0           # Левая грань\n",
    "    x[border_selector == 1] = bounds['x']   # Правая грань\n",
    "    y[border_selector == 2] = 0.0           # Нижняя грань\n",
    "    y[border_selector == 3] = bounds['y']   # Верхняя грань\n",
    "\n",
    "    u_b = torch.zeros_like(t, device=device)\n",
    "    return x, y, t, u_b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Обучение модели\n",
    "\n",
    "Чтобы достичь воспроизводимости результатов, зафиксируем случайную генерацию. Для ускорения последующих вычислений определим доступное устройство (GPU, если доступно, или CPU). \n",
    "\n",
    "Объявим параметры задачи, такие как число точек коллокации различных типов, границы расчетной области по пространству и времени, а также коэффициент температура проводности"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "np.random.seed(42)\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "N_F, N_B, N_I = 10000, 5000, 2000\n",
    "BOUNDS = {'x': 0.1, 'y': 0.1, 't': 12.0}\n",
    "ALPHA = 1.12 * 1e-4\n",
    "\n",
    "model = SinPINN([3, 64, 1]).to(device)\n",
    "mse_loss = nn.MSELoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "x_f, y_f, t_f = generate_pde_points(N_F, BOUNDS, device)\n",
    "x_i, y_i, t_i, u_i = generate_initial_condition(N_I, BOUNDS, device)\n",
    "x_b, y_b, t_b, u_b = generate_boundary_conditions(N_B, BOUNDS, device)\n",
    "\n",
    "w_pde, w_bc, w_ic = 10.0, 1.0, 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Перейдем к обучению модели. Основной цикл обучения проходит в течение заданного количества эпох, где на каждой эпохе:\n",
    "- Обнуляются градиенты.\n",
    "- Вычисляются компоненты функции потерь.\n",
    "- Производится обратное распространение ошибки.\n",
    "- Обновляются параметры модели.\n",
    "- Сохраняется история значений компонент функции потерь.\n",
    "- Периодически выводится значение потерь для контроля процесса обучения."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 20000\n",
    "\n",
    "loss_history = {'pde': [],\n",
    "                'ic': [],\n",
    "                'bc': []}\n",
    "\n",
    "for epoch in range(epochs+1):\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    u_pred_i = model(x_i, y_i, t_i)\n",
    "    loss_ic = mse_loss(u_pred_i, u_i)\n",
    "    \n",
    "    u_pred_b = model(x_b, y_b, t_b)\n",
    "    loss_bc = mse_loss(u_pred_b, u_b)\n",
    "    \n",
    "    loss_pde = mse_loss(pde_residual(model, x_f, y_f, t_f, ALPHA), \n",
    "                        torch.zeros_like(x_f))\n",
    "    \n",
    "    loss = w_ic * loss_ic + w_bc * loss_bc + w_pde * loss_pde\n",
    "    \n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    loss_history['ic'].append(loss_ic.item())\n",
    "    loss_history['bc'].append(loss_bc.item())\n",
    "    loss_history['pde'].append(loss_pde.item())\n",
    "    \n",
    "    if epoch % 1000 == 0:\n",
    "        print(f\"Epoch {epoch}, Loss: {loss.item():.3e}, \"\n",
    "              f\"IC: {loss_ic.item():.3e}, \"\n",
    "              f\"BC: {loss_bc.item():.3e}, \"\n",
    "              f\"PDE: {loss_pde.item():.3e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Для визуальной оценки качества процесса обучения можно провести визуализацию значений компонент функции потерь от эпох в логарифмической шкале."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_loss_history(loss_history: dict, \n",
    "                      step: int=200):\n",
    "    \"\"\"\n",
    "    Функция для построения графика функции потерь.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    loss_history : dict\n",
    "        Словарь, содержащий данные для потерь.\n",
    "    step : int, optional\n",
    "        Шаг выборки данных для отображения, по умолчанию 100.\n",
    "    \"\"\"\n",
    "    epochs = range(0, len(loss_history['pde']), step)\n",
    "\n",
    "    plt.plot(epochs, loss_history['pde'][::step], label='PDE Loss')\n",
    "    plt.plot(epochs, loss_history['bc'][::step], label='BC Loss')\n",
    "    plt.plot(epochs, loss_history['ic'][::step], label='IC Loss')\n",
    "    \n",
    "    plt.yscale('log')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.grid()\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "plot_loss_history(loss_history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В дополнение к оптимизации методом `Adam` часто используется квазиньютоновский метод `LBFGS`. Он вычислительно сложнее, но хорошо подходит для квадратичных функций потерь, какую мы и используем, позволяя найти минимум за несколько итераций. Однако `LBFGS` склонен попадать в локальные минимумы при спуске, в то время как у `Adam` не такого недостатка. По этим причинам они используются совместно и именно в указанной последовательности.\n",
    "\n",
    "Из-за несколько отличающегося интерфейса взаимодействия в реализации метода `LBFGS` в библиотеке `Pytorch` от остальных оптимизаторов, необходимо объявить вспомогательную функцию `closure`, которая будет вызываться при работе оптимизатора."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def closure():\n",
    "    \"\"\"\n",
    "    Вспомогательная функция для оптимизатора LBFGS.\n",
    "    \"\"\"\n",
    "    optimizer_lbfgs.zero_grad()\n",
    "\n",
    "    loss_ic = mse_loss(model(x_i, y_i, t_i), u_i)\n",
    "\n",
    "    loss_bc = mse_loss(model(x_b, y_b, t_b), u_b)\n",
    "\n",
    "    loss_pde = mse_loss(pde_residual(model, x_f, y_f, t_f, ALPHA), \n",
    "                           torch.zeros_like(x_f))\n",
    "    \n",
    "    loss = w_ic * loss_ic + w_bc * loss_bc + w_pde * loss_pde\n",
    "\n",
    "    loss.backward()\n",
    "    \n",
    "    return loss\n",
    "\n",
    "optimizer_lbfgs = optim.LBFGS(model.parameters(), lr=1.0, \n",
    "                              max_iter=500, history_size=50)\n",
    "optimizer_lbfgs.step(closure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Таким образом мы получили обученную модель для решения поставленной задачи. Проведем оценку полученного решения с помощью PINN сравнив его с объявленным ранее аналитическим решением."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Экспорт модели в формат ONNX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "После обучения модель переводится в режим двойной точности с помощью метода `double()`. Затем, используя шаблон входного тензора, модель экспортируется в формат ONNX. При этом задаются имена входных и выходных параметров для корректной интеграции.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.double()\n",
    "\n",
    "x_temp = torch.randn(1, 1, device=device, dtype=torch.float64) * BOUNDS['x']\n",
    "y_temp = torch.randn(1, 1, device=device, dtype=torch.float64) * BOUNDS['y']\n",
    "t_temp = torch.randn(1, 1, device=device, dtype=torch.float64) * BOUNDS['t']\n",
    "\n",
    "input_names = ['x', 'y', 't']\n",
    "output_names = ['u']\n",
    "\n",
    "torch.onnx.export(model,\n",
    "                  (x_temp, y_temp, t_temp),\n",
    "                  \"[DEMO] 2DheatModel.onnx\",\n",
    "                  opset_version=12,\n",
    "                  input_names=input_names,\n",
    "                  output_names=output_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Добавление метаданных о границах парамеров\n",
    "\n",
    "К сожалению, функция `torch.onnx.export` не позволяет передать дополнительные данные о граничных значениях параметров, в пределах которых проводился расчет. Для их добавления необходимо загрузить ранее экспортированную модель из файла, создать словарь метаданных, содержащий информацию о границах."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_model = onnx.load(\"[DEMO] 2DheatModel.onnx\")\n",
    "\n",
    "metadata_dict = {\"x\": f\"{BOUNDS['x']}\",\n",
    "                 \"y\": f\"{BOUNDS['y']}\",\n",
    "                 \"t\": f\"{BOUNDS['t']}\"}\n",
    "\n",
    "metadata_json = json.dumps(metadata_dict)\n",
    "\n",
    "metadata_prop = onnx.StringStringEntryProto()\n",
    "metadata_prop.key = \"bounds\"\n",
    "metadata_prop.value = metadata_json\n",
    "\n",
    "onnx_model.metadata_props.append(metadata_prop)\n",
    "onnx.save(onnx_model, \"[DEMO] 2DheatModel.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Проверка файла ONNX\n",
    "\n",
    "Опционально, можно проверить правильность структуры модели, параметров и метаданных. Для этого может использоваться библиотека `Netron`, которая визуализирует структуру модели."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_model_path = \"[DEMO] 2DheatModel.onnx\"\n",
    "\n",
    "url = netron.start(onnx_model_path, browse=False)\n",
    "IFrame(src=f\"http://localhost:8080/\", width=600, height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "netron.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Интеграция PINN модели в проект REPEAT\n",
    "\n",
    "Теперь всё готово для использования PINN модели в проекте на [платформе REPEAT](https://app.repeatlab.ru) с помощью блока **Аппроксимационная модель**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}