/
stream_sim_util.py
108 lines (104 loc) · 4.17 KB
/
stream_sim_util.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
# stream_sim_util.py: some utility functions copied from stream-stream
import os, os.path
import csv
import numpy
import copy
from galpy.potential import LogarithmicHaloPotential
from galpy.util import bovy_coords, multi
from galpy.actionAngle import actionAngleIsochroneApprox
R0= 8.
V0= 220.
lp= LogarithmicHaloPotential(normalize=1.,q=0.9)
def rectangular_to_cylindrical(xv):
R,phi,Z= bovy_coords.rect_to_cyl(xv[:,0],xv[:,1],xv[:,2])
vR,vT,vZ= bovy_coords.rect_to_cyl_vec(xv[:,3],xv[:,4],xv[:,5],
R,phi,Z,cyl=True)
out= numpy.empty_like(xv)
# Preferred galpy arrangement of cylindrical coordinates
out[:,0]= R
out[:,1]= vR
out[:,2]= vT
out[:,3]= Z
out[:,4]= vZ
out[:,5]= phi
return out
def calc_aA_sim(RvR,filename,snap_gc):
# Calculate the action angle variables for a simulation and store
if not os.path.exists(filename):
aAI= actionAngleIsochroneApprox(pot=lp,b=0.8)
nbatch= 20
multiOut= multi.parallel_map(\
lambda x: aAI.actionsFreqsAngles(RvR[x*nbatch:(x+1)*nbatch,0]/R0,
RvR[x*nbatch:(x+1)*nbatch,1]/V0,
RvR[x*nbatch:(x+1)*nbatch,2]/V0,
RvR[x*nbatch:(x+1)*nbatch,3]/R0,
RvR[x*nbatch:(x+1)*nbatch,4]/V0,
RvR[x*nbatch:(x+1)*nbatch,5]),
range(len(snap_gc)//nbatch),
numcores=25)
acfs= numpy.reshape(numpy.swapaxes(numpy.array(multiOut),0,1),
(9,numpy.prod(numpy.array(multiOut).shape)//9))
# Write to file
csvfile= open(filename,'w')
writer= csv.writer(csvfile,delimiter=',')
for jj in range(len(acfs[0])):
writer.writerow([acfs[0][jj],acfs[1][jj],acfs[2][jj],
acfs[3][jj],acfs[4][jj],acfs[5][jj],
acfs[6][jj],acfs[7][jj],acfs[8][jj]])
csvfile.flush()
csvfile.close()
else:
acfs= numpy.loadtxt(filename,delimiter=',').T
return acfs
def calc_apar(acfs,angle=None,freq=False,debrisThreshold=6.):
# Calculate the parallel angle offset,
# of angle if set (otherwise of the entire simulation),
# angle is a frequency if freq
thetar= acfs[6]
thetap= acfs[7]
thetaz= acfs[8]
if not angle is None:
if not freq:
angle[0]= (numpy.pi+(angle[0]-numpy.median(thetar))) % (2.*numpy.pi)
angle[1]= (numpy.pi+(angle[1]-numpy.median(thetap))) % (2.*numpy.pi)
angle[2]= (numpy.pi+(angle[2]-numpy.median(thetaz))) % (2.*numpy.pi)
thetap= (numpy.pi+(thetap-numpy.median(thetap))) % (2.*numpy.pi)
debrisIndx= numpy.fabs(thetap-numpy.pi) > (debrisThreshold*numpy.median(numpy.fabs(thetap-numpy.median(thetap))))
if angle is None:
thetar= (numpy.pi+(thetar-numpy.median(thetar))) % (2.*numpy.pi)
thetaz= (numpy.pi+(thetaz-numpy.median(thetaz))) % (2.*numpy.pi)
#center around 0 (instead of pi)
thetar-= numpy.pi
thetap-= numpy.pi
thetaz-= numpy.pi
elif freq:
thetar= angle[0]
thetap= angle[1]
thetaz= angle[2]
else:
thetar= angle[0]-numpy.pi
thetap= angle[1]-numpy.pi
thetaz= angle[2]-numpy.pi
#Frequencies
Or= acfs[3]
Op= acfs[4]
Oz= acfs[5]
dOr= Or[debrisIndx]-numpy.median(Or)
dOp= Op[debrisIndx]-numpy.median(Op)
dOz= Oz[debrisIndx]-numpy.median(Oz)
dO= numpy.vstack((dOr,dOp,dOz))
dO4dir= copy.copy(dO)
dO4dir[:,dO4dir[0] < 0.]*= -1.
dOdir= numpy.median(dO4dir,axis=1)
dOdir/= numpy.sqrt(numpy.sum(dOdir**2.))
#parallel angle
dangle= numpy.vstack((thetar,thetap,thetaz))
return numpy.dot(dangle.T,dOdir)
def calc_apars_dm(thetar,thetap,thetaz,aa_dm,oparDir_dm,t=0.):
dangle= numpy.vstack((thetar-aa_dm[3]*t,
thetap-aa_dm[4]*t,
thetaz-aa_dm[5]*t))
return numpy.dot(dangle.T,oparDir_dm)
def calc_opars_dm(Or,Op,Oz,oparDir_dm):
dfreq= numpy.vstack((Or,Op,Oz))
return numpy.dot(dfreq.T,oparDir_dm)