In [13]:
from __future__ import division, print_function
# Third-party
from astropy import log as logger
import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
import vegas
%matplotlib inline
# Custom
import streams.coordinates as stc
import streamteam.io as io
import streamteam.dynamics as sd
import streamteam.potential as sp
from streamteam.units import galactic
from streams.inference.rewinder_likelihood import rewinder_likelihood
In [2]:
scf = io.APWSCFReader("/Users/adrian/projects/morphology/simulations/runs/spherical/")
ppars = scf.read_potential(units=galactic)
potential = sp.LeeSutoNFWPotential(v_h=ppars['v_h'], r_h=ppars['r_h'],
a=1., b=1., c=1., units=galactic)
# read all particles
tbl = scf.read_snap("SNAP006", units=galactic)
ptcl = io.tbl_to_w(tbl)
unbound_tbl = tbl[tbl['tub']!=0]
unbound_stars = ptcl[tbl['tub']!=0]
prog = np.median(ptcl[tbl['tub']==0], axis=0).reshape(1,6)
np.random.seed(42)
stars_ix = np.random.randint(len(unbound_stars), size=32)
stars_tbl = tbl[tbl['tub']!=0][stars_ix]
stars = unbound_stars[stars_ix]
stars.shape
Out[2]:
In [11]:
potential.c_instance.tidal_radius(2.5E8, np.array([[41.,0.,0.]]))
Out[11]:
In [9]:
Menc = 41 * potential.value(np.array([[41,0,0.]])) / potential.c_instance.G
41.*(2.5E8/np.abs(Menc))**(0.333333333333)
Out[9]:
In [3]:
betas = np.ones(len(stars_ix))
betas[stars[:,1] > (stars[:,0] + 22.)] = -1.
fig = sd.plot_orbits(ptcl, linestyle='none', alpha=0.4)
fig = sd.plot_orbits(prog, linestyle='none', marker='+', markeredgewidth=2.,
markersize=15, axes=fig.axes)
fig = sd.plot_orbits(stars[betas>0], linestyle='none', ms=10,
marker='o', color='g', axes=fig.axes)
fig = sd.plot_orbits(stars[betas<0], linestyle='none', ms=10,
marker='o', color='r', axes=fig.axes)
lims = -50,50
for ax in fig.axes:
ax.set_xlim(lims)
ax.set_ylim(lims)
x = np.linspace(-50,50,100)
fig.axes[0].plot(x, x+22, marker=None)
Out[3]:
In [4]:
def change_basis(w, prog_w, theta=0.):
x1_hat = prog_w[...,:3]
x3_hat = np.cross(prog_w[...,:3], prog_w[...,3:])
x2_hat = np.cross(x3_hat, x1_hat)
x1_hat = x1_hat / np.linalg.norm(x1_hat, axis=-1)[...,None]
x2_hat = x2_hat / np.linalg.norm(x2_hat, axis=-1)[...,None]
x3_hat = x3_hat / np.linalg.norm(x3_hat, axis=-1)[...,None]
x1_hat[...,0] = x1_hat[...,0]*np.cos(theta) + x1_hat[...,1]*np.sin(theta)
x1_hat[...,1] = -x1_hat[...,0]*np.sin(theta) + x1_hat[...,1]*np.cos(theta)
x2_hat[...,0] = x2_hat[...,0]*np.cos(theta) + x2_hat[...,1]*np.sin(theta)
x2_hat[...,1] = -x2_hat[...,0]*np.sin(theta) + x2_hat[...,1]*np.cos(theta)
dx = w[...,:3] - prog_w[...,:3]
dv = w[...,3:] - prog_w[...,3:]
x1 = dx[...,0]*x1_hat[...,0] + dx[...,1]*x1_hat[...,1] + dx[...,2]*x1_hat[...,2]
x2 = dx[...,0]*x2_hat[...,0] + dx[...,1]*x2_hat[...,1] + dx[...,2]*x2_hat[...,2]
x3 = dx[...,0]*x3_hat[...,0] + dx[...,1]*x3_hat[...,1] + dx[...,2]*x3_hat[...,2]
x = np.array([x1.T,x2.T,x3.T]).T
return x
In [6]:
# plot stars around tub
ub_t,ub_w = potential.integrate_orbit(np.vstack((prog, unbound_stars)), dt=-1., nsteps=6000)
xprime = change_basis(ub_w[:,1:], ub_w[:,0][:,np.newaxis], theta=-0.3)
ixes = (6000 - unbound_tbl['tub']).astype(int)
at_tub = np.array([xprime[ix,i] for i,ix in enumerate(ixes)])
fig = sd.plot_orbits(at_tub, marker='.', linestyle='none')
fig = sd.plot_orbits(at_tub[stars_ix], marker='o', color='r', linestyle='none', axes=fig.axes)
for ax in fig.axes:
ax.set_xlim(-3,3)
ax.set_ylim(-3,3)
In [8]:
%timeit ll = rewinder_likelihood(6000., 0., -1., \
potential.c_instance, \
prog, stars[:4], \
2.5E6, 0., \
1., betas, -0.3)
In [354]:
%timeit ll = rewinder_likelihood(6000., 0., -1., \
potential.c_instance, \
prog, stars, \
2.5E6, 0., \
1., betas, -0.3)
In [163]:
stars_tbl['tub']
Out[163]:
In [164]:
# plt.plot(ll, marker=None)
i = 2
plt.plot(ll[:,i], marker=None)
plt.axvline(6000 - stars_tbl['tub'][i])
plt.ylim(0,40)
# ll.max(axis=0)
Out[164]:
In [169]:
def trapezoid(X, dt):
return dt*(X[0] + np.sum(2*X[1:-1], axis=0) + X[-1])/2.
vhs = np.linspace(0.3,0.7,25)
l = []
for v_h in vhs:
p = sp.LeeSutoNFWPotential(v_h=v_h, r_h=20, a=1, b=1, c=1, units=galactic)
ll = rewinder_likelihood(6000., 0., -1.5, \
p.c_instance, \
prog, stars, \
2.5E6, 0., \
1., betas, -0.3)
l.append(np.log(trapezoid(np.exp(ll), dt=1.)).sum())
l = np.array(l)
fig,axes = plt.subplots(1,2,figsize=(10,5))
axes[0].plot(vhs, l, marker='o')
axes[1].plot(vhs, np.exp(l - l.max()), marker='o')
# plt.ylim(0,250)
Out[169]:
In [14]:
stars_hel = stc.gal_to_hel(stars)
sigma = np.zeros_like(stars_hel)
sigma[:,:2] = 1E-8
sigma[:,2] = stars_hel[:,2] * 0.02
sigma[:,3:5] = stars_hel[:,3:5]*0.1
sigma[:,5] = (10*u.km/u.s).to(u.kpc/u.Myr).value
sigma = np.abs(sigma)
obs_stars_hel = np.random.normal(stars_hel, sigma)
obs_stars = stc.hel_to_gal(obs_stars_hel)
In [43]:
test_gal = stc.hel_to_gal(np.random.normal(obs_stars_hel[0], sigma[0], size=(1000,6)))
rt = potential.c_instance.tidal_radius(2.5E6, prog)
# dPhi_dr_rt = potential.value(prog[:,:3] + rt*prog[:,:3]/np.linalg.norm(prog[:,:3])) - potential.value(prog[:,:3] - rt*prog[:,:3]/np.linalg.norm(prog[:,:3]))
r = np.linalg.norm(prog)
rh = potential.parameters['r_h']
vh = potential.parameters['v_h']
dPhi_dr = vh**2 / r * (np.log(1 + r/rh)/(r/rh) - rh/(r+rh))
dPhi_dr_rt = dPhi_dr * rt
print(dPhi_dr_rt)
plt.hist(np.abs(potential.energy(test_gal[:,:3], test_gal[:,3:]) - potential.energy(prog[:,:3], prog[:,3:])))
plt.axvline(3*dPhi_dr_rt)
Out[43]:
In [352]:
def integrand(w):
w = np.array(w).reshape(1,6)
log_p_D_W = norm.logpdf(w[0], loc=obs_stars_hel[0], scale=sigma[0]).sum()
x = stc.hel_to_gal(w)
ll = rewinder_likelihood(6000., 0., -1., \
potential.c_instance, \
prog, x, \
2.5E6, 0., \
1., betas, -0.3)
ll = np.log(trapezoid(np.exp(ll), dt=1.))
L = ll + log_p_D_W
return np.exp(L)
# def integrand(w):
# w = np.array(w).reshape(1,6)
# p_D_W = np.prod(norm.pdf(w[0], loc=obs_stars_hel[0], scale=sigma[0]))
# #x = stc.hel_to_gal(x)
# return p_D_W
print(integrand(stars_hel[0]))
In [353]:
bounds = zip(obs_stars_hel[0] - 5*sigma[0], obs_stars_hel[0] + 5*sigma[0])
integ = vegas.Integrator(bounds)
integ(integrand, nitn=10, neval=1000) # adapt the grid
result = integ(integrand, nitn=1, neval=10000) # estimate the integral
print(result.summary())
print('result = %s Q = %.2f' % (result, result.Q))
In [341]:
def f(w):
return np.prod(norm.pdf(w, loc=[15., 1., 0., 0., 0., 0.], scale=[0.01, 0.05, 0.01, 0.01, 0.01, 0.01]))
integ = vegas.Integrator([[14, 16], [0, 2], [-1, 1], [-1, 1], [-1, 1], [-1, 1]])
# adapt
result = integ(f, nitn=10, neval=10000)
# run
result = integ(f, nitn=1, neval=10000)
print(result.summary())
print('result = %s Q = %.2f' % (result, result.Q))
In [ ]: