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

Completed

Push — master ( daea06...2c8f60 )

by Ben

created 2016-12-09 07:10 UTC

Voids.xyz_rot() A

↳ Parent: Voids

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes	1
Bugs	0	Features	0

Metric	Value
cc	2
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 .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):
        """
        """
        try:
            return self._xyz
        except AttributeError:
            self._xyz = self.get_xyz()
        return self._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 radius(self):
        """
        Radius of the clump (kpc)
        """
        return self._data['rc']

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

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


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

benbaror / ne2001

Push — master ( daea06...2c8f60 )

Voids.xyz_rot() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like