NEobjects.data() - Code Metrics - Inspection of "Improved DRY and run-time" - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

Completed

Push — master ( 603af1...643330 )

by Ben

created 2016-12-15 04:18 UTC

NEobjects.data() A

↳ Parent: NEobjects

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes

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

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


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


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


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


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


def set_xyz_sun(xyz_sun):
    global XYZ_SUN
    global RSUN

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


def thick_disk(xyz, radius, height):
    """
    Calculate the contribution of the thick disk to the free electron density
     at x, y, z = `xyz`
    """
    r_ratio = sqrt(rad2d2(xyz))/radius
    return (cos(r_ratio*pi/2)/cos(RSUN*pi/2/radius) /
            cosh(xyz[-1]/height)**2 *
            (r_ratio < 1))


def thin_disk(xyz, radius, height):
    """
    Calculate the contribution of the thin disk to the free electron density
     at x, y, z = `xyz`
    """
    rad2 = sqrt(rad2d2(xyz))
    return (exp(-(radius - rad2)**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_ratio2 = rad2d2(xyz)/radius**2

    # ????
    # Cordes and Lazio 2011 (0207156v3) (Table 2)
    # return ne_gc0*exp(-(r2d/rgc)**2 - (xyz[-1]/hgc)**2)
    # ????

    # Constant ne (form NE2001 code)
    return (r_ratio2 + (xyz[-1]/height)**2 < 1)*(r_ratio2 <= 1)


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

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

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

    def __add__(self, other):
        return Add(self, other)

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

    def DM(self, l, b, d,
           epsrel=1e-4, epsabs=1e-6, integrator=quad, step_size=0.001,
           *arg, **kwargs):
        """
        Calculate the dispersion measure at location `xyz`
        """
        xyz = galactic_to_galactocentric(l, b, d, [0, 0, 0])

        dfinal = sqrt(rad3d2(xyz))
        if integrator.__name__ is 'quad':
            return integrator(lambda x: self.ne(XYZ_SUN + x*xyz),
                              0, 1, *arg, epsrel=epsrel, epsabs=epsabs,
                              **kwargs)[0]*dfinal*1000
        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()

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

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

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

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

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

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

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

    @lzproperty
    def edge(self):
        """
        The edge of the object
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        return np.array(self._data['edge'])

    def get_xyz(self):
        """
        Get the location in Galactocentric coordinates
        """
        # xyz = SkyCoord(frame="galactic", l=self.gl, b=self.gb,
        #                distance=self.distance,
        #                z_sun = z_sun*us.kpc,
        #                unit="deg, deg, kpc").galactocentric.
        #                                      cartesian.xyz.value
        # return xyz
        return galactic_to_galactocentric(l=self.gl, b=self.gb,
                                          distance=self.distance,
                                          xyz_sun=XYZ_SUN)

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

        q2 = rad3d2(xyz) / self._radius2
        # NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
        # TODO: check this
        q5 = (q2 <= 5)*(self.edge == 0)
        res = np.zeros_like(q2)
        res[(q2 <= 1)*(self.edge == 1)] = 1
        res[q5] = exp(-q2[q5])
        return res

    def electron_density(self, xyz):
        """
        The contribution of the object to the free
        electron density at x, y, z = `xyz`
        """
        return (self._factor(xyz)*self.ne0).sum(axis=-1)


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

    def __init__(self, clumps_file=None):
        """
        """
        if not clumps_file:
            this_dir, _ = os.path.split(__file__)
            clumps_file = os.path.join(this_dir, "data", "neclumpN.NE2001.dat")
        super().__init__(clumps_file)


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

    def __init__(self, voids_file=None):
        """
        """
        if not voids_file:
            this_dir, _ = os.path.split(__file__)
            voids_file = os.path.join(this_dir, "data", "nevoidN.NE2001.dat")
        super().__init__(voids_file)

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

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

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

    @lzproperty
    def rotation(self):
        """
        Rotation and rescaling matrix
        """
        return np.array([
            (rotation(thetaz*pi/180, -1).dot(
                rotation(thetay*pi/180, 1)).T/abc).T
            for thetaz, thetay, abc
            in zip(self._data['theta_z'], self._data['theta_y'],
                   self.ellipsoid_abc.T)
        ])

    def _factor(self, xyz):
        """
        Clump edge
        0 => use exponential rolloff out to 5 clump radii
        1 => uniform and truncated at 1/e clump radius
        """
        if xyz.ndim == 1:
            return object_factor(matmul(self.rotation, xyz),
                                 self.xyz_rot,
                                 self._radius2, self.edge)
        else:
            xyz = (self.rotation.dot(xyz).T - self.xyz_rot).T

            q2 = np.sum(xyz**2, axis=1).T
            # NOTE: In the original NE2001 code q2 <= 5
            # is used instead of q <= 5.
            # TODO: check thisif xyz.ndim == 1:
            return (q2 <= 1)*self.edge + (q2 <= 5)*(1-self.edge)*exp(-q2)


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

    def __init__(self, clumps_file=None, voids_file=None,
                 **params):
        """
        """
        self._params = params
        self._thick_disk = NEobject(thick_disk, **params['thick_disk'])
        self._thin_disk = NEobject(thin_disk, **params['thin_disk'])
        self._galactic_center = NEobject(gc, **params['galactic_center'])
        self._lism = LocalISM(**params)
        self._clumps = Clumps(clumps_file=clumps_file)
        self._voids = Voids(voids_file=voids_file)
        self._combined = ((self._voids |
                          (self._lism |
                           (self._thick_disk +
                            self._thin_disk +
                            self._galactic_center))) +
                          self._clumps)

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


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

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

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

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

        xyz = matmul(self.transform, xyz)

        return rad3d2(xyz) <= 1


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


def in_cylinder(xyz, center, cylinder, theta):
    """
    Test if xyz in the cylinder
    Theta in radians
    """
    xyz0 = xyz
    try:
        xyz = xyz - center
    except ValueError:
        xyz = xyz - center[:, None]
        cylinder = np.vstack([cylinder]*xyz.shape[-1]).T
    xyz[1] -= tan(theta)*xyz0[-1]
    cylinder_p = cylinder
    z_c = (center[-1] - cylinder[-1])
    izz = (xyz0[-1] <= 0)*(xyz0[-1] >= z_c)
    cylinder_p[0] = (0.001 +
                     (cylinder[0] - 0.001) *
                     (1 - xyz0[-1]/z_c))*izz + cylinder[0]*(~izz)
    xyz_p = xyz/cylinder_p

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


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


def object_factor(xyz, xyz0, r2, edge):
    """
    edge
    0 => use exponential rolloff out to 5 clump radii
    1 => uniform and truncated at 1/e clump radius
    """
    xyz = (xyz - xyz0.T).T
    q2 = rad3d2(xyz) / r2
    # NOTE: In the original NE2001 code q2 <= 5 is used instead of q <= 5.
    # TODO: check this
    q5 = (q2 <= 5)*(edge == 0)
    res = 1.0*(q2 <= 1)*(edge == 1)
    res[q5] = exp(-q2[q5])
    return res


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

benbaror / ne2001

Push — master ( 603af1...643330 )

NEobjects.data() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like