NEobjects.distance() - Code Metrics - benbaror/ne2001 - Measure and Improve Code Quality continuously with Scrutinizer

NEobjects.distance() A
last analyzed 2017-04-25 12:30 UTC

↳ Parent: NEobjects

Complexity

Conditions

Size

Total Lines

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
cc	1
dl	0
loc	6
rs	9.4285
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 import units as u
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 . import ne_io
from .spiral_arms import ne_spiral_arm
from .utils import galactic_to_galactocentric
from .utils import lzproperty
from .utils import matmul
from .utils import parse_DM
from .utils import parse_lbd
from .utils import rad2d2
from .utils import rad3d2
from .utils import rotation


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

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


def set_xyz_sun(xyz_sun):
    global XYZ_SUN
    global RSUN

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


def thick_disk(xyz, radius, height):
    """ Calculate the contribution of the thick disk to the free electron density
    at x, y, z = `xyz`
    Parameters
    ----------
    xyz
    radius
    height

    Returns
    -------
    dens : ndarray
      Density values

    """
    r_ratio = sqrt(rad2d2(xyz))/radius
    dens = (cos(r_ratio*pi/2)/cos(RSUN*pi/2/radius) /
            cosh(xyz[-1]/height)**2 *
            (r_ratio < 1))
    # Return
    return dens


def thin_disk(xyz, radius, height):
    """
    Calculate the contribution of the thin disk to the free electron density
     at x, y, z = `xyz`
    """
    rad2 = sqrt(rad2d2(xyz))
    dens = np.zeros_like(rad2)
    # Avoid floating point
    zeros = np.any([np.abs(radius-rad2) > 20., np.abs(xyz[-1]) > 40], axis=0)
    ok = ~zeros
    dens[ok] = (exp(-(radius - rad2[ok])**2/1.8**2) /
                cosh(xyz[-1][ok]/height)**2)  # Why 1.8?
    return dens


def gc(xyz, center, radius, height):
    """
    Calculate the contribution of the Galactic center to the free
    electron density at x, y, z = `xyz`
    """
    # Here I'm using the expression in the NE2001 code which is inconsistent
    # with Cordes and Lazio 2011 (0207156v3) (See Table 2)
    try:
        xyz = xyz - center
    except ValueError:
        xyz = xyz - center[:, None]

    r_ratio2 = rad2d2(xyz)/radius**2

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

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


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

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

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

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

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

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

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

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

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

        dfinal = sqrt(rad3d2(xyz))
        if integrator.__name__ is 'quad':
            return integrator(lambda x: self.ne(XYZ_SUN + x*xyz),
                              0, 1, *arg, epsrel=epsrel, epsabs=epsabs,
                              **kwargs)[0]*dfinal*1000 * DM_unit
        else:   # Assuming sapling integrator
            nsamp = max(1000, dfinal/step_size)
            x = np.linspace(0, 1, nsamp + 1)
            xyz = galactic_to_galactocentric(l, b, x*dfinal, XYZ_SUN)
            ne = self.ne(xyz)
            return integrator(ne)*dfinal*1000*x[1] * DM_unit

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

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

        Returns
        -------

        """
        # Parse
        DM = parse_DM(DM)

        # Initial guess
        dist0 = DM/self.params['thick_disk']['e_density']/1000

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

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

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

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


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

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

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


class Add(NEobject):
    """
    Return A + B where A and B are instance of
    and the combined electron density is ne_A + ne_B.
    """

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

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


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

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

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

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


class NEobjects(NEobject):
    """
    Read objects from file
    """

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


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

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


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

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

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

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

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

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

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

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


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

    def __init__(self, clumps_file=None, voids_file=None,
                 **params):
        """
        """
        self._params = ne_io.Params(**params)
        self._thick_disk = NEobject(thick_disk, **self.params['thick_disk'])
        self._thin_disk = NEobject(thin_disk, **self.params['thin_disk'])
        self._galactic_center = NEobject(gc, **self.params['galactic_center'])
        self._lism = LocalISM(**self.params)
        self._spiral_arms = NEobject(ne_spiral_arm,
                                     **self.params['spiral_arms'])
        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._spiral_arms +
                            self._galactic_center))) +
                          self._clumps)

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

    @property
    def params(self):
        return self._params


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

benbaror / ne2001

NEobjects.distance() A last analyzed 2017-04-25 12:30 UTC

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like

NEobjects.distance() A
last analyzed 2017-04-25 12:30 UTC