def integ(z, TheTa1, TheTa2):
TheTa = np.sqrt(TheTa1**2 + TheTa2**2)
R = D_d * TheTa
NFW_p = NFWPotential(amp=nfw_M0, a=nfw_a, normalize=False)
Densidad = NFW_p.dens(R, z)
Kappa = 2 * Densidad
return Kappa / (SIGMA_CRIT**2)
def POTDEF1(z, TheTa2):
TheTa = np.sqrt(TheTa2**2 + theta1[l]**2)
R = D_d * TheTa
NFW_p = NFWPotential(amp=M_0, a=r_s, normalize=False)
Sigma = NFW_p.dens(R, z)
kappa = Sigma / SIGMA_CRIT
return (4 / theta2[l]) * TheTa2 * kappa / SIGMA_CRIT**2
def minimize_function(x, rho_nfw_target, rho_midplane_target,
density_conversion):
"""
Computes the chi^2 penalty for a galactic potential for the purpose of finding an NFWPotential and MiyamotoNagaiPotential
circular velocity normalization that yeild the desired midplane and nfw physical densities
:param x: numpy array of proposed disk and nfw circular velocity normalizations (in that order)
:param rho_nfw_target: desired nfw_normalization in physical M_sun / pc^2
:param rho_midplane_target: desired midplane density in physical M_sun / pc^2
:param density_conversion: a conversion factor between galpy internal density units and physical M_sun / pc^2
:return: chi^2 penalty
"""
galactic_potential = [
PowerSphericalPotentialwCutoff(normalize=0.05, alpha=1.8, rc=1.9 / 8.),
MiyamotoNagaiPotential(a=3. / 8., b=0.28 / 8., normalize=x[0]),
NFWPotential(a=2., normalize=x[1])
]
nfw_potential = NFWPotential(a=2., normalize=x[1])
rho = evaluateDensities(galactic_potential, R=1.,
z=0.) * density_conversion
rho_nfw = evaluateDensities(nfw_potential, R=1, z=0.) * density_conversion
dx = (rho - rho_midplane_target)**2 / 0.000001**2 + (
rho_nfw - rho_nfw_target)**2 / 0.000001**2
return dx**0.5
def test_leapfrog_nfw_galpy(self): # Compare a more complicated set of initial conditions to the galpy # outputs (which are slow but known to work). # Set the parameters for a slightly offset circular orbit rho_0 = 0.1 r_scale = 1 G = 0.8962419740798497 dt = 0.001 num_dt = int(5e4) pos_nfw_array = np.tile(np.zeros(3,dtype=np.float64),(num_dt,1)) save_pos = np.zeros((num_dt+1,3)) save_vel = np.zeros((num_dt+1,3)) # Change the rho_0 and r_scale to a fixed array in time rho_0_array = rho_0*np.ones(num_dt,dtype=np.float64) r_scale_array = r_scale*np.ones(num_dt,dtype=np.float64) r_init = 20 pos_init = np.array([r_init,0,0],dtype=np.float64) M_r = 4*np.pi*rho_0*r_scale**3*(np.log((r_scale+r_init)/r_scale)+ r_scale/(r_scale+r_init) - 1) v_r = np.sqrt(G*M_r/r_init) v_kick = 1e-1 vel_init = np.array([v_kick,v_r,0],dtype=np.float64) integrator.leapfrog_int_nfw(pos_init,vel_init,rho_0_array, r_scale_array,pos_nfw_array,dt,save_pos,save_vel) # Compare to galpy orbits o=Orbit([r_init,v_kick,v_r,0,0,0]) nfw = NFWPotential(a=r_scale,amp=rho_0*G*4*np.pi*r_scale**3) ts = np.linspace(0,dt*num_dt,num_dt+1) o.integrate(ts,nfw,method='leapfrog',dt=dt) np.testing.assert_almost_equal(o.x(ts),save_pos[:,0]) np.testing.assert_almost_equal(o.y(ts),save_pos[:,1]) np.testing.assert_almost_equal(o.z(ts),save_pos[:,2]) # Make the kick bigger and see what happens v_kick = 5e-1 pos_init = np.array([r_init,0,0],dtype=np.float64) vel_init = np.array([v_kick,v_r,v_kick],dtype=np.float64) integrator.leapfrog_int_nfw(pos_init,vel_init,rho_0_array, r_scale_array,pos_nfw_array,dt,save_pos,save_vel) # Do the same for galpy o=Orbit([r_init,v_kick,v_r,0,v_kick,0]) nfw = NFWPotential(a=r_scale,amp=rho_0*G*4*np.pi*r_scale**3) ts = np.linspace(0,dt*num_dt,num_dt+1) o.integrate(ts,nfw,method='leapfrog',dt=dt) np.testing.assert_almost_equal(o.x(ts),save_pos[:,0]) np.testing.assert_almost_equal(o.y(ts),save_pos[:,1]) np.testing.assert_almost_equal(o.z(ts),save_pos[:,2])
def get_halo_potentials(masses,
mass_definition,
concentration_at_8=18.0,
physical_off=False):
halo_potentials = []
for m in masses:
if mass_definition == 'HERNQUIST':
c = sample_concentration_herquist(m, concentration_at_8)
pot = HernquistPotential(amp=0.5 * m * apu.solMass, a=c * apu.kpc)
elif mass_definition == 'NFW':
c = sample_concentration_nfw(m, concentration_at_8)
pot = NFWPotential(mvir=m / 10**12, conc=c)
else:
raise Exception('mass definition must be HERNQUIST or NFW')
if physical_off:
galpy.potential.turn_physical_off(pot)
halo_potentials.append(pot)
return halo_potentials
def setup(self):
nfw_norm = [0.2, 0.4, 0.6]
disk_norm = [0.15, 0.3, 0.45]
pot_list = []
for normnfw in nfw_norm:
for normdisk in disk_norm:
pot = [
PowerSphericalPotentialwCutoff(normalize=0.05,
alpha=1.8,
rc=1.9 / 8.),
MiyamotoNagaiPotential(a=3. / 8.,
b=0.28 / 8.,
normalize=normdisk),
NFWPotential(a=2., normalize=normnfw)
]
pot_extension = PotentialExtension(pot, 2, 120, 10)
pot_list.append(pot_extension)
self.disk_norm = disk_norm
self.nfw_norm = nfw_norm
self.pot_list = pot_list
self.tabulated_potential = TabulatedPotential(self.pot_list,
self.nfw_norm,
self.disk_norm)
def setup(self):
nfw_norm = [0.2, 0.4, 0.6]
disk_norm = [0.15, 0.3, 0.45]
scale_height = [0.25, 0.27, 0.29]
pot_list = []
for normnfw in nfw_norm:
for normdisk in disk_norm:
for scale_h in scale_height:
pot = [
PowerSphericalPotentialwCutoff(normalize=0.05,
alpha=1.8,
rc=1.9 / 8.),
MiyamotoNagaiPotential(a=3. / 8.,
b=scale_h / 8.,
normalize=normdisk),
NFWPotential(a=2., normalize=normnfw)
]
pot_extension = PotentialExtension(
pot, 2, 120, 3, compute_action_angle=True)
pot_list.append(pot_extension)
self.disk_norm = disk_norm
self.nfw_norm = nfw_norm
self.scale_height = scale_height
self.pot_list = pot_list
self.tabulated_potential = TabulatedPotential3D(
self.pot_list, self.nfw_norm, self.disk_norm, self.scale_height)
def render_subhalos(mlow,
mhigh,
f_sub,
log_slope,
m_host,
via_lactea_kde,
c8,
galactic_potential,
global_potential_extension,
time_Gyr,
mdef='HERNQUIST',
dr_max=8,
pass_through_disk_limit=3):
N_halos = normalization(f_sub, log_slope, m_host, mlow, mhigh)
sample_orbits_0 = generate_sample_orbits_kde(N_halos, via_lactea_kde,
galactic_potential, time_Gyr)
# print('generated ' + str(N_halos) + ' halos... ')
rs_host, r_core = 25, 25.
args, func = (rs_host, r_core), core_nfw_pdf
inds_keep = filter_orbits_NFW(sample_orbits_0, time_Gyr, func, args)
sample_orbits_1 = [sample_orbits_0[idx] for idx in inds_keep]
nearby_orbits_1_inds, _ = passed_near_solar_neighorhood(
sample_orbits_1,
time_Gyr,
global_potential_extension,
R_solar=8,
dr_max=dr_max,
pass_through_disk_limit=pass_through_disk_limit,
tdep=True)
nearby_orbits_1 = [sample_orbits_0[idx] for idx in nearby_orbits_1_inds]
n_nearby_1 = len(nearby_orbits_1)
halo_masses_1 = sample_mass_function(n_nearby_1, log_slope, mlow, mhigh)
halo_potentials = []
if mdef == 'HERNQUIST':
concentrations = sample_concentration_herquist(halo_masses_1, c8)
for m, c in zip(halo_masses_1, concentrations):
pot = HernquistPotential(amp=0.5 * m * apu.solMass, a=c * apu.kpc)
halo_potentials.append(pot)
elif mdef == 'NFW':
concentrations = sample_concentration_nfw(halo_masses_1, c8)
for m, c in zip(halo_masses_1, concentrations):
pot = NFWPotential(mvir=m / 10**12, conc=c)
halo_potentials.append(pot)
subhalo_orbit_list = []
for orb in nearby_orbits_1:
orb.turn_physical_off()
subhalo_orbit_list.append(orb)
return subhalo_orbit_list, halo_potentials
def galpy_nfw_orbit():
# Setting up the potential
nfw = NFWPotential(conc=C, mvir=M_200, H=70.0, wrtcrit=True, overdens=200)
nfw.turn_physical_on()
vxvv = [
8.0 * units.kpc,
0.0 * units.km / units.s,
240.0 * units.km / units.s,
0.0 * units.pc,
5.0 * units.km / units.s,
]
# Calculating the orbit
ts = np.linspace(0.0, 0.58, 1000) * units.Gyr
o = Orbit(vxvv=vxvv)
o.integrate(ts, nfw, method="odeint")
return o
def obj(x,data,bp,dp,hp):
#x=[fd,fh,vc,rs,hdisk]
if x[0] > 1. or x[0] < 0.: return numpy.finfo(numpy.dtype(numpy.float64)).max
if x[1] > 1. or x[1] < 0.: return numpy.finfo(numpy.dtype(numpy.float64)).max
if (1.-x[0]-x[1]) > 1. or (1.-x[0]-x[1]) < 0.: return numpy.finfo(numpy.dtype(numpy.float64)).max
if x[2] < 0. or x[3] < 0. or x[4] < 0.: return numpy.finfo(numpy.dtype(numpy.float64)).max
if x[2] > 2. or x[3] > 10. or x[4] > 2.: return numpy.finfo(numpy.dtype(numpy.float64)).max
if False:
#Renormalize potentials, intially normalized to 1./3. each
bp= copy.deepcopy(bp)
dp= copy.deepcopy(dp)
hp= copy.deepcopy(hp)
bp._amp*= (1.-x[0]-x[1])**2.*9.
dp._amp*= x[0]**2.*9.
hp._amp*= x[1]**2.*9.
#Re-define disk scale length and halo scale length
if _dexp:
dp._hr= x[4]
else:
dp._a= x[4]
hp._a= x[3]
else:
#Set-up
bp= HernquistPotential(a=0.6/_REFR0,normalize=1.-x[0]-x[1])
if _dexp:
dp= DoubleExponentialDiskPotential(normalize=x[0],
hr=x[4],
hz=0.3/_REFR0)
else:
dp= MiyamotoNagaiPotential(normalize=x[0],
a=x[4],
b=0.3/_REFR0)
hp= NFWPotential(normalize=x[1],
a=x[3])
#Re-normalize data
vcircdata= copy.copy(data)
vcircdata[:,1]/= x[2]
vcircdata[:,2]/= x[2]
#Vcirc chi2
vcmodel= numpy.zeros(vcircdata.shape[0])
for ii in range(vcircdata.shape[0]):
vcmodel[ii]= numpy.sqrt(vcircdata[ii,0]\
*numpy.fabs(evaluateRforces(vcircdata[ii,0],
0.,[bp,dp,hp])))
#print vcircdata[:,0], vcmodel
chi2= numpy.sum((vcircdata[:,1]-vcmodel)**2./vcircdata[:,2]**2.)
#Add scale length measurement
chi2+= (x[4]-_diskscale)**2./_diskscaleerr**2.
#Add dark matter density at the Solar radius
#print hp.dens(1.,0.),_rhodm*x[2]**2.
#chi2+= (hp.dens(1.,0.)-_rhodm*x[2]**2.)**2./_rhodmerr**2./x[2]**4.
return chi2
def sample_galactic_potential(sigma_norm):
sigma_scale = 1.
mwpot = [
PowerSphericalPotentialwCutoff(normalize=sigma_norm * 0.05,
alpha=1.8,
rc=sigma_scale * 1.9 / 8.),
MiyamotoNagaiPotential(a=sigma_scale * 3. / 8.,
b=sigma_scale * 0.28 / 8.,
normalize=sigma_norm * 0.6),
NFWPotential(a=sigma_scale * 2., normalize=sigma_norm * 0.35)
]
return mwpot
def test_overview():
from galpy.potential import NFWPotential
np= NFWPotential(normalize=1.)
from galpy.orbit import Orbit
o= Orbit(vxvv=[1.,0.1,1.1,0.1,0.02,0.])
from galpy.actionAngle import actionAngleSpherical
aA= actionAngleSpherical(pot=np)
js= aA(o)
assert numpy.fabs((js[0]-0.00980542)/js[0]) < 10.**-3., 'Action calculation in the overview section has changed'
assert numpy.fabs((js[1]-1.1)/js[0]) < 10.**-3., 'Action calculation in the overview section has changed'
assert numpy.fabs((js[2]-0.00553155)/js[0]) < 10.**-3., 'Action calculation in the overview section has changed'
from galpy.df import quasiisothermaldf
qdf= quasiisothermaldf(1./3.,0.2,0.1,1.,1.,
pot=np,aA=aA)
assert numpy.fabs((qdf(o)-61.57476085)/61.57476085) < 10.**-3., 'qdf calculation in the overview section has changed'
return None
def MWPotential(Ms=0.76, rs=24.8, c=1., T=True):
'''
Milky Way potential from Marchetti 2017b -- see galpy for the definitions of the potential components
Parameters
----------
Ms : float
NFW profile scale mass in units of e12 Msun
rs : float
Radial profile in units of kpc
c : float
Axis ratio
T : bool
If True, use triaxialNFWPotential
'''
# NFW profile
Ms = Ms * 1e12 * u.Msun
rs = rs * u.kpc
#Disk
Md = 1e11 * u.Msun
ad = 6.5 * u.kpc
bd = 260. * u.pc
#Bulge
Mb = 3.4 * 1e10 * u.Msun
Rb = 0.7 * u.kpc
#BH mass in 1e6 Msun
Mbh = 4e6 * u.Msun
if (T):
halop = TriaxialNFWPotential(amp=Ms, a=rs, c=c, normalize=False)
else:
halop = NFWPotential(amp=Ms, a=rs, normalize=False)
diskp = MiyamotoNagaiPotential(amp=Md, a=ad, b=bd, normalize=False)
bulgep = HernquistPotential(
amp=2 * Mb, a=Rb,
normalize=False) #Factor 2 because of the galpy definition
bh = KeplerPotential(amp=Mbh, normalize=False)
return [halop, diskp, bulgep, bh]
def test_minimize_func(self):
disk_norm, nfw_norm = solve_normalizations(self.rho_nfw_target,
self.rho_midplane_target,
self.density_conversion)
pot = [
PowerSphericalPotentialwCutoff(normalize=0.05,
alpha=1.8,
rc=1.9 / 8.),
MiyamotoNagaiPotential(a=3. / 8., b=0.28 / 8.,
normalize=disk_norm),
NFWPotential(a=2., normalize=nfw_norm)
]
rho = evaluateDensities(pot, R=1., z=0) * self.density_conversion
npt.assert_almost_equal(rho, self.rho_midplane_target, 3)
disk_norm, nfw_norm = solve_normalizations(self.rho_nfw_target,
self.rho_nfw_target * 0.5,
self.density_conversion)
npt.assert_equal(True, disk_norm < 0)
def calc_potentials_vs_time(self, differential=False, method='astropy'):
"""Calculates the gravitational potentials of each component for all redshift steps using galpy
The gas and stars are represented by a double exponential profile by default.
The DM is represented by a NFW profile.
Can construct potential in both astropy units (method=='astropy') or galpy's 'natural' units (method=='natural') for the purposes of interpolation.
Stars can be calculated using a differential mass profile with varying scale radii, but in this case it is best to interpolate the potentials first
"""
if method == 'astropy':
print(
"Calculating galactic potentials at each redshift using astropy units...\n"
)
elif method == 'natural':
print(
"Calculating galactic potentials at each redshift using galpy's natural units...\n"
)
else:
raise NameError(
'Method {0:s} for constructing the potential not recognized!'.
format(method))
# lists for saving potentials at each step
stars_potentials = []
gas_potentials = []
dm_potentials = []
full_potentials = []
if self.disk_profile not in [
'RazorThinExponential', 'DoubleExponential'
]:
raise NameError('Disk profile {0:s} not recognized!'.format(
self.disk_profile))
# iterate over each time-step until when the sgrb occurred
for ii, zz in enumerate(self.redz):
# --- get the gas and DM masses at this step
if method == 'astropy':
mgas = self.mass_gas[ii]
mdm = self.mass_dm[ii]
elif method == 'natural':
mgas = utils.Mphys_to_nat(self.mass_gas[ii])
mdm = utils.Mphys_to_nat(self.mass_dm[ii])
# --- if differential stellar potential not being used, just take the total stellar mass at each timestep
# --- this is also done for the first differential timestep
if (differential == False) or (ii == 0):
if method == 'astropy':
mstar = self.mass_stars[ii]
elif method == 'natural':
mstar = utils.Mphys_to_nat(self.mass_stars[ii])
else:
if method == 'astropy':
mstar = self.mass_stars[ii] - self.mass_stars[ii - 1]
elif method == 'natural':
mstar = utils.Mphys_to_nat(self.mass_stars[ii] -
self.mass_stars[ii - 1])
# --- get the scale lengths for the baryons and halo at this redshift step
if method == 'astropy':
rs_baryons = self.Rscale_baryons[ii]
rs_dm = self.Rscale_dm[ii]
# if galaxy hasn't formed yet, give the potentials neglible scale sizes to avoid dividing by 0
if rs_baryons == 0:
rs_baryons = 1e-10 * rs_baryons.unit
if rs_dm == 0:
rs_dm = 1e-10 * rs_dm.unit
elif method == 'natural':
rs_baryons = utils.Rphys_to_nat(self.Rscale_baryons[ii])
rs_dm = utils.Rphys_to_nat(self.Rscale_dm[ii])
# if galaxy hasn't formed yet, give the potentials neglible scale sizes to avoid dividing by 0
if rs_baryons == 0:
rs_baryons = 1e-10
if rs_dm == 0:
rs_dm = 1e-10
# --- construct the stellar and gas potentials
if self.disk_profile == 'RazorThinExponential':
# for a razor-thin disk, the amplitude is mdisk / (2 * pi * rs**2)
amp_stars = mstar / (4 * np.pi * rs_baryons**2)
amp_gas = mgas / (4 * np.pi * rs_baryons**2)
# construct the potentials at this timestep
stars_potential = RazorThinExponentialDiskPotential(
amp=amp_stars, hr=rs_baryons)
gas_potential = RazorThinExponentialDiskPotential(
amp=amp_gas, hr=rs_baryons)
elif self.disk_profile == 'DoubleExponential':
# for a double exponential disk, the amplitude is mdisk / (2 * pi * rs**2 * 2*rz)
amp_stars = mstar / (4 * np.pi * rs_baryons**2 *
(self.z_scale * rs_baryons))
amp_gas = mgas / (4 * np.pi * rs_baryons**2 *
(self.z_scale * rs_baryons))
stars_potential = DoubleExponentialDiskPotential(
amp=amp_stars, hr=rs_baryons, hz=self.z_scale * rs_baryons)
gas_potential = DoubleExponentialDiskPotential(
amp=amp_gas, hr=rs_baryons, hz=self.z_scale * rs_baryons)
# --- construct the DM potentials
amp_dm = mdm
dm_potential = NFWPotential(amp=amp_dm, a=rs_dm)
# --- add the potentials to the lists for each step
stars_potentials.append(stars_potential)
gas_potentials.append(gas_potential)
dm_potentials.append(dm_potential)
# --- if differential is specified, we use *all* the stellar profiles up to this point
if differential == True:
combined_potentials = stars_potentials[:]
combined_potentials.extend([gas_potential, dm_potential])
else:
combined_potentials = [
stars_potential, gas_potential, dm_potential
]
full_potentials.append(combined_potentials)
if method == 'astropy':
self.stars_potentials = stars_potentials
self.gas_potentials = gas_potentials
self.dm_potentials = dm_potentials
self.full_potentials = full_potentials
if method == 'natural':
self.stars_potentials_natural = stars_potentials
self.gas_potentials_natural = gas_potentials
self.dm_potentials_natural = dm_potentials
self.full_potentials_natural = full_potentials
return
def __init__(self,amp=1.,a=2.,b=1.,c=1.,zvec=None,pa=None,
normalize=False,
conc=None,mvir=None,
glorder=50,vo=None,ro=None,
H=70.,Om=0.3,overdens=200.,wrtcrit=False):
"""
NAME:
__init__
PURPOSE:
Initialize a triaxial NFW potential
INPUT:
amp - amplitude to be applied to the potential (default: 1); can be a Quantity with units of mass or Gxmass
a - scale radius (can be Quantity)
b - y-to-x axis ratio of the density
c - z-to-x axis ratio of the density
zvec= (None) If set, a unit vector that corresponds to the z axis
pa= (None) If set, the position angle of the x axis
glorder= (50) if set, compute the relevant force and potential integrals with Gaussian quadrature of this order
normalize - if True, normalize such that vc(1.,0.)=1., or, if given as a number, such that the force is this fraction of the force necessary to make vc(1.,0.)=1.
Alternatively, NFW potentials can be initialized using
conc= concentration
mvir= virial mass in 10^12 Msolar
in which case you also need to supply the following keywords
H= (default: 70) Hubble constant in km/s/Mpc
Om= (default: 0.3) Omega matter
overdens= (200) overdensity which defines the virial radius
wrtcrit= (False) if True, the overdensity is wrt the critical density rather than the mean matter density
ro=, vo= distance and velocity scales for translation into internal units (default from configuration file)
OUTPUT:
(none)
HISTORY:
2016-05-30 - Written - Bovy (UofT)
"""
Potential.__init__(self,amp=amp,ro=ro,vo=vo,amp_units='mass')
if _APY_LOADED and isinstance(a,units.Quantity):
a= a.to(units.kpc).value/self._ro
self.alpha= 1
self.beta= 3
self._b= b
self._b2= self._b**2.
self._c= c
self._c2= self._c**2.
self._setup_gl(glorder)
self._setup_zvec_pa(zvec,pa)
self._force_hash= None
if conc is None:
self.a= a
if normalize or \
(isinstance(normalize,(int,float)) \
and not isinstance(normalize,bool)):
self.normalize(normalize)
else:
from galpy.potential import NFWPotential
dum= NFWPotential(mvir=mvir,conc=conc,ro=self._ro,vo=self._vo,
H=H,Om=Om,wrtcrit=wrtcrit,overdens=overdens)
self.a= dum.a
self._amp= dum._amp
self._scale= self.a
self.hasC= not self._glorder is None
self.hasC_dxdv= False
if not self._aligned or numpy.fabs(self._b-1.) > 10.**-10.:
self.isNonAxi= True
return None
class MessyStreamData:
nfwp = NFWPotential(normalize=1., a=14. / 10.)
G = 4.302e-3 #pc*M_sun^-1*(km/s)^2
def __init__(self,
mass,
rstream,
rsub,
impact,
subvel,
t,
nstars,
sigmav,
std,
psi0=[-20. / 180. * np.pi, 20 / 180. * np.pi]):
self.nump = nstars
self.psi0 = np.sort(np.random.uniform(psi0[0], psi0[1], size=nstars))
self.vysigma = sigmav
self.dxstd = std[0]
self.dxdotstd = std[1]
self.m = mass
self.r0 = rstream
self.rs = rsub
self.b = impact
self.wvec = subvel
self.t = t * 1.0227 # Myr*1.02(pc/(km/s)/Myr) - converts to time in units pc/(km/s)
self.vy = self.nfwp.vcirc(self.r0 / 10000.) * 168.2
self.noisyvy = self.vy + np.random.normal(scale=self.vysigma,
size=nstars)
self.wperp = np.sqrt(subvel[0]**2 + subvel[2]**2)
self.wpar = self.noisyvy - subvel[1]
self.w = np.sqrt(self.wperp**2 + self.wpar**2)
self.gamma = self.calc_gamma()
'''
print('vy:',np.mean(self.noisyvy),np.std(self.noisyvy))
print('gamma:',np.mean(self.gamma),np.std(self.gamma))
print('wpar:',np.mean(self.wpar),np.std(self.wpar))
print('wperp:',np.mean(self.wperp),np.std(self.wperp))
print('w:',np.mean(self.w),np.std(self.w))
print('tau:',np.mean(self.w*self.r0**2/2./self.G/self.m), np.std(self.w*self.r0**2/2./self.G/self.m))
'''
self.dv = self.calc_dv(self.psi0)
self.psi, self.rho = self.calc_psi_rho(self.psi0)
plt.figure()
plt.plot(self.psi0, self.psi, '.')
plt.xlabel(r'$\psi_o$ (rad)')
plt.ylabel(r'$\psi$ (rad)')
plt.show()
self.dx = self.calc_dx(self.psi0, self.dxstd)
self.dxdot = self.calc_dxdot(self.psi0, self.dxdotstd)
sorteddata = np.array([
(self.psi[i], self.dx[0, i], self.dx[1, i], self.dx[2, i],
self.dxdot[0, i], self.dxdot[1, i], self.dxdot[2, i])
for i in np.argsort(self.psi)
])
self.data = [
sorteddata[:, 0],
np.transpose(sorteddata[:, 1:4]),
np.transpose(sorteddata[:, 4:7])
]
def gT(self):
angle = self.gamma * self.noisyvy * self.t / self.r0
return angle
def calc_gamma(self):
g = 3. + (self.r0 / 10000.)**2 * 168.2**2 * self.nfwp.R2deriv(
self.r0 / 10000., 0.) / self.noisyvy**2.
return np.sqrt(g)
def calc_dv(self, psi0):
M = self.m
r0 = self.r0
wperp = self.wperp
wpar = self.wpar
w = self.w
wvec = self.wvec
b = self.b
rs = self.rs
G = self.G
nump = len(psi0)
deltav = np.zeros([3, nump])
deltav[0] = 2. * G * M / r0 / wperp**2 / w * (
b * w**2 * wvec[2] / wperp -
psi0 * wpar * wvec[0]) / (psi0**2 +
(b**2 + rs**2 / r0**2) * w**2 / wperp**2)
deltav[1] = -2. * G * M * psi0 / r0 / w / (
psi0**2 + (b**2 + rs**2 / r0**2) * w**2 / wperp**2)
deltav[2] = -2. * G * M / r0 / wperp**2 / w * (
b * w**2 * wvec[0] / wperp +
psi0 * wpar * wvec[2]) / (psi0**2 +
(b**2 + rs**2 / r0**2) * w**2 / wperp**2)
return deltav
def calc_psi_rho(self, psi0):
gam = self.calc_gamma()
gT = self.gT()
t = self.t
r0 = self.r0
vy = self.noisyvy
wperp = self.wperp
wpar = self.wpar
wvec = self.wvec
w = self.w
b = self.b
rs = self.rs
tau = w * r0**2 / 2. / self.G / self.m
f = (4. - gam**2) / gam**2 * t / tau - 4. * np.sin(
gT) / gam**3 * r0 / vy / tau + 2. * (1. - np.cos(
gT)) / gam**2 * wpar * wvec[0] / wperp**2 * r0 / vy / tau
g = 2. * (1 - np.cos(gT)) * b * w**2 * wvec[2] * r0 / (
gam**2 * wperp**3 * vy * tau)
B2 = (b**2 + rs**2 / r0**2) * w**2 / wperp**2
psi = psi0 + (f * psi0 - g) / (psi0**2 + B2) - (self.vy -
self.noisyvy) * t / r0
rho = (1. + (f * B2 - f * psi0**2 + 2. * g * psi0) /
(psi0**2 + B2)**2)**(-1)
return psi, rho
def calc_dx(self, psi0, std):
r0 = self.r0
vy = self.noisyvy
dv = self.calc_dv(psi0)
gam = self.calc_gamma()
gT = self.gT()
nump = len(psi0)
noisedx = np.zeros([3, nump])
noisedx[0] = np.random.normal(scale=std[0], size=nump)
noisedx[2] = np.random.normal(scale=std[2], size=nump)
deltax = np.zeros([3, nump])
deltax[0] = 2. * r0 * dv[1] / vy * (
1. - np.cos(gT)
) / gam**2 + r0 * dv[0] / vy * np.sin(gT) / self.gamma + noisedx[0]
deltax[2] = dv[2] / vy * np.sin(vy * self.t / r0) + noisedx[2]
return deltax
def calc_dxdot(self, psi0, std):
vy = self.noisyvy
dv = self.calc_dv(psi0)
gamma = self.calc_gamma()
gT = self.gT()
nump = len(psi0)
noisedxdot = np.zeros([3, nump])
for i in range(3):
noisedxdot[i] = np.random.normal(scale=std[i], size=nump)
deltaxdot = np.zeros([3, nump])
deltaxdot[0] = 2. * dv[1] / gamma * np.sin(gT) + dv[0] * np.cos(
gT) + noisedxdot[0]
deltaxdot[1] = -dv[1] * (
2. - gamma**2) / gamma**2 + 2. * dv[1] * np.cos(
gT) / gamma**2 - dv[0] * np.sin(gT) / gamma + noisedxdot[1]
deltaxdot[2] = dv[2] * np.cos(vy * self.t / self.r0) + noisedxdot[2]
return deltaxdot
return gas_dens(R,z)+stellar_dens(R,z)+bulge_dens(R,z)+NFW_dens(R,z)
def sech(x):
return 1./np.cosh(x)
#dicts used in DiskSCFPotential
sigmadict = [{'type':'exp','h':Rd_HI,'amp':Sigma0_HI, 'Rhole':Rm_HI},
{'type':'exp','h':Rd_H2,'amp':Sigma0_H2, 'Rhole':Rm_H2},
{'type':'exp','h':Rd_thin,'amp':Sigma0_thin, 'Rhole':0.},
{'type':'exp','h':Rd_thick,'amp':Sigma0_thick, 'Rhole':0.}]
hzdict = [{'type':'sech2', 'h':zd_HI},
{'type':'sech2', 'h':zd_H2},
{'type':'exp', 'h':0.3/ro},
{'type':'exp', 'h':0.9/ro}]
#generate separate disk and halo potential - and combined potential
McMillan_bulge=\
mySCFPotential(Acos=scf_compute_coeffs_axi(bulge_dens,20,10,a=0.1)[0],
a=0.1,ro=ro,vo=vo)
McMillan_disk = myDiskSCFPotential(dens=lambda R,z: gas_stellar_dens(R,z),
Sigma=sigmadict, hz=hzdict,
a=2.5, N=30, L=30,ro=ro,vo=vo)
McMillan_halo = NFWPotential(amp = rho0_halo*(4*np.pi*rh**3),
a = rh,ro=ro,vo=vo)
McMillan2017 = [McMillan_disk,McMillan_halo,McMillan_bulge]
def test_calc_neg_grad_nfw(self): # Test that calculating the nfw negative gradient returns the same # results as galpy (but faster hopefully :)) r_scale = 2 rho_0 = 0.1 G = 0.8962419740798497 nfw = NFWPotential(a=r_scale,amp=rho_0*G*4*np.pi*r_scale**3) pos_nfw = np.zeros(3,dtype=np.float64) pos = np.zeros(3,dtype=np.float64) # Start just by testing a couple of different values of radius 1,2, and # 3. thetas = np.random.rand(10).astype(np.float64)*2*np.pi rs = np.array([1,2,3],dtype=np.float64) for theta in thetas: for r in rs: # Update the position pos[0] = r*np.cos(theta) pos[1] = r*np.sin(theta) # Compare both magnitudes r_force = nfw.Rforce(r,0) neg_grad = integrator.calc_neg_grad_nfw(rho_0,r_scale,pos_nfw, pos) self.assertAlmostEqual(np.abs(r_force), np.sqrt(np.sum(np.square(neg_grad))),places=5) # Ensure the direction is correct np.testing.assert_almost_equal( neg_grad/np.sqrt(np.sum(np.square(neg_grad))), -pos/r) # Check that the force softening scale is being used. thetas = np.random.rand(10).astype(np.float64)*2*np.pi rs = np.array([1,2,3],dtype=np.float64) force_softening = 1 for theta in thetas: for r in rs: # Update the position pos[0] = r*np.cos(theta) pos[1] = r*np.sin(theta) # Compare both magnitudes r_force = nfw.Rforce(r,0) neg_grad = integrator.calc_neg_grad_nfw(rho_0,r_scale,pos_nfw, pos,force_softening=force_softening) self.assertGreater(np.abs(r_force), np.sqrt(np.sum(np.square(neg_grad)))) # Ensure the direction is still correct np.testing.assert_almost_equal( neg_grad/np.sqrt(np.sum(np.square(neg_grad))), -pos/r) # Check that the box boundary is being used correctly box_length = 4 pos[0] = 3 pos[1] = 1 r_force = nfw.Rforce(np.sqrt(2),0) neg_grad = integrator.calc_neg_grad_nfw(rho_0,r_scale,pos_nfw,pos, box_length=box_length) # Test that it's calculating the radius correctly self.assertAlmostEqual(np.abs(r_force), np.sqrt(np.sum(np.square(neg_grad))),places=5) # Test that the direction is correct np.testing.assert_almost_equal( neg_grad/np.sqrt(np.sum(np.square(neg_grad))), np.array([1.0,-1.0,0.0])/np.sqrt(2))
def calc_es():
savefilename= 'myes.sav'
if os.path.exists(savefilename):
savefile= open(savefilename,'rb')
mye= pickle.load(savefile)
e= pickle.load(savefile)
savefile.close()
else:
#Read data
dialect= csv.excel
dialect.skipinitialspace=True
reader= csv.reader(open('../data/Dierickx-etal-tab2.txt','r'),delimiter=' ',dialect=dialect)
vxvs= []
es= []
vphis= []
vxs= []
vys= []
vzs= []
ls= []
for row in reader:
thisra= float(row[3])
thisdec= float(row[4])
thisd= float(row[17])/1000.
thispmra= float(row[13])
thispmdec= float(row[15])
thisvlos= float(row[11])
thise= float(row[26])
vxvs.append([thisra,thisdec,thisd,thispmra,thispmdec,thisvlos])
es.append(thise)
vphis.append(float(row[25]))
vxs.append(float(row[19]))
vys.append(float(row[21]))
vzs.append(float(row[23]))
ls.append(float(row[5]))
vxvv= nu.array(vxvs)
e= nu.array(es)
vphi= nu.array(vphis)
vx= nu.array(vxs)
vy= nu.array(vys)
vz= nu.array(vzs)
l= nu.array(ls)
#Define potential
lp= LogarithmicHaloPotential(normalize=1.)
mp= MiyamotoNagaiPotential(a=0.5,b=0.0375,amp=1.,normalize=.6)
np= NFWPotential(a=4.5,normalize=.35)
hp= HernquistPotential(a=0.6/8,normalize=0.05)
ts= nu.linspace(0.,20.,10000)
mye= nu.zeros(len(e))
for ii in range(len(e)):
#Integrate the orbit
o= Orbit(vxvv[ii,:],radec=True,vo=220.,ro=8.)
o.integrate(ts,lp)
mye[ii]= o.e()
#Save
savefile= open(savefilename,'wb')
pickle.dump(mye,savefile)
pickle.dump(e,savefile)
savefile.close()
#plot
plot.print()
plot.plot(nu.array([0.,1.]),nu.array([0.,1.]),'k-',
xlabel=r'$\mathrm{Dierickx\ et\ al.}\ e$',
ylabel=r'$\mathrm{galpy}\ e$')
plot.plot(e,mye,'k,',overplot=True)
plot.end_print('myee.png')
plot.print()
plot.hist(e,bins=30,xlabel=r'$\mathrm{Dierickx\ et\ al.}\ e$')
plot.end_print('ehist.png')
plot.print()
plot.hist(mye,bins=30,xlabel=r'$\mathrm{galpy}\ e$')
plot.end_print('myehist.png')
def fitMass():
#Read data
vcircdata= numpy.loadtxt('vcirc.txt',comments='#',delimiter='|')
vcircdata[:,0]/= _REFR0
vcircdata[:,1]/= _REFV0
vcircdata[:,2]/= _REFV0
#Set-up
bp= HernquistPotential(a=0.6/_REFR0,normalize=1./3.)
if _dexp:
dp= DoubleExponentialDiskPotential(normalize=1./3.,
hr=3.25/_REFR0,
hz=0.3/_REFR0)
else:
dp= MiyamotoNagaiPotential(normalize=1./3.,
a=3.25/_REFR0,
b=0.3/_REFR0)
hp= NFWPotential(normalize=1./3.,
a=5./_REFR0)
init= numpy.array([0.6,0.35,218./_REFV0,15./_REFR0,3.25/_REFR0])
#print _convertrho
out= optimize.fmin_powell(obj,init,args=(vcircdata,bp,dp,hp),
callback=cb)
print out
#Calculate mass
#Halo mass
halo= halomass(out)
bulge= halomass(out)
disk= diskmass(out)
print halo, disk, bulge, totalmass(out)
#Sample
samples= bovy_mcmc.markovpy(out,0.05,
(lambda x: -obj(x,vcircdata,bp,dp,hp)),
(),
isDomainFinite=[[True,True],
[True,True],
[True,True],
[True,True],
[True,True]],
domain=[[0.,1.],
[0.,1.],
[0.,2.],
[0.,10.],
[0.,2.]],
nwalkers=8,
nsamples=10000)
print "Done with sampling ..."
print numpy.mean(numpy.array(samples),axis=0)
print numpy.std(numpy.array(samples),axis=0)
samples= numpy.random.permutation(samples)[0:500]
#halo
totalmasssamples= []
for s in samples:
totalmasssamples.append(halomass(s))
totalmasssamples= numpy.array(totalmasssamples)
print "halo mass: ", numpy.mean(totalmasssamples), numpy.std(totalmasssamples)
print sixtyeigthinterval(halomass(out),totalmasssamples,quantile=.68)
#total
totalmasssamples= []
for s in samples:
totalmasssamples.append(totalmass(s))
totalmasssamples= numpy.array(totalmasssamples)
print "total mass: ", numpy.mean(totalmasssamples), numpy.std(totalmasssamples)
print sixtyeigthinterval(totalmass(out),totalmasssamples,quantile=.68)
bovy_plot.bovy_print()
bovy_plot.bovy_hist(totalmasssamples,bins=16,range=[0.,2.])
bovy_plot.bovy_end_print('totalmass.png')
return None
#disk
totalmasssamples= []
for s in samples:
totalmasssamples.append(diskmass(s))
totalmasssamples= numpy.array(totalmasssamples)
print "disk mass: ", numpy.mean(totalmasssamples), numpy.std(totalmasssamples)
#bulge
totalmasssamples= []
for s in samples:
totalmasssamples.append(bulgemass(s))
totalmasssamples= numpy.array(totalmasssamples)
print "bulge mass: ", numpy.mean(totalmasssamples), numpy.std(totalmasssamples)
return None
def halomass(out):
hp= NFWPotential(normalize=out[1],
a=out[3])
return integrate.quad((lambda x: hp.dens(x,0.)*x**2.),0.,250./_REFR0)[0]*4.*numpy.pi*_convertmass/out[2]**2.
# 'corr_colnames':None,
}
cart_colnames = {
# 'main_colnames':None,
# 'error_colnames':None,
# 'corr_colnames':None,
}
from galpy.potential import PowerSphericalPotentialwCutoff,\
MiyamotoNagaiPotential, NFWPotential, verticalfreq, MWPotential2014
scale_height_factor = 2.0
bp = PowerSphericalPotentialwCutoff(alpha=1.8, rc=1.9 / 8., normalize=0.05)
mp = MiyamotoNagaiPotential(a=3. / 8., b=scale_height_factor * 0.28 / 8.,
normalize=.6)
np = NFWPotential(a=16 / 8., normalize=.35)
my_mwpotential2014 = [bp, mp, np]
orbit = {
'potential': my_mwpotential2014, # TC: varied params randomly...
}
special = {
'component':'sphere', # parameterisation for the origin
}
advanced = {
'burnin_steps':500, # emcee parameters, number of steps for each burnin iteraton
'sampling_steps':500,
}
def calcj(rotcurve):
if rotcurve == 'flat':
savefilename = 'myjs.sav'
elif rotcurve == 'power':
savefilename = 'myjs_power.sav'
if os.path.exists(savefilename):
savefile = open(savefilename, 'rb')
myjr = pickle.load(savefile)
myjp = pickle.load(savefile)
mywr = pickle.load(savefile)
mywp = pickle.load(savefile)
mye = pickle.load(savefile)
myzmax = pickle.load(savefile)
e = pickle.load(savefile)
zmax = pickle.load(savefile)
savefile.close()
else:
dialect = csv.excel
dialect.skipinitialspace = True
reader = csv.reader(open('../data/gcs.tsv', 'r'),
delimiter='|',
dialect=dialect)
vxvs = []
es = []
zmaxs = []
for row in reader:
if row[0][0] == '#':
continue
thisra = row[0]
thisdec = row[1]
thisd = read_float(row[2]) / 1000.
if thisd > 0.2: continue
thisu = read_float(row[3])
thisv = read_float(row[4])
thisw = read_float(row[5])
thise = read_float(row[6])
thiszmax = read_float(row[7])
if thisd == -9999 or thisu == -9999 or thisv == -9999 or thisw == -9999:
continue
vxvs.append([
hms_to_rad(thisra),
dms_to_rad(thisdec), thisd, thisu, thisv, thisw
])
es.append(thise)
zmaxs.append(thiszmax)
vxvv = nu.array(vxvs)
e = nu.array(es)
zmax = nu.array(zmaxs)
#Define potential
lp = LogarithmicHaloPotential(normalize=1.)
pp = PowerSphericalPotential(normalize=1., alpha=-2.)
mp = MiyamotoNagaiPotential(a=0.5, b=0.0375, amp=1., normalize=.6)
np = NFWPotential(a=4.5, normalize=.35)
hp = HernquistPotential(a=0.6 / 8, normalize=0.05)
ts = nu.linspace(0., 20., 10000)
myjr = nu.zeros(len(e))
myjp = nu.zeros(len(e))
mywr = nu.zeros(len(e))
mywp = nu.zeros(len(e))
mye = nu.zeros(len(e))
myzmax = nu.zeros(len(e))
for ii in range(len(e)):
#Integrate the orbit
o = Orbit(vxvv[ii, :], radec=True, uvw=True, vo=220., ro=8.)
#o.integrate(ts,[mp,np,hp])
if rotcurve == 'flat':
o.integrate(ts, lp)
mye[ii] = o.e()
myzmax[ii] = o.zmax() * 8.
print e[ii], mye[ii], zmax[ii], myzmax[ii]
o = o.toPlanar()
myjr[ii] = o.jr(lp)[0]
else:
o = o.toPlanar()
myjr[ii] = o.jr(pp)[0]
myjp[ii] = o.jp()[0]
mywr[ii] = o.wr()[0]
mywp[ii] = o.wp()
#Save
savefile = open(savefilename, 'wb')
pickle.dump(myjr, savefile)
pickle.dump(myjp, savefile)
pickle.dump(mywr, savefile)
pickle.dump(mywp, savefile)
pickle.dump(mye, savefile)
pickle.dump(myzmax, savefile)
pickle.dump(e, savefile)
pickle.dump(zmax, savefile)
savefile.close()
#plot
if rotcurve == 'flat':
plot.bovy_print()
plot.bovy_plot(nu.array([0., 1.]),
nu.array([0., 1.]),
'k-',
xlabel=r'$\mathrm{Holmberg\ et\ al.}\ e$',
ylabel=r'$\mathrm{galpy}\ e$')
plot.bovy_plot(e, mye, 'k,', overplot=True)
plot.bovy_end_print('myee.png')
plot.bovy_print()
plot.bovy_plot(
nu.array([0., 2.5]),
nu.array([0., 2.5]),
'k-',
xlabel=r'$\mathrm{Holmberg\ et\ al.}\ z_{\mathrm{max}}$',
ylabel=r'$\mathrm{galpy}\ z_{\mathrm{max}}$')
plot.bovy_plot(zmax, myzmax, 'k,', overplot=True)
plot.bovy_end_print('myzmaxzmax.png')
plot.bovy_print()
plot.bovy_plot(myjp,
myjr / 2. / nu.pi,
'k,',
xlabel=r'$J_{\phi}$',
ylabel=r'$J_R/2\pi$',
xrange=[0.7, 1.3],
yrange=[0., 0.05])
if rotcurve == 'flat':
plot.bovy_end_print('jrjp.png')
else:
plot.bovy_end_print('jrjp_power.png')
class Stream:
'''
These are the variables common to all streams.
G - The gravitational constant
nfwp - NFW potential for the host galaxy of the streams and subhalos
nump - number of psi0 data points at which to calculate the different variables
psi0 - the range of initial angles to the different points along the stream.
psi0 = 0 occurs at the point of closest approach of the subhalo along the stream.
'''
nfwp = NFWPotential(normalize=1., a=14. / 10.)
G = 4.302e-3 #pc*M_sun^-1*(km/s)^2
nump = 1000
'''
There are several functions which can be called on the stream.
__init__ - initializes a stream with the parameters passe to it and calculates several important parameters
Input parameters:
mass (Solar masses) - mass of the subhalo
rstream (pc) - radius of the circular orbit of the stream around the galaxy centre
rsub (pc) - Radius of the subhalo
impact - impact parameter of the subhalo
subvel (km/s) - Velocity of the subhalo in coordinates specified by the stream
t (Myr) - Time elapsed since time of closest approach of the subhalo
Calculated parameters:
wperp (km/s) - Perpendicular component of the relative velocity to points in the stream
wpar (km/s) - Parallel component of relative velocity
w (km/s) - Total magnitude of relative velocity between the subha and points in the stream
dv (km/s) - Calcultes the initial kick to velocity as a function of angle. Not time dependent.
psi (radians) - Calculates the new angles to stars after a time, t.
rho - Calculates the density of the stream after a time, t.
dx (dx[0] (pc), dx[2] (radians)) - Change to position of stars in the stream after a time, t.
dxdot (km/s) - Change in velocity of stars in the stream after a time, t.
gT (radians) - calculates a very commonly used angle in several functions
gamma -
'''
def __init__(self,
mass,
rstream,
rsub,
impact,
subvel,
t,
psi0=[-20. / 180. * np.pi, 20 / 180. * np.pi]):
self.psi0 = np.linspace(psi0[0], psi0[1], self.nump)
self.m = mass
self.r0 = rstream
self.rs = rsub
self.b = impact
self.wvec = subvel
self.t = t * 1.02 # Myr*1.02(pc/(km/s)/Myr) - converts to time in units pc/(km/s)
self.vy = self.nfwp.vcirc(self.r0 / 10000.) * 168.2
self.wperp = np.sqrt(subvel[0]**2 + subvel[2]**2)
self.wpar = self.vy - subvel[1]
self.w = np.sqrt(subvel[0]**2 + subvel[2]**2 +
(self.vy - subvel[1])**2)
self.gamma = self.calc_gamma()
self.dv = self.calc_dv(self.psi0)
self.psi, self.rho = self.calc_psi_rho(self.psi0)
self.dx = self.calc_dx(self.psi0)
self.dxdot = self.calc_dxdot(self.psi0)
def gT(self):
angle = self.gamma * self.vy * self.t / self.r0
return angle
def calc_gamma(self):
g = 3. + (self.r0 / 10000.)**2 * 168.2**2 * self.nfwp.R2deriv(
self.r0 / 10000., 0.) / self.vy**2.
return np.sqrt(g)
def calc_dv(self, psi0):
M = self.m
r0 = self.r0
wperp = self.wperp
wpar = self.wpar
w = self.w
wvec = self.wvec
b = self.b
rs = self.rs
G = self.G
nump = len(psi0)
deltav = np.zeros([3, nump])
deltav[0] = 2. * G * M / r0 / wperp**2 / w * (
b * w**2 * wvec[2] / wperp -
psi0 * wpar * wvec[0]) / (psi0**2 +
(b**2 + rs**2 / r0**2) * w**2 / wperp**2)
deltav[1] = -2. * G * M * psi0 / r0 / w / (
psi0**2 + (b**2 + rs**2 / r0**2) * w**2 / wperp**2)
deltav[2] = -2. * G * M / r0 / wperp**2 / w * (
b * w**2 * wvec[0] / wperp +
psi0 * wpar * wvec[2]) / (psi0**2 +
(b**2 + rs**2 / r0**2) * w**2 / wperp**2)
return deltav
def calc_psi_rho(self, psi0):
gam = self.calc_gamma()
gT = self.gT()
t = self.t
r0 = self.r0
vy = self.vy
wperp = self.wperp
wpar = self.wpar
wvec = self.wvec
w = self.w
b = self.b
rs = self.rs
tau = w * r0**2 / 2. / self.G / self.m
f = (4. - gam**2) / gam**2 * t / tau - 4. * np.sin(
gT) / gam**3 * r0 / vy / tau + 2. * (1. - np.cos(
gT)) / gam**2 * wpar * wvec[0] / wperp**2 * r0 / vy / tau
g = 2. * (1 - np.cos(gT)) * b * w**2 * wvec[2] * r0 / (
gam**2 * wperp**3 * vy * tau)
B2 = (b**2 + rs**2 / r0**2) * w**2 / wperp**2
psi = psi0 + (f * psi0 - g) / (psi0**2 + B2)
rho = (1. + (f * B2 - f * psi0**2 + 2. * g * psi0) /
(psi0**2 + B2)**2)**(-1)
return psi, rho
def calc_dx(self, psi0):
r0 = self.r0
vy = self.vy
dv = self.calc_dv(psi0)
gam = self.calc_gamma()
gT = self.gT()
nump = len(psi0)
deltax = np.zeros([3, nump])
deltax[0] = 2. * r0 * dv[1] / vy * (1. - np.cos(
gT)) / gam**2 + r0 * dv[0] / vy * np.sin(gT) / self.gamma
deltax[2] = dv[2] / vy * np.sin(vy * self.t / r0)
return deltax
def calc_dxdot(self, psi0):
vy = self.vy
dv = self.calc_dv(psi0)
gamma = self.calc_gamma()
gT = self.gT()
nump = len(psi0)
deltaxdot = np.zeros([3, nump])
deltaxdot[0] = 2. * dv[1] / gamma * np.sin(gT) + dv[0] * np.cos(gT)
deltaxdot[1] = -dv[1] * (2. - gamma**2) / gamma**2 + 2. * dv[
1] * np.cos(gT) / gamma**2 - dv[0] * np.sin(gT) / gamma
deltaxdot[2] = dv[2] * np.cos(vy * self.t / self.r0)
return deltaxdot