<jupyter_start><jupyter_text>You'll need a few packages for this to work. You can install everything you need by running these two commands:    conda install astropy numpy h5py matplotlib tqdm    pip install astro-gala pyia<jupyter_code>import sys
from os import path
import warnings

# Third-party
from astropy.table import Table
from astropy.io import fits
import astropy.coordinates as coord
import astropy.units as u
import h5py
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline
from tqdm import tqdm

import gala.dynamics as gd
import gala.integrate as gi
import gala.potential as gp
from gala.units import galactic
from pyia import GaiaData<jupyter_output>/Users/adrian/anaconda/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
  from ._conv import register_converters as _register_converters<jupyter_text>We'll use a simple model for the Milky Way potential that is implemented in `Gala` that contains components for the Galactic disk, bulge, nucleus, and dark matter halo:<jupyter_code>mw = gp.MilkyWayPotential()
H = gp.Hamiltonian(mw)
print(mw.keys())<jupyter_output>odict_keys(['disk', 'bulge', 'nucleus', 'halo'])<jupyter_text>We'll need to transform observed positions and velocities to the Galactocentric rest frame. To do that, we have to make assumptions about the solar velocity and position. We'll assume the Sun is at:$$\begin{align}\boldsymbol{r}_{\odot} &= (-8, 0, 0)~{\rm kpc}\\\boldsymbol{v}_{\odot} &= (11.1, 232.24, 7.25)~{\rm kpc}\end{align}$$(but feel free to play with the definitions if you prefer other values).<jupyter_code>rsun = 8 * u.kpc
vsun = [11.1, 232.24, 7.25] * u.km/u.s
gc_frame = coord.Galactocentric(galcen_distance=rsun,
                                galcen_v_sun=coord.CartesianDifferential(*vsun))<jupyter_output><empty_output><jupyter_text>We next need to load some Gaia data. For now, I'll load a mock (simulated_ dataset with DR2-like uncertainties for all stars within 200 pc of the Sun:<jupyter_code>g = GaiaData('/Users/adrian/data/GaiaDR2-mock/Gaia-DR2-mock-200pc.fits')<jupyter_output><empty_output><jupyter_text>We can use this object to get an Astropy sky coordinate object, which has the sky positions, distance (parallax), proper motions, and radial velocity automatically filled:<jupyter_code>c = g.skycoord
c<jupyter_output><empty_output><jupyter_text>**Note: not all Gaia DR2 stars will have radial velocities, so you may have to filter out those stars here**Next, we will transform these heliocentric positions/velocities into Galactocentric values, then pass in to a Gala class that will handle computing dynamical quantities from the Galactocentric Cartesian phase-space positions:<jupyter_code>w = gd.PhaseSpacePosition(c.transform_to(gc_frame).cartesian)<jupyter_output><empty_output><jupyter_text>With this object, we can do things like compute the angular momentum:<jupyter_code>L = w.angular_momentum()
L.shape<jupyter_output><empty_output><jupyter_text>Or, given a model for the gravitational potential, compute the energy of the stellar orbits:<jupyter_code>E = H.energy(w)<jupyter_output><empty_output><jupyter_text>Let's now plot the energy vs. $z$-component of the angular momentum (index 2):<jupyter_code>L_unit = u.km/u.s * u.kpc
E_unit = u.km/u.s * u.kpc/u.Myr

fig, ax = plt.subplots(1, 1, figsize=(6, 5))

ax.plot(L[2].to(L_unit).value,
        E.to(E_unit).value, 
        linestyle='none', marker=',', alpha=0.2)

ax.set_xlim(-3000, -500)
ax.set_ylim(-160, -100)
ax.set_xlabel('$L_z$ [{0:latex_inline}]'.format(L_unit))
ax.set_ylabel('$E$ [{0:latex_inline}]'.format(E_unit))

fig.tight_layout()<jupyter_output><empty_output><jupyter_text>That's a lot of points! Let's instead make a histogram so we can look at the log-density of points:<jupyter_code>E_grid = np.linspace(-160, -100, 128)
Lz_grid = np.linspace(-3000, -500, 128)
H, xedg, yedg = np.histogram2d(L[2].to(L_unit).value, E.to(E_unit).value,
                               bins=(Lz_grid, E_grid))
norm = mpl.colors.LogNorm(vmin=1e-1, vmax=1E5)

fig, ax = plt.subplots(1, 1, figsize=(6, 5))

ax.pcolormesh(xedg, yedg, H.T, norm=norm, cmap='Blues')

ax.set_xlabel('$L_z$ [{0:latex_inline}]'.format(L_unit))
ax.set_ylabel('$E$ [{0:latex_inline}]'.format(E_unit))

fig.tight_layout()<jupyter_output><empty_output><jupyter_text>That's a very smooth distribution! The real Galaxy probably won't look like that...but we'll see! We can also compute other quantities for these stars, like the actions. These are other integrals of motion that are useful because they are adiabatically invariant, and because they are conserved in a static potential (unlike the angular momentum components which can vary if the orbit is not planar and the potential is non-spherical). The problem is that, except in very simple potential models, computing the actions has to be done numerically. There are many algorithms out there for estimating actions (see papers by <NAME>, <NAME>, <NAME>). Here, with Gala, we'll use a method that requires numerically integrating orbits in order to compute the actions. This makes it quite slow to run for millions of stars, but we can at least run for a subset of stars as a demo. In practice, you can parallelize this or run on batches (subsets) of stars.  How do we choose an integration timestep?One thing we have to choose when numerically estimating the actions is the timestep and length of orbit integration. We'll set the length to 5 Gyr (~20 complete orbits of a sun-like orbit around the Galaxy), and vary the timestep to see if the value of the actions converges:<jupyter_code>all_actions = []
dts = [0.1, 0.2, 0.4, 0.8, 1., 2., 4, 8] * u.Myr

with warnings.catch_warnings(record=True):
    warnings.simplefilter("ignore")
    for dt in tqdm(dts):
        orbit = mw.integrate_orbit(w[0], dt=dt, t1=0*u.Gyr, t2=5*u.Gyr,
                                   Integrator=gi.DOPRI853Integrator)
        res = gd.actionangle.find_actions(orbit, N_max=8)
        all_actions.append(res['actions'])
act = u.Quantity(all_actions)

plt.figure(figsize=(6, 5))

for k in range(3):
    plt.plot(dts[1:], np.abs((act[1:, k]-act[0, k])/act[0, k]), 
             label='$J_{0}$'.format(k+1))
    
plt.xscale('log')
plt.yscale('log')
plt.xlabel('timestep [Myr]')
plt.ylabel('fractional error')
plt.legend(loc='best', fontsize=16)
plt.tight_layout()<jupyter_output><empty_output><jupyter_text>From this, it looks like we can set the timestep to 2 Myr and only suffer a fractional error of $10^{-5}$. How long does computing the actions for one orbit take?<jupyter_code>%%time
orbit = mw.integrate_orbit(w[0], dt=2*u.Myr, t1=0*u.Gyr, t2=5*u.Gyr,
                           Integrator=gi.DOPRI853Integrator)
res = gd.actionangle.find_actions(orbit, N_max=8)<jupyter_output>CPU times: user 594 ms, sys: 28.5 ms, total: 623 ms
Wall time: 429 ms<jupyter_text>~0.5 seconds! Let's run on a subset (128) of the orbits:<jupyter_code>some_w = w[:128]
orbits = mw.integrate_orbit(some_w, dt=1*u.Myr, t1=0*u.Gyr, t2=4*u.Gyr,
                            Integrator=gi.DOPRI853Integrator)

with warnings.catch_warnings(record=True):
    warnings.simplefilter("ignore")

    all_actions = []
    for orbit in tqdm(orbits.orbit_gen(), total=some_w.shape[0]):
        res = gd.actionangle.find_actions(orbit, N_max=8)
        all_actions.append(res['actions'])
    all_actions = u.Quantity(all_actions)<jupyter_output>100%|██████████| 128/128 [01:20<00:00,  1.59it/s]<jupyter_text>This is a simulated data set, and I don't remember what kind of age-action relations were put in, but we generally expect the vertical action, $J_z$, to increase with stellar age. Let's see if that's the case for these simulated stars:<jupyter_code>plt.figure(figsize=(6, 5))
plt.scatter(g.age[:some_w.shape[0]],  
            all_actions[:, 2].to(L_unit))
plt.xscale('log')
plt.yscale('log')
plt.xlim(1E-2, 20)
plt.ylim(1E-4, 1e2)
plt.xlabel('age [Gyr]')
plt.ylabel('$J_z$ [{0:latex_inline}]'.format(L_unit))<jupyter_output><empty_output><jupyter_text>Older stars definitely tend to have larger values of vertical action! Another way to look at this is by looking at the maximum height a star reaches above the plane, $|z_{\rm max}|$:<jupyter_code>plt.figure(figsize=(6, 5))
plt.scatter(g.age[:some_w.shape[0]],  
            orbits.zmax(approximate=True).to(u.pc).value)
plt.xscale('log')
plt.yscale('log')
plt.xlim(1E-2, 20)
plt.ylim(10, 2e3)
plt.xlabel('age [Gyr]')
plt.ylabel(r'$\left|z_{\rm max}\right|$ ' + '[{0:latex_inline}]'.format(L_unit))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12158752007341134, 'fraction_numerical': 0.02557926129846295, 'mean_word_length': 4.895199459093982, 'pattern_counts': {'":': 0, '<': 61, '<?xml version=': 0, '>': 60, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '15', 'dir': 'jupyter-structured-clean-dedup', 'id': '544415', 'license': 'MIT', 'n_tokens_mistral': 3235, 'n_tokens_neox': 2998, 'n_words': 888, 'path': 'jay-z/Vertical-action-demo.ipynb', 'repo_name': 'adrn/dr2-zero-day-exploits', 'size': '8732'}