Creating_Plots_with_gas-Copy1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "##########################################################################\n",
    "##########################################################################\n",
    "#This file is designed to plot animation of all 17 clusters using the file total_data_all_clusters_all_snapshots.pkl that contains tracked data of all clusters\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import os\n",
    "import pickle\n",
    "import matplotlib\n",
    "import matplotlib.colors as colors\n",
    "from matplotlib import rc #to use Latex math symbols like 'phi'\n",
    "import astropy\n",
    "from astropy.io import ascii\n",
    "\n",
    "matplotlib.rcParams['pdf.fonttype'] = 42\n",
    "matplotlib.rcParams['xtick.labelsize'] = 16\n",
    "matplotlib.rcParams['ytick.labelsize'] = 16\n",
    "matplotlib.rc('text', usetex=False)\n",
    "#############################################################################\n",
    "#constants\n",
    "#############################################################################\n",
    "\n",
    "MsunToGm = 1.99e33\n",
    "KpcToCm = 3.086e21\n",
    "mp = 1.67e-24\n",
    "#bin_edge = 10.\n",
    "bin_edge = 30.\n",
    "\n",
    "bins = np.arange(-25,25,0.1)\n",
    "\n",
    "############################################################################\n",
    "#read in sim files and find relevant particles\n",
    "############################################################################\n",
    "#STAMPEDE\n",
    "simname = 'm12i_res7100_mhdcv'\n",
    "simtype=\"sf-fire3-alpha01\"\n",
    "#simdir = '/scratch/projects/xsede/GalaxiesOnFIRE/mhdcv/m12i_res7100_mhdcv/1Myr/sf-fire3/'\n",
    "datapath=\"./\"+simtype+\"data_pkl/\" #this is teh path where the data of our tracked clusters is\n",
    "gas_datapath=\"./\"+simtype+\"_gas_data_pkl/\"\n",
    "if not os.path.exists(gas_datapath):\n",
    "  os.makedirs(gas_datapath)\n",
    "\n",
    "plot_path=\"./\"+simtype+\"_plots_with_gas/\" #creating a path to store the plots only if it does not exist\n",
    "if not os.path.exists(plot_path):\n",
    "  os.makedirs(plot_path)\n",
    "\n",
    "snap=596 #this is the snapshot at which the clusters were taken from using the fof algorithm I am using this to extract information from the snapshot where the clusters were first seen\n",
    "snapshot_start=596  #snapshot to begin creating the figure\n",
    "snapshot_end=696    #snapshot to stop at\n",
    "\n",
    "\n",
    "############################\n",
    "############################\n",
    "#Doing all this trouble to store the initial mass of the clusters\n",
    "fn=simtype+\"_clusters_at_snapshot_\"+str(snap)+\".pkl\"  #snap determines where the cluster was taken from\n",
    "with open(datapath+fn, \"rb\") as input:\n",
    "  initial_cluster_data = pickle.load(input)\n",
    "\n",
    "clustermass=[]   #it stores the mass of each cluster that was in the beginning\n",
    "\n",
    "for c in range(len(initial_cluster_data)):\n",
    "    m=np.sum(initial_cluster_data[c+1][\"mass_tracked\"])\n",
    "    clustermass.append(m)\n",
    "print(\"These are the initial mass of each star in the cluster\",clustermass)\n",
    "############################\n",
    "############################\n",
    "\n",
    "\n",
    "################################\n",
    "################################\n",
    "#This section is just to create a list of clusters and display a print like this: These are the clusters groups we have tracked data of:\n",
    "#['snapshot660_cluster_group1', 'snapshot660_cluster_group2', 'snapshot660_cluster_group3', 'snapshot660_cluster_group4', 'snapshot660_cluster_group5', \n",
    "#'snapshot660_cluster_group6', 'snapshot660_cluster_group7', 'snapshot660_cluster_group8', 'snapshot660_cluster_group9', 'snapshot660_cluster_group10']\n",
    "cluster_groupid=[]\n",
    "total_clusters=len(initial_cluster_data)\n",
    "for i in range(total_clusters):\n",
    "    cluster_groupid.append(\"snapshot\"+str(snapshot_start)+\"_cluster_group\"+str(i+1))\n",
    "\n",
    "print(\"These are the clusters groups we have tracked data of:\\n\",cluster_groupid)\n",
    "##############################\n",
    "##############################\n",
    "n=snapshot_end-snapshot_start+1 # total no. of snapshots to generate image \n",
    "\n",
    "#colors=['cyan','blue','green','magenta','yellow','teal','brown','darkslategray','lime','red','orange','purple','rosybrown','pink','navy','olive','cornflowerblue','blueviolet','darkgreen'] #19 colors\n",
    "#colors=['cyan','blue','green','magenta','yellow','teal','brown','darkslategray','lime','red','orange','purple','rosybrown','pink','navy','olive','cornflowerblue'] #17 colors\n",
    "#colors=['cyan','blue','green','magenta','yellow','teal','brown','tan','lime','red'] #10 colors\n",
    "\n",
    "#total_clusters=10 #just in case you want to change the no. of clusters. Setting 10 means select the first 10\n",
    "############################\n",
    "############################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated the plot and saved the figure as:  sf-fire3-alpha01_gas_and_clusters696.png\r"
     ]
    }
   ],
   "source": [
    "plot_count=0       #counts the plot as we move forward with each snapshot\n",
    "snapshot_list=np.arange(snapshot_start,snapshot_end+1) #create a list of snapshot numbers to plot to plot eg. [596,597, ...]\n",
    "time=snapshot_list-snapshot_start #time starts from zero here where t=0 is at the snapshot where we start from eg. [0,1,2,3,....]\n",
    "#time=[[0],[20],[40],[60],[80],[100]] #another way of providing fixed time stamps to look at\n",
    "\n",
    "tracked_gas_all_snaps={}\n",
    "for i in range(len(time)):                              \n",
    "    snapnumber=time[i]+snapshot_start      #this is the true snapshot number each time\n",
    "    tracked_gas={} \n",
    "    ###################\n",
    "    \n",
    "    cmap = plt.get_cmap(\"gist_ncar\")\n",
    "    colors = iter(cmap(np.linspace(0,1,total_clusters))) #number of clusters\n",
    "    file_name=simtype+\"_clusters_at_snapshot_\"+str(snapnumber)+\".pkl\"\n",
    "    #print(\"Reading cluster data from the file: \",file_name,end=\"\\r\",flush=True)\n",
    "    with open(datapath+file_name, \"rb\") as input:\n",
    "      importdata = pickle.load(input)\n",
    "    #loading data from each snapshot into importdata\n",
    "    ##################\n",
    "    ##################\n",
    "    fig1=plt.figure()\n",
    "    fig1.set_size_inches(7,7)\n",
    "    ax=fig1.add_axes([0.17, 0.185, 0.65, 0.65]) #left, bottom, width, height\n",
    "    cluster_count=0   #reset cluster_count to 0 as the j loop scans through all clusters to create a plot for given snapshot but i loop is for next snapshot\n",
    "    for j in range(total_clusters):  #total_clusters tells us how many clusters to look for\n",
    "        \n",
    "        x=importdata[cluster_count+1][\"x_tracked\"]\n",
    "        y=importdata[cluster_count+1][\"y_tracked\"]\n",
    "        xcm=importdata[cluster_count+1][\"xcm\"]\n",
    "        ycm=importdata[cluster_count+1][\"ycm\"]\n",
    "        z=importdata[cluster_count+1][\"z_tracked\"]\n",
    "        #cluster_mass=np.sum(importdata[cluster_count+1][\"mass_tracked\"]) #do this to take into account the mass in each cluster\n",
    "        #cluster_mass_initial=np.sum(initial_cluster_data[cluster_count+1][\"mass_tracked\"])\n",
    "        #s1=ax.scatter(x,y,label=f'{clustermass[cluster_count]:.2e}'+' M$_{{\\odot}}$',c=colors[cluster_count])\n",
    "        s1=ax.scatter(x,y,label=f'{clustermass[cluster_count]:.2e}'+' M$_{{\\odot}}$',c=next(colors).reshape(1,-1)) #use this if no. of clusters is large\n",
    "        #ax.scatter(np.abs(xcm),np.abs(ycm),c=\"black\")\n",
    "        ax.minorticks_on()\n",
    "        ax.tick_params(labelsize=12)\n",
    "        ax.set_xlabel(\"x (kpc)\",fontsize=12,labelpad=3)\n",
    "        ax.set_ylabel(\"y (kpc)\",fontsize=12,labelpad=-5)\n",
    "        #time_label= 'time = ' + f'{snaptime:.3f}' + ' Gyr'\n",
    "        #ax.set_title(\"Clusters at T=\"+str(time[plot_count])+\" in Myr\")\n",
    "        #ax.set_title(\"Clusters at \"+time_label)\n",
    "        cluster_count+=1\n",
    "    \n",
    "    handles,labels = ax.get_legend_handles_labels()\n",
    "    sorted_handles= [x for _,x in sorted(zip(clustermass,handles),reverse=True)] #sort the handles (the colors next to the labels) based on clustermass\n",
    "    sorted_legends= [x for _,x in sorted(zip(clustermass,labels),reverse=True)] #sort the labels based on the clustermass which is a list\n",
    "    #ax.legend(sorted_handles,sorted_legends,bbox_to_anchor=(1,0.5), loc='center left')\n",
    "    \n",
    "    #plt.tight_layout()\n",
    "    '''\n",
    "    part = gizmo.io.Read.read_snapshots(['all'],'snapshot_index', time[i]+snapshot_start, simulation_name=simname, simulation_directory=simdir, assign_hosts_rotation=True, assign_hosts=True)  \n",
    "    t = np.max(part['star'].prop('form.time'))  \n",
    "    \n",
    "    rGas = part['gas'].prop('host.distance.principal.cylindrical')[:,0]\n",
    "    zGas = part['gas'].prop('host.distance.principal.cylindrical')[:,1]\n",
    "    \n",
    "    xGas = part['gas'].prop('host.distance.principal.cartesian')[:,0]\n",
    "    yGas = part['gas'].prop('host.distance.principal.cartesian')[:,1]\n",
    "    zGas = part['gas'].prop('host.distance.principal.cartesian')[:,2]\n",
    "    \n",
    "    vxGas = part['gas'].prop('host.velocity.principal.cartesian')[:,0]\n",
    "    vyGas = part['gas'].prop('host.velocity.principal.cartesian')[:,1]\n",
    "    vzGas = part['gas'].prop('host.velocity.principal.cartesian')[:,2]\n",
    "    \n",
    "    mGas = part['gas']['mass']\n",
    "    rhoGas = part['gas']['density']\n",
    "    tGas = part['gas']['temperature']\n",
    "    idGas = part['gas']['id']\n",
    "    \n",
    "    i_gas = np.where((rGas <= bin_edge) & (np.fabs(zGas) <= 1.5) & (part['gas']['density']*((MsunToGm/KpcToCm**3)/mp) >= 10.) & (tGas <= 1e4))\n",
    "    \n",
    "    x = xGas[i_gas]\n",
    "    y = yGas[i_gas]\n",
    "    z = zGas[i_gas]\n",
    "    vx = vxGas[i_gas]\n",
    "    vy = vyGas[i_gas]\n",
    "    vz = vzGas[i_gas]\n",
    "    m = mGas[i_gas]\n",
    "    rho = part['gas'].prop('number.density')[i_gas]\n",
    "    id = part['gas']['id'][i_gas]\n",
    "    \n",
    "    ###########################################################################\n",
    "    #gas image (2d histogram)\n",
    "    ###########################################################################\n",
    "    #cold (< 10^4 K) gas in the midplane (|z| <= 1.5 kpc within bin_edge\n",
    "    v =  np.where((rGas <= bin_edge) & (np.fabs(zGas) <= 1.5) & (tGas <= 1e4))\n",
    "    face, xh, yh = np.histogram2d(part['gas'].prop('host.distance.principal.cartesian')[v,1][0],part['gas'].prop('host.distance.principal.cartesian')[v,0][0],bins=[bins,bins], weights=part['gas']['mass'][v])\n",
    "\n",
    "    ###########################################################################\n",
    "\n",
    "    \n",
    "    tracked_gas={\"v\":v,\"face\":face,\"xh\":xh,\"yh\":yh}\n",
    "    \n",
    "    file_name=simtype+\"gas_data_snapshot_\"+str(snapnumber)+\".pkl\"\n",
    "    with open(gas_datapath+file_name, 'wb') as output:\n",
    "      pickle.dump(tracked_gas, output)\n",
    "    print(\"\\n Stored the gas data for background plot of in the snapshot no.\",snap,\"to filename:\",file_name,\"\\n#####\\n\")\n",
    "    print(tracked_gas)\n",
    "    tracked_gas_all_snaps.update({snapnumber:tracked_gas})\n",
    "    '''\n",
    "    ##################\n",
    "    ##################\n",
    "    #Loading the gas data\n",
    "    file_name=simtype+\"_gas_data\"+str(snapnumber)+\".pkl\"\n",
    "    #print(\"Reading gas data from the file: \",file_name,end=\"\\r\",flush=True)\n",
    "    with open(gas_datapath+file_name, \"rb\") as input:\n",
    "      import_gasdata = pickle.load(input)\n",
    "    \n",
    "    v=import_gasdata[\"v\"]\n",
    "    face=import_gasdata[\"face\"]\n",
    "    xh=import_gasdata[\"xh\"]\n",
    "    yh=import_gasdata[\"yh\"]\n",
    "    snaptime=import_gasdata[\"snaptime\"]\n",
    "    norm = matplotlib.colors.LogNorm(vmin=1, vmax=1000) #the color range plotted\n",
    "    im = ax.imshow(face/(((bins[1]-bins[0])*1000)**2),origin='lower',interpolation='nearest',norm=norm,extent=(-25,25,-25,25),cmap='binary') \n",
    "    \n",
    "\n",
    "    #colorbar for the background gas density\n",
    "    cmap_gray = matplotlib.cm.get_cmap('binary')\n",
    "    norm1 = matplotlib.colors.LogNorm(vmin=1,vmax=1000)\n",
    "    cbar_ax1 = fig1.add_axes([0.04, 0.185, 0.04, 0.65]) # position of gray colorbar (left, bottom, width, height)\n",
    "    cb1 = fig1.colorbar(im, cax=cbar_ax1, ticklocation='left')\n",
    "    cb1.set_label('$\\Sigma$ (M$_{{\\odot}}$/pc$^2$)', labelpad=-5, fontsize=12)\n",
    "    \n",
    "    #plot a scale bar 5 kpc long\n",
    "    ax.plot([-22.5,-17.5], [22.5,22.5], 'k-', linewidth=5)\n",
    "    label1 = \"5 kpc\"\n",
    "    ax.text(-22.5, 20, label1, fontsize=12.5)\n",
    "    \n",
    "    #label the name of the galaxy on plot \n",
    "    label2=\"m12i mhdcv \"+simtype.upper() \n",
    "    ax.text(-23.5,-23.5,label2,fontsize=13.5) \n",
    "    \n",
    "    #label the time from the snapshot_times.txt file\n",
    "    label3 = 'time = ' + f'{snaptime:.3f}' + ' Gyr'\n",
    "    ax.text(3.5,21,label3,fontsize=13.5)  \n",
    "    plot_name=simtype+\"_gas_and_clusters\"+str(plot_count+snapshot_start)+\".png\"\n",
    "    fig1.savefig(plot_path+plot_name,bbox_inches='tight',dpi=200)\n",
    "    print(\"Generated the plot and saved the figure as: \",plot_name,end=\"\\r\",flush=True)\n",
    "    #print(\"###################\\n\")\n",
    "    plot_count+=1\n",
    "    fig1.clf()\n",
    "    plt.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Scanned image from the File:  sf-fire3-alpha01_gas_and_clusters690.png\n",
      "Saved the GIF into the File:  sf-fire3-alpha01.gif\n"
     ]
    }
   ],
   "source": [
    "from PIL import Image\n",
    "#import glob  #use it if you want to read all certain file type\n",
    "    \n",
    "imgs=[]\n",
    "for i in range(596,691): #startig and ending value+1 of the index that identifies different file names or imgs = glob.glob(\"*.png\") can do done as well\n",
    "    filename=simtype+\"_gas_and_clusters\"+str(i)+\".png\"\n",
    "    imgs.append(plot_path+filename)\n",
    "    print(\"Scanned image from the File: \",filename,end=\"\\r\",flush=True)    \n",
    "\n",
    "frames = []\n",
    "for i in imgs:\n",
    "    new_frame = Image.open(i)\n",
    "    frames.append(new_frame)\n",
    "# Save into a GIF file that loops forever\n",
    "gifname=simtype+\".gif\"\n",
    "frames[0].save(plot_path+gifname, format='GIF',\n",
    "               append_images=frames[1:],\n",
    "               save_all=True,\n",
    "               duration=200, loop=0)    #duration =200 looked good\n",
    "print(\"\\nSaved the GIF into the File: \",gifname)"
   ]
  }
 ],
 "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}