{
 "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 pandas as pd\n",
    "import repeat as rp\n",
    "import repeat.parametric_study as ps\n",
    "import repeat_onnx as rx\n",
    "from pprint import pp"
   ]
  },
  {
   "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": "5a53a925-553f-4edd-ae0d-4a2ac49fe948",
   "metadata": {},
   "outputs": [],
   "source": [
    "user = rp.User(\n",
    "    \"enter a unique API key from REPEAT user cabinet instead of this comment\"\n",
    ")\n",
    "app = rp.Application(user)\n",
    "# enter a project ID from REPEAT application instead of 12345\n",
    "project = 12345\n",
    "app = rp.Application(user)\n",
    "t_interval = rp.TimeInterval(start = 900, end = 1000)\n",
    "model = app.get_exploration_model(project, t_interval)\n",
    "print(model, model.existing_variables, sep = \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22dda19a-ef77-4b5d-a5d7-0a502c70034a",
   "metadata": {},
   "outputs": [],
   "source": [
    "variable_bounds = dict(\n",
    "        PRESS = [6.5e6, 9.5e6],\n",
    "        TEMP  = [10., 30.],\n",
    "    )\n",
    "input_data = ps.get_latin_hypercube(\n",
    "    model, \n",
    "    variable_bounds,\n",
    "    seed = 42,\n",
    "    sample_size = 800,\n",
    "    strength = 1,\n",
    "    optimization = \"lloyd\"\n",
    ")\n",
    "figure, axis = plt.subplots(dpi = 100)\n",
    "input_data.plot.scatter(x = \"TEMP\", y = \"PRESS\", c = \"tab:green\", marker = \".\", ax = axis)\n",
    "axis.set_xlabel(\"TEMP $[^{\\circ}C]$\")\n",
    "axis.set_ylabel(\"PRESS $[Pa]$\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaa72b41-090f-432a-8d87-dafc84244c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "result_names = [*model_info.potentials, *node_2.flows]\n",
    "design_cases = input_data.copy(deep = True)\n",
    "for param in result_names:\n",
    "    design_cases[param] = pd.NA\n",
    "\n",
    "design_cases.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "009c9309-eebc-4d6a-b485-fc29ebf15f22",
   "metadata": {},
   "outputs": [],
   "source": [
    "def solve(input_data, design_cases):\n",
    "    for case_num, variables in input_data.iterrows():\n",
    "        with model as md:\n",
    "            md.run(variables)\n",
    "            for param in result_names:\n",
    "                design_cases.loc[case_num, param] = md.get_results(param).iloc[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd0456e1-795d-4d66-9b70-5210fc77185d",
   "metadata": {},
   "outputs": [],
   "source": [
    "solve(input_data, design_cases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "274f5893-44ce-4567-8ed1-2462d66f6b35",
   "metadata": {},
   "outputs": [],
   "source": [
    "design_cases[node_1.flows] = - design_cases[node_2.flows]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e383088-f6d8-41b6-b42b-b272dc769c37",
   "metadata": {},
   "outputs": [],
   "source": [
    "design_cases.to_csv(\"gas_flow_data.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32db166c-8734-44a4-93bf-22cf32eae9ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_3d_scatter(table, column = \"flowrate_2\", units = \"$[kg/s]$\", c = \"tab:blue\"):\n",
    "    fig = plt.figure(dpi = 200)\n",
    "    ax = fig.add_subplot(projection = \"3d\")\n",
    "    ax.scatter(\n",
    "        table[\"pressure_2\"], \n",
    "        table[\"temperature_2\"], \n",
    "        table[column], \n",
    "        marker = \".\",\n",
    "        c = c\n",
    "      )\n",
    "    ax.view_init(azim = 240)\n",
    "    ax.set_xlabel(\"pressure_2 $[Pa]$\")\n",
    "    ax.set_ylabel(\"temperature_2 $[^{\\circ}C]$\")\n",
    "    ax.set_zlabel(f\"{column} {units}\")\n",
    "    fig.savefig(f\"results_{column}.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b5eb8aa-5831-4ae6-961e-97b1c5243c40",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_3d_scatter(design_cases)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "677f52a3-683b-460e-87c0-21b4758c3540",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_3d_scatter(design_cases, \"energy_2\", \"$[W]$\", \"tab:orange\")"
   ]
  }
 ],
 "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
}
