def mass_enclosed(self, x):
"""
Estimate the mass enclosed within the given position by assumine the potential
is spherical. This is basic, and assumes spherical symmetry.
Parameters
----------
x : array_like, numeric
Position to estimate the enclossed mass.
"""
# Fractional step-size in radius
h = 0.01
# Radius
r = np.sqrt(np.sum(x**2, axis=-1))
epsilon = h*x/r[...,np.newaxis]
dPhi_dr_plus = self.value(x + epsilon)
dPhi_dr_minus = self.value(x - epsilon)
diff = dPhi_dr_plus - dPhi_dr_minus
if self.units is None:
raise ValueError("No units specified when creating potential object.")
Gee = G.decompose(self.units).value
return np.abs(r*r * diff / Gee / (2.*h))
def _read_units(self):
"""
Read and parse the SCFPAR file containing simulation parameters
and initial conditions. Right now, only parse out the simulation units.
"""
pars = dict()
parfile = os.path.join(self.path, "SCFPAR")
with open(parfile) as f:
lines = f.readlines()
# find what G is set to
for i,line in enumerate(lines):
if line.split()[1].strip() == "G":
break
pars['G'] = float(lines[i].split()[0])
pars['length'] = float(lines[i+10].split()[0])
pars['mass'] = float(lines[i+11].split()[0])
self.x0 = np.array(map(float, lines[19].split()))*u.kpc
self.v0 = np.array(map(float, lines[20].split()))*u.km/u.s
_G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
X = (_G / pars['length']**3 * pars['mass'])**-0.5
length_unit = u.Unit("{0} kpc".format(pars['length']))
mass_unit = u.Unit("{0} M_sun".format(pars['mass']))
time_unit = u.Unit("{:08f} Myr".format(X))
# units = dict(length=length_unit,
# mass=mass_unit,
# time=time_unit,
# speed=length_unit/time_unit,
# dimensionless=u.dimensionless_unscaled)
return UnitSystem(length_unit, mass_unit, time_unit, length_unit/time_unit, u.radian)
def test_against_gary():
cos_coeff = np.array([[[1.509, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [-2.606, 0.0, 0.665, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [6.406, 0.0, -0.66, 0.0, 0.044, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [-5.859, 0.0, 0.984, 0.0, -0.03, 0.0, 0.001]], [[-0.086, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [-0.221, 0.0, 0.129, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.295, 0.0, -0.14, 0.0, -0.012, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], [[-0.033, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [-0.001, 0.0, 0.006, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]], [[-0.02, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]])
sin_coeff = np.zeros_like(cos_coeff)
nmax = sin_coeff.shape[0]-1
lmax = sin_coeff.shape[1]-1
max_x = 10.
grid = np.linspace(-max_x,max_x,8)
grid = grid[grid != 0.]
xyz = np.ascontiguousarray(np.vstack(map(np.ravel, np.meshgrid(grid,grid,grid))).T)
# xyz = np.array([[1.,1.,1.]])
grad = np.zeros_like(xyz)
acceleration(xyz, grad, sin_coeff, cos_coeff, nmax, lmax)
grad *= -1.
val = np.zeros(len(grad))
value(xyz, val, sin_coeff, cos_coeff, nmax, lmax)
# gary
import gary.potential as gp
from gary.units import galactic
G = _G.decompose(galactic).value
potential = gp.SCFPotential(m=1/G, r_s=1.,
sin_coeff=sin_coeff, cos_coeff=cos_coeff,
units=galactic)
gary_val = potential.value(xyz)
gary_grad = potential.gradient(xyz)
np.testing.assert_allclose(gary_val, val, rtol=1E-5)
np.testing.assert_allclose(gary_grad, grad, rtol=1E-5)
def __init__(self, m, a, units):
self.parameters = OrderedDict()
self.parameters['m'] = m
self.parameters['a'] = a
super(KuzminPotential, self).__init__(units=units)
self.G = G.decompose(units).value
def __init__(self, parameters, origin=None, R=None,
ndim=3, units=None):
self._units = self._validate_units(units)
self.parameters = self._prepare_parameters(parameters, self.units)
try:
self.G = G.decompose(self.units).value
except u.UnitConversionError:
self.G = 1. # TODO: this is a HACK and could lead to user confusion
self.ndim = ndim
if origin is None:
origin = np.zeros(self.ndim)
self.origin = self._remove_units(origin)
if R is not None and ndim not in [2, 3]:
raise NotImplementedError('Gala potentials currently only support '
'rotations when ndim=2 or ndim=3.')
if R is not None:
if isinstance(R, Rotation):
R = R.as_dcm()
R = np.array(R)
if R.shape != (ndim, ndim):
raise ValueError('Rotation matrix passed to potential {0} has '
'an invalid shape: expected {1}, got {2}'
.format(self.__class__.__name__,
(ndim, ndim), R.shape))
self.R = R
def test_sech2():
Sigma = 65. * u.Msun / u.pc**2
hz = 250 * u.pc
_G = G.decompose([u.pc, u.Msun, u.Myr]).value
rho0 = Sigma / (4 * hz)
with u.set_enabled_equivalencies(u.dimensionless_angles()):
assert quantity_allclose(sech2_density(0. * u.pc, Sigma, hz), rho0)
# check numerical derivatives of potential against functions
rnd = np.random.RandomState(42)
for z0 in rnd.uniform(100, 500, size=16): # check a few random places
num_grad = derivative(sech2_potential,
z0,
dx=1e-2,
n=1,
args=(Sigma.value, hz.value, _G))
grad = sech2_gradient(z0, Sigma.value, hz.value, _G)
assert np.allclose(grad, num_grad)
num_dens = derivative(
sech2_potential,
z0,
dx=1e-2,
n=2,
args=(Sigma.value, hz.value, _G)) / (4 * np.pi * _G)
dens = sech2_density(z0, Sigma.value, hz.value)
assert np.allclose(dens, num_dens)
def _setup_potential(self, parameters, origin=None, R=None, units=None):
self._units = self._validate_units(units)
self.parameters = self._prepare_parameters(parameters, self.units)
try:
self.G = G.decompose(self.units).value
except u.UnitConversionError:
# TODO: this is a convention that and could lead to confusion!
self.G = 1.
if origin is None:
origin = np.zeros(self.ndim)
self.origin = self._remove_units(origin)
if R is not None and self.ndim not in [2, 3]:
raise NotImplementedError('Gala potentials currently only support '
'rotations when ndim=2 or ndim=3.')
if R is not None:
if isinstance(R, Rotation):
R = R.as_matrix()
R = np.array(R)
if R.shape != (self.ndim, self.ndim):
raise ValueError(
'Rotation matrix passed to potential {0} has '
'an invalid shape: expected {1}, got {2}'.format(
self.__class__.__name__, (self.ndim, self.ndim),
R.shape))
self.R = R
def _units_from_file(scfpar):
""" Generate a unit system from an SCFPAR file. """
with open(scfpar) as f:
lines = f.readlines()
length = float(lines[16].split()[0])
mass = float(lines[17].split()[0])
GG = G.decompose(bases=[u.kpc, u.M_sun, u.Myr]).value
X = (GG / length**3 * mass)**-0.5
length_unit = u.Unit("{0} kpc".format(length))
mass_unit = u.Unit("{0} M_sun".format(mass))
time_unit = u.Unit("{:08f} Myr".format(X))
return dict(length=length_unit, mass=mass_unit, time=time_unit)
def mass_enclosed(self, q, t=0.):
"""
Estimate the mass enclosed within the given position by assuming the potential
is spherical.
Parameters
----------
q : `~gala.dynamics.PhaseSpacePosition`, `~astropy.units.Quantity`, array_like
Position(s) to estimate the enclossed mass.
Returns
-------
menc : `~astropy.units.Quantity`
Mass enclosed at the given position(s). If the input position
has shape ``q.shape``, the output energy will have shape
``q.shape[1:]``.
"""
q = self._remove_units_prepare_shape(q)
orig_shape,q = self._get_c_valid_arr(q)
t = self._validate_prepare_time(t, q)
# small step-size in direction of q
h = 1E-3 # MAGIC NUMBER
# Radius
r = np.sqrt(np.sum(q**2, axis=1))
epsilon = h*q/r[:,np.newaxis]
dPhi_dr_plus = self._energy(q + epsilon, t=t)
dPhi_dr_minus = self._energy(q - epsilon, t=t)
diff = (dPhi_dr_plus - dPhi_dr_minus)
if isinstance(self.units, DimensionlessUnitSystem):
Gee = 1.
else:
Gee = G.decompose(self.units).value
Menc = np.abs(r*r * diff / Gee / (2.*h))
Menc = Menc.reshape(orig_shape[1:])
sgn = 1.
if 'm' in self.parameters and self.parameters['m'] < 0:
sgn = -1.
return sgn * Menc * self.units['mass']
def mass_enclosed(self, q, t=0.):
"""
Estimate the mass enclosed within the given position by assuming the potential
is spherical.
Parameters
----------
q : `~gala.dynamics.PhaseSpacePosition`, `~astropy.units.Quantity`, array_like
Position(s) to estimate the enclossed mass.
Returns
-------
menc : `~astropy.units.Quantity`
Mass enclosed at the given position(s). If the input position
has shape ``q.shape``, the output energy will have shape
``q.shape[1:]``.
"""
q = self._remove_units_prepare_shape(q)
orig_shape, q = self._get_c_valid_arr(q)
t = self._validate_prepare_time(t, q)
# small step-size in direction of q
h = 1E-3 # MAGIC NUMBER
# Radius
r = np.sqrt(np.sum(q**2, axis=1))
epsilon = h*q/r[:, np.newaxis]
dPhi_dr_plus = self._energy(q + epsilon, t=t)
dPhi_dr_minus = self._energy(q - epsilon, t=t)
diff = (dPhi_dr_plus - dPhi_dr_minus)
if isinstance(self.units, DimensionlessUnitSystem):
Gee = 1.
else:
Gee = G.decompose(self.units).value
Menc = np.abs(r*r * diff / Gee / (2.*h))
Menc = Menc.reshape(orig_shape[1:])
sgn = 1.
if 'm' in self.parameters and self.parameters['m'] < 0:
sgn = -1.
return sgn * Menc * self.units['mass']
def _units_from_file(scfpar):
""" Generate a unit system from an SCFPAR file. """
with open(scfpar) as f:
lines = f.readlines()
length = float(lines[16].split()[0])
mass = float(lines[17].split()[0])
GG = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
X = (GG / length**3 * mass)**-0.5
length_unit = u.Unit("{0} kpc".format(length))
mass_unit = u.Unit("{0} M_sun".format(mass))
time_unit = u.Unit("{:08f} Myr".format(X))
return dict(length=length_unit,
mass=mass_unit,
time=time_unit)
def main():
# The unit system we'll use:
units = [u.Myr, u.kpc, u.Msun]
_G = G.decompose(units).value
# Initial conditions
x0 = [10., 0, 0]
v0 = [0, 0.15, 0]
# Parameters of the Hernquist potential model
m = 1E11 # Msun
c = 1. # kpc
# Timestep in Myr
dt = 1.
n_steps = 10000 # 10 Gyr
t, x, v = leapfrog_integrate(x0,
v0,
dt,
n_steps,
hernquist_args=(_G, m, c))
# Plot the orbit
fig, axes = plt.subplots(1, 2, figsize=(10, 5))
axes[0].plot(x[:, 0], x[:, 1], marker='.', linestyle='none', alpha=0.1)
axes[0].set_xlim(-12, 12)
axes[0].set_ylim(-12, 12)
axes[0].set_xlabel('$x$')
axes[0].set_ylabel('$y$')
# ---
axes[1].plot(v[:, 0], v[:, 1], marker='.', linestyle='none', alpha=0.1)
axes[1].set_xlim(-0.35, 0.35)
axes[1].set_ylim(-0.35, 0.35)
axes[1].set_xlabel('$v_x$')
axes[1].set_ylabel('$v_y$')
fig.tight_layout()
plt.show()
def __init__(self, parameters, origin=None, parameter_physical_types=None,
ndim=3, units=None):
self.units = self._validate_units(units)
if parameter_physical_types is None:
parameter_physical_types = dict()
self._ptypes = parameter_physical_types
self.parameters = self._prepare_parameters(parameters, self._ptypes,
self.units)
try:
self.G = G.decompose(self.units).value
except u.UnitConversionError:
self.G = 1. # TODO: this is a HACK and could lead to user confusion
self.ndim = ndim
if origin is None:
origin = np.zeros(self.ndim)
self.origin = self._remove_units(origin)
def __init__(self, parameters, units=None):
# make sure the units specified are a UnitSystem instance
if units is not None and not isinstance(units, UnitSystem):
units = UnitSystem(*units)
elif units is None:
units = DimensionlessUnitSystem()
self.units = units
# in case the user specified an ordered dict
self.parameters = OrderedDict()
for k, v in parameters.items():
if hasattr(v, 'unit'):
self.parameters[k] = v.decompose(self.units)
else:
self.parameters[k] = v * u.one
try:
self.G = G.decompose(self.units).value
except u.UnitConversionError:
self.G = 1. # TODO: this is a HACK and could lead to user confusion
def __init__(self, parameters, units=None):
# make sure the units specified are a UnitSystem instance
if units is not None and not isinstance(units, UnitSystem):
units = UnitSystem(*units)
elif units is None:
units = DimensionlessUnitSystem()
self.units = units
# in case the user specified an ordered dict
self.parameters = OrderedDict()
for k,v in parameters.items():
if hasattr(v, 'unit'):
self.parameters[k] = v.decompose(self.units)
else:
self.parameters[k] = v*u.one
try:
self.G = G.decompose(self.units).value
except u.UnitConversionError:
self.G = 1. # TODO: this is a HACK and could lead to user confusion
def mass_enclosed(self, q, t=0.):
"""
Estimate the mass enclosed within the given position by assuming the potential
is spherical.
Parameters
----------
q : array_like, numeric
Position to estimate the enclossed mass.
Returns
-------
menc : `~astropy.units.Quantity`
Mass enclosed at the given position(s). If the input position
has shape ``q.shape``, the output energy will have shape
``q.shape[1:]``.
"""
q = self._prefilter_pos(q)
# small step-size in direction of q
h = 1E-3 # MAGIC NUMBER
# Radius
r = np.sqrt(np.sum(q**2, axis=0))
epsilon = h * q / r[np.newaxis]
dPhi_dr_plus = self._value(q + epsilon, t=t)
dPhi_dr_minus = self._value(q - epsilon, t=t)
diff = (dPhi_dr_plus - dPhi_dr_minus)
if isinstance(self.units, DimensionlessUnitSystem):
Gee = 1.
else:
Gee = G.decompose(self.units).value
return np.abs(r * r * diff / Gee / (2. * h)) * self.units['mass']
def mass_enclosed(self, q, t=0.):
"""
Estimate the mass enclosed within the given position by assuming the potential
is spherical.
Parameters
----------
x : array_like, numeric
Position to estimate the enclossed mass.
Returns
-------
menc : `~astropy.units.Quantity`
The potential energy or value of the potential. If the input
position has shape ``q.shape``, the output energy will have
shape ``q.shape[1:]``.
"""
q = self._prefilter_pos(q)
# Fractional step-size in radius
h = 0.01
# Radius
r = np.sqrt(np.sum(q**2, axis=0))
epsilon = h*q/r[np.newaxis]
dPhi_dr_plus = self.value(q + epsilon, t=t)
dPhi_dr_minus = self.value(q - epsilon, t=t)
diff = dPhi_dr_plus - dPhi_dr_minus
if isinstance(self.units, DimensionlessUnitSystem):
raise ValueError("No units specified when creating potential object.")
Gee = G.decompose(self.units).value
return np.abs(r*r * diff / Gee / (2.*h)) * self.units['mass']
def _read_units(self):
""" Read and parse the SCFPAR file containing simulation parameters
and initial conditions. Right now, only parse out the simulation
units.
"""
pars = dict()
parfile = os.path.join(self.path, "SCFPAR")
with open(parfile) as f:
lines = f.readlines()
# find what G is set to
for i,line in enumerate(lines):
if line.split()[1].strip() == "G":
break
pars['G'] = float(lines[i].split()[0])
pars['length'] = float(lines[i+10].split()[0])
pars['mass'] = float(lines[i+11].split()[0])
self.x0 = np.array(map(float, lines[19].split()))*u.kpc
self.v0 = np.array(map(float, lines[20].split()))*u.km/u.s
_G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
X = (_G / pars['length']**3 * pars['mass'])**-0.5
length_unit = u.Unit("{0} kpc".format(pars['length']))
mass_unit = u.Unit("{0} M_sun".format(pars['mass']))
time_unit = u.Unit("{:08f} Myr".format(X))
units = dict(length=length_unit,
mass=mass_unit,
time=time_unit,
speed=length_unit/time_unit,
dimensionless=u.dimensionless_unscaled)
return units
def mass_enclosed(self, q, t=0.):
"""
Estimate the mass enclosed within the given position by assuming the potential
is spherical.
Parameters
----------
x : array_like, numeric
Position to estimate the enclossed mass.
Returns
-------
menc : `~numpy.ndarray`
The mass. Will have the same shape as the input position,
array, ``q``, but without the coordinate axis, ``axis=0``
"""
q = np.ascontiguousarray(atleast_2d(q, insert_axis=1))
# Fractional step-size in radius
h = 0.01
# Radius
r = np.sqrt(np.sum(q**2, axis=0))
epsilon = h*q/r[np.newaxis]
dPhi_dr_plus = self.value(q + epsilon, t=t)
dPhi_dr_minus = self.value(q - epsilon, t=t)
diff = dPhi_dr_plus - dPhi_dr_minus
if self.units is None:
raise ValueError("No units specified when creating potential object.")
Gee = G.decompose(self.units).value
return np.abs(r*r * diff / Gee / (2.*h))
# Third-party
import astropy.units as u
from astropy.constants import G as _G
import numpy as np
import pytest
# Project
from gala._cconfig import GSL_ENABLED
from gala.units import galactic
import gala.potential as gp
from gala.potential.potential.tests.helpers import PotentialTestBase
from gala.potential.potential.io import load
from .. import _bfe_class
G = _G.decompose(galactic).value
if not GSL_ENABLED:
pytest.skip("skipping SCF tests: they depend on GSL",
allow_module_level=True)
def test_hernquist():
nmax = 6
lmax = 2
M = 1E10
r_s = 3.5
cos_coeff = np.zeros((nmax+1,lmax+1,lmax+1))
sin_coeff = np.zeros((nmax+1,lmax+1,lmax+1))
def test_compose_into_arbitrary_units():
# Issue #1438
from astropy.constants import G
G.decompose([u.kg, u.km, u.Unit("15 s")])
def test_constants():
usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun)
assert np.allclose(usys.get_constant('G'), G.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value)
assert np.allclose(usys.get_constant('c'), c.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value)
def isochrone_aa_to_xv(actions, angles, potential):
"""
Transform the input actions and angles to ordinary phase space (position
and velocity) in cartesian coordinates. See Section 3.5.2 in
Binney & Tremaine (2008), and be aware of the errata entry for
Eq. 3.225.
.. note::
This function is included as a method of the :class:`~gary.potential.IsochronePotential`
and it is recommended to call :meth:`~gary.potential.IsochronePotential.action_angle()`
instead.
# TODO: should accept quantities
Parameters
----------
actions : array_like
Action variables. Must have shape (N,3) or (3,).
angles : array_like
Angle variables. Must have shape (N,3) or (3,). Should be in radians.
potential : :class:`gary.potential.IsochronePotential`
An instance of the potential to use for computing the transformation
to angle-action coordinates.
Returns
-------
x : :class:`numpy.ndarray`
An array of cartesian positions computed from the input angles and actions. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
v : :class:`numpy.ndarray`
An array of cartesian velocities computed from the input angles and actions. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
"""
actions = np.atleast_2d(actions)
angles = np.atleast_2d(angles)
_G = G.decompose(potential.units).value
GM = _G * potential.parameters['m']
b = potential.parameters['b']
# actions
Jr = actions[:, 0]
Lz = actions[:, 1]
L = actions[:, 2] + np.abs(Lz)
# angles
theta_r, theta_phi, theta_theta = angles.T
# get longitude of ascending node
theta_1 = theta_phi - np.sign(Lz) * theta_theta
Omega = theta_1
# Ly = -np.cos(Omega) * np.sqrt(L**2 - Lz**2)
# Lx = np.sqrt(L**2 - Ly**2 - Lz**2)
cosi = Lz / L
sini = np.sqrt(1 - cosi**2)
# Hamiltonian (energy)
H = -2. * GM**2 / (2. * Jr + L + np.sqrt(4. * b * GM + L**2))**2
if np.any(H > 0.):
raise ValueError("Unbound particle. (E = {})".format(H))
# Eq. 3.240
c = -GM / (2. * H) - b
e = np.sqrt(1 - L * L * (1 + b / c) / GM / c)
# solve for eta
theta_3 = theta_r
eta_func = lambda x: x - e * c / (b + c) * np.sin(x) - theta_3
eta_func_prime = lambda x: 1 - e * c / (b + c) * np.cos(x)
# use newton's method to find roots
niter = 100
eta = np.ones_like(theta_3) * np.pi / 2.
for i in range(niter):
eta -= eta_func(eta) / eta_func_prime(eta)
# TODO: when to do this???
eta -= 2 * np.pi
r = c * np.sqrt((1 - e * np.cos(eta)) * (1 - e * np.cos(eta) + 2 * b / c))
vr = np.sqrt(GM / (b + c)) * (c * e * np.sin(eta)) / r
theta_2 = theta_theta
Omega_23 = 0.5 * (1 + L / np.sqrt(L**2 + 4 * GM * b))
a = np.sqrt((1 + e) / (1 - e))
ap = np.sqrt((1 + e + 2 * b / c) / (1 - e + 2 * b / c))
def F(x, y):
z = np.zeros_like(x)
ix = y > np.pi / 2.
z[ix] = np.pi / 2. - np.arctan(
np.tan(np.pi / 2. - 0.5 * y[ix]) / x[ix])
ix = y < -np.pi / 2.
z[ix] = -np.pi / 2. + np.arctan(
np.tan(np.pi / 2. + 0.5 * y[ix]) / x[ix])
ix = (y <= np.pi / 2) & (y >= -np.pi / 2)
z[ix] = np.arctan(x[ix] * np.tan(0.5 * y[ix]))
return z
theta_2[Lz < 0] -= 2 * np.pi
theta_3 -= 2 * np.pi
A = Omega_23 * theta_3 - F(
a, eta) - F(ap, eta) / np.sqrt(1 + 4 * GM * b / L / L)
psi = theta_2 - A
# theta
theta = np.arccos(np.sin(psi) * sini)
vtheta = L * sini * np.cos(psi) / np.cos(theta)
vtheta = -L * sini * np.cos(psi) / np.sin(theta) / r
vphi = Lz / (r * np.sin(theta))
# phi
sinu = np.sin(psi) * cosi / np.sin(theta)
uu = np.arcsin(sinu)
uu[sinu > 1.] = np.pi / 2.
uu[sinu < -1.] = -np.pi / 2.
uu[vtheta > 0.] = np.pi - uu[vtheta > 0.]
sinu = cosi / sini * np.cos(theta) / np.sin(theta)
phi = (uu + Omega) % (2 * np.pi)
# We now need to convert from spherical polar coord to cart. coord.
pos = coord.PhysicsSphericalRepresentation(r=r * u.dimensionless_unscaled,
phi=phi * u.rad,
theta=theta * u.rad)
x = pos.represent_as(coord.CartesianRepresentation).xyz.T.value
v = physicsspherical_to_cartesian(pos, [vr, vphi, vtheta] *
u.dimensionless_unscaled).T.value
return x, v
# Standard library
from collections import OrderedDict
# Third party
import pytest
import numpy as np
from astropy.constants import G
import astropy.units as u
from matplotlib import cm
# This package
from ..core import PotentialBase, CompositePotential
from ....units import UnitSystem
units = [u.kpc, u.Myr, u.Msun, u.radian]
G = G.decompose(units)
def test_new_simple():
class MyPotential(PotentialBase):
def __init__(self, units=None):
super(MyPotential, self).__init__(parameters={},
units=units)
def _energy(self, r, t=0.):
return -1/r
def _gradient(self, r, t=0.):
return r**-2
p = MyPotential()
# Third-party
from astropy.constants import G
import astropy.time as at
import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
from gala.units import UnitSystem
# Project
from .celestialmechanics import rv_from_elements
from .util import find_t0
from .units import usys
__all__ = ['RVOrbit', 'SimulatedRVOrbit']
_G = G.decompose(usys).value
EPOCH = 55555. # Magic Number: used for un-modding phi0
class RVOrbit(object):
"""
Parameters
----------
P : `~astropy.units.Quantity` [time]
Orbital period.
a_sin_i : `~astropy.units.Quantity` [length]
Semi-major axis times sine of inclination angle.
ecc : numeric
Eccentricity.
omega : `~astropy.units.Quantity` [angle]
Argument of perihelion.
def q_p(**kwargs):
filename = os.path.join(plot_path, "q_p.pdf")
fig, axes = plt.subplots(2, 4, figsize=(14, 7.5), sharex=True, sharey=True)
bins = np.linspace(0., 10, 40)
nparticles = 5000
for kk, _m in enumerate(range(6, 9 + 1)):
mass = "2.5e{}".format(_m)
m = float(mass)
print(mass)
sgr = SgrSimulation(sgr_path.format(_m), snapfile)
p = sgr.particles(n=nparticles, expr="(tub!=0)") #" & (tub<400)")
tub = p.tub
s = sgr.satellite()
potential = LawMajewski2010()
X = np.vstack((s._X[..., :3], p._X[..., :3].copy()))
V = np.vstack((s._X[..., 3:], p._X[..., 3:].copy()))
integrator = LeapfrogIntegrator(potential._acceleration_at,
np.array(X),
np.array(V),
args=(X.shape[0], np.zeros_like(X)))
ts, rs, vs = integrator.run(t1=sgr.t1, t2=sgr.t2, dt=-1.)
s_orbit = np.vstack(
(rs[:, 0][:, np.newaxis].T, vs[:, 0][:, np.newaxis].T)).T
p_orbits = np.vstack((rs[:, 1:].T, vs[:, 1:].T)).T
t_idx = np.array([np.argmin(np.fabs(ts - t)) for t in p.tub])
p_x = np.array([p_orbits[jj, ii] for ii, jj in enumerate(t_idx)])
s_x = np.array([s_orbit[jj, 0] for jj in t_idx])
#############################################
# determine tail_bit
diff = p_x - s_x
norm_r = s_x[:, :3] / np.sqrt(np.sum(s_x[:, :3]**2,
axis=-1))[:, np.newaxis]
norm_diff_r = diff[:, :3] / np.sqrt(np.sum(diff[:, :3]**2,
axis=-1))[:, np.newaxis]
dot_prod_r = np.sum(norm_diff_r * norm_r, axis=-1)
tail_bit = (dot_prod_r > 0.).astype(int) * 2 - 1
#############################################
r_tide = potential._tidal_radius(m, s_orbit[..., :3]) #*0.69336
s_R_orbit = np.sqrt(np.sum(s_orbit[..., :3]**2, axis=-1))
a_pm = (s_R_orbit + r_tide * tail_bit) / s_R_orbit
q = np.sqrt(np.sum((p_x[:, :3] - s_x[:, :3])**2, axis=-1))
f = r_tide / s_R_orbit
s_V = np.sqrt(np.sum(s_orbit[..., 3:]**2, axis=-1))
vdisp = s_V * f / 1.4
p = np.sqrt(np.sum((p_x[:, 3:] - s_x[..., 3:])**2, axis=-1))
fig, axes = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
axes[0].plot(tub, q, marker='.', alpha=0.5, color='#666666')
axes[0].plot(ts,
r_tide * 1.4,
linewidth=2.,
alpha=0.8,
color='k',
linestyle='-',
marker=None)
axes[0].set_ylim(0., max(r_tide) * 4)
axes[1].plot(tub, (p * u.kpc / u.Myr).to(u.km / u.s).value,
marker='.',
alpha=0.5,
color='#666666')
axes[1].plot(ts, (vdisp * u.kpc / u.Myr).to(u.km / u.s).value,
color='k',
linewidth=2.,
alpha=0.75,
linestyle='-',
marker=None)
M_enc = potential._enclosed_mass(s_R_orbit)
#delta_E = 4/3.*G.decompose(usys).value**2*m*(M_enc / s_V)**2*r_tide**2/s_R_orbit**4
delta_v2 = 4/3.*G.decompose(usys).value**2*(M_enc / s_V)**2*\
np.mean(r_tide**2)/s_R_orbit**4
delta_v = (np.sqrt(2 * delta_v2) * u.kpc / u.Myr).to(u.km / u.s).value
axes[1].plot(ts,
delta_v,
linewidth=2.,
color='#2166AC',
alpha=0.75,
linestyle='--',
marker=None)
axes[1].set_ylim(0.,
max((vdisp * u.kpc / u.Myr).to(u.km / u.s).value) * 4)
axes[0].set_xlim(min(ts), max(ts))
fig.savefig(os.path.join(plot_path, "q_p_{}.png".format(mass)),
transparent=True)
def __init__(self, galcen, element_data, abundance_names,
frozen_pars=None, usys=None,
metals_deg=3,
marginalize=None):
"""TODO
Parameters
----------
galcen : `astropy.coordinates.CartesianRepresentation`
element_data : TODO
abundance_names : iterable
metals_deg : int (optional)
The degree of the polynomial used to represent the abundance
distribution mean as a function of invariants.
usys : `gala.units.UnitSystem` (optional)
The unit system to work in. Defaults to (pc, Myr, Msun, km/s).
"""
# Make sure the unit system has a default
if usys is None:
usys = UnitSystem(u.pc, u.Myr, u.Msun, u.radian, u.km/u.s)
self.usys = usys
self.marginalize_all = False
self.marginalize_alpha = False
if marginalize == 'all':
self.marginalize_all = True
elif marginalize == 'alpha':
self.marginalize_alpha = True
# kinematic and abundance data
self.galcen = galcen
# put requested abundances into a dictionary
self.element_data = dict()
for name in abundance_names:
self.element_data[name] = get_abundance_data(element_data, name)
self.abundance_names = abundance_names
# if enforce_finite:
# elem_mask = np.ones(len(galcen), dtype=bool)
# for name in abundance_names:
# elem_mask &= np.isfinite(self.element_data[name])
#
# for name in abundance_names:
# self.element_data[name] = self.element_data[name][elem_mask]
# self.galcen = self.galcen[elem_mask]
if frozen_pars is None:
frozen_pars = dict()
self.frozen_pars = frozen_pars
# Convert data to units we expect
self._z = self.galcen.z.decompose(self.usys).value
self._vz = self.galcen.v_z.decompose(self.usys).value
# TODO HACK: hard-coded for now!
self._potential = sech2_potential
self._G = G.decompose(self.usys).value
# Cache array used below
self.metals_deg = int(metals_deg)
self._AT = np.ones((self.metals_deg+1, len(self.galcen)))
# coding: utf-8
from __future__ import division, print_function
__author__ = "adrn <*****@*****.**>"
# Third-party
import astropy.units as u
from astropy.constants import G as _G
G = _G.decompose([u.kpc,u.Myr,u.Msun]).value
import numpy as np
import gary.potential as gp
from gary.units import galactic
# Project
from .._bfe import SCFPotential
def test_hernquist():
nmax = 6
lmax = 2
M = 1E10
r_s = 3.5
cos_coeff = np.zeros((nmax+1,lmax+1,lmax+1))
sin_coeff = np.zeros((nmax+1,lmax+1,lmax+1))
cos_coeff[0,0,0] = 1.
scf_potential = SCFPotential(m=M, r_s=r_s,
Snlm=cos_coeff, Tnlm=sin_coeff,
units=galactic)
def isochrone_aa_to_xv(actions, angles, potential):
"""
Transform the input actions and angles to ordinary phase space (position
and velocity) in cartesian coordinates. See Section 3.5.2 in
Binney & Tremaine (2008), and be aware of the errata entry for
Eq. 3.225.
.. note::
This function is included as a method of the :class:`~gary.potential.IsochronePotential`
and it is recommended to call :meth:`~gary.potential.IsochronePotential.action_angle()`
instead.
# TODO: should accept quantities
Parameters
----------
actions : array_like
Action variables. Must have shape (N,3) or (3,).
angles : array_like
Angle variables. Must have shape (N,3) or (3,). Should be in radians.
potential : :class:`gary.potential.IsochronePotential`
An instance of the potential to use for computing the transformation
to angle-action coordinates.
Returns
-------
x : :class:`numpy.ndarray`
An array of cartesian positions computed from the input angles and actions. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
v : :class:`numpy.ndarray`
An array of cartesian velocities computed from the input angles and actions. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
"""
actions = np.atleast_2d(actions)
angles = np.atleast_2d(angles)
_G = G.decompose(potential.units).value
GM = _G*potential.parameters['m']
b = potential.parameters['b']
# actions
Jr = actions[:,0]
Lz = actions[:,1]
L = actions[:,2] + np.abs(Lz)
# angles
theta_r,theta_phi,theta_theta = angles.T
# get longitude of ascending node
theta_1 = theta_phi - np.sign(Lz)*theta_theta
Omega = theta_1
# Ly = -np.cos(Omega) * np.sqrt(L**2 - Lz**2)
# Lx = np.sqrt(L**2 - Ly**2 - Lz**2)
cosi = Lz/L
sini = np.sqrt(1 - cosi**2)
# Hamiltonian (energy)
H = -2. * GM**2 / (2.*Jr + L + np.sqrt(4.*b*GM + L**2))**2
if np.any(H > 0.):
raise ValueError("Unbound particle. (E = {})".format(H))
# Eq. 3.240
c = -GM / (2.*H) - b
e = np.sqrt(1 - L*L*(1 + b/c) / GM / c)
# solve for eta
theta_3 = theta_r
eta_func = lambda x: x - e*c/(b+c)*np.sin(x) - theta_3
eta_func_prime = lambda x: 1 - e*c/(b+c)*np.cos(x)
# use newton's method to find roots
niter = 100
eta = np.ones_like(theta_3)*np.pi/2.
for i in range(niter):
eta -= eta_func(eta)/eta_func_prime(eta)
# TODO: when to do this???
eta -= 2*np.pi
r = c*np.sqrt((1-e*np.cos(eta)) * (1-e*np.cos(eta) + 2*b/c))
vr = np.sqrt(GM/(b+c))*(c*e*np.sin(eta))/r
theta_2 = theta_theta
Omega_23 = 0.5*(1 + L / np.sqrt(L**2 + 4*GM*b))
a = np.sqrt((1+e) / (1-e))
ap = np.sqrt((1 + e + 2*b/c) / (1 - e + 2*b/c))
def F(x, y):
z = np.zeros_like(x)
ix = y>np.pi/2.
z[ix] = np.pi/2. - np.arctan(np.tan(np.pi/2.-0.5*y[ix])/x[ix])
ix = y<-np.pi/2.
z[ix] = -np.pi/2. + np.arctan(np.tan(np.pi/2.+0.5*y[ix])/x[ix])
ix = (y<=np.pi/2) & (y>=-np.pi/2)
z[ix] = np.arctan(x[ix]*np.tan(0.5*y[ix]))
return z
theta_2[Lz < 0] -= 2*np.pi
theta_3 -= 2*np.pi
A = Omega_23*theta_3 - F(a,eta) - F(ap,eta)/np.sqrt(1 + 4*GM*b/L/L)
psi = theta_2 - A
# theta
theta = np.arccos(np.sin(psi)*sini)
vtheta = L*sini*np.cos(psi)/np.cos(theta)
vtheta = -L*sini*np.cos(psi)/np.sin(theta)/r
vphi = Lz / (r*np.sin(theta))
# phi
sinu = np.sin(psi)*cosi/np.sin(theta)
uu = np.arcsin(sinu)
uu[sinu > 1.] = np.pi/2.
uu[sinu < -1.] = -np.pi/2.
uu[vtheta > 0.] = np.pi - uu[vtheta > 0.]
sinu = cosi/sini * np.cos(theta)/np.sin(theta)
phi = (uu + Omega) % (2*np.pi)
# We now need to convert from spherical polar coord to cart. coord.
pos = coord.PhysicsSphericalRepresentation(r=r*u.dimensionless_unscaled,
phi=phi*u.rad, theta=theta*u.rad)
x = pos.represent_as(coord.CartesianRepresentation).xyz.T.value
v = physicsspherical_to_cartesian(pos, [vr,vphi,vtheta]*u.dimensionless_unscaled).T.value
return x,v
def total_rv():
filenamer = os.path.join(plot_path, "rel_r.png")
filenamev = os.path.join(plot_path, "rel_v.png")
figr,axesr = plt.subplots(4,1,figsize=(10,14),
sharex=True)
figv,axesv = plt.subplots(4,1,figsize=(10,14),
sharex=True)
nparticles = 2000
for k,_m in enumerate(range(6,9+1)):
mass = "2.5e{}".format(_m)
m = float(mass)
print(mass)
sgr = SgrSimulation(sgr_path.format(_m),snapfile)
p = sgr.particles(n=nparticles, expr=expr)
s = sgr.satellite()
X = np.vstack((s._X[...,:3], p._X[...,:3].copy()))
V = np.vstack((s._X[...,3:], p._X[...,3:].copy()))
integrator = LeapfrogIntegrator(sgr.potential._acceleration_at,
np.array(X), np.array(V),
args=(X.shape[0], np.zeros_like(X)))
ts, rs, vs = integrator.run(t1=sgr.t1, t2=sgr.t2, dt=-1.)
s_orbit = np.vstack((rs[:,0][:,np.newaxis].T, vs[:,0][:,np.newaxis].T)).T
p_orbits = np.vstack((rs[:,1:].T, vs[:,1:].T)).T
t_idx = np.array([np.argmin(np.fabs(ts - t)) for t in p.tub])
m_t = (-s.mdot*ts + s.m0)[:,np.newaxis]
s_R = np.sqrt(np.sum(s_orbit[...,:3]**2, axis=-1))
s_V = np.sqrt(np.sum(s_orbit[...,3:]**2, axis=-1))
r_tide = sgr.potential._tidal_radius(m_t, s_orbit[...,:3])
v_disp = s_V * r_tide / s_R
# cartesian basis to project into
x_hat = s_orbit[...,:3] / np.sqrt(np.sum(s_orbit[...,:3]**2, axis=-1))[...,np.newaxis]
_y_hat = s_orbit[...,3:] / np.sqrt(np.sum(s_orbit[...,3:]**2, axis=-1))[...,np.newaxis]
z_hat = np.cross(x_hat, _y_hat)
y_hat = -np.cross(x_hat, z_hat)
# translate to satellite position
rel_orbits = p_orbits - s_orbit
rel_pos = rel_orbits[...,:3]
rel_vel = rel_orbits[...,3:]
# project onto each
X = np.sum(rel_pos * x_hat, axis=-1)
Y = np.sum(rel_pos * y_hat, axis=-1)
Z = np.sum(rel_pos * z_hat, axis=-1)
RR = np.sqrt(X**2 + Y**2 + Z**2)
VX = np.sum(rel_vel * x_hat, axis=-1)
VY = np.sum(rel_vel * y_hat, axis=-1)
VZ = np.sum(rel_vel * z_hat, axis=-1)
VV = (np.sqrt(VX**2 + VY**2 + VZ**2)*u.kpc/u.Myr).to(u.km/u.s).value
v_disp = (v_disp*u.kpc/u.Myr).to(u.km/u.s).value
_tcross = r_tide / np.sqrt(G.decompose(usys).value*m/r_tide)
for ii,jj in enumerate(t_idx):
#tcross = r_tide[jj,0] / _v[jj,ii]
tcross = _tcross[jj]
bnd = int(tcross / 2)
ix1,ix2 = jj-bnd, jj+bnd
if ix1 < 0: ix1 = 0
if ix2 > max(sgr.t1,sgr.t2): ix2 = -1
axesr[k].plot(ts[ix1:ix2],
RR[ix1:ix2,ii],
linestyle='-', alpha=0.1, marker=None, color='#555555', zorder=-1)
axesv[k].plot(ts[ix1:ix2],
VV[ix1:ix2,ii],
linestyle='-', alpha=0.1, marker=None, color='#555555', zorder=-1)
axesr[k].plot(ts, r_tide*2., marker=None)
axesr[k].set_xlim(ts.min(), ts.max())
axesv[k].set_xlim(ts.min(), ts.max())
axesr[k].set_ylim(0,max(r_tide)*7)
axesv[k].set_ylim(0,max(v_disp)*7)
# axes[1,k].set_xlabel(r"$x_1$")
# if k == 0:
# axes[0,k].set_ylabel(r"$x_2$")
# axes[1,k].set_ylabel(r"$x_3$")
axesr[k].text(3000, max(r_tide)*5, r"$2.5\times10^{}M_\odot$".format(_m))
axesv[k].text(3000, max(v_disp)*5, r"$2.5\times10^{}M_\odot$".format(_m))
axesr[-1].set_xlabel("time [Myr]")
axesv[-1].set_xlabel("time [Myr]")
figr.suptitle("Relative distance", fontsize=26)
figr.tight_layout()
figr.subplots_adjust(top=0.92, hspace=0.025, wspace=0.1)
figr.savefig(filenamer)
figv.suptitle("Relative velocity", fontsize=26)
figv.tight_layout()
figv.subplots_adjust(top=0.92, hspace=0.025, wspace=0.1)
figv.savefig(filenamev)
def Lpts():
np.random.seed(42)
potential = LawMajewski2010()
filename = os.path.join(plot_path, "Lpts_r.{}".format(ext))
filename2 = os.path.join(plot_path, "Lpts_v.{}".format(ext))
fig,axes = plt.subplots(2,4,figsize=grid_figsize,
sharex=True, sharey=True)
fig2,axes2 = plt.subplots(2,4,figsize=grid_figsize,
sharex=True, sharey=True)
bins = np.linspace(-3,3,50)
nparticles = 2000
for k,_m in enumerate(range(6,9+1)):
mass = "2.5e{}".format(_m)
m = float(mass)
print(mass)
sgr = SgrSimulation(sgr_path.format(_m),snapfile)
p = sgr.particles(n=nparticles, expr=expr)
s = sgr.satellite()
dt = -1.
coord, r_tide, v_disp = particles_x1x2x3(p, s,
sgr.potential,
sgr.t1, sgr.t2, dt,
at_tub=False)
(x1,x2,x3,vx1,vx2,vx3) = coord
ts = np.arange(sgr.t1,sgr.t2+dt,dt)
t_idx = np.array([np.argmin(np.fabs(ts - t)) for t in p.tub])
_tcross = r_tide / np.sqrt(G.decompose(usys).value*m/r_tide)
for ii,jj in enumerate(t_idx):
#tcross = r_tide[jj,0] / _v[jj,ii]
tcross = _tcross[jj]
bnd = int(tcross / 2)
ix1,ix2 = jj-bnd, jj+bnd
if ix1 < 0: ix1 = 0
if ix2 > max(sgr.t1,sgr.t2): ix2 = -1
axes[0,k].set_rasterization_zorder(1)
axes[0,k].plot(x1[jj-bnd:jj+bnd,ii]/r_tide[jj-bnd:jj+bnd,0],
x2[jj-bnd:jj+bnd,ii]/r_tide[jj-bnd:jj+bnd,0],
linestyle='-', alpha=0.1, marker=None, color='#555555',
zorder=-1)
axes[1,k].set_rasterization_zorder(1)
axes[1,k].plot(x1[jj-bnd:jj+bnd,ii]/r_tide[jj-bnd:jj+bnd,0],
x3[jj-bnd:jj+bnd,ii]/r_tide[jj-bnd:jj+bnd,0],
linestyle='-', alpha=0.1, marker=None, color='#555555',
zorder=-1)
circ = Circle((0,0), radius=1., fill=False, alpha=0.75,
edgecolor='k', linestyle='solid')
axes[0,k].add_patch(circ)
circ = Circle((0,0), radius=1., fill=False, alpha=0.75,
edgecolor='k', linestyle='solid')
axes[1,k].add_patch(circ)
axes[0,k].axhline(0., color='k', alpha=0.75)
axes[1,k].axhline(0., color='k', alpha=0.75)
axes[0,k].set_xlim(-5,5)
axes[0,k].set_ylim(axes[0,k].get_xlim())
axes[1,k].set_xlabel(r"$x_1/r_{\rm tide}$")
if k == 0:
axes[0,k].set_ylabel(r"$x_2/r_{\rm tide}$")
axes[1,k].set_ylabel(r"$x_3/r_{\rm tide}$")
_tcross = r_tide / np.sqrt(G.decompose(usys).value*m/r_tide)
for ii,jj in enumerate(t_idx):
#tcross = r_tide[jj,0] / _v[jj,ii]
tcross = _tcross[jj]
bnd = int(tcross / 2)
ix1,ix2 = jj-bnd, jj+bnd
if ix1 < 0: ix1 = 0
if ix2 > max(sgr.t1,sgr.t2): ix2 = -1
axes2[0,k].set_rasterization_zorder(1)
axes2[0,k].plot(vx1[jj-bnd:jj+bnd,ii]/v_disp[jj-bnd:jj+bnd,0],
vx2[jj-bnd:jj+bnd,ii]/v_disp[jj-bnd:jj+bnd,0],
linestyle='-', alpha=0.1, marker=None, color='#555555',
zorder=-1)
axes2[1,k].set_rasterization_zorder(1)
axes2[1,k].plot(vx1[jj-bnd:jj+bnd,ii]/v_disp[jj-bnd:jj+bnd,0],
vx3[jj-bnd:jj+bnd,ii]/v_disp[jj-bnd:jj+bnd,0],
linestyle='-', alpha=0.1, marker=None, color='#555555',
zorder=-1)
circ = Circle((0,0), radius=1., fill=False, alpha=0.75,
edgecolor='k', linestyle='solid')
axes2[0,k].add_patch(circ)
circ = Circle((0,0), radius=1., fill=False, alpha=0.75,
edgecolor='k', linestyle='solid')
axes2[1,k].add_patch(circ)
axes2[0,k].axhline(0., color='k', alpha=0.75)
axes2[1,k].axhline(0., color='k', alpha=0.75)
axes2[1,k].set_xlim(-5,5)
axes2[1,k].set_ylim(axes2[1,k].get_xlim())
axes2[1,k].set_xlabel(r"$v_{x_1}/\sigma_v$")
if k == 0:
axes2[0,k].set_ylabel(r"$v_{x_2}/\sigma_v$")
axes2[1,k].set_ylabel(r"$v_{x_3}/\sigma_v$")
axes[0,k].text(0.5, 1.05, r"$2.5\times10^{}M_\odot$".format(_m),
horizontalalignment='center',
fontsize=24,
transform=axes[0,k].transAxes)
axes2[0,k].text(0.5, 1.05, r"$2.5\times10^{}M_\odot$".format(_m),
horizontalalignment='center',
fontsize=24,
transform=axes2[0,k].transAxes)
fig.tight_layout()
fig.subplots_adjust(top=0.92, hspace=0.025, wspace=0.1)
fig.savefig(filename)
fig2.tight_layout()
fig2.subplots_adjust(top=0.92, hspace=0.025, wspace=0.1)
fig2.savefig(filename2)
elif args.quiet:
logger.setLevel(logging.ERROR)
else:
logger.setLevel(logging.INFO)
# Contains user-specified parameters for SCF
scfpars = dict()
if args.rscale is None:
ru = 0.43089*(args.mass/2.5e9)**(1/3.)
else:
ru = args.rscale
scfpars['rscale'] = ru
scfpars['mass'] = args.mass
_G = G.decompose(bases=[u.kpc,u.M_sun,u.Myr]).value
X = (_G / ru**3 * args.mass)**-0.5
length_unit = u.Unit("{0} kpc".format(ru))
mass_unit = u.Unit("{0} M_sun".format(args.mass))
time_unit = u.Unit("{:08f} Myr".format(X))
scfpars['dt'] = args.dt / (1*time_unit).to(u.Myr).value
scfpars['nsteps'] = args.nsteps
scfpars['ncen'] = args.ncen
scfpars['nsnap'] = args.nsnap
scfpars['ntide'] = args.ntide
if args.scfbi_path is None:
scfpars['SCFBIpath'] = os.path.abspath(os.path.join(project_path, 'src', 'SCFBI'))
main(name=args.name, pos=args.x, vel=args.v, scfpars=scfpars,
elif args.quiet:
logger.setLevel(logging.ERROR)
else:
logger.setLevel(logging.INFO)
# Contains user-specified parameters for SCF
scfpars = dict()
if args.rscale is None:
ru = 0.43089 * (args.mass / 2.5e9)**(1 / 3.)
else:
ru = args.rscale
scfpars['rscale'] = ru
scfpars['mass'] = args.mass
_G = G.decompose(bases=[u.kpc, u.M_sun, u.Myr]).value
X = (_G / ru**3 * args.mass)**-0.5
length_unit = u.Unit("{0} kpc".format(ru))
mass_unit = u.Unit("{0} M_sun".format(args.mass))
time_unit = u.Unit("{:08f} Myr".format(X))
scfpars['dt'] = args.dt / (1 * time_unit).to(u.Myr).value
scfpars['nsteps'] = args.nsteps
scfpars['ncen'] = args.ncen
scfpars['nsnap'] = args.nsnap
scfpars['ntide'] = args.ntide
if args.scfbi_path is None:
scfpars['SCFBIpath'] = os.path.abspath(
os.path.join(project_path, 'src', 'SCFBI'))
def test_compose_into_arbitrary_units():
# Issue #1438
from astropy.constants import G
G.decompose([u.kg, u.km, u.Unit("15 s")])
from ..core import CompositePotential
from ..cbuiltin import *
from ..io import load
from ...units import galactic, solarsystem
# HACK: bad solution is to do this:
# python setup.py build_ext --inplace
top_path = "plots/"
plot_path = os.path.join(top_path, "tests/potential/cpotential")
if not os.path.exists(plot_path):
os.makedirs(plot_path)
units = [u.kpc,u.Myr,u.Msun,u.radian]
G = G.decompose(units)
print()
color_print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~", "yellow")
color_print("To view plots:", "green")
print(" open {}".format(plot_path))
color_print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~", "yellow")
niter = 1000
nparticles = 1000
class PotentialTestBase(object):
name = None
units = galactic
def setup(self):
def __init__(self, m, a, units):
self.parameters = dict(m=m, a=a)
self.units = units
self.G = G.decompose(units).value
def q_p(**kwargs):
filename = os.path.join(plot_path, "q_p.pdf")
fig,axes = plt.subplots(2,4,figsize=(14,7.5),
sharex=True, sharey=True)
bins = np.linspace(0.,10,40)
nparticles = 5000
for kk,_m in enumerate(range(6,9+1)):
mass = "2.5e{}".format(_m)
m = float(mass)
print(mass)
sgr = SgrSimulation(sgr_path.format(_m), snapfile)
p = sgr.particles(n=nparticles, expr="(tub!=0)")#" & (tub<400)")
tub = p.tub
s = sgr.satellite()
potential = LawMajewski2010()
X = np.vstack((s._X[...,:3], p._X[...,:3].copy()))
V = np.vstack((s._X[...,3:], p._X[...,3:].copy()))
integrator = LeapfrogIntegrator(potential._acceleration_at,
np.array(X), np.array(V),
args=(X.shape[0], np.zeros_like(X)))
ts, rs, vs = integrator.run(t1=sgr.t1, t2=sgr.t2, dt=-1.)
s_orbit = np.vstack((rs[:,0][:,np.newaxis].T, vs[:,0][:,np.newaxis].T)).T
p_orbits = np.vstack((rs[:,1:].T, vs[:,1:].T)).T
t_idx = np.array([np.argmin(np.fabs(ts - t)) for t in p.tub])
p_x = np.array([p_orbits[jj,ii] for ii,jj in enumerate(t_idx)])
s_x = np.array([s_orbit[jj,0] for jj in t_idx])
#############################################
# determine tail_bit
diff = p_x-s_x
norm_r = s_x[:,:3] / np.sqrt(np.sum(s_x[:,:3]**2, axis=-1))[:,np.newaxis]
norm_diff_r = diff[:,:3] / np.sqrt(np.sum(diff[:,:3]**2, axis=-1))[:,np.newaxis]
dot_prod_r = np.sum(norm_diff_r*norm_r, axis=-1)
tail_bit = (dot_prod_r > 0.).astype(int)*2 - 1
#############################################
r_tide = potential._tidal_radius(m, s_orbit[...,:3])#*0.69336
s_R_orbit = np.sqrt(np.sum(s_orbit[...,:3]**2, axis=-1))
a_pm = (s_R_orbit + r_tide*tail_bit) / s_R_orbit
q = np.sqrt(np.sum((p_x[:,:3] - s_x[:,:3])**2,axis=-1))
f = r_tide / s_R_orbit
s_V = np.sqrt(np.sum(s_orbit[...,3:]**2, axis=-1))
vdisp = s_V * f / 1.4
p = np.sqrt(np.sum((p_x[:,3:] - s_x[...,3:])**2,axis=-1))
fig,axes = plt.subplots(2,1,figsize=(10,6),sharex=True)
axes[0].plot(tub, q, marker='.', alpha=0.5, color='#666666')
axes[0].plot(ts, r_tide*1.4, linewidth=2., alpha=0.8, color='k',
linestyle='-', marker=None)
axes[0].set_ylim(0., max(r_tide)*4)
axes[1].plot(tub, (p*u.kpc/u.Myr).to(u.km/u.s).value,
marker='.', alpha=0.5, color='#666666')
axes[1].plot(ts, (vdisp*u.kpc/u.Myr).to(u.km/u.s).value, color='k',
linewidth=2., alpha=0.75, linestyle='-', marker=None)
M_enc = potential._enclosed_mass(s_R_orbit)
#delta_E = 4/3.*G.decompose(usys).value**2*m*(M_enc / s_V)**2*r_tide**2/s_R_orbit**4
delta_v2 = 4/3.*G.decompose(usys).value**2*(M_enc / s_V)**2*\
np.mean(r_tide**2)/s_R_orbit**4
delta_v = (np.sqrt(2*delta_v2)*u.kpc/u.Myr).to(u.km/u.s).value
axes[1].plot(ts, delta_v, linewidth=2., color='#2166AC',
alpha=0.75, linestyle='--', marker=None)
axes[1].set_ylim(0., max((vdisp*u.kpc/u.Myr).to(u.km/u.s).value)*4)
axes[0].set_xlim(min(ts), max(ts))
fig.savefig(os.path.join(plot_path, "q_p_{}.png".format(mass)),
transparent=True)
from matplotlib import animation
import matplotlib.pyplot as plt
import numpy as np
from astropy.constants import G
import astropy.units as u
from scipy.signal import argrelmin
from streamteam.integrate import DOPRI853Integrator
from streamteam.dynamics import lyapunov
# Create logger
logger = logging.getLogger(__name__)
plot_path = "plots"
usys = [u.kpc, u.Myr, u.radian, u.M_sun]
_G = G.decompose(usys).value
# standard prolate:
prolate_params = (_G, 1.63E11, 3., 6., 0.2,
5.8E9, 0.25,
0.204542433, 0.5, 8.5)
# standard oblate:
oblate_params = (_G, 1.63E11, 3., 6., 0.2,
5.8E9, 0.25,
0.204542433, 1.5, 8.5)
def zotos_potential(R, z, G, Md, alpha, b, h, Mn, cn, v0, beta, ch):
Rsq = R*R
Vd = -G*Md / np.sqrt(b**2 + Rsq + (alpha + np.sqrt(h**2+z**2))**2)
Vn = -G*Mn / np.sqrt(Rsq+z**2+cn**2)
import matplotlib.pyplot as plt
import numpy as np
import scipy.interpolate as inter
from astropy import log as logger
logger.setLevel(logging.DEBUG)
from astropy.constants import G
# Stream-Team
import streamteam.integrate as si
from streamteam.potential.lm10 import LM10Potential
from streamteam.potential.apw import PW14Potential
from streamteam.dynamics.actionangle import find_actions, fit_isochrone
from streamteam.potential import IsochronePotential
from streamteam.units import galactic
G = G.decompose(galactic).value
cache_path = "/home/spearson/Hessian/stream-team/hessian"
#---------------------- We want to check both LM10() triax & LM10(q1=1,q2=1,q3=1)-------------------------
#---------------------Hessian test - APW-----------------#
ppars = {'a': 6.5,
'b': 0.26,
'c': 0.3,
'm_disk': 65000000000.0,
'm_spher': 20000000000.0,
'phi': 1.570796,
'psi': 1.570796,
'q1': 1.3,
'q2': 1.0,
def isochrone_xv_to_aa(x, v, potential):
"""
Transform the input cartesian position and velocity to action-angle
coordinates in the Isochrone potential. See Section 3.5.2 in
Binney & Tremaine (2008), and be aware of the errata entry for
Eq. 3.225.
This transformation is analytic and can be used as a "toy potential"
in the Sanders & Binney (2014) formalism for computing action-angle
coordinates in any potential.
.. note::
This function is included as a method of the :class:`~gary.potential.IsochronePotential`
and it is recommended to call :meth:`~gary.potential.IsochronePotential.phase_space()`
instead.
# TODO: should accept quantities
Parameters
----------
x : array_like
Cartesian positions. Must have shape (N,3) or (3,).
v : array_like
Cartesian velocities. Must have shape (N,3) or (3,).
potential : :class:`gary.potential.IsochronePotential`
An instance of the potential to use for computing the transformation
to angle-action coordinates.
Returns
-------
actions : :class:`numpy.ndarray`
An array of actions computed from the input positions and velocities. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
angles : :class:`numpy.ndarray`
An array of angles computed from the input positions and velocities. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
"""
x = np.atleast_2d(x)
v = np.atleast_2d(v)
_G = G.decompose(potential.units).value
GM = _G * potential.parameters['m']
b = potential.parameters['b']
E = potential.total_energy(x, v)
x = x * u.dimensionless_unscaled
v = v * u.dimensionless_unscaled
if np.any(E > 0.):
raise ValueError("Unbound particle. (E = {})".format(E))
# convert position, velocity to spherical polar coordinates
x_rep = coord.CartesianRepresentation(x.T)
x_rep = x_rep.represent_as(coord.PhysicsSphericalRepresentation)
v_sph = cartesian_to_physicsspherical(x_rep, v.T)
r, phi, theta = x_rep.r.value, x_rep.phi.value, x_rep.theta.value
vr, vphi, vtheta = v_sph.value
# ----------------------------
# Actions
# ----------------------------
L_vec = angular_momentum(x, v)
Lz = L_vec[:, 2]
L = np.linalg.norm(L_vec, axis=1)
# Radial action
Jr = GM / np.sqrt(-2 * E) - 0.5 * (L + np.sqrt(L * L + 4 * GM * b))
# compute the three action variables
actions = np.array([Jr, Lz, L - np.abs(Lz)]).T
# ----------------------------
# Angles
# ----------------------------
c = GM / (-2 * E) - b
e = np.sqrt(1 - L * L * (1 + b / c) / GM / c)
# Compute theta_r using eta
tmp1 = r * vr / np.sqrt(-2. * E)
tmp2 = b + c - np.sqrt(b * b + r * r)
eta = np.arctan2(tmp1, tmp2)
thetar = eta - e * c * np.sin(eta) / (c + b) # same as theta3
# Compute theta_z
psi = np.arctan2(np.cos(theta), -np.sin(theta) * r * vtheta / L)
psi[np.abs(vtheta) <= 1e-10] = np.pi / 2. # blows up for small vtheta
omega = 0.5 * (1 + L / np.sqrt(L * L + 4 * GM * b))
a = np.sqrt((1 + e) / (1 - e))
ap = np.sqrt((1 + e + 2 * b / c) / (1 - e + 2 * b / c))
def F(x, y):
z = np.zeros_like(x)
ix = y > np.pi / 2.
z[ix] = np.pi / 2. - np.arctan(
np.tan(np.pi / 2. - 0.5 * y[ix]) / x[ix])
ix = y < -np.pi / 2.
z[ix] = -np.pi / 2. + np.arctan(
np.tan(np.pi / 2. + 0.5 * y[ix]) / x[ix])
ix = (y <= np.pi / 2) & (y >= -np.pi / 2)
z[ix] = np.arctan(x[ix] * np.tan(0.5 * y[ix]))
return z
A = omega * thetar - F(a,
eta) - F(ap, eta) / np.sqrt(1 + 4 * GM * b / L / L)
thetaz = psi + A
LR = Lz / L
sinu = (LR / np.sqrt(1. - LR * LR) / np.tan(theta))
sinu = sinu.value
uu = np.arcsin(sinu)
uu[sinu > 1.] = np.pi / 2.
uu[sinu < -1.] = -np.pi / 2.
uu[vtheta > 0.] = np.pi - uu[vtheta > 0.]
thetap = phi - uu + np.sign(Lz) * thetaz
angles = np.array([thetar, thetap, thetaz]).T
angles %= (2 * np.pi)
return actions, angles
def __init__(self, m, a, units):
self.parameters = dict(m=m, a=a)
self.units = units
self.G = G.decompose(units).value
def test_constants():
usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun)
assert np.allclose(usys.get_constant('G'),
G.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value)
assert np.allclose(usys.get_constant('c'),
c.decompose([u.kpc, u.Myr, u.radian, u.Msun]).value)
def isochrone_xv_to_aa(x, v, potential):
"""
Transform the input cartesian position and velocity to action-angle
coordinates in the Isochrone potential. See Section 3.5.2 in
Binney & Tremaine (2008), and be aware of the errata entry for
Eq. 3.225.
This transformation is analytic and can be used as a "toy potential"
in the Sanders & Binney (2014) formalism for computing action-angle
coordinates in any potential.
.. note::
This function is included as a method of the :class:`~gary.potential.IsochronePotential`
and it is recommended to call :meth:`~gary.potential.IsochronePotential.phase_space()`
instead.
# TODO: should accept quantities
Parameters
----------
x : array_like
Cartesian positions. Must have shape (N,3) or (3,).
v : array_like
Cartesian velocities. Must have shape (N,3) or (3,).
potential : :class:`gary.potential.IsochronePotential`
An instance of the potential to use for computing the transformation
to angle-action coordinates.
Returns
-------
actions : :class:`numpy.ndarray`
An array of actions computed from the input positions and velocities. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
angles : :class:`numpy.ndarray`
An array of angles computed from the input positions and velocities. Will
always have shape (N,3) -- if input coordinates are 1D, the output shape will
be (1,3).
"""
x = np.atleast_2d(x)
v = np.atleast_2d(v)
_G = G.decompose(potential.units).value
GM = _G*potential.parameters['m']
b = potential.parameters['b']
E = potential.total_energy(x, v)
x = x * u.dimensionless_unscaled
v = v * u.dimensionless_unscaled
if np.any(E > 0.):
raise ValueError("Unbound particle. (E = {})".format(E))
# convert position, velocity to spherical polar coordinates
x_rep = coord.CartesianRepresentation(x.T)
x_rep = x_rep.represent_as(coord.PhysicsSphericalRepresentation)
v_sph = cartesian_to_physicsspherical(x_rep, v.T)
r,phi,theta = x_rep.r.value, x_rep.phi.value, x_rep.theta.value
vr,vphi,vtheta = v_sph.value
# ----------------------------
# Actions
# ----------------------------
L_vec = angular_momentum(x,v)
Lz = L_vec[:,2]
L = np.linalg.norm(L_vec, axis=1)
# Radial action
Jr = GM / np.sqrt(-2*E) - 0.5*(L + np.sqrt(L*L + 4*GM*b))
# compute the three action variables
actions = np.array([Jr, Lz, L - np.abs(Lz)]).T
# ----------------------------
# Angles
# ----------------------------
c = GM / (-2*E) - b
e = np.sqrt(1 - L*L*(1 + b/c) / GM / c)
# Compute theta_r using eta
tmp1 = r*vr / np.sqrt(-2.*E)
tmp2 = b + c - np.sqrt(b*b + r*r)
eta = np.arctan2(tmp1,tmp2)
thetar = eta - e*c*np.sin(eta) / (c + b) # same as theta3
# Compute theta_z
psi = np.arctan2(np.cos(theta), -np.sin(theta)*r*vtheta/L)
psi[np.abs(vtheta) <= 1e-10] = np.pi/2. # blows up for small vtheta
omega = 0.5 * (1 + L/np.sqrt(L*L + 4*GM*b))
a = np.sqrt((1+e) / (1-e))
ap = np.sqrt((1 + e + 2*b/c) / (1 - e + 2*b/c))
def F(x, y):
z = np.zeros_like(x)
ix = y>np.pi/2.
z[ix] = np.pi/2. - np.arctan(np.tan(np.pi/2.-0.5*y[ix])/x[ix])
ix = y<-np.pi/2.
z[ix] = -np.pi/2. + np.arctan(np.tan(np.pi/2.+0.5*y[ix])/x[ix])
ix = (y<=np.pi/2) & (y>=-np.pi/2)
z[ix] = np.arctan(x[ix]*np.tan(0.5*y[ix]))
return z
A = omega*thetar - F(a,eta) - F(ap,eta)/np.sqrt(1 + 4*GM*b/L/L)
thetaz = psi + A
LR = Lz/L
sinu = (LR/np.sqrt(1.-LR*LR)/np.tan(theta))
sinu = sinu.value
uu = np.arcsin(sinu)
uu[sinu > 1.] = np.pi/2.
uu[sinu < -1.] = -np.pi/2.
uu[vtheta > 0.] = np.pi - uu[vtheta > 0.]
thetap = phi - uu + np.sign(Lz)*thetaz
angles = np.array([thetar, thetap, thetaz]).T
angles %= (2*np.pi)
return actions, angles
