{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6173c502-e4f2-4eef-a993-ea532ef8415c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import onnx\n",
    "import pandas as pd\n",
    "import repeat_onnx as rx\n",
    "import torch\n",
    "from pprint import pp\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from torch import nn \n",
    "from torch import optim\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ba5f97a-f768-452a-8cee-9ca31a2b0121",
   "metadata": {},
   "outputs": [],
   "source": [
    "node_1 = rx.NodeP(\n",
    "    \"Input_unit\", \"pressure_1\", \"temperature_1\", \"flowrate_1\", \"energy_1\"\n",
    ")\n",
    "node_2 = rx.NodeP(\n",
    "    \"Output_unit\", \"pressure_2\", \"temperature_2\",\"flowrate_2\", \"energy_2\"\n",
    ")\n",
    "model_info = rx.ModelP([node_1, node_2])\n",
    "pp(model_info)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef4f4c35-5804-4d7f-bee5-005067880aad",
   "metadata": {},
   "outputs": [],
   "source": [
    "TENSOR_SETTINGS = {\n",
    "    \"device\" : \"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
    "    \"dtype\" :  torch.double,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4ebbd50-12ed-4d91-9015-daca847d04b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_tensor(data)->torch.tensor:\n",
    "    return torch.tensor(data, **TENSOR_SETTINGS)\n",
    "\n",
    "def to_column_tensor(data):\n",
    "    return to_tensor(data).reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4578da8-3ac8-41b9-995d-cdf5e58b7951",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_torch_tensors(tables : list[pd.DataFrame]):\n",
    "    total = []\n",
    "    for table in tables:\n",
    "        tensors = {}\n",
    "        for key, values in table.items():\n",
    "            tensors[key] = to_column_tensor(values.to_numpy())\n",
    "        \n",
    "        total.append(tensors)\n",
    "    \n",
    "    return total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6db36483-a1b5-4bd9-8faa-deac415b4fbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential_train, potential_test, flow_train, flow_test = train_test_split(\n",
    "  design_cases[model_info.potentials], design_cases[model_info.flows], \n",
    "  test_size = 0.3, random_state = 42\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6beec320-78f9-4ccc-8f85-0e1435e712c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential_train_tensor, potential_test_tensor = get_torch_tensors(\n",
    "  (potential_train, potential_test)\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d5a6e0-6848-4205-8977-51a9025f356a",
   "metadata": {},
   "outputs": [],
   "source": [
    "potential_scaler = StandardScaler()\n",
    "potential_scaler.fit(potential_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2518eea-b0d8-4979-8c9c-c7188526b178",
   "metadata": {},
   "outputs": [],
   "source": [
    "flow_scaler = StandardScaler()\n",
    "flow_scaler.fit(flow_train)\n",
    "flow_train_scaled = to_tensor(flow_scaler.transform(flow_train))\n",
    "flow_test_scaled = to_tensor(flow_scaler.transform(flow_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "406ba7c2-81c1-4395-9cfd-94efe535dd99",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary = dict(\n",
    "    potential_mean = to_tensor(potential_scaler.mean_),\n",
    "    potential_scale = to_tensor(potential_scaler.scale_),\n",
    "    flow_mean = to_tensor(flow_scaler.mean_),\n",
    "    flow_scale = to_tensor(flow_scaler.scale_),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "980c1939-d54a-4fe4-94f3-16d435998bdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ModelOrifice(nn.Module):\n",
    "    def __init__(self, potential_mean, potential_scale, flow_mean, flow_scale,\n",
    "                 hidden_size = 128):\n",
    "        super().__init__()\n",
    "        self.common_part = nn.Sequential(\n",
    "            nn.Linear(4, hidden_size),\n",
    "            nn.LayerNorm(hidden_size),\n",
    "            nn.Tanh(),\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.LayerNorm(hidden_size),\n",
    "            nn.Tanh()\n",
    "        )\n",
    "        self.flowrate_part = nn.Linear(hidden_size, 1, bias = False)\n",
    "        self.energy_part = nn.Linear(hidden_size, 1, bias = False)\n",
    "        self.potential_mean = potential_mean\n",
    "        self.potential_scale = potential_scale\n",
    "        self.flow_mean = flow_mean\n",
    "        self.flow_scale = flow_scale\n",
    "    \n",
    "    def forward(self, pressure_1, temperature_1, pressure_2, temperature_2):\n",
    "        x = torch.cat([pressure_1, temperature_1, pressure_2, temperature_2], \n",
    "                      dim = 1)\n",
    "        x_scaled = (x - self.potential_mean) / self.potential_scale\n",
    "        common_results = self.common_part(x_scaled)\n",
    "        flowrate_2 = self.flowrate_part(common_results)\n",
    "        energy_2 = self.energy_part(common_results)\n",
    "        flowrate_1 = - flowrate_2\n",
    "        energy_1 = - energy_2\n",
    "        flows = torch.cat([flowrate_1, energy_1, flowrate_2, energy_2], \n",
    "                          dim = 1)\n",
    "        if self.training: \n",
    "            return flows\n",
    "        \n",
    "        return self.postprocessing(flows)\n",
    "    \n",
    "    def postprocessing(self, flows):\n",
    "        orig_flows = flows * self.flow_scale + self.flow_mean\n",
    "        columns = range(flows.shape[1])\n",
    "        \n",
    "        return [self.get_column_vector(orig_flows, i) for i in columns]\n",
    "\n",
    "    def get_column_vector(self, tensor, column):\n",
    "        return tensor[:, column].reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d02f90da-a8d7-47c7-8ab1-13bca68b47c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "model = ModelOrifice(hidden_size = 180, **summary).double().to(TENSOR_SETTINGS[\"device\"])\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.AdamW(model.parameters(), lr = 1e-3, weight_decay = 1e-3)\n",
    "scheduler = ReduceLROnPlateau(optimizer, \"min\", factor = 0.5, patience = 50) \n",
    "num_epochs = 3000\n",
    "train_losses, test_losses = [], []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8ba537e-b087-4c78-ac72-60ea3fab7ee2",
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in tqdm(range(num_epochs)):\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "    outputs_train = model(**potential_train_tensor)\n",
    "    loss_train = criterion(outputs_train, flow_train_scaled)\n",
    "    loss_train.backward()\n",
    "    optimizer.step()\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        model.train() \n",
    "        outputs_test = model(**potential_test_tensor)\n",
    "        model.eval()\n",
    "        loss_test = criterion(outputs_test, flow_test_scaled)\n",
    "    \n",
    "    scheduler.step(loss_test)\n",
    "    train_losses.append(loss_train.item())\n",
    "    test_losses.append(loss_test.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcf9f288-0e2b-4978-aa94-0d5811230aa5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.plot(train_losses, c = \"tab:blue\")\n",
    "plt.plot(test_losses, c = \"tab:orange\")\n",
    "plt.xlabel('epoch')\n",
    "plt.ylabel('MSE Loss')\n",
    "plt.yscale('log')\n",
    "plt.grid(True)\n",
    "plt.savefig(\"ml_losses.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68449154-9639-4912-aba9-bf27465f36e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "onnx_file = \"model_orifice.onnx\"\n",
    "model.eval()\n",
    "torch.onnx.export(\n",
    "    model,\n",
    "    args = (to_column_tensor(6456249.5),\n",
    "            to_column_tensor(19.912780),\n",
    "            to_column_tensor(5065394.0),\n",
    "            to_column_tensor(16.410051),\n",
    "            ),\n",
    "    f = onnx_file,\n",
    "    opset_version = 17,\n",
    "    input_names = model_info.potentials,\n",
    "    output_names = model_info.flows,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fbab66c-3c62-4d78-a4db-617ba5afd779",
   "metadata": {},
   "outputs": [],
   "source": [
    "properties = onnx.StringStringEntryProto()\n",
    "properties.key = \"repeat_contract\"\n",
    "properties.value = model_info.to_onnx()\n",
    "model_onnx = onnx.load(onnx_file)\n",
    "model_onnx.metadata_props.append(properties)\n",
    "onnx.save(model_onnx, onnx_file)"
   ]
  }
 ],
 "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": 5
}