doc/source/analyzing/Particle_Trajectories.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One can create particle trajectories from a `DatasetSeries` object for a specified list of particles identified by their unique indices using the `particle_trajectories` method. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import glob\n",
    "from os.path import join\n",
    "\n",
    "import yt\n",
    "from yt.config import ytcfg\n",
    "\n",
    "path = ytcfg.get(\"yt\", \"test_data_dir\")\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's start off with a FLASH dataset containing only two particles in a mutual circular orbit. We can get the list of filenames this way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "my_fns = glob.glob(join(path, \"Orbit\", \"orbit_hdf5_chk_00[0-9][0-9]\"))\n",
    "my_fns.sort()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And let's define a list of fields that we want to include in the trajectories. The position fields will be included by default, so let's just ask for the velocity fields:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fields = [\"particle_velocity_x\", \"particle_velocity_y\", \"particle_velocity_z\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are only two particles, but for consistency's sake let's grab their indices from the dataset itself:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ds = yt.load(my_fns[0])\n",
    "dd = ds.all_data()\n",
    "indices = dd[\"all\", \"particle_index\"].astype(\"int\")\n",
    "print(indices)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which is what we expected them to be. Now we're ready to create a `DatasetSeries` object and use it to create particle trajectories: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ts = yt.DatasetSeries(my_fns)\n",
    "# suppress_logging=True cuts down on a lot of noise\n",
    "trajs = ts.particle_trajectories(indices, fields=fields, suppress_logging=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `ParticleTrajectories` object `trajs` is essentially a dictionary-like container for the particle fields along the trajectory, and can be accessed as such:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(trajs[\"all\", \"particle_position_x\"])\n",
    "print(trajs[\"all\", \"particle_position_x\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that each field is a 2D NumPy array with the different particle indices along the first dimension and the times along the second dimension. As such, we can access them individually by indexing the field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(trajs[\"all\", \"particle_position_x\"][0], trajs[\"all\", \"particle_position_y\"][0])\n",
    "plt.plot(trajs[\"all\", \"particle_position_x\"][1], trajs[\"all\", \"particle_position_y\"][1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can plot the velocity fields as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(trajs[\"all\", \"particle_velocity_x\"][0], trajs[\"all\", \"particle_velocity_y\"][0])\n",
    "plt.plot(trajs[\"all\", \"particle_velocity_x\"][1], trajs[\"all\", \"particle_velocity_y\"][1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to access the time along the trajectory, we use the key `\"particle_time\"`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(trajs[\"particle_time\"], trajs[\"particle_velocity_x\"][1])\n",
    "plt.plot(trajs[\"particle_time\"], trajs[\"particle_velocity_y\"][1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, if we know the particle index we'd like to examine, we can get an individual trajectory corresponding to that index:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "particle1 = trajs.trajectory_from_index(1)\n",
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(particle1[\"all\", \"particle_time\"], particle1[\"all\", \"particle_position_x\"])\n",
    "plt.plot(particle1[\"all\", \"particle_time\"], particle1[\"all\", \"particle_position_y\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at a more complicated (and fun!) example. We'll use an Enzo cosmology dataset. First, we'll find the maximum density in the domain, and obtain the indices of the particles within some radius of the center. First, let's have a look at what we're getting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ds = yt.load(\"enzo_tiny_cosmology/DD0046/DD0046\")\n",
    "slc = yt.SlicePlot(\n",
    "    ds,\n",
    "    \"x\",\n",
    "    [(\"gas\", \"density\"), (\"gas\", \"dark_matter_density\")],\n",
    "    center=\"max\",\n",
    "    width=(3.0, \"Mpc\"),\n",
    ")\n",
    "slc.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, so good--it looks like we've centered on a galaxy cluster. Let's grab all of the dark matter particles within a sphere of 0.5 Mpc (identified by `\"particle_type == 1\"`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "sp = ds.sphere(\"max\", (0.5, \"Mpc\"))\n",
    "indices = sp[\"all\", \"particle_index\"][sp[\"all\", \"particle_type\"] == 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we'll get the list of datasets we want, and create trajectories for these particles:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "my_fns = glob.glob(join(path, \"enzo_tiny_cosmology/DD*/*.hierarchy\"))\n",
    "my_fns.sort()\n",
    "ts = yt.DatasetSeries(my_fns)\n",
    "trajs = ts.particle_trajectories(indices, fields=fields, suppress_logging=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matplotlib can make 3D plots, so let's pick three particle trajectories at random and look at them in the volume:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8.0, 8.0))\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "ax.plot(\n",
    "    trajs[\"all\", \"particle_position_x\"][100],\n",
    "    trajs[\"all\", \"particle_position_y\"][100],\n",
    "    trajs[\"all\", \"particle_position_z\"][100],\n",
    ")\n",
    "ax.plot(\n",
    "    trajs[\"all\", \"particle_position_x\"][8],\n",
    "    trajs[\"all\", \"particle_position_y\"][8],\n",
    "    trajs[\"all\", \"particle_position_z\"][8],\n",
    ")\n",
    "ax.plot(\n",
    "    trajs[\"all\", \"particle_position_x\"][25],\n",
    "    trajs[\"all\", \"particle_position_y\"][25],\n",
    "    trajs[\"all\", \"particle_position_z\"][25],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It looks like these three different particles fell into the cluster along different filaments. We can also look at their x-positions only as a function of time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(trajs[\"all\", \"particle_time\"], trajs[\"all\", \"particle_position_x\"][100])\n",
    "plt.plot(trajs[\"all\", \"particle_time\"], trajs[\"all\", \"particle_position_x\"][8])\n",
    "plt.plot(trajs[\"all\", \"particle_time\"], trajs[\"all\", \"particle_position_x\"][25])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we wanted to know the gas density along the particle trajectory, but there wasn't a particle field corresponding to that in our dataset. Never fear! If the field exists as a grid field, yt will interpolate this field to the particle positions and add the interpolated field to the trajectory. To add such a field (or any field, including additional particle fields) we can call the `add_fields` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "trajs.add_fields([(\"gas\", \"density\")])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also could have included `\"density\"` in our original field list. Now, plot up the gas density for each particle as a function of time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6, 6))\n",
    "plt.plot(trajs[\"all\", \"particle_time\"], trajs[\"gas\", \"density\"][100])\n",
    "plt.plot(trajs[\"all\", \"particle_time\"], trajs[\"gas\", \"density\"][8])\n",
    "plt.plot(trajs[\"all\", \"particle_time\"], trajs[\"gas\", \"density\"][25])\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the particle trajectories can be written to disk. Two options are provided: ASCII text files with a column for each field and the time, and HDF5 files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "trajs.write_out(\n",
    "    \"halo_trajectories\"\n",
    ")  # This will write a separate file for each trajectory\n",
    "trajs.write_out_h5(\n",
    "    \"halo_trajectories.h5\"\n",
    ")  # This will write all trajectories to a single file"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}