set_xyz_sun() - Code Metrics - Inspection of "Defined Sun's location as global variable" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( dfa698...603af1 )

by Ben

created 2016-12-15 01:46 UTC

set_xyz_sun() A

↳ Parent: Project

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

"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}}

XYZ_SUN = np.array([0, 8.5, 0])
RSUN = sqrt(XYZ_SUN[0]**2 + XYZ_SUN[1]**2)


def set_xyz_sun(xyz_sun):
    global XYZ_SUN
    global RSUN

    XYZ_SUN = xyz_sun
    RSUN = sqrt(XYZ_SUN[0]**2 + XYZ_SUN[1]**2)


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

benbaror / ne2001

Push — master ( dfa698...603af1 )

set_xyz_sun() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like