Mercurial > repos > astroteam > crbeam_astro_tool

comparison mc_rotate.py @ 0:f40d05521dca draft default tip

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
planemo upload for repository https://github.com/esg-epfl-apc/tools-astro/tree/main/tools commit de01e3c02a26cd6353a6b9b6f8d1be44de8ccd54
author astroteam
date Fri, 25 Apr 2025 19:33:20 +0000
parents
children
comparison
equal deleted inserted replaced
-1:000000000000 0:f40d05521dca
1 import math
2
3 import numpy as np
4
5
6 def mc_rotate(
7 mc_file, jet_half_size_deg, jet_direction_deg, psf_deg=180, rotated_file=None
8 ):
9 if rotated_file is None:
10 rotated_file = f"{mc_file}_{jet_half_size_deg}_{jet_direction_deg}_{psf_deg}"
11
12 jet_half_size = (
13 np.longdouble(jet_half_size_deg) / 180.0 * math.pi
14 ) # jet openning angle
15 jet_direction = (
16 np.longdouble(jet_direction_deg) / 180.0 * math.pi
17 ) # angle between direction to source and jet axes
18
19 if jet_half_size > math.pi or jet_half_size <= 0:
20 jet_half_size = math.pi
21 # here we assume that jet axis is located in y-z plane if angle=0, then jet points towards observer
22 jet_axes = np.array(
23 [0, np.sin(jet_direction), np.cos(jet_direction)], dtype=np.longdouble
24 )
25
26 if jet_direction > math.pi or jet_direction < 0:
27 raise ValueError(
28 "invalid jet_direction_deg param [0, 180] range expected")
29
30 with open(mc_file, "r") as lines, open(rotated_file, "w") as rotated:
31 column_names = []
32 for line in lines:
33 columns = line.split()
34 if len(columns) == 0:
35 continue
36 if columns[0] == "#":
37 if "weight" in columns:
38 column_names = (
39 columns[1:3]
40 + ["Theta", "Phi", "dT_raw/T", "dT_calc/T"]
41 + columns[9:]
42 )
43 print("# " + "\t".join(column_names), file=rotated)
44 else:
45 print(line.strip(), file=rotated)
46 else:
47 particle = np.array(list(map(np.longdouble, columns)))
48 beta = phi = 0 #dTcalc = 0
49
50 r = particle[2:5]
51 dz = r[
52 2
53 # dz is roughly equal to time delay in units of full path (assuming very small deflection angle)
54 ]
55
56 # negative dz may apper in data as a result of rounding errors
57 dz = max(dz, 0) # replacing negative values by zeros
58
59 r[2] = 1.0 - dz # now r is position of particle
60
61 # direction of momentum
62 n = particle[5:8]
63 n[2] = 1.0 - n[2]
64 # compensate rounding errors for n
65 mod_n = np.dot(n, n) ** 0.5
66 n = 1 / mod_n * n
67 t = 0 # calculated delay
68 if dz > 0: # there is some delay
69 r2 = np.dot(r, r)
70 # assert r2 < 1
71 if r2 > 1:
72 r *= 1.0 / np.sqrt(r2)
73 r2 = 1.0
74 # we continue trajectory in the direction of momentun until it reaches the sphere of radius 1.
75 # to find the delay t, we need to solve the following quadratic equation with respect to variable t:
76 # [(x + V_x)*t]^2 + [(y + V_y)*t]^2 + [(z + V_z)*t]^2 = 1
77 # with extra condition t > 0
78 # below is the solution:
79 rV = np.dot(n, r)
80 t = np.sqrt(1.0 - r2 + rV * rV) - rV
81
82 # calculating the end point of the trajectory
83 r = r + t * n
84
85 # normalizing r to 1 to fix the rounding errors
86 modR = np.dot(r, r) ** 0.5
87 r = 1.0 / modR * r
88
89 # now we will construct the rotation of coordinates which converts vector r to 0,0,1
90 # to find the starting direction of the particle momentum
91 dr2 = r[0] ** 2 + r[1] ** 2
92 if dr2 > 0:
93 # building rotation matrix to move x,y,z to 0,0,1
94 norm = 1.0 / (dr2**0.5)
95 # rotation axes:
96 ux = norm * r[1]
97 uy = -norm * r[0]
98 # uz=0
99 cos = r[2]
100 sin = (1 - cos**2) ** 0.5
101
102 toCenter = np.array(
103 [
104 [cos + ux * ux * (1 - cos), ux *
105 uy * (1 - cos), uy * sin],
106 [ux * uy * (1 - cos), cos + uy *
107 uy * (1 - cos), -ux * sin],
108 [-uy * sin, ux * sin, cos],
109 ]
110 )
111
112 zAxes = np.array([0.0, 0.0, 1.0], dtype=np.longdouble)
113 # test
114 error = np.fabs(np.dot(toCenter, r) - zAxes).max()
115 assert error < 1e-7
116 # finding starting direction of the momentum
117 nInitial = np.dot(toCenter, zAxes)
118 norm_nInitial = np.sqrt(np.dot(nInitial, nInitial))
119 nInitial *= (
120 1.0 / norm_nInitial
121 ) # normalize after rotation to avoid rounding errors
122 # calculating angle between starting direction and jet axis
123 startCos = np.dot(nInitial, jet_axes)
124 startAngle = 0
125 if (
126 startCos < 1
127 ): # avoid roundout errors if angle is very close to zero
128 startAngle = np.arccos(startCos)
129 # TODO: here instead of using step function we can modify weight depending
130 # on angular distance from jet axes
131 if (
132 startAngle < jet_half_size
133 ): # if the particle is inside the cone (passes cut)
134 # calculate angle between arrival direction and direction to source
135 cosBeta = np.dot(
136 r, n
137 ) # this angle is invariant (doesn't depend on turns)
138 if (
139 cosBeta < 1
140 ): # avoid rounding errors if angle is very close to zero
141 beta = np.arccos(cosBeta)
142 nFinal = np.dot(toCenter, n).transpose()
143 cosPhi = nFinal[0] / np.sin(
144 beta
145 # second angle (could be useful in assymetric tasks) if beta=0 it is undefined (could be any)
146 )
147 if cosPhi < 1 and cosPhi > -1:
148 phi = np.arccos(
149 cosPhi
150 ) # returns value in interval [0,pi]
151 if nFinal[1] < 0:
152 phi = 2 * math.pi - phi
153
154 # s = math.sin(startAngle+beta)
155 # if s>0.:
156 # dTcalc = (math.sin(startAngle)+math.sin(beta))/s - 1.
157 # dTcalc = (1.-min(modR, 1.))/cosBeta # condition modR > 1 can only occur due to rounding error
158 else:
159 continue
160 elif (
161 jet_direction > jet_half_size
162 ): # no deflection: check if jet is pointing towards observer
163 continue
164
165 if beta * 180.0 / math.pi > psf_deg:
166 continue
167
168 output = [
169 particle[0],
170 particle[1],
171 beta * 180.0 / math.pi,
172 phi * 180.0 / math.pi,
173 dz,
174 t,
175 ] + list(particle[8:])
176 output = ["{0:g}".format(i) for i in output]
177 print(*output, file=rotated)
178
179 return rotated_file