Source code for irfpy.util.vdf

""" Represents velocity distribution function.

Implementation of the velocity distribution function.
A similar functionality (differential flux, count rate) is implemented
in :mod:`irfpy.util.energyspectrum` module.
:mod:`irfpy.util.fluxtools` provides the conversion in between.

.. codeauthor:: Yoshifumi Futaana

.. autosummary::
    
    VelocityDistributionGrid
    maxwellVelocityDistributionGrid
"""
import numpy as np

from irfpy.util.polarframe import RegularGridVolumeContainer
from irfpy.util import maxwell


class VelocityDistributionGrid(RegularGridVolumeContainer):
    """Represents the velocity distribution function on grids in polar frame

    This object represents the velocity distribution function
    on grids in the polar frame (velocity-theta-phi).
    Therefore, it is useful to represents a data from electrostatic analyzer.
    
    This can be converted to the object of
    :class:`irfpy.energyspectrum.DifferentialFluxGrid`
    and :class:`irfpy.energyspectrum.CountRateGrid`
    using functionality in :mod:`irfpy.fluxtools`.
    """

    UNIT = "s^3 / m^6"
    """ Unit of the VDF
    """

    def __init__(self, vdf, velocity_grid, theta_grid, phi_grid,
                    velocity_grid_size='log', theta_grid_size='linear', phi_grid_size='linear',
                    mass=1.67262178e-27,):
        """ Initilizer of velocity distribution function..

        :param vdf: Velocity distribution function. Unit should be MKSA (s^3/m^6). Should be a shape of (ne, nt, np).
        :param velocity_grid: Velocity grid in MKSA (m/s).  Ascending order.
        :param theta_grid: Theta (polar angle) angle array in radians.  Ascending order.  0 to pi.
        :param phi_grid: Phi (azimuth) angle array in radians. Ascending order.
        :keyword mass: Mass in kg.
        :keyword velocity_grid_size: Specify the velocity grid size.
            Use 'linear', 'log', or np.array with the same shape as *velocity_grid*.
        """
        RegularGridVolumeContainer.__init__(self,
            np.ma.masked_array(vdf),
            velocity_grid, theta_grid, phi_grid,
            rgridsize=velocity_grid_size, tgridsize=theta_grid_size, pgridsize=phi_grid_size)

        self.mass = mass


def maxwellVelocityDistributionGrid(n, vb, vth, vlist, tlist, plist):
    """ Create :class:`VelocityDistributionGrid` object with Maxwell distribution.

    :param n: Density in m^-3
    :param vb: Bulk velocity in m/s. (3,) shaped list.
    :param vth: Thermal velocity in m/s. Scalar.
    :param vlist: List of velocity grid to be passed to :class:`VelocityDistributionGrid`
    :param tlist: List of theta grid to be passed to :class:`VelocityDistributionGrid`
    :param plist: List of phi grid to be passed to :class:`VelocityDistributionGrid`
    :return: :class:`VelocityDistributionGrid` object
    """
    fmaxwell = maxwell.mkfunc(n, vb, vth)

    ve, t, p = np.meshgrid(vlist, tlist, plist, indexing='ij')
    vx = ve * np.sin(t) * np.cos(p)    # (Ne, Nt, Np) shape
    vy = ve * np.sin(t) * np.sin(p)
    vz = ve * np.cos(t)

    v = np.array([vx, vy, vz])    # (3, Ne, Nt, Np) shape.
    fval = fmaxwell(v)

    f = VelocityDistributionGrid(fval, vlist, tlist, plist)

    return f