hydrodynamics_with_deep_learning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TUrmzAiQIJqV"
      },
      "source": [
        "# Загружаем архив и датасет:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kp4LERrO3i7M"
      },
      "outputs": [],
      "source": [
        "# подключаем гугл диск\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZRmti7D_6rOs"
      },
      "outputs": [],
      "source": [
        "# проверяем, что у нас есть архив\n",
        "!ls /content/drive/MyDrive/GNN/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MbEd5zPWysLO"
      },
      "outputs": [],
      "source": [
        "# рахархивируем архив\n",
        "!unzip -q /content/drive/MyDrive/GNN/deepmind-research-master.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xPCIxQWS9lKI"
      },
      "outputs": [],
      "source": [
        "# создаем папки для загрузки датасета\n",
        "!mkdir -p /tmp/rollous\n",
        "!mkdir -p /tmp/datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hkPzW5t79-FP"
      },
      "outputs": [],
      "source": [
        "# загружаем датасет\n",
        "!bash /content/deepmind-research-master/learning_to_simulate/download_dataset.sh WaterRamps /tmp/datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MPdU_w0AIYuh"
      },
      "source": [
        "## Извлекаем данные из TFrecords"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gIIfGNWaJpQW"
      },
      "outputs": [],
      "source": [
        "# переходим в дерикторию архива \n",
        "%cd /content/deepmind-research-master"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "__dqBJJHJMhs"
      },
      "outputs": [],
      "source": [
        "import functools\n",
        "import os\n",
        "import json\n",
        "import pickle\n",
        "\n",
        "import tensorflow.compat.v1 as tf\n",
        "import numpy as np\n",
        "\n",
        "from learning_to_simulate import reading_utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_hZoj8-zJMkM"
      },
      "outputs": [],
      "source": [
        "# Определяем путь к данным и имя файла\n",
        "data_path = '/tmp/datasets/WaterRamps'\n",
        "filename = 'train.tfrecord'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_fIviKJsxCsr"
      },
      "source": [
        "### Metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uvK_5_ycJMmy"
      },
      "outputs": [],
      "source": [
        "# Читаем metadata\n",
        "def _read_metadata(data_path):\n",
        "    with open(os.path.join(data_path, 'metadata.json'), 'rt') as fp:\n",
        "        return json.loads(fp.read())\n",
        "\n",
        "metadata = _read_metadata(data_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GJsl_KWFJMpY"
      },
      "outputs": [],
      "source": [
        "for key in metadata:\n",
        "  print(key, metadata[key])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4s04UTlGWsxl"
      },
      "outputs": [],
      "source": [
        "# Определим глобальные переменные\n",
        "CON_RAD = metadata['default_connectivity_radius'] ** 2\n",
        "SEQ_LEN = metadata['sequence_length']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "szGid7VixYo-"
      },
      "source": [
        "### Data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qmqOVNGQb8GB"
      },
      "source": [
        "В массиве positons находятся позиции точек. positions.shape = [t_steps_num, nodes_num, dim]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KEBuMUUhJMsJ"
      },
      "outputs": [],
      "source": [
        "ds_org = tf.data.TFRecordDataset([os.path.join(data_path, filename)])\n",
        "ds = ds_org.map(functools.partial(reading_utils.parse_serialized_simulation_example, metadata=metadata))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y2qFa3DoLcU1"
      },
      "outputs": [],
      "source": [
        "particle_types = []\n",
        "keys = []\n",
        "positions = []\n",
        "for _ds in ds:\n",
        "    context, features = _ds\n",
        "    particle_types.append(context[\"particle_type\"].numpy().astype(np.int64))\n",
        "    keys.append(context[\"key\"].numpy().astype(np.int64))\n",
        "    positions.append(features[\"position\"].numpy().astype(np.float32))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MOaxgfrbHxrg"
      },
      "outputs": [],
      "source": [
        "res_set = set(particle_types[0])\n",
        "for seq in particle_types:\n",
        "  cur_set = set(seq)\n",
        "  res_set = res_set | cur_set\n",
        "print('Different values in patrical_types:')\n",
        "print(res_set)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rDY1KvnYiAM4"
      },
      "outputs": [],
      "source": [
        "for i in range(len(particle_types)):\n",
        "  for j in range(len(particle_types[i])):\n",
        "    if particle_types[i][j] == 5:\n",
        "      particle_types[i][j] = True\n",
        "    else:\n",
        "      particle_types[i][j] = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vD2nZZfs3l3M"
      },
      "outputs": [],
      "source": [
        "print('Shape of each element in positions:')\n",
        "print(positions[0].shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uvhSA4zKnq5K"
      },
      "outputs": [],
      "source": [
        "def get_borders_features(positions, borders):\n",
        "  return np.concatenate(((positions[:, 0]-borders[0][0]).reshape(-1, 1), (borders[0][1]-positions[:, 0]).reshape(-1, 1), \n",
        "                         (positions[:, 1]-borders[1][0]).reshape(-1, 1), (borders[1][1]-positions[:, 1]).reshape(-1, 1)), axis=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Валидационная выборка"
      ],
      "metadata": {
        "id": "E8YsO6ZvoXuZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Определяем путь к данным и имя файла\n",
        "val_data_path = '/tmp/datasets/WaterRamps'\n",
        "val_filename = 'valid.tfrecord'"
      ],
      "metadata": {
        "id": "4g3SVEb0oZif"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "ds_org = tf.data.TFRecordDataset([os.path.join(val_data_path, val_filename)])\n",
        "ds = ds_org.map(functools.partial(reading_utils.parse_serialized_simulation_example, metadata=metadata))"
      ],
      "metadata": {
        "id": "ujNtQbBRoZkd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "val_particle_types = []\n",
        "val_keys = []\n",
        "val_positions = []\n",
        "for _ds in ds:\n",
        "    context, features = _ds\n",
        "    val_particle_types.append(context[\"particle_type\"].numpy().astype(np.int64))\n",
        "    val_keys.append(context[\"key\"].numpy().astype(np.int64))\n",
        "    val_positions.append(features[\"position\"].numpy().astype(np.float32))"
      ],
      "metadata": {
        "id": "IJNuiKljobiB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(len(val_particle_types)):\n",
        "  for j in range(len(val_particle_types[i])):\n",
        "    if val_particle_types[i][j] == 5:\n",
        "      val_particle_types[i][j] = True\n",
        "    else:\n",
        "      val_particle_types[i][j] = False"
      ],
      "metadata": {
        "id": "ieL9Tt04odrW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74lLt5EkHYkc"
      },
      "source": [
        "# GNN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QSjuffI1x5rn"
      },
      "source": [
        "## Загружаем torch.geometric, определяем библиотеки"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZOY1VOgAJhY7"
      },
      "outputs": [],
      "source": [
        "!pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0+cu113.html\n",
        "!pip install torch-sparse -f https://data.pyg.org/whl/torch-1.10.0+cu113.html\n",
        "!pip install torch-cluster -f https://data.pyg.org/whl/torch-1.10.0+cu113.html\n",
        "!pip install torch-geometric"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yoK8TTYuKG-w"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import Tensor\n",
        "from torch_geometric.nn import GCNConv, MessagePassing, EdgeConv\n",
        "from torch_cluster import knn_graph\n",
        "import torch.nn as nn\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import clear_output\n",
        "from tqdm.notebook import trange\n",
        "from tqdm import tqdm\n",
        "from torch.optim import lr_scheduler"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SMP2Ed9m_VHG"
      },
      "outputs": [],
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BMDF68aWT5HB"
      },
      "source": [
        "## Message passing net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gpt6BMTS7uVc"
      },
      "source": [
        "### CMPNV3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-Y420qGm7zNv"
      },
      "outputs": [],
      "source": [
        "class CMPNV3(MessagePassing):\n",
        "    def __init__(self, in_channels, hidden_channels, k=10):\n",
        "        super().__init__(aggr='add')\n",
        "\n",
        "        self.gnn1_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, hidden_channels//4),\n",
        "                       nn.BatchNorm1d(hidden_channels//4),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//4, hidden_channels//2),\n",
        "                       nn.BatchNorm1d(hidden_channels//2),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//2, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "\n",
        "        self.gnn1 = EdgeConv(nn=self.gnn1_mlp, aggr='add')\n",
        "\n",
        "        self.gnn2_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "        self.gnn2 = EdgeConv(nn=self.gnn2_mlp, aggr='add')\n",
        "\n",
        "        self.gnn3_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "        self.gnn3 = EdgeConv(nn=self.gnn3_mlp, aggr='add')\n",
        "\n",
        "        self.mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, hidden_channels)\n",
        "                       )\n",
        "        \n",
        "        #self.conv = GCNConv(in_channels, conv_channels)\n",
        "\n",
        "        self.lin = nn.Sequential(\n",
        "                       nn.Linear(in_channels+hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, hidden_channels//2),\n",
        "                       nn.BatchNorm1d(hidden_channels//2),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//2, hidden_channels//4),\n",
        "                       nn.BatchNorm1d(hidden_channels//4),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//4, 2)\n",
        "                       )\n",
        "        \n",
        "        self.k = k\n",
        "\n",
        "    def forward(self, x, nodes_mask):\n",
        "        # x has shape [N, in_channels]\n",
        "        # edge_index has shape [2, E]\n",
        "        edge_index = knn_graph(x[:, :2], self.k, loop=True, flow=self.flow)\n",
        "      \n",
        "        new_embedding = x + self.gnn1(x, edge_index)\n",
        "        new_embedding += self.gnn2(new_embedding, edge_index)\n",
        "        new_embedding += self.gnn3(new_embedding, edge_index)\n",
        "        #new_embedding += self.gnn4(new_embedding, edge_index)\n",
        "        #new_embedding += self.gnn5(new_embedding, edge_index)\n",
        "\n",
        "        \n",
        "        return self.propagate(edge_index=edge_index, x=new_embedding, nodes_mask=nodes_mask, first_input=x)\n",
        "\n",
        "    def message(self, x_i, x_j, nodes_mask):\n",
        "        # x_i has shape [E, in_channels]\n",
        "        # x_j has shape [E, in_channels]\n",
        "\n",
        "        return self.mlp(torch.cat([x_i, x_i-x_j], dim=1))\n",
        "      \n",
        "    def update(self, aggr_out, nodes_mask, first_input):\n",
        "        # aggr_out has shape [N, latent_size]\n",
        "\n",
        "        new_embedding = torch.cat([first_input[nodes_mask], aggr_out[nodes_mask]], dim=1)\n",
        "        new_embedding = self.lin(new_embedding)\n",
        "\n",
        "        result = first_input[:, :2]\n",
        "        result[nodes_mask] = new_embedding\n",
        "        \n",
        "        return result"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "net_arch = '''\n",
        "class CMPNV3(MessagePassing):\n",
        "    def __init__(self, in_channels, hidden_channels, k=10):\n",
        "        super().__init__(aggr='add')\n",
        "\n",
        "        self.gnn1_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, hidden_channels//4),\n",
        "                       nn.BatchNorm1d(hidden_channels//4),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//4, hidden_channels//2),\n",
        "                       nn.BatchNorm1d(hidden_channels//2),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//2, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "\n",
        "        self.gnn1 = EdgeConv(nn=self.gnn1_mlp, aggr='add')\n",
        "\n",
        "        self.gnn2_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "        self.gnn2 = EdgeConv(nn=self.gnn2_mlp, aggr='add')\n",
        "\n",
        "        self.gnn3_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "        self.gnn3 = EdgeConv(nn=self.gnn3_mlp, aggr='add')\n",
        "\n",
        "        self.gnn4_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "        self.gnn4 = EdgeConv(nn=self.gnn4_mlp, aggr='add')\n",
        "\n",
        "        self.gnn5_mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, in_channels)\n",
        "                       )\n",
        "        \n",
        "        self.gnn5 = EdgeConv(nn=self.gnn5_mlp, aggr='add')\n",
        "\n",
        "        self.mlp = nn.Sequential(\n",
        "                       nn.Linear(2*in_channels, 2*hidden_channels),\n",
        "                       nn.BatchNorm1d(2*hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(2*hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, hidden_channels)\n",
        "                       )\n",
        "        \n",
        "        #self.conv = GCNConv(in_channels, conv_channels)\n",
        "\n",
        "        self.lin = nn.Sequential(\n",
        "                       nn.Linear(in_channels+hidden_channels, hidden_channels),\n",
        "                       nn.BatchNorm1d(hidden_channels),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels, hidden_channels//2),\n",
        "                       nn.BatchNorm1d(hidden_channels//2),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//2, hidden_channels//4),\n",
        "                       nn.BatchNorm1d(hidden_channels//4),\n",
        "                       nn.ReLU(),\n",
        "                       nn.Linear(hidden_channels//4, 2)\n",
        "                       )\n",
        "        \n",
        "        self.k = k\n",
        "\n",
        "    def forward(self, x, nodes_mask):\n",
        "        # x has shape [N, in_channels]\n",
        "        # edge_index has shape [2, E]\n",
        "        edge_index = knn_graph(x[:, :2], self.k, loop=True, flow=self.flow)\n",
        "      \n",
        "        new_embedding = x + self.gnn1(x, edge_index)\n",
        "        new_embedding += self.gnn2(new_embedding, edge_index)\n",
        "        new_embedding += self.gnn3(new_embedding, edge_index)\n",
        "        new_embedding += self.gnn4(new_embedding, edge_index)\n",
        "        new_embedding += self.gnn5(new_embedding, edge_index)\n",
        "\n",
        "        \n",
        "        return self.propagate(edge_index=edge_index, x=new_embedding, nodes_mask=nodes_mask, first_input=x)\n",
        "\n",
        "    def message(self, x_i, x_j, nodes_mask):\n",
        "        # x_i has shape [E, in_channels]\n",
        "        # x_j has shape [E, in_channels]\n",
        "\n",
        "        return self.mlp(torch.cat([x_i, x_i-x_j], dim=1))\n",
        "      \n",
        "    def update(self, aggr_out, nodes_mask, first_input):\n",
        "        # aggr_out has shape [N, latent_size]\n",
        "\n",
        "        new_embedding = torch.cat([first_input[nodes_mask], aggr_out[nodes_mask]], dim=1)\n",
        "        new_embedding = self.lin(new_embedding)\n",
        "\n",
        "        result = first_input[:, :2]\n",
        "        result[nodes_mask] = new_embedding\n",
        "        \n",
        "        return result\n",
        "'''"
      ],
      "metadata": {
        "id": "pqdHwX6foPXx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lz605a6S40Ec"
      },
      "source": [
        "## Train function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C2o4HPkTlyEC"
      },
      "source": [
        "### Train per step"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yUGQE6Ic-Npp"
      },
      "outputs": [],
      "source": [
        "def get_borders_features(positions, borders):\n",
        "  return np.concatenate(((positions[:, 0]-borders[0][0]).reshape(-1, 1), (borders[0][1]-positions[:, 0]).reshape(-1, 1), \n",
        "                         (positions[:, 1]-borders[1][0]).reshape(-1, 1), (borders[1][1]-positions[:, 1]).reshape(-1, 1)), axis=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DduXPZYipQfl"
      },
      "outputs": [],
      "source": [
        "def build_scene(data, particle_types, inds, metadata, input_len, noise_scale):\n",
        "\n",
        "  i, j = inds\n",
        "\n",
        "  noise = np.random.normal(loc=1.0, scale=noise_scale, size=data[i][j].shape)\n",
        "\n",
        "  features = torch.tensor(data[i][j]*noise, dtype=torch.float)\n",
        "\n",
        "  for l in range(input_len):\n",
        "    features = torch.cat([features, torch.tensor(data[i][j-l] - data[i][j-l-1])], dim=1)\n",
        "\n",
        "  features = torch.cat([features, torch.tensor(get_borders_features(data[i][j], metadata['bounds'])), \n",
        "                        torch.tensor(particle_types[i]).view(-1, 1)], dim=1)\n",
        "  \n",
        "  node_mask = torch.BoolTensor(particle_types[i])\n",
        "  next_pos = torch.tensor(data[i][j+1])\n",
        "\n",
        "  return features, node_mask, next_pos"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V61KgA3R2XOu"
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "import random\n",
        "import time\n",
        "\n",
        "def train_vel(model, optimizer, scheduler, criterion, epochs, train_data, particle_types_train, val_data, particle_types_val):\n",
        "\n",
        "  model.eval()\n",
        "  model.load_state_dict(torch.load('/content/drive/MyDrive/GNN/tets_1/state_dict_model.pt'))\n",
        "\n",
        "  train_losses = []\n",
        "  val_losses = []\n",
        "\n",
        "  train_inds = []\n",
        "  for i in range(len(train_data)):\n",
        "    for j in range(5, metadata['sequence_length']):\n",
        "      train_inds.append((i, j))\n",
        "\n",
        "  val_inds = []\n",
        "  for i in range(len(val_data)):\n",
        "    for j in range(5, metadata['sequence_length']):\n",
        "      val_inds.append((i, j))\n",
        "\n",
        "  for epoch in range(epochs):\n",
        "\n",
        "    start_epoch = time.time()\n",
        "\n",
        "    print('Epoch', epoch+1)\n",
        "\n",
        "    model.train()\n",
        "    seq_train_losses = []\n",
        "\n",
        "    print('Train')\n",
        "\n",
        "    # train_data.shape = [batches, SEQ_LEN, nodes_num, 2]\n",
        "    random.shuffle(train_inds)\n",
        "    for inds in tqdm(train_inds):\n",
        "\n",
        "      features, node_mask, next_pos = build_scene(train_data, particle_types_train, inds, metadata, 5, 0.0003)\n",
        "\n",
        "      features = features.to(device)\n",
        "      node_mask = node_mask.to(device)\n",
        "      next_pos = next_pos.to(device)\n",
        "\n",
        "      optimizer.zero_grad()\n",
        "      y_pred = model(features, node_mask)\n",
        "      loss = criterion(y_pred[node_mask], next_pos[node_mask])\n",
        "      seq_train_losses.append(loss.item())\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "\n",
        "    train_losses.append(np.mean(seq_train_losses))\n",
        "    print('Train loss %f' % train_losses[-1])\n",
        "\n",
        "    with open(\"/content/drive/MyDrive/GNN/tets_1/train_losses.txt\", \"a\") as file:\n",
        "      print(train_losses[-1], file=file)\n",
        "\n",
        "    # validation\n",
        "    model.eval()\n",
        "    seq_val_losses = []\n",
        "\n",
        "    print('Validation')\n",
        "\n",
        "    random.shuffle(val_inds)\n",
        "    with torch.no_grad():\n",
        "\n",
        "      for inds in tqdm(val_inds):\n",
        "\n",
        "        features, node_mask, next_pos = build_scene(val_data, particle_types_val, inds, metadata, 5, 0.0003)\n",
        "\n",
        "        features = features.to(device)\n",
        "        node_mask = node_mask.to(device)\n",
        "        next_pos = next_pos.to(device)\n",
        "\n",
        "        y_pred = model(features, node_mask)\n",
        "        loss = criterion(y_pred[node_mask], next_pos[node_mask])\n",
        "        seq_val_losses.append(loss.item())\n",
        "\n",
        "    scheduler.step()\n",
        "\n",
        "    val_losses.append(np.mean(seq_val_losses))\n",
        "    print('Validation loss %f' % val_losses[-1])\n",
        "    print('')\n",
        "\n",
        "    with open(\"/content/drive/MyDrive/GNN/tets_1/val_losses.txt\", \"a\") as file:\n",
        "      print(val_losses[-1], file=file)\n",
        "\n",
        "    torch.save(model.state_dict(), '/content/drive/MyDrive/GNN/tets_1/state_dict_model.pt')\n",
        "\n",
        "    end_epoch = time.time()\n",
        "    with open(\"/content/drive/MyDrive/GNN/tets_1/time.txt\", \"a\") as file:\n",
        "      print(end_epoch-start_epoch, file=file)\n",
        "\n",
        "\n",
        "  return train_losses, val_losses"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wtQCMObNzm-R"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3aYprCtmDy0"
      },
      "source": [
        "### Data define"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "izAyiZM-5Thl"
      },
      "outputs": [],
      "source": [
        "train_data = positions\n",
        "particle_types_train = particle_types\n",
        "\n",
        "val_data = val_positions\n",
        "particle_types_val = val_particle_types"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZOtIa3-uIV9K"
      },
      "source": [
        "### Define CMPNV3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MhQW6bKLYflm"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z2uoWtSvIQ3D"
      },
      "outputs": [],
      "source": [
        "cmpnv3 = CMPNV3(in_channels=2 + 2*5 + 4 + 1, hidden_channels=128).to(device)\n",
        "optimizer = torch.optim.Adam(cmpnv3.parameters(), lr=1e-3)\n",
        "criterion = torch.nn.MSELoss()\n",
        "scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=[3, 5, 7], gamma=1)\n",
        "epochs = 10"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#with open(\"/content/drive/MyDrive/GNN/tets_1/net_info.txt\", \"w\") as file:\n",
        "#  print(net_arch, file=file)"
      ],
      "metadata": {
        "id": "eI7y0j_noErb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#total_params = sum(p.numel() for p in cmpnv3.parameters())\n",
        "#with open(\"/content/drive/MyDrive/GNN/tets_1/net_info.txt\", \"a\") as file:\n",
        "#  print('TOTAL_PARAMS:', total_params, file=file)"
      ],
      "metadata": {
        "id": "3VH4TTC8oFA2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#with open(\"/content/drive/MyDrive/GNN/tets_1/net_info.txt\", \"a\") as file:\n",
        "#  print('LR_strat:', 'Adam', 'lr=1e-3', file=file)"
      ],
      "metadata": {
        "id": "VK97ncRdoFDb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#torch.save(cmpnv3.state_dict(), '/content/drive/MyDrive/GNN/tets_1/state_dict_model.pt')"
      ],
      "metadata": {
        "id": "ipzHeoZ1oFF0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YRsmMZvj1mV7"
      },
      "source": [
        "### Train per step"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_losses, val_losses = train_vel(cmpnv3, optimizer, scheduler, criterion, epochs, train_data, particle_types_train, val_data, particle_types_val)"
      ],
      "metadata": {
        "id": "06iaS5Gq5mco"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ccI5cpy4EaAG"
      },
      "outputs": [],
      "source": [
        "train_losses, val_losses = train_vel(cmpnv3, optimizer, scheduler, criterion, epochs, train_data, particle_types_train, val_data, particle_types_val)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ps8vZ6btmT29"
      },
      "source": [
        "### Figure"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZiLYadlHFeiA"
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(10, 7))\n",
        "ax = fig.add_subplot(1, 1, 1)\n",
        "\n",
        "ax.plot(range(1, epochs+1), train_losses, color='blue', label='train losses')\n",
        "ax.plot(range(1, epochs+1), val_losses, color='red', label='validation losses')\n",
        "ax.set_title('Losses per epoch:')\n",
        "ax.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vyjCyatR6TxZ"
      },
      "source": [
        "# Посмотрим на предсказанную эволюцию "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IjcKkV_KADUZ"
      },
      "source": [
        "## Предсказание следующего шага"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1bIE-lpuAHdS"
      },
      "outputs": [],
      "source": [
        "test_data = positions[1]\n",
        "particle_types_test = particle_types[1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KLRTv3JDBINg"
      },
      "outputs": [],
      "source": [
        "torch.tensor(positions[1][0]).shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HSREUcnWBBjn"
      },
      "outputs": [],
      "source": [
        "torch.tensor(positions[1][0]).shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CdU3uc-C4PCj"
      },
      "outputs": [],
      "source": [
        "edge_index = knn_graph(torch.tensor(positions[1][0]), 15, loop=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0p_7Vyw9BUcw"
      },
      "outputs": [],
      "source": [
        "edge_index.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "75MTy4e1Ezvn"
      },
      "outputs": [],
      "source": [
        "edge_index[-1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_sATwvFDBYk1"
      },
      "outputs": [],
      "source": [
        "res = np.array(edge_index)\n",
        "for i in res:\n",
        "  ans = i[-15:]\n",
        "  break\n",
        "print(ans)\n",
        "point = ans[0]\n",
        "print(point)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VorJNNhOD-xO"
      },
      "outputs": [],
      "source": [
        "test_data[1].shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jmdc4ep1DTNN"
      },
      "outputs": [],
      "source": [
        "fig = plt.figure(figsize=(24, 12))\n",
        "ax1 = fig.add_subplot(1, 2, 1)\n",
        "ax1.set_xlim(metadata['bounds'][0][0], metadata['bounds'][0][1])\n",
        "ax1.set_ylim(metadata['bounds'][1][0], metadata['bounds'][1][1])\n",
        "\n",
        "X = [test_data[1][j][0] for j in range(test_data[1].shape[0])]\n",
        "Y = [test_data[1][j][1] for j in range(test_data[1].shape[0])]\n",
        "\n",
        "X1 = [test_data[1][j][0] for j in ans]\n",
        "Y1 = [test_data[1][j][1] for j in ans]\n",
        "\n",
        "X2 = test_data[1][point][0]\n",
        "Y2 = test_data[1][point][1]\n",
        "\n",
        "ax1.scatter(X, Y, color='blue')\n",
        "ax1.scatter(X1, Y1, color='red')\n",
        "ax1.scatter(X2, Y2, color='green')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zYCWnfld4Cy0"
      },
      "outputs": [],
      "source": [
        "for i in range(1, 100):\n",
        "\n",
        "  fig = plt.figure(figsize=(24, 12))\n",
        "  ax1 = fig.add_subplot(1, 2, 1)\n",
        "  ax1.set_xlim(metadata['bounds'][0][0], metadata['bounds'][0][1])\n",
        "  ax1.set_ylim(metadata['bounds'][1][0], metadata['bounds'][1][1])\n",
        "\n",
        "  X = [test_data[i][j][0] for j in range(test_data[i].shape[0]) if particle_types_test[j]==1]\n",
        "  Y = [test_data[i][j][1] for j in range(test_data[i].shape[0]) if particle_types_test[j]==1]\n",
        "\n",
        "  X_borders = [test_data[i][j][0] for j in range(test_data[i].shape[0]) if particle_types_test[j]==0]\n",
        "  Y_borders = [test_data[i][j][1] for j in range(test_data[i].shape[0]) if particle_types_test[j]==0]\n",
        "\n",
        "  edge_index = knn_graph(x[:, :2], self.k, loop=True, flow=self.flow)\n",
        "\n",
        "  ax1.scatter(X, Y, color='blue')\n",
        "  ax1.scatter(X_borders, Y_borders, color='red')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-1NKdr0_ALh8"
      },
      "outputs": [],
      "source": [
        "for i in range(1, 100):\n",
        "\n",
        "  fig = plt.figure(figsize=(24, 12))\n",
        "  ax1 = fig.add_subplot(1, 2, 1)\n",
        "  ax1.set_xlim(metadata['bounds'][0][0], metadata['bounds'][0][1])\n",
        "  ax1.set_ylim(metadata['bounds'][1][0], metadata['bounds'][1][1])\n",
        "\n",
        "  X = [test_data[i][j][0] for j in range(test_data[i].shape[0]) if particle_types_test[j]==1]\n",
        "  Y = [test_data[i][j][1] for j in range(test_data[i].shape[0]) if particle_types_test[j]==1]\n",
        "\n",
        "  X_borders = [test_data[i][j][0] for j in range(test_data[i].shape[0]) if particle_types_test[j]==0]\n",
        "  Y_borders = [test_data[i][j][1] for j in range(test_data[i].shape[0]) if particle_types_test[j]==0]\n",
        "\n",
        "  ax1.scatter(X, Y, color='blue')\n",
        "  ax1.scatter(X_borders, Y_borders, color='red')\n",
        "\n",
        "  with torch.no_grad():\n",
        "\n",
        "    nodes_mask = torch.BoolTensor(particle_types_test).to(device)\n",
        "    input = torch.cat([torch.tensor(test_data[i-1], device=device), torch.tensor(particle_types_test, device=device).view(-1, 1)], dim=1)\n",
        "\n",
        "    pred_img = cmpn.forward(input, nodes_mask).cpu()\n",
        "\n",
        "    ax2 = fig.add_subplot(1, 2, 2)\n",
        "    ax2.set_xlim(metadata['bounds'][0][0], metadata['bounds'][0][1])\n",
        "    ax2.set_ylim(metadata['bounds'][1][0], metadata['bounds'][1][1])\n",
        "\n",
        "    X = [pred_img[j][0].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==1]\n",
        "    Y = [pred_img[j][1].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==1]\n",
        "\n",
        "    X_borders = [pred_img[j][0].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==0]\n",
        "    Y_borders = [pred_img[j][1].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==0]\n",
        "    \n",
        "    ax2.scatter(X, Y, color='blue')\n",
        "    ax2.scatter(X_borders, Y_borders, color='red')\n",
        "\n",
        "    clear_output(wait=True)\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lB5WOc5F_-Ir"
      },
      "source": [
        "## Предсказания на основе предыдущих"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Pxov5_MC3vSS"
      },
      "outputs": [],
      "source": [
        "test_data = positions[53]\n",
        "particle_types_test = particle_types[53]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W45EKGRnxcEp"
      },
      "outputs": [],
      "source": [
        "last_out = torch.tensor(test_data[0], device=device)\n",
        "last_out = torch.cat((last_out, torch.tensor(particle_types_test, device=device).view(-1, 1)), dim=1)\n",
        "\n",
        "for i in range(100):\n",
        "\n",
        "  fig = plt.figure(figsize=(24, 12))\n",
        "  ax1 = fig.add_subplot(1, 2, 1)\n",
        "  ax1.set_xlim(metadata['bounds'][0][0], metadata['bounds'][0][1])\n",
        "  ax1.set_ylim(metadata['bounds'][1][0], metadata['bounds'][1][1])\n",
        "\n",
        "  X = [test_data[i][j][0] for j in range(test_data[i].shape[0]) if particle_types_test[j]==1]\n",
        "  Y = [test_data[i][j][1] for j in range(test_data[i].shape[0]) if particle_types_test[j]==1]\n",
        "\n",
        "  X_borders = [test_data[i][j][0] for j in range(test_data[i].shape[0]) if particle_types_test[j]==0]\n",
        "  Y_borders = [test_data[i][j][1] for j in range(test_data[i].shape[0]) if particle_types_test[j]==0]\n",
        "\n",
        "  ax1.scatter(X, Y, color='blue')\n",
        "  ax1.scatter(X_borders, Y_borders, color='red')\n",
        "\n",
        "  with torch.no_grad():\n",
        "    pred_img = cmpn.forward(last_out, torch.BoolTensor(particle_types_test)).cpu()\n",
        "\n",
        "    last_out = torch.cat((pred_img.to(device), torch.tensor(particle_types_test, device=device).view(-1, 1)), dim=1)\n",
        "\n",
        "    ax2 = fig.add_subplot(1, 2, 2)\n",
        "    ax2.set_xlim(metadata['bounds'][0][0], metadata['bounds'][0][1])\n",
        "    ax2.set_ylim(metadata['bounds'][1][0], metadata['bounds'][1][1])\n",
        "\n",
        "    X = [pred_img[j][0].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==1]\n",
        "    Y = [pred_img[j][1].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==1]\n",
        "\n",
        "    X_borders = [pred_img[j][0].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==0]\n",
        "    Y_borders = [pred_img[j][1].item() for j in range(pred_img.shape[0]) if particle_types_test[j]==0]\n",
        "    \n",
        "    ax2.scatter(X, Y, color='blue')\n",
        "    ax2.scatter(X_borders, Y_borders, color='red')\n",
        "\n",
        "    clear_output(wait=True)\n",
        "    plt.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "toc_visible": true,
      "provenance": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}