{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "080cd826",
   "metadata": {},
   "source": [
    "# Working with Simulations\n",
    "\n",
    "Sometimes, you need to work with many snapshots across an entire simulation instead of a single snapshot at once. This notebook describes the routines included in `snapAnalysis.sim` for extracting useful data from entire simualtions. We'll work with the isolated dark matter halo snapshots included in the test suite. There's only two snapshots in this test simulation (for space reasons), but these recipes translate directly to simulations including thousands of snapshots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "03e48c8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snapanalysis import sim\n",
    "from snapanalysis.snap import snapshot\n",
    "import astropy.units as u"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a58fdf2",
   "metadata": {},
   "source": [
    "## Obtaining a list of available snapshots\n",
    "\n",
    "Sometimes, it's convenient to get a list of all the snapshot files in a particular simulation directory. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "de13eeb0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['../tests/example_snaps/cdm_snaps/snap_000.hdf5'\n",
      " '../tests/example_snaps/cdm_snaps/snap_001.hdf5']\n"
     ]
    }
   ],
   "source": [
    "sim_dir = \"../tests/example_snaps/cdm_snaps\"\n",
    "snaps = sim.get_snaps(sim_dir)\n",
    "print(snaps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2bcb947",
   "metadata": {},
   "source": [
    "## Extracting the center-of-mass orbit of a galaxy\n",
    "\n",
    "To extract the center-of-mass orbit of a galaxy, use the `orbit_com` function. In a simulation that includes multiple galaxies, you must specify the particle ID range that belongs to the galaxy of interest. You can also supply dictionaries of kwargs to the position and velocity center-finding routines. The columns in the output array are time, x, y, z, vx, vy, vz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9a6fa20e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.      ,  0.120152, -0.287273,  0.114837, -2.411327,  0.092466,\n",
       "        -1.306028],\n",
       "       [ 0.197182, -0.426582, -0.400833, -0.421794, -0.098579, -0.782643,\n",
       "        -0.2411  ]])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim.orbit_com(sim_dir, part_type=1, select_IDs=None, com_kwargs={\"vol_dec\": 1.1, \"r_start\":15.0*u.kpc, \"guess\": [0,0,0]*u.kpc})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f4e6100",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "3ca911f9",
   "metadata": {},
   "source": [
    "## Extracting rotations\n",
    "\n",
    "To extract the rotation matrices needed at each time to align a galaxy's angular momentum to the $z$-axis, use the `get_alignment_rotations` function. Again, you can specify an ID range to select a particular galaxy. By default, `get_alignment_rotations` will center each snapshot before calculating its angular momentum, but you can also specify a file of precomputed centers from `orbit_com` if you have one. The resulting array will be the stacked rotation matrices required to align each snapshot's angular momentum with $z$, of the shape N_snapshots x 3 x 3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0466beaf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/haydenfoote/Documents/gradschool/research/code/snapAnalysis/src/snapanalysis/snap.py:468: UserWarning: COM: Minimum number of particles reached before COM converged.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[[-0.95041252,  0.05663625, -0.30579141],\n",
       "        [-0.0594857 , -0.99822916,  0.        ],\n",
       "        [-0.3052499 ,  0.01819022,  0.95209853]],\n",
       "\n",
       "       [[ 0.53087062,  0.66868122, -0.52061676],\n",
       "        [-0.78319123,  0.62178091,  0.        ],\n",
       "        [ 0.32370956,  0.40774248,  0.85379048]]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sim.get_alignment_rotations(sim_dir, part_type=1, r_max=10*u.kpc, select_IDs=None, use_centers=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e310b65b",
   "metadata": {},
   "source": [
    "## Extracting particle orbits\n",
    "To extract the orbit of a specific particle(s), use `get_particle_orbit`. Again, by default this function will center and align the snapshots for you, but you may specify precomputed centers and rotations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0ac81e9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# get 10 random particle ID from the first snapshot\n",
    "s = snapshot(snaps[0], 1)\n",
    "s.load_particle_data([\"ParticleIDs\"])\n",
    "rand_ids = np.random.choice(s.data_fields[\"ParticleIDs\"], size=10, replace=False)\n",
    "\n",
    "# and extract their orbits\n",
    "orbits = sim.get_particle_orbit(sim_dir, part_type=1, ids=rand_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ebc0416",
   "metadata": {},
   "source": [
    "This function returns an `orbit` object, which keeps track of particle orbits. `orbit` objects are designed for ease of use with the the [Gala package](https://gala.adrian.pw/en/latest/index.html) for galactric dynamics. You can access the data as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c94bd29e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.         0.19718176] Gyr\n"
     ]
    }
   ],
   "source": [
    "# The times of each stored point in the orbit are stored in\n",
    "print(orbits.t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "95742ca0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(2, 3, 10)\n",
      "(2, 3, 10)\n"
     ]
    }
   ],
   "source": [
    "# positions and velocities of each particle are stored in orbits.pos and orbits.vel.\n",
    "# these have shape (N_times, 3, N_particles) for consistency with Gala.\n",
    "print(orbits.pos.shape)\n",
    "print(orbits.vel.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9eea01da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'1881': 0, '3490': 1, '4334': 2, '4556': 3, '4881': 4, '6363': 5, '6548': 6, '7777': 7, '8855': 8, '9644': 9}\n"
     ]
    }
   ],
   "source": [
    "# The IDs attribute is a dictionary whose keys are the particle IDs, \n",
    "# and whose values are the indices of those particles in the data arrays.\n",
    "print(orbits.ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f50f9d15",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[-14.05476249 -11.31392662 -15.47666254]\n",
      " [ 11.23177433   2.7399446  -12.70862759]] kpc\n"
     ]
    }
   ],
   "source": [
    "# Therefore, to access the position of a particular particle at all times, \n",
    "id = str(int(rand_ids[0]))\n",
    "print(orbits.pos[:, :, orbits.ids[id]])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}