def test_careful_traceback_and_forward():
"""Step by step, project orbit forward, then backward"""
bovy_times = np.array([0., np.pi / 3.])
chron_times = torb.convert_bovytime2myr(bovy_times)
init_pos_chron = np.array([
4000, 8000. * np.sqrt(3) / 2, 0,
np.sin(np.pi / 3) * 220., -np.cos(np.pi / 3) * 220., 0
])
init_pos_galpy = torb.convert_cart2galpycoords(init_pos_chron, ts=0.)
assert np.allclose(np.array([1., 0, 1, 0, 0, np.pi / 3.]),
init_pos_galpy)
o = Orbit(vxvv=init_pos_galpy, ro=8., vo=220.)
o.integrate(bovy_times, MWPotential2014, method='odeint')
orbit_galpy = o.getOrbit()
assert np.allclose(init_pos_galpy, orbit_galpy[0])
assert np.allclose(
init_pos_galpy + np.array([0., 0., 0., 0., 0., bovy_times[-1]]),
orbit_galpy[-1])
orbit_chron = torb.convert_galpycoords2cart(orbit_galpy, ts=bovy_times)
assert np.allclose(init_pos_chron, orbit_chron[0])
assert np.allclose(init_pos_chron, orbit_chron[-1])
# Setup for backwards time integration
# Currently at time of PI/3
back_init_pos_chron = orbit_chron[-1]
back_init_pos_galpy = torb.convert_cart2galpycoords(
back_init_pos_chron,
bovy_times=bovy_times[-1],
)
assert np.allclose(
back_init_pos_galpy,
torb.convert_cart2galpycoords(back_init_pos_chron,
bovy_times=bovy_times[-1]))
back_o = Orbit(vxvv=back_init_pos_galpy, ro=8., vo=220.)
back_o.integrate(-1 * bovy_times, MWPotential2014, method='odeint')
back_orbit_galpy = back_o.getOrbit()
assert np.allclose(back_init_pos_galpy, back_orbit_galpy[0])
assert np.allclose(
back_init_pos_galpy -
np.array([0., 0., 0., 0., 0., bovy_times[-1]]),
back_orbit_galpy[-1])
assert np.allclose(init_pos_galpy, back_orbit_galpy[-1])
back_orbit_chron = torb.convert_galpycoords2cart(
back_orbit_galpy,
ts=bovy_times[::-1],
)
assert np.allclose(init_pos_chron, back_orbit_chron[-1])
def test_galpy_moving_conversions():
"""Check if gaply conversions behave as expected where time
is allowed to vary."""
lsr_chron = np.zeros(6)
lsr_galpy = np.array([1.,0,1,0,0,0])
# Incorporate positive time into lsr position checks
NSTEPS = 10
galpy_times = np.linspace(0., 2*np.pi, NSTEPS)
lsrs_chron = np.repeat(lsr_chron, NSTEPS).reshape(6,-1).T
lsrs_galpy = np.repeat(lsr_galpy, NSTEPS).reshape(6,-1).T
lsrs_galpy[:,-1] = galpy_times
chron_times = torb.convert_bovytime2myr(galpy_times)
assert np.allclose(
lsrs_chron,
torb.convert_galpycoords2cart(lsrs_galpy, ts=galpy_times))
assert np.allclose(
lsrs_galpy,
torb.convert_cart2galpycoords(lsrs_chron, ts=chron_times)
)
# Incorporate negative time into lsr position checks
galpy_times = np.linspace(0., -2*np.pi, NSTEPS)
lsrs_chron = np.repeat(lsr_chron, NSTEPS).reshape(6,-1).T
lsrs_galpy = np.repeat(lsr_galpy, NSTEPS).reshape(6,-1).T
lsrs_galpy[:,-1] = galpy_times
chron_times = torb.convert_bovytime2myr(galpy_times)
assert np.allclose(
lsrs_chron,
torb.convert_galpycoords2cart(lsrs_galpy, ts=galpy_times))
assert np.allclose(
lsrs_galpy,
torb.convert_cart2galpycoords(lsrs_chron, ts=chron_times)
)
# Test random positions with random times
SPREAD = int(1e4) # pc
NSAMPLES = 100
many_pos_chron = (np.random.rand(NSAMPLES,6) - 0.5) * SPREAD # uniform between -10 and 10
many_chron_times = np.random.rand(NSAMPLES) * 100 #Myr
many_pos_galpy = torb.convert_cart2galpycoords(
many_pos_chron, ts=many_chron_times
)
many_galpy_times = torb.convert_myr2bovytime(many_chron_times)
for i in range(NSAMPLES):
assert np.allclose(many_pos_chron[i],
torb.convert_galpycoords2cart(
many_pos_galpy[i], ts=many_galpy_times[i]
),
atol=1e-2)
def test_galpy_moving_conversions():
"""Check if gaply conversions behave as expected where time
is allowed to vary."""
lsr_chron = np.zeros(6)
lsr_galpy = np.array([1., 0, 1, 0, 0, 0])
# Incorporate positive time into lsr position checks
NSTEPS = 10
galpy_times = np.linspace(0., 2 * np.pi, NSTEPS)
lsrs_chron = np.repeat(lsr_chron, NSTEPS).reshape(6, -1).T
lsrs_galpy = np.repeat(lsr_galpy, NSTEPS).reshape(6, -1).T
lsrs_galpy[:, -1] = galpy_times
chron_times = torb.convert_bovytime2myr(galpy_times)
assert np.allclose(
lsrs_chron,
torb.convert_galpycoords2cart(lsrs_galpy, ts=galpy_times))
assert np.allclose(
lsrs_galpy,
torb.convert_cart2galpycoords(lsrs_chron, ts=chron_times)
)
# Incorporate negative time into lsr position checks
galpy_times = np.linspace(0., -2 * np.pi, NSTEPS)
lsrs_chron = np.repeat(lsr_chron, NSTEPS).reshape(6, -1).T
lsrs_galpy = np.repeat(lsr_galpy, NSTEPS).reshape(6, -1).T
lsrs_galpy[:, -1] = galpy_times
chron_times = torb.convert_bovytime2myr(galpy_times)
assert np.allclose(
lsrs_chron,
torb.convert_galpycoords2cart(lsrs_galpy, ts=galpy_times))
assert np.allclose(
lsrs_galpy,
torb.convert_cart2galpycoords(lsrs_chron, ts=chron_times)
)
# Test random positions with random times
SPREAD = int(1e4) # pc
NSAMPLES = 100
many_pos_chron = (np.random.rand(NSAMPLES, 6) - 0.5) * SPREAD # uniform between -10 and 10
many_chron_times = np.random.rand(NSAMPLES) * 100 # Myr
many_pos_galpy = torb.convert_cart2galpycoords(
many_pos_chron, ts=many_chron_times
)
many_galpy_times = torb.convert_myr2bovytime(many_chron_times)
for i in range(NSAMPLES):
assert np.allclose(many_pos_chron[i],
torb.convert_galpycoords2cart(
many_pos_galpy[i], ts=many_galpy_times[i]
),
atol=1e-2)
def test_galpy_stationary_conversions():
"""Check if gaply conversions behave as expected where everything
is at time 0"""
# Test LSR
lsr_chron = np.zeros(6)
lsr_galpy = np.array([1., 0, 1, 0, 0, 0])
assert np.allclose(lsr_chron,
torb.convert_galpycoords2cart(lsr_galpy, ts=0.))
assert np.allclose(lsr_galpy,
torb.convert_cart2galpycoords(lsr_chron, ts=0.))
# Test galactic centre
gc_chron = np.array([
8000.,
0,
0,
0,
-220.,
0,
])
gc_galpy = np.ones(6) * 1e-15
assert np.allclose(gc_chron,
torb.convert_galpycoords2cart(gc_galpy, ts=0.))
assert np.allclose(gc_galpy,
torb.convert_cart2galpycoords(gc_chron, ts=0.))
# Test simple, off origin point
off_chron = np.array([
4000, 8000. * np.sqrt(3) / 2, 0,
np.sin(np.pi / 3) * 220., -np.cos(np.pi / 3) * 220., 0
])
off_galpy = np.array([1., 0, 1, 0, 0, np.pi / 3.])
assert np.allclose(off_galpy,
torb.convert_cart2galpycoords(off_chron, ts=0.))
assert np.allclose(off_chron,
torb.convert_galpycoords2cart(off_galpy, ts=0.))
# Test random positions
SPREAD = 100000
NSAMPLES = int(1e6)
many_pos_chron = (np.random.rand(NSAMPLES, 6) -
0.5) * SPREAD # uniform between -10 and 10
many_pos_galpy = torb.convert_cart2galpycoords(many_pos_chron, ts=0.)
assert np.allclose(many_pos_chron,
torb.convert_galpycoords2cart(many_pos_galpy,
ts=0.),
atol=1e-2)
def get_guiding_radius(xyzuvw=None, galpy_coords=None):
"""
Returns radius of guiding centre, in pc
Accepts as input either Chronostar xyzuvw, with
units pc and km/s
or galpy [R, vR, vT, Z, vZ, phi]
Exploits conservation of angular momentum, and `psi` to construct and then
solve a quadratic equation for R_g. See latex document for derivation.
"""
if galpy_coords is None:
galpy_coords = torb.convert_cart2galpycoords(xyzuvw)
R, _, V_T, _, _, _ = galpy_coords
# Scale up from galpy dimensionless values
R *= R0
V_T *= V0
psi = A + B
# Calculate the coefficients from ax^2 + bx + c, where x is R_L
# See section 1.1 in latex file for derivation
a_coeff = psi # myr-1
b_coeff = -(V0 + psi * R0) # pc/myr
c_coeff = R * V_T # pc^2/myr
results = quadratic_formula(a_coeff, b_coeff, c_coeff)
# Return the closest value to initial radius, because the other value
# is irrelevant
diff = np.abs(results - R)
return results[np.argmin(diff)]
def test_misc():
POS_SPAN = 10
VEL_SPAN = 10
for i in range(100):
# Generate random stars with chron coords between [SPAN, SPAN]
rand_xyz = 2 * POS_SPAN * np.random.rand(3) - POS_SPAN
rand_uvw = 2 * VEL_SPAN * np.random.rand(3) - VEL_SPAN
rand_xyzuvw = np.hstack((rand_xyz, rand_uvw))
rand_galpy = torb.convert_cart2galpycoords(rand_xyzuvw)
rand_epi = eg.convert_galpy2epi(rand_galpy)
galpy_res = eg.convert_epi2galpy(rand_epi)
xyzuvw_res = torb.convert_galpycoords2cart(galpy_res)
try:
# Positions shouldn't vary by more than 3 pc
pos_atol = 2
# Velocities shouldn't vary by more than 0.05 km/s
vel_atol = 0.03
assert np.allclose(rand_xyzuvw[:3], xyzuvw_res[:3], atol=pos_atol)
assert np.allclose(rand_xyzuvw[3:], xyzuvw_res[3:], atol=vel_atol)
# assert np.allclose(rand_galpy, eg.convert_epi2galpy(rand_epi), rtol=rtol)
except AssertionError:
print(i)
import pdb
pdb.set_trace()
def test_galpy_stationary_conversions():
"""Check if gaply conversions behave as expected where everything
is at time 0"""
# Test LSR
lsr_chron = np.zeros(6)
lsr_galpy = np.array([1.,0,1,0,0,0])
assert np.allclose(lsr_chron,
torb.convert_galpycoords2cart(lsr_galpy, ts=0.))
assert np.allclose(lsr_galpy,
torb.convert_cart2galpycoords(lsr_chron, ts=0.))
# Test galactic centre
gc_chron = np.array([8000.,0,0,0,-220.,0,])
gc_galpy = np.ones(6) * 1e-15
assert np.allclose(gc_chron,
torb.convert_galpycoords2cart(gc_galpy, ts=0.))
assert np.allclose(gc_galpy,
torb.convert_cart2galpycoords(gc_chron, ts=0.))
# Test simple, off origin point
off_chron = np.array([4000, 8000.*np.sqrt(3)/2, 0,
np.sin(np.pi/3)*220.,
-np.cos(np.pi/3)*220.,
0])
off_galpy = np.array([1.,0,1,0,0,np.pi/3.])
assert np.allclose(off_galpy,
torb.convert_cart2galpycoords(off_chron, ts=0.))
assert np.allclose(off_chron,
torb.convert_galpycoords2cart(off_galpy, ts=0.))
# Test random positions
SPREAD = 100000
NSAMPLES = int(1e6)
many_pos_chron = (np.random.rand(NSAMPLES,6) - 0.5) * SPREAD # uniform between -10 and 10
many_pos_galpy = torb.convert_cart2galpycoords(many_pos_chron, ts=0.)
assert np.allclose(many_pos_chron,
torb.convert_galpycoords2cart(many_pos_galpy, ts=0.),
atol=1e-2)
def test_galpy2chron2galpy_stationary():
"""
Check that converting from Galpy to Chronostar is internally
consistent (you can convert back and forth)
Time is fixed at 0
"""
for i in range(100):
xyzuvw_start = np.random.rand(6)
galpy_start = torb.convert_cart2galpycoords(xyzuvw_start)
xyzuvw_res = torb.convert_galpycoords2cart(galpy_start)
assert np.allclose(xyzuvw_start, xyzuvw_res)
def test_evolve_chronspace():
"""
Compare orbits evolved with both galpy and epicyclic, comparing in chron space
"""
time = 50. #Myr
btime = torb.convert_myr2bovytime(time)
chron_start = np.array([0., 0., 10., -2., 0., 0.])
galpy_start = torb.convert_cart2galpycoords(chron_start)
epi_start = eg.convert_galpy2epi(galpy_start)
# Just make sure the starting point can be transformed back and forth
assert np.allclose(galpy_start, eg.convert_epi2galpy(epi_start))
assert np.allclose(
chron_start,
torb.convert_galpycoords2cart(eg.convert_epi2galpy(epi_start)))
epi_end = eg.evolve_epi(epi_start, time)[0]
galpy_end = torb.trace_galpy_orbit(galpy_start,
times=time,
single_age=True)
chron_end = torb.trace_cartesian_orbit(chron_start,
times=time,
single_age=True)
# Chronostar orbit end point, in galpy units
chron_end_gu = torb.convert_cart2galpycoords(chron_end, ts=time)
# This should be exact, because algorithmically it's the same thing.
assert np.allclose(galpy_end, chron_end_gu)
# Epicyclic orbit end point, in chronostar units
epi_end_chron = torb.convert_galpycoords2cart(
eg.convert_epi2galpy(epi_end), ts=btime)
# assert position accurate within 15 pc
# NOTE this is quite large... offset mainly in X. Maybe something to
# investigate.
assert np.allclose(chron_end[:3], epi_end_chron[:3], atol=15.)
# assert velocity accurate within 0.5 km/s
assert np.allclose(chron_end[3:], epi_end_chron[3:], atol=.5)
def test_galpy2chron2galpy_moving():
"""
Check that converting from Galpy to Chronostar is internally
consistent (you can convert back and forth)
Time is allowed to vary
"""
# Test first LSR
xyzuvw_start = np.zeros(6)
time = 1.
galpy_start = torb.convert_cart2galpycoords(xyzuvw_start, bovy_times=time)
# import pdb; pdb.set_trace()
xyzuvw_res = torb.convert_galpycoords2cart(galpy_start, ts=time)
# import pdb; pdb.set_trace()
assert np.allclose(xyzuvw_start, xyzuvw_res)
# Now test slightly rotated LSR
time = np.pi/4
galpy_stat_start = np.array([1., 0., 1., 0., 0., time])
xyzuvw_start = torb.convert_galpycoords2cart(galpy_stat_start)
galpy_start = torb.convert_cart2galpycoords(xyzuvw_start, bovy_times=time)
xyzuvw_res = torb.convert_galpycoords2cart(galpy_start, ts=time)
assert np.allclose(xyzuvw_start, xyzuvw_res)
# Now test points randomly dispersed near the LSR
# but at a time range corresponding to LSR rotation of [0-1] rad
for i in range(100):
xyzuvw_start = np.random.rand(6)
# Time is in galpy units
time = np.random.rand()
galpy_start = torb.convert_cart2galpycoords(xyzuvw_start, bovy_times=time)
# import pdb; pdb.set_trace()
xyzuvw_res = torb.convert_galpycoords2cart(galpy_start, ts=time)
# import pdb; pdb.set_trace()
assert np.allclose(xyzuvw_start, xyzuvw_res)
- https://galpy.readthedocs.io/en/v1.4.0/actionAngle.html#actionanglestaeckel
- https://galpy.readthedocs.io/en/v1.4.0/orbit.html#fastchar
- https://galpy.readthedocs.io/en/v1.4.0/actionAngle.html#accessing-action-angle-coordinates-for-orbit-instances
"""
import numpy as np
from galpy.orbit import Orbit
from galpy.potential import MWPotential2014
from galpy.actionAngle import actionAngleTorus
from galpy.potential import MWPotential2014
import sys
sys.path.insert(0, '..')
import chronostar.traceorbit as torb
xyzuvw_start = [0., 0., 25., 0., 0., 0.]
print('xyzuvw start: {}'.format(xyzuvw_start))
galpy_coords = torb.convert_cart2galpycoords(xyzuvw_start)
print('galpy start: {}'.format(galpy_coords))
aAT = actionAngleTorus(pot=MWPotential2014)
o = Orbit(vxvv=galpy_coords, ro=8., vo=220.)
o.e(analytic=True, type='staeckel', pot=MWPotential2014)
aAS = actionAngleStaeckel(pot=mp, delta=0.4)
# Om= aAT.Freqs(jr,lz,jz)
def test_careful_traceback_and_forward():
"""Step by step, project orbit forward, then backward"""
bovy_times = np.array([0., np.pi/3.])
chron_times = torb.convert_bovytime2myr(bovy_times)
init_pos_chron = np.array([
4000, 8000.*np.sqrt(3)/2, 0,
np.sin(np.pi/3)*220.,
-np.cos(np.pi/3)*220.,
0
])
init_pos_galpy = torb.convert_cart2galpycoords(init_pos_chron, ts=0.)
assert np.allclose(np.array([1.,0,1,0,0,np.pi/3.]),
init_pos_galpy)
o = Orbit(vxvv=init_pos_galpy, ro=8., vo=220.)
o.integrate(bovy_times, MWPotential2014, method='odeint')
orbit_galpy = o.getOrbit()
assert np.allclose(init_pos_galpy, orbit_galpy[0])
assert np.allclose(init_pos_galpy
+ np.array([0.,0.,0.,0.,0.,bovy_times[-1]]),
orbit_galpy[-1])
orbit_chron = torb.convert_galpycoords2cart(orbit_galpy,
ts=bovy_times)
assert np.allclose(init_pos_chron, orbit_chron[0])
assert np.allclose(init_pos_chron,
orbit_chron[-1])
# Setup for backwards time integration
# Currently at time of PI/3
back_init_pos_chron = orbit_chron[-1]
back_init_pos_galpy = torb.convert_cart2galpycoords(
back_init_pos_chron,
bovy_times=bovy_times[-1],
)
assert np.allclose(back_init_pos_galpy,
torb.convert_cart2galpycoords(
back_init_pos_chron,
bovy_times=bovy_times[-1]
))
back_o = Orbit(vxvv=back_init_pos_galpy, ro=8., vo=220.)
back_o.integrate(-1*bovy_times, MWPotential2014, method='odeint')
back_orbit_galpy = back_o.getOrbit()
assert np.allclose(back_init_pos_galpy, back_orbit_galpy[0])
assert np.allclose(back_init_pos_galpy
- np.array([0.,0.,0.,0.,0.,bovy_times[-1]]),
back_orbit_galpy[-1])
assert np.allclose(init_pos_galpy, back_orbit_galpy[-1])
back_orbit_chron = torb.convert_galpycoords2cart(
back_orbit_galpy,
ts=bovy_times[::-1],
)
assert np.allclose(init_pos_chron, back_orbit_chron[-1])
def test_w_offset():
w_offset = torb.convert_cart2galpycoords([0., 0., 0., 0., 0., 1.])
assert np.allclose(w_offset,
eg.convert_epi2galpy(eg.convert_galpy2epi(w_offset)))
def test_v_offset():
rtol = 1e-4
v_offset = torb.convert_cart2galpycoords([0., 0., 0., 0., 1., 0.])
assert np.allclose(v_offset,
eg.convert_epi2galpy(eg.convert_galpy2epi(v_offset)),
rtol=rtol)
def test_y_offset():
y_offset = torb.convert_cart2galpycoords([0., 10., 0., 0., 0., 0.])
y_res = eg.convert_galpy2epi(y_offset)
assert np.allclose(y_offset,
eg.convert_epi2galpy(eg.convert_galpy2epi(y_offset)))
def test_x_offset():
x_offset = torb.convert_cart2galpycoords([10., 0., 0., 0., 0., 0.])
assert np.allclose(x_offset,
eg.convert_epi2galpy(eg.convert_galpy2epi(x_offset)))