Clumps.radius2() - Code Metrics - Inspection of "Improved run-time with numba" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( 2c8f60...04a3aa )

by Ben

created 2016-12-09 16:26 UTC

Clumps.radius2() A

↳ Parent: Clumps

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	6
rs	9.4285

"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 sqrt
from numpy import tan
from scipy.integrate import quad as integrate
from numba import jit

from .utils import ClassOperation
from .utils import galactic_to_galactocentric
from .utils import lzproperty
from .utils import rotation

# 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]),
           epsrel=1e-4, epsabs=1e-6,
           *arg, **kwargs):
        """
        Calculate the dispersion measure at location `xyz`
        """
        xyz = xyz - xyz_sun

        dfinal = sqrt(np.sum(xyz**2, axis=0))
        return integrate(lambda x: self.ne(xyz_sun + x*xyz),
                         0, 1, *arg, epsrel=epsrel, epsabs=epsabs,
                         **kwargs)[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)


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

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

    @lzproperty
    def xyz(self):
        """
        """
        return self.get_xyz()

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

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

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

    @lzproperty
    def radius2(self):
        """
        Radius of the clump (kpc)
        """
        return self._data['rc']**2

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

    @lzproperty
    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:
            return clump_factor(xyz, self.xyz, self.radius2, self.edge)
        else:
            xyz = xyz[:, :, None] - self.xyz[:, None, :]

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

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

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

    @lzproperty
    def xyz(self):
        """
        """
        return self.get_xyz()

    @lzproperty
    def xyz_rot(self):
        """
        """
        return np.array([R.dot(xyzi) for R,
                         xyzi in zip(self.rotation, self.xyz.T)]).T

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

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

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

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

    @lzproperty
    def rotation(self):
        """
        Rotation and rescaling matrix
        """
        return np.array([
            (rotation(thetaz*pi/180, -1).dot(
                rotation(thetay*pi/180, 1)).T/abc).T
            for thetaz, thetay, abc
            in zip(self._data['thvz'], self._data['thvy'],
                   self.ellipsoid_abc.T)
        ])

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

    @lzproperty
    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
        """
        xyz = (self.rotation.dot(xyz).T - self.xyz_rot).T

        q2 = np.sum(xyz**2, axis=1).T
        # 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)


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._clumps)

    def electron_density(self, xyz):
        return self._combined.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)


@jit
def clump_factor(xyz, xyz0, r2, edge):
    """
    Clump edge
    0 => use exponential rolloff out to 5 clump radii
    1 => uniform and truncated at 1/e clump radius
    """
    xyz = (xyz - xyz0.T).T

    q2 = (xyz[0]**2 + xyz[1]**2 + xyz[2]**2) / r2
    # NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
    # TODO: check this
    q5 = (q2 <= 5)*(edge == 0)
    res = np.zeros_like(q2)
    res[(q2 <= 1)*(edge == 1)] = 1
    res[q5] = exp(-q2[q5])
    return res


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

benbaror / ne2001

Push — master ( 2c8f60...04a3aa )

Clumps.radius2() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like