utils.py 3.5 KB
 DavidWalz committed May 04, 2017 1 2 3 4 5 6 ``````from __future__ import division import numpy as np def rectangular_array(n=11): `````` DavidWalz committed May 04, 2017 7 `````` """ Return x,y coordinates for rectangular array with n^2 stations. """ `````` DavidWalz committed May 04, 2017 8 `````` n0 = (n - 1) / 2 `````` DavidWalz committed May 04, 2017 9 `````` return (np.mgrid[0:n, 0:n].astype(float) - n0) `````` DavidWalz committed May 04, 2017 10 11 12 `````` def triangular_array(n=11, offset=True): `````` DavidWalz committed May 04, 2017 13 `````` """ Return x,y coordinates for triangular array with n^2 stations. """ `````` DavidWalz committed May 04, 2017 14 15 16 17 18 19 `````` n0 = (n - 1) / 2 x, y = np.mgrid[0:n, 0:n].astype(float) - n0 if offset: # offset coordinates x += 0.5 * (y % 2) else: # axial coordinates x += 0.5 * y `````` DavidWalz committed May 04, 2017 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 `````` y *= np.sin(np.pi / 3) return x, y def station_coordinates(n=11, layout='axial', spacing=1500, height=1400): """ Return array of n^2*(x,y,z) coordinates of SD stations for given layout. """ if layout == 'axial': x, y = triangular_array(n, offset=False) elif layout == 'offset': x, y = triangular_array(n, offset=True) elif layout == 'cartesian': x, y = rectangular_array(n) else: raise ValueError('layout must be one of axial, offset, cartesian') x = x.reshape(n**2) * spacing y = y.reshape(n**2) * spacing z = np.ones_like(x) * height return np.c_[x, y, z] `````` DavidWalz committed May 04, 2017 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 `````` def rand_zenith(N=1, zmax=np.pi / 3): """ Sample random zenith angles z for air shower surface detector. Returns zenith angles z sampled pdf f(z) ~ sin(z) cos(z) in range [0, zmax]. """ b = 1 / (1 - np.cos(zmax)**2) r = np.random.rand(N) return np.arccos(np.sqrt(1 - r / b)) def ang2vec(phi, zenith): """ Get 3-vector from spherical angles. Args: phi (array): azimuth (pi, -pi), 0 points in x-direction, pi/2 in y-direction zenith (array): zenith (0, pi), 0 points in z-direction Returns: array of 3-vectors """ x = np.sin(zenith) * np.cos(phi) y = np.sin(zenith) * np.sin(phi) z = np.cos(zenith) return np.array([x, y, z]).T def vec2ang(v): """ Get spherical angles phi and zenith from 3-vector Args: array of 3-vectors Returns: phi, zenith phi (array): azimuth (pi, -pi), 0 points in x-direction, pi/2 in y-direction zenith (array): zenith (0, pi), 0 points in z-direction """ x, y, z = v.T phi = np.arctan2(y, x) zenith = np.pi / 2 - np.arctan2(z, (x**2 + y**2)**.5) return phi, zenith `````` DavidWalz committed May 04, 2017 78 ``````def distance2showerplane(v, vc, va): `````` DavidWalz committed May 04, 2017 79 80 81 82 `````` """ Get shortest distance to shower plane Args: v (Nx3 array): array of positions va (3 array): shower axis = normal vector of the shower plane `````` DavidWalz committed May 04, 2017 83 `````` vc (3 array): shower core `````` DavidWalz committed May 04, 2017 84 `````` """ `````` DavidWalz committed May 04, 2017 85 `````` return np.dot(v - vc, va) `````` DavidWalz committed May 04, 2017 86 87 `````` `````` DavidWalz committed May 04, 2017 88 ``````def distance2showeraxis(v, vc, va): `````` DavidWalz committed May 04, 2017 89 90 91 92 `````` """ Shortest distance to shower axis. Args: v (Nx3 array): array of positions va (3 array): shower axis `````` DavidWalz committed May 04, 2017 93 `````` vc (3 array): shower core `````` DavidWalz committed May 04, 2017 94 `````` """ `````` DavidWalz committed May 04, 2017 95 96 `````` d = distance2showerplane(v, vc, va) vp = v - vc - np.outer(d, va) `````` DavidWalz committed May 04, 2017 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 `````` return np.linalg.norm(vp, axis=-1) def distance2showermaximum(v, vm): """ Shortest distance to shower maximum Args: v (Nx3 array): array of vectors vm (3 array): position of the shower maximum """ return np.linalg.norm(v - vm, axis=-1) def arrival_time_planar(v, vc, va): """ Get arrival times for a planar wavefront. Note: The shower core is not considered here and as it only adds a constant time offset to all stations. Args: v (N x 3 array) vc (3 array): shower core va (3 array): shower axis, pointing upwards Return: array: arrival times [s] """ `````` DavidWalz committed May 04, 2017 119 `````` d = distance2showerplane(v, vc, -va) `````` DavidWalz committed May 04, 2017 120 `` return d / 3E8``