NEobject.F() - Code Metrics - Inspection of "Defined a new class to combine NEobject(s)" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( cc9cbc...fb6854 )

by Ben

created 2016-12-08 19:14 UTC

NEobject.F() A

↳ Parent: NEobject

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
c	0
b	0
f	0
dl	0
loc	4
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

# 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(object):
    """
    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(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
        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 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
        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 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]])


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])


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

benbaror / ne2001

Push — master ( cc9cbc...fb6854 )

NEobject.F() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like