NEobject.dist() - Code Metrics - Inspection of "Implemented DM -> dist" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( d8c699...d287d3 )

by Ben

created 2016-12-09 20:24 UTC

NEobject.dist() A

↳ Parent: NEobject

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	2
c	0
b	0
f	0
dl	0
loc	20
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 cumtrapz
from scipy.integrate import quad

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._params = params
        self._fparam = params.pop('F')
        self._ne0 = params.pop('e_density')
        try:
            self._func = func(**params)
        except TypeError:
            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, integrator=quad, step_size=0.001,
           *arg, **kwargs):
        """
        Calculate the dispersion measure at location `xyz`
        """
        xyz = xyz - xyz_sun

        dfinal = sqrt(np.sum(xyz**2, axis=0))
        if integrator.__name__ is 'quad':
            return integrator(lambda x: self.ne(xyz_sun + x*xyz),
                              0, 1, *arg, epsrel=epsrel, epsabs=epsabs,
                              **kwargs)[0]*dfinal*1000
        else:   # Assuming sapling integrator
            nsamp = max(1000, dfinal/step_size)
            x = np.linspace(0, 1, nsamp)
            xyz = xyz_sun[:, None] + x*xyz[:, None]
            ne = self.ne(xyz)
            return integrator(ne)*dfinal*1000/nsamp

    def dist(self, l, b, DM, rsun=8.5, step_size=0.0001):
        """
        Estimate the distance to an object with dispersion measure `DM`
        located at the direction `l ,b'
        """

        # Initial guess
        dist0 = DM/PARAMS['thick_disk']['e_density']/1000

        while self.DM(galactic_to_galactocentric(l, b, dist0, rsun)) < DM:
            dist0 *= 2

        nsamp = max(1000, dist0/step_size)
        d_samp = np.linspace(0, dist0, nsamp)
        step_size = np.diff(d_samp[:2])
        ne_samp = self.ne(galactic_to_galactocentric(l, b, d_samp, rsun))
        d_samp = (d_samp[1:] + d_samp[:-1])/2
        dm_samp = cumtrapz(ne_samp)*1000*step_size
        dist = np.interp(DM, dm_samp, d_samp)
        return dist

    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 np.array(self._data['flag']) == 0

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

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

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

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

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

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

    @lzproperty
    def ne0_use(self):
        """
        Electron density of each clump (cm^{-3})
        """
        return self.ne0*self.use_clump

    @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 np.array(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_use, 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 np.array(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 np.array(self._data['l'])

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

    @lzproperty
    def distance(self):
        """
        Distance from the sun (kpc)
        """
        return np.array(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 np.array(self._data['nev'])

    @lzproperty
    def ne0_use(self):
        """
        Electron density of each void (cm^{-3})
        """
        return self.ne0*self.use_void

    @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 np.array(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 + (q2 <= 5)*(1-self.edge)*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_use, 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)


class Ellipsoid(object):
    """
    """

    def __init__(self, center, ellipsoid, theta):
        """
        """
        self.center = center
        self.ellipsoid = ellipsoid
        self.theta = theta

    @lzproperty
    def transform(self):
        "Rotation and rescaling matrix"
        return (rotation(self.theta, -1).T/self.ellipsoid).T

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

        xyz = self.transform.dot(xyz)

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


def in_ellipsoid(center, ellipsoid, theta):
    return Ellipsoid(center, ellipsoid, theta).in_ellipsoid


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)


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 cumtrapz
14			from scipy.integrate import quad
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._params = params
167			self._fparam = params.pop('F')
168			self._ne0 = params.pop('e_density')
169			try:
170			self._func = func(**params)
171			except TypeError:
172			self._func = partial(func, **params)
173			self._params = params
174
175			def DM(self, xyz, xyz_sun=np.array([0, 8.5, 0]),
176			epsrel=1e-4, epsabs=1e-6, integrator=quad, step_size=0.001,
177			arg, *kwargs):
178			"""
179			Calculate the dispersion measure at location `xyz`
180			"""
181			xyz = xyz - xyz_sun
182
183			dfinal = sqrt(np.sum(xyz**2, axis=0))
184			if integrator.__name__ is 'quad':
185			return integrator(lambda x: self.ne(xyz_sun + x*xyz),
186			0, 1, *arg, epsrel=epsrel, epsabs=epsabs,
187			*kwargs)[0]dfinal*1000
188			else: # Assuming sapling integrator
189			nsamp = max(1000, dfinal/step_size)
190			x = np.linspace(0, 1, nsamp)
191			xyz = xyz_sun[:, None] + x*xyz[:, None]
192			ne = self.ne(xyz)
193			return integrator(ne)dfinal1000/nsamp
194
195			def dist(self, l, b, DM, rsun=8.5, step_size=0.0001):
196			"""
197			Estimate the distance to an object with dispersion measure `DM`
198			located at the direction `l ,b'
199			"""
200
201			# Initial guess
202			dist0 = DM/PARAMS['thick_disk']['e_density']/1000
203
204			while self.DM(galactic_to_galactocentric(l, b, dist0, rsun)) < DM:
205			dist0 *= 2
206
207			nsamp = max(1000, dist0/step_size)
208			d_samp = np.linspace(0, dist0, nsamp)
209			step_size = np.diff(d_samp[:2])
210			ne_samp = self.ne(galactic_to_galactocentric(l, b, d_samp, rsun))
211			d_samp = (d_samp[1:] + d_samp[:-1])/2
212			dm_samp = cumtrapz(ne_samp)1000step_size
213			dist = np.interp(DM, dm_samp, d_samp)
214			return dist
215
216			def ne(self, *args):
217			return self.electron_density(*args)
218
219			def electron_density(self, xyz):
220			"Electron density at the location `xyz`"
221			return self._ne0*self._func(xyz)
222
223			@property
224			def F(self, xyz):
225			"Fluctuation parameter"
226			return (self.ne(xyz) > 0)*self._fparam
227
228
229			class NEcombine(NEobject):
230			"""
231			"""
232
233			def __init__(self, object1, object2):
234			"""
235			"""
236			self._object1 = object1
237			self._object2 = object2
238
239			def electron_density(self, *args):
240			ne1 = self._object1.ne(*args)
241			ne2 = self._object2.ne(*args)
242			return ne1 + ne2*(ne1 <= 0)
243
244
245			class LocalISM(NEobject):
246			"""
247			Calculate the contribution of the local ISM
248			to the free electron density at x, y, z = `xyz`
249			"""
250
251			def __init__(self, **params):
252			"""
253			"""
254			self.ldr = NEobject(in_ellipsoid, **params['ldr'])
255			self.lsb = NEobject(in_ellipsoid, **params['lsb'])
256			self.lhb = NEobject(in_cylinder, **params['lhb'])
257			self.loop_in = NEobject(in_half_sphere, **params['loop_in'])
258			self.loop_out = NEobject(in_half_sphere, **params['loop_out'])
259
260			self.loop = NEcombine(self.loop_in, self.loop_out)
261			self._lism = NEcombine(self.lhb,
262			NEcombine(self.loop,
263			NEcombine(self.lsb, self.ldr)))
264
265			def electron_density(self, xyz):
266			"""
267			Calculate the contribution of the local ISM to the free
268			electron density at x, y, z = `xyz`
269			"""
270			return self._lism.ne(xyz)
271
272
273			class Clumps(NEobject):
274			"""
275			"""
276
277			def __init__(self, clumps_file=None):
278			"""
279			"""
280			if not clumps_file:
281			this_dir, _ = os.path.split(__file__)
282			clumps_file = os.path.join(this_dir, "data", "neclumpN.NE2001.dat")
283			self._data = Table.read(clumps_file, format='ascii')
284
285			@lzproperty
286			def use_clump(self):
287			"""
288			"""
289			return np.array(self._data['flag']) == 0
290
291			@lzproperty
292			def xyz(self):
293			"""
294			"""
295			return self.get_xyz()
296
297			@lzproperty
298			def gl(self):
299			"""
300			Galactic longitude (deg)
301			"""
302			return np.array(self._data['l'])
303
304			@lzproperty
305			def gb(self):
306			"""
307			Galactic latitude (deg)
308			"""
309			return np.array(self._data['b'])
310
311			@lzproperty
312			def distance(self):
313			"""
314			Distance from the sun (kpc)
315			"""
316			return np.array(self._data['dc'])
317
318			@lzproperty
319			def radius2(self):
320			"""
321			Radius of the clump (kpc)
322			"""
323			return np.array(self._data['rc']**2)
324
325			@lzproperty
326			def ne0(self):
327			"""
328			Electron density of each clump (cm^{-3})
329			"""
330			return np.array(self._data['nec'])
331
332			@lzproperty
333			def ne0_use(self):
334			"""
335			Electron density of each clump (cm^{-3})
336			"""
337			return self.ne0*self.use_clump
338
339			@lzproperty
340			def edge(self):
341			"""
342			Clump edge
343			0 => use exponential rolloff out to 5 clump radii
344			1 => uniform and truncated at 1/e clump radius
345			"""
346			return np.array(self._data['edge'])
347
348			def get_xyz(self, rsun=8.5):
349			"""
350			"""
351			# xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
352			# distance=self.distance,
353			# z_sun = z_sun*us.kpc,
354			# unit="deg, deg, kpc").galactocentric.
355			# cartesian.xyz.value
356			# return xyz
357			return galactic_to_galactocentric(l=self.gl, b=self.gb,
358			distance=self.distance, rsun=rsun)
359
360			def clump_factor(self, xyz):
361			"""
362			Clump edge
363			0 => use exponential rolloff out to 5 clump radii
364			1 => uniform and truncated at 1/e clump radius
365			"""
366			if xyz.ndim == 1:
367			return clump_factor(xyz, self.xyz, self.radius2, self.edge)
368			else:
369			xyz = xyz[:, :, None] - self.xyz[:, None, :]
370
371			q2 = (np.sum(xyz**2, axis=0) /
372			self.radius2)
373			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
374			# TODO: check this
375			q5 = (q2 <= 5)*(self.edge == 0)
376			res = np.zeros_like(q2)
377			res[(q2 <= 1)*(self.edge == 1)] = 1
378			res[q5] = exp(-q2[q5])
379			return res
380
381			def electron_density(self, xyz):
382			"""
383			The contribution of the clumps to the free
384			electron density at x, y, z = `xyz`
385			"""
386			return np.sum(self.clump_factor(xyz)*self.ne0_use, axis=-1)
387
388
389			class Voids(NEobject):
390			"""
391			"""
392
393			def __init__(self, voids_file=None):
394			"""
395			"""
396			if not voids_file:
397			this_dir, _ = os.path.split(__file__)
398			voids_file = os.path.join(this_dir, "data", "nevoidN.NE2001.dat")
399			self._data = Table.read(voids_file, format='ascii')
400
401			@lzproperty
402			def use_void(self):
403			"""
404			"""
405			return np.array(self._data['flag'] == 0)
406
407			@lzproperty
408			def xyz(self):
409			"""
410			"""
411			return self.get_xyz()
412
413			@lzproperty
414			def xyz_rot(self):
415			"""
416			"""
417			return np.array([R.dot(xyzi) for R,
418			xyzi in zip(self.rotation, self.xyz.T)]).T
419
420			@lzproperty
421			def gl(self):
422			"""
423			Galactic longitude (deg)
424			"""
425			return np.array(self._data['l'])
426
427			@lzproperty
428			def gb(self):
429			"""
430			Galactic latitude (deg)
431			"""
432			return np.array(self._data['b'])
433
434			@lzproperty
435			def distance(self):
436			"""
437			Distance from the sun (kpc)
438			"""
439			return np.array(self._data['dv'])
440
441			@lzproperty
442			def ellipsoid_abc(self):
443			"""
444			Void axis
445			"""
446			return np.array([self._data['aav'],
447			self._data['bbv'],
448			self._data['ccv']])
449
450			@lzproperty
451			def rotation(self):
452			"""
453			Rotation and rescaling matrix
454			"""
455			return np.array([
456			(rotation(thetaz*pi/180, -1).dot(
457			rotation(thetay*pi/180, 1)).T/abc).T
458			for thetaz, thetay, abc
459			in zip(self._data['thvz'], self._data['thvy'],
460			self.ellipsoid_abc.T)
461			])
462
463			@lzproperty
464			def ne0(self):
465			"""
466			Electron density of each void (cm^{-3})
467			"""
468			return np.array(self._data['nev'])
469
470			@lzproperty
471			def ne0_use(self):
472			"""
473			Electron density of each void (cm^{-3})
474			"""
475			return self.ne0*self.use_void
476
477			@lzproperty
478			def edge(self):
479			"""
480			Void edge
481			0 => use exponential rolloff out to 5 clump radii
482			1 => uniform and truncated at 1/e clump radius
483			"""
484			return np.array(self._data['edge'])
485
486			def get_xyz(self, rsun=8.5):
487			"""
488			"""
489			# xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
490			# distance=self.distance,
491			# z_sun = z_sun*us.kpc,
492			# unit="deg, deg, kpc").galactocentric.
493			# cartesian.xyz.value
494			# return xyz
495			return galactic_to_galactocentric(l=self.gl, b=self.gb,
496			distance=self.distance, rsun=rsun)
497
498			def void_factor(self, xyz):
499			"""
500			Clump edge
501			0 => use exponential rolloff out to 5 clump radii
502			1 => uniform and truncated at 1/e clump radius
503			"""
504			xyz = (self.rotation.dot(xyz).T - self.xyz_rot).T
505
506			q2 = np.sum(xyz**2, axis=1).T
507			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
508			# TODO: check this
509			return (q2 <= 1)self.edge + (q2 <= 5)(1-self.edge)*exp(-q2)
510
511			def electron_density(self, xyz):
512			"""
513			The contribution of the clumps to the free
514			electron density at x, y, z = `xyz`
515			"""
516			return np.sum(self.void_factor(xyz)*self.ne0_use, axis=-1)
517
518
519			class ElectronDensity(NEobject):
520			"""
521			A class holding all the elements which contribute to free electron density
522			"""
523
524			def __init__(self, r_sun=8.5, clumps_file=None, voids_file=None,
525			**params):
526			"""
527			"""
528			self._params = params
529			self._thick_disk = NEobject(thick_disk, r_sun=r_sun,
530			**params['thick_disk'])
531			self._thin_disk = NEobject(thin_disk, **params['thin_disk'])
532			self._galactic_center = NEobject(gc, **params['galactic_center'])
533			self._lism = LocalISM(**params)
534			self._clumps = Clumps(clumps_file=clumps_file)
535			self._voids = Voids(voids_file=voids_file)
536			self._combined = (NEcombine(self._voids,
537			NEcombine(self._lism,
538			self._thick_disk +
539			self._thin_disk +
540			self._galactic_center)) +
541			self._clumps)
542
543			def electron_density(self, xyz):
544			return self._combined.ne(xyz)
545
546
547			class Ellipsoid(object):
548			"""
549			"""
550
551			def __init__(self, center, ellipsoid, theta):
552			"""
553			"""
554			self.center = center
555			self.ellipsoid = ellipsoid
556			self.theta = theta
557
558			@lzproperty
559			def transform(self):
560			"Rotation and rescaling matrix"
561			return (rotation(self.theta, -1).T/self.ellipsoid).T
562
563			def in_ellipsoid(self, xyz):
564			"""
565			Test if xyz in the ellipsoid
566			Theta in radians
567			"""
568			try:
569			xyz = xyz - self.center
570			except ValueError:
571			xyz = xyz - self.center[:, None]
572
573			xyz = self.transform.dot(xyz)
574
575			return np.sum(xyz**2, axis=0) <= 1
576
577
578			def in_ellipsoid(center, ellipsoid, theta):
579			return Ellipsoid(center, ellipsoid, theta).in_ellipsoid
580
581
582			def in_cylinder(xyz, center, cylinder, theta):
583			"""
584			Test if xyz in the cylinder
585			Theta in radians
586			"""
587			xyz0 = xyz.copy()
588			try:
589			xyz = xyz - center
590			except ValueError:
591			xyz = xyz - center[:, None]
592			cylinder = np.vstack([cylinder]*xyz.shape[-1]).T
593			xyz[1] -= tan(theta)*xyz0[-1]
594
595			cylinder_p = cylinder.copy()
596			z_c = (center[-1] - cylinder[-1])
597			izz = (xyz0[-1] <= 0)*(xyz0[-1] >= z_c)
598			cylinder_p[0] = (0.001 +
599			(cylinder[0] - 0.001) *
600			(1 - xyz0[-1]/z_c))izz + cylinder[0](~izz)
601			xyz_p = xyz/cylinder_p
602
603			return (xyz_p[0]2 + xyz_p[1]2 <= 1) * (xyz_p[-1]**2 <= 1)
604
605
606			def in_half_sphere(xyz, center, radius):
607			"Test if `xyz` in the sphere with radius r_sphere centerd at `xyz_center`"
608			xyz0 = xyz.copy()
609			try:
610			xyz = xyz - center
611			except ValueError:
612			xyz = xyz - center[:, None]
613			distance = sqrt(np.sum(xyz**2, axis=0))
614			return (distance <= radius)*(xyz0[-1] >= 0)
615
616
617			def clump_factor(xyz, xyz0, r2, edge):
618			"""
619			Clump edge
620			0 => use exponential rolloff out to 5 clump radii
621			1 => uniform and truncated at 1/e clump radius
622			"""
623			xyz = (xyz - xyz0.T).T
624
625			q2 = (xyz[0]2 + xyz[1]2 + xyz[2]**2) / r2
626			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
627			# TODO: check this
628			q5 = (q2 <= 5)*(edge == 0)
629			res = np.zeros_like(q2)
630			res[(q2 <= 1)*(edge == 1)] = 1
631			res[q5] = exp(-q2[q5])
632			return res
633

benbaror / ne2001

Push — master ( d8c699...d287d3 )

NEobject.dist() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like