{
 "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/tutorials/blocklib/externalModels/Approximation/springMassOscillation/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Загрузка необходимых библиотек"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install torch pandas onnx netron"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Импорт библиотек\n",
    "\n",
    "Для работы потребуются:\n",
    "- torch - библиотека для работы с тензорами и нейронными сетями.\n",
    "- torch.nn - модуль для создания нейронных сетей.\n",
    "- onnx - библиотека для работы с моделями в формате ONNX.\n",
    "- pandas - пакет для выполнения анализа результатов моделирования.\n",
    "- netron - инструмент для визуализации моделей нейронных сетей.\n",
    "- json - модуль для работы с JSON-данными.\n",
    "- IPython.display - модуль для отображения изображений и других медиа в Jupyter-ноутбуках."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import pandas as pd\n",
    "import onnx\n",
    "import json\n",
    "import netron\n",
    "from IPython.display import IFrame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Определение модели PINN\n",
    "\n",
    "`PINN` - класс модели, представляющий полносвязную физически-информированную нейронную сеть.\n",
    "- `__init__` - конструктор класса, который инициализирует слои сети. Сеть состоит из нескольких линейных слоев, соединенных функцией активации `Tanh`.\n",
    "- `forward` - метод, который определяет, как данные проходят через сеть. Он принимает входные данные (время) и возвращает предсказанные значения."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PINN(nn.Module):\n",
    "    \"\"\"\n",
    "    Физически-информированная нейронная сеть (PINN) для моделирования динамики колебаний.\n",
    "\n",
    "    Arguments\n",
    "    ---------\n",
    "    layers: nn.ModuleList\n",
    "        Список линейных слоев нейронной сети.\n",
    "    activation: callable\n",
    "        Функция активации, используемая в скрытых слоях (Tanh).\n",
    "\n",
    "    Methods\n",
    "    -------\n",
    "    forward: torch.tensor -> torch.tensor\n",
    "        Прямой проход через сеть, возвращающий предсказанные значения для входных данных t.\n",
    "    \"\"\"\n",
    "    def __init__(self, layers: list):\n",
    "        super(PINN, self).__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "        for i in range(len(layers) - 1):\n",
    "            self.layers.append(nn.Linear(layers[i], layers[i + 1]))\n",
    "        self.activation = nn.Tanh()\n",
    "\n",
    "    def forward(self, t: torch.Tensor) -> torch.Tensor:\n",
    "        \"\"\"\n",
    "        Прямой проход через сеть.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        t: torch.Tensor\n",
    "            Входные данные (время), для которых необходимо получить предсказания.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        x: torch.Tensor\n",
    "            Предсказанные значения координат для входных значений t.\n",
    "        \"\"\"\n",
    "        inputs = t\n",
    "        for layer in self.layers[:-1]:\n",
    "            inputs = self.activation(layer(inputs))\n",
    "        x = self.layers[-1](inputs)\n",
    "        \n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Параметры системы\n",
    "\n",
    "- m - масса системы, которая влияет на динамику колебаний.\n",
    "- k - коэффициент упругости, определяющий жесткость пружины.\n",
    "- r - коэффициент демпфирования, который влияет на затухание колебаний.\n",
    "- x_0 - начальное смещение системы от равновесного положения.\n",
    "- v_0 - начальная скорость системы.\n",
    "- T - Общее время моделирования, в течение которого будет происходить обучение и предсказание."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = 1.0 # [кг]\n",
    "k = 50.0 # [Н/м]\n",
    "r = 2.0 # [Н·с/м]\n",
    "x_0 = 2.0 # [м]\n",
    "v_0 = 0.0 # [м/c]\n",
    "T = 10.0 # [с]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Определение функции потерь\n",
    "\n",
    "Функция потерь включает три компонента:\n",
    "- **Потери на начальных условиях** – сравнение предсказанного значения координаты и скорости в момент времени $ t=0 $.\n",
    "- **Потери для уравнения** – невязка, посчитанная по дифференциальному уравнению $ \\ddot{x}(t) + \\frac{r}{m} \\, \\dot{x}(t) + \\frac{k}{m} \\, x(t) = 0 $.\n",
    "\n",
    "Для начала рассмотрим компонент потерь для уравнения и напишем функцию для его расчета."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss_function(model: PINN,\n",
    "                  t_f: torch.Tensor,\n",
    "                  t_i: torch.Tensor) -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Вычисляет функцию потерь для модели PINN по уравнению колебаний.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    model: PINN\n",
    "        Объект модели PINN.\n",
    "    t_f: torch.Tensor\n",
    "        Входные данные для вычисления потерь на основе уравнений движения.\n",
    "    t_i: torch.Tensor\n",
    "        Входные данные для начальных условий.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple: Кортеж, содержащий:\n",
    "    - loss (torch.Tensor): Значение функции потерь.\n",
    "    - (loss_x0, loss_v0, loss_f) (tuple): Значения потерь для начального смещения, начальной скорости и уравнения движения.\n",
    "    \"\"\"\n",
    "    x_f = model(t_f)\n",
    "    x_i = model(t_i)\n",
    "\n",
    "    x_f_t = torch.autograd.grad(x_f, t_f, \n",
    "                                grad_outputs=torch.ones_like(x_f),\n",
    "                                create_graph=True)[0]\n",
    "    \n",
    "    x_f_tt = torch.autograd.grad(x_f_t, t_f, \n",
    "                                 grad_outputs=torch.ones_like(x_f_t),\n",
    "                                 create_graph=True)[0]\n",
    "\n",
    "    x_i_t = torch.autograd.grad(x_i, t_i, \n",
    "                                grad_outputs=torch.ones_like(x_i),\n",
    "                                create_graph=True)[0]\n",
    "    \n",
    "    global r, k, m, x_0, v_0\n",
    "    \n",
    "    f = x_f_tt + (r/m) * x_f_t + (k/m) * x_f\n",
    "\n",
    "    loss_f = torch.mean(torch.squeeze(f)**2)\n",
    "    loss_x0 = torch.mean((torch.squeeze(x_i) - x_0)**2)\n",
    "    loss_v0 = torch.mean((torch.squeeze(x_i_t) - v_0)**2)\n",
    "\n",
    "    coeff0, coeff1, coeff2 = 1, 1e-1, 1e-3\n",
    "    loss = coeff0 * loss_x0 + coeff1 * loss_v0 + coeff2 * loss_f\n",
    "    \n",
    "    return loss, (loss_x0.item(), \n",
    "                  loss_v0.item(), \n",
    "                  loss_f.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Обучение модели\n",
    "\n",
    "Определим, возможно ли использовать GPU (если доступен) или CPU для выполнения обучения. Это позволяет ускорить обучение модели."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Определим архитектуру нейронной сети, состоящую из входного слоя, нескольких скрытых слоев и выходного слоя.\n",
    "\n",
    "Создадим экземпляр модели PINN и переносем её на выбранное устройство (CPU или GPU).\n",
    "\n",
    "Инициализирует оптимизатор Adam для обновления параметров модели с заданной скоростью обучения."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "layers = [1, 32, 32, 32, 32, 1]\n",
    "\n",
    "model = PINN(layers).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=4e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Проведем подготовку данных для обучения. Сама задача может быть отнесена к категории обучения без учителя, однако необходимо задать точки, в которых будут вычисляться остатки уравнения и начальных условий.\n",
    "\n",
    "Для этого определим количество точек времени, используемых для обучения модели. Сгенерируем их линейно распределенными от момента времени 0 до до момента конца моделирования $ T $, а также переносим их на устройство. Эти точки будут использоваться для вычисления потерь. Для учета начальных данных определим начальное время (0) и также перенесем его на устройство."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_f = 1000\n",
    "\n",
    "t_f = torch.linspace(0, T, N_f, device=device).view(-1, 1).requires_grad_(True)\n",
    "t_i = torch.tensor(0., device=device).view(-1, 1).requires_grad_(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Опишем основной цикл обучения. Зададим **num_epochs** - количество эпох для обучения модели.\n",
    "\n",
    "В цикле происходит:\n",
    "- Обнуление градиентов оптимизатора.\n",
    "- Вычисление функции потерь с использованием ранее объявленной функции `loss_function`.\n",
    "- Обратное распространение ошибки для вычисления градиентов.\n",
    "- Обновление параметров модели с помощью оптимизатора.\n",
    "- Вывод информация о текущих потерях."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 20000\n",
    "\n",
    "for epoch in range(num_epochs+1):\n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    loss, losses = loss_function(model, t_f, t_i)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if epoch % 1000 == 0:\n",
    "        print(f'Epoch {epoch},',\n",
    "              f'Losses: x_i = {round(losses[0], 4)},', \n",
    "              f'v_i = {round(losses[1], 4)},', \n",
    "              f'f = {round(losses[2], 4)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Визуализация результатов PINN модели\n",
    "\n",
    "На следующем шаге выполним вывод графика аппроксимированных перемещений точечной массы при затухающих колебаниях, полученный с помощью обученной PINN модели."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "displacement = model(t_f)\n",
    "approx_data = pd.DataFrame(dict(\n",
    "    time = t_f.detach().numpy().flatten(),\n",
    "    displacement = displacement.detach().numpy().flatten(),\n",
    "))\n",
    "axis = approx_data.plot(\n",
    "    y = \"displacement\", \n",
    "    x = \"time\", \n",
    "    color = \"tab:red\", \n",
    "    label = \"Перемещение $x [м]$\"\n",
    "    )\n",
    "axis.set_ylabel(\"Перемещение $x [м]$\")\n",
    "axis.set_xlabel(\"Время $t [с]$\")\n",
    "axis.set_xlim(0, 10)\n",
    "axis.grid(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Экспорт модели в формат ONNX\n",
    "\n",
    "Для последующего использования модели в [REPEAT](https://app.repeatlab.ru) необходимо перевести ее в режим двойной точности, изменив тип данных составляющих модель тензоров. Для этого необходимо вызвать метод `double()` у объекта модели. \n",
    "\n",
    "Кроме того, для корректной интеграции необходимо обозначить входные и выходные параметры модели. Для этого создается шаблон единичного тензора входных данных (случайное число из временного диапазона моделирования), которому, как и выходному параметру, необходимо присвоить корректное имя, отражающее его физический смысл. После чего, можно экспортировать модель, вызвав функцию `torch.onnx.export`, передав в нее все необходимые параметры и указав имя файла модели, например `pinn_oscillator.onnx`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.double()\n",
    "\n",
    "x = torch.randn(1, 1, device=device, dtype=torch.float64) * T\n",
    "\n",
    "input_names = ['t']\n",
    "output_names = ['x']\n",
    "\n",
    "torch.onnx.export(model,\n",
    "                  (x),\n",
    "                  \"[DEMO] OscillatorModel.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` не позволяет передать дополнительные данные о граничных значениях параметров, в пределах которых проводился расчет. Для их добавления необходимо загрузить ранее экспортированную модель из файла, создать словарь метаданных, содержащий информацию о границах, в данном случае — значение $ t $, которое ограничено временем моделирования $ T $.\n",
    "\n",
    "Далее создается объект `StringStringEntryProto`, который будет содержать ключ и значение метаданных, которых добавляется в список свойств модели и ее обновленная версия с добавленными метаданными сохраняется, например, в тот же файл."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_model = onnx.load(\"[DEMO] OscillatorModel.onnx\")\n",
    "\n",
    "metadata_dict = {\"t\": f\"{T}\"}\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",
    "\n",
    "onnx.save(onnx_model, \"[DEMO] OscillatorModel.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Проверка файла ONNX\n",
    "\n",
    "Опционально, можно проверить правильность структуры модели, параметров и метаданных. Для этого может использоваться библиотека `Netron`, которая визуализирует структуру модели."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "onnx_model_path = \"[DEMO] OscillatorModel.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": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "netron.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Интеграция PINN модели в проект REPEAT\n",
    "\n",
    "Теперь всё готово для использования PINN модели в проекте на [платформе REPEAT](https://app.repeatlab.ru) с помощью блока **Аппроксимационная модель**."
   ]
  }
 ],
 "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
}