209 gary/dynamics/mockstream/core.py
View file
@@ -4,10 +4,6 @@

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os
import sys

# Third-party
from astropy import log as logger
import astropy.units as u
@@ -19,7 +15,7 @@
from ...integrate import DOPRI853Integrator, LeapfrogIntegrator
from ._mockstream import _mock_stream_dop853#, _mock_stream_leapfrog

__all__ = ['mock_stream']
__all__ = ['mock_stream', 'streakline_stream', 'fardal_stream', 'dissolved_fardal_stream']

def mock_stream(potential, w0, prog_mass, k_mean, k_disp,
t_f, dt=1., t_0=0., release_every=1,
@@ -59,6 +55,12 @@ def mock_stream(potential, w0, prog_mass, k_mean, k_disp,
Integrator to use.
Integrator_kwargs : dict (optional)
Any extra keyword argumets to pass to the integrator function.
Returns
-------
prog_orbit : `~gary.dynamics.CartesianOrbit`
stream : `~gary.dynamics.CartesianPhaseSpacePosition`
"""

# ------------------------------------------------------------------------
@@ -111,3 +113,200 @@ def mock_stream(potential, w0, prog_mass, k_mean, k_disp,
raise RuntimeError("Should never get here...")

return prog_orbit, CartesianPhaseSpacePosition.from_w(w=stream_w.T, units=potential.units)

def streakline_stream(potential, w0, prog_mass, t_f, dt=1., t_0=0., release_every=1,
Integrator=LeapfrogIntegrator, Integrator_kwargs=dict()):
"""
Generate a mock stellar stream in the specified potential with a
progenitor system that ends up at the specified position.
This uses the Streakline method from Kuepper et al. (2012).
Parameters
----------
potential : `~gary.potential.PotentialBase`
The gravitational potential.
w0 : `~gary.dynamics.PhaseSpacePosition`, array_like
Initial conditions.
prog_mass : numeric, array_like
A single mass or an array of masses if the progenitor mass evolves
with time.
t_f : numeric
The final time for integrating.
t_0 : numeric (optional)
The initial time for integrating -- the time at which ``w0`` is the position.
dt : numeric (optional)
The time-step.
release_every : int (optional)
Release particles at the Lagrange points every X timesteps.
Integrator : `~gary.integrate.Integrator` (optional)
Integrator to use.
Integrator_kwargs : dict (optional)
Any extra keyword argumets to pass to the integrator function.
Returns
-------
prog_orbit : `~gary.dynamics.CartesianOrbit`
stream : `~gary.dynamics.CartesianPhaseSpacePosition`
"""
k_mean = np.zeros(6)
k_disp = np.zeros(6)

k_mean[0] = 1. # R
k_disp[0] = 0.

k_mean[1] = 0. # phi
k_disp[1] = 0.

k_mean[2] = 0. # z
k_disp[2] = 0.

k_mean[3] = 0. # vR
k_disp[3] = 0.

k_mean[4] = 1. # vt
k_disp[4] = 0.

k_mean[5] = 0. # vz
k_disp[5] = 0.

return mock_stream(potential=potential, w0=w0, prog_mass=prog_mass,
k_mean=k_mean, k_disp=k_disp,
t_f=t_f, dt=dt, t_0=t_0, release_every=release_every,
Integrator=Integrator, Integrator_kwargs=Integrator_kwargs)

def fardal_stream(potential, w0, prog_mass, t_f, dt=1., t_0=0., release_every=1,
Integrator=LeapfrogIntegrator, Integrator_kwargs=dict()):
"""
Generate a mock stellar stream in the specified potential with a
progenitor system that ends up at the specified position.
This uses the prescription from Fardal et al. (2015).
Parameters
----------
potential : `~gary.potential.PotentialBase`
The gravitational potential.
w0 : `~gary.dynamics.PhaseSpacePosition`, array_like
Initial conditions.
prog_mass : numeric, array_like
A single mass or an array of masses if the progenitor mass evolves
with time.
t_f : numeric
The final time for integrating.
t_0 : numeric (optional)
The initial time for integrating -- the time at which ``w0`` is the position.
dt : numeric (optional)
The time-step.
release_every : int (optional)
Release particles at the Lagrange points every X timesteps.
Integrator : `~gary.integrate.Integrator` (optional)
Integrator to use.
Integrator_kwargs : dict (optional)
Any extra keyword argumets to pass to the integrator function.
Returns
-------
prog_orbit : `~gary.dynamics.CartesianOrbit`
stream : `~gary.dynamics.CartesianPhaseSpacePosition`
"""
k_mean = np.zeros(6)
k_disp = np.zeros(6)

k_mean[0] = 2. # R
k_disp[0] = 0.5

k_mean[1] = 0. # phi
k_disp[1] = 0.

k_mean[2] = 0. # z
k_disp[2] = 0.5

k_mean[3] = 0. # vR
k_disp[3] = 0.

k_mean[4] = 0.3 # vt
k_disp[4] = 0.5

k_mean[5] = 0. # vz
k_disp[5] = 0.5

return mock_stream(potential=potential, w0=w0, prog_mass=prog_mass,
k_mean=k_mean, k_disp=k_disp,
t_f=t_f, dt=dt, t_0=t_0, release_every=release_every,
Integrator=Integrator, Integrator_kwargs=Integrator_kwargs)

def dissolved_fardal_stream(potential, w0, prog_mass, t_disrupt, t_f, dt=1., t_0=0.,
release_every=1, Integrator=LeapfrogIntegrator, Integrator_kwargs=dict()):
"""
Generate a mock stellar stream in the specified potential with a
progenitor system that ends up at the specified position.
This uses the prescription from Fardal et al. (2015), but at a specified
time the progenitor completely dissolves and the radial offset of the
tidal radius is reduced to 0 at fixed dispersion.
Parameters
----------
potential : `~gary.potential.PotentialBase`
The gravitational potential.
w0 : `~gary.dynamics.PhaseSpacePosition`, array_like
Initial conditions.
prog_mass : numeric, array_like
A single mass or an array of masses if the progenitor mass evolves
with time.
t_disrupt : numeric
The time that the progenitor completely disrupts.
t_f : numeric
The final time for integrating.
t_0 : numeric (optional)
The initial time for integrating -- the time at which ``w0`` is the position.
dt : numeric (optional)
The time-step.
release_every : int (optional)
Release particles at the Lagrange points every X timesteps.
Integrator : `~gary.integrate.Integrator` (optional)
Integrator to use.
Integrator_kwargs : dict (optional)
Any extra keyword argumets to pass to the integrator function.
Returns
-------
prog_orbit : `~gary.dynamics.CartesianOrbit`
stream : `~gary.dynamics.CartesianPhaseSpacePosition`
"""

# the time index closest to when the disruption happens
nsteps = int(round(np.abs((t_f-t_0)/dt)))
t = np.linspace(t_0, t_f, nsteps)
disrupt_ix = np.abs(t - t_disrupt).argmin()

k_mean = np.zeros((nsteps,6))
k_disp = np.zeros((nsteps,6))

k_mean[:,0] = 2. # R
k_mean[disrupt_ix:,0] = 0.
k_disp[:,0] = 0.5

k_mean[:,1] = 0. # phi
k_disp[:,1] = 0.

k_mean[:,2] = 0. # z
k_disp[:,2] = 0.5

k_mean[:,3] = 0. # vR
k_disp[:,3] = 0.

k_mean[:,4] = 0.3 # vt
k_disp[:,4] = 0.5

k_mean[:,5] = 0. # vz
k_disp[:,5] = 0.5

return mock_stream(potential=potential, w0=w0, prog_mass=prog_mass,
k_mean=k_mean, k_disp=k_disp,
t_f=t_f, dt=dt, t_0=t_0, release_every=release_every,
Integrator=Integrator, Integrator_kwargs=Integrator_kwargs)
117 gary/dynamics/mockstream/tests/test_mockstream.py
View file
@@ -8,26 +8,27 @@
import astropy.units as u
import matplotlib.pyplot as pl
import numpy as np
import pytest

# Custom
import gary.dynamics as gd
import gary.integrate as gi
import gary.potential as gp
from gary.units import galactic
from ....potential import SphericalNFWPotential
from ....dynamics import CartesianPhaseSpacePosition
from ....integrate import DOPRI853Integrator
from ....units import galactic

# Project
from ..core import mock_stream
from ..core import mock_stream, streakline_stream, fardal_stream, dissolved_fardal_stream

def test_dumb():
potential = gp.SphericalNFWPotential(v_c=0.2, r_s=20., units=galactic)
def test_mock_stream():
potential = SphericalNFWPotential(v_c=0.2, r_s=20., units=galactic)

w0 = gd.CartesianPhaseSpacePosition(pos=[0.,15.,0]*u.kpc,
vel=[-0.13,0,0]*u.kpc/u.Myr)
w0 = CartesianPhaseSpacePosition(pos=[0.,15.,0]*u.kpc,
vel=[-0.13,0,0]*u.kpc/u.Myr)
k_mean = [1.,0.,0.,0.,1.,0.]
k_disp = [0.,0.,0.,0.,0.,0.]
prog,stream = mock_stream(potential, w0, k_mean=k_mean, k_disp=k_disp,
prog_mass=1E4, t_f=-2048., dt=-2.,
Integrator=gi.DOPRI853Integrator)
Integrator=DOPRI853Integrator)

# fig = prog.plot(subplots_kwargs=dict(sharex=False,sharey=False))
# fig = stream.plot(color='#ff0000', alpha=0.5, axes=fig.axes)
@@ -46,87 +47,23 @@ def test_dumb():
w0 = [0.,15.,0,-0.13,0,0]
prog,stream = mock_stream(potential, w0, k_mean=k_mean, k_disp=k_disp,
prog_mass=1E4, t_f=-2048., dt=-2.,
Integrator=gi.DOPRI853Integrator)
Integrator=DOPRI853Integrator)

# def test_streakline():
# # potential = gp.LogarithmicPotential(v_c=0.2, r_h=20., q1=1., q2=1., q3=1., units=galactic)
# # potential = gp.MiyamotoNagaiPotential(m=1E11, a=6.5, b=0.26, units=galactic)
# potential = gp.SphericalNFWPotential(v_c=0.2, r_s=20., units=galactic)
# release_every = 1
# time = 2048
# dt = 2
# nsteps = time // dt
mock_funcs = [streakline_stream, fardal_stream, dissolved_fardal_stream]
all_extra_args = [dict(), dict(), dict(t_disrupt=-250.)]
@pytest.mark.parametrize("mock_func, extra_kwargs", zip(mock_funcs, all_extra_args))
def test_each_type(mock_func, extra_kwargs):
potential = SphericalNFWPotential(v_c=0.2, r_s=20., units=galactic)

# # w0 = np.array([0.,15.,0,-0.2,0,0]) # circular
# w0 = np.array([0.,15.,0,-0.13,0,0]) # eccentric
w0 = CartesianPhaseSpacePosition(pos=[0.,15.,0]*u.kpc,
vel=[-0.13,0,0]*u.kpc/u.Myr)
prog,stream = mock_func(potential, w0, prog_mass=1E4, t_f=-2048., dt=-2.,
Integrator=DOPRI853Integrator, **extra_kwargs)

# print("integrating parent orbit")
# # TODO: this should return an orbit! But i need to test...
# raise NotImplementedError("FIX THE DAMN TODO ADRIAN")

# torig,w = potential.integrate_orbit(w0, dt=-dt, nsteps=nsteps, Integrator=gi.DOPRI853Integrator)
# t,w = potential.integrate_orbit(w[-1], dt=dt, t1=torig[-1], nsteps=nsteps, Integrator=gi.DOPRI853Integrator)
# ww = w[:,0].copy()
# print("done")
# r = np.sqrt(np.sum(ww[:,:3]**2, axis=-1))
# ecc = (r.max() - r.min()) / (r.min() + r.max())
# print("eccentricity", ecc)

# fig = gd.plot_orbits(w, marker=None)
# # pl.show()
# # return

# prog_mass = np.zeros_like(t) + 1E4

# stream = streakline_stream(potential.c_instance, t, ww,
# release_every=release_every, G=potential.G,
# prog_mass=prog_mass)

# ixs = [0,1]
# pl.figure(figsize=(6,6))
# pl.plot(ww[:,ixs[0]], ww[:,ixs[1]], marker=None, alpha=0.5)
# pl.scatter(stream[::2,ixs[0]], stream[::2,ixs[1]], c=t[:-1:release_every], vmin=t.min(), vmax=t.max(), cmap='coolwarm', s=4)
# pl.scatter(stream[1::2,ixs[0]], stream[1::2,ixs[1]], c=t[:-1:release_every], vmin=t.min(), vmax=t.max(), cmap='coolwarm', s=4)
# # pl.xlim(-0.6, 0.6)
# # pl.ylim(14.9,15.1)
# pl.show()

# def test_fardal():
# # potential = gp.LogarithmicPotential(v_c=0.2, r_h=20., q1=1., q2=1., q3=1., units=galactic)
# # potential = gp.MiyamotoNagaiPotential(m=1E11, a=6.5, b=0.26, units=galactic)
# potential = gp.SphericalNFWPotential(v_c=0.2, r_s=20., units=galactic)
# release_every = 1
# time = 2048
# dt = 2
# nsteps = time // dt

# # w0 = np.array([0.,15.,0,-0.2,0,0]) # circular
# w0 = np.array([0.,15.,0,-0.13,0,0]) # eccentric

# print("integrating parent orbit")
# torig,w = potential.integrate_orbit(w0, dt=-dt, nsteps=nsteps, Integrator=gi.DOPRI853Integrator)
# t,w = potential.integrate_orbit(w[-1], dt=dt, t1=torig[-1], nsteps=nsteps, Integrator=gi.DOPRI853Integrator)
# ww = w[:,0].copy()
# print("done")
# r = np.sqrt(np.sum(ww[:,:3]**2, axis=-1))
# ecc = (r.max() - r.min()) / (r.min() + r.max())
# print("eccentricity", ecc)

# fig = gd.plot_orbits(w, marker=None)
# # pl.show()
# # return

# prog_mass = np.zeros_like(t) + 1E4

# stream = fardal_stream(potential.c_instance, t, ww,
# release_every=release_every, G=potential.G,
# prog_mass=prog_mass)
# fig = prog.plot(subplots_kwargs=dict(sharex=False,sharey=False))
# fig = stream.plot(color='#ff0000', alpha=0.5, axes=fig.axes)
# fig = stream.plot()
# pl.show()

# ixs = [0,1]
# pl.figure(figsize=(6,6))
# pl.plot(ww[:,ixs[0]], ww[:,ixs[1]], marker=None, alpha=0.5)
# pl.scatter(stream[::2,ixs[0]], stream[::2,ixs[1]], c=t[:-1:release_every], vmin=t.min(), vmax=t.max(), cmap='coolwarm', s=4)
# pl.scatter(stream[1::2,ixs[0]], stream[1::2,ixs[1]], c=t[:-1:release_every], vmin=t.min(), vmax=t.max(), cmap='coolwarm', s=4)
# # pl.xlim(-0.6, 0.6)
# # pl.ylim(14.9,15.1)
# pl.show()
assert prog.t.shape == (1024,)
assert stream.pos.shape == (3,2048) # two particles per step