{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "42979432",
   "metadata": {},
   "source": [
    "# Multiple Plates In a Fuel Rod\n",
    "The example below attempts to model a single fuel rod, as alternating flow channels and fuel plates, where all flow channels are identical, and all fuel plates are identical."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4878d03",
   "metadata": {},
   "source": [
    "We start by importing everything needed for the example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89a8da58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib import rc\n",
    "\n",
    "rc(\"font\", family=\"serif\", size=15)\n",
    "rc(\"savefig\", dpi=600)\n",
    "rc(\"figure\", figsize=(8, 6))\n",
    "rc(\"mathtext\", fontset=\"dejavuserif\")\n",
    "rc(\"lines\", linewidth=3)\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "from more_itertools import interleave_longest\n",
    "\n",
    "from stream import State\n",
    "from stream.calculations import ChannelAndContacts, Fuel\n",
    "from stream.calculations.heat_diffusion import Solid\n",
    "from stream.calculations.ideal import HeatExchanger, Pump\n",
    "from stream.calculations.kirchhoff import Junction\n",
    "from stream.composition import FlowGraph, flow_edge, uniform_x_power_shape\n",
    "from stream.composition.mtr_geometry import rod, x_boundaries\n",
    "from stream.jacobians import ALG_jacobian\n",
    "from stream.physical_models.heat_transfer_coefficient import (\n",
    "    regime_dependent_q_scb,\n",
    "    spl_htc,\n",
    "    wall_heat_transfer_coeff,\n",
    ")\n",
    "from stream.physical_models.pressure_drop import (\n",
    "    friction_factor,\n",
    "    pressure_diff,\n",
    "    rectangular_laminar_correction,\n",
    ")\n",
    "from stream.pipe_geometry import EffectivePipe\n",
    "from stream.substances import light_water\n",
    "from stream.units import cm, mm\n",
    "from stream.utilities import cosine_shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bdf9eeb",
   "metadata": {},
   "source": [
    "## Entering data for fuel plate modeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad7a3941",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Spatial discretization of the fuel plate\n",
    "z_N, fuel_N, clad_N = 25, 3, 2\n",
    "\n",
    "# Geometrical dimensions\n",
    "meat_height = 60.4 * cm\n",
    "meat_width = 63 * mm\n",
    "meat_depth = 0.51 * mm\n",
    "clad_depth = 0.38 * mm\n",
    "\n",
    "# Material properties\n",
    "clad = Solid(density=2700, specific_heat=9000, conductivity=2500)  # numbers are approximate\n",
    "fuel = Solid(density=3500, specific_heat=7500, conductivity=1000)  # numbers are approximate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "163576bd",
   "metadata": {},
   "source": [
    "## Initializing the input to the fuel model\n",
    "Modeling a fuel plate requires defining the x and z boundaries, plate materials, plate width, and a matrix that defines the distribution of power in the plate meat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49eefe7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fuel-plate-matrix shape. based on the spatial discretization defined above\n",
    "shape = np.array([z_N, fuel_N + 2 * clad_N])\n",
    "\n",
    "# Mask matrix that defines which cells are meat and which cells are cladding\n",
    "meat = np.zeros(shape, dtype=bool)\n",
    "meat[:, clad_N:-clad_N] = True\n",
    "\n",
    "# Plate materials. Define a solid object from  ``fuel``  and  ``clad``\n",
    "materials = np.empty(shape, dtype=object)\n",
    "materials[meat] = fuel\n",
    "materials[~meat] = clad\n",
    "material = Solid.from_array(materials)\n",
    "\n",
    "# Calculate a power distribution in the plate. The distribution is uniform along the x axis\n",
    "# and follows a cosine shape along the z axis.  ``uniform_x_power_shape``  is used for this purpose\n",
    "plate_power_shape = uniform_x_power_shape(\n",
    "    z_N,\n",
    "    fuel_N,\n",
    "    clad_N,\n",
    "    meat_width,\n",
    "    meat_width,\n",
    "    meat_height,\n",
    "    partial(cosine_shape, ppf=np.pi / 2),\n",
    ")\n",
    "\n",
    "# Just making sure the total power shape sums to 1 as it should\n",
    "print(np.sum(plate_power_shape))\n",
    "\n",
    "# Fuel model requires the x boundaries of the spatial cells\n",
    "x_bounds = x_boundaries(clad_N=clad_N, fuel_N=fuel_N, clad_w=clad_depth, meat_w=meat_depth)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c1c5137",
   "metadata": {},
   "source": [
    "## Defining the fuel plate model\n",
    "Note that `name='Fuel'` is temporary, since an aggregator requires each calculation to have a unique name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b7ae626",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = Fuel(\n",
    "    z_boundaries=np.linspace(0, meat_height, z_N + 1),\n",
    "    x_boundaries=x_bounds,\n",
    "    material=material,\n",
    "    y_length=meat_width,\n",
    "    meat_indices=meat,\n",
    "    power_shape=plate_power_shape,\n",
    "    name=\"Fuel\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50b05bf0",
   "metadata": {},
   "source": [
    "## Entering data for channel modeling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "059beca6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Geometrical dimensions\n",
    "channel_width = 66.6 * mm\n",
    "channel_depth = 2.1 * mm\n",
    "\n",
    "# Laminar and turbulent flow limits, based on Reynold's number\n",
    "re_bounds = laminar_re_max, turbulent_re_min = 2500, 4000"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63935223",
   "metadata": {},
   "source": [
    "## Initializing the input to the channel model\n",
    "Modeling a flow channel requires defining z boundaries, fluid, pipe dimensions, heat transfer coefficient function, and a pressure difference function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f36b09a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Laminar Darcy friction factor correction for a rectangular pipe (based on the pipe's aspect ratio)\n",
    "a_dp = channel_depth / channel_width\n",
    "k_R = rectangular_laminar_correction(a_dp)\n",
    "\n",
    "# Pressure difference function\n",
    "comp_friction = friction_factor(\"regime_dependent\", re_bounds=re_bounds, k_R=k_R)\n",
    "pressure_diff_func = partial(pressure_diff, f=comp_friction)\n",
    "\n",
    "# Wall heat transfer coefficient function\n",
    "comp_h_spl = spl_htc(\"regime_dependent\", re_bounds=re_bounds, aspect_ratio=a_dp, Lh=meat_height)\n",
    "comp_scb = partial(regime_dependent_q_scb, re_bounds=re_bounds)\n",
    "htc = partial(wall_heat_transfer_coeff, h_spl=comp_h_spl, q_scb=comp_scb)\n",
    "\n",
    "# Channel properties are held in an  ``EffectivePipe``  object\n",
    "fluid = light_water\n",
    "pipe = EffectivePipe.rectangular(length=meat_height, edge1=channel_depth, edge2=channel_width, heated_edge=meat_width)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae527240",
   "metadata": {},
   "source": [
    "## Defining the flow channel model\n",
    "Again, `name='Channel'` is temporary, and will be renamed later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c68a59ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "C = ChannelAndContacts(\n",
    "    z_boundaries=np.linspace(0, meat_height, z_N + 1),\n",
    "    fluid=fluid,\n",
    "    pipe=pipe,\n",
    "    h_wall_func=htc,\n",
    "    pressure_func=pressure_diff_func,\n",
    "    name=\"Channel\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749c8938",
   "metadata": {},
   "source": [
    "## Let us define some approximate parameters of IRR-1\n",
    "I used some typical parameters just for the sake of this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76479d03",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_power = 15e6  # Watt\n",
    "fuel_assembly_flow = 15  # m^3/h\n",
    "core_rods = 24\n",
    "plates_per_rod = 7\n",
    "channels_per_rod = plates_per_rod + 1\n",
    "\n",
    "plates_in_core = plates_per_rod * core_rods\n",
    "\n",
    "# Assuming the power is generated evenly across the core\n",
    "power_per_plate = total_power / plates_in_core\n",
    "\n",
    "# mdot is split evenly between all the channels\n",
    "mdot = fuel_assembly_flow * (1000 / 3600) / channels_per_rod  # Kg/s\n",
    "Tin = 35\n",
    "p_abs = 1.67e5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c0a9fa",
   "metadata": {},
   "source": [
    "## Coupling channels and fuel plates using ``rod``\n",
    "``rod`` couples each adjacent fuel plate calculation and flow channel calculation correctly. This just means that the calculations are coupled by  ``T_left , T_right , h_left , h_right``."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d5fb90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ``rod`` conveniently returns the calculation graph, as well as lists of the channels and fuels used\n",
    "calculation_graph, channels, fuels = rod(C, channels_per_rod, F, plates_per_rod)\n",
    "\n",
    "# As said before, each calculation must have a unique name\n",
    "# Also, here we insert initial conditions for all calculations in the form of external funcs\n",
    "calculation_graph.funcs = {}\n",
    "for i, c in enumerate(channels):\n",
    "    c.name = f\"channel#{i + 1}\"\n",
    "    calculation_graph.funcs |= {c: dict(p_abs=p_abs)}\n",
    "for i, f in enumerate(fuels):\n",
    "    f.name = f\"fuel#{i + 1}\"\n",
    "    calculation_graph.funcs |= {f: dict(power=power_per_plate)}\n",
    "\n",
    "# Finally, an aggregator can be initialized from the calculation graph\n",
    "agr = calculation_graph.to_aggregator()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a5edba6",
   "metadata": {},
   "source": [
    "Note that this example can be extended to differing power distributions across the fuel plates, and even differing flow rates across the flow channels."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9379680",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "## Adding kirchhoff calculations with a pump and heat exchanger\n",
    "Another layer for this example is the use of kirchhoff calculations. Here a closed flow loop is created, and the flow rate (``mdot``) is controlled by the pressure of the pump."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab9eeaca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pump and HeatExchanger are ideal calculations.\n",
    "# Pump propagates the same water forward with a given pressure difference\n",
    "# HeatExchanger changes the water temperature without affecting anything else\n",
    "pressure = 1e5\n",
    "P = Pump(pressure=pressure)\n",
    "HX = HeatExchanger(Tin)\n",
    "\n",
    "# Junctions are needed to define ``flow_edge``s and a ``FlowGraph``\n",
    "A = Junction(\"A\")\n",
    "B = Junction(\"B\")\n",
    "\n",
    "# This flow graph describes a system of parallel channels,\n",
    "# connected to a heat exchanger and a pump in series.\n",
    "fg = FlowGraph(*[flow_edge((A, B), c) for c in channels], flow_edge((B, A), P, HX))\n",
    "\n",
    "# Aggregator addition connects all the relevant calculations\n",
    "agr_with_k = agr + fg.aggregator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a3ee5fa",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Report is a nice tool that helps make sure our aggregator has everything it needs\n",
    "from stream.analysis.report import report\n",
    "\n",
    "report(agr_with_k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c255b1fa",
   "metadata": {},
   "source": [
    "Now that we have no variables in the `Missing` column, all that is left to do is supply a guess for the steady state solution and solve.\n",
    "\n",
    "For an initial guess here, we use the analytical solution for a single symmetrical plate and channel system, in every single fuel plate and channel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dedce5e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# All fuel plates and channels are identical, so using the parameters of the first ones is fine for a guess\n",
    "f = fuels[0]\n",
    "c = channels[0]\n",
    "\n",
    "power_mat = np.zeros(f.shape)\n",
    "power_mat[f.meat == 1] = power_per_plate * f.power_shape\n",
    "cp = c.fluid.specific_heat(Tin)\n",
    "p_z = np.sum(power_mat, 1)\n",
    "q2t_z = p_z / (c.pipe.heated_perimeter * c.dz)\n",
    "_tc0 = Tin + np.cumsum(p_z / (np.abs(mdot) * cp))\n",
    "tc0 = _tc0 if mdot >= 0 else _tc0[::-1]\n",
    "tw0 = tc0\n",
    "\n",
    "initial_guess_iterations = 5\n",
    "for _ in range(initial_guess_iterations):\n",
    "    dp0 = c.pressure(T=tc0, Tw=tw0, mdot=mdot)\n",
    "    h0 = c.h_wall(T_wall=tw0, T_cool=tc0, mdot=mdot, pressure=p_abs - dp0)\n",
    "    tw0 = (q2t_z / h0) + tc0\n",
    "T0 = np.tile(tw0, f.shape[1]).reshape(f.shape[::-1]).T.flatten()\n",
    "\n",
    "# Creating a user friendly dictionary of initial guesses for all calculation\n",
    "# To be later converted to an array the aggregator can read by the ``agr.load`` method\n",
    "y0_guess = {}\n",
    "for f in fuels:\n",
    "    y0_guess |= {f.name: dict(T=T0, T_wall_left=tw0, T_wall_right=tw0)}\n",
    "for c in channels:\n",
    "    y0_guess |= {c.name: dict(T_cool=tc0, h_left=h0, h_right=h0, pressure=np.sum(dp0))}\n",
    "\n",
    "# We now need to provide an initial guess for the kirchhoff element of the aggregator.\n",
    "# `` guess_steady_state `` gives a good primitive guess.\n",
    "k_guess = fg.guess_steady_state({c: mdot for c in channels} | {P: mdot * channels_per_rod}, temperature=Tin)\n",
    "\n",
    "# Combining two states can be done with `` State.merge ``\n",
    "y0_with_k = agr_with_k.load(State.merge(k_guess, y0_guess))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "866cfcbe",
   "metadata": {},
   "source": [
    "## Running a steady state calculation\n",
    "Solving for the steady state of the described system above,and saving the result in a user friendly dictionary.\n",
    "\n",
    "Note that the use of `jac=ALG_jacobian(agr)` is meant to accelerate the proccess of calculating a jacobian for each step of root finding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46b6abb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "state = agr_with_k.save(agr_with_k.solve_steady(y0_with_k, jac=ALG_jacobian(agr_with_k), tol=1e-2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bc21e54",
   "metadata": {},
   "source": [
    "## Analysis of the resulting state\n",
    "After solving for the steady state, we can present the solution in a few ways and analyze the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bfd90da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the temperature of each flow channel along the z axis\n",
    "for c in channels:\n",
    "    plt.plot(c.centers, state[c.name][\"T_cool\"], label=c.name)\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel(\"Distance from top [m]\")\n",
    "plt.ylabel(r\"Temperature [$\\degree$C]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a51d192",
   "metadata": {},
   "source": [
    "Note that flow channels are emulated as one dimensional single phase liquid channels, where the heat is transfered via convection only. Since the channels are one dimensional, we assume the liquid temperature varies along the z axis only, meaning we get a one dimensional array of temperatures as seen above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de51e1c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the temperature of each fuel plate along the z axis, in the middle of the plate\n",
    "for f in fuels:\n",
    "    T_mid = state[f.name][\"T\"][:, int(state[f.name][\"T\"].shape[1] / 2)]\n",
    "    plt.plot(channels[0].centers, T_mid, label=f\"{f.name}\")\n",
    "plt.legend()\n",
    "plt.grid()\n",
    "plt.xlabel(\"Distance from top [m]\")\n",
    "plt.ylabel(r\"Temperature [$\\degree$C]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdbca5c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store the temperatures of the channels and plates in a single matrix for presentation\n",
    "total_matrix = np.zeros((len(channels[0].centers), 1 * len(channels) + fuels[0].shape[1] * len(fuels)))\n",
    "\n",
    "i = 0\n",
    "for calc in interleave_longest(channels, fuels):\n",
    "    if isinstance(calc, ChannelAndContacts):\n",
    "        total_matrix[:, i] = state[calc.name][\"T_cool\"]\n",
    "        i += 1\n",
    "    elif isinstance(calc, Fuel):\n",
    "        total_matrix[:, i : i + calc.shape[1]] = state[calc.name][\"T\"]\n",
    "        i += calc.shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29899a06",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot a meshgrid of the temperatures in the core\n",
    "depth = np.array([0])\n",
    "\n",
    "d = 0\n",
    "for calc in interleave_longest(channels, fuels):\n",
    "    if isinstance(calc, ChannelAndContacts):\n",
    "        depth = np.concatenate((depth, [d + channel_depth]))\n",
    "        d += channel_depth\n",
    "    elif isinstance(calc, Fuel):\n",
    "        bounds = x_boundaries(clad_N, fuel_N, clad_depth, meat_depth)\n",
    "        depth = np.concatenate((depth, d + 0.5 * (bounds[1:] + bounds[:-1])))\n",
    "        d += meat_depth + 2 * clad_depth\n",
    "\n",
    "height = np.linspace(0, meat_height, z_N + 1)\n",
    "\n",
    "X, Y = np.meshgrid(depth, height)\n",
    "\n",
    "plt.pcolormesh(X, Y, total_matrix, shading=\"flat\", cmap=\"hot\")\n",
    "plt.colorbar()\n",
    "plt.ylabel(\"Distance from top [m]\")\n",
    "plt.xlabel(\"Distance from left [m]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a293c32c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the hottest point in each fuel plate\n",
    "T_max = np.zeros(len(fuels))\n",
    "for i, f in enumerate(fuels):\n",
    "    T_max[i] = np.max(state[f.name][\"T\"])\n",
    "plt.scatter(range(1, len(fuels) + 1), T_max)\n",
    "plt.xticks(range(1, len(fuels) + 1))\n",
    "plt.grid()\n",
    "plt.xlabel(\"Fuel plate number\")\n",
    "plt.ylabel(r\"Temperature [$\\degree$C]\")\n",
    "plt.show()"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}