LocalISM.electron_density() - Code Metrics - Inspection of "Rearranged density functions into classes" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( de38d6...c9e5b6 )

by Ben

created 2016-12-07 07:43 UTC

LocalISM.electron_density() A

↳ Parent: LocalISM

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	2
c	0
b	0
f	0
dl	0
loc	16
rs	9.4285

"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}

# 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}

PARAMS = {
    'thick_disk': {'e_density': 0.033/0.97,
                   'height': 0.97,
                   'radius': 17.5,
                   'F': 0.18},

    'thin_disk': {'e_density': 0.08,
                  'height': 0.15,
                  'radius': 3.8,
                  'F': 120},

    'galactic_center': {'e_density': 10.0,
                        'center': np.array([-0.01, 0.0, -0.020]),
                        'radius': 0.145,
                        'height': 0.026,
                        'F': 0.6e5},

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

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

    'lhb': {'cylinder': np.array([.0850, .1000, .330]),
            'center': np.array([0.01, 8.45, 0.17]),
            'theta': 15*pi/180,
            'e_density': 0.005,
            'F': 0.01},

    'loop_in': {'center': np.array([-0.045, 8.40, 0.07]),
                'radius': 0.120,
                'e_density': 0.0125,
                'F': 0.2},
    'loop_out': {'center': np.array([-0.045, 8.40, 0.07]),
                 'radius': 0.120 + 0.060,
                 'e_density': 0.0125,
                 'F': 0.01}}


def thick_disk(xyz, r_sun, radius, height):
    """
    Calculate the contribution of the thick disk to the free electron density
     at x, y, z = `xyz`
    """
    r_ratio = sqrt(xyz[0]**2 + xyz[1]**2)/radius
    return (cos(r_ratio*pi/2)/cos(r_sun*pi/2/radius) /
            cosh(xyz[-1]/height)**2 *
            (r_ratio < 1))


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


def gc(xyz, center, radius, height):
    """
    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 - center
    except ValueError:
        xyz = xyz - center[:, None]

    r_ratio = sqrt(xyz[0]**2 + xyz[1]**2)/radius

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

    # Constant ne (form NE2001 code)
    return (r_ratio**2 + (xyz[-1]/height)**2 < 1)*(r_ratio <= 1)


class Class_Operation(object):
    """
    Class Operation
    """

    def __init__(self, operation, cls1, cls2):
        """
        """
        self.cls1 = cls1
        self.cls2 = cls2
        self._operation = operation

    def __getattr__(self, arg):
        return getattr(getattr(self.cls1, arg),
                       self._operation)(getattr(self.cls2, arg))


class NEobject(object):
    """
    A general electron density object
    """

    def __init__(self, xyz, func, **params):
        """

        Arguments:
        - `xyz`: Location where the electron density is calculated
        - `func`: Electron density function
        - `**params`: Model parameter
        """
        self._xyz = xyz
        self._func = func
        self._fparam = params.pop('F')
        self._ne0 = params.pop('e_density')
        self._params = params

    def __add__(self, other):
        return Class_Operation('__add__', self, other)

    def __sub__(self, other):
        def sub(arg):
            return getattr(self, arg) - getattr(other, arg)
        return sub

    @property
    def xyz(self):
        "Location where the electron density will be calculated"
        return self._xyz

    @property
    def electron_density(self):
        "Electron density at the location `xyz`"
        try:
            return self._ne
        except AttributeError:
            self._ne = self._ne0*self._func(self.xyz, **self._params)
        return self._ne

    @property
    def wight(self):
        """
        Is this object contributing to the electron density
        at the location `xyz`
        """
        return self.electron_density > 0

    @property
    def w(self):
        return self.wight

    @property
    def ne(self):
        return self.electron_density

    @wight.setter
    def wight(self, wight):
        """
        Is this object contributing to the electron density
        at the location `xyz`
        """
        self._ne = self.ne*wight

    @property
    def F(self):
        "Fluctuation parameter"
        return self.wight*self._fparam


class LocalISM(object):
    """
    Calculate the contribution of the local ISM
    to the free electron density at x, y, z = `xyz`
    """

    def __init__(self, xyz, **params):
        """
        """
        self.xyz = xyz
        self.ldr = NEobject(xyz, in_ellipsoid, **params['ldr'])
        self.lsb = NEobject(xyz, in_ellipsoid, **params['lsb'])
        self.lhb = NEobject(xyz, in_cylinder, **params['lhb'])
        self.loop_in = NEobject(xyz, in_half_sphere, **params['loop_in'])
        self.loop_out = NEobject(xyz, in_half_sphere, **params['loop_out'])
        self.loop_out.wight = ~self.loop_in.w
        self.loop = self.loop_in + self.loop_out

    @property
    def electron_density(self):
        """
        Calculate the contribution of the local ISM to the free
        electron density at x, y, z = `xyz`
        """

        try:
            return self._nelism
        except AttributeError:
            self._nelism = (self.lhb.ne +
                            (self.loop.ne +
                             (self.lsb.ne + self.ldr.ne*~self.lsb.w) *
                             ~self.loop.w)*~self.lhb.w)

        return self._nelism

    @property
    def flism(self):
        try:
            return self._flism
        except AttributeError:
            self._flism = (self.lhb.F + ~self.lhb.w *
                           (self.loop.F + ~self.loop.w *
                            (self.lsb.F + self.ldr.F*~self.lsb.w)))

        return self._flism

    @property
    def wlism(self):
        # This should be equivalent to ne>0
        # TODO: Check this!
        return np.maximum(self.loop.w,
                          np.maximum(self.ldr.w,
                                     np.maximum(self.lsb.w, self.lhb.w)))


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

    rot = rotation(theta, -1)
    xyz = rot.dot(xyz)

    xyz_p = xyz/ellipsoid

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


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

    cylinder_p = cylinder.copy()
    z_c = (center[-1] - cylinder[-1])
    izz = (xyz[-1] <= 0)*(xyz[-1] <= z_c)
    cylinder_p[0] = (0.001 +
                     (cylinder[0] - 0.001) *
                     (1 - xyz[-1]/z_c))*izz + cylinder[0]*(~izz)

    xyz_p = xyz/cylinder_p

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


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


class Voids(object):
    """
    """

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

    @property
    def use_void(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['dv']

    @property
    def ellipsoid_abc(self):
        """
        Void axis
        """
        return np.array([self._data['aav'],
                         self._data['bbv'],
                         self._data['ccv']])

    @property
    def rotation_y(self):
        """
        Rotation around the y axis
        """
        return [rotation(theta*pi/180, 1) for theta in self._data['thvy']]

    @property
    def rotation_z(self):
        """
        Rotation around the z axis
        """
        return [rotation(theta*pi/180, -1) for theta in self._data['thvz']]

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

    @property
    def edge(self):
        """
        Void 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, 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 void_factor(self, xyz):
        """
        Clump edge
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        if xyz.ndim == 1:
            xyz = xyz[:, None] - self.xyz
            ellipsoid_abc = self.ellipsoid_abc
        else:
            xyz = xyz[:, :, None] - self.xyz[:, None, :]
            ellipsoid_abc = self.ellipsoid_abc[:, None, :]

        xyz = np.array([Rz.dot(Ry).dot(XYZ.T).T
                        for Rz, Ry, XYZ in
                        zip(self.rotation_z, self.rotation_y, xyz.T)]).T

        q2 = np.sum(xyz**2 / ellipsoid_abc**2, axis=0)
        # 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_voids(self, xyz):
        """
        The contribution of the clumps to the free
        electron density at x, y, z = `xyz`
        """
        return np.sum(self.void_factor(xyz)*self.ne*self.use_void, axis=-1)


def rotation(theta, axis=-1):
    """
    Return a rotation matrix around axis
    0:x, 1:y, 2:z
    """
    ct = cos(theta)
    st = sin(theta)

    if axis in (0, -3):
        return np.array([[1, 0, 0],
                         [0, ct, st],
                         [0, -st, ct]])

    if axis in (1, -2):
        return np.array([[ct, 0, st],
                         [0, 1, 0],
                         [-st, 0, ct]])

    if axis in (2, -1):
        return np.array([[ct, st, 0],
                         [-st, ct, 0],
                         [0, 0, 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			# spiral_arms_params = {'na': 0.028,
32			# 'ha': 0.23,
33			# 'wa': 0.65,
34			# 'Aa': 10.5,
35			# 'Fa': 5,
36			# 'narm1': 0.5,
37			# 'narm2': 1.2,
38			# 'narm3': 1.3,
39			# 'narm4': 1.0,
40			# 'narm5': 0.25,
41			# 'warm1': 1.0,
42			# 'warm2': 1.5,
43			# 'warm3': 1.0,
44			# 'warm4': 0.8,
45			# 'warm5': 1.0,
46			# 'harm1': 1.0,
47			# 'harm2': 0.8,
48			# 'harm3': 1.3,
49			# 'harm4': 1.5,
50			# 'harm5': 1.0,
51			# 'farm1': 1.1,
52			# 'farm2': 0.3,
53			# 'farm3': 0.4,
54			# 'farm4': 1.5,
55			# 'farm5': 0.3}
56
57			PARAMS = {
58			'thick_disk': {'e_density': 0.033/0.97,
59			'height': 0.97,
60			'radius': 17.5,
61			'F': 0.18},
62
63			'thin_disk': {'e_density': 0.08,
64			'height': 0.15,
65			'radius': 3.8,
66			'F': 120},
67
68			'galactic_center': {'e_density': 10.0,
69			'center': np.array([-0.01, 0.0, -0.020]),
70			'radius': 0.145,
71			'height': 0.026,
72			'F': 0.6e5},
73
74			'ldr': {'ellipsoid': np.array([1.50, .750, .50]),
75			'center': np.array([1.36, 8.06, 0.0]),
76			'theta': -24.2*pi/180,
77			'e_density': 0.012,
78			'F': 0.1},
79
80			'lsb': {'ellipsoid': np.array([1.050, .4250, .3250]),
81			'center': np.array([-0.75, 9.0, -0.05]),
82			'theta': 139.*pi/180,
83			'e_density': 0.016,
84			'F': 0.01},
85
86			'lhb': {'cylinder': np.array([.0850, .1000, .330]),
87			'center': np.array([0.01, 8.45, 0.17]),
88			'theta': 15*pi/180,
89			'e_density': 0.005,
90			'F': 0.01},
91
92			'loop_in': {'center': np.array([-0.045, 8.40, 0.07]),
93			'radius': 0.120,
94			'e_density': 0.0125,
95			'F': 0.2},
96			'loop_out': {'center': np.array([-0.045, 8.40, 0.07]),
97			'radius': 0.120 + 0.060,
98			'e_density': 0.0125,
99			'F': 0.01}}
100
101
102			def thick_disk(xyz, r_sun, radius, height):
103			"""
104			Calculate the contribution of the thick disk to the free electron density
105			at x, y, z = `xyz`
106			"""
107			r_ratio = sqrt(xyz[0]2 + xyz[1]2)/radius
108			return (cos(r_ratiopi/2)/cos(r_sunpi/2/radius) /
109			cosh(xyz[-1]/height)*2
110			(r_ratio < 1))
111
112
113			def thin_disk(xyz, radius, height):
114			"""
115			Calculate the contribution of the thin disk to the free electron density
116			at x, y, z = `xyz`
117			"""
118			r_ratio = sqrt(xyz[0]2 + xyz[1]2)/radius
119			return (exp(-(1 - r_ratio)*2radius2/1.82) /
120			cosh(xyz[-1]/height)**2) # Why 1.8?
121
122
123			def gc(xyz, center, radius, height):
124			"""
125			Calculate the contribution of the Galactic center to the free
126			electron density at x, y, z = `xyz`
127			"""
128			# Here I'm using the expression in the NE2001 code which is inconsistent
129			# with Cordes and Lazio 2011 (0207156v3) (See Table 2)
130			try:
131			xyz = xyz - center
132			except ValueError:
133			xyz = xyz - center[:, None]
134
135			r_ratio = sqrt(xyz[0]2 + xyz[1]2)/radius
136
137			# ????
138			# Cordes and Lazio 2011 (0207156v3) (Table 2)
139			# return ne_gc0exp(-(r2d/rgc)2 - (xyz[-1]/hgc)*2)
140			# ????
141
142			# Constant ne (form NE2001 code)
143			return (r_ratio2 + (xyz[-1]/height)2 < 1)*(r_ratio <= 1)
144
145
146			class Class_Operation(object):
147			"""
148			Class Operation
149			"""
150
151			def __init__(self, operation, cls1, cls2):
152			"""
153			"""
154			self.cls1 = cls1
155			self.cls2 = cls2
156			self._operation = operation
157
158			def __getattr__(self, arg):
159			return getattr(getattr(self.cls1, arg),
160			self._operation)(getattr(self.cls2, arg))
161
162
163			class NEobject(object):
164			"""
165			A general electron density object
166			"""
167
168			def __init__(self, xyz, func, **params):
169			"""
170
171			Arguments:
172			- `xyz`: Location where the electron density is calculated
173			- `func`: Electron density function
174			- `**params`: Model parameter
175			"""
176			self._xyz = xyz
177			self._func = func
178			self._fparam = params.pop('F')
179			self._ne0 = params.pop('e_density')
180			self._params = params
181
182			def __add__(self, other):
183			return Class_Operation('__add__', self, other)
184
185			def __sub__(self, other):
186			def sub(arg):
187			return getattr(self, arg) - getattr(other, arg)
188			return sub
189
190			@property
191			def xyz(self):
192			"Location where the electron density will be calculated"
193			return self._xyz
194
195			@property
196			def electron_density(self):
197			"Electron density at the location `xyz`"
198			try:
199			return self._ne
200			except AttributeError:
201			self._ne = self._ne0self._func(self.xyz, *self._params)
202			return self._ne
203
204			@property
205			def wight(self):
206			"""
207			Is this object contributing to the electron density
208			at the location `xyz`
209			"""
210			return self.electron_density > 0
211
212			@property
213			def w(self):
214			return self.wight
215
216			@property
217			def ne(self):
218			return self.electron_density
219
220			@wight.setter
221			def wight(self, wight):
222			"""
223			Is this object contributing to the electron density
224			at the location `xyz`
225			"""
226			self._ne = self.ne*wight
227
228			@property
229			def F(self):
230			"Fluctuation parameter"
231			return self.wight*self._fparam
232
233
234			class LocalISM(object):
235			"""
236			Calculate the contribution of the local ISM
237			to the free electron density at x, y, z = `xyz`
238			"""
239
240			def __init__(self, xyz, **params):
241			"""
242			"""
243			self.xyz = xyz
244			self.ldr = NEobject(xyz, in_ellipsoid, **params['ldr'])
245			self.lsb = NEobject(xyz, in_ellipsoid, **params['lsb'])
246			self.lhb = NEobject(xyz, in_cylinder, **params['lhb'])
247			self.loop_in = NEobject(xyz, in_half_sphere, **params['loop_in'])
248			self.loop_out = NEobject(xyz, in_half_sphere, **params['loop_out'])
249			self.loop_out.wight = ~self.loop_in.w
250			self.loop = self.loop_in + self.loop_out
251
252			@property
253			def electron_density(self):
254			"""
255			Calculate the contribution of the local ISM to the free
256			electron density at x, y, z = `xyz`
257			"""
258
259			try:
260			return self._nelism
261			except AttributeError:
262			self._nelism = (self.lhb.ne +
263			(self.loop.ne +
264			(self.lsb.ne + self.ldr.ne~self.lsb.w)
265			~self.loop.w)*~self.lhb.w)
266
267			return self._nelism
268
269			@property
270			def flism(self):
271			try:
272			return self._flism
273			except AttributeError:
274			self._flism = (self.lhb.F + ~self.lhb.w *
275			(self.loop.F + ~self.loop.w *
276			(self.lsb.F + self.ldr.F*~self.lsb.w)))
277
278			return self._flism
279
280			@property
281			def wlism(self):
282			# This should be equivalent to ne>0
283			# TODO: Check this!
284			return np.maximum(self.loop.w,
285			np.maximum(self.ldr.w,
286			np.maximum(self.lsb.w, self.lhb.w)))
287
288
289			def in_ellipsoid(xyz, center, ellipsoid, theta):
290			"""
291			Test if xyz in the ellipsoid
292			Theta in radians
293			"""
294			try:
295			xyz = xyz - center
296			except ValueError:
297			xyz = xyz - center[:, None]
298			ellipsoid = ellipsoid[:, None]
299
300			rot = rotation(theta, -1)
301			xyz = rot.dot(xyz)
302
303			xyz_p = xyz/ellipsoid
304
305			return np.sum(xyz_p**2, axis=0) <= 1
306
307
308			def in_cylinder(xyz, center, cylinder, theta):
309			"""
310			Test if xyz in the cylinder
311			Theta in radians
312			"""
313			try:
314			xyz = xyz - center
315			except ValueError:
316			xyz = xyz - center[:, None]
317			cylinder = np.vstack([cylinder]*xyz.shape[-1]).T
318			xyz[2] -= tan(theta)*xyz[-1]
319
320			cylinder_p = cylinder.copy()
321			z_c = (center[-1] - cylinder[-1])
322			izz = (xyz[-1] <= 0)*(xyz[-1] <= z_c)
323			cylinder_p[0] = (0.001 +
324			(cylinder[0] - 0.001) *
325			(1 - xyz[-1]/z_c))izz + cylinder[0](~izz)
326
327			xyz_p = xyz/cylinder_p
328
329			return (xyz_p[0]2 + xyz_p[1]2 <= 1) * (xyz_p[-1]**2 <= 1)
330
331
332			def in_half_sphere(xyz, center, radius):
333			"Test if `xyz` in the sphere with radius r_sphere centerd at `xyz_center`"
334			try:
335			xyz = xyz - center
336			except ValueError:
337			xyz = xyz - center[:, None]
338			distance = sqrt(np.sum(xyz**2, axis=0))
339			return (distance <= radius)*(xyz[-1] >= 0)
340
341
342			class Clumps(object):
343			"""
344			"""
345
346			def __init__(self, clumps_file=None):
347			"""
348			"""
349			if not clumps_file:
350			this_dir, _ = os.path.split(__file__)
351			clumps_file = os.path.join(this_dir, "data", "neclumpN.NE2001.dat")
352			self._data = Table.read(clumps_file, format='ascii')
353
354			@property
355			def use_clump(self):
356			"""
357			"""
358			return self._data['flag'] == 0
359
360			@property
361			def xyz(self):
362			"""
363			"""
364			try:
365			return self._xyz
366			except AttributeError:
367			self._xyz = self.get_xyz()
368			return self._xyz
369
370			@property
371			def gl(self):
372			"""
373			Galactic longitude (deg)
374			"""
375			return self._data['l']
376
377			@property
378			def gb(self):
379			"""
380			Galactic latitude (deg)
381			"""
382			return self._data['b']
383
384			@property
385			def distance(self):
386			"""
387			Distance from the sun (kpc)
388			"""
389			return self._data['dc']
390
391			@property
392			def radius(self):
393			"""
394			Radius of the clump (kpc)
395			"""
396			return self._data['rc']
397
398			@property
399			def ne(self):
400			"""
401			Electron density of each clump (cm^{-3})
402			"""
403			return self._data['nec']
404
405			@property
406			def edge(self):
407			"""
408			Clump edge
409			0 => use exponential rolloff out to 5 clump radii
410			1 => uniform and truncated at 1/e clump radius
411			"""
412			return self._data['edge']
413
414			def get_xyz(self, rsun=8.5):
415			"""
416			"""
417			# xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
418			# distance=self.distance,
419			# z_sun = z_sun*us.kpc,
420			# unit="deg, deg, kpc").galactocentric.
421			# cartesian.xyz.value
422			# return xyz
423
424			slc = sin(self.gl/180*pi)
425			clc = cos(self.gl/180*pi)
426			sbc = sin(self.gb/180*pi)
427			cbc = cos(self.gb/180*pi)
428			rgalc = self.distance*cbc
429			xc = rgalc*slc
430			yc = rsun-rgalc*clc
431			zc = self.distance*sbc
432			return np.array([xc, yc, zc])
433
434			def clump_factor(self, xyz):
435			"""
436			Clump edge
437			0 => use exponential rolloff out to 5 clump radii
438			1 => uniform and truncated at 1/e clump radius
439			"""
440			if xyz.ndim == 1:
441			xyz = xyz[:, None] - self.xyz
442			else:
443			xyz = xyz[:, :, None] - self.xyz[:, None, :]
444
445			q2 = (np.sum(xyz**2, axis=0) /
446			self.radius**2)
447			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
448			# TODO: check this
449			return (q2 <= 1)(self.edge == 1) + (q2 <= 5)(self.edge == 0)*exp(-q2)
450
451			def ne_clumps(self, xyz):
452			"""
453			The contribution of the clumps to the free
454			electron density at x, y, z = `xyz`
455			"""
456			return np.sum(self.clump_factor(xyz)self.neself.use_clump, axis=-1)
457
458
459			class Voids(object):
460			"""
461			"""
462
463			def __init__(self, voids_file=None):
464			"""
465			"""
466			if not voids_file:
467			this_dir, _ = os.path.split(__file__)
468			voids_file = os.path.join(this_dir, "data", "nevoidN.NE2001.dat")
469			self._data = Table.read(voids_file, format='ascii')
470
471			@property
472			def use_void(self):
473			"""
474			"""
475			return self._data['flag'] == 0
476
477			@property
478			def xyz(self):
479			"""
480			"""
481			try:
482			return self._xyz
483			except AttributeError:
484			self._xyz = self.get_xyz()
485			return self._xyz
486
487			@property
488			def gl(self):
489			"""
490			Galactic longitude (deg)
491			"""
492			return self._data['l']
493
494			@property
495			def gb(self):
496			"""
497			Galactic latitude (deg)
498			"""
499			return self._data['b']
500
501			@property
502			def distance(self):
503			"""
504			Distance from the sun (kpc)
505			"""
506			return self._data['dv']
507
508			@property
509			def ellipsoid_abc(self):
510			"""
511			Void axis
512			"""
513			return np.array([self._data['aav'],
514			self._data['bbv'],
515			self._data['ccv']])
516
517			@property
518			def rotation_y(self):
519			"""
520			Rotation around the y axis
521			"""
522			return [rotation(theta*pi/180, 1) for theta in self._data['thvy']]
523
524			@property
525			def rotation_z(self):
526			"""
527			Rotation around the z axis
528			"""
529			return [rotation(theta*pi/180, -1) for theta in self._data['thvz']]
530
531			@property
532			def ne(self):
533			"""
534			Electron density of each void (cm^{-3})
535			"""
536			return self._data['nev']
537
538			@property
539			def edge(self):
540			"""
541			Void edge
542			0 => use exponential rolloff out to 5 clump radii
543			1 => uniform and truncated at 1/e clump radius
544			"""
545			return self._data['edge']
546
547			def get_xyz(self, rsun=8.5):
548			"""
549			"""
550			# xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
551			# distance=self.distance,
552			# z_sun = z_sun*us.kpc,
553			# unit="deg, deg, kpc").galactocentric.
554			# cartesian.xyz.value
555			# return xyz
556
557			slc = sin(self.gl/180*pi)
558			clc = cos(self.gl/180*pi)
559			sbc = sin(self.gb/180*pi)
560			cbc = cos(self.gb/180*pi)
561			rgalc = self.distance*cbc
562			xc = rgalc*slc
563			yc = rsun-rgalc*clc
564			zc = self.distance*sbc
565			return np.array([xc, yc, zc])
566
567			def void_factor(self, xyz):
568			"""
569			Clump edge
570			0 => use exponential rolloff out to 5 clump radii
571			1 => uniform and truncated at 1/e clump radius
572			"""
573			if xyz.ndim == 1:
574			xyz = xyz[:, None] - self.xyz
575			ellipsoid_abc = self.ellipsoid_abc
576			else:
577			xyz = xyz[:, :, None] - self.xyz[:, None, :]
578			ellipsoid_abc = self.ellipsoid_abc[:, None, :]
579
580			xyz = np.array([Rz.dot(Ry).dot(XYZ.T).T
581			for Rz, Ry, XYZ in
582			zip(self.rotation_z, self.rotation_y, xyz.T)]).T
583
584			q2 = np.sum(xyz2 / ellipsoid_abc2, axis=0)
585			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
586			# TODO: check this
587			return (q2 <= 1)(self.edge == 1) + (q2 <= 5)(self.edge == 0)*exp(-q2)
588
589			def ne_voids(self, xyz):
590			"""
591			The contribution of the clumps to the free
592			electron density at x, y, z = `xyz`
593			"""
594			return np.sum(self.void_factor(xyz)self.neself.use_void, axis=-1)
595
596
597			def rotation(theta, axis=-1):
598			"""
599			Return a rotation matrix around axis
600			0:x, 1:y, 2:z
601			"""
602			ct = cos(theta)
603			st = sin(theta)
604
605			if axis in (0, -3):
606			return np.array([[1, 0, 0],
607			[0, ct, st],
608			[0, -st, ct]])
609
610			if axis in (1, -2):
611			return np.array([[ct, 0, st],
612			[0, 1, 0],
613			[-st, 0, ct]])
614
615			if axis in (2, -1):
616			return np.array([[ct, st, 0],
617			[-st, ct, 0],
618			[0, 0, 1]])
619

benbaror / ne2001

Push — master ( de38d6...c9e5b6 )

LocalISM.electron_density() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like