Clumps.xyz() - Code Metrics - Inspection of "Clumps and Voids inherent NEojects for DRY" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( 5148a4...954152 )

by Ben

created 2016-12-11 05:01 UTC

Clumps.xyz() A

↳ Parent: NEobjects

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes	3
Bugs	0	Features	0

Metric	Value
cc	1
c	3
b	0
f	0
dl	0
loc	5
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 ClassOperation
from .utils import galactic_to_galactocentric
from .utils import lzproperty
from .utils import rotation

# import astropy.units as us
# from astropy.coordinates import SkyCoord

# Configuration
# TODO: use to config file
# input parameters for large-scale components of NE2001 30 June '02
# flags = {'wg1': 1,
#          'wg2': 1,
#          'wga': 1,
#          'wggc': 1,
#          'wglism': 1,
#          'wgcN': 1,
#          'wgvN': 1}

# solar_params = {'Rsun': 8.3}

# spiral_arms_params = {'na': 0.028,
#                       'ha': 0.23,
#                       'wa': 0.65,
#                       'Aa': 10.5,
#                       'Fa': 5,
#                       'narm1': 0.5,
#                       'narm2': 1.2,
#                       'narm3': 1.3,
#                       'narm4': 1.0,
#                       'narm5': 0.25,
#                       'warm1': 1.0,
#                       'warm2': 1.5,
#                       'warm3': 1.0,
#                       'warm4': 0.8,
#                       'warm5': 1.0,
#                       'harm1': 1.0,
#                       'harm2': 0.8,
#                       'harm3': 1.3,
#                       'harm4': 1.5,
#                       'harm5': 1.0,
#                       'farm1': 1.1,
#                       'farm2': 0.3,
#                       'farm3': 0.4,
#                       'farm4': 1.5,
#                       'farm5': 0.3}

PARAMS = {
    'thick_disk': {'e_density': 0.033/0.97,
                   'height': 0.97,
                   'radius': 17.5,
                   'F': 0.18},

    'thin_disk': {'e_density': 0.08,
                  'height': 0.15,
                  'radius': 3.8,
                  'F': 120},

    'galactic_center': {'e_density': 10.0,
                        'center': np.array([-0.01, 0.0, -0.020]),
                        'radius': 0.145,
                        'height': 0.026,
                        'F': 0.6e5},

    'ldr': {'ellipsoid': np.array([1.50, .750, .50]),
            'center': np.array([1.36, 8.06, 0.0]),
            'theta': -24.2*pi/180,
            'e_density': 0.012,
            'F': 0.1},

    'lsb': {'ellipsoid': np.array([1.050, .4250, .3250]),
            'center': np.array([-0.75, 9.0, -0.05]),
            'theta': 139.*pi/180,
            'e_density': 0.016,
            'F': 0.01},

    'lhb': {'cylinder': np.array([.0850, .1000, .330]),
            'center': np.array([0.01, 8.45, 0.17]),
            'theta': 15*pi/180,
            'e_density': 0.005,
            'F': 0.01},

    'loop_in': {'center': np.array([-0.045, 8.40, 0.07]),
                'radius': 0.120,
                'e_density': 0.0125,
                'F': 0.2},

    'loop_out': {'center': np.array([-0.045, 8.40, 0.07]),
                 'radius': 0.120 + 0.060,
                 'e_density': 0.0125,
                 'F': 0.01},

    '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(ClassOperation):
    """
    A general electron density object
    """

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

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

    def DM(self, 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)

    @property
    def F(self, xyz):
        "Fluctuation parameter"
        return (self.ne(xyz) > 0)*self._fparam


class NEcombine(NEobject):
    """
    """

    def __init__(self, object1, object2):
        """
        """
        self._object1 = object1
        self._object2 = object2

    def electron_density(self, *args):
        ne1 = self._object1.ne(*args)
        ne2 = self._object2.ne(*args)
        return ne1 + ne2*(ne1 <= 0)


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

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

        self.loop = NEcombine(self.loop_in, self.loop_out)
        self._lism = NEcombine(self.lhb,
                               NEcombine(self.loop,
                                         NEcombine(self.lsb, self.ldr)))

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


class 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 = (NEcombine(self._voids,
                                    NEcombine(self._lism,
                                              self._thick_disk +
                                              self._thin_disk +
                                              self._galactic_center)) +
                          self._clumps)

    def electron_density(self, xyz):
        return self._combined.ne(xyz)


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

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

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

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

        xyz = self.transform.dot(xyz)

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


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


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

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

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


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


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

benbaror / ne2001

Push — master ( 5148a4...954152 )

Clumps.xyz() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like