{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The First Stars: Formation under X-ray Feedback -- Structure Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook generates all simulation results figures for my current paper.  All simulations run on stampede.tacc.utexas.edu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "from matplotlib import pyplot as plt\n",
    "mpl.rc('font', size=20.)\n",
    "mpl.rc('font', family='serif')\n",
    "mpl.rc('text', usetex=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import pyGadget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Density structure zoom-in to central minihalo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the final snapshot of the vanilla simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sims = ['vanilla', 'xr_tau_J0', 'xr_tau_J1', 'xr_tau2_J2', 'xr_tau_J3', 'XR_sfr_1e-1', 'XR_sfr_1e-2', 'XR_sfr_1e-3']\n",
    "n0 = [355, 306, 327, 227, 235, 200, 201, 269]\n",
    "n5k = [1857, 1546, 1852, 1758, 1687, 1616, 1900]\n",
    "t0 = '_t0'\n",
    "t5k = '_t5k'\n",
    "n, tag = n0, t0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "i = 0\n",
    "sim = pyGadget.sim.Simulation('stampede/'+sims[i])\n",
    "sim.refine_by_mass(False)\n",
    "sim.set_coordinate_system('physical')\n",
    "snap = sim.load_snapshot(n[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "imzoom = []\n",
    "for scale in ['5376pc', '1000pc', '10pc', '100pc']:\n",
    "    imzoom.append(pyGadget.visualize.project(snap, 'ndensity', scale, 'xz', centering='avg'))\n",
    "#snap.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "\n",
    "scales = ['140 kpc (comoving)', '1 kpc (physical)', '10 pc (physical)', '100 pc (physical)']\n",
    "ratio = [.1788, .1, None, .1]\n",
    "zoom = ['right', 'down', None, 'left']\n",
    "clims = [(-2.5,1.5), (-2.,2.), (1.5,7.5), (-0.5,5.)]\n",
    "ticks = [(-2,-1,0,1), (-1,0,1), (2,3,4,5,6,7), (0,1,2,3,4)]\n",
    "cpad = [-17, -17, -15, -16]\n",
    "clabel = [False, True, False, True]\n",
    "bbox_props = dict(boxstyle=\"round\", fc=\"k\", ec=\"k\", alpha=0.5)\n",
    "zc = 'w'\n",
    "zls = '--'\n",
    "zlw = 1.5\n",
    "\n",
    "fig = plt.figure(1, (12., 12.), dpi=600)\n",
    "grid = ImageGrid(fig, 111, # similar to subplot(111)\n",
    "                nrows_ncols = (2, 2), # creates 2x2 grid of axes\n",
    "                axes_pad=0.0, # pad between axes in inch.\n",
    "                cbar_mode = 'each', cbar_size='7%', cbar_pad=0.\n",
    "                )\n",
    "\n",
    "for i in range(4):\n",
    "    x = imzoom[i][0]\n",
    "    y = imzoom[i][1]\n",
    "    im = imzoom[i][2]\n",
    "    ax = grid[i]\n",
    "    img = ax.imshow(im, cmap=plt.cm.bone, origin='lower')\n",
    "    ax.xaxis.set_visible(False)\n",
    "    ax.yaxis.set_visible(False)\n",
    "    img.set_clim(clims[i])\n",
    "    \n",
    "    cb = plt.colorbar(img, cax=grid.cbar_axes[i])\n",
    "    cb.set_ticks(ticks[i])\n",
    "    cb.ax.tick_params(left='on', pad=cpad[i],\n",
    "                      labelsize=15, labelcolor='k', labelleft='on', labelright='off')\n",
    "    if clabel[i]: cb.set_label('Log Number Density [cm$^{-3}$]')\n",
    "        \n",
    "    ax.text(0.5, 0.025, scales[i], color='w', ha='center', va='bottom', size=12, \n",
    "            transform=grid[i].transAxes, bbox=bbox_props)\n",
    "    \n",
    "    if ratio[i]:\n",
    "        axmin, axmax = ax.get_xlim()\n",
    "        axlength = axmax - axmin\n",
    "        mid = axlength/2\n",
    "        s = ratio[i] * axlength\n",
    "        s00 = [mid - s/2, mid - s/2]\n",
    "        s01 = [mid - s/2, mid + s/2]\n",
    "        s11 = [mid + s/2, mid + s/2]\n",
    "        ax.add_line(plt.Line2D(s00, s01, c=zc, lw=zlw))\n",
    "        ax.add_line(plt.Line2D(s11, s01, c=zc, lw=zlw))\n",
    "        ax.add_line(plt.Line2D(s01, s00, c=zc, lw=zlw))\n",
    "        ax.add_line(plt.Line2D(s01, s11, c=zc, lw=zlw))\n",
    "        if zoom[i] == 'right':\n",
    "            ax.add_line(plt.Line2D([mid+s/2, axmax], [mid+s/2, axmax], c=zc, lw=zlw, ls=zls))\n",
    "            ax.add_line(plt.Line2D([mid+s/2, axmax], [mid-s/2, axmin], c=zc, lw=zlw, ls=zls))\n",
    "        elif zoom[i] == 'down':\n",
    "            ax.add_line(plt.Line2D([mid-s/2, axmin], [mid-s/2, axmin], c=zc, lw=zlw, ls=zls))\n",
    "            ax.add_line(plt.Line2D([mid+s/2, axmax], [mid-s/2, axmin], c=zc, lw=zlw, ls=zls))\n",
    "        elif zoom[i] == 'left':\n",
    "            ax.add_line(plt.Line2D([mid-s/2, axmin], [mid+s/2, axmax], c=zc, lw=zlw, ls=zls))\n",
    "            ax.add_line(plt.Line2D([mid-s/2, axmin], [mid-s/2, axmin], c=zc, lw=zlw, ls=zls))\n",
    "plt.show()\n",
    "#fig.savefig('figures/structure/structure-'+sim.name.split('/')[-1]+tag+'.png', bbox_inches='tight', dpi=100)\n",
    "fig.savefig('figures/structure/structure-'+sim.name.split('/')[-1]+tag+'.pdf', bbox_inches='tight', dpi=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4 Panel Simulation Comparison Plots"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For each sim, picking the snapshot just prior to the formation of the first sink."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "simV = pyGadget.sim.Simulation('stampede/vanilla', track_sinks=True)\n",
    "sim1 = pyGadget.sim.Simulation('stampede/xr_tau2_J0', track_sinks=True)\n",
    "sim2 = pyGadget.sim.Simulation('stampede/xr_tau2_J1', track_sinks=True)\n",
    "sim3 = pyGadget.sim.Simulation('stampede/xr_tau2_J2', track_sinks=True)\n",
    "snapV = simV.load_snapshot(1900)\n",
    "snap1 = sim1.load_snapshot(1778)\n",
    "snap2 = sim2.load_snapshot(1669)\n",
    "snap3 = sim3.load_snapshot(1727)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Disk Structure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Density structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "snaplist = [snapV, snap1, snap2, snap3]\n",
    "imlist = []\n",
    "sinklist = []\n",
    "scale = '15000AU'\n",
    "\n",
    "#shifty = [None, None, 5000, -2000, None, None, None, None]\n",
    "shifty = [None, None, None, None, None, None, None, None]\n",
    "face = [[('x', 0.29518), ('z', 0.825), ('x',np.pi/2)],\n",
    "        [('y', 0.6346), ('z', 2.03), ('x',np.pi/2)],\n",
    "        [('x', 1.865), ('z', 2.919), ('x',np.pi/2)],\n",
    "        [('x', 1.7136), ('z', 0.18), ('x',np.pi/2)]]\n",
    "edge = [[('x', 0.29518), ('z', 0.825)],\n",
    "        [('y', 0.6346), ('z', 2.03)],\n",
    "        [('x', 1.865), ('z', 2.919)],\n",
    "        [('x', 1.7136), ('z', 0.18)]\n",
    "        ]\n",
    "ocount = 0\n",
    "for view in [face, edge]:\n",
    "    count = 0\n",
    "    for snap in snaplist:\n",
    "        imlist.append(pyGadget.visualize.project(snap, 'ndensity', scale, view[count], centering='avg', \n",
    "                                                 depth=2., shifty=shifty[ocount], dens_lim=None))\n",
    "        sinklist.append(copy.deepcopy(snap.sinks))\n",
    "        count += 1\n",
    "        ocount += 1\n",
    "#    snap.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "\n",
    "sim = ['J = 0', 'J = J$_{0}$', 'J = 10 J$_{0}$', 'J = 100 J$_{0}$']\n",
    "bbox_props = dict(boxstyle=\"round\", fc=\"k\", ec=\"k\", alpha=0.5)\n",
    "ticks = [(7,8,9,10,11),(6,7,8,9,10)]\n",
    "\n",
    "fig = plt.figure(1, (20, 8), dpi=600)\n",
    "grid = ImageGrid(fig, 111, # similar to subplot(111)\n",
    "                nrows_ncols = (2, 4), # creates 4x2 grid of axes\n",
    "                axes_pad=0.0, # pad between axes in inch.\n",
    "                cbar_mode = 'edge', cbar_location = 'right', cbar_size='7%', cbar_pad=0.0\n",
    "                )\n",
    "for i in range(8):\n",
    "    x = imlist[i][0]\n",
    "    y = imlist[i][1]\n",
    "    im = imlist[i][2]\n",
    "    ax = grid[i]\n",
    "    img = ax.imshow(im, extent=[x.min(),x.max(),y.min(),y.max()], cmap=plt.cm.bone, origin='lower')\n",
    "    ax.xaxis.set_visible(False)\n",
    "    ax.yaxis.set_visible(False)\n",
    "    img.set_clim((6.5,10.75))\n",
    "    \n",
    "    cb = plt.colorbar(img, cax=grid.cbar_axes[i])\n",
    "    #cb.set_ticks(ticks[i/4])\n",
    "    cb.set_ticks((7,8,9,10))\n",
    "    cb.ax.tick_params(left='on', labelsize=15, labelcolor='k')\n",
    "    #if clabel[i]: \n",
    "    cb.set_label('Log Number Density \\n    [cm$^{-3}$]')\n",
    "    \n",
    "    for sink in sinklist[i]:\n",
    "        #mscale = sink.mass*6./27. + .33\n",
    "        mscale = np.log(sink.mass) +1\n",
    "        ax.plot(sink.x, sink.y, 'ko', ms=mscale, mew=1)\n",
    "    ax.set_xlim(x.min(), x.max())\n",
    "    ax.set_ylim(y.min(), y.max())\n",
    "\n",
    "    if i > 3:\n",
    "        cb.set_ticks(ticks[i/4])\n",
    "        ax.text(0.5, 0.025, sim[i-4], color='w', ha='center', va='bottom', size=18, \n",
    "                transform=grid[i].transAxes, bbox=bbox_props)\n",
    "\n",
    "plt.show()\n",
    "fig.savefig('figures/structure/disks.png', bbox_inches='tight', dpi=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Density / Temperature / HD Fraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "reload(pyGadget.visualize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "snaplist = [snapV, snap1, snap2, snap3]\n",
    "imlist = []\n",
    "sinklist = []\n",
    "scale = '15000AU'\n",
    "imscale=['log','log','linear']\n",
    "\n",
    "shifty = [None, None, 5000, -2000]\n",
    "\n",
    "face = [[('x', 0.29518), ('z', 0.825), ('x',np.pi/2)],\n",
    "        [('y', 0.6346), ('z', 2.03), ('x',np.pi/2)],\n",
    "        [('x', 1.865), ('z', 2.919), ('x',np.pi/2)],\n",
    "        [('x', 1.7136), ('z', 0.18), ('x',np.pi/2)]]\n",
    "\n",
    "for property in ['ndensity', 'temp', 'h2frac']:\n",
    "    count = 0\n",
    "    for snap in snaplist:\n",
    "        imlist.append(pyGadget.visualize.project(snap, property, scale, face[count], centering='avg', depth=.05,\n",
    "                                                 shifty=shifty[count], imscale=imscale[count/4]))\n",
    "                                                 \n",
    "        sinklist.append(copy.deepcopy(snap.sinks))\n",
    "        count += 1\n",
    "#    snap.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from mpl_toolkits.axes_grid1 import ImageGrid\n",
    "\n",
    "sim = ['J = 0', 'J = J$_{0}$', 'J = 10 J$_{0}$', 'J = 100 J$_{0}$']\n",
    "colormap = [plt.cm.RdGy_r, plt.cm.afmhot, plt.cm.Blues]\n",
    "color_lims = [(6.5,11), (1.8,3.8), (-3,.5)]\n",
    "labels = ['Log Number Density \\n    [cm$^{-3}$]', 'Log Gas Temperature\\n     [K]', 'H$_2$ Fraction']\n",
    "bbox_props = dict(boxstyle=\"round\", fc=\"k\", ec=\"k\", alpha=0.5)\n",
    "ticks = [(7,8,9,10,11),(2.0,2.4,2.8,3.2,3.6,4),(-3,-2,-1,0)]\n",
    "\n",
    "fig = plt.figure(1, (20, 12), dpi=600)\n",
    "grid = ImageGrid(fig, 111, # similar to subplot(111)\n",
    "                nrows_ncols = (3, 4), # creates 4x2 grid of axes\n",
    "                axes_pad=0.0, # pad between axes in inch.\n",
    "                cbar_mode = 'each', cbar_location = 'right', cbar_size='7%', cbar_pad=0.0\n",
    "                )\n",
    "for i in range(12):\n",
    "    x = imlist[i][0]\n",
    "    y = imlist[i][1]\n",
    "    im = imlist[i][2]\n",
    "    ax = grid[i]\n",
    "    img = ax.imshow(im, extent=[x.min(),x.max(),y.min(),y.max()], cmap=colormap[i/4], origin='lower')\n",
    "    ax.xaxis.set_visible(False)\n",
    "    ax.yaxis.set_visible(False)\n",
    "    img.set_clim(color_lims[i/4])\n",
    "    \n",
    "    cb = plt.colorbar(img, cax=grid.cbar_axes[i])\n",
    "    cb.set_ticks(ticks[i/4])\n",
    "    cb.ax.tick_params(left='on', labelsize=15, labelcolor='k')\n",
    "    if (i+1) % 4 == 0:\n",
    "        cb.set_label(labels[i/4])\n",
    "    else:\n",
    "        plt.setp(cb.ax.get_yticklabels(), visible=False)\n",
    "        #cb.ax.set_axis_off()\n",
    "    \n",
    "    for sink in sinklist[i]:\n",
    "        mscale = np.log(sink.mass) +1\n",
    "        ax.plot(sink.x, sink.y, 'ko', ms=mscale, mew=1)\n",
    "    ax.set_xlim(x.min(), x.max())\n",
    "    ax.set_ylim(y.min(), y.max())\n",
    "    \n",
    "    if i > 7:\n",
    "        ax.text(0.5, 0.025, sim[i/3], color='w', ha='center', va='bottom', size=18, \n",
    "                transform=grid[i].transAxes, bbox=bbox_props)\n",
    "    \n",
    "plt.show()\n",
    "fig.savefig('figures/diskprops.png', bbox_inches='tight', dpi=300)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}