coord.py 22.2 KB
Newer Older
DavidWalz's avatar
DavidWalz committed
1
2
3
"""
Tools for coordinate transformation and vector calculations
"""
4
import numpy as np
DavidWalz's avatar
changes    
DavidWalz committed
5

6
# Rotation matrix for the conversion : x_galactic = R * x_equatorial (J2000)
DavidWalz's avatar
changes    
DavidWalz committed
7
# http://adsabs.harvard.edu/abs/1989A&A...218..325M
8
9
10
_RGE = np.array([[-0.054875539, -0.873437105, -0.483834992],
                 [+0.494109454, -0.444829594, +0.746982249],
                 [-0.867666136, -0.198076390, +0.455983795]])
DavidWalz's avatar
changes    
DavidWalz committed
11

12
# Rotation matrix for the conversion : x_supergalactic = R * x_galactic
DavidWalz's avatar
changes    
DavidWalz committed
13
# http://iopscience.iop.org/0004-637X/530/2/625
14
15
16
_RSG = np.array([[-0.73574257480437488, +0.67726129641389432, 0.0000000000000000],
                 [-0.07455377836523376, -0.08099147130697673, 0.9939225903997749],
                 [+0.67314530210920764, +0.73127116581696450, 0.1100812622247821]])
DavidWalz's avatar
changes    
DavidWalz committed
17

18
19
# Rotation matrix for the conversion: x_equatorial = R * x_ecliptic
# http://en.wikipedia.org/wiki/Ecliptic_coordinate_system
Martin Urban's avatar
Martin Urban committed
20
ECLIPTIC = np.deg2rad(23.44)
21
_REE = np.array([[1., 0., 0.],
Martin Urban's avatar
Martin Urban committed
22
23
                 [0., np.cos(ECLIPTIC), -np.sin(ECLIPTIC)],
                 [0., np.sin(ECLIPTIC), np.cos(ECLIPTIC)]])
24

DavidWalz's avatar
PEP8    
DavidWalz committed
25

26
def eq2gal(v):
DavidWalz's avatar
changes    
DavidWalz committed
27
    """
28
    Rotate equatorial to galactical coordinates (same origin)
marcus's avatar
marcus committed
29

30
    :param v: cartesian vector(s) of shape (3, n) in equatorial coordinate system
31
    :return: vector(s) of shape (3, n) in galactic coordinate system
DavidWalz's avatar
changes    
DavidWalz committed
32
    """
33
    return np.dot(_RGE, np.asarray(v))
DavidWalz's avatar
PEP8    
DavidWalz committed
34

35

36
def gal2eq(v):
DavidWalz's avatar
changes    
DavidWalz committed
37
    """
38
    Rotate galactic to equatorial coordinates (same origin)
marcus's avatar
marcus committed
39

40
    :param v: cartesian vector(s) of shape (3, n) in galactic coordinate system
41
    :return: vector(s) of shape (3, n) in equatorial coordinate system
DavidWalz's avatar
changes    
DavidWalz committed
42
    """
43
    return np.dot(_RGE.transpose(), np.asarray(v))
DavidWalz's avatar
PEP8    
DavidWalz committed
44

45

46
def gal2sgal(v):
47
48
    """
    Rotate galactic to supergalactic coordinates (same origin)
marcus's avatar
marcus committed
49

50
    :param v: cartesian vector(s) of shape (3, n) in galactic coordinate system
51
    :return: vector(s) of shape (3, n) in supergalactic coordinate system
52
    """
53
    return np.dot(_RSG, np.asarray(v))
54

55

56
def sgal2gal(v):
57
58
    """
    Rotate supergalactic to galactic coordinates (same origin)
marcus's avatar
marcus committed
59

60
    :param v: cartesian vector(s) of shape (3, n) in supergalactic coordinate system
61
    :return: vector(s) of shape (3, n) in galactic coordinate system
62
    """
63
    return np.dot(_RSG.transpose(), np.asarray(v))
DavidWalz's avatar
PEP8    
DavidWalz committed
64

65

66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def sgal2eq(v):
    """
    Rotate supergalactic to equatorial coordinates (same origin)

    :param v: cartesian vector(s) of shape (3, n) in supergalactic coordinate system
    :return: vector(s) of shape (3, n) in equatorial coordinate system
    """
    return gal2eq(sgal2gal(v))


def eq2sgal(v):
    """
    Rotate equatorial to supergalactic coordinates (same origin)

    :param v: cartesian vector(s) of shape (3, n) in equatorial coordinate system
    :return: vector(s) of shape (3, n) in supergalactic coordinate system
    """
    return gal2sgal(eq2gal(v))


86
def ecl2eq(v):
87
88
    """
    Rotate ecliptic to equatorial coordinates (same origin)
marcus's avatar
marcus committed
89

90
    :param v: cartesian vector(s) of shape (3, n) in ecliptic coordinate system
91
    :return: vector(s) of shape (3, n) in equatorial coordinate system
92
    """
93
    return np.dot(_REE, np.asarray(v))
DavidWalz's avatar
PEP8    
DavidWalz committed
94

95

96
def eq2ecl(v):
97
98
    """
    Rotate equatorial to ecliptic coordinates (same origin)
marcus's avatar
marcus committed
99

100
    :param v: cartesian vector(s) of shape (3, n) in equatorial coordinate system
101
    :return: vector(s) of shape (3, n) in ecliptic coordinate system
102
    """
103
    return np.dot(_REE.transpose(), np.asarray(v))
104

DavidWalz's avatar
PEP8    
DavidWalz committed
105

106
107
108
def dms2rad(degree, minutes, seconds):
    """
    Transform declination (degree, minute, second) to radians
marcus's avatar
marcus committed
109

110
111
112
    :param degree: integer angle of declination in degree
    :param minutes: arc minute of the declination angle
    :param seconds: arc second of the declination angle
113
    :return: declination angle in radians
114
    """
115
116
    return np.sign(degree) * (np.fabs(degree) + 1. / 60 * minutes + 1. / 3600 * seconds) / 180. * np.pi

117
118
119
120

def hms2rad(hour, minutes, seconds):
    """
    Transform right ascension (hour, minute, second) to radians
marcus's avatar
marcus committed
121

122
123
124
    :param hour: integer hour of right ascension
    :param minutes: arc minute of the right ascension angle
    :param seconds: arc second of the right ascension angle
125
    :return: right ascension angle in radians
126
127
    """
    return (hour / 12. + minutes / 720. + seconds / 43200.) * np.pi
DavidWalz's avatar
changes    
DavidWalz committed
128

DavidWalz's avatar
PEP8    
DavidWalz committed
129

130
def get_hour_angle(ra, lst):
131
    """
132
    Returns the hour angle (in radians) for a specific right ascension and local sidereal time
marcus's avatar
marcus committed
133

134
135
    :param ra: right ascension in radians
    :param lst: local sidereal time in radians
136
    :return: hour angle in radians
137
    """
138
139
140
    return (lst - ra) % (2 * np.pi)


141
def alt2zen(elevation):
142
    """
Marcus Wirtz's avatar
Marcus Wirtz committed
143
    Transforms an elevation angle [radians] in zenith angles
144

Marcus Wirtz's avatar
Marcus Wirtz committed
145
    :param elevation: elevation angle in radians
146
147
    :return: zenith angle in degrees
    """
Marcus Wirtz's avatar
Marcus Wirtz committed
148
    return np.pi / 2. - elevation
149
150


Marcus Wirtz's avatar
Marcus Wirtz committed
151
def date_to_julian_day(year, month, day):
152
    """
Marcus Wirtz's avatar
Marcus Wirtz committed
153
154
    Returns the Julian day number of a date from:
    http://code-highlights.blogspot.de/2013/01/julian-date-in-python.html
155
    http://code.activestate.com/recipes/117215/
156
157
158
159

    :param year: year (measured after christ)
    :param month: month (january: 1, december: 12)
    :param day: day in the month (1 - 31)
160
    """
Marcus Wirtz's avatar
Marcus Wirtz committed
161
162
163
164
    a = (14 - month) // 12
    y = year + 4800 - a
    m = month + 12 * a - 3
    return day + ((153 * m + 2) // 5) + 365 * y + y // 4 - y // 100 + y // 400 - 32045
165
166
167


def get_greenwich_siderial_time(time):
168
169
    """
    Convert civil time to (mean, wrt the mean equinox) Greenwich sidereal time.
170
171
172
    uncertainty of not taking the apparent time (wrt true equinox) is less then 0.01 deg
    time must be a datetime object
    adapted from http://infohost.nmt.edu/tcc/help/lang/python/examples/sidereal/ims/SiderealTime-gst.html
173
174

    :param time: class instance of datetime.date
175
    :return: Greenwich sidereal time
176
177
178
    """
    import datetime
    # [ nDays  :=  number of days between January 0.0 and utc ]
Martin Urban's avatar
Martin Urban committed
179
180
181
    date_ord = time.toordinal()
    jan1_ord = datetime.date(time.year, 1, 1).toordinal()
    n_days = date_ord - jan1_ord + 1
182

Martin Urban's avatar
Martin Urban committed
183
    jan_dt = datetime.datetime(time.year, 1, 1)
Marcus Wirtz's avatar
Marcus Wirtz committed
184
    jan_jd = date_to_julian_day(jan_dt.year, jan_dt.month, jan_dt.day) - 1.5
Martin Urban's avatar
Martin Urban committed
185
    s = jan_jd - 2415020.0
186
187
188
    t = s / 36525.0
    r = (0.00002581 * t + 2400.051262) * t + 6.6460656
    u = r - 24 * (time.year - 1900)
Martin Urban's avatar
Martin Urban committed
189
    factor_b = 24.0 - u
190

Martin Urban's avatar
Martin Urban committed
191
192
193
194
195
    sidereal_a = 0.0657098
    t0 = (n_days * sidereal_a - factor_b)
    dec_utc = time.hour + 1. / 60. * time.minute + 1. / 3600. * (time.second + time.microsecond * 1.e-6)
    sidereal_c = 1.002738
    gst = (dec_utc * sidereal_c + t0) % 24.0
196
197
198
    return gst


199
def get_local_sidereal_time(time, ra):
200
201
202
203
    """
    Convert civil time to local sidereal time

    :param time: class instance of datetime.date
204
    :param ra: right ascension in radians
205
    :return: Local sidereal time (in radians)
206
    """
207
    gst = get_greenwich_siderial_time(time)
208
209
    gst *= 2 * np.pi / 24.
    return (gst + ra) % (2 * np.pi)
210
211


212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
def get_longitude(v, coord_system='gal'):
    """
    Return galactic longitude angle.

    :param v: vector(s) of shape (3, n)
    :param coord_system: coordinate system of the input vectors
    :return: longitude angle (-pi, pi)
    """
    if coord_system != 'gal':
        v = globals()['%s2gal' % coord_system](v)
    return vec2ang(v)[0]


def get_latitude(v, coord_system='gal'):
    """
    Return galactic latitude angle.

    :param v: vector(s) of shape (3, n)
    :param coord_system: coordinate system of the input vectors
    :return: latitude angle (-pi/2, pi/2)
    """
    if coord_system != 'gal':
        v = globals()['%s2gal' % coord_system](v)
    return vec2ang(v)[1]


def get_right_ascension(v, coord_system='gal'):
    """
    Return equatorial right ascension angle.

    :param v: vector(s) of shape (3, n)
    :param coord_system: coordinate system of the input vectors
    :return: declination angle (-pi, pi)
    """
    if coord_system != 'eq':
        v = globals()['%s2eq' % coord_system](v)
    return vec2ang(v)[0]


def get_declination(v, coord_system='gal'):
    """
    Return equatorial declination angle.

    :param v: vector(s) of shape (3, n)
    :param coord_system: coordinate system of the input vectors
    :return: declination angle (-pi/2, pi/2)
    """
    if coord_system != 'eq':
        v = globals()['%s2eq' % coord_system](v)
    return vec2ang(v)[1]


def get_exposure(v, coord_system='gal', **kwargs):
    """
    Returns exposure values of direction.

    :param v: vector(s) of shape (3, n)
    :param coord_system: coordinate system of the input vectors
    :param kwargs: Additionally named keyword arguments passed to exposure_equatorial()
    :return: exposure values
    """
    decs = get_declination(v, coord_system=coord_system)
    return exposure_equatorial(decs, **kwargs)


277
278
279
280
281
282
283
284
def atleast_kd(array, k):
    """
    Extends numpy.atleast_1d() to arbitrary number of dimensions.
    """
    array = np.asarray(array)
    return array.reshape(array.shape + (1,) * (k - array.ndim))


285
286
287
def normed(v):
    """
    Return the normalized (lists of) vectors.
marcus's avatar
marcus committed
288

289
    :param v: vector(s) of shape (3, n)
290
    :return: corresponding normalized vectors of shape (3, n)
291
292
293
    """
    return np.asarray(v) / np.linalg.norm(v, axis=0)

294

295
296
297
def distance(v1, v2):
    """
    Linear distance between each pair from two (lists of) vectors.
marcus's avatar
marcus committed
298

299
300
    :param v1: vector(s) of shape (3, n)
    :param v2: vector(s) of shape (3, n)
301
    :return: linear distance between each pair
302
303
304
    """
    return np.linalg.norm(np.asarray(v1) - np.asarray(v2), axis=0)

305

306
307
def angle(v1, v2, each2each=False):
    """
308
    Angular distance in radians for each pair from two (lists of) vectors.
309
    Use each2each=True to calculate every combination.
marcus's avatar
marcus committed
310

311
312
313
    :param v1: vector(s) of shape (3, n)
    :param v2: vector(s) of shape (3, n)
    :param each2each: if true calculates every combination of the two lists v1, v2
314
    :return: angular distance in radians
DavidWalz's avatar
changes    
DavidWalz committed
315
    """
316
317
318
319
320
    a = normed(v1)
    b = normed(v2)
    if each2each:
        d = np.outer(a[0], b[0]) + np.outer(a[1], b[1]) + np.outer(a[2], b[2])
    else:
DavidWalz's avatar
DavidWalz committed
321
        if len(a.shape) == 1:
322
            a = a.reshape(3, 1)
DavidWalz's avatar
DavidWalz committed
323
        if len(b.shape) == 1:
324
            b = b.reshape(3, 1)
325
326
327
        d = np.sum(a * b, axis=0)
    return np.arccos(np.clip(d, -1., 1.))

328

329
330
331
def vec2ang(v):
    """
    Get spherical angles (phi, theta) from a (list of) vector(s).
marcus's avatar
marcus committed
332

333
    :param v: vector(s) of shape (3, n)
334
    :return: tuple consisting of
marcus's avatar
marcus committed
335
336
337

             - phi [range (pi, -pi), 0 points in x-direction, pi/2 in y-direction],
             - theta [range (pi/2, -pi/2), pi/2 points in z-direction]
DavidWalz's avatar
changes    
DavidWalz committed
338
    """
339
    x, y, z = np.asarray(v)
340
    phi = np.arctan2(y, x)
341
    theta = np.arctan2(z, (x * x + y * y) ** .5)
342
    return phi, theta
DavidWalz's avatar
changes    
DavidWalz committed
343

344

345
def ang2vec(phi, theta):
DavidWalz's avatar
changes    
DavidWalz committed
346
    """
347
    Get vector from spherical angles (phi, theta)
marcus's avatar
marcus committed
348

349
    :param phi: range (pi, -pi), 0 points in x-direction, pi/2 in y-direction
350
    :param theta: range (pi/2, -pi/2), pi/2 points in z-direction
351
    :return: vector of shape (3, n)
DavidWalz's avatar
changes    
DavidWalz committed
352
    """
353
    assert np.ndim(phi) == np.ndim(theta), "Inputs phi and theta in 'coord.ang2vec()' must have same shape!"
354
355
356
    x = np.cos(theta) * np.cos(phi)
    y = np.cos(theta) * np.sin(phi)
    z = np.sin(theta)
357
    return np.array([x, y, z])
DavidWalz's avatar
PEP8    
DavidWalz committed
358

359

360
def sph_unit_vectors(phi, theta):
DavidWalz's avatar
DavidWalz committed
361
    """
362
    Get spherical unit vectors e_r, e_phi, e_theta from spherical angles phi theta.
marcus's avatar
marcus committed
363

364
365
    :param phi: range (pi, -pi), 0 points in x-direction, pi/2 in y-direction, float or array_like
    :param theta: range (pi/2, -pi/2), pi/2 points in z-direction, float or array_like
366
    :return: shperical unit vectors e_r, e_phi, e_theta; each with shape (3, N)
DavidWalz's avatar
DavidWalz committed
367
    """
368
    assert np.ndim(phi) == np.ndim(theta), "Inputs phi and theta in 'coord.sph_unit_vectors()' must have same shape!"
369
370
    sp, cp = np.sin(phi), np.cos(phi)
    st, ct = np.sin(theta), np.cos(theta)
DavidWalz's avatar
PEP8    
DavidWalz committed
371

372
373
    e_r = np.array([ct * cp, ct * sp, st])
    e_phi = np.array([-sp, cp, np.zeros_like(phi)])
374
375
    e_theta = np.array([-st * cp, -st * sp, ct])
    return np.array([e_r, e_phi, e_theta])
376

DavidWalz's avatar
PEP8    
DavidWalz committed
377

marcus's avatar
marcus committed
378
def rotation_matrix(rotation_axis, rotation_angle):
DavidWalz's avatar
changes    
DavidWalz committed
379
    """
380
381
    Rotation matrix for given rotation axis and angle.
    See http://en.wikipedia.org/wiki/Euler-Rodrigues_parameters
marcus's avatar
marcus committed
382

383
384
385
    :param rotation_axis: rotation axis, either np.array([x, y, z]) or ndarray with shape (3, n)
    :param rotation_angle: rotation angle in radians, either float or array size n
    :return: rotation matrix R, either shape (3, 3) or (3, 3, n)
DavidWalz's avatar
PEP8    
DavidWalz committed
386

387
    Example:
388
    R = rotation_matrix( np.array([4,4,1]), 1.2 )
389
390
    v1 = np.array([3,5,0])
    v2 = np.dot(R, v1)
DavidWalz's avatar
DavidWalz committed
391
    """
392
    assert np.ndim(rotation_axis) > np.ndim(rotation_angle), "Shape of rotation axis and angle do not not match"
marcus's avatar
marcus committed
393
394
395
    rotation_axis = normed(rotation_axis)
    a = np.cos(rotation_angle / 2.)
    b, c, d = - rotation_axis * np.sin(rotation_angle / 2.)
Martin Urban's avatar
Martin Urban committed
396
    r = np.array([[a * a + b * b - c * c - d * d, 2 * (b * c - a * d), 2 * (b * d + a * c)],
397
398
                  [2 * (b * c + a * d), a * a + c * c - b * b - d * d, 2 * (c * d - a * b)],
                  [2 * (b * d - a * c), 2 * (c * d + a * b), a * a + d * d - b * b - c * c]])
marcus's avatar
marcus committed
399

Martin Urban's avatar
Martin Urban committed
400
    return np.squeeze(r)
401

402

marcus's avatar
marcus committed
403
def rotate(v, rotation_axis, rotation_angle):
DavidWalz's avatar
DavidWalz committed
404
    """
405
    Perform rotation for given rotation axis and angle(s).
marcus's avatar
marcus committed
406

407
408
409
410
    :param v: vector that is supposed to be rotated, either np.array([x, y, z]) or ndarray with shape (3, n)
    :param rotation_axis: rotation axis, either np.array([x, y, z]) or ndarray with shape (3, n)
    :param rotation_angle: rotation angle in radians, either float or array size n
    :return: rotated vector, same shape as input v
DavidWalz's avatar
DavidWalz committed
411
    """
412
    shape = v.shape
413
    v, rotation_axis, rotation_angle = np.squeeze(v), np.squeeze(rotation_axis), np.squeeze(rotation_angle)
414
    rotation_axis = atleast_kd(rotation_axis, k=max(np.ndim(v), np.ndim(rotation_angle)+1))
Martin Urban's avatar
Martin Urban committed
415
    r = rotation_matrix(rotation_axis, rotation_angle)
416
    if rotation_axis.ndim == 1:
417
        return np.dot(r, v).reshape(shape)
418

419
    rotated_vector = np.sum(atleast_kd(r, k=v.ndim+1) * atleast_kd(v[np.newaxis], r.ndim), axis=1)
420
421
422
    if rotated_vector.size == v.size:
        rotated_vector = rotated_vector.reshape(shape)
    return rotated_vector
DavidWalz's avatar
DavidWalz committed
423

DavidWalz's avatar
PEP8    
DavidWalz committed
424

425
426
427
428
429
430
431
432
433
434
435
def rotate_zaxis_to_x(v, x0):
    """
    Transfers the relative orientation between vectors v and the z-axis towards
    v and the reference vectors x0. Mathematically, the scalar products z_axis*v
    before the rotation and x0*v after rotation are the same (see e.g. unit test).

    :param v: vectors that should be rotated, shape: (3) or (3, n)
    :param x0: reference vectors of same shape like v
    """
    # defines rotation axis by the cross-product with z-axis
    u = np.array([x0[1], -x0[0], np.zeros_like(x0[0])])
436
    u[0, np.sum(u**2, axis=0) < 1e-10] = 1      # chose x-axis in case of z-axis for x0
437
    angles = angle(x0, atleast_kd(np.array([0, 0, 1]), k=np.ndim(x0)))
438
439
440
    return rotate(v, normed(u), angles)


441
def exposure_equatorial(dec, a0=-35.25, zmax=60):
442
    """
443
    Relative exposure per solid angle of a detector at latitude a0 (-90, 90 degrees, default: Auger),
marcus's avatar
marcus committed
444
    with maximum acceptance zenith angle zmax (0, 90 degrees, default: 60) and for given equatorial declination
445
    dec (-pi/2, pi/2) or vectoe. See astro-ph/0004016
marcus's avatar
marcus committed
446

447
    :param vec_or_dec: value(s) of declination in radians (-pi/2, pi/2)
448
449
    :param a0: latitude of detector (-90, 90) degrees (default: Auger)
    :param zmax: maximum acceptance zenith angle (0, 90) degrees
450
451
    :return: exposure value(s) for input declination value(s)
    """
452
    # noinspection PyTypeChecker,PyUnresolvedReferences
453
    if (abs(dec) > np.pi / 2).any():
454
        raise Exception('exposure_equatorial: declination not in range (-pi/2, pi/2)')
455
    if (zmax < 0) or (zmax > 90):
456
        raise Exception('exposure_equatorial: zmax not in range (0, 90 degrees)')
457
    if (a0 < -90) or (a0 > 90):
458
        raise Exception('exposure_equatorial: a0 not in range (-90, 90 degrees)')
459

460
461
    zmax *= np.pi / 180
    a0 *= np.pi / 180
462
463

    xi = (np.cos(zmax) - np.sin(a0) * np.sin(dec)) / np.cos(a0) / np.cos(dec)
464
    am = np.arccos(np.clip(xi, -1, 1))
465
466

    cov = np.cos(a0) * np.cos(dec) * np.sin(am) + am * np.sin(a0) * np.sin(dec)
DavidWalz's avatar
PEP8    
DavidWalz committed
467
468
    return cov / np.pi  # normalize the maximum possible value to 1

469

470
def rand_fisher(kappa, n=1, seed=None):
471
472
473
    """
    Random angles to the center of a Fisher distribution with concentration parameter kappa.

474
475
    :param kappa: concentration parameter, kappa=1/sigma^2 (sigma: smearing angle in radians)
                  either float, or np.array (n is set to kappa.shape[0] then)
476
477
478
    :param n: number of vectors drawn from fisher distribution
    :return: theta values (angle towards the mean direction)
    """
479
480
    if seed is not None:
        np.random.seed(seed)
481
    if np.ndim(kappa) > 0:
482
        n = kappa.shape
483
    return np.arccos(1 + np.log(1 - np.random.random(n) * (1 - np.exp(-2 * kappa))) / kappa)
484
485


486
def rand_phi(n=1, seed=None):
487
488
    """
    Random uniform phi (-pi, pi).
marcus's avatar
marcus committed
489

490
    :param n: number of samples that are drawn
491
    :return: random numbers in range (-pi, pi)
492
    """
493
494
    if seed is not None:
        np.random.seed(seed)
495
    return (np.random.random(n) * 2 - 1) * np.pi
496

497

498
def rand_theta(n=1, theta_min=-np.pi/2, theta_max=np.pi/2, seed=None):
499
500
    """
    Random theta (pi/2, -pi/2) from uniform cos(theta) distribution.
marcus's avatar
marcus committed
501

502
    :param n: number of samples that are drawn
503
    :return: random theta in range (-pi/2, pi/2) from cos(theta) distribution
504
    """
505
506
    if seed is not None:
        np.random.seed(seed)
507
    assert theta_max > theta_min
508
    u = np.sin(theta_min) + (np.sin(theta_max) - np.sin(theta_min)) * np.random.random(n)
509
    return np.pi / 2 - np.arccos(u)
510

511

512
def rand_theta_plane(n=1, seed=None):
513
514
515
516
517
518
    """
    Random theta (pi/2, 0) on a planar surface from uniform cos(theta)^2 distribution.

    :param n: number of samples that are drawn
    :return: random theta on plane in range (pi/2, 0)
    """
519
520
    if seed is not None:
        np.random.seed(seed)
521
    return np.pi / 2 - np.arccos(np.sqrt(np.random.random(n)))
522
523


524
def rand_vec(n=1, seed=None):
DavidWalz's avatar
DavidWalz committed
525
526
    """
    Random spherical unit vectors.
marcus's avatar
marcus committed
527

528
    :param n: number of vectors that are drawn
529
    :return: random unit vectors of shape (3, n)
DavidWalz's avatar
DavidWalz committed
530
    """
531
532
533
    if seed is not None:
        np.random.seed(seed)
    return ang2vec(rand_phi(n, seed=seed), rand_theta(n, seed=seed))
DavidWalz's avatar
DavidWalz committed
534

535

536
def rand_vec_on_surface(x0, n=1, seed=None):
537
538
539
540
541
542
543
    """
    Given unit normal vectors x0 orthogonal on a surface, samples one isotropic
    direction for each given vector x0 from a cos(theta)*sin(theta) distribution

    :param x0: ortogonal unit vector on the surface, shape: (3, N)
    :return: isotropic directions for the respective normal vectors x0
    """
544
545
    if seed is not None:
        np.random.seed(seed)
546
547
    if np.ndim(np.squeeze(x0)) > 1:
        n = x0.shape[1]
548
    v = ang2vec(rand_phi(n, seed=seed), rand_theta_plane(n, seed=seed))   # produce random vecs on plane through z-axis
549
    return rotate_zaxis_to_x(v, x0)                 # rotation to respective surface vector x0
550
551


552
def rand_vec_on_sphere(n, r=1, seed=None):
553
554
555
556
557
558
559
560
    """
    Calculates n random positions x and correpsonding directions on the surface of the
    sphere as they would arise from an isotropic distribution of cosmic rays in the universe.

    :param n: number of drawn samples, int
    :param r: size of the sphere, which scales the position vectors, default: 1
    :return: position vectors, directions
    """
561
562
563
564
    if seed is not None:
        np.random.seed(seed)
    x = rand_vec(n, seed=seed)
    v = rand_vec_on_surface(x, seed=seed)
565
566
567
    return r*x, v


568
def rand_exposure_vec(a0=-35.25, zmax=60, n=1, coord_system='gal', res_theta=5000, seed=None):
569
    """
570
    Random vecs following the exposure of an experiment (equatorial coordinates).
571
    If you need other coordinates, change 'coord_system' keyword to eg. 'eq' or 'sgal'. This
572
573
    method bins theta and samples from corresponding probabilities as the
    corresponding probability function is not integratable and invertable.
marcus's avatar
marcus committed
574

575
576
577
    :param a0: latitude of detector (-90, 90) in degrees (default: Auger)
    :param zmax: maximum acceptance zenith angle (0, 90) degrees
    :param n: number of samples that are drawn
578
    :param coord_system: coordinate system of the returned vectors (default: 'gal')
579
580
581
    :param res_theta: resolution of theta, number of bins in (-pi/2, pi/2)
    :return: random unit vectors from exposure of shape (3, n), equatorial coordinates
    """
582
583
    if seed is not None:
        np.random.seed(seed)
584
585
586
587
    eps = 1. / res_theta
    theta_bin = np.linspace(-np.pi/2.+eps, np.pi/2.-eps, num=res_theta)
    p = np.cos(theta_bin) * exposure_equatorial(theta_bin, a0, zmax)
    thetas = np.random.choice(theta_bin, n, p=p/p.sum()) + np.random.uniform(-eps, eps, n)
588
    v = ang2vec(rand_phi(n, seed=seed), thetas)
589
590
591
592
593

    if coord_system != 'eq':
        v = globals()['eq2%s' % coord_system](v)

    return v
594

595

596
def rand_fisher_vec(vmean, kappa, n=1, seed=None):
597
598
    """
    Random Fisher distributed vectors with mean direction vmean and concentration parameter kappa.
marcus's avatar
marcus committed
599

600
    :param vmean: mean direction of the fisher distribution, (x, y, z), either shape (3) or (3, n)
601
    :param kappa: concentration parameter, translates to 1/sigma^2 (sigma: smearing angle in radians)
602
    :param n: number of vectors drawn from fisher distribution, becomes m if vmean has shape (3, m)
603
    :return: vectors from fisher distribution of shape (3, n)
604
    """
605
606
    if seed is not None:
        np.random.seed(seed)
607
    vmean = atleast_kd(vmean, k=np.ndim(kappa)+1)
608
    if np.ndim(kappa) > 0:
609
        n = kappa.shape
610
611
    elif np.ndim(vmean) > 1:
        n = vmean.shape[1:]
612
    # create random fisher distributed directions around z-axis (0, 0, 1)
613
614
    theta = np.pi / 2 - rand_fisher(kappa, n, seed=seed)
    phi = rand_phi(theta.shape, seed=seed)
615
616
    # rotate these to reference vector vmean
    return rotate_zaxis_to_x(ang2vec(phi, theta), vmean)
617
618


619
def equatorial_scrambling(v, n=1, coord_system='gal', seed=None):
620
621
622
623
624
625
626
627
628
    """
    Performs a scrambling of the input vectors v around the equatorial z-axis. In this sense it keeps
    the declination while assigning new uniform azimuth angles.

    :param v: Input vectors in equatoial coordinates and shape (3, ncrs)
    :param n: Number of scrambled output sets
    :param coord_system: coordinate system for input vectors
    :return: scrambled vectors in shape (3, n, ncrs)
    """
629
630
    if seed is not None:
        np.random.seed(seed)
631
    decs = get_declination(v, coord_system)
632
    v_scrambled = ang2vec(rand_phi(v.shape[1] * n, seed=seed), np.tile(decs, n))
633
634
635
636

    if coord_system != 'eq':
        v_scrambled = globals()['eq2%s' % coord_system](v_scrambled)

637
    return np.reshape(v_scrambled, (3, n, -1)) if n > 1 else np.reshape(v_scrambled, (3, -1))