#!/usr/bin/env python
# coding: utf-8

# # Pynbody Demo

# This notebook highlights some of the basic `pynbody` functionality to get you on the way to analyzing your own simulations. The `pynbody` webpage is [here]( http://pynbody.github.io/) and you can find the documentation with additional tutorials [here](http://pynbody.github.io/pynbody/). If you find that things are broken *please please please* let us know by submitting a [bug report on github](https://github.com/pynbody/pynbody/issues). If you want to run the notebook on your own machine, you need to get the [testdata tarball](https://github.com/pynbody/pynbody/releases) and change the path in the `load` function below accordingly. 

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
from matplotlib.pylab import *

rcParams['figure.figsize'] = (10,6)
rcParams['font.size'] = 18


# ## Basic data loading/exploration

# loading a simulation output using the `pynbody.load()` function, which tries to automatically determine which type of code output you have:

# In[2]:


import pynbody

s = pynbody.load('../nose/testdata/gasoline_ahf/g15784.lr.01024.gz')


# Note that the above assumes you have downloaded and unpacked the test data (in the nose folder). You can obtain it
# [here](http://star.ucl.ac.uk/~app/testdata.tar.gz), if you haven't already. Or, of course, you can just substitute your favourite cosmological simulation.
# A quick look at some basic information about the run we just opened:

# In[3]:


s


# In[4]:


len(s)


# In[5]:


len(s.stars)


# `stars`, `gas`, `dark` also available as `s`, `g`, `d`

# In[6]:


len(s.g), len(s.gas), len(s.dark)


# the `properties` attribute of a `SimSnap` tells us some more basic info 

# In[7]:


s.properties


# In[8]:


s.properties['time'].in_units('Gyr')


# which quantities do we have available?

# In[9]:


s.keys()


# None! Because pynbody "lazy-loads" data... so lets see which data is actually on-disk:

# In[10]:


s.loadable_keys()


# ## Accessing data

# to access any of these arrays or vectors, you access them like a python dictionary: 

# In[11]:


s['pos'].in_units('kpc')


# Note that each array has units attached...

# In[12]:


s['mass'], s['vel'], s.g['HeI'], s.s['posform']


# by default everything is in system units, but most of the time thinking in physical units is easier:

# In[13]:


s.physical_units()


# We have defined many useful quantities that are automatically calculated for you. For example, the radial and tangential velocities are simply obtained by

# In[14]:


s['vt'],s['vr']


# you can get a list of all available derivable quantities 

# In[15]:


s.derivable_keys()[0:10]


# ## Loading and centering/aligning halos

# if we have information about substructure (i.e. from the Amiga Halo Finder) we can load in the halos: 

# In[16]:


h = s.halos()


# We are not really near the main halo, so lets center the simulation there using the "shrinking sphere" method

# In[17]:


pynbody.analysis.halo.center(h[1],mode='ssc')


# we're going to work with the main halo a lot, so it's useful to save a reference to it in a variable to save some typing:

# In[18]:


h1 = h[1]
h1


# this new object is a `SubSnap` and behaves exactly like any other `SimSnap`:

# In[19]:


len(h1.g), len(h1.s)


# In[20]:


h[1]['mass'].sum().in_units('1e12 Msol')


# ## Rendering images

# lets look at a picture of the dark matter distribution of a part of the whole simulation...

# In[21]:


im = pynbody.plot.image(s.d,width='50 Mpc', units='Msol kpc^-2', cmap=cm.Greys)


# In[22]:


pynbody.plot.image(h[1].g,width='200 kpc', av_z=True, cmap=cm.Blues)


# Typically we want to make sure the system is aligned in some predictable way, so pynbody allows us to easily make the disk angular momentum axis the z-axis:

# In[23]:


pynbody.analysis.angmom.faceon(h[1],cen=(0,0,0))


# we use `cen=(0,0,0)` because we have already centered it in the previous step... if you don't specify a center, it is first centered and then aligned. Now lets look at a slice of the gas disk edge-on: 

# In[24]:


s.rotate_x(90)
pynbody.plot.image(h[1].g,width='50 kpc',cmap=cm.Blues)


# By default, `rho` (i.e. density) is used to render the image, but you can *any* valid array is fair game:

# In[25]:


pynbody.plot.image(h1.g,qty='temp',width='40 kpc', cmap=cm.YlOrRd)
s.rotate_x(-90)


# ## Profiles

# In[26]:


p = pynbody.analysis.profile.Profile(h1,max='200 kpc',min='.1 kpc',type='log',nbins=100)
ps = pynbody.analysis.profile.Profile(h1.s,max='50 kpc',nbins=50)
pg = pynbody.analysis.profile.Profile(h1.g,max='50 kpc',nbins=50)
plot(p['rbins'], p['density'], label = 'all')
plot(ps['rbins'], ps['density'], label = 'stars')
plot(pg['rbins'], pg['density'], label = 'gas')
semilogy()
xlim(0,50)
legend()
xlabel('$R$ [kpc]')
ylabel(r'$\Sigma$ [M$_{\odot}$ kpc$^{-2}$]')


# The `Profile` class has lots of built-in functionality... including calculating a rotation curve:

# In[27]:


plot(p['rbins'], p['v_circ'], label = 'all')
plot(ps['rbins'], ps['v_circ'], label = 'stars')
plot(pg['rbins'], pg['v_circ'], label = 'gas')
xlim(0,50)
xlabel('$R$ [kpc]')
ylabel('$V_{circ}$')
legend()


# `Profile` can use *any* loaded (or loadable/derivable!) array: 

# In[28]:


plot(ps['rbins'],ps['age'])
xlabel('$R$ [kpc]')
ylabel('Age [Gyr]')


# By default, `Profile` calculates the means of quantities, but it trivially calculate derivatives and root-mean-square values as well: 

# In[29]:


plot(ps['rbins'],ps['d_v_circ'])
xlabel('$R$ [kpc]')
ylabel('d$v_{c}$/d$R$')


# In[30]:


plot(ps['rbins'],ps['vr_rms'])
xlabel('R [kpc]')
ylabel('$v_{rms}$')


# ## Filters

# You can define filters to easily select particles based on their quantities. For example, to get stars that are 1-2 Gyr old:

# In[31]:


agefilt = pynbody.filt.BandPass('age', '1 Gyr','2 Gyr')
agefilt


# In[32]:


sphere = pynbody.filt.Sphere('50 kpc')
h1[sphere].s


# In[33]:


h1.s[sphere]['mass'].sum().in_units('1e12 Msol')


# you can combine filters by using simple logic operators -- so to get the total mass of all 1-2 Gyr old stars within 50 kpc, you can do

# In[34]:


h1.s[sphere&agefilt]['mass'].sum().in_units('1e12 Msol')


# ## Units

# You've noticed by now that we've been asking for things in various kinds of units. `Pynbody` has a unit-handling system built in that automatically ensures that arrays in different units are compatible and performs the necessary conversions for calculations. Adding mass and position together makes no sense, and pynbody knows this: 

# In[35]:


s['mass'] + s['x']


# but, it handles this just fine and returns an array with the right units: 

# In[ ]:


s['mass']*s['x']


# This concludes the basic overview of what ``pynbody`` can do -- for more tutorials and in-depth explanation of the inner-workings of the framework, head over to the [walkthroughs](http://pynbody.github.io/pynbody/tutorials/tutorials.html#walkthroughs) and [cookbook examples](http://pynbody.github.io/pynbody/tutorials/tutorials.html#cookbook-recipes). If you're stuck implementing a particular analysis workflow, stop by our [Github issues](https://github.com/pynbody/pynbody/issues) and post your question, we'll be happy to help.