{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "94f05818",
   "metadata": {},
   "source": [
    "# Saving and Loading Model Objects\n",
    "\n",
    "In this tutorial we will repeat the models of 2b, but this time we will leave all the plotting for the end. Instead, we will store the models into a `.h5` file, delete the objects, and in the end reload them so we can perform the same analysis as was done in 2b. For details on what models are being created, we refer you back to tutorial 2b, here we will only discuss the technical components of saving and loading models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e5fdae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uclchem\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3753ed85",
   "metadata": {},
   "source": [
    "# The Prestellar Core\n",
    "We will run both the initial conditions (Phase 1) and the science model (Phase 2) in one cell, storing the p_core object into the file `\"../examples/test-models/models.h5\"` with a name of `\"prestellar_core\"` so we can easily recall it later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc4d2cbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_file = Path(\"output_2c/models.h5\")\n",
    "\n",
    "if not os.path.exists(save_file.parent):\n",
    "    os.makedirs(save_file.parent)\n",
    "else:\n",
    "    # Remove existing file to avoid appending to old data\n",
    "    if os.path.exists(save_file):\n",
    "        os.remove(save_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d17ef455",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set a parameter dictionary for cloud collapse model\n",
    "param_dict = {\n",
    "    \"endAtFinalDensity\": False,  # stop at finalTime\n",
    "    \"freefall\": True,  # increase density in freefall\n",
    "    \"initialDens\": 1e2,  # starting density\n",
    "    \"finalDens\": 1e6,  # final density\n",
    "    \"initialTemp\": 10.0,  # temperature of gas\n",
    "    \"finalTime\": 6.0e6,  # final time\n",
    "    \"rout\": 0.1,  # radius of cloud in pc\n",
    "    \"baseAv\": 1.0,  # visual extinction at cloud edge.\n",
    "}\n",
    "\n",
    "cloud = uclchem.model.Cloud(param_dict=param_dict)\n",
    "cloud.save_model(file=save_file, name=\"cloud\", overwrite=True)\n",
    "\n",
    "# change other bits of input to set up phase 2\n",
    "param_dict[\"initialDens\"] = 1e6\n",
    "param_dict[\"finalTime\"] = 1e6\n",
    "param_dict[\"freefall\"] = False\n",
    "\n",
    "# freeze out is completely overwhelmed by thermal desorption\n",
    "# so turning it off has no effect on abundances but speeds up integrator.\n",
    "param_dict[\"freezeFactor\"] = 0.0\n",
    "\n",
    "param_dict[\"abstol_factor\"] = 1e-18\n",
    "param_dict[\"reltol\"] = 1e-12\n",
    "\n",
    "p_core = uclchem.model.PrestellarCore(\n",
    "    temp_indx=3, max_temperature=300.0, param_dict=param_dict, previous_model=cloud\n",
    ")\n",
    "p_core.save_model(file=save_file, name=\"prestellar_core\", overwrite=True)\n",
    "\n",
    "del cloud\n",
    "del p_core"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5afe1249",
   "metadata": {},
   "source": [
    "We will delete the cloud and p_core object to demonstrate loading models later when we plot the results of p_core.\n",
    "\n",
    "# Shocks\n",
    "Now let's run the shock models, again, for the science explanations, we refer to tutorial 2b, instead here we simply run the models and save them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "274a8fb9",
   "metadata": {},
   "outputs": [],
   "source": [
    "param_dict = {\n",
    "    \"endAtFinalDensity\": False,  # stop at finalTime\n",
    "    \"freefall\": True,  # increase density in freefall\n",
    "    \"initialDens\": 1e2,  # starting density\n",
    "    \"finalDens\": 1e4,  # final density\n",
    "    \"initialTemp\": 10.0,  # temperature of gas\n",
    "    \"finalTime\": 6.0e6,  # final time\n",
    "    \"rout\": 0.1,  # radius of cloud in pc\n",
    "    \"baseAv\": 1.0,  # visual extinction at cloud edge.\n",
    "    \"abundSaveFile\": \"../examples/test-output/shockstart.dat\",\n",
    "}\n",
    "\n",
    "shock_start = uclchem.model.Cloud(param_dict=param_dict)\n",
    "shock_start.save_model(file=save_file, name=\"shock_start\", overwrite=True)\n",
    "\n",
    "param_dict[\"initialDens\"] = 1e4\n",
    "param_dict[\"finalTime\"] = 1e6\n",
    "\n",
    "cshock = uclchem.model.CShock(\n",
    "    shock_vel=40, param_dict=param_dict, previous_model=shock_start\n",
    ")\n",
    "cshock.save_model(file=save_file, name=\"cshock\", overwrite=True)\n",
    "\n",
    "param_dict[\"initialDens\"] = 1e3\n",
    "param_dict[\"freefall\"] = False  # lets remember to turn it off this time\n",
    "param_dict[\"reltol\"] = 1e-12\n",
    "\n",
    "jshock = uclchem.model.JShock(\n",
    "    shock_vel=10.0, param_dict=param_dict, previous_model=shock_start, timepoints=1500\n",
    ")\n",
    "jshock.save_model(file=save_file, name=\"jshock\", overwrite=True)\n",
    "\n",
    "del shock_start\n",
    "del cshock\n",
    "del jshock"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95dc690d",
   "metadata": {},
   "source": [
    "Having run all models and deleted the objects from the notebook, we can now show how to load models, before showcasing that the loaded objects perform exactly as regular versions of the object.\n",
    "\n",
    "# Load the Prestellar Core\n",
    "\n",
    "If this notebook has been run on your machine before, and the `\"../examples/test-models/models.h5\"` still contains the models from this tutorial, the notebook could be started from this point onward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aebbf08a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uclchem\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "save_file = \"../../examples/test-models/models.h5\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e969fb2c",
   "metadata": {},
   "source": [
    "We begin by loading the prestellar core model. We will check the conservation of elements, as well as the success flag and plot it as was done in tutorial 2b."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c0beda8",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_p_core = uclchem.model.load_model(\n",
    "    file=save_file, name=\"prestellar_core\"\n",
    ")\n",
    "print(f\"Success Flag = {loaded_p_core.success_flag}\")\n",
    "loaded_p_core.check_conservation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b86c7f11",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_p_core = loaded_p_core.get_dataframes()\n",
    "species = [\"CO\", \"H2O\", \"CH3OH\", \"#CO\", \"#H2O\", \"#CH3OH\", \"@H2O\", \"@CO\", \"@CH3OH\"]\n",
    "fig, [ax, ax2] = plt.subplots(1, 2, figsize=(16, 9))\n",
    "ax = loaded_p_core.plot_species(ax, species)\n",
    "settings = ax.set(\n",
    "    yscale=\"log\",\n",
    "    xlim=(1e2, 1e6),\n",
    "    ylim=(1e-10, 1e-2),\n",
    "    xlabel=\"Time / years\",\n",
    "    ylabel=\"Fractional Abundance\",\n",
    "    xscale=\"log\",\n",
    ")\n",
    "\n",
    "ax2.plot(df_p_core[\"Time\"], df_p_core[\"Density\"], color=\"black\")\n",
    "ax2.set(xscale=\"log\")\n",
    "ax3 = ax2.twinx()\n",
    "ax3.plot(df_p_core[\"Time\"], df_p_core[\"gasTemp\"], color=\"red\")\n",
    "ax2.set(xlabel=\"Time / year\", ylabel=\"Density\")\n",
    "ax3.set(ylabel=\"Temperature\", facecolor=\"red\", xlim=(1e2, 1e6))\n",
    "ax3.tick_params(axis=\"y\", colors=\"red\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2237b87",
   "metadata": {},
   "source": [
    "The way in which the loaded models are used, is equivalent to run models, and the output matches that of tutorial 2b. With that, let us look at the shock models next.\n",
    "\n",
    "# Loading Shocks\n",
    "\n",
    "We start again with checking the success flag, and the element conservation check for each model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b7b19f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_cshock = uclchem.model.load_model(file=save_file, name=\"cshock\")\n",
    "print(f\"Success Flag for cshock = {loaded_cshock.success_flag}\")\n",
    "print(\"cshock conservation check:\")\n",
    "loaded_cshock.check_conservation()\n",
    "\n",
    "loaded_jshock = uclchem.model.load_model(file=save_file, name=\"jshock\")\n",
    "print(f\"Success Flag for jshock = {loaded_cshock.success_flag}\")\n",
    "print(\"jshock conservation check:\")\n",
    "loaded_jshock.check_conservation()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e3554df",
   "metadata": {},
   "source": [
    "Now we can plot the cshock, just as was done in tutorial 2b."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea17df29",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_cshock = loaded_cshock.get_dataframes()\n",
    "species = [\"CO\", \"H2O\", \"CH3OH\", \"NH3\", \"$CO\", \"$H2O\", \"$CH3OH\", \"$NH3\"]\n",
    "\n",
    "fig, [ax, ax2] = plt.subplots(1, 2, figsize=(16, 9))\n",
    "ax = loaded_cshock.plot_species(ax, species)\n",
    "settings = ax.set(\n",
    "    yscale=\"log\",\n",
    "    xlim=(1, 20 * loaded_cshock.dissipation_time),\n",
    "    ylim=(1e-10, 1e-2),\n",
    "    xlabel=\"Time / years\",\n",
    "    ylabel=\"Fractional Abundance\",\n",
    "    xscale=\"log\",\n",
    ")\n",
    "\n",
    "ax2.plot(df_cshock[\"Time\"], df_cshock[\"Density\"], color=\"black\")\n",
    "ax2.set(xscale=\"log\")\n",
    "ax3 = ax2.twinx()\n",
    "ax3.plot(df_cshock[\"Time\"], df_cshock[\"gasTemp\"], color=\"red\")\n",
    "ax2.set(xlabel=\"Time / year\", ylabel=\"Density\")\n",
    "ax3.set(\n",
    "    ylabel=\"Temperature\", facecolor=\"red\", xlim=(1, 20 * loaded_cshock.dissipation_time)\n",
    ")\n",
    "ax3.tick_params(axis=\"y\", colors=\"red\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10a20ec7",
   "metadata": {},
   "source": [
    "Followed by the plot for the jshock."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb7656e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_jshock = loaded_jshock.get_dataframes()\n",
    "species = [\"CO\", \"H2O\", \"CH3OH\", \"NH3\", \"$CO\", \"$H2O\", \"$CH3OH\", \"$NH3\"]\n",
    "\n",
    "fig, [ax, ax2] = plt.subplots(1, 2, figsize=(16, 9))\n",
    "ax = loaded_jshock.plot_species(ax, species)\n",
    "settings = ax.set(\n",
    "    yscale=\"log\",\n",
    "    xlim=(1e-7, 1e6),\n",
    "    ylim=(1e-10, 1e-2),\n",
    "    xlabel=\"Time / years\",\n",
    "    ylabel=\"Fractional Abundance\",\n",
    "    xscale=\"log\",\n",
    ")\n",
    "\n",
    "ax2.plot(df_jshock[\"Time\"], df_jshock[\"Density\"], color=\"black\")\n",
    "ax2.set(xscale=\"log\", yscale=\"log\")\n",
    "ax3 = ax2.twinx()\n",
    "ax3.plot(df_jshock[\"Time\"], df_jshock[\"gasTemp\"], color=\"red\")\n",
    "ax2.set(xlabel=\"Time / year\", ylabel=\"Density\")\n",
    "ax3.set(ylabel=\"Temperature\", facecolor=\"red\", xlim=(1e-7, 1e6))\n",
    "ax3.tick_params(axis=\"y\", colors=\"red\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7ad514d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent",
   "main_language": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}