--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/mc_rotate.py	Fri Apr 25 19:33:20 2025 +0000
@@ -0,0 +1,179 @@
+import math
+
+import numpy as np
+
+
+def mc_rotate(
+    mc_file, jet_half_size_deg, jet_direction_deg, psf_deg=180, rotated_file=None
+):
+    if rotated_file is None:
+        rotated_file = f"{mc_file}_{jet_half_size_deg}_{jet_direction_deg}_{psf_deg}"
+
+    jet_half_size = (
+        np.longdouble(jet_half_size_deg) / 180.0 * math.pi
+    )  # jet openning angle
+    jet_direction = (
+        np.longdouble(jet_direction_deg) / 180.0 * math.pi
+    )  # angle between direction to source and jet axes
+
+    if jet_half_size > math.pi or jet_half_size <= 0:
+        jet_half_size = math.pi
+    # here we assume that jet axis is located in y-z plane if angle=0, then jet points towards observer
+    jet_axes = np.array(
+        [0, np.sin(jet_direction), np.cos(jet_direction)], dtype=np.longdouble
+    )
+
+    if jet_direction > math.pi or jet_direction < 0:
+        raise ValueError(
+            "invalid jet_direction_deg param [0, 180] range expected")
+
+    with open(mc_file, "r") as lines, open(rotated_file, "w") as rotated:
+        column_names = []
+        for line in lines:
+            columns = line.split()
+            if len(columns) == 0:
+                continue
+            if columns[0] == "#":
+                if "weight" in columns:
+                    column_names = (
+                        columns[1:3]
+                        + ["Theta", "Phi", "dT_raw/T", "dT_calc/T"]
+                        + columns[9:]
+                    )
+                    print("# " + "\t".join(column_names), file=rotated)
+                else:
+                    print(line.strip(), file=rotated)
+            else:
+                particle = np.array(list(map(np.longdouble, columns)))
+                beta = phi = 0 #dTcalc = 0
+
+                r = particle[2:5]
+                dz = r[
+                    2
+                    # dz is roughly equal to time delay in units of full path (assuming very small deflection angle)
+                ]
+
+                # negative dz may apper in data as a result of rounding errors
+                dz = max(dz, 0)  # replacing negative values by zeros
+
+                r[2] = 1.0 - dz  # now r is position of particle
+
+                # direction of momentum
+                n = particle[5:8]
+                n[2] = 1.0 - n[2]
+                # compensate rounding errors for n
+                mod_n = np.dot(n, n) ** 0.5
+                n = 1 / mod_n * n
+                t = 0  # calculated delay
+                if dz > 0:  # there is some delay
+                    r2 = np.dot(r, r)
+                    # assert r2 < 1
+                    if r2 > 1:
+                        r *= 1.0 / np.sqrt(r2)
+                        r2 = 1.0
+                    # we continue trajectory in the direction of momentun until it reaches the sphere of radius 1.
+                    # to find the delay t, we need to solve the following quadratic equation with respect to variable t:
+                    # [(x + V_x)*t]^2 + [(y + V_y)*t]^2 + [(z + V_z)*t]^2 = 1
+                    # with extra condition t > 0
+                    # below is the solution:
+                    rV = np.dot(n, r)
+                    t = np.sqrt(1.0 - r2 + rV * rV) - rV
+
+                    # calculating the end point of the trajectory
+                    r = r + t * n
+
+                # normalizing r to 1 to fix the rounding errors
+                modR = np.dot(r, r) ** 0.5
+                r = 1.0 / modR * r
+
+                # now we will construct the rotation of coordinates which converts vector r to 0,0,1
+                # to find the starting direction of the particle momentum
+                dr2 = r[0] ** 2 + r[1] ** 2
+                if dr2 > 0:
+                    # building rotation matrix to move x,y,z to 0,0,1
+                    norm = 1.0 / (dr2**0.5)
+                    # rotation axes:
+                    ux = norm * r[1]
+                    uy = -norm * r[0]
+                    # uz=0
+                    cos = r[2]
+                    sin = (1 - cos**2) ** 0.5
+
+                    toCenter = np.array(
+                        [
+                            [cos + ux * ux * (1 - cos), ux *
+                             uy * (1 - cos), uy * sin],
+                            [ux * uy * (1 - cos), cos + uy *
+                             uy * (1 - cos), -ux * sin],
+                            [-uy * sin, ux * sin, cos],
+                        ]
+                    )
+
+                    zAxes = np.array([0.0, 0.0, 1.0], dtype=np.longdouble)
+                    # test
+                    error = np.fabs(np.dot(toCenter, r) - zAxes).max()
+                    assert error < 1e-7
+                    # finding starting direction of the momentum
+                    nInitial = np.dot(toCenter, zAxes)
+                    norm_nInitial = np.sqrt(np.dot(nInitial, nInitial))
+                    nInitial *= (
+                        1.0 / norm_nInitial
+                    )  # normalize after rotation to avoid rounding errors
+                    # calculating angle between starting direction and jet axis
+                    startCos = np.dot(nInitial, jet_axes)
+                    startAngle = 0
+                    if (
+                        startCos < 1
+                    ):  # avoid roundout errors if angle is very close to zero
+                        startAngle = np.arccos(startCos)
+                    # TODO: here instead of using step function we can modify weight depending
+                    # on angular distance from jet axes
+                    if (
+                        startAngle < jet_half_size
+                    ):  # if the particle is inside the cone (passes cut)
+                        # calculate angle between arrival direction and direction to source
+                        cosBeta = np.dot(
+                            r, n
+                        )  # this angle is invariant (doesn't depend on turns)
+                        if (
+                            cosBeta < 1
+                        ):  # avoid rounding errors if angle is very close to zero
+                            beta = np.arccos(cosBeta)
+                            nFinal = np.dot(toCenter, n).transpose()
+                            cosPhi = nFinal[0] / np.sin(
+                                beta
+                                # second angle (could be useful in assymetric tasks) if beta=0 it is undefined (could be any)
+                            )
+                            if cosPhi < 1 and cosPhi > -1:
+                                phi = np.arccos(
+                                    cosPhi
+                                )  # returns value in interval [0,pi]
+                                if nFinal[1] < 0:
+                                    phi = 2 * math.pi - phi
+
+                        # s = math.sin(startAngle+beta)
+                        # if s>0.:
+                        #     dTcalc = (math.sin(startAngle)+math.sin(beta))/s - 1.
+                        # dTcalc = (1.-min(modR, 1.))/cosBeta  # condition modR > 1 can only occur due to rounding error
+                    else:
+                        continue
+                elif (
+                    jet_direction > jet_half_size
+                ):  # no deflection: check if jet is pointing towards observer
+                    continue
+
+                if beta * 180.0 / math.pi > psf_deg:
+                    continue
+
+                output = [
+                    particle[0],
+                    particle[1],
+                    beta * 180.0 / math.pi,
+                    phi * 180.0 / math.pi,
+                    dz,
+                    t,
+                ] + list(particle[8:])
+                output = ["{0:g}".format(i) for i in output]
+                print(*output, file=rotated)
+
+    return rotated_file