def add_rotation(cluster, qrot):
"""Add a degree of rotation to an already generated StarCluster
Parameters
----------
cluster : class
StarCluster
qrot : float
fraction of stars with v_phi < 0 that are switched to vphi > 0
Returns
-------
x,y,z,vx,vy,vz : float
stellar positions and velocities (now with rotation)
History
-------
2019 - Written - Webb (UofT)
"""
r, theta, z = bovy_coords.rect_to_cyl(cluster.x, cluster.y, cluster.z)
vr, vtheta, vz = bovy_coords.rect_to_cyl_vec(cluster.vx, cluster.vy,
cluster.vz, cluster.x,
cluster.y, cluster.z)
indx = vtheta < 0.0
rindx = np.random.rand(cluster.ntot) < qrot
vtheta[indx * rindx] = np.sqrt(vtheta[indx * rindx] * vtheta[indx * rindx])
x, y, z = bovy_coords.cyl_to_rect(r, theta, z)
vx, vy, vz = bovy_coords.cyl_to_rect_vec(vr, vtheta, vz, theta)
return x, y, z, vx, vy, vz
def sample(self, n, returndt=False, integrate=True, xy=False, lb=False):
"""
NAME:
sample
PURPOSE:
sample from the DF
INPUT:
n - number of points to return
returndt= (False) if True, also return the time since the star was stripped
integrate= (True) if True, integrate the orbits to the present time, if False, return positions at stripping (probably want to combine with returndt=True then to make sense of them!)
xy= (False) if True, return Galactocentric rectangular coordinates
lb= (False) if True, return Galactic l,b,d,vlos,pmll,pmbb coordinates
OUTPUT:
(R,vR,vT,z,vz,phi) of points on the stream in 6,N array
HISTORY:
2018-07-31 - Written - Bovy (IAS)
"""
if xy or lb:
raise NotImplementedError(
"xy=True and lb=True options currently not implemented")
# First sample times
dt = numpy.random.uniform(size=n) * self._tdisrupt
# Build all rotation matrices
rot, rot_inv = self._setup_rot(dt)
# Compute progenitor position in the instantaneous frame
xyzpt = numpy.einsum(
'ijk,ik->ij', rot,
numpy.array([
self._progenitor.x(-dt),
self._progenitor.y(-dt),
self._progenitor.z(-dt)
]).T)
vxyzpt = numpy.einsum(
'ijk,ik->ij', rot,
numpy.array([
self._progenitor.vx(-dt),
self._progenitor.vy(-dt),
self._progenitor.vz(-dt)
]).T)
Rpt, phipt, Zpt = bovy_coords.rect_to_cyl(xyzpt[:, 0], xyzpt[:, 1],
xyzpt[:, 2])
vRpt, vTpt, vZpt = bovy_coords.rect_to_cyl_vec(vxyzpt[:, 0],
vxyzpt[:, 1],
vxyzpt[:, 2],
Rpt,
phipt,
Zpt,
cyl=True)
# Sample positions and velocities in the instantaneous frame
k = self._meankvec + numpy.random.normal(
size=n)[:, numpy.newaxis] * self._sigkvec
try:
rtides = rtide(self._rtpot,
Rpt,
Zpt,
phi=phipt,
t=-dt,
M=self._progenitor_mass,
use_physical=False)
vcs = numpy.sqrt(-Rpt * evaluateRforces(
self._rtpot, Rpt, Zpt, phi=phipt, t=-dt, use_physical=False))
except (ValueError, TypeError):
rtides = numpy.array([
rtide(self._rtpot,
Rpt[ii],
Zpt[ii],
phi=phipt[ii],
t=-dt[ii],
M=self._progenitor_mass,
use_physical=False) for ii in range(len(Rpt))
])
vcs = numpy.array([
numpy.sqrt(-Rpt[ii] * evaluateRforces(self._rtpot,
Rpt[ii],
Zpt[ii],
phi=phipt[ii],
t=-dt[ii],
use_physical=False))
for ii in range(len(Rpt))
])
rtides_as_frac = rtides / Rpt
RpZst = numpy.array([
Rpt + k[:, 0] * rtides, phipt + k[:, 5] * rtides_as_frac,
k[:, 3] * rtides_as_frac
]).T
vRTZst = numpy.array([
vRpt * (1. + k[:, 1]), vTpt + k[:, 2] * vcs * rtides_as_frac,
k[:, 4] * vcs * rtides_as_frac
]).T
# Now rotate these back to the galactocentric frame
xst, yst, zst = bovy_coords.cyl_to_rect(RpZst[:, 0], RpZst[:, 1],
RpZst[:, 2])
vxst, vyst, vzst = bovy_coords.cyl_to_rect_vec(vRTZst[:, 0], vRTZst[:,
1],
vRTZst[:, 2], RpZst[:,
1])
xyzs = numpy.einsum('ijk,ik->ij', rot_inv,
numpy.array([xst, yst, zst]).T)
vxyzs = numpy.einsum('ijk,ik->ij', rot_inv,
numpy.array([vxst, vyst, vzst]).T)
Rs, phis, Zs = bovy_coords.rect_to_cyl(xyzs[:, 0], xyzs[:, 1], xyzs[:,
2])
vRs, vTs, vZs = bovy_coords.rect_to_cyl_vec(vxyzs[:, 0],
vxyzs[:, 1],
vxyzs[:, 2],
Rs,
phis,
Zs,
cyl=True)
out = numpy.empty((6, n))
if integrate:
# Now integrate the orbits
for ii in range(n):
o = Orbit(
[Rs[ii], vRs[ii], vTs[ii], Zs[ii], vZs[ii], phis[ii]])
o.integrate(numpy.linspace(-dt[ii], 0., 10001), self._pot)
o = o(0.)
out[:, ii] = [o.R(), o.vR(), o.vT(), o.z(), o.vz(), o.phi()]
else:
out[0] = Rs
out[1] = vRs
out[2] = vTs
out[3] = Zs
out[4] = vZs
out[5] = phis
if returndt:
return (out, dt)
else:
return out
def _determine_impact_coordtransform(self,deltaAngleTrackImpact,
nTrackChunksImpact,
timpact,impact_angle):
"""Function that sets up the transformation between (x,v) and (O,theta)"""
# Integrate the progenitor backward to the time of impact
self._gap_progenitor_setup()
# Sign of delta angle tells us whether the impact happens to the
# leading or trailing arm, self._sigMeanSign contains this info
if impact_angle > 0.:
self._gap_leading= True
else:
self._gap_leading= False
if (self._gap_leading and not self._leading) \
or (not self._gap_leading and self._leading):
raise ValueError('Modeling leading (trailing) impact for trailing (leading) arm; this is not allowed because it is nonsensical in this framework')
self._impact_angle= numpy.fabs(impact_angle)
self._gap_sigMeanSign= 1.
if (self._gap_leading and self._progenitor_Omega_along_dOmega/self._sigMeanSign < 0.) \
or (not self._gap_leading and self._progenitor_Omega_along_dOmega/self._sigMeanSign > 0.):
self._gap_sigMeanSign= -1.
# Determine how much orbital time is necessary for the progenitor's orbit at the time of impact to cover the part of the stream near the impact; we cover the whole leading (or trailing) part of the stream
if nTrackChunksImpact is None:
#default is floor(self._deltaAngleTrackImpact/0.15)+1
self._nTrackChunksImpact= int(numpy.floor(self._deltaAngleTrackImpact/0.15))+1
else:
self._nTrackChunksImpact= nTrackChunksImpact
dt= self._deltaAngleTrackImpact\
/self._progenitor_Omega_along_dOmega\
/self._sigMeanSign*self._gap_sigMeanSign
self._gap_trackts= numpy.linspace(0.,2*dt,2*self._nTrackChunksImpact-1) #to be sure that we cover it
#Instantiate an auxiliaryTrack, which is an Orbit instance at the mean frequency of the stream, and zero angle separation wrt the progenitor; prog_stream_offset is the offset between this track and the progenitor at zero angle (same as in streamdf, but just done at the time of impact rather than the current time)
prog_stream_offset=\
_determine_stream_track_single(self._aA,
self._gap_progenitor,
self._timpact, # around the t of imp
self._progenitor_angle-self._timpact*self._progenitor_Omega,
self._gap_sigMeanSign,
self._dsigomeanProgDirection,
lambda da: super(streamgapdf,self).meanOmega(da,offset_sign=self._gap_sigMeanSign,tdisrupt=self._tdisrupt-self._timpact),
0.) #angle = 0
auxiliaryTrack= Orbit(prog_stream_offset[3])
if dt < 0.:
self._gap_trackts= numpy.linspace(0.,-2.*dt,2.*self._nTrackChunksImpact-1)
#Flip velocities before integrating
auxiliaryTrack= auxiliaryTrack.flip()
auxiliaryTrack.integrate(self._gap_trackts,self._pot)
if dt < 0.:
#Flip velocities again
auxiliaryTrack._orb.orbit[:,1]= -auxiliaryTrack._orb.orbit[:,1]
auxiliaryTrack._orb.orbit[:,2]= -auxiliaryTrack._orb.orbit[:,2]
auxiliaryTrack._orb.orbit[:,4]= -auxiliaryTrack._orb.orbit[:,4]
#Calculate the actions, frequencies, and angle for this auxiliary orbit
acfs= self._aA.actionsFreqs(auxiliaryTrack(0.),maxn=3)
auxiliary_Omega= numpy.array([acfs[3],acfs[4],acfs[5]]).reshape(3\
)
auxiliary_Omega_along_dOmega= \
numpy.dot(auxiliary_Omega,self._dsigomeanProgDirection)
# compute the transformation using _determine_stream_track_single
allAcfsTrack= numpy.empty((self._nTrackChunksImpact,9))
alljacsTrack= numpy.empty((self._nTrackChunksImpact,6,6))
allinvjacsTrack= numpy.empty((self._nTrackChunksImpact,6,6))
thetasTrack= numpy.linspace(0.,self._deltaAngleTrackImpact,
self._nTrackChunksImpact)
ObsTrack= numpy.empty((self._nTrackChunksImpact,6))
ObsTrackAA= numpy.empty((self._nTrackChunksImpact,6))
detdOdJps= numpy.empty((self._nTrackChunksImpact))
if self._multi is None:
for ii in range(self._nTrackChunksImpact):
multiOut= _determine_stream_track_single(self._aA,
auxiliaryTrack,
self._gap_trackts[ii]*numpy.fabs(self._progenitor_Omega_along_dOmega/auxiliary_Omega_along_dOmega), #this factor accounts for the difference in frequency between the progenitor and the auxiliary track, no timpact bc gap_tracks is relative to timpact
self._progenitor_angle-self._timpact*self._progenitor_Omega,
self._gap_sigMeanSign,
self._dsigomeanProgDirection,
lambda da: super(streamgapdf,self).meanOmega(da,offset_sign=self._gap_sigMeanSign,tdisrupt=self._tdisrupt-self._timpact),
thetasTrack[ii])
allAcfsTrack[ii,:]= multiOut[0]
alljacsTrack[ii,:,:]= multiOut[1]
allinvjacsTrack[ii,:,:]= multiOut[2]
ObsTrack[ii,:]= multiOut[3]
ObsTrackAA[ii,:]= multiOut[4]
detdOdJps[ii]= multiOut[5]
else:
multiOut= multi.parallel_map(\
(lambda x: _determine_stream_track_single(self._aA,
auxiliaryTrack,
self._gap_trackts[x]*numpy.fabs(self._progenitor_Omega_along_dOmega/auxiliary_Omega_along_dOmega), #this factor accounts for the difference in frequency between the progenitor and the auxiliary track, no timpact bc gap_tracks is relative to timpact
self._progenitor_angle-self._timpact*self._progenitor_Omega,
self._gap_sigMeanSign,
self._dsigomeanProgDirection,
lambda da: super(streamgapdf,self).meanOmega(da,offset_sign=self._gap_sigMeanSign,tdisrupt=self._tdisrupt-self._timpact),
thetasTrack[x])),
range(self._nTrackChunksImpact),
numcores=numpy.amin([self._nTrackChunksImpact,
multiprocessing.cpu_count(),
self._multi]))
for ii in range(self._nTrackChunksImpact):
allAcfsTrack[ii,:]= multiOut[ii][0]
alljacsTrack[ii,:,:]= multiOut[ii][1]
allinvjacsTrack[ii,:,:]= multiOut[ii][2]
ObsTrack[ii,:]= multiOut[ii][3]
ObsTrackAA[ii,:]= multiOut[ii][4]
detdOdJps[ii]= multiOut[ii][5]
#Repeat the track calculation using the previous track, to get closer to it
for nn in range(self.nTrackIterations):
if self._multi is None:
for ii in range(self._nTrackChunksImpact):
multiOut= _determine_stream_track_single(self._aA,
Orbit(ObsTrack[ii,:]),
0.,
self._progenitor_angle-self._timpact*self._progenitor_Omega,
self._gap_sigMeanSign,
self._dsigomeanProgDirection,
lambda da: super(streamgapdf,self).meanOmega(da,offset_sign=self._gap_sigMeanSign,tdisrupt=self._tdisrupt-self._timpact),
thetasTrack[ii])
allAcfsTrack[ii,:]= multiOut[0]
alljacsTrack[ii,:,:]= multiOut[1]
allinvjacsTrack[ii,:,:]= multiOut[2]
ObsTrack[ii,:]= multiOut[3]
ObsTrackAA[ii,:]= multiOut[4]
detdOdJps[ii]= multiOut[5]
else:
multiOut= multi.parallel_map(\
(lambda x: _determine_stream_track_single(self._aA,Orbit(ObsTrack[x,:]),0.,
self._progenitor_angle-self._timpact*self._progenitor_Omega,
self._gap_sigMeanSign,
self._dsigomeanProgDirection,
lambda da: super(streamgapdf,self).meanOmega(da,offset_sign=self._gap_sigMeanSign,tdisrupt=self._tdisrupt-self._timpact),
thetasTrack[x])),
range(self._nTrackChunksImpact),
numcores=numpy.amin([self._nTrackChunksImpact,
multiprocessing.cpu_count(),
self._multi]))
for ii in range(self._nTrackChunksImpact):
allAcfsTrack[ii,:]= multiOut[ii][0]
alljacsTrack[ii,:,:]= multiOut[ii][1]
allinvjacsTrack[ii,:,:]= multiOut[ii][2]
ObsTrack[ii,:]= multiOut[ii][3]
ObsTrackAA[ii,:]= multiOut[ii][4]
detdOdJps[ii]= multiOut[ii][5]
#Store the track
self._gap_thetasTrack= thetasTrack
self._gap_ObsTrack= ObsTrack
self._gap_ObsTrackAA= ObsTrackAA
self._gap_allAcfsTrack= allAcfsTrack
self._gap_alljacsTrack= alljacsTrack
self._gap_allinvjacsTrack= allinvjacsTrack
self._gap_detdOdJps= detdOdJps
self._gap_meandetdOdJp= numpy.mean(self._gap_detdOdJps)
self._gap_logmeandetdOdJp= numpy.log(self._gap_meandetdOdJp)
#Also calculate _ObsTrackXY in XYZ,vXYZ coordinates
self._gap_ObsTrackXY= numpy.empty_like(self._gap_ObsTrack)
TrackX= self._gap_ObsTrack[:,0]*numpy.cos(self._gap_ObsTrack[:,5])
TrackY= self._gap_ObsTrack[:,0]*numpy.sin(self._gap_ObsTrack[:,5])
TrackZ= self._gap_ObsTrack[:,3]
TrackvX, TrackvY, TrackvZ=\
bovy_coords.cyl_to_rect_vec(self._gap_ObsTrack[:,1],
self._gap_ObsTrack[:,2],
self._gap_ObsTrack[:,4],
self._gap_ObsTrack[:,5])
self._gap_ObsTrackXY[:,0]= TrackX
self._gap_ObsTrackXY[:,1]= TrackY
self._gap_ObsTrackXY[:,2]= TrackZ
self._gap_ObsTrackXY[:,3]= TrackvX
self._gap_ObsTrackXY[:,4]= TrackvY
self._gap_ObsTrackXY[:,5]= TrackvZ
return None
def _interpolate_stream_track_kick(self):
"""Build interpolations of the stream track near the kick"""
if hasattr(self,'_kick_interpolatedThetasTrack'): #pragma: no cover
return None #Already did this
# Setup the trackpoints where the kick will be computed, covering the
# full length of the stream
self._kick_interpolatedThetasTrack= \
numpy.linspace(self._gap_thetasTrack[0],
self._gap_thetasTrack[-1],
self._nKickPoints)
TrackX= self._gap_ObsTrack[:,0]*numpy.cos(self._gap_ObsTrack[:,5])
TrackY= self._gap_ObsTrack[:,0]*numpy.sin(self._gap_ObsTrack[:,5])
TrackZ= self._gap_ObsTrack[:,3]
TrackvX, TrackvY, TrackvZ=\
bovy_coords.cyl_to_rect_vec(self._gap_ObsTrack[:,1],
self._gap_ObsTrack[:,2],
self._gap_ObsTrack[:,4],
self._gap_ObsTrack[:,5])
#Interpolate
self._kick_interpTrackX=\
interpolate.InterpolatedUnivariateSpline(self._gap_thetasTrack,
TrackX,k=3)
self._kick_interpTrackY=\
interpolate.InterpolatedUnivariateSpline(self._gap_thetasTrack,
TrackY,k=3)
self._kick_interpTrackZ=\
interpolate.InterpolatedUnivariateSpline(self._gap_thetasTrack,
TrackZ,k=3)
self._kick_interpTrackvX=\
interpolate.InterpolatedUnivariateSpline(self._gap_thetasTrack,
TrackvX,k=3)
self._kick_interpTrackvY=\
interpolate.InterpolatedUnivariateSpline(self._gap_thetasTrack,
TrackvY,k=3)
self._kick_interpTrackvZ=\
interpolate.InterpolatedUnivariateSpline(self._gap_thetasTrack,
TrackvZ,k=3)
#Now store an interpolated version of the stream track
self._kick_interpolatedObsTrackXY= numpy.empty((len(self._kick_interpolatedThetasTrack),6))
self._kick_interpolatedObsTrackXY[:,0]=\
self._kick_interpTrackX(self._kick_interpolatedThetasTrack)
self._kick_interpolatedObsTrackXY[:,1]=\
self._kick_interpTrackY(self._kick_interpolatedThetasTrack)
self._kick_interpolatedObsTrackXY[:,2]=\
self._kick_interpTrackZ(self._kick_interpolatedThetasTrack)
self._kick_interpolatedObsTrackXY[:,3]=\
self._kick_interpTrackvX(self._kick_interpolatedThetasTrack)
self._kick_interpolatedObsTrackXY[:,4]=\
self._kick_interpTrackvY(self._kick_interpolatedThetasTrack)
self._kick_interpolatedObsTrackXY[:,5]=\
self._kick_interpTrackvZ(self._kick_interpolatedThetasTrack)
#Also in cylindrical coordinates
self._kick_interpolatedObsTrack= \
numpy.empty((len(self._kick_interpolatedThetasTrack),6))
tR,tphi,tZ= bovy_coords.rect_to_cyl(self._kick_interpolatedObsTrackXY[:,0],
self._kick_interpolatedObsTrackXY[:,1],
self._kick_interpolatedObsTrackXY[:,2])
tvR,tvT,tvZ=\
bovy_coords.rect_to_cyl_vec(self._kick_interpolatedObsTrackXY[:,3],
self._kick_interpolatedObsTrackXY[:,4],
self._kick_interpolatedObsTrackXY[:,5],
tR,tphi,tZ,cyl=True)
self._kick_interpolatedObsTrack[:,0]= tR
self._kick_interpolatedObsTrack[:,1]= tvR
self._kick_interpolatedObsTrack[:,2]= tvT
self._kick_interpolatedObsTrack[:,3]= tZ
self._kick_interpolatedObsTrack[:,4]= tvZ
self._kick_interpolatedObsTrack[:,5]= tphi
# Also store (x,v) for the point of closest approach
self._kick_ObsTrackXY_closest= numpy.array([\
self._kick_interpTrackX(self._impact_angle),
self._kick_interpTrackY(self._impact_angle),
self._kick_interpTrackZ(self._impact_angle),
self._kick_interpTrackvX(self._impact_angle),
self._kick_interpTrackvY(self._impact_angle),
self._kick_interpTrackvZ(self._impact_angle)])
return None
