ElectronDensity.electron_density() - Code Metrics - Inspection of "Added ElectronDensity class to handle the combined..." - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( de1306...06a50e )

by Ben

created 2016-12-09 02:36 UTC

ElectronDensity.electron_density() A

↳ Parent: ElectronDensity

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	2
rs	10

"Free electron density model"
import os
from functools import partial

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
from scipy.integrate import quad

from .utils import ClassOperation

# 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 NEobject(ClassOperation):
    """
    A general electron density object
    """

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

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

    def DM(self, xyz, xyz_sun=np.array([0, 8.5, 0])):
        """
        Calculate the dispersion measure at location `xyz`
        """
        xyz = xyz - xyz_sun

        dfinal = sqrt(np.sum(xyz**2, axis=0))

        return quad(lambda x: self.ne(xyz_sun + x*xyz),
                    0, 1)[0]*dfinal*1000

    def ne(self, *args):
        return self.electron_density(*args)

    def electron_density(self, xyz):
        "Electron density at the location `xyz`"
        return self._ne0*self._func(xyz)

    @property
    def F(self, xyz):
        "Fluctuation parameter"
        return (self.ne(xyz) > 0)*self._fparam


class NEcombine(NEobject):
    """
    """

    def __init__(self, object1, object2):
        """
        """
        self._object1 = object1
        self._object2 = object2

    def electron_density(self, *args):
        ne1 = self._object1.ne(*args)
        ne2 = self._object2.ne(*args)
        return ne1 + ne2*(ne1 <= 0)


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

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

        self.loop = NEcombine(self.loop_in, self.loop_out)
        self._lism = NEcombine(self.lhb,
                               NEcombine(self.loop,
                                         NEcombine(self.lsb, self.ldr)))

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


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
    """
    xyz0 = xyz.copy()
    try:
        xyz = xyz - center
    except ValueError:
        xyz = xyz - center[:, None]
        cylinder = np.vstack([cylinder]*xyz.shape[-1]).T
    xyz[1] -= tan(theta)*xyz0[-1]

    cylinder_p = cylinder.copy()
    z_c = (center[-1] - cylinder[-1])
    izz = (xyz0[-1] <= 0)*(xyz0[-1] >= z_c)
    cylinder_p[0] = (0.001 +
                     (cylinder[0] - 0.001) *
                     (1 - xyz0[-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`"
    xyz0 = xyz.copy()
    try:
        xyz = xyz - center
    except ValueError:
        xyz = xyz - center[:, None]
    distance = sqrt(np.sum(xyz**2, axis=0))
    return (distance <= radius)*(xyz0[-1] >= 0)


class Clumps(NEobject):
    """
    """

    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 ne0(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
        return galactic_to_galactocentric(l=self.gl, b=self.gb,
                                          distance=self.distance, rsun=rsun)

    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 electron_density(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.ne0*self.use_clump, axis=-1)


class Voids(NEobject):
    """
    """

    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 ne0(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
        return galactic_to_galactocentric(l=self.gl, b=self.gb,
                                          distance=self.distance, rsun=rsun)

    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 electron_density(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.ne0*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]])


def galactic_to_galactocentric(l, b, distance, rsun):
    slc = sin(l/180*pi)
    clc = cos(l/180*pi)
    sbc = sin(b/180*pi)
    cbc = cos(b/180*pi)
    rgalc = distance*cbc
    xc = rgalc*slc
    yc = rsun-rgalc*clc
    zc = distance*sbc
    return np.array([xc, yc, zc])


class ElectronDensity(NEobject):
    """
    A class holding all the elements which contribute to free electron density
    """

    def __init__(self, r_sun=8.5, clumps_file=None, voids_file=None,
                 **params):
        """
        """
        self._params = params
        self._thick_disk = NEobject(thick_disk, r_sun=r_sun,
                                    **params['thick_disk'])
        self._thin_disk = NEobject(thin_disk, **params['thin_disk'])
        self._galactic_center = NEobject(gc, **params['galactic_center'])
        self._lism = LocalISM(**params)
        self._clumps = Clumps(clumps_file=clumps_file)
        self._voids = Voids(voids_file=voids_file)
        self._combined = (NEcombine(self._voids,
                                    NEcombine(self._lism,
                                              self._thick_disk +
                                              self._thin_disk +
                                              self._galactic_center)) +
                          self._galactic_center)

    def electron_density(self, xyz):
        return self._combined.ne(xyz)


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

benbaror / ne2001

Push — master ( de1306...06a50e )

ElectronDensity.electron_density() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like