{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# An example for custom reference states from MESA\n",
    "\n",
    "This script will take a MESA stellar evolution profile and convert it into a format that can be read in as a custom reference state in Rayleigh. You will need the `rayleigh_diagnostics.py`, `reference_tools.py`, and `mesa.py` files. You will also need a suitable MESA profile file, such as `profile_mesa.data`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.interpolate as spi\n",
    "import scipy.integrate as spint\n",
    "import scipy.signal as spsig\n",
    "import sys, os\n",
    "\n",
    "sys.path.insert(0, os.path.abspath('../../'))\n",
    "\n",
    "import post_processing.rayleigh_diagnostics as rd\n",
    "import post_processing.reference_tools as rt\n",
    "import mesa\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def interp(r, v):\n",
    "    prad = p.rmid[::-1] * mesa.rsol\n",
    "    #You can also use 10**p.logR[::-1] or p.radius[::-1] instead of rmid[::-1], but rmid is the most accurate choice\n",
    "    return np.interp(r, prad, v[::-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the `work_dir` variable to the location of the Python files listed above and MESA profile you would like to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "work_dir = './'\n",
    "sys.path.append(work_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = mesa.profile('profile_mesa.data')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Choose a suitable number of radial grid points. They do not need to be regularly spaced. You should err on the side of high resolution since Rayleigh's Chebyshev domains have very fine grid spacing at the top and bottom of the domain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nr = 5000\n",
    "r0 = 5.1e10 # in cm\n",
    "r1 = 6.8e10 # in cm\n",
    "radius = np.linspace(r0, r1, nr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the MESA model, Rayleigh will need the density, buoyancy function $\\rho g / C_P$, temperature, vicoscity, thermal diffusion, electrical resistivity (for magnetic cases), heating profile (for cases with $Q \\ne 0$), entropy gradient (for cases with reference state advection). Note that MESA radial indicies start at the bottom, while Rayleigh radial indicies start at the top."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_MESA = p.rmid*mesa.rsol\n",
    "density = interp(radius, 10**p.logRho)\n",
    "temperature = interp(radius, 10**p.logT)\n",
    "grav = interp(radius, p.grav)\n",
    "cp = interp(radius, p.cp)\n",
    "buoy = density * grav / cp\n",
    "nu = 1e14 * np.ones_like(radius)\n",
    "kappa = 1e14 * np.ones_like(radius)\n",
    "eta = 1e14 * np.ones_like(radius)\n",
    "hprofile = np.zeros_like(radius)\n",
    "dsdr = np.zeros_like(radius)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*WARNING* You should be very careful how you think about the entropy gradient when moving from MESA to Rayleigh due to the differing equations of state. Here we have chosen to simply set the reference state entropy gradient to zero and let the convection establish it's own entropy gradient. This may or may not be satisfactory for your application.\n",
    "\n",
    "*ANOTHER WARNING* You should be very careful with radiative luminosity and/or nuclear energy generation. There are a number of ways to compute the heating functions you need. For this example, we have chosen to simply compute the luminosity profile needed for flux balance if the convective transport matches the values from MESA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_rad = -np.gradient((p.luminosity - p.conv_L_div_L*p.luminosity)*mesa.solarlum, r_MESA)/(4.0*np.pi*r_MESA**2)\n",
    "heatingp = interp(radius, q_rad)\n",
    "luminosity = np.trapz(4.0*np.pi*radius**2*heatingp, radius)\n",
    "hprofile = heatingp/luminosity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the density from MESA and the newly interpolated density that will be fed into Rayleigh to make sure they are consistent. We also plot the heating profile as a sanity check."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(radius, density,'or')\n",
    "plt.plot(p.rmid[::-1] * mesa.rsol, 10**p.logRho[::-1], '-b')\n",
    "plt.xlabel('Radius (cm)')\n",
    "plt.ylabel(r'Desnity (g/cm$^3$)')\n",
    "plt.xlim([r0,r1])\n",
    "plt.ylim([np.min(density), np.max(density)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(radius, luminosity*hprofile, '-r')\n",
    "plt.xlabel('Radius (cm)')\n",
    "plt.ylabel(r'$Q$ (erg/(cm$^3$ s))')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now create the data structure that will be written to a file that Rayleigh can read, and then load in the needed radial functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_ref = rt.equation_coefficients(radius)\n",
    "\n",
    "my_ref.set_function(density,'density')\n",
    "my_ref.set_function(buoy,'buoy')\n",
    "my_ref.set_constant(1.0, 'buoy_fact')\n",
    "my_ref.set_function(nu,'nu')\n",
    "my_ref.set_constant(1.0, 'visc_fact')\n",
    "my_ref.set_function(temperature,'temperature')\n",
    "my_ref.set_function(kappa,'kappa')\n",
    "my_ref.set_constant(1.0, 'diff_fact')\n",
    "my_ref.set_constant(1.0, 'p_fact')\n",
    "my_ref.set_function(hprofile,'heating')\n",
    "my_ref.set_constant(1.0, 'luminosity')\n",
    "my_ref.set_function(eta,'eta')\n",
    "my_ref.set_constant(1.0, 'resist_fact')\n",
    "my_ref.set_function(dsdr,'ds_dr')\n",
    "\n",
    "my_ref.set_constant(luminosity,'luminosity')\n",
    "print(my_ref.constants)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_write='cref_from_MESA.dat'\n",
    "my_ref.write(file_write)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can use this file to run a Rayleigh model. Once your Rayleigh model has run you can use the ``equation_coefficients`` file to check how your specified reference state looks when transfered into Rayleigh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Once you're run for one time step, set have_run = True\n",
    "radius1 = radius\n",
    "gravity = grav\n",
    "try:\n",
    "    cref = rd.PDE_Coefficients()\n",
    "    have_run = True\n",
    "except:\n",
    "    have_run = False\n",
    "    \n",
    "if (have_run):\n",
    "\n",
    "    fig, ax = plt.subplots(ncols=3,nrows=3, figsize=(9,3*3))\n",
    "    # Density variables\n",
    "    ax[0][0].plot(cref.radius,cref.density,'yo')\n",
    "    ax[0][0].plot(radius1,density)\n",
    "    ax[0][0].set_xlabel('Radius')\n",
    "    ax[0][0].set_title('Density')\n",
    "    \n",
    "    ax[0][1].plot(cref.radius, cref.nu,'yo')\n",
    "    ax[0][1].plot(radius1, nu)\n",
    "    ax[0][1].set_xlabel('Radius')\n",
    "    ax[0][1].set_title(r'$\\nu$')\n",
    "    \n",
    "    ax[0][2].plot(cref.radius,cref.kappa,'yo')\n",
    "    ax[0][2].plot(radius1,kappa)\n",
    "    ax[0][2].set_xlabel('Radius')\n",
    "    ax[0][2].set_title(r'$\\kappa$')\n",
    "        \n",
    "    ax[1][1].plot(cref.radius,cref.temperature,'yo')\n",
    "    ax[1][1].plot(radius1,temperature)\n",
    "    ax[1][1].set_xlabel('Radius')\n",
    "    ax[1][1].set_title('Temperature')\n",
    "    \n",
    "    ''''\n",
    "    #Activate this if your case is magnetic\n",
    "    ax[1][0].plot(cref.radius, cref.eta,'yo')\n",
    "    ax[1][0].plot(radius1, eta)\n",
    "    ax[1][0].set_xlabel('Radius')\n",
    "    ax[1][0].set_title(r'$\\eta$')\n",
    "    '''\n",
    "    \n",
    "    ax[2][1].plot(cref.radius, cref.dsdr,'yo')\n",
    "    ax[2][1].plot(radius1, dsdr)\n",
    "    ax[2][1].set_xlabel('Radius')\n",
    "    ax[2][1].set_title('Log entropy gradient')   \n",
    "    \n",
    "    ax[1][2].plot(cref.radius, cref.functions[:,1]*cref.constants[1],'yo')\n",
    "    ax[1][2].plot(radius1, gravity*density/cp)\n",
    "    ax[1][2].set_xlabel('Radius')\n",
    "    ax[1][2].set_title('Gravity')   \n",
    "\n",
    "    ax[2][0].plot(cref.radius, cref.heating*cref.rho*cref.T,'yo')\n",
    "    ax[2][0].plot(radius1, hprofile*luminosity)\n",
    "    ax[2][0].set_xlabel('Radius')\n",
    "    ax[2][0].set_title('Heating')      \n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}