utils.py 3.5 KB
Newer Older
DavidWalz's avatar
DavidWalz committed
1 2 3 4 5 6
from __future__ import division

import numpy as np


def rectangular_array(n=11):
DavidWalz's avatar
DavidWalz committed
7
    """ Return x,y coordinates for rectangular array with n^2 stations. """
DavidWalz's avatar
DavidWalz committed
8
    n0 = (n - 1) / 2
DavidWalz's avatar
DavidWalz committed
9
    return (np.mgrid[0:n, 0:n].astype(float) - n0)
DavidWalz's avatar
DavidWalz committed
10 11 12


def triangular_array(n=11, offset=True):
DavidWalz's avatar
DavidWalz committed
13
    """ Return x,y coordinates for triangular array with n^2 stations. """
DavidWalz's avatar
DavidWalz committed
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's avatar
DavidWalz committed
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's avatar
DavidWalz committed
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's avatar
DavidWalz committed
78
def distance2showerplane(v, vc, va):
DavidWalz's avatar
DavidWalz committed
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's avatar
DavidWalz committed
83
        vc (3 array): shower core
DavidWalz's avatar
DavidWalz committed
84
    """
DavidWalz's avatar
DavidWalz committed
85
    return np.dot(v - vc, va)
DavidWalz's avatar
DavidWalz committed
86 87


DavidWalz's avatar
DavidWalz committed
88
def distance2showeraxis(v, vc, va):
DavidWalz's avatar
DavidWalz committed
89 90 91 92
    """ Shortest distance to shower axis.
    Args:
        v (Nx3 array): array of positions
        va (3 array): shower axis
DavidWalz's avatar
DavidWalz committed
93
        vc (3 array): shower core
DavidWalz's avatar
DavidWalz committed
94
    """
DavidWalz's avatar
DavidWalz committed
95 96
    d = distance2showerplane(v, vc, va)
    vp = v - vc - np.outer(d, va)
DavidWalz's avatar
DavidWalz committed
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's avatar
DavidWalz committed
119
    d = distance2showerplane(v, vc, -va)
DavidWalz's avatar
DavidWalz committed
120
    return d / 3E8