Completed
Push — master ( f94699...de38d6 )
by Ben
01:27
created

Voids.rotation_y()   A

Complexity

Conditions 2

Size

Total Lines 6

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 2
c 1
b 0
f 0
dl 0
loc 6
rs 9.4285
1
"Free electron density model"
2
import os
3
4
import numpy as np
5
from astropy.table import Table
6
from numpy import cos
7
from numpy import cosh
8
from numpy import exp
9
from numpy import pi
10
from numpy import sin
11
from numpy import sqrt
12
from numpy import tan
13
14
15
# import astropy.units as us
16
# from astropy.coordinates import SkyCoord
17
18
# Configuration
19
# TODO: use to config file
20
# input parameters for large-scale components of NE2001 30 June '02
21
# flags = {'wg1': 1,
22
#          'wg2': 1,
23
#          'wga': 1,
24
#          'wggc': 1,
25
#          'wglism': 1,
26
#          'wgcN': 1,
27
#          'wgvN': 1}
28
29
# solar_params = {'Rsun': 8.3}
30
31
# thick_disk_params = {'n1h1': 0.033,
32
#                      'h1': 0.97,
33
#                      'A1': 17.5,
34
#                      'F1': 0.18}
35
36
# thin_disk_params = {'n2': 0.08,
37
#                     'h2': 0.15,
38
#                     'A2': 3.8,
39
#                     'F2': 120}
40
41
# galactic_center_params = {'xgc': -0.01,
42
#                           'ygc': 0.0,
43
#                           'zgc': -0.020,
44
#                           'rgc': 0.145,
45
#                           'hgc': 0.026,
46
#                           'negc0': 10.0,
47
#                           'Fgc0': 0.6e5}
48
49
50
# spiral_arms_params = {'na': 0.028,
51
#                       'ha': 0.23,
52
#                       'wa': 0.65,
53
#                       'Aa': 10.5,
54
#                       'Fa': 5,
55
#                       'narm1': 0.5,
56
#                       'narm2': 1.2,
57
#                       'narm3': 1.3,
58
#                       'narm4': 1.0,
59
#                       'narm5': 0.25,
60
#                       'warm1': 1.0,
61
#                       'warm2': 1.5,
62
#                       'warm3': 1.0,
63
#                       'warm4': 0.8,
64
#                       'warm5': 1.0,
65
#                       'harm1': 1.0,
66
#                       'harm2': 0.8,
67
#                       'harm3': 1.3,
68
#                       'harm4': 1.5,
69
#                       'harm5': 1.0,
70
#                       'farm1': 1.1,
71
#                       'farm2': 0.3,
72
#                       'farm3': 0.4,
73
#                       'farm4': 1.5,
74
#                       'farm5': 0.3}
75
76
77
ldr_params = {'abc': np.array([1.50, .750, .50]),
78
              'center': np.array([1.36, 8.06, 0.0]),
79
              'theta': -24.2*pi/180,
80
              'ne': 0.012,
81
              'F': 0.1}
82
83
lsb_params = {'abc': np.array([1.050, .4250, .3250]),
84
              'center': np.array([-0.75, 9.0, -0.05]),
85
              'theta': 139.*pi/180,
86
              'ne': 0.016,
87
              'F':  0.01}
88
89
lhb_params = {'abc': np.array([.0850, .1000, .330]),
90
              'center': np.array([0.01, 8.45, 0.17]),
91
              'theta': 15*pi/180,
92
              'ne': 0.005,
93
              'F': 0.01}
94
95
loop_params = {'center': np.array([-0.045, 8.40, 0.07]),
96
               'r': 0.120,
97
               'dr': 0.060,
98
               'ne1': 0.0125,
99
               'ne2': 0.0125,
100
               'F1': 0.2,
101
               'F2': 0.01}
102
103
104
def ne_thick_disk(xyz, ne_disk, rdisk, hdisk, r_sun):
105
    """
106
    Calculate the contribution of the thick disk to the free electron density
107
     at x, y, z = `xyz`
108
    """
109
    r2d = sqrt(xyz[0]**2 + xyz[1]**2)
110
    k = pi/2/rdisk
111
    return ne_disk * (cos(r2d*k)/cos(r_sun*k) /
112
                      cosh(xyz[-1]/hdisk)**2 *
113
                      (r2d < rdisk))
114
115
116
def ne_thin_disk(xyz, ne_disk, rdisk, hdisk):
117
    """
118
    Calculate the contribution of the thin disk to the free electron density
119
     at x, y, z = `xyz`
120
    """
121
    r2d = sqrt(xyz[0]**2 + xyz[1]**2)
122
    return ne_disk * (exp(-(r2d - rdisk)**2/1.8**2) /
123
                      cosh(xyz[-1]/hdisk)**2)  # Why 1.8?
124
125
126
def ne_gc(xyz, ne_gc0, rgc, hgc, xyz_gc):
127
    """
128
    Calculate the contribution of the Galactic center to the free
129
    electron density at x, y, z = `xyz`
130
    """
131
    # Here I'm using the expression in the NE2001 code which is inconsistent
132
    # with Cordes and Lazio 2011 (0207156v3) (See Table 2)
133
    try:
134
        xyz = xyz - xyz_gc
135
    except ValueError:
136
        xyz = xyz - xyz_gc[:, None]
137
138
    r2d = sqrt(xyz[0]**2 + xyz[1]**2)
139
140
    # ????
141
    # Cordes and Lazio 2011 (0207156v3) (Table 2)
142
    # return ne_gc0*exp(-(r2d/rgc)**2 - (xyz[-1]/hgc)**2)
143
    # ????
144
145
    # Constant ne (form NE2001 code)
146
    return ne_gc0*((r2d/rgc)**2 + (xyz[-1]/hgc)**2 < 1)*(r2d < rgc)
147
148
149
def ne_local_ism(xyz, ldr_params, lsb_params, lhb_params, loop_params):
150
    """
151
    Calculate the contribution of the local ISM to the free
152
    electron density at x, y, z = `xyz`
153
    """
154
    # low density region in Q1
155
    neldr = ldr_params['ne']*in_ellisoid(xyz, ldr_params['center'],
156
                                         ldr_params['abc'],
157
                                         ldr_params['theta'])
158
    # Local Super Bubble
159
    nelsb = lsb_params['ne']*in_ellisoid(xyz, lsb_params['center'],
160
                                         lsb_params['abc'],
161
                                         lsb_params['theta'])
162
163
    # Local Hot Bubble
164
    nelhb = lhb_params['ne']*in_cylinder(xyz, lhb_params['center'],
165
                                         lhb_params['abc'],
166
                                         lhb_params['theta'])
167
    # Loop I
168
    irr1 = in_half_sphere(xyz, loop_params['center'], loop_params['r'])
169
    irr2 = in_half_sphere(xyz, loop_params['center'],
170
                          loop_params['r'] + loop_params['dr'])
171
    neloop = loop_params['ne1'] * irr1 + loop_params['ne2'] * irr2*(~irr1)
172
    wlhb, wloop, wlsb, wldr = (nelhb > 0,
173
                               neloop > 0,
174
                               nelsb > 0,
175
                               neldr > 0)
176
    ne_lism = ((1 - wlhb) *
177
               ((1 - wloop) * (wlsb*nelsb + (1-wlsb) * neldr) +
178
                wloop*neloop) +
179
               wlhb*nelhb)
180
181
    wlism = np.maximum(wloop, np.maximum(wldr, np.maximum(wlsb, wlhb)))
182
    return ne_lism, wlism
183
184
185
def in_ellisoid(xyz, xyz_center, abc_ellipsoid, theta):
186
    """
187
    Test if xyz in the ellipsoid
188
    Theta in radians
189
    """
190
    try:
191
        xyz = xyz - xyz_center
192
    except ValueError:
193
        xyz = xyz - xyz_center[:, None]
194
        abc_ellipsoid = abc_ellipsoid[:, None]
195
196
    rot = rotation(theta, -1)
197
    xyz = rot.dot(xyz)
198
199
    xyz_p = xyz/abc_ellipsoid
200
201
    return np.sum(xyz_p**2, axis=0) <= 1
202
203
204
def in_cylinder(xyz, xyz_center, abc_cylinder, theta):
205
    """
206
    Test if xyz in the cylinder
207
    Theta in radians
208
    """
209
    try:
210
        xyz = xyz - xyz_center
211
    except ValueError:
212
        xyz = xyz - xyz_center[:, None]
213
        abc_cylinder = np.vstack([abc_cylinder]*xyz.shape[-1]).T
214
    xyz[2] -= tan(theta)*xyz[-1]
215
216
    abc_cylinder_p = abc_cylinder.copy()
217
    z_c = (xyz_center[-1] - abc_cylinder[-1])
218
    izz = (xyz[-1] <= 0)*(xyz[-1] <= z_c)
219
    abc_cylinder_p[0] = (0.001 +
220
                         (abc_cylinder[0] - 0.001) *
221
                         (1 - xyz[-1]/z_c))*izz + abc_cylinder[0]*(~izz)
222
223
    xyz_p = xyz/abc_cylinder_p
224
225
    return (xyz_p[0]**2 + xyz_p[1]**2 <= 1) * (xyz_p[-1]**2 <= 1)
226
227
228
def in_half_sphere(xyz, xyz_center, r_sphere):
229
    "Test if `xyz` in the sphere with radius r_sphere  centerd at `xyz_center`"
230
    try:
231
        xyz = xyz - xyz_center
232
    except ValueError:
233
        xyz = xyz - xyz_center[:, None]
234
    distance = sqrt(np.sum(xyz**2, axis=0))
235
    return (distance <= r_sphere)*(xyz[-1] >= 0)
236
237
238
class Clumps(object):
239
    """
240
    """
241
242
    def __init__(self, clumps_file=None):
243
        """
244
        """
245
        if not clumps_file:
246
            this_dir, _ = os.path.split(__file__)
247
            clumps_file = os.path.join(this_dir, "data", "neclumpN.NE2001.dat")
248
        self._data = Table.read(clumps_file, format='ascii')
249
250
    @property
251
    def use_clump(self):
252
        """
253
        """
254
        return self._data['flag'] == 0
255
256
    @property
257
    def xyz(self):
258
        """
259
        """
260
        try:
261
            return self._xyz
262
        except AttributeError:
263
            self._xyz = self.get_xyz()
264
        return self._xyz
265
266
    @property
267
    def gl(self):
268
        """
269
        Galactic longitude (deg)
270
        """
271
        return self._data['l']
272
273
    @property
274
    def gb(self):
275
        """
276
        Galactic latitude (deg)
277
        """
278
        return self._data['b']
279
280
    @property
281
    def distance(self):
282
        """
283
        Distance from the sun (kpc)
284
        """
285
        return self._data['dc']
286
287
    @property
288
    def radius(self):
289
        """
290
        Radius of the clump (kpc)
291
        """
292
        return self._data['rc']
293
294
    @property
295
    def ne(self):
296
        """
297
        Electron density of each clump (cm^{-3})
298
        """
299
        return self._data['nec']
300
301
    @property
302
    def edge(self):
303
        """
304
        Clump edge
305
        0 => use exponential rolloff out to 5 clump radii
306
        1 => uniform and truncated at 1/e clump radius
307
        """
308
        return self._data['edge']
309
310
    def get_xyz(self, rsun=8.5):
311
        """
312
        """
313
        # xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
314
        #                distance=self.distance,
315
        #                z_sun = z_sun*us.kpc,
316
        #                unit="deg, deg, kpc").galactocentric.cartesian.xyz.value
317
        # return xyz
318
319
        slc = sin(self.gl/180*pi)
320
        clc = cos(self.gl/180*pi)
321
        sbc = sin(self.gb/180*pi)
322
        cbc = cos(self.gb/180*pi)
323
        rgalc = self.distance*cbc
324
        xc = rgalc*slc
325
        yc = rsun-rgalc*clc
326
        zc = self.distance*sbc
327
        return np.array([xc, yc, zc])
328
329
    def clump_factor(self, xyz):
330
        """
331
        Clump edge
332
        0 => use exponential rolloff out to 5 clump radii
333
        1 => uniform and truncated at 1/e clump radius
334
        """
335
        if xyz.ndim == 1:
336
            xyz = xyz[:, None] - self.xyz
337
        else:
338
            xyz = xyz[:, :, None] - self.xyz[:, None, :]
339
340
        q2 = (np.sum(xyz**2, axis=0) /
341
              self.radius**2)
342
        # NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
343
        # TODO: check this
344
        return (q2 <= 1)*(self.edge == 1) + (q2 <= 5)*(self.edge == 0)*exp(-q2)
345
346
    def ne_clumps(self, xyz):
347
        """
348
        The contribution of the clumps to the free
349
        electron density at x, y, z = `xyz`
350
        """
351
        return np.sum(self.clump_factor(xyz)*self.ne*self.use_clump, axis=-1)
352
353
354
class Voids(object):
355
    """
356
    """
357
358
    def __init__(self, voids_file=None):
359
        """
360
        """
361
        if not voids_file:
362
            this_dir, _ = os.path.split(__file__)
363
            voids_file = os.path.join(this_dir, "data", "nevoidN.NE2001.dat")
364
        self._data = Table.read(voids_file, format='ascii')
365
366
    @property
367
    def use_void(self):
368
        """
369
        """
370
        return self._data['flag'] == 0
371
372
    @property
373
    def xyz(self):
374
        """
375
        """
376
        try:
377
            return self._xyz
378
        except AttributeError:
379
            self._xyz = self.get_xyz()
380
        return self._xyz
381
382
    @property
383
    def gl(self):
384
        """
385
        Galactic longitude (deg)
386
        """
387
        return self._data['l']
388
389
    @property
390
    def gb(self):
391
        """
392
        Galactic latitude (deg)
393
        """
394
        return self._data['b']
395
396
    @property
397
    def distance(self):
398
        """
399
        Distance from the sun (kpc)
400
        """
401
        return self._data['dv']
402
403
    @property
404
    def ellipsoid_abc(self):
405
        """
406
        Void axis
407
        """
408
        return np.array([self._data['aav'],
409
                         self._data['bbv'],
410
                         self._data['ccv']])
411
412
    @property
413
    def rotation_y(self):
414
        """
415
        Rotation around the y axis
416
        """
417
        return [rotation(theta*pi/180, 1) for theta in self._data['thvy']]
418
419
    @property
420
    def rotation_z(self):
421
        """
422
        Rotation around the z axis
423
        """
424
        return [rotation(theta*pi/180, -1) for theta in self._data['thvz']]
425
426
    @property
427
    def ne(self):
428
        """
429
        Electron density of each void (cm^{-3})
430
        """
431
        return self._data['nev']
432
433
    @property
434
    def edge(self):
435
        """
436
        Void edge
437
        0 => use exponential rolloff out to 5 clump radii
438
        1 => uniform and truncated at 1/e clump radius
439
        """
440
        return self._data['edge']
441
442
    def get_xyz(self, rsun=8.5):
443
        """
444
        """
445
        # xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
446
        #                distance=self.distance,
447
        #                z_sun = z_sun*us.kpc,
448
        #                unit="deg, deg, kpc").galactocentric.cartesian.xyz.value
449
        # return xyz
450
451
        slc = sin(self.gl/180*pi)
452
        clc = cos(self.gl/180*pi)
453
        sbc = sin(self.gb/180*pi)
454
        cbc = cos(self.gb/180*pi)
455
        rgalc = self.distance*cbc
456
        xc = rgalc*slc
457
        yc = rsun-rgalc*clc
458
        zc = self.distance*sbc
459
        return np.array([xc, yc, zc])
460
461
    def void_factor(self, xyz):
462
        """
463
        Clump edge
464
        0 => use exponential rolloff out to 5 clump radii
465
        1 => uniform and truncated at 1/e clump radius
466
        """
467
        if xyz.ndim == 1:
468
            xyz = xyz[:, None] - self.xyz
469
            ellipsoid_abc = self.ellipsoid_abc
470
        else:
471
            xyz = xyz[:, :, None] - self.xyz[:, None, :]
472
            ellipsoid_abc = self.ellipsoid_abc[:, None, :]
473
474
        xyz = np.array([Rz.dot(Ry).dot(XYZ.T).T
475
                        for Rz, Ry, XYZ in
476
                        zip(self.rotation_z, self.rotation_y, xyz.T)]).T
477
478
        q2 = np.sum(xyz**2 / ellipsoid_abc**2, axis=0)
479
        # NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
480
        # TODO: check this
481
        return (q2 <= 1)*(self.edge == 1) + (q2 <= 5)*(self.edge == 0)*exp(-q2)
482
483
    def ne_voids(self, xyz):
484
        """
485
        The contribution of the clumps to the free
486
        electron density at x, y, z = `xyz`
487
        """
488
        return np.sum(self.void_factor(xyz)*self.ne*self.use_void, axis=-1)
489
490
491
def rotation(theta, axis=-1):
492
    """
493
    Return a rotation matrix around axis
494
    0:x, 1:y, 2:z
495
    """
496
    ct = cos(theta)
497
    st = sin(theta)
498
499
    if axis in (0, -3):
500
        return np.array([[1, 0, 0],
501
                         [0, ct, st],
502
                         [0, -st, ct]])
503
504
    if axis in (1, -2):
505
        return np.array([[ct, 0, st],
506
                         [0, 1, 0],
507
                         [-st, 0, ct]])
508
509
    if axis in (2, -1):
510
        return np.array([[ct, st, 0],
511
                         [-st, ct, 0],
512
                         [0, 0, 1]])
513