galactic_to_galactocentric() - Code Metrics - Inspection of "Added DM to NEobject" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( c9e5b6...b5ca57 )

by Ben

created 2016-12-07 20:17 UTC

galactic_to_galactocentric() A

↳ Parent: Project

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

"Free electron density model"
import os

import numpy as np
from astropy.table import Table
from numpy import cos
from numpy import cosh
from numpy import exp
from numpy import pi
from numpy import sin
from numpy import sqrt
from numpy import tan
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 Class_Operation(object):
    """
    Class Operation
    """

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

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


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

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

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

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

    def __sub__(self, other):
        return Class_Operation('__sub_', self, other)

    def __mul__(self, other):
        return Class_Operation('__mul_', self, other)

    def DM(self, xyz_sun, filter=None):
        """
        Calculate the dispersion measure at location `xyz`
        """
        n = 1000
        try:
            xyz = self.xyz - xyz_sun
        except ValueError:
            xyz = self.xyz - xyz_sun[:, None]

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

        if filter is None:
            return quad(lambda x: self._func(xyz_sun + x*xyz, **self._params),
                        0, 1)[0]*dfinal*1000*self._ne0

        else:
            return (dfinal*1000*self._ne0 *
                    sum([quad(lambda x: self._func(xyz_sun + x*xyz,
                                                   **self._params),
                              ii/n, (ii+1)/n)[0] for ii in range(n)
                         if filter(xyz_sun + (2*ii + 1)*xyz/n/2)]))

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

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

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

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

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

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

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


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

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

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

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

        return self._nelism

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

        return self._flism

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


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

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

    xyz_p = xyz/ellipsoid

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


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

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

    xyz_p = xyz/cylinder_p

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


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


class Clumps(object):
    """
    """

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

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

    @property
    def xyz(self):
        """
        """
        try:
            return self._xyz
        except AttributeError:
            self._xyz = self.get_xyz()
        return self._xyz

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

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

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

    @property
    def radius(self):
        """
        Radius of the clump (kpc)
        """
        return self._data['rc']

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

    @property
    def edge(self):
        """
        Clump edge
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        return self._data['edge']

    def get_xyz(self, rsun=8.5):
        """
        """
        # xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
        #                distance=self.distance,
        #                z_sun = z_sun*us.kpc,
        #                unit="deg, deg, kpc").galactocentric.
        #                                      cartesian.xyz.value
        # return xyz

        slc = sin(self.gl/180*pi)
        clc = cos(self.gl/180*pi)
        sbc = sin(self.gb/180*pi)
        cbc = cos(self.gb/180*pi)
        rgalc = self.distance*cbc
        xc = rgalc*slc
        yc = rsun-rgalc*clc
        zc = self.distance*sbc
        return np.array([xc, yc, zc])

    def clump_factor(self, xyz):
        """
        Clump edge
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        if xyz.ndim == 1:
            xyz = xyz[:, None] - self.xyz
        else:
            xyz = xyz[:, :, None] - self.xyz[:, None, :]

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

    def ne_clumps(self, xyz):
        """
        The contribution of the clumps to the free
        electron density at x, y, z = `xyz`
        """
        return np.sum(self.clump_factor(xyz)*self.ne*self.use_clump, axis=-1)


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

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

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

    @property
    def xyz(self):
        """
        """
        try:
            return self._xyz
        except AttributeError:
            self._xyz = self.get_xyz()
        return self._xyz

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

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

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

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

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

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

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

    @property
    def edge(self):
        """
        Void edge
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        return self._data['edge']

    def get_xyz(self, rsun=8.5):
        """
        """
        # xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
        #                distance=self.distance,
        #                z_sun = z_sun*us.kpc,
        #                unit="deg, deg, kpc").galactocentric.
        #                cartesian.xyz.value
        # return xyz

        slc = sin(self.gl/180*pi)
        clc = cos(self.gl/180*pi)
        sbc = sin(self.gb/180*pi)
        cbc = cos(self.gb/180*pi)
        rgalc = self.distance*cbc
        xc = rgalc*slc
        yc = rsun-rgalc*clc
        zc = self.distance*sbc
        return np.array([xc, yc, zc])

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

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

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

    def ne_voids(self, xyz):
        """
        The contribution of the clumps to the free
        electron density at x, y, z = `xyz`
        """
        return np.sum(self.void_factor(xyz)*self.ne*self.use_void, axis=-1)


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

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

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

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


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

benbaror / ne2001

Push — master ( c9e5b6...b5ca57 )

galactic_to_galactocentric() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like