TNG_extracting_data_example

TNG_extracting_data_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we try to use local TNG data to extract galaxy (subhalo) properties.\n",
    "\n",
    "Ref. https://github.com/duncandc/Illustris_Shapes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"/work8/astro/projects/TNG/notebooks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from illustris_python.snapshot import loadSubhalo\n",
    "from illustris_python.groupcat import loadSubhalos\n",
    "from simulation_props import sim_prop_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_particles(center, coords, Lbox):\n",
    "    \"\"\"\n",
    "    center the particle coordinates on (0,0,0) and account for PBCs\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    center : array_like\n",
    "        array of shape (3,) storing the coordinates of the\n",
    "        center of the particle distribution\n",
    "\n",
    "    coords :  array_like\n",
    "        array of shape (nptcls, 3) storing the coordinates of the\n",
    "        particles\n",
    "\n",
    "    Lbox : array_like\n",
    "        length 3-array giving the box size in each dimension\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    coords : numpy.array\n",
    "        array of shape (nptcls, 3) storing the centered particle coordinates\n",
    "    \"\"\"\n",
    "\n",
    "    dx = coords[:,0] - center[0]\n",
    "    dy = coords[:,1] - center[1]\n",
    "    dz = coords[:,2] - center[2]\n",
    "\n",
    "    # x-coordinate\n",
    "    mask = (dx > Lbox[0]/2.0)\n",
    "    dx[mask] = dx[mask] - Lbox[0]\n",
    "    mask = (dx < -Lbox[0]/2.0)\n",
    "    dx[mask] = dx[mask] + Lbox[0]\n",
    "\n",
    "    # y-coordinate\n",
    "    mask = (dy > Lbox[1]/2.0)\n",
    "    dy[mask] = dy[mask] - Lbox[1]\n",
    "    mask = (dy < -Lbox[1]/2.0)\n",
    "    dy[mask] = dy[mask] + Lbox[1]\n",
    "\n",
    "    # z-coordinate\n",
    "    mask = (dz > Lbox[2]/2.0)\n",
    "    dz[mask] = dz[mask] - Lbox[2]\n",
    "    mask = (dz < -Lbox[2]/2.0)\n",
    "    dz[mask] = dz[mask] + Lbox[2]\n",
    "\n",
    "    # format coordinates\n",
    "    coords = np.vstack((dx,dy,dz)).T\n",
    "\n",
    "    return coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def particle_selection(gal_id, ptcl_coords, galaxy_table, basePath, snapNum, radial_mask):\n",
    "    \"\"\"\n",
    "    Apply the selection criteria to galaxy particles\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    mask : numpy.array\n",
    "       boolean array of shape (nptcls,)\n",
    "    \"\"\"\n",
    "\n",
    "    # don't use wind particles\n",
    "    sf_time = loadSubhalo(basePath, snapNum, gal_id, 4, fields=['GFM_StellarFormationTime'])\n",
    "#smhr_start    \n",
    "    #star_mask = (sf_time>=0.0) # don't use wind particles\n",
    "    star_mask = (sf_time>0.0) # don't use wind particles\n",
    "#smhr_end\n",
    "    # get the half mass radius\n",
    "    gal_rhalfs = galaxy_table['SubhaloHalfmassRadType'][:,4]/1000.0\n",
    "    gal_rhalf = gal_rhalfs[gal_id]\n",
    "\n",
    "    # use only particles within 2 * R_half\n",
    "    r = np.sqrt(np.sum(ptcl_coords**2, axis=1))/gal_rhalf\n",
    "    if radial_mask:\n",
    "        radial_mask = (r<=2.0) # <------ radial criterion\n",
    "    else:\n",
    "        radial_mask = np.array([True]*len(star_mask))\n",
    "\n",
    "    return (radial_mask) & (star_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def galaxy_center(gal_id, galaxy_table):\n",
    "    \"\"\"\n",
    "    Return the coordinates of the center of the galaxy\n",
    "    \"\"\"\n",
    "\n",
    "    # load position of the most bound particle (of any type)\n",
    "    coords = galaxy_table['SubhaloPos']/1000.0\n",
    "    coord = coords[gal_id]\n",
    "\n",
    "    return coord"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def galaxy_load(gal_id, galaxy_table, basePath, snapNum, Lbox, radial_mask):\n",
    "    \"\"\"\n",
    "    Parameters\n",
    "    ----------\n",
    "    gal_id : int\n",
    "\n",
    "    basepath : string\n",
    "\n",
    "    snapNum : int\n",
    "\n",
    "    Lbox : array_like\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    xyz, mass\n",
    "    \"\"\"\n",
    "\n",
    "    # choose a 'center' for each galaxy\n",
    "    gal_position = galaxy_center(gal_id, galaxy_table)\n",
    "\n",
    "    # load stellar particle positions and masses\n",
    "    ptcl_coords = loadSubhalo(basePath, snapNum, gal_id, 4, fields=['Coordinates'])/1000.0\n",
    "    ptcl_masses = loadSubhalo(basePath, snapNum, gal_id, 4, fields=['Masses'])*10.0**10\n",
    "\n",
    "    # center and account for PBCs\n",
    "    ptcl_coords = format_particles(gal_position, ptcl_coords, Lbox)\n",
    "\n",
    "    ptcl_mask = particle_selection(gal_id, ptcl_coords, galaxy_table, basePath, snapNum, radial_mask)\n",
    "\n",
    "    return ptcl_coords[ptcl_mask], ptcl_masses[ptcl_mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def galaxy_selection(min_mstar, basePath, snapNum):\n",
    "    \"\"\"\n",
    "    make a cut on galaxy properties\n",
    "    \"\"\"\n",
    "\n",
    "    # make selection\n",
    "    galaxy_table = loadSubhalos(basePath, snapNum, fields=['SubhaloGrNr', 'SubhaloMassType'])\n",
    "\n",
    "    gal_ids = np.arange(0,len(galaxy_table['SubhaloGrNr']))\n",
    "\n",
    "    # total mass of stellar particles\n",
    "    mstar = galaxy_table['SubhaloMassType'][:,4]\n",
    "    mstar = mstar*10**10\n",
    "\n",
    "    mask = (mstar >= min_mstar)\n",
    "\n",
    "    return mask, gal_ids[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use above functions to extract galaxy data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim_name = 'TNG50-1'\n",
    "snapNum = 99"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# make galaxy selection according to their stellar mass\n",
    "h = 0.6774\n",
    "basePath = '/work8/astro/TNG/TNG50-1/output'\n",
    "m_dm = 3.07367708626464e-05*10**10\n",
    "Lbox = np.array([35.0]*3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the minimum stellar mass criterion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_mstar = 10**10 # We set 10^10 solar mass for the minimum."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask, gal_ids = galaxy_selection(min_mstar*h, basePath, snapNum)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of galaxies in selection: 903\n"
     ]
    }
   ],
   "source": [
    "# number of galaxies in selection\n",
    "Ngals = gal_ids.size\n",
    "print(\"number of galaxies in selection: {0}\".format(Ngals))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load galaxy table\n",
    "fields = ['SubhaloGrNr', 'SubhaloMassType', 'SubhaloPos', 'SubhaloHalfmassRadType']\n",
    "galaxy_table = loadSubhalos(basePath, snapNum, fields=fields)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As an example we load subhalo 13. We can set radial mask to just consider stars inside a factor of half-mass radius. See function `particle_selection`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "xyz, mass = galaxy_load(13, galaxy_table, basePath, snapNum, Lbox, radial_mask=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(366978, 3)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xyz.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(366978,)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mass.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "xyz, mass = galaxy_load(13, galaxy_table, basePath, snapNum, Lbox, radial_mask=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(260842, 3)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xyz.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(260842,)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mass.shape"
   ]
  }
 ],
 "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
}