{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "28e6aaa4-3704-474d-a257-62ef4d0c052f",
   "metadata": {},
   "source": [
    "# Example: Predict UV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af8b8894-0d32-4158-ac8d-87fad3c7ad93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sunbather\n",
    "sunbather.firstrun(quiet=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1846bd0",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## License\n",
    "\n",
    "The code in this notebook is free to be used, edited and redistributed by anyone free of charge. Please cite Linssen et al. (in prep) when making use of _sunbather_ and/or the code in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3b01a8a",
   "metadata": {},
   "source": [
    "## Example goal\n",
    "\n",
    "In this example notebook, we predict the NUV spectrum of WASP-52 b to see if the signatures of escaping metal species are strong enough to be potentially observable with HST/STIS. In the *fit_helium.ipynb* example notebook, we constrained the mass-loss rate and thermospheric temperature of the planet based on metastable helium observations (the examples can be done in any order). Here, we use those constrained values and make NUV predictions. Since we do not know the (upper) atmospheric metallicity of WASP-52 b, we will explore three different models; one assuming solar metallicity, one assuming 10x solar metallicity, and one assuming 10x solar metallicity but a 100x solar magnesium abundance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8603e0ac",
   "metadata": {},
   "source": [
    "## Example layout\n",
    "\n",
    "The analysis consists of three main steps:\n",
    "\n",
    "- Step 1. Generate Parker wind models at different atmospheric metallicities.\n",
    "- Step 2. Run the Parker wind models through Cloudy to obtain nonisothermal temperature structures and the chemical state of the atmosphere.\n",
    "- Step 3. Make synthetic NUV transit spectra of the different models.\n",
    "\n",
    "Some of the steps we need to take to go through this analysis have to be taken outside of *sunbather* (e.g. storing the planet parameters in the planets.txt file, making the stellar SED available to *Cloudy*). Some steps are part of the \"core functionality\" of the *sunbather* package, which means they can be executed with a single command-line call to a Python script (e.g. generating the Parker wind profiles, running them through the Cloudy algorithm). And some of the steps are performed inside a user-made Python file/Jupyter notebook (e.g. calculating the transmission spectra). In this notebook, we will provide all neccessary code."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52a4caf3",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "\n",
    "For this exercise, we assume you have all codes set-up. That is; you have downloaded *sunbather* and installed its dependencies (the Python packages, including *p-winds*). You have installed *Cloudy* and have the path to it stored as your _\\$CLOUDY_PATH_ environmental variable. You have created your \"project\" folder, and have the path to it stored as your _\\$SUNBATHER_PROJECT_PATH_ environmental variable. You have copied the _planets.txt_ file to the project path. These steps are described in more detail in the \"installation\" section of the _sunbather_ wiki.\n",
    "\n",
    "Before *sunbather* can create Parker wind profiles, we need to make sure the parameters of the system are available to the code. The parameters are stored in the _\\$SUNBATHER_PROJECT_PATH/planets.txt_ file, and the parameters of the WASP-52 b system have already been added. If you want to model additional planets, you can simply add lines in the _planets.txt_ file with their parameters, there is no need to replace previous planet parameters. The last column of the _planets.txt_ file specifies the **name** of the stellar SED that we want to use. The SED with exactly this name must be available to *Cloudy*, so it must be placed in its source folder, specifically: _\\$CLOUDY_PATH/data/SED/_. In the _/sunbather/stellar_SEDs/_ folder, we have provided the SED that we are going to use for WASP-52 b. This is the MUSCLES spectrum (France et al. 2016; Youngblood et al. 2016; Loyd et al. 2016) of eps Eri, which is a similar spectral type to WASP-52. The code is very specific about the format of the spectrum, so we refer to the wiki on how to prep your stellar SED for *sunbather*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f3a65c4-8c80-4dc2-9da8-de8bb27e564c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The only step you need to take here, is to ensure the eps_Eri_binned.spec file\n",
    "# is in Cloudy's SED folder:\n",
    "import os.path\n",
    "cloudypath = sunbather.tools.get_cloudy_path()\n",
    "os.path.exists(f\"{cloudypath}/data/SED/eps_Eri_binned.spec\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b85abb-2118-4b17-8be6-eafe91bad824",
   "metadata": {},
   "source": [
    "## Step 1: Create Parker wind profiles with p-winds/Cloudy\n",
    "\n",
    "This step can be done by calling `sunbather.construct_parker.run` with the proper keyword arguments. Calling `help(sunbather.construct_parker.run)` will give an overview of the available arguments. In this module, the atmospheric composition/metallicity only affects the structure through the mean molecular weight $\\mu$. At solar metallicity, the metal content is low enough to not significantly affect $\\mu$, so we will make the Parker wind profile with *p-winds* standalone, assuming a pure H/He composition of 90/10. At 10x solar metallicity, $\\mu$ may be significantly affected by the metal content, so we will make the Parker wind profile with a hybrid *p-winds*/_Cloudy_ calculation. In Step 2 of this example, we will use this profile also for the run with an enhanced magnesium abundance, since $\\mu$ (and hence the structure we calculate here in step 1) will not change much when increasing the abundance of just one metal species.\n",
    "\n",
    "To make the profile with a 90% hydrogen, 10% helium composition, we pass `fraction_hydrogen=0.9`. We will only make a profile for the constrained temperature of 9200 K and mass-loss rate of $\\dot{M}=10^{11.3}$ g/s, so we pass `temp=9200` and  `mdot=11.3`. Since we explore different compositions, `construct_parker` always expects you to give a folder name `pdir` where we want to store our Parker profiles. We reccommend using a descriptive name, so in this case we will go with *fH_0.9* and the path where our profiles will be saved is then _$SUNBATHER_PROJECT_PATH/parker_profiles/WASP52b/fH_0.9/_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76969515-c0a5-432e-af53-eee9d48934b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The full command to create our first Parker wind model thus becomes - go ahead and run it:\n",
    "import sunbather.construct_parker\n",
    "sunbather.construct_parker.run(\n",
    "    plname=\"WASP52b\", pdir=\"fH_0.9\", temp=9200, mdot=11.3, fraction_hydrogen=0.9,  # overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e94c6ec-3826-48fd-b910-bd13c1f8571d",
   "metadata": {},
   "source": [
    "For this set of parameters, this command should take on the order of 1 second (but it depends on your machine).\n",
    "\n",
    "To make the 10x solar metallicity profile, instead of the `fraction_hydrogen` argument, we use `z=10`. We don't want this profile to overwrite the file we just saved for H/He only, so we provide a different `pdir`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bab02aa-13cf-4485-a6fd-47aca1cdef46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The full command to create our second Parker wind model thus becomes - go ahead and run it:\n",
    "sunbather.construct_parker.run(\n",
    "    plname=\"WASP52b\", pdir=\"z_10\", temp=9200, mdot=11.3, z=10,  # overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2530de2e-176f-4015-a50e-37767a62e074",
   "metadata": {},
   "source": [
    "For this set of parameters, this command should take on the order of 3 minutes (but it depends on your machine). In the _$SUNBATHER_PROJECT_PATH/parker_profiles/WASP52b/_ folder, there should now be 2 subfolders, each containing a .txt file with the isothermal Parker wind structure, feel free to inspect them!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7f792ef-e15b-4374-8672-454fd471f168",
   "metadata": {},
   "source": [
    "## Step 2: Run the Parker wind profiles through Cloudy\n",
    "\n",
    "This step can be done by calling `sunbather.convergeT_parker.run` with the proper arguments The $T_0$ and $\\dot{M}$ commands are the same as in Step 1 of this example. We need to specify a folder name where we want to save our *Cloudy* simulations. For the solar composition we will use `dir=\"z_1\"`. We also need to specify the folder where we want to read the Parker wind profiles from, so `pdir=\"fH_0.9\"`. The last thing we need to think about, is for which atomic/ionic species we want to save *Cloudy's* output. Since many different metal species absorb in the UV, we will save everything that's available, which is the default behavior of the `save_sp` argument (so we do not need to specify it here), but does result in a rather large file size of ~5 MB."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4785b7d6-0386-4abf-88e1-30bc2c95e53e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The command to run our solar composition Parker wind model through *Cloudy* thus becomes\n",
    "# - go ahead and run it:\n",
    "import sunbather.convergeT_parker\n",
    "sunbather.convergeT_parker.run(\n",
    "    plname=\"WASP52b\", workingdir=\"z_1\", pdir=\"fH_0.9\", temp=9200, mdot=11.3, overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a58d9f3-0018-48fb-af55-c59415a53ec2",
   "metadata": {},
   "source": [
    "For the 10x solar metallicity model, we need to make sure we specify `pdir=\"z_10\"` to read the Parker wind profile from the correct folder. We will save the _Cloudy_ runs in `workingdir=\"z_10\"`, and we also need to tell _Cloudy_ to actually use a 10x solar metallicity with `z=10`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aba44b46-fd01-4726-8253-fd171bbfaeb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# **The command to run our 10x solar metallicity Parker wind model through *Cloudy* thus becomes - go ahead and run it**: <br><br>\n",
    "sunbather.convergeT_parker.run(\n",
    "    plname=\"WASP52b\", workingdir=\"z_10\", pdir=\"z_10\", temp=9200, mdot=11.3, z=10, overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1aec36da-a707-4f61-879e-6469ad98f0f2",
   "metadata": {},
   "source": [
    "Finally for the model with 10x solar metallicity, but 100x solar magnesium abundance, we choose a different output folder `workingdir=\"z_10_Mg10\"`, but we can use the same `pdir` as before (see explanation under Step 1 of this example). Since we already have an overall scaling factor of 10 for all the metals, we have to pass `zelem={\"Mg\":10}` to scale magnesium by another factor of 10 to get a 100x solar abundance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea96e73a-aa59-4cb8-9977-2753e6c0a5df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# **The command to run our 10x solar metallicity with enhanced magnesium Parker wind model through *Cloudy* thus becomes - go ahead and run it**: <br><br>\n",
    "sunbather.convergeT_parker.run(\n",
    "    plname=\"WASP52b\", workingdir=\"z_10_Mg10\", pdir=\"z_10\", temp=9200, mdot=11.3,\n",
    "    z=10, zelem={\"Mg\": 10}, save_sp=\"all\", overwrite=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d63ec328-f8f6-4ca6-8b59-b5abb22749d3",
   "metadata": {},
   "source": [
    "For these parameters, all commands together should take on the order of 35 minutes (but it depends on your machine). In the *$SUNBATHER_PROJECT_PATH/sims/1D/WASP52b/fH_0.9/* and */.../z_10/* and */.../z_10_Mg10/* folders, there should now be a sub-folder (or more if you also did the _fit_helium.ipynb_ example notebook), with the output of the _Cloudy_ simulation, feel free to inspect the files! The _converged.png_ file shows the converged temperature structure, and the other _converged.*_ files are the _Cloudy_ output files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84fa611b",
   "metadata": {},
   "source": [
    "## Step 3: Make NUV transit spectra\n",
    "\n",
    "To make transit spectra, we can make use of the `FinFout()` function in the `RT.py` module.\n",
    "\n",
    "The `FinFout()` function takes a `Sim` object, which is a class defined in `tools.py`. For more information on how to use this class, we refer to the wiki. The `FinFout()` function takes three required arguments: the `Sim` object, a wavelength array in angstrom and vacuum, and a list of atomic/ionic species which to include in the calculations. There are a few optional arguments (see wiki), for example to specify a limb-darkening law or transit phase, but we will not use those here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6622346",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from sunbather import tools\n",
    "from sunbather import RT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f9f8e7f-1d2e-4141-a877-fee93db5168d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the project path, and keep this constant during the example\n",
    "projectpath = tools.get_sunbather_project_path()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e014ac29",
   "metadata": {},
   "source": [
    "Let's start by making the spectrum at solar composition. As a heads-up, there will be a few warnings from _sunbather_. This is expected."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e1e0593-e082-455d-b69e-77f659b671fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavs = RT.constantR_wavs(\n",
    "    2000, 3000, 100000\n",
    ")  # set up wavelength grid at spectral resolution of 100,000\n",
    "W52b_solar = tools.Sim(\n",
    "    f\"{projectpath}/sims/1D/WASP52b/z_1/parker_9200_11.300/converged\"\n",
    ")  # load simulation\n",
    "# we can quickly get a list of all species as follows:\n",
    "all_species = tools.get_specieslist(\n",
    "    max_ion=2\n",
    ")  # this includes all species up to doubly ionized - higher is not neccessary here\n",
    "transit_spectrum_solar, _, _ = RT.FinFout(W52b_solar, wavs, all_species)  # do RT\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(12, 7))\n",
    "ax.plot(wavs, transit_spectrum_solar)\n",
    "ax.set_xlabel(\"Wavelength [Å]\")\n",
    "ax.set_ylabel(r\"$F_{in} / F_{out}$\")\n",
    "ax.set_title(r\"WASP-52 b with T=9200 and $\\dot{M}=10^{11.3}$, solar composition\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c334b40",
   "metadata": {},
   "source": [
    "Now let's do 10x solar metallicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d25a07ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "W52b_z_10 = tools.Sim(\n",
    "    projectpath + \"/sims/1D/WASP52b/z_10/parker_9200_11.300/converged\"\n",
    ")  # load simulation\n",
    "transit_spectrum_z_10, _, _ = RT.FinFout(W52b_z_10, wavs, all_species)  # do RT\n",
    "\n",
    "fig, ax = plt.subplots(1, figsize=(12, 7))\n",
    "ax.plot(wavs, transit_spectrum_z_10)\n",
    "ax.set_xlabel(\"Wavelength [Å]\")\n",
    "ax.set_ylabel(r\"$F_{in} / F_{out}$\")\n",
    "ax.set_title(r\"WASP-52 b with T=9200 and $\\dot{M}=10^{11.3}$, 10x solar metallicity\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bea17ca8",
   "metadata": {},
   "source": [
    "And finally a zoom-in on the magnesium doublet for all three models together: solar composition, 10x solar composition, and the model with a 100x solar magnesium abundance. \n",
    "\n",
    "A technical detail that is important here: To save computational time, `RT.FinFout()` function calculates each individual spectral line in a narrow wavelength window (i.e. only a subset of `wavs`) around the rest-frame wavelength of that line. In the absolute majority of cases, this window is large enough to encompass the whole spectral line, but not in special cases where the line is very strong, such as commonly for the Ly-$\\alpha$ line, or here for the Mg II doublet in the 100x solar abundance model. You could visually see this in the resulting spectrum as a discrete \"jump\" from the line wing to the continuum. In such a case, we have to manually increase the computation window by a certain factor with the `width_fac` argument. We will set it to 20 here to catch the full Lorentzian wing of the line (try without it to check the difference!)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c2128fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavs2 = np.linspace(2795, 2805, num=500)  # set up a wav-grid around the Mg II doublet\n",
    "W52b_z_10_Mg10 = tools.Sim(\n",
    "    projectpath + \"/sims/1D/WASP52b/z_10_Mg10/parker_9200_11.300/converged\"\n",
    ")  # load simulation\n",
    "transit_spectrum_z_10_Mg10, _, _ = RT.FinFout(\n",
    "    W52b_z_10_Mg10, wavs2, \"Mg+\", width_fac=20.0\n",
    ")  # do RT\n",
    "\n",
    "fig, ax = plt.subplots(1)\n",
    "ax.plot(wavs, transit_spectrum_solar, label=\"Solar\")\n",
    "ax.plot(wavs, transit_spectrum_z_10, label=\"Z=10\")\n",
    "ax.plot(wavs2, transit_spectrum_z_10_Mg10, label=\"Z=10, Mg=100\")\n",
    "ax.set_xlim(wavs2[0], wavs2[-1])\n",
    "ax.legend(loc=\"best\")\n",
    "ax.set_title(r\"WASP-52 b Mg II doublet at different metallicities\")\n",
    "ax.set_xlabel(\"Wavelength [Å]\")\n",
    "ax.set_ylabel(r\"$F_{in} / F_{out}$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195add5e",
   "metadata": {},
   "source": [
    "To predict the actual observability of these lines in WASP-52 b, we would first need to convolve the spectrum down to the HST/STIS resolution and then do a S/N calculation based on the system magnitude, exposure time, etc. This is outside the scope of this example problem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8385ea04",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "France, K., Loyd, R. O. P., Youngblood, A., et al. 2016, ApJ, 820, 89\n",
    "\n",
    "Loyd, R. O. P., France, K., Youngblood, A., et al. 2016, ApJ, 824, 102\n",
    "\n",
    "Youngblood, A., France, K., Loyd, R. O. P., et al. 2016, ApJ, 824, 101"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}