In [1]:
import numpy
import tqdm
import wendy
%pylab inline
import matplotlib.animation as animation
from matplotlib import cm
from IPython.display import HTML
import copy
_SAVE_GIFS= False
rcParams.update({'axes.labelsize': 17.,
              'font.size': 12.,
              'legend.fontsize': 17.,
              'xtick.labelsize':15.,
              'ytick.labelsize':15.,
              'text.usetex': _SAVE_GIFS,
              'figure.figsize': [5,5],
              'xtick.major.size' : 4,
              'ytick.major.size' : 4,
              'xtick.minor.size' : 2,
              'ytick.minor.size' : 2,
              'legend.numpoints':1})
numpy.random.seed(2)
Populating the interactive namespace from numpy and matplotlib

The Gaia phase-space spiral

One of the amazing early discoveries in the Gaia DR2 data set is the Gaia phase-space spiral. This is a spiral feature in the vertical phase-space distribution function $f(z,v_z)$ first found by Antoja et al. (2018). In this example, we investigate how such a phase-space spiral can form from a simple perturbation to the Milky Way's disk. We also use it as an opportunity to showcase the support for arbitrary external forces and for different sorting algorithms in wendy's approximate N-body solution.

A simple model for the phase-space spiral is that it results from the disk's equilibrium $f(z,v_z)$ being offset from $\langle v_z \rangle =0$ (Darling & Widrow 2019). As a highly simplified model for this, we initialize a self-gravitating $\mathrm{sech}^2$ disk and only treat a fraction $\alpha$ of that as self-gravitating masses, filling in the rest as a static, external potential. This simplification makes it straightforward to setup the equilibrium solution, because that is just that of a fully self-gravitating, $\mathrm{sech}^2$ disk.

First we sample $N$ points from the disk:

In [2]:
# Initial disk
N= 100000
# compute zh based on sigma and totmass
totmass= 1. # Sigma above
sigma= 1.
zh= sigma**2./totmass # twopiG = 1. in our units
tdyn= zh/sigma
x= numpy.arctanh(2.*numpy.random.uniform(size=N)-1)*zh*2.
v= numpy.random.normal(size=N)*sigma
v-= numpy.mean(v) # stabilize
m= totmass*numpy.ones_like(x)/N

We then assign only $\alpha$ of that disk to be self-gravitating and define the external force:

In [3]:
alpha= 0.3 # "live" fraction
# Adjust masses to only represent alpha of the mass
m*= alpha
# 1-alpha in the mass is then given by the external force
sigma2= sigma**2.
def iso_force(x,t):
    return -(1.-alpha)*sigma2*numpy.tanh(0.5*x/zh)/zh

We then run the $N$-body simulation, using the approximate algorithm with an external force and using a fast quicksort implementation to calculate the $N$-body forces. For a simple external force implemented using numpy functions as we have done here, numba is used to compile C byte-code which is directly called in the underlying C code, allowing this external force to be very efficiently added to the system (don't worry, external forces that cannot be automatically compiled to C code are also supported, but they are slightly slower).

In [4]:
g= wendy.nbody(x,v,m,0.05*tdyn,nleap=10,approx=True,sort='quick',ext_force=iso_force)
In [5]:
nt= 1000
xt= numpy.empty((N,nt+1))
vt= numpy.empty((N,nt+1))
xt[:,0]= x
vt[:,0]= v
x_init= copy.copy(x)
v_init= copy.copy(v)
for ii in tqdm.trange(nt):
    tx,tv= next(g)
    xt[:,ii+1]= tx
    vt[:,ii+1]= tv
100%|██████████| 1000/1000 [01:44<00:00,  9.54it/s]

We check that the original disk is indeed in equilibrium:

In [6]:
figsize(6,4)
fig, ax= subplots()
ii= 0
a= ax.hist(xt[:,ii],bins=31,histtype='step',lw=1.,color='k',range=[-8.,8.],weights=31./16./N*numpy.ones(N))
xs= numpy.linspace(-8.,8.,101)
ax.plot(xs,totmass/4./zh/numpy.cosh(xs/2./zh)**2.,'b--',lw=2.,zorder=0)
ax.set_xlim(-8.,8.)
ax.set_ylim(10.**-3.,1.)
ax.set_xlabel(r'$x$')
ax.set_ylabel(r'$\rho(x)$')
ax.set_yscale('log')
ax.annotate(r'$t=0$',(0.95,0.95),xycoords='axes fraction',
             horizontalalignment='right',verticalalignment='top',size=18.)
subsamp= 4
def animate(ii):
    ax.clear()
    a= ax.hist(xt[:,ii*subsamp],bins=31,histtype='step',lw=1.,color='k',range=[-8.,8.],weights=31./16./N*numpy.ones(N))
    xs= numpy.linspace(-8.,8.,101)
    ax.plot(xs,totmass/4./zh/numpy.cosh(xs/2./zh)**2.,'b--',lw=2.,zorder=0)
    ax.set_xlim(-8.,8.)
    ax.set_ylim(10.**-3.,1.)
    ax.set_xlabel(r'$x$')
    ax.set_ylabel(r'$\rho(x)$')
    ax.set_yscale('log')
    ax.annotate(r'$t=%.0f$' % (ii*subsamp/20.),
                (0.95,0.95),xycoords='axes fraction',
                horizontalalignment='right',verticalalignment='top',size=18.)
    return a[2]
anim = animation.FuncAnimation(fig,animate,#init_func=init_anim_frame,
                               frames=nt//subsamp,interval=40,blit=True,repeat=True)
# The following is necessary to just get the movie, and not an additional initial frame
plt.close()
out= HTML(anim.to_html5_video())
plt.close()
out
Out[6]:

Indeed, the disk is in equilibrium!

Next, we offset all of the initial velocities by $1\sigma$ and run the simulation again to study the effect of this perturbation. This time we use timsort, a version of Python's own sorting algorithm (typically sort='quick' is in fact the fastest method; wendy also supports sort='merge' for a mergesort, sort='qsort' for the C standard library's own sorting algorithm, and sort='parallel' for a parallel implementation of mergesort).

In [7]:
x= copy.copy(x_init)
v= copy.copy(v_init)+sigma
g= wendy.nbody(x,v,m,0.05*tdyn,nleap=10,approx=True,sort='tim',ext_force=iso_force)
In [8]:
nt= 1000
xt= numpy.empty((N,nt+1))
vt= numpy.empty((N,nt+1))
xt[:,0]= x
vt[:,0]= v
for ii in tqdm.trange(nt):
    tx,tv= next(g)
    xt[:,ii+1]= tx
    vt[:,ii+1]= tv
100%|██████████| 1000/1000 [01:41<00:00,  9.82it/s]

Now we plot the evolution of the phase-space distribution of time, color-coding the points by their initial energy in the unperturbed gravitational field, and we see that a strong spiral quickly develops and winds up over time. The spiral develops, because the starts of different energies have different frequencies and they therefore orbit on different times. The frequency goes down with increasing energy (or the period goes up), so the result is a winding spiral pattern:

In [9]:
def init_anim_frame():
    line1= plot([],[])
    xlabel(r'$x$')
    ylabel(r'$v$')
    xlim(-7.99,7.99)
    ylim(-4.99,4.99)
    return (line1[0],)
figsize(6,4)
fig, ax= subplots()
# Directly compute the initial energy from the known sech^2 disk potential
c= v_init**2./2.+2.*sigma2*numpy.log(numpy.cosh(0.5*x_init/zh))
s= 5.*((c-numpy.amin(c))/(numpy.amax(c)-numpy.amin(c))*2.+1.)
line= ax.scatter(x,v,c=c,s=s,edgecolors='None',cmap=cm.jet_r)
txt= ax.annotate(r'$t=%.0f$' % (0.),
                 (0.95,0.95),xycoords='axes fraction',
                 horizontalalignment='right',verticalalignment='top',size=18.)
subsamp= 4
def animate(ii):
    line.set_offsets(numpy.array([xt[:,ii*subsamp],vt[:,ii*subsamp]]).T)
    txt.set_text(r'$t=%.0f$' % (ii*subsamp/20.))
    return (line,)
anim = animation.FuncAnimation(fig,animate,init_func=init_anim_frame,
                               frames=nt//subsamp,interval=40,blit=True,repeat=True)
if _SAVE_GIFS:
    anim.save('phasespiral_phasespace.gif',writer='imagemagick',dpi=80)
# The following is necessary to just get the movie, and not an additional initial frame
plt.close()
out= HTML(anim.to_html5_video())
plt.close()
out
Out[9]: