NEobject.F() - Code Metrics - Inspection of "Simplified NEobject operations" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( 954152...dfa698 )

by Ben

created 2016-12-15 01:18 UTC

NEobject.F() A

↳ Parent: NEobject

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
c	0
b	0
f	0
dl	0
loc	4
rs	10

"Free electron density model"
from __future__ import division
import os
from builtins import super
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 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},

    'sun_location': [0, 8.5, 0]}


def thick_disk(xyz, rsun, 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(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`
    """
    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(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, xyz_sun=PARAMS['sun_location'],
           epsrel=1e-4, epsabs=1e-6, integrator=quad, step_size=0.001,
           *arg, **kwargs):
        """
        Calculate the dispersion measure at location `xyz`
        """
        xyz = galactic_to_galactocentric(l, b, d, [0, 0, 0])

        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 + 1)
            xyz = galactic_to_galactocentric(l, b, x*dfinal, xyz_sun)
            ne = self.ne(xyz)
            return integrator(ne)*dfinal*1000*x[1]

    def dist(self, l, b, DM, xyz_sun=PARAMS['sun_location'], step_size=0.001):
        """
        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(l, b, dist0) < 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=xyz_sun))
        dm_samp = cumtrapz(ne_samp, dx=d_samp[1])*1000
        return np.interp(DM, dm_samp, d_samp[1:])

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

    @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 _ne0_use(self):
        """
        """
        return self.ne0*self.use_flag

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

benbaror / ne2001

Push — master ( 954152...dfa698 )

NEobject.F() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like