NEobject.dist() - Code Metrics - Inspection of "Merge pull request #7 from benbaror/units" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( 233f70...185362 )

unknown

created 2017-03-03 14:12 UTC

NEobject.dist() B

↳ Parent: NEobject

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	32
rs	8.8571

"Free electron density model"
from __future__ import division
import os
from builtins import super
from functools import partial
import pdb


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 astropy import units as u

from .utils import galactic_to_galactocentric
from .utils import lzproperty
from .utils import rotation
from .utils import parse_lbd
from .utils import parse_DM


# 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 rad3d2(xyz):
    return xyz[0]**2 + xyz[1]**2 + xyz[-1]**2


def rad2d2(xyz):
    return xyz[0]**2 + xyz[1]**2


def matmul(a, b):
    try:
        return a.__matmul__(b)
    except AttributeError:
        return np.matmul(a, b)



# Units
DM_unit = u.pc / u.cm**3
d_unit = u.kpc

# Sun
XYZ_SUN = np.array([0, 8.5, 0])
RSUN = sqrt(rad2d2(XYZ_SUN))

def set_xyz_sun(xyz_sun):
    global XYZ_SUN
    global RSUN

    XYZ_SUN = xyz_sun
    RSUN = sqrt(rad2d2(XYZ_SUN))


def thick_disk(xyz, radius, height):
    """
    Calculate the contribution of the thick disk to the free electron density
     at x, y, z = `xyz`
    """
    r_ratio = sqrt(rad2d2(xyz))/radius
    return (cos(r_ratio*pi/2)/cos(RSUN*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`
    """
    rad2 = sqrt(rad2d2(xyz))
    dens = np.zeros_like(rad2)
    zeros = np.any([np.abs(radius-rad2) > 20., np.abs(xyz[-1]) > 40],axis=0)  # Avoid floating point
    # Ok
    ok = ~zeros
    dens[ok] = (exp(-(radius - rad2[ok])**2/1.8**2) /
            cosh(xyz[-1][ok]/height)**2)  # Why 1.8?
    # Return
    return dens


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_ratio2 = rad2d2(xyz)/radius**2

    # ????
    # 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_ratio2 + (xyz[-1]/height)**2 < 1)*(r_ratio2 <= 1)


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

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

        Arguments:
        - `xyz`: Location where the electron density is calculated
        - `func`: Electron density function
        - `**params`: Model parameter
        """
        self._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 __add__(self, other):
        return Add(self, other)

    def __or__(self, other):
        return OR(self, other)

    def DM(self, l, b, d,
           epsrel=1e-4, epsabs=1e-6, integrator=quad, step_size=0.001,
           *arg, **kwargs):
        """ Calculate the dispersion measure towards direction l,b

        Parameters
        ----------
        l : float or Angle
          Galactic longitude; assumed deg if unitless
        b : float
          Galactic latitude; assumed deg if unitless
        d : float or Quantity
          Distance to source; assumed kpc if unitless
        epsrel : float, optional
        epsabs : float, optional
        integrator : method
        step_size : float, optional

        Returns
        -------
        DM : Quantity
          Dispersion Measure with units pc cm**-3

        """
        # Convert to floats
        l,b,d = parse_lbd(l,b,d)
        #
        xyz = galactic_to_galactocentric(l, b, d, [0, 0, 0])

        dfinal = sqrt(rad3d2(xyz))
        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 * DM_unit
        else:   # Assuming sapling integrator
            nsamp = max(1000, dfinal/step_size)
            x = np.linspace(0, 1, nsamp + 1)
            xyz = galactic_to_galactocentric(l, b, x*dfinal, XYZ_SUN)
            ne = self.ne(xyz)
            return integrator(ne)*dfinal*1000*x[1] * DM_unit

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

        Parameters
        ----------
        l : float or Angle
          Galactic longitude; assumed deg if unitless
        b : float
          Galactic latitude; assumed deg if unitless
        DM : float or Quantity
          Dispersion Measure;  assumed pc cm**^-3 if unitless
        step_size

        Returns
        -------

        """
        # Parse
        DM = parse_DM(DM)

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

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

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

    def ne(self, xyz):
        "Electron density at the location `xyz`"
        return self.electron_density(xyz)

    def electron_density(self, xyz):
        "Electron density at the location `xyz`"
        return self._ne0*self._func(xyz)


class OR(NEobject):
    """
    Return A or B where A and B are instance of
    and the combined electron density is ne_A
    for all ne_A > 0 and ne_B otherwise.
    """

    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 Add(NEobject):
    """
    Return A + B where A and B are instance of
    and the combined electron density is ne_A + ne_B.
    """

    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


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 = self.loop_in | self.loop_out
        self._lism = (self.lhb |
                      (self.loop |
                       (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 NEobjects(NEobject):
    """
    Read objects from file
    """

    def __init__(self, objects_file):
        """
        """
        self._data = Table.read(objects_file, format='ascii')

    @lzproperty
    def use_flag(self):
        """
        A list of flags which determine which objects to use
        """
        return np.array(self._data['flag']) == 0

    @lzproperty
    def xyz(self):
        """
        The locations of the objects in Galactocentric coordinates (kpc)
        """
        return self.get_xyz()

    @property
    def data(self):
        return self._data[self.use_data]

    @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['dist'])

    @lzproperty
    def _radius2(self):
        """
        Radius^2 of each object (kpc)
        """
        return self.radius**2

    @lzproperty
    def radius(self):
        """
        Radius of each object (kpc)
        """
        return np.array(self._data['radius'])

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

    @lzproperty
    def edge(self):
        """
        The edge of the object
        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):
        """
        Get the location in Galactocentric coordinates
        """
        # 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,
                                          xyz_sun=XYZ_SUN)

    def _factor(self, xyz):
        """
        """
        if xyz.ndim == 1:
            return object_factor(xyz, self.xyz, self._radius2, self.edge)
        else:
            xyz = xyz[:, :, None] - self.xyz[:, None, :]

        q2 = rad3d2(xyz) / 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 object to the free
        electron density at x, y, z = `xyz`
        """
        return (self._factor(xyz)*self.ne0).sum(axis=-1)


class Clumps(NEobjects):
    """
    """

    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")
        super().__init__(clumps_file)


class Voids(NEobjects):
    """
    """

    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")
        super().__init__(voids_file)

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

    @lzproperty
    def ellipsoid_abc(self):
        """
        Void axis
        """
        return np.array([self._data['aa'],
                         self._data['bb'],
                         self._data['cc']])

    @property
    def radius(self):
        """
        """
        return 1

    @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['theta_z'], self._data['theta_y'],
                   self.ellipsoid_abc.T)
        ])

    def _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 object_factor(matmul(self.rotation, xyz),
                                 self.xyz_rot,
                                 self._radius2, self.edge)
        else:
            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 thisif xyz.ndim == 1:
            return (q2 <= 1)*self.edge + (q2 <= 5)*(1-self.edge)*exp(-q2)


class ElectronDensity(NEobject):
    """
    A class holding all the elements which contribute to free electron density
    """

    def __init__(self, clumps_file=None, voids_file=None,
                 **params):
        """
        """
        self._params = params
        self._thick_disk = NEobject(thick_disk, **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 = ((self._voids |
                          (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 = matmul(self.transform, xyz)

        return rad3d2(xyz) <= 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
    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
    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 (rad2d2(xyz_p) <= 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
    try:
        xyz = xyz - center
    except ValueError:
        xyz = xyz - center[:, None]
    distance2 = rad3d2(xyz)
    return (distance2 <= radius**2)*(xyz0[-1] >= 0)


def object_factor(xyz, xyz0, r2, edge):
    """
    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 = rad3d2(xyz) / r2
    # NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
    # TODO: check this
    q5 = (q2 <= 5)*(edge == 0)
    res = 1.0*(q2 <= 1)*(edge == 1)
    res[q5] = exp(-q2[q5])
    return res


1			"Free electron density model"
2			from __future__ import division
3			import os
4			from builtins import super
5			from functools import partial
6			import pdb
7
8
9			import numpy as np
10			from astropy.table import Table
11			from numpy import cos
12			from numpy import cosh
13			from numpy import exp
14			from numpy import pi
15			from numpy import sqrt
16			from numpy import tan
17			from scipy.integrate import cumtrapz
18			from scipy.integrate import quad
19
20			from astropy import units as u
21
22			from .utils import galactic_to_galactocentric
23			from .utils import lzproperty
24			from .utils import rotation
25			from .utils import parse_lbd
26			from .utils import parse_DM
27
28
29			# import astropy.units as us
30			# from astropy.coordinates import SkyCoord
31
32			# Configuration
33			# TODO: use to config file
34			# input parameters for large-scale components of NE2001 30 June '02
35			# flags = {'wg1': 1,
36			# 'wg2': 1,
37			# 'wga': 1,
38			# 'wggc': 1,
39			# 'wglism': 1,
40			# 'wgcN': 1,
41			# 'wgvN': 1}
42
43			# solar_params = {'Rsun': 8.3}
44
45			# spiral_arms_params = {'na': 0.028,
46			# 'ha': 0.23,
47			# 'wa': 0.65,
48			# 'Aa': 10.5,
49			# 'Fa': 5,
50			# 'narm1': 0.5,
51			# 'narm2': 1.2,
52			# 'narm3': 1.3,
53			# 'narm4': 1.0,
54			# 'narm5': 0.25,
55			# 'warm1': 1.0,
56			# 'warm2': 1.5,
57			# 'warm3': 1.0,
58			# 'warm4': 0.8,
59			# 'warm5': 1.0,
60			# 'harm1': 1.0,
61			# 'harm2': 0.8,
62			# 'harm3': 1.3,
63			# 'harm4': 1.5,
64			# 'harm5': 1.0,
65			# 'farm1': 1.1,
66			# 'farm2': 0.3,
67			# 'farm3': 0.4,
68			# 'farm4': 1.5,
69			# 'farm5': 0.3}
70
71			PARAMS = {
72			'thick_disk': {'e_density': 0.033/0.97,
73			'height': 0.97,
74			'radius': 17.5,
75			'F': 0.18},
76
77			'thin_disk': {'e_density': 0.08,
78			'height': 0.15,
79			'radius': 3.8,
80			'F': 120},
81
82			'galactic_center': {'e_density': 10.0,
83			'center': np.array([-0.01, 0.0, -0.020]),
84			'radius': 0.145,
85			'height': 0.026,
86			'F': 0.6e5},
87
88			'ldr': {'ellipsoid': np.array([1.50, .750, .50]),
89			'center': np.array([1.36, 8.06, 0.0]),
90			'theta': -24.2*pi/180,
91			'e_density': 0.012,
92			'F': 0.1},
93
94			'lsb': {'ellipsoid': np.array([1.050, .4250, .3250]),
95			'center': np.array([-0.75, 9.0, -0.05]),
96			'theta': 139.*pi/180,
97			'e_density': 0.016,
98			'F': 0.01},
99
100			'lhb': {'cylinder': np.array([.0850, .1000, .330]),
101			'center': np.array([0.01, 8.45, 0.17]),
102			'theta': 15*pi/180,
103			'e_density': 0.005,
104			'F': 0.01},
105
106			'loop_in': {'center': np.array([-0.045, 8.40, 0.07]),
107			'radius': 0.120,
108			'e_density': 0.0125,
109			'F': 0.2},
110
111			'loop_out': {'center': np.array([-0.045, 8.40, 0.07]),
112			'radius': 0.120 + 0.060,
113			'e_density': 0.0125,
114			'F': 0.01}}
115
116
117			def rad3d2(xyz):
118			return xyz[0]2 + xyz[1]2 + xyz[-1]**2
119
120
121			def rad2d2(xyz):
122			return xyz[0]2 + xyz[1]2
123
124
125			def matmul(a, b):
126			try:
127			return a.__matmul__(b)
128			except AttributeError:
129			return np.matmul(a, b)
130
131
132
133			# Units
134			DM_unit = u.pc / u.cm**3
135			d_unit = u.kpc
136
137			# Sun
138			XYZ_SUN = np.array([0, 8.5, 0])
139			RSUN = sqrt(rad2d2(XYZ_SUN))
140
141			def set_xyz_sun(xyz_sun):
142			global XYZ_SUN
143			global RSUN
144
145			XYZ_SUN = xyz_sun
146			RSUN = sqrt(rad2d2(XYZ_SUN))
147
148
149			def thick_disk(xyz, radius, height):
150			"""
151			Calculate the contribution of the thick disk to the free electron density
152			at x, y, z = `xyz`
153			"""
154			r_ratio = sqrt(rad2d2(xyz))/radius
155			return (cos(r_ratiopi/2)/cos(RSUNpi/2/radius) /
156			cosh(xyz[-1]/height)*2
157			(r_ratio < 1))
158
159
160			def thin_disk(xyz, radius, height):
161			"""
162			Calculate the contribution of the thin disk to the free electron density
163			at x, y, z = `xyz`
164			"""
165			rad2 = sqrt(rad2d2(xyz))
166			dens = np.zeros_like(rad2)
167			zeros = np.any([np.abs(radius-rad2) > 20., np.abs(xyz[-1]) > 40],axis=0) # Avoid floating point
168			# Ok
169			ok = ~zeros
170			dens[ok] = (exp(-(radius - rad2[ok])2/1.82) /
171			cosh(xyz[-1][ok]/height)**2) # Why 1.8?
172			# Return
173			return dens
174
175
176			def gc(xyz, center, radius, height):
177			"""
178			Calculate the contribution of the Galactic center to the free
179			electron density at x, y, z = `xyz`
180			"""
181			# Here I'm using the expression in the NE2001 code which is inconsistent
182			# with Cordes and Lazio 2011 (0207156v3) (See Table 2)
183			try:
184			xyz = xyz - center
185			except ValueError:
186			xyz = xyz - center[:, None]
187
188			r_ratio2 = rad2d2(xyz)/radius**2
189
190			# ????
191			# Cordes and Lazio 2011 (0207156v3) (Table 2)
192			# return ne_gc0exp(-(r2d/rgc)2 - (xyz[-1]/hgc)*2)
193			# ????
194
195			# Constant ne (form NE2001 code)
196			return (r_ratio2 + (xyz[-1]/height)*2 < 1)(r_ratio2 <= 1)
197
198
199			class NEobject(object):
200			"""
201			A general electron density object
202			"""
203
204			def __init__(self, func, **params):
205			"""
206
207			Arguments:
208			- `xyz`: Location where the electron density is calculated
209			- `func`: Electron density function
210			- `**params`: Model parameter
211			"""
212			self._params = params
213			self._fparam = params.pop('F')
214			self._ne0 = params.pop('e_density')
215			try:
216			self._func = func(**params)
217			except TypeError:
218			self._func = partial(func, **params)
219			self._params = params
220
221			def __add__(self, other):
222			return Add(self, other)
223
224			def __or__(self, other):
225			return OR(self, other)
226
227			def DM(self, l, b, d,
228			epsrel=1e-4, epsabs=1e-6, integrator=quad, step_size=0.001,
229			arg, *kwargs):
230			""" Calculate the dispersion measure towards direction l,b
231
232			Parameters
233			----------
234			l : float or Angle
235			Galactic longitude; assumed deg if unitless
236			b : float
237			Galactic latitude; assumed deg if unitless
238			d : float or Quantity
239			Distance to source; assumed kpc if unitless
240			epsrel : float, optional
241			epsabs : float, optional
242			integrator : method
243			step_size : float, optional
244
245			Returns
246			-------
247			DM : Quantity
248			Dispersion Measure with units pc cm**-3
249
250			"""
251			# Convert to floats
252			l,b,d = parse_lbd(l,b,d)
253			#
254			xyz = galactic_to_galactocentric(l, b, d, [0, 0, 0])
255
256			dfinal = sqrt(rad3d2(xyz))
257			if integrator.__name__ is 'quad':
258			return integrator(lambda x: self.ne(XYZ_SUN + x*xyz),
259			0, 1, *arg, epsrel=epsrel, epsabs=epsabs,
260			*kwargs)[0]dfinal1000 DM_unit
261			else: # Assuming sapling integrator
262			nsamp = max(1000, dfinal/step_size)
263			x = np.linspace(0, 1, nsamp + 1)
264			xyz = galactic_to_galactocentric(l, b, x*dfinal, XYZ_SUN)
265			ne = self.ne(xyz)
266			return integrator(ne)dfinal1000x[1] DM_unit
267
268			def dist(self, l, b, DM, step_size=0.001):
269			""" Estimate the distance to an object with dispersion measure `DM`
270			Located at the direction `l ,b'
271
272			Parameters
273			----------
274			l : float or Angle
275			Galactic longitude; assumed deg if unitless
276			b : float
277			Galactic latitude; assumed deg if unitless
278			DM : float or Quantity
279			Dispersion Measure; assumed pc cm**^-3 if unitless
280			step_size
281
282			Returns
283			-------
284
285			"""
286			# Parse
287			DM = parse_DM(DM)
288
289			# Initial guess
290			dist0 = DM/PARAMS['thick_disk']['e_density']/1000
291
292			while self.DM(l, b, dist0).value < DM:
293			dist0 *= 2
294
295			nsamp = max(1000, dist0/step_size)
296			d_samp = np.linspace(0, dist0, nsamp + 1)
297			ne_samp = self.ne(galactic_to_galactocentric(l, b, d_samp, XYZ_SUN))
298			dm_samp = cumtrapz(ne_samp, dx=d_samp[1])*1000
299			return np.interp(DM, dm_samp, d_samp[1:]) * d_unit
300
301			def ne(self, xyz):
302			"Electron density at the location `xyz`"
303			return self.electron_density(xyz)
304
305			def electron_density(self, xyz):
306			"Electron density at the location `xyz`"
307			return self._ne0*self._func(xyz)
308
309
310			class OR(NEobject):
311			"""
312			Return A or B where A and B are instance of
313			and the combined electron density is ne_A
314			for all ne_A > 0 and ne_B otherwise.
315			"""
316
317			def __init__(self, object1, object2):
318			"""
319			"""
320			self._object1 = object1
321			self._object2 = object2
322
323			def electron_density(self, *args):
324			ne1 = self._object1.ne(*args)
325			ne2 = self._object2.ne(*args)
326			return ne1 + ne2*(ne1 <= 0)
327
328
329			class Add(NEobject):
330			"""
331			Return A + B where A and B are instance of
332			and the combined electron density is ne_A + ne_B.
333			"""
334
335			def __init__(self, object1, object2):
336			"""
337			"""
338			self._object1 = object1
339			self._object2 = object2
340
341			def electron_density(self, *args):
342			ne1 = self._object1.ne(*args)
343			ne2 = self._object2.ne(*args)
344			return ne1 + ne2
345
346
347			class LocalISM(NEobject):
348			"""
349			Calculate the contribution of the local ISM
350			to the free electron density at x, y, z = `xyz`
351			"""
352
353			def __init__(self, **params):
354			"""
355			"""
356			self.ldr = NEobject(in_ellipsoid, **params['ldr'])
357			self.lsb = NEobject(in_ellipsoid, **params['lsb'])
358			self.lhb = NEobject(in_cylinder, **params['lhb'])
359			self.loop_in = NEobject(in_half_sphere, **params['loop_in'])
360			self.loop_out = NEobject(in_half_sphere, **params['loop_out'])
361
362			self.loop = self.loop_in \| self.loop_out
363			self._lism = (self.lhb \|
364			(self.loop \|
365			(self.lsb \| self.ldr)))
366
367			def electron_density(self, xyz):
368			"""
369			Calculate the contribution of the local ISM to the free
370			electron density at x, y, z = `xyz`
371			"""
372			return self._lism.ne(xyz)
373
374
375			class NEobjects(NEobject):
376			"""
377			Read objects from file
378			"""
379
380			def __init__(self, objects_file):
381			"""
382			"""
383			self._data = Table.read(objects_file, format='ascii')
384
385			@lzproperty
386			def use_flag(self):
387			"""
388			A list of flags which determine which objects to use
389			"""
390			return np.array(self._data['flag']) == 0
391
392			@lzproperty
393			def xyz(self):
394			"""
395			The locations of the objects in Galactocentric coordinates (kpc)
396			"""
397			return self.get_xyz()
398
399			@property
400			def data(self):
401			return self._data[self.use_data]
402
403			@lzproperty
404			def gl(self):
405			"""
406			Galactic longitude (deg)
407			"""
408			return np.array(self._data['l'])
409
410			@lzproperty
411			def gb(self):
412			"""
413			Galactic latitude (deg)
414			"""
415			return np.array(self._data['b'])
416
417			@lzproperty
418			def distance(self):
419			"""
420			Distance from the sun (kpc)
421			"""
422			return np.array(self._data['dist'])
423
424			@lzproperty
425			def _radius2(self):
426			"""
427			Radius^2 of each object (kpc)
428			"""
429			return self.radius**2
430
431			@lzproperty
432			def radius(self):
433			"""
434			Radius of each object (kpc)
435			"""
436			return np.array(self._data['radius'])
437
438			@lzproperty
439			def ne0(self):
440			"""
441			Electron density of each object (cm^{-3})
442			"""
443			return np.array(self._data['ne'])
444
445			@lzproperty
446			def edge(self):
447			"""
448			The edge of the object
449			0 => use exponential rolloff out to 5 clump radii
450			1 => uniform and truncated at 1/e clump radius
451			"""
452			return np.array(self._data['edge'])
453
454			def get_xyz(self):
455			"""
456			Get the location in Galactocentric coordinates
457			"""
458			# xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
459			# distance=self.distance,
460			# z_sun = z_sun*us.kpc,
461			# unit="deg, deg, kpc").galactocentric.
462			# cartesian.xyz.value
463			# return xyz
464			return galactic_to_galactocentric(l=self.gl, b=self.gb,
465			distance=self.distance,
466			xyz_sun=XYZ_SUN)
467
468			def _factor(self, xyz):
469			"""
470			"""
471			if xyz.ndim == 1:
472			return object_factor(xyz, self.xyz, self._radius2, self.edge)
473			else:
474			xyz = xyz[:, :, None] - self.xyz[:, None, :]
475
476			q2 = rad3d2(xyz) / self._radius2
477			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
478			# TODO: check this
479			q5 = (q2 <= 5)*(self.edge == 0)
480			res = np.zeros_like(q2)
481			res[(q2 <= 1)*(self.edge == 1)] = 1
482			res[q5] = exp(-q2[q5])
483			return res
484
485			def electron_density(self, xyz):
486			"""
487			The contribution of the object to the free
488			electron density at x, y, z = `xyz`
489			"""
490			return (self._factor(xyz)*self.ne0).sum(axis=-1)
491
492
493			class Clumps(NEobjects):
494			"""
495			"""
496
497			def __init__(self, clumps_file=None):
498			"""
499			"""
500			if not clumps_file:
501			this_dir, _ = os.path.split(__file__)
502			clumps_file = os.path.join(this_dir, "data", "neclumpN.NE2001.dat")
503			super().__init__(clumps_file)
504
505
506			class Voids(NEobjects):
507			"""
508			"""
509
510			def __init__(self, voids_file=None):
511			"""
512			"""
513			if not voids_file:
514			this_dir, _ = os.path.split(__file__)
515			voids_file = os.path.join(this_dir, "data", "nevoidN.NE2001.dat")
516			super().__init__(voids_file)
517
518			@lzproperty
519			def xyz_rot(self):
520			"""
521			Rotated xyz
522			"""
523			return np.array([R.dot(xyzi) for R,
524			xyzi in zip(self.rotation, self.xyz.T)]).T
525
526			@lzproperty
527			def ellipsoid_abc(self):
528			"""
529			Void axis
530			"""
531			return np.array([self._data['aa'],
532			self._data['bb'],
533			self._data['cc']])
534
535			@property
536			def radius(self):
537			"""
538			"""
539			return 1
540
541			@lzproperty
542			def rotation(self):
543			"""
544			Rotation and rescaling matrix
545			"""
546			return np.array([
547			(rotation(thetaz*pi/180, -1).dot(
548			rotation(thetay*pi/180, 1)).T/abc).T
549			for thetaz, thetay, abc
550			in zip(self._data['theta_z'], self._data['theta_y'],
551			self.ellipsoid_abc.T)
552			])
553
554			def _factor(self, xyz):
555			"""
556			Clump edge
557			0 => use exponential rolloff out to 5 clump radii
558			1 => uniform and truncated at 1/e clump radius
559			"""
560			if xyz.ndim == 1:
561			return object_factor(matmul(self.rotation, xyz),
562			self.xyz_rot,
563			self._radius2, self.edge)
564			else:
565			xyz = (self.rotation.dot(xyz).T - self.xyz_rot).T
566
567			q2 = np.sum(xyz**2, axis=1).T
568			# NOTE: In the original NE2001 code q2 <= 5
569			# is used instead of q <= 5.
570			# TODO: check thisif xyz.ndim == 1:
571			return (q2 <= 1)self.edge + (q2 <= 5)(1-self.edge)*exp(-q2)
572
573
574			class ElectronDensity(NEobject):
575			"""
576			A class holding all the elements which contribute to free electron density
577			"""
578
579			def __init__(self, clumps_file=None, voids_file=None,
580			**params):
581			"""
582			"""
583			self._params = params
584			self._thick_disk = NEobject(thick_disk, **params['thick_disk'])
585			self._thin_disk = NEobject(thin_disk, **params['thin_disk'])
586			self._galactic_center = NEobject(gc, **params['galactic_center'])
587			self._lism = LocalISM(**params)
588			self._clumps = Clumps(clumps_file=clumps_file)
589			self._voids = Voids(voids_file=voids_file)
590			self._combined = ((self._voids \|
591			(self._lism \|
592			(self._thick_disk +
593			self._thin_disk +
594			self._galactic_center))) +
595			self._clumps)
596
597			def electron_density(self, xyz):
598			return self._combined.ne(xyz)
599
600
601			class Ellipsoid(object):
602			"""
603			"""
604
605			def __init__(self, center, ellipsoid, theta):
606			"""
607			"""
608			self.center = center
609			self.ellipsoid = ellipsoid
610			self.theta = theta
611
612			@lzproperty
613			def transform(self):
614			"Rotation and rescaling matrix"
615			return (rotation(self.theta, -1).T/self.ellipsoid).T
616
617			def in_ellipsoid(self, xyz):
618			"""
619			Test if xyz in the ellipsoid
620			Theta in radians
621			"""
622			try:
623			xyz = xyz - self.center
624			except ValueError:
625			xyz = xyz - self.center[:, None]
626
627			xyz = matmul(self.transform, xyz)
628
629			return rad3d2(xyz) <= 1
630
631
632			def in_ellipsoid(center, ellipsoid, theta):
633			return Ellipsoid(center, ellipsoid, theta).in_ellipsoid
634
635
636			def in_cylinder(xyz, center, cylinder, theta):
637			"""
638			Test if xyz in the cylinder
639			Theta in radians
640			"""
641			xyz0 = xyz
642			try:
643			xyz = xyz - center
644			except ValueError:
645			xyz = xyz - center[:, None]
646			cylinder = np.vstack([cylinder]*xyz.shape[-1]).T
647			xyz[1] -= tan(theta)*xyz0[-1]
648			cylinder_p = cylinder
649			z_c = (center[-1] - cylinder[-1])
650			izz = (xyz0[-1] <= 0)*(xyz0[-1] >= z_c)
651			cylinder_p[0] = (0.001 +
652			(cylinder[0] - 0.001) *
653			(1 - xyz0[-1]/z_c))izz + cylinder[0](~izz)
654			xyz_p = xyz/cylinder_p
655
656			return (rad2d2(xyz_p) <= 1) * (xyz_p[-1]**2 <= 1)
657
658
659			def in_half_sphere(xyz, center, radius):
660			"Test if `xyz` in the sphere with radius r_sphere centerd at `xyz_center`"
661			xyz0 = xyz
662			try:
663			xyz = xyz - center
664			except ValueError:
665			xyz = xyz - center[:, None]
666			distance2 = rad3d2(xyz)
667			return (distance2 <= radius*2)(xyz0[-1] >= 0)
668
669
670			def object_factor(xyz, xyz0, r2, edge):
671			"""
672			edge
673			0 => use exponential rolloff out to 5 clump radii
674			1 => uniform and truncated at 1/e clump radius
675			"""
676			xyz = (xyz - xyz0.T).T
677			q2 = rad3d2(xyz) / r2
678			# NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
679			# TODO: check this
680			q5 = (q2 <= 5)*(edge == 0)
681			res = 1.0(q2 <= 1)(edge == 1)
682			res[q5] = exp(-q2[q5])
683			return res
684

benbaror / ne2001

Push — master ( 233f70...185362 )

NEobject.dist() B

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like