{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ab427265",
   "metadata": {},
   "source": [
    "# Running a Grid\n",
    "\n",
    "A common task is to run UCLCHEM over a grid of parameter combinations. This notebook shows how to use the built-in ```GridRunner``` class to doing so for regular grids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39a1752f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f61ea5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac9e035",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7983f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "import uclchem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfb0f513",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "# Ensure output directory exists\n",
    "if not os.path.exists(\"./output_3\"):\n",
    "    os.makedirs(\"./output_3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d8671eb",
   "metadata": {},
   "source": [
    "## A Simple Grid\n",
    "### Define Parameter Space\n",
    "First, we define our parameter space. We do this by using numpy arrays or functions to produce a table of all possible combinations of some parameters of interest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3718d4e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "ParameterDictionary = {\n",
    "    \"param_dict\": {\n",
    "        \"endatfinaldensity\": False,\n",
    "        \"freefall\": False,\n",
    "        \"initialDens\": np.logspace(4, 6, 3),\n",
    "        \"initialTemp\": np.linspace(10, 50, 3),\n",
    "        \"zeta\": np.logspace(1, 3, 3),\n",
    "        \"finalTime\": 1.0e6,\n",
    "        \"baseAv\": 10,\n",
    "        \"abstol_min\": 1e-22,\n",
    "    }\n",
    "}\n",
    "grid = uclchem.model.GridRunner(\n",
    "    model_type=\"Cloud\",\n",
    "    full_parameters=ParameterDictionary,\n",
    "    max_workers=4,\n",
    "    grid_file=\"./output_3/grid_basic.h5\",\n",
    "    model_name_prefix=\"\",\n",
    "    delay_run=True,\n",
    ")\n",
    "print(f\"Total models to run: {len(grid.flat_grids)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a92c7894",
   "metadata": {},
   "source": [
    "# Run Grid\n",
    "\n",
    "Using the ```GridRunner``` class, we can take advantage of setting up our varying parameter space with numpy arrays. We can define the type of model to run, how many concurrent workers should be used, where to store the output models as well as the naming convention of the models in the stored .h5 file. By default, instantiating a ```uclchem.model.GridRunner``` object will call the ```run()``` method, but for the sake of demonstration, we have delayed the run using ```delay_run=True``` in the call to instantiate the object ```grid```. To run the models, we now call the ```run()``` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9cc6dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3cffc02",
   "metadata": {},
   "source": [
    "# Checking Your Grid\n",
    "After running, we can do two things. First, we can validate element conservation using the ```check_conservation``` method. Second, we can check the status of the individual models by inspecting the list of models that were run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebcdc2bc",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "grid.check_conservation()\n",
    "grid.models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b32e0630",
   "metadata": {},
   "source": [
    "Each model is stored in the \"./output_3/grid_basic.h5\" under the name listed in the key 'Model' for each of the entries of ```grid.models```. Each of these models can be loaded individually using the ```load_model``` function. If we wanted to load the model '9' with initialDens=10000, initialTemp=30 and zeta=10 we could do the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "280ea6bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "cloud = uclchem.model.load_model(file=\"./output_3/grid_basic.h5\", name=\"9\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d122a89e",
   "metadata": {},
   "source": [
    "Now ```cloud``` behaves as any model object should, allowing us to perform the same analyses and plotting as done in previous notebooks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5a6823a",
   "metadata": {},
   "source": [
    "# Complex Grid\n",
    "The above was straightforward enough but what about a modeling a grid of shocks? Not only do we want to loop over relevant parameters, we also need to run preliminary models to give ourselves starting abundances. We can do this by taking advantage of the ```SequentialRunner``` class.\n",
    "\n",
    "This class can be used in isolation, in the following way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a9ec88",
   "metadata": {},
   "outputs": [],
   "source": [
    "sequential_model_params: list[dict[str, Any]] = [{\n",
    "    \"Cloud\": {\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",
    "    {\"CShock\": {\n",
    "        \"param_dict\": {\n",
    "            \"endAtFinalDensity\": False,\n",
    "            \"freefall\": False,\n",
    "            \"initialTemp\": 10.0,\n",
    "            \"finalTime\": 1.0e5,\n",
    "            \"abstol_factor\": 1e-14,\n",
    "            \"abstol_min\": 1e-20,\n",
    "            \"reltol\": 1e-6,\n",
    "            \"baseAv\": 1,\n",
    "        },\n",
    "        \"shock_vel\": 10,\n",
    "    }}\n",
    "]\n",
    "\n",
    "SequentialCShock = uclchem.model.SequentialRunner(\n",
    "    sequenced_model_parameters=sequential_model_params,\n",
    "    parameters_to_match=[\"finalDens\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ec640e4",
   "metadata": {},
   "source": [
    "Running this in a grid, works the same way as it does for normal models, except we can pass the additional ```parameters_to_match``` through the ```full_parameters``` dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07db8773",
   "metadata": {},
   "outputs": [],
   "source": [
    "models_to_run = [\n",
    "    {\"Cloud\": {\n",
    "        \"param_dict\": {\n",
    "            \"endAtFinalDensity\": False,  # stop at finalTime\n",
    "            \"freefall\": True,  # increase density in freefall\n",
    "            \"initialDens\": 1e2,  # starting density\n",
    "            \"finalDens\": np.logspace(4, 6, 3),  # 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",
    "    {\"CShock\": {\n",
    "        \"param_dict\": {\n",
    "            \"endAtFinalDensity\": False,\n",
    "            \"freefall\": False,\n",
    "            \"initialTemp\": 10.0,\n",
    "            \"finalTime\": 1.0e5,\n",
    "            \"abstol_factor\": 1e-14,\n",
    "            \"abstol_min\": 1e-20,\n",
    "            \"reltol\": 1e-6,\n",
    "            \"baseAv\": 1,\n",
    "        },\n",
    "        \"shock_vel\": np.linspace(10, 50, 3),\n",
    "    }},\n",
    "]\n",
    "\n",
    "complex_grid = uclchem.model.GridRunner(\n",
    "    model_type=\"SequentialRunner\",\n",
    "    full_parameters=models_to_run,\n",
    "    max_workers=10,\n",
    "    grid_file=\"./output_3/grid_complex.h5\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed54d46a",
   "metadata": {},
   "outputs": [],
   "source": [
    "complex_grid.models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cb32dd5",
   "metadata": {},
   "source": [
    "In the case of SequentialRunners being run in a grid, each individual model is saved using the naming convention of \"<Model name>_<Model Class>_<Order in Sequence>\" so the 6th model in our grid would consist of both \"6_Cloud_0\" and \"6_CShock_1\"."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75e0ff68",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "There are many ways to run grids of models and users will naturally develop their own methods. This notebook is just a simple example of how to run UCLCHEM for many parameter combinations whilst producing a useful output (the model_table) to keep track of all the combinations that were run. In a real script, we'd save the model file to csv at the end.\n",
    "\n",
    "For much larger grids, it's recommended that you find a way to make your script robust to failure. Over a huge range of parameters, it is quite likely UCLCHEM will hit integration trouble for at least a few parameter combinations. Very occasionally, UCLCHEM will get caught in a loop where it fails to integrate and cannot adjust its strategy to manage it. This isn't a problem for small grids as the model can be stopped and the tolerances adjusted. However, for very large grids, you may end up locking all threads as they each get stuck on a different model. The best solution we've found for this case is to add a check so that models in your dataframe are skipped if their file already exists, this allows you to stop and restart the grid script as needed.\n"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:light"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}