{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from parameters_U235_pump_transient import *\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from jitcdde import t\n",
    "from msrDynamics import Node, System\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MSRE = System()\n",
    "\n",
    "# radiator\n",
    "T_out_rc = Node(m = mn_rp, scp = mcp_rpn/mn_rp, W = W_rp, y0 = T0_rp)\n",
    "T_out_air = Node(m = mn_rs, scp = mcp_rsn/mn_rs, W = W_rs, y0 = T0_rs)\n",
    "\n",
    "# heat exchanger\n",
    "T_hf1 = Node(m = mn_p, scp = mcp_pn/mn_p, W = W_p, y0 = T0_p1)\n",
    "T_hf2 = Node(m = mn_p, scp = mcp_pn/mn_p, W = W_p, y0 = T0_p2)\n",
    "T_hf3 = Node(m = mn_p, scp = mcp_pn/mn_p, W = W_p, y0 = T0_p3)\n",
    "T_hf4 = Node(m = mn_p, scp = mcp_pn/mn_p, W = W_p, y0 = T0_p4)\n",
    "T_ht1 = Node(m = m_tn, scp = scp_t, y0 = T0_t1)\n",
    "T_ht2 = Node(m = m_tn, scp = scp_t, y0 = T0_t2)\n",
    "T_hc1 = Node(m = mn_s, scp = mcp_sn/mn_s, W = W_s, y0 = T0_s1)\n",
    "T_hc2 = Node(m = mn_s, scp = mcp_sn/mn_s, W = W_s, y0 = T0_s2)\n",
    "T_hc3 = Node(m = mn_s, scp = mcp_sn/mn_s, W = W_s, y0 = T0_s3)\n",
    "T_hc4 = Node(m = mn_s, scp = mcp_sn/mn_s, W = W_s, y0 = T0_s4)\n",
    "\n",
    "# core \n",
    "n = Node(y0 = n_frac0)\n",
    "C1 = Node(y0 = C0[0])\n",
    "C2 = Node(y0 = C0[1])\n",
    "C3 = Node(y0 = C0[2])\n",
    "C4 = Node(y0 = C0[3])\n",
    "C5 = Node(y0 = C0[4])\n",
    "C6 = Node(y0 = C0[5])\n",
    "rho = Node(y0 = rho_0)\n",
    "\n",
    "T_cg = Node(m = mcp_g1/scp_g, scp = scp_g, y0 = T0_g1)\n",
    "T_cf1 = Node(m = mn_f, scp = scp_f, W = W_f, y0 = T0_f1)\n",
    "T_cf2 = Node(m = mn_f, scp = scp_f, W = W_f, y0 = T0_f1)\n",
    "\n",
    "MSRE.add_nodes([T_out_rc,T_out_air,T_hf1,T_hf2,T_hf3,T_hf4,T_ht1,T_ht2,T_hc1,\n",
    "               T_hc2,T_hc3,T_hc4,n,C1,C2,C3,C4,C5,C6,T_cg,T_cf1,T_cf2,rho])\n",
    "\n",
    "\n",
    "# dynamics \n",
    "\n",
    "# radiator\n",
    "T_out_rc.set_dTdt_advective(source = T_hc4.y(t-tau_hx_r))\n",
    "T_out_rc.set_dTdt_convective(source = [T_out_air.y()], hA = [hA_rpn])\n",
    "\n",
    "T_out_air.set_dTdt_advective(source = Trs_in)\n",
    "T_out_air.set_dTdt_convective(source = [T_out_rc.y()], hA = [hA_rsn])\n",
    "\n",
    "# heat exchanger\n",
    "T_hf1.set_dTdt_advective(source = T_cf2.y(t-tau_c_hx))\n",
    "T_hf1.set_dTdt_convective(source = [T_ht1.y()], hA = [hA_pn])\n",
    "\n",
    "T_hf2.set_dTdt_advective(source = T_hf1.y())\n",
    "T_hf2.dTdt_convective = T_hf1.dTdt_convective\n",
    "\n",
    "T_hf3.set_dTdt_advective(source = T_hf2.y())\n",
    "T_hf3.set_dTdt_convective(source = [T_ht2.y()], hA = [hA_pn])\n",
    "\n",
    "T_hf4.set_dTdt_advective(source = T_hf3.y())\n",
    "# T_hf4.set_dTdt_convective(source = [T_ht2.y()], hA = [hA_pn])\n",
    "T_hf4.dTdt_convective = T_hf3.dTdt_convective\n",
    "\n",
    "# T_ht1.set_dTdt_convective(source = [T_hf1.y(),T_hf2.y(),T_hc3.y(),T_hc4.y()], hA = [hA_pn,hA_pn,hA_sn,hA_sn])\n",
    "# T_ht2.set_dTdt_convective(source = [T_hf3.y(),T_hf4.y(),T_hc1.y(),T_hc2.y()], hA = [hA_pn,hA_pn,hA_sn,hA_sn])\n",
    "T_ht1.set_dTdt_convective(source = [T_hf1.y(),T_hf1.y(),T_hc3.y(),T_hc3.y()], hA = [hA_pn,hA_pn,hA_sn,hA_sn])\n",
    "T_ht2.set_dTdt_convective(source = [T_hf3.y(),T_hf3.y(),T_hc1.y(),T_hc1.y()], hA = [hA_pn,hA_pn,hA_sn,hA_sn])\n",
    "\n",
    "T_hc1.set_dTdt_advective(source = T_out_rc.y(t-tau_r_hx))\n",
    "T_hc1.set_dTdt_convective(source = [T_ht2.y()], hA = [hA_sn])\n",
    "\n",
    "T_hc2.set_dTdt_advective(source = T_hc1.y())\n",
    "T_hc2.dTdt_convective = T_hc1.dTdt_convective\n",
    "\n",
    "T_hc3.set_dTdt_advective(source = T_hc2.y())\n",
    "T_hc3.set_dTdt_convective(source = [T_ht1.y()], hA = [hA_sn])\n",
    "\n",
    "T_hc4.set_dTdt_advective(source = T_hc3.y())\n",
    "T_hc4.dTdt_convective = T_hc3.dTdt_convective\n",
    "\n",
    "# core\n",
    "n.set_dndt(r = rho.y(), beta_eff = beta_t, Lambda = Lam, lam = lam, C = [C1.y(),C2.y(),C3.y(),C4.y(),C5.y(),C6.y()])\n",
    "C1.set_dcdt(n.y(),beta = beta[0],Lambda = Lam,lam = lam[0],t_c=tau_c,t_l = tau_l, flow = True)\n",
    "C2.set_dcdt(n.y(), beta = beta[1],Lambda = Lam,lam = lam[1],t_c=tau_c,t_l = tau_l, flow = True)\n",
    "C3.set_dcdt(n.y(),beta = beta[2],Lambda = Lam,lam = lam[2],t_c=tau_c,t_l = tau_l, flow = True)\n",
    "C4.set_dcdt(n.y(),beta = beta[3],Lambda = Lam,lam = lam[3],t_c=tau_c,t_l = tau_l, flow = True)\n",
    "C5.set_dcdt(n.y(),beta = beta[4],Lambda = Lam,lam = lam[4],t_c=tau_c,t_l = tau_l, flow = True)\n",
    "C6.set_dcdt(n.y(),beta = beta[5],Lambda = Lam,lam = lam[5],t_c=tau_c,t_l = tau_l, flow = True)\n",
    "\n",
    "T_cg.set_dTdt_convective(source = [T_cf1.y()], hA = [hA_fg])\n",
    "T_cg.set_dTdt_internal(source = [n.y()], k = [k_g*P])\n",
    "\n",
    "T_cf1.set_dTdt_advective(source = T_hf4.y(t-tau_hx_c))\n",
    "T_cf1.set_dTdt_convective(source = [T_cg.y()], hA = [k_1*hA_fg])\n",
    "T_cf1.set_dTdt_internal(source = [n.y()], k = [k_f1*P])\n",
    "\n",
    "T_cf2.set_dTdt_advective(source = T_cf1.y())\n",
    "T_cf2.dTdt_convective = T_cf1.dTdt_convective\n",
    "T_cf2.set_dTdt_internal(source = [n.y()], k = [k_f2*P])\n",
    "\n",
    "rho.set_drdt(sources = [T_cf1.dydt, T_cg.dydt], coeffs = [a_f/2,a_g])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finalizing integrator...\n",
      "integrating...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Integration progress:   0%|          | 0/500000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating, compiling, and loading C code.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Integration progress:   0%|          | 1/500000 [00:00<111:03:32,  1.25it/s]/home/luke/miniconda3/envs/onion_dynamics_py312/lib/python3.12/site-packages/jitcdde/_jitcdde.py:820: UserWarning: The target time is smaller than the current time. No integration step will happen. The returned state will be extrapolated from the interpolating Hermite polynomial for the last integration step. You may see this because you try to integrate backwards in time, in which case you did something wrong. You may see this just because your sampling step is small, in which case there is no need to worry (though you should think about increasing your sampling time).\n",
      "  warn(\"The target time is smaller than the current time. No integration step will happen. The returned state will be extrapolated from the interpolating Hermite polynomial for the last integration step. You may see this because you try to integrate backwards in time, in which case you did something wrong. You may see this just because your sampling step is small, in which case there is no need to worry (though you should think about increasing your sampling time).\")\n",
      "Integration progress: 100%|██████████| 500000/500000 [00:01<00:00, 255515.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "populating nodes objects solution vectors...\n"
     ]
    }
   ],
   "source": [
    "sol_jit = MSRE.solve(T, max_delay = 30.0, abs_tol = 1e-14, rel_tol = 1e-12, populate_nodes = True)\n",
    "\n",
    "# T, sol_jit = MSRE.equilibrium_search(dT = 0.01, \n",
    "#                                      max_delay = tau_l, \n",
    "#                                      populate_nodes = False, \n",
    "#                                      md_step = 0.0001, \n",
    "#                                      abs_tol_eq = 1.0e-12, \n",
    "#                                      rel_tol_eq = 1.0e-10,\n",
    "#                                      show_conv_metrics = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig, axs = plt.subplots(3,2, figsize = (8,8))\n",
    "\n",
    "# axs[0,0].set_xlim([0,20])\n",
    "# axs[0,0].plot(T, T_hf4.y_out, label = 'core inlet')\n",
    "# axs[0,0].plot(T, T_cf2.y_out, label = 'core outlet')\n",
    "# axs[0,0].plot(T, T_cg.y_out, label = 'graphite')\n",
    "# axs[0,0].set_ylim([647,649])\n",
    "# axs[0,0].legend()\n",
    "\n",
    "# axs[1,0].set_xlim([0,20])\n",
    "# axs[1,0].plot(T, T_cf2.y_out, label = 'hx fuel inlet')\n",
    "# axs[1,0].plot(T, T_hf4.y_out, label = 'hx fuel outlet')\n",
    "# axs[1,0].plot(T, (T_ht1.y_out+T_ht1.y_out)/2, label = 'hx tubes')\n",
    "# axs[1,0].legend()\n",
    "\n",
    "# axs[2,0].set_xlim([0,20])\n",
    "# axs[2,0].plot(T, T_out_rc.y_out, label = 'hx coolant inlet')\n",
    "# axs[2,0].plot(T, T_hc4.y_out, label = 'hx coolant outlet')\n",
    "# axs[2,0].plot(T, (T_ht1.y_out+T_ht1.y_out)/2, label = 'hx tubes')\n",
    "# axs[2,0].legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x750c0c9332c0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 61828 characters]",
      "text/plain": [
       "<Figure size 640x480 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(3,1)\n",
    "\n",
    "axs[0].set_ylim([0,20])\n",
    "# axs[0].set_xlim([0,100])\n",
    "axs[0].plot(T, P*n.y_out*(1e6), label = 'Power (W)')\n",
    "\n",
    "# axs[1].set_ylim([600,700])\n",
    "# axs[1].set_xlim([0,100])\n",
    "axs[1].plot(T, T_hf4.y_out, label = 'Core Fuel Inlet')\n",
    "axs[1].plot(T, T_cf2.y_out, label = 'Core Fuel Outlet')\n",
    "axs[1].plot(T, T_cg.y_out, label = 'Graphite Temp')\n",
    "axs[1].plot(T, T_hf2.y_out, label = 'Hx Fuel Temp')\n",
    "axs[1].plot(T, T_hc2.y_out, label = 'Hx Coolant Temp')\n",
    "axs[1].legend()\n",
    "\n",
    "axs[2].plot(T, rho.y_out, label = 'feedback')\n",
    "axs[2].legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig,axs = plt.subplots(2,3,figsize=(18,12))\n",
    "\n",
    "# # Set a professional color scheme\n",
    "# colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple', 'tab:brown']\n",
    "\n",
    "# # Function to update the style of each axis\n",
    "# def update_axis_style(ax, title, x_label='', y_label='', x_ticks=True):\n",
    "#     ax.set_xlim([t0,tf])\n",
    "#     ax.set_title(title)\n",
    "#     ax.set_xlabel(x_label)\n",
    "#     ax.set_ylabel(y_label)\n",
    "#     ax.grid(True, which='both', linestyle='--', linewidth=0.5)\n",
    "#     ax.tick_params(axis='x', which='both', bottom=x_ticks, top=False, labelbottom=x_ticks)\n",
    "#     ax.tick_params(axis='y', which='both', left=True, right=False, labelleft=True)\n",
    "\n",
    "# # Applying the updated style to the subplots\n",
    "# # Fuel temperatures\n",
    "# sol_jit = np.array(sol_jit)\n",
    "# update_axis_style(axs[0, 0], \"Fuel Node Temperatures (C)\", y_label=r\"$^\\circ$ C\")\n",
    "# axs[0,0].plot(T,[s[20] for s in sol_jit],label=\"core 1\",color=colors[0]) \n",
    "# axs[0,0].plot(T,[s[21] for s in sol_jit],label=\"core 2\",color=colors[1]) \n",
    "# axs[0,0].plot(T,[s[2] for s in sol_jit],label=\"hx 1\",color=colors[2])  \n",
    "# axs[0,0].plot(T,[s[3] for s in sol_jit],label=\"hx 2\",color=colors[3])\n",
    "# axs[0,0].plot(T,[s[4] for s in sol_jit],label=\"hx 3\",color=colors[4])\n",
    "# axs[0,0].plot(T,[s[5] for s in sol_jit],label=\"hx 4\",color=colors[5]) \n",
    "\n",
    "# # Coolant temperatures\n",
    "# update_axis_style(axs[0, 1], \"Coolant Node Temperatures (C)\", y_label=r\"$^\\circ$ C\")\n",
    "# axs[0, 1].plot(T, sol_jit[:, 6], label=f\"hx 1\", color=colors[0])\n",
    "# axs[0, 1].plot(T, sol_jit[:, 7], label=f\"hx 1\", color=colors[1])\n",
    "# axs[0, 1].plot(T, sol_jit[:, 8], label=f\"hx 1\", color=colors[2])\n",
    "# axs[0, 1].plot(T, sol_jit[:, 9], label=f\"hx 1\", color=colors[3])\n",
    "# axs[0, 1].plot(T, sol_jit[:, 0], label=f\"r 1\", color=colors[4])\n",
    "\n",
    "# # Tube node temperatures\n",
    "# update_axis_style(axs[0, 2], \"Tube Node Temperatures (C)\", y_label=r\"$^\\circ$ C\")\n",
    "# axs[0,2].plot(T,[s[6] for s in sol_jit],label=\"hx 1\",color=colors[0]) \n",
    "# axs[0,2].plot(T,[s[7] for s in sol_jit],label=\"hx 2\",color=colors[1]) \n",
    "\n",
    "# # Precursor concentrations\n",
    "# update_axis_style(axs[1, 2], \"Precursor Concentrations\", x_label=\"t (s)\", y_label=r\"concentration (1/cm$^3$)\")\n",
    "# for i in range(6):\n",
    "#     axs[1, 2].plot(T, sol_jit[:, i+13], label=f\"C{i+1}\", color=colors[i])\n",
    "# axs[1, 2].set_yscale(\"log\")\n",
    "\n",
    "# # Multiplication factor temp\n",
    "# update_axis_style(axs[1, 0], r\"$n$\", x_label=\"t (s)\", y_label=r\"$\\frac{n}{n_0}$\")\n",
    "# axs[1, 0].plot(T, sol_jit[:, 12], label=\"n\", color='tab:blue')\n",
    "\n",
    "# # Reactivity\n",
    "# update_axis_style(axs[1, 1], r\"$\\rho$\", x_label=\"t (s)\")\n",
    "# axs[1, 1].plot(T, sol_jit[:, 22], label=\"n\", color='tab:orange')\n",
    "\n",
    "# # Adding legends\n",
    "# for ax in axs.flat:\n",
    "#     ax.legend()\n",
    "\n",
    "# plt.tight_layout()\n",
    "# plt.show()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "onion_dynamics_py312",
   "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.12.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
