{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e76dc5bf-454d-4699-bb92-b5868cea7d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install repeat torch scikit-learn pandas onnx matplotlib tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d54552be-8524-4bcf-bdda-1abf52c4b374",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tqdm.notebook import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import onnx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c2e396f-704c-44b7-b964-b426c0e14475",
   "metadata": {},
   "outputs": [],
   "source": [
    "design_cases = pd.read_csv(\"heat_transfer_data.csv\", index_col=0)\n",
    "design_cases.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1d2f6af-e52d-4731-8940-b6b2d86aa2ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = design_cases[[\"TEMPERATURE_1\", \"TEMPERATURE_2\"]]\n",
    "heat_flux = design_cases[[\"heat_flux_1\", \"heat_flux_2\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7e8eccc-7019-441d-b524-99dde002253c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "axis = input_data.plot.scatter(x = 'TEMPERATURE_1', y = 'TEMPERATURE_2', c = 'tab:red')\n",
    "axis.set_xlabel(\"TEMPERATURE_1 $[^{\\circ}C]$\")\n",
    "axis.set_ylabel(\"TEMPERATURE_2 $[^{\\circ}C]$\")\n",
    "fig = axis.get_figure()\n",
    "fig.savefig(\"input_data.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0213238e-e9a1-4dd6-a7ac-912ec3b51b15",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_torch_tensors(tables : pd.DataFrame):\n",
    "    total = []\n",
    "    for table in tables:\n",
    "        tensors = {}\n",
    "        for key, values in table.items():\n",
    "            tensors[key] = torch.tensor(values.to_numpy(), dtype = torch.double).reshape(-1, 1)\n",
    "        \n",
    "        total.append(tensors)    \n",
    "        \n",
    "    return total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9129c7d-b6c5-4c60-9ce7-f4acea1f7260",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "input_data_train, input_data_test, heat_flux_train, heat_flux_test = train_test_split(\n",
    "    input_data, heat_flux, test_size = 0.33, random_state = 42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7861734c-9fc9-45a9-973c-68130ecefc97",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary = input_data_train.describe()\n",
    "temperature_mean = torch.tensor(summary.loc[\"mean\"].to_numpy(), dtype = torch.double)\n",
    "temperature_std = torch.tensor(summary.loc[\"std\"].to_numpy(), dtype = torch.double)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c96963-8eed-407f-a719-998ca60fcbbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data_train, input_data_test, heat_flux_train, heat_flux_test = get_torch_tensors((\n",
    "    input_data_train, input_data_test, heat_flux_train, heat_flux_test\n",
    "))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1383cf2-af6e-4293-ab56-640cbcc09910",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearModel(nn.Module):\n",
    "    def __init__(self, temperature_mean, temperature_std):\n",
    "        super().__init__()\n",
    "        self.linear = nn.Linear(\n",
    "            in_features = 2, \n",
    "            out_features = 1 \n",
    "        )\n",
    "        self.temperature_mean = temperature_mean\n",
    "        self.temperature_std = temperature_std\n",
    "        \n",
    "    def forward(self, TEMPERATURE_1, TEMPERATURE_2):\n",
    "        x = torch.cat([TEMPERATURE_1, TEMPERATURE_2], dim = 1)\n",
    "        x = (x - self.temperature_mean) / self.temperature_std\n",
    "        heat_flux_2 = self.linear(x)\n",
    "        heat_flux_1 = - heat_flux_2\n",
    "        return heat_flux_1, heat_flux_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67278028-b057-4bb6-af78-eae11c2dbdcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(42)\n",
    "model = LinearModel(temperature_mean, temperature_std).double()\n",
    "criterion = nn.MSELoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr = 0.1)\n",
    "num_epochs = 100\n",
    "model.train()\n",
    "for epoch in tqdm(range(num_epochs)):\n",
    "    optimizer.zero_grad()\n",
    "    outputs = model(**input_data_train)\n",
    "    loss = criterion(outputs[1], heat_flux_train[\"heat_flux_2\"])\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f2a96e-23cf-4cce-a3cc-ecf181228bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    predictions = model(**input_data_test)\n",
    "    mse = criterion(predictions[1], heat_flux_test[\"heat_flux_2\"])\n",
    "    print(f\"Mean Squared Error (MSE): {mse.item():.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98237366-a177-43b8-9403-7b455d4d21df",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.onnx.export(\n",
    "    model,\n",
    "    args = (\n",
    "        torch.tensor([[5]], dtype = torch.double), \n",
    "        torch.tensor([[-15]], dtype = torch.double)\n",
    "    ),\n",
    "    f = \"ml_heat_transfer.onnx\",\n",
    "    input_names = [\"TEMPERATURE_1\", \"TEMPERATURE_2\"],\n",
    "    output_names = [\"heat_flux_1\", \"heat_flux_2\"],    \n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdd41e93-1e43-4970-bb60-56b5049d8a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.plot(input_data_test[\"TEMPERATURE_1\"], predictions[1], 'ro', markersize=5, label = \"Predicted\")\n",
    "ax.plot(input_data_test[\"TEMPERATURE_1\"], heat_flux_test[\"heat_flux_2\"], 'b+', markersize=10, label = \"Test\")\n",
    "ax.legend()"
   ]
  }
 ],
 "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
}