{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Flexible model creation and molecule formation investigation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook will explore how to create and explore a custom model, and use that to compare the stardis molecular solver to the molecular solver implemented in [Korg](https://ajwheeler.github.io/Korg.jl/stable/). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Basic needed imports\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "\n",
    "from stardis.io.base import parse_config_to_model\n",
    "from stardis.plasma import create_stellar_plasma\n",
    "from astropy import units as u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config, adata, stellar_model = parse_config_to_model('basic_config.yml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stellar_model.temperatures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These temperatures come from the model we read in, which represent the sun. However, for demonstration purposes, let's overwrite these temperatures to go from 3,000 to 10,000 K. This can be done by simply overwriting the object's temperatures attribute. For ease of use, let's match the same number of depth points as the existing one, so we don't have to worry about shape mismatches for other parts of the model. We'll also overwrite the density so we can examine how molecular formation changes purely as a function of temperature."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stellar_model.temperatures = np.linspace(3000, 10000, len(stellar_model.temperatures)) * u.K\n",
    "stellar_model.composition.density = np.ones_like(stellar_model.composition.density) * 3e-7 #This is a reasonable density above the solar photosphere"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stellar_plasma = create_stellar_plasma(stellar_model, adata, config) #Then we go ahead and create the stellar plasma. This will solve the ionization and molecular balance equations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can investigate each of the molecular densities in the plasma. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stellar_plasma.molecule_number_density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "korg_densities = pd.read_csv('korg_comparison_number_densities.csv') #This is a file with output number densities gathered from Korg, using the same densities and temperatures defined here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "korg_densities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 8))\n",
    "\n",
    "ax1 = plt.subplot(211)\n",
    "\n",
    "plt.plot(stellar_model.temperatures, stellar_plasma.molecule_number_density.loc['C2'], label='stardis')\n",
    "plt.plot(korg_densities['T'], korg_densities['C2'], label='Korg')\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('C2 Number Density $[cm^{-3}]$')\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()\n",
    "\n",
    "ax2 = plt.subplot(212, sharex= ax1)\n",
    "plt.plot(stellar_model.temperatures, (stellar_plasma.molecule_number_density.loc['C2'] - korg_densities['C2']) / korg_densities['C2'], label='Fractional Difference')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('C2 Fractional Difference')\n",
    "plt.ylim(1e-4, 1e4)\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 8))\n",
    "\n",
    "ax1 = plt.subplot(211)\n",
    "\n",
    "plt.plot(stellar_model.temperatures, stellar_plasma.molecule_number_density.loc['H2'], label='stardis')\n",
    "plt.plot(korg_densities['T'], korg_densities['H2'], label='Korg')\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('H2 Number Density $[cm^{-3}]$')\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()\n",
    "\n",
    "ax2 = plt.subplot(212, sharex= ax1)\n",
    "plt.plot(stellar_model.temperatures, (stellar_plasma.molecule_number_density.loc['H2'] - korg_densities['H2']) / korg_densities['H2'], label='Fractional Difference')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('H2 Fractional Difference')\n",
    "plt.ylim(1e-4, 1)\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 8))\n",
    "\n",
    "ax1 = plt.subplot(211)\n",
    "\n",
    "plt.plot(stellar_model.temperatures, stellar_plasma.molecule_number_density.loc['H2+'], label='stardis')\n",
    "plt.plot(korg_densities['T'], korg_densities['HHplus'], label='Korg')\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('H2+ Number Density $[cm^{-3}]$')\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()\n",
    "\n",
    "ax2 = plt.subplot(212, sharex= ax1)\n",
    "plt.plot(stellar_model.temperatures, (stellar_plasma.molecule_number_density.loc['H2+'] - korg_densities['HHplus']) / korg_densities['HHplus'], label='Fractional Difference')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('H2+ Fractional Difference')\n",
    "plt.ylim(1e-4, 1)\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 8))\n",
    "\n",
    "ax1 = plt.subplot(211)\n",
    "\n",
    "plt.plot(stellar_model.temperatures, stellar_plasma.molecule_number_density.loc['O2'], label='stardis')\n",
    "plt.plot(korg_densities['T'], korg_densities['O2'], label='Korg')\n",
    "\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('O2 Number Density $[cm^{-3}]$')\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()\n",
    "\n",
    "ax2 = plt.subplot(212, sharex= ax1)\n",
    "plt.plot(stellar_model.temperatures, (stellar_plasma.molecule_number_density.loc['O2'] - korg_densities['O2']) / korg_densities['O2'], label='Fractional Difference')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('T [K]')\n",
    "plt.ylabel('O2 Fractional Difference')\n",
    "plt.ylim(1e-4, 1e4)\n",
    "plt.xlim(4000, 10000)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This was a small comparison of choice molecules between stardis and korg, but the stellar plasma object contains much more information. We could also choose to look at any given ionization or excitation level of each atomic species."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stellar_plasma.ion_number_density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stellar_plasma.level_number_density"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "stardis",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}