ne_local_ism() - Code Metrics - Inspection of "Added more density models and tests" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( f99e41...be615f )

by Ben

created 2016-12-05 23:02 UTC

ne_local_ism() B

↳ Parent: Project

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes	1
Bugs	0	Features	0

Metric	Value
cc	1
c	1
b	0
f	0
dl	0
loc	34
rs	8.8571

"Free electron density model"
import os

import numpy as np
from astropy.table import Table
from numpy import cos
from numpy import cosh
from numpy import exp
from numpy import pi
from numpy import sin
from numpy import sqrt
from numpy import tan


# import astropy.units as us
# from astropy.coordinates import SkyCoord

# Configuration
# TODO: use to config file
# input parameters for large-scale components of NE2001 30 June '02
# flags = {'wg1': 1,
#          'wg2': 1,
#          'wga': 1,
#          'wggc': 1,
#          'wglism': 1,
#          'wgcN': 1,
#          'wgvN': 1}

# solar_params = {'Rsun': 8.3}

# thick_disk_params = {'n1h1': 0.033,
#                      'h1': 0.97,
#                      'A1': 17.5,
#                      'F1': 0.18}

# thin_disk_params = {'n2': 0.08,
#                     'h2': 0.15,
#                     'A2': 3.8,
#                     'F2': 120}

# galactic_center_params = {'xgc': -0.01,
#                           'ygc': 0.0,
#                           'zgc': -0.020,
#                           'rgc': 0.145,
#                           'hgc': 0.026,
#                           'negc0': 10.0,
#                           'Fgc0': 0.6e5}


# spiral_arms_params = {'na': 0.028,
#                       'ha': 0.23,
#                       'wa': 0.65,
#                       'Aa': 10.5,
#                       'Fa': 5,
#                       'narm1': 0.5,
#                       'narm2': 1.2,
#                       'narm3': 1.3,
#                       'narm4': 1.0,
#                       'narm5': 0.25,
#                       'warm1': 1.0,
#                       'warm2': 1.5,
#                       'warm3': 1.0,
#                       'warm4': 0.8,
#                       'warm5': 1.0,
#                       'harm1': 1.0,
#                       'harm2': 0.8,
#                       'harm3': 1.3,
#                       'harm4': 1.5,
#                       'harm5': 1.0,
#                       'farm1': 1.1,
#                       'farm2': 0.3,
#                       'farm3': 0.4,
#                       'farm4': 1.5,
#                       'farm5': 0.3}


ldr_params = {'abc': np.array([1.50, .750, .50]),
              'center': np.array([1.36, 8.06, 0.0]),
              'theta': -24.2*pi/180,
              'ne': 0.012,
              'F': 0.1}

lsb_params = {'abc': np.array([1.050, .4250, .3250]),
              'center': np.array([-0.75, 9.0, -0.05]),
              'theta': 139.*pi/180,
              'ne': 0.016,
              'F':  0.01}

lhb_params = {'abc': np.array([.0850, .1000, .330]),
              'center': np.array([0.01, 8.45, 0.17]),
              'theta': 15*pi/180,
              'ne': 0.005,
              'F': 0.01}

loop_params = {'center': np.array([-0.045, 8.40, 0.07]),
               'r': 0.120,
               'dr': 0.060,
               'ne1': 0.0125,
               'ne2': 0.0125,
               'F1': 0.2,
               'F2': 0.01}


def ne_thick_disk(xyz, ne_disk, rdisk, hdisk, r_sun):
    """
    Calculate the contribution of the thick disk to the free electron density
     at x, y, z = `xyz`
    """
    r2d = sqrt(xyz[0]**2 + xyz[1]**2)
    k = pi/2/rdisk
    return ne_disk * (cos(r2d*k)/cos(r_sun*k) /
                      cosh(xyz[-1]/hdisk)**2 *
                      (r2d < rdisk))


def ne_thin_disk(xyz, ne_disk, rdisk, hdisk):
    """
    Calculate the contribution of the thin disk to the free electron density
     at x, y, z = `xyz`
    """
    r2d = sqrt(xyz[0]**2 + xyz[1]**2)
    return ne_disk * (exp(-(r2d - rdisk)**2/1.8**2) /
                      cosh(xyz[-1]/hdisk)**2)  # Why 1.8?


def ne_gc(xyz, ne_gc0, rgc, hgc, xyz_gc):
    """
    Calculate the contribution of the Galactic center to the free
    electron density at x, y, z = `xyz`
    """
    # Here I'm using the expression in the NE2001 code which is inconsistent
    # with Cordes and Lazio 2011 (0207156v3) (See Table 2)
    try:
        xyz = xyz - xyz_gc
    except ValueError:
        xyz = xyz - xyz_gc[:, None]

    r2d = sqrt(xyz[0]**2 + xyz[1]**2)

    # ????
    # Cordes and Lazio 2011 (0207156v3) (Table 2)
    # return ne_gc0*exp(-(r2d/rgc)**2 - (xyz[-1]/hgc)**2)
    # ????

    # Constant ne (form NE2001 code)
    return ne_gc0*((r2d/rgc)**2 + (xyz[-1]/hgc)**2 < 1)*(r2d < rgc)


def ne_local_ism(xyz, ldr_params, lsb_params, lhb_params, loop_params):
    """
    Calculate the contribution of the local ISM to the free
    electron density at x, y, z = `xyz`
    """
    # low density region in Q1
    neldr = ldr_params['ne']*in_ellisoid(xyz, ldr_params['center'],
                                         ldr_params['abc'],
                                         ldr_params['theta'])
    # Local Super Bubble
    nelsb = lsb_params['ne']*in_ellisoid(xyz, lsb_params['center'],
                                         lsb_params['abc'],
                                         lsb_params['theta'])

    # Local Hot Bubble
    nelhb = lhb_params['ne']*in_cylinder(xyz, lhb_params['center'],
                                         lhb_params['abc'],
                                         lhb_params['theta'])
    # Loop I
    irr1 = in_half_sphere(xyz, loop_params['center'], loop_params['r'])
    irr2 = in_half_sphere(xyz, loop_params['center'],
                          loop_params['r'] + loop_params['dr'])
    neloop = loop_params['ne1'] * irr1 + loop_params['ne2'] * irr2*(~irr1)
    wlhb, wloop, wlsb, wldr = (nelhb > 0,
                               neloop > 0,
                               nelsb > 0,
                               neldr > 0)
    ne_lism = ((1 - wlhb) *
               ((1 - wloop) * (wlsb*nelsb + (1-wlsb) * neldr) +
                wloop*neloop) +
               wlhb*nelhb)

    wlism = np.maximum(wloop, np.maximum(wldr, np.maximum(wlsb, wlhb)))
    return ne_lism, wlism


def in_ellisoid(xyz, xyz_center, abc_ellipsoid, theta):
    """
    Test if xyz in the ellipsoid
    Theta in radians
    """
    try:
        xyz = xyz - xyz_center
    except ValueError:
        xyz = xyz - xyz_center[:, None]
        abc_ellipsoid = abc_ellipsoid[:, None]

    rot = np.array([[cos(theta), sin(theta)],
                    [-sin(theta), cos(theta)]])
    xyz[:2] = (xyz[:2].T.dot(rot)).T

    xyz_p = xyz/abc_ellipsoid

    return np.sum(xyz_p**2, axis=0) <= 1


def in_cylinder(xyz, xyz_center, abc_cylinder, theta):
    """
    Test if xyz in the cylinder
    Theta in radians
    """
    try:
        xyz = xyz - xyz_center
    except ValueError:
        xyz = xyz - xyz_center[:, None]
        abc_cylinder = np.vstack([abc_cylinder]*xyz.shape[-1]).T
    xyz[2] -= tan(theta)*xyz[-1]

    abc_cylinder_p = abc_cylinder.copy()
    z_c = (xyz_center[-1] - abc_cylinder[-1])
    izz = (xyz[-1] <= 0)*(xyz[-1] <= z_c)
    abc_cylinder_p[0] = (0.001 +
                         (abc_cylinder[0] - 0.001) *
                         (1 - xyz[-1]/z_c))*izz + abc_cylinder[0]*(~izz)

    xyz_p = xyz/abc_cylinder_p

    return (xyz_p[0]**2 + xyz_p[1]**2 <= 1) * (xyz_p[-1]**2 <= 1)


def in_half_sphere(xyz, xyz_center, r_sphere):
    "Test if `xyz` in the sphere with radius r_sphere  centerd at `xyz_center`"
    try:
        xyz = xyz - xyz_center
    except ValueError:
        xyz = xyz - xyz_center[:, None]
    distance = sqrt(np.sum(xyz**2, axis=0))
    return (distance <= r_sphere)*(xyz[-1] >= 0)


class Clumps(object):
    """
    """

    def __init__(self, clumps_file=None):
        """
        """
        if not clumps_file:
            this_dir, _ = os.path.split(__file__)
            clumps_file = os.path.join(this_dir, "data", "neclumpN.NE2001.dat")
        self._data = Table.read(clumps_file, format='ascii')

    @property
    def use_clump(self):
        """
        """
        return self._data['flag'] == 0

    @property
    def xyz(self):
        """
        """
        try:
            return self._xyz
        except AttributeError:
            self._xyz = self.get_xyz()
        return self._xyz

    @property
    def gl(self):
        """
        Galactic longitude (deg)
        """
        return self._data['l']

    @property
    def gb(self):
        """
        Galactic latitude (deg)
        """
        return self._data['b']

    @property
    def distance(self):
        """
        Distance from the sun (kpc)
        """
        return self._data['dc']

    @property
    def radius(self):
        """
        Radius of the clump (kpc)
        """
        return self._data['rc']

    @property
    def ne(self):
        """
        Electron density of each clump (cm^{-3})
        """
        return self._data['nec']

    @property
    def edge(self):
        """
        Clump edge
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        return self._data['edge']

    def get_xyz(self, z_sun=0, rsun=8.5):
        """
        """
        # xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
        #                distance=self.distance,
        #                z_sun = z_sun*us.kpc,
        #                unit="deg, deg, kpc").galactocentric.cartesian.xyz.value
        # return xyz

        slc = sin(self.gl/180*pi)
        clc = cos(self.gl/180*pi)
        sbc = sin(self.gb/180*pi)
        cbc = cos(self.gb/180*pi)
        rgalc = self.distance*cbc
        xc = rgalc*slc
        yc = rsun-rgalc*clc
        zc = self.distance*sbc
        return np.array([xc, yc, zc])

    def clump_factor(self, xyz):
        """
        Clump edge
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        try:
            xyz = xyz[:, None] - self.xyz
        except ValueError:
            xyz = xyz[:, :, None] - self.xyz[:, None, :]

        q2 = (np.sum(xyz**2, axis=0) /
              self.radius**2)
        # NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
        # TODO: check this
        return (q2 <= 1)*(self.edge == 1) + (q2 <= 5)*(self.edge == 0)*exp(-q2)

    def ne_clumps(self, xyz):
        """
        The contribution of the clumps to the free
        electron density at x, y, z = `xyz`
        """
        return np.sum(self.clump_factor(xyz)*self.ne*self.use_clump, axis=-1)


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(r2dk)/cos(r_sunk) /
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.82) /
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_gc0exp(-(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) * (wlsbnelsb + (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 = np.array([[cos(theta), sin(theta)],
197			[-sin(theta), cos(theta)]])
198			xyz[:2] = (xyz[:2].T.dot(rot)).T
199
200			xyz_p = xyz/abc_ellipsoid
201
202			return np.sum(xyz_p**2, axis=0) <= 1
203
204
205			def in_cylinder(xyz, xyz_center, abc_cylinder, theta):
206			"""
207			Test if xyz in the cylinder
208			Theta in radians
209			"""
210			try:
211			xyz = xyz - xyz_center
212			except ValueError:
213			xyz = xyz - xyz_center[:, None]
214			abc_cylinder = np.vstack([abc_cylinder]*xyz.shape[-1]).T
215			xyz[2] -= tan(theta)*xyz[-1]
216
217			abc_cylinder_p = abc_cylinder.copy()
218			z_c = (xyz_center[-1] - abc_cylinder[-1])
219			izz = (xyz[-1] <= 0)*(xyz[-1] <= z_c)
220			abc_cylinder_p[0] = (0.001 +
221			(abc_cylinder[0] - 0.001) *
222			(1 - xyz[-1]/z_c))izz + abc_cylinder[0](~izz)
223
224			xyz_p = xyz/abc_cylinder_p
225
226			return (xyz_p[0]2 + xyz_p[1]2 <= 1) * (xyz_p[-1]**2 <= 1)
227
228
229			def in_half_sphere(xyz, xyz_center, r_sphere):
230			"Test if `xyz` in the sphere with radius r_sphere centerd at `xyz_center`"
231			try:
232			xyz = xyz - xyz_center
233			except ValueError:
234			xyz = xyz - xyz_center[:, None]
235			distance = sqrt(np.sum(xyz**2, axis=0))
236			return (distance <= r_sphere)*(xyz[-1] >= 0)
237
238
239			class Clumps(object):
240			"""
241			"""
242
243			def __init__(self, clumps_file=None):
244			"""
245			"""
246			if not clumps_file:
247			this_dir, _ = os.path.split(__file__)
248			clumps_file = os.path.join(this_dir, "data", "neclumpN.NE2001.dat")
249			self._data = Table.read(clumps_file, format='ascii')
250
251			@property
252			def use_clump(self):
253			"""
254			"""
255			return self._data['flag'] == 0
256
257			@property
258			def xyz(self):
259			"""
260			"""
261			try:
262			return self._xyz
263			except AttributeError:
264			self._xyz = self.get_xyz()
265			return self._xyz
266
267			@property
268			def gl(self):
269			"""
270			Galactic longitude (deg)
271			"""
272			return self._data['l']
273
274			@property
275			def gb(self):
276			"""
277			Galactic latitude (deg)
278			"""
279			return self._data['b']
280
281			@property
282			def distance(self):
283			"""
284			Distance from the sun (kpc)
285			"""
286			return self._data['dc']
287
288			@property
289			def radius(self):
290			"""
291			Radius of the clump (kpc)
292			"""
293			return self._data['rc']
294
295			@property
296			def ne(self):
297			"""
298			Electron density of each clump (cm^{-3})
299			"""
300			return self._data['nec']
301
302			@property
303			def edge(self):
304			"""
305			Clump edge
306			0 => use exponential rolloff out to 5 clump radii
307			1 => uniform and truncated at 1/e clump radius
308			"""
309			return self._data['edge']
310
311			def get_xyz(self, z_sun=0, rsun=8.5):
312			"""
313			"""
314			# xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
315			# distance=self.distance,
316			# z_sun = z_sun*us.kpc,
317			# unit="deg, deg, kpc").galactocentric.cartesian.xyz.value
318			# return xyz
319
320			slc = sin(self.gl/180*pi)
321			clc = cos(self.gl/180*pi)
322			sbc = sin(self.gb/180*pi)
323			cbc = cos(self.gb/180*pi)
324			rgalc = self.distance*cbc
325			xc = rgalc*slc
326			yc = rsun-rgalc*clc
327			zc = self.distance*sbc
328			return np.array([xc, yc, zc])
329
330			def clump_factor(self, xyz):
331			"""
332			Clump edge
333			0 => use exponential rolloff out to 5 clump radii
334			1 => uniform and truncated at 1/e clump radius
335			"""
336			try:
337			xyz = xyz[:, None] - self.xyz
338			except ValueError:
339			xyz = xyz[:, :, None] - self.xyz[:, None, :]
340
341			q2 = (np.sum(xyz**2, axis=0) /
342			self.radius**2)
343			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
344			# TODO: check this
345			return (q2 <= 1)(self.edge == 1) + (q2 <= 5)(self.edge == 0)*exp(-q2)
346
347			def ne_clumps(self, xyz):
348			"""
349			The contribution of the clumps to the free
350			electron density at x, y, z = `xyz`
351			"""
352			return np.sum(self.clump_factor(xyz)self.neself.use_clump, axis=-1)
353

benbaror / ne2001

Push — master ( f99e41...be615f )

ne_local_ism() B

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like