{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom reference state for non-dimensional anelastic MHD formulation in a convective spherical shell\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We non-dimensionalize the  dimensional anelastic MHD equations (see Rayleigh documentation for the full set of equations, and the definitions of the variables and parameters used here) using the width of the convection zone $[l]=L=r_o-r_{cz}$ as the lengthscale (where $r_o$ is the outer radius and $r_{cz}$ is the radius at the bottom of the convection zone (CZ)), the viscous timescale $[t]=L^2/\\nu$, and the velocity scale $[u]=\\nu/L$. For the density and temperature scales, we choose\n",
    "$[\\rho]=\\tilde{\\rho}$, and $[T]=\\tilde{T}$, respectively, where the tildes correspond to the volume average of the respective quantity such that $\\tilde{q}=\\dfrac{1}{V}\\int_V q dV$, where $V$ is the volume of the convective region of the spherical shell (and $q\\rightarrow \\rho, T$). For the magnetic field, we use $[B]=\\sqrt{\\tilde{\\rho}\\mu\\eta\\Omega_o}$ and finally for the entropy $S$ we use a scaling related to the thermal energy flux $\\tilde{F}$ such that $[S]=\\dfrac{L\\tilde{F}}{\\tilde{\\rho}\\tilde{T}\\kappa}$.\n",
    "Then, we can write the non-dimensional anelastic MHD equations as:\n",
    "\n",
    "\\begin{equation}\n",
    "\\dfrac{\\partial{\\vec{u}}}{\\partial{t}} +\\vec{u}\\cdot\\nabla\\vec{u}+\\dfrac{1}{\\rm{E}}2\\hat{z}\\times\\vec{u}=\\dfrac{\\rm{Ra}}{\\rm{Pr}}\\dfrac{g(r)}{\\tilde{g}}S\\hat{r}-\\nabla(p/\\bar{\\rho})+\\dfrac{1}{\\bar{\\rho}}\\nabla\\cdot \\vec{D}+\n",
    "\\dfrac{1}{4\\pi\\bar{\\rho}}\\dfrac{1}{\\rm{P_m E}}(\\nabla\\times\\vec{B})\\times\\vec{B}.\n",
    "\\end{equation}\n",
    "\n",
    "\\begin{equation}\n",
    "\\nabla\\cdot(\\bar{\\rho}\\vec{u})=0,\n",
    "\\end{equation}\n",
    "\\begin{equation}\n",
    "\\nabla\\cdot\\vec{B}=0.\n",
    "\\end{equation}\n",
    "\n",
    "\\begin{equation}\n",
    "\\bar{\\rho}\\bar{T}\\left( \\dfrac{dS}{dt}+\\vec{u}\\cdot\\nabla S \\right )=\\dfrac{1}{\\rm{Pr}}\\nabla\\cdot(\\bar{\\rho}\\bar{T}\\nabla S)+\\dfrac{1}{\\rm{Pr}}Q_{nd}+\\dfrac{\\rm{Di Pr}}{\\rm{Ra}}\\Phi+\\dfrac{1}{4\\pi}\\dfrac{\\rm{Di Pr}}{\\rm{E P_m^2 Ra}}|\\nabla\\times\\vec{B}|^2,\n",
    "\\end{equation}\n",
    " where $Q_{nd}=\\dfrac{L}{\\tilde{F}}Q$, \n",
    "and\n",
    "\\begin{equation}\n",
    "\\dfrac{\\partial{\\vec{B}}}{\\partial{t}}-\\nabla\\times(\\vec{u}\\times\\vec{B})=-\\nabla\\times\\left(\\dfrac{1}{\\rm{P_m}}\\nabla\\times\\vec{B}\\right)\\left\\{=\\dfrac{1}{\\rm{P_m}}\\nabla^2\\vec{B}\\right\\},\n",
    "\\end{equation}\n",
    "where we have assumed that $\\kappa$, $\\nu$ and $\\eta$ are constants.\n",
    " We now have five non-dimensional numbers: the flux Rayleigh number Ra, the Prandtl number Pr, the magnetic Prandtl number $\\rm{P_m}$, the Ekman number E and the  dissipation number Di, which are defined respectively as:\n",
    "${\\rm{Ra}}=\\dfrac{\\tilde{g}\\tilde{F}L^4}{c_p\\tilde{\\rho}\\tilde{T}\\kappa^2\\nu},\\quad\n",
    "{\\rm{Pr}}=\\dfrac{\\nu}{\\kappa}, \\quad{\\rm{P_m}}=\\dfrac{\\nu}{\\eta},\\quad\n",
    "{\\rm{E}}=\\dfrac{\\nu}{\\Omega_oL^2},\\quad\\text{and}\\quad \n",
    "{\\rm{Di}}=\\dfrac{\\tilde{g}L}{c_p\\tilde{T}}$.\n",
    "Then the functions $f$ are:\n",
    "\n",
    "$f_1(r)\\rightarrow \\bar{\\rho}$,\n",
    "\n",
    "$f_2(r)\\rightarrow \\bar{\\rho}g(r)/\\tilde{g}$,\n",
    "\n",
    "$f_3(r)\\rightarrow 1$,\n",
    "\n",
    "$f_4(r)\\rightarrow \\bar{T}$,\n",
    "\n",
    "$f_5(r)\\rightarrow 1$,\n",
    "\n",
    "$f_6(r)\\rightarrow Q_{nd}(r)$, \n",
    "\n",
    "$f_7(r)\\rightarrow 1$, \n",
    "\n",
    "and the constants are \n",
    "\n",
    "$c_1\\rightarrow 2/{\\rm{E}}$, \n",
    "\n",
    "$c_2\\rightarrow {\\rm{Ra/Pr}}$, \n",
    "\n",
    "$c_3\\rightarrow 1$,\n",
    "\n",
    "$c_4\\rightarrow \\dfrac{1}{4\\pi}\\dfrac{1}{\\rm{P_m E}}$, \n",
    "\n",
    "$c_5\\rightarrow 1$, \n",
    "\n",
    "$c_6\\rightarrow 1/{\\rm{Pr}}$, \n",
    "\n",
    "$c_7\\rightarrow 1/{\\rm{P_m}}$, \n",
    "\n",
    "$c_8\\rightarrow \\rm{Di Pr/Ra}$, \n",
    "\n",
    "$c_9\\rightarrow\\dfrac{1}{4\\pi}\\dfrac{\\rm{Di Pr}}{\\rm{E P_m^2 Ra}}$, and \n",
    "\n",
    "$c_{10}\\rightarrow \\dfrac{1}{\\rm{Pr}}$.\n",
    "\n",
    "\n",
    " The non-dimensional reference  density and  temperature profiles are equal to their dimensional profile divided by its volume average in the CZ, i.e. $\\bar{\\rho}=\\bar{\\rho}_{dim}/\\tilde{\\rho}$, and $\\bar{T}=\\bar{T}_{dim}/\\tilde{T}$.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the reference state is based on a polytrope, similarly to the dimensional case. For more info on that, \n",
    "check the dimensional anelastic notebook version!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#######################\n",
    "import numpy\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import plot, draw, show\n",
    "\n",
    "import os, sys\n",
    "sys.path.insert(0, os.path.abspath('../../'))\n",
    "\n",
    "import post_processing.reference_tools as rt # You will need the refernce_tools.py to run this notebook\n",
    "import post_processing.rayleigh_diagnostics as rdiag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Grid Parameters\n",
    "nr    = 512         # Number of radial points\n",
    "ri    = 5e10        # Inner boundary of radial domain\n",
    "ro    = 6.586e10    # Outer boundary of radial domain\n",
    "rcz   = 5e10        # Base of the CZ\n",
    "\n",
    "# apsect ratio\n",
    "beta=ri/ro\n",
    "\n",
    "# shell depth according to non-dimensionalization\n",
    "d=1.e0\n",
    "\n",
    "## In the main_input file, we should use the following radial boundaries or else set the right aspect ratio and shell depth \n",
    "\n",
    "ri_nd=beta*d/(1-beta) # non-dimensional inner boundary\n",
    "\n",
    "ro_nd=d/(1-beta) # non-dimensional outer boundary\n",
    "\n",
    "print(ri_nd,ro_nd)\n",
    "\n",
    "# Polytrope Parameters\n",
    "# In this example, we use solar values\n",
    "\n",
    "ncz   = 1.5         # polytropic index of convection zone\n",
    "nrho  = 3.          # number of density scaleheights across convection zone (not full domain)\n",
    "mass  = 1.98891e33  # Mass of the star\n",
    "G     = 6.67e-8     # gravitational constant\n",
    "rhoi  = 0.18053428  # density at the base of convection zone\n",
    "cp    = 3.5e8       # specific heat at constant pressure\n",
    "gamma = 5.0/3.0     # Ratio of specific heats for ideal gas\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Radial Grid\n",
    "\n",
    "radius = numpy.linspace(ri,ro,nr)           # Full domain radial grid for reference state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Compute  CZ polytrope -- see reference_tools.py script\n",
    "\n",
    "poly1 = rt.gen_poly(radius,ncz,nrho,mass,rhoi,G,cp,rcz)\n",
    "gas_constant = cp*(1.0-1.0/gamma)  # R\n",
    "\n",
    "#Dimensional profiles for reference state functions\n",
    "temperature1 = poly1.temperature # T\n",
    "density1 = poly1.density         # rho\n",
    "pressure1 = poly1.pressure       # P\n",
    "dsdr1 = poly1.entropy_gradient   # dS/dr\n",
    "s1= poly1.entropy                # S\n",
    "dpdr1 = poly1.pressure_gradient  # dP/dr\n",
    "gravity1= mass*G/radius**2       # g\n",
    "\n",
    "# We use the volume average of rho,T and g in the CZ to non-dimensionalize our reference state functions\n",
    "def vol_av(f,radius):\n",
    "    int1=radius**2.\n",
    "    int2=f*radius**2.\n",
    "    vol = numpy.trapz(int1,x=radius)\n",
    "    f_av1 = numpy.trapz(int2,x=radius)\n",
    "    f_av=f_av1/vol\n",
    "    return f_av\n",
    "  \n",
    "temperature_av=vol_av(temperature1,radius) # CZ volume average of T\n",
    "density_av=vol_av(density1,radius)         # CZ volume average of rho\n",
    "gravity_av=vol_av(gravity1,radius)         # CZ volume average of g\n",
    "\n",
    "\n",
    "print(temperature_av,density_av,gravity_av)\n",
    "\n",
    "# Here we define the non-dimensional T, rho, g, S, r, dS/dr and P\n",
    "\n",
    "# The non-dimensional temperature: T=T_dimensional/T_av\n",
    "temperature=temperature1/temperature_av\n",
    "\n",
    "# The non-dimensional density: rho=rho_dimensional/rho_av\n",
    "density=density1/density_av\n",
    "\n",
    "# The non-dimensional gravity: g=g_dimensional/g_av\n",
    "gravity=gravity1/gravity_av\n",
    "\n",
    "\n",
    "# Here we define our non-dimensional radius \n",
    "radius1=radius/(ro-ri) # OR: radius1=numpy.linspace(ri_nd,ro_nd,nr)     \n",
    "\n",
    "#The pressure profile won't matter, we set it equal to rho*T as a reference\n",
    "pressure=temperature*density\n",
    "\n",
    "# For a purely convective region, we use dS/dr=0\n",
    "dsdr=dsdr1\n",
    "\n",
    "# Entropy won't matter, set it to something-- here I use the non-dimensional S i.e. S_nondim=[S_dim*(L^3*g_av)]/Ra*cp*kappa*nu]\n",
    "entropy=numpy.ones(len(radius))*s1*(1.586e10**3.*gravity_av)/(13303.43109666924*cp*(8e12)**2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we calculate the derivatives of lnrho and lnT based on the non-dimensional radius1, since\n",
    "# we want their non-dimensional profiles\n",
    "\n",
    "d_density_dr = numpy.gradient(density,radius1, edge_order=2)\n",
    "dlnrho = d_density_dr/density\n",
    "d2lnrho = numpy.gradient(dlnrho,radius1, edge_order=2)\n",
    "dtdr = numpy.gradient(temperature,radius1, edge_order =2)\n",
    "dlnt = dtdr/temperature\n",
    "\n",
    "\n",
    "### Here, we plot our non-dimensional profiles for our reference state functions, i.e. for rho, T, etc.\n",
    "\n",
    "fig, ax = plt.subplots(nrows =3,ncols = 3, figsize=(15,10) )\n",
    "\n",
    "ax[0][0].plot(radius1,density,'r')\n",
    "ax[0][0].set_xlabel('Radius')\n",
    "ax[0][0].set_ylabel('Density')\n",
    "\n",
    "ax[0][1].plot(radius1,entropy)\n",
    "ax[0][1].set_xlabel('Radius')\n",
    "ax[0][1].set_ylabel('Entropy')\n",
    "\n",
    "\n",
    "ax[0][2].plot(radius1,temperature)\n",
    "ax[0][2].set_xlabel('Radius')\n",
    "ax[0][2].set_ylabel('Temperature')\n",
    "\n",
    "ax[1][0].plot(radius1,dsdr)\n",
    "ax[1][0].set_xlabel('Radius')\n",
    "ax[1][0].set_ylabel('Entropy Gradient')\n",
    "\n",
    "ax[1][1].plot(radius1,pressure)\n",
    "ax[1][1].set_xlabel('Radius')\n",
    "ax[1][1].set_ylabel('Pressure')\n",
    "\n",
    "ax[1][2].plot(radius1,gravity)\n",
    "ax[1][2].set_xlabel('Radius')\n",
    "ax[1][2].set_ylabel('Gravity')\n",
    "\n",
    "ax[2][0].plot(radius1,dlnrho)\n",
    "ax[2][0].set_xlabel('Radius')\n",
    "ax[2][0].set_ylabel('dlnrho')\n",
    "\n",
    "ax[2][1].plot(radius1,d2lnrho)\n",
    "ax[2][1].set_xlabel('Radius')\n",
    "ax[2][1].set_ylabel('d2lnrho')\n",
    "\n",
    "ax[2][2].plot(radius1,dlnt)\n",
    "ax[2][2].set_xlabel('Radius')\n",
    "ax[2][2].set_ylabel('dlnt')\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "plt.show()\n",
    "print(density[0],temperature[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## This is where we define a heating function, if we want one in our model.\n",
    "\n",
    "# Units are energy / volume / time such that {rho_hat T_hat dS/dt} = hprofile(r)\n",
    "\n",
    "\n",
    "hprofile = numpy.zeros(nr,dtype='float64')\n",
    "hprofile[:] = pressure1[:]\n",
    "\n",
    "# Next, we need to integrate the heating profile\n",
    "# We normalize it such that its integral over the volume is 1\n",
    "# This way, we can set the luminosity via a constant in the input file for the dimensional case!\n",
    "\n",
    "integrand= numpy.pi*4*radius*radius*hprofile\n",
    "hint = numpy.trapz(integrand,x=radius)\n",
    "hprofile = hprofile/hint\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "###########################################################\n",
    "# Plot the integrated luminosity as a function of radius\n",
    "# (should integrate to 1 at r = r_top)\n",
    "###########################################################\n",
    "\n",
    "\n",
    "# We then need to calculate the non-dimensional heating function\n",
    "\n",
    "lq1 = numpy.zeros(nr)\n",
    "lq1[0]=0\n",
    "lq2 = numpy.zeros(nr)\n",
    "lq2[0]=0\n",
    "lq3 = numpy.zeros(nr)\n",
    "lq3[0]=0\n",
    "\n",
    "lsun= 3.846e33 # solar luminosity\n",
    "integrand1= lsun*radius*radius*hprofile # Luminosity*r^2*heating\n",
    "for i in range(1,nr):\n",
    "    lq1[i] =(1/(radius[i]**2.))*numpy.trapz(integrand1[0:i+1],x=radius[0:i+1])\n",
    "\n",
    "integrand2=4*numpy.pi*radius**2.\n",
    "lq2= numpy.trapz(integrand2,x=radius)\n",
    "integrand3=lq1*4*numpy.pi*radius**2.\n",
    "lq3=numpy.trapz(integrand3,x=radius)/lq2 # That is the volume Flux F_tilde\n",
    "\n",
    "\n",
    "# The  non-dimensional heating profile is hprofile_nd = lsun*hprofile*L/F_tilde=Q_dim*(r_o-r_cz)/F_tilde\n",
    "hprofile_nd=lsun*hprofile*(ro-ri)/lq3 \n",
    "#print((ro-ri)/lq3) # This comes out of the non-dimensionalization used (L/F_tilde)\n",
    "\n",
    "nu=8.e12\n",
    "kappa=8.e12\n",
    "\n",
    "Ra=gravity_av*lq3*(ro-ri)**4./(cp*density_av*temperature_av*nu*kappa**2.) # Ra is defined earlier in the notes\n",
    "Di=gravity_av*(ro-ri)/(cp*temperature_av) # Dissipation number is defined in the notes above\n",
    "print(Ra,Di/Ra)\n",
    "\n",
    "fig, ax = plt.subplots(ncols=2,figsize=(12,4))\n",
    "\n",
    "ax[0].plot(radius1,hprofile_nd,'ob')\n",
    "ax[0].set_title('Non-Dimensional Heating Profile')\n",
    "ax[0].set_xlabel('Radius')\n",
    "\n",
    "ax[1].plot(radius,hprofile,'.b')\n",
    "ax[1].set_xlabel('Radius (cm)')\n",
    "ax[1].set_title('Dimensional Heating Profile')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Have everything in terms of the non-dimensional radius\n",
    "my_ref = rt.equation_coefficients(radius1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we define all the functions and constants that will be written in our data file and\n",
    "# read by Rayleigh if we choose the custom reference state (=4)\n",
    "\n",
    "unity = numpy.ones(nr,dtype='float64')\n",
    "buoy =density*gravity # buoyancy term calculation\n",
    "\n",
    "my_ref.set_function(density,1) # density rho\n",
    "my_ref.set_function(buoy,2)    # buoyancy term\n",
    "my_ref.set_function(unity,3)  # nu(r) -- can be overwritten via nu_type in Rayleigh\n",
    "my_ref.set_function(temperature,4) # temperature T\n",
    "my_ref.set_function(unity,5)  # kappa(r) -- works like nu\n",
    "my_ref.set_function(hprofile_nd,6)  # heating profile\n",
    "\n",
    "my_ref.set_function(dlnrho,8) # dlnrho/dr\n",
    "my_ref.set_function(d2lnrho,9) # d^2lnrho/dr^2\n",
    "my_ref.set_function(dlnt,10)   # dlnT/dr\n",
    "my_ref.set_function(unity,7)  # eta -- works like nu and kappa\n",
    "my_ref.set_function(dsdr,14) # dS/dr\n",
    "\n",
    "\n",
    "# The constants can all be set/overridden in the input file\n",
    "# NOTE that they default to ZERO, but we want\n",
    "# most of them to be UNITY.  These constants will explicitly depend on the non-dimensionalization chosen.\n",
    "\n",
    "# The comments corresponding to each one of the constants are generic but  we also specify \n",
    "# what they are exactly in our example here, assuming the non-dimensionalization we used.\n",
    "\n",
    "# This is a non-magnetic, non-rotating example that reproduces the results from Featherstone & Hindman (2016)\n",
    "# for the case with Nrho=3,and  kappa=nu=8e12\n",
    "\n",
    "my_ref.set_constant(1.0,1)  # multiplies the Coriolis force, here it is: 2/E  \n",
    "my_ref.set_constant(13303.43109666924,2)  # multiplies buoyancy -- Ra/Pr here\n",
    "my_ref.set_constant(1.0,3)  # multiplies the pressure gradient\n",
    "my_ref.set_constant(0.0 , 4)  # multiplies the lorentz force, here it is: (1/(4*pi)*1/(Pm*E))\n",
    "my_ref.set_constant(1.0,5)  # multiplies viscosity\n",
    "my_ref.set_constant(1.0,6)  # multiplies entropy diffusion , here it is 1/Pr\n",
    "my_ref.set_constant(0.0,7)  # multiplies magnetic diffusion in induction equation, here: 1/Pm\n",
    "my_ref.set_constant(0.00012954929449971041,8)  # multiplies viscous heating, here it is: Di*Pr/Ra\n",
    "my_ref.set_constant(1.0,9)  # multiplies ohmic heating, here it is: (1/(4*pi)*Di*Pr/(E*Pm^2*Ra))\n",
    "my_ref.set_constant(1.0,10) # multiplies the heating function -- here it is 1/Pr (if normalized to 1, this is the luminosity)\n",
    "my_ref.write('CZtest.dat') # Here we write our data file to be used to run our simulation with Rayleigh!\n",
    "print(my_ref.fset)\n",
    "print(my_ref.cset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Once you've run for one time step, set have_run = True\n",
    "\n",
    "# Here we check if the output reference state is the same as the one we used as an input (sanity check)!\n",
    "\n",
    "# NOTE: We need the output file \"equation_coefficients\" to run this, as well as the PDE_Coefficients \n",
    "# from rayleigh_diagnostics.py\n",
    "\n",
    "try:\n",
    "    cref = rdiag.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=4, figsize=(16,4*4))\n",
    "    \n",
    "    # Density and derivatives of lnrho\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.dlnrho,'yo')\n",
    "    ax[0][1].plot(radius1, dlnrho)\n",
    "    ax[0][1].set_xlabel('Radius')\n",
    "    ax[0][1].set_title('Log density gradient')\n",
    "    \n",
    "    ax[0][2].plot(cref.radius,cref.d2lnrho,'yo')\n",
    "    ax[0][2].plot(radius1,d2lnrho)\n",
    "    ax[0][2].set_xlabel('Radius')\n",
    "    ax[0][2].set_title('d_dr{Log density gradient}')\n",
    "        \n",
    "    # Temperature  and derivative of lnT    \n",
    "    ax[1][0].plot(cref.radius,cref.temperature,'yo')\n",
    "    ax[1][0].plot(radius1,temperature)\n",
    "    ax[1][0].set_xlabel('Radius')\n",
    "    ax[1][0].set_title('Temperature')\n",
    "    \n",
    "    ax[1][1].plot(cref.radius, cref.dlnT,'yo')\n",
    "    ax[1][1].plot(radius1, dlnt)\n",
    "    ax[1][1].set_xlabel('Radius')\n",
    "    ax[1][1].set_title('Log temperature gradient')\n",
    "\n",
    "    # dS/dr\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('Entropy gradient')   \n",
    "    \n",
    "    # Buoyancy, Heating\n",
    "    # Note that you must build the buoyancy term from the functions/constants\n",
    "    ax[3][1].plot(cref.radius, cref.functions[:,1]*cref.constants[1],'yo')\n",
    "    ax[3][1].plot(radius1, gravity*density*Ra)\n",
    "    ax[3][1].set_xlabel('Radius')\n",
    "    ax[3][1].set_title('Buoyancy')\n",
    "    \n",
    "    # Note that the output heating (cref.heating) is hprofile/density/temperature\n",
    "    ax[3][2].plot(cref.radius, cref.heating,'yo')\n",
    "    ax[3][2].plot(radius1, hprofile_nd/density/temperature)\n",
    "    ax[3][2].set_xlabel('Radius')\n",
    "    ax[3][2].set_title('Heating')      \n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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
}