# Source code for irfpy.util.maxwell

r''' Module collecting maxwell distribution function related functions.

In plasma physics, Maxwell distribution function appears everywhere.
In this module, Maxwell distribution functions are implemented.

.. codeauthor:: Yoshifumi Futaana

.. autosummary::

mkfunc
mkfunc_u
mk1d
mkslice
fit1d
getRandom
getRandom1D
flux_integrated
flux_precip
irfpy.util.energyspectrum.maxwellDifferentialFluxGrid
irfpy.util.vdf.maxwellVelocityDistributionGrid

The theory starts from:

.. math::

f(\vec{v})=n_0\left(\frac{m}{2\pi k}\right)^\frac{3}{2}
\frac{1}{\sqrt{\mathrm{det} \textbf{T}}}
\exp-\frac{m}{2k}(\vec{v}-\vec{v_b})^T \textbf{T}^{-1} (\vec{v}-\vec{v_b})

When the temperature is isotropic, the above temperature tensor can be
treated as a scalar, becoming

.. math::

f(\vec{v})=n_0\left(\frac{m}{2\pi kT}\right)^\frac{3}{2}
\exp-\frac{m}{2kT}(\vec{v}-\vec{v_b})^2

Using the definition of the thermal speed,

.. math::

v_\mathrm{th}^2 = \frac{kT}{m}

.. math::

f(v_x,v_y,v_z) = n_0 ( 2\pi v_\mathrm{th}^2 ) ^{-\frac{3}{2}}
\exp -\frac{ (v_x-v_{x0})^2 + (v_y-v_{y0})^2 + (v_z-v_{z0})^2 }
{2 v_\mathrm{th}^2}

Cookbook
========

Calculate Maxwellian VDF
------------------------

Making 3-D Maxwellian VDF with plasma parameters 5/cc, vx = -400km/s and vth=40km/s

>>> fsw = mkfunc(5e6, [-400e3, 0, 0], 40e3)

You can calculate the value of VDF as

>>> print('%.2e' % fsw([-400e3, 0, 0]))    # The peak
4.96e-09

Be careful with the unit system. Always use MKSA for :func:mkfunc.
Alternatively you can use :func:mkfunc_u to handle the unit.

>>> import numpy as np
>>> fsw_u = mkfunc_u(5 / u.cc, np.array([-400, 0, 0]) * u.km/u.s, 40 * u.km / u.s)
>>> fsw_peak = fsw_u(np.array([-400, 0, 0]) * u.km / u.s)
>>> print("{:.2e}".format(fsw_peak))
4.96e-09 s**3/m**6

Maxwellian flux object
----------------------

Using :class:MaxwellDistributionG, the Maxwellian VDF and DFlux
is mapped onto energy/velocity spaces.

Fitting by Maxwellian
---------------------
Fitting of Maxwellian from a dataset :math:x=\{x_i\}, y=\{y_i\} can be
calculated via :meth:fit1d function.

.. code-block:: py

x = some_array
y = some_array
(A,B,C), success = fit1d(x,y)           # A * ( 2pi B^2)  exp - (v-C)^2/B^2
ffit = mk1d(A,B,C)
yfit = ffit(x)
plot(x,y,'-', x,yfit, 'o')

Creating a random value set
---------------------------

Sometimes, you want to get set of particles whose distribution function
follows the given Maxwellian.

.. code-block:: py

vels = getRandom([-400e3, 0, 0], 40e3, n=100)

Calculating integral flux
-------------------------

The net flux of Maxwellian gas but with a given lower velocity threshold
can be calculated analitycally.
This is useful for example to calculate the Langmuir probe or Retarding Potential Analyzer performance,
and to estimate the outflowing flux (not including inward flux).

.. math::

F_{v_z>v_{zmin}} = \int_{-\infty}^{\infty}dv_x\int_{-\infty}^{\infty}dv_y\int_{v_{zmin}}^{\infty}d_z[v_zf(v_x, v_y, v_z)]

where the precipitating direction is assumed (without losing generality)
as :math:+v_z direction.

The integral will result in

.. math::

F = \frac{n_0v_{th}}{\sqrt{2\pi}}\exp-\frac{(v_{zmin}-v_{z0})^2}{2v_{th}^2} + \frac{n_0v_{z0}}{2}\mathrm{erfc}\frac{v_{zmin}-v_{z0}}{\sqrt{2}{v_{th}}}

This is implemented in :func:flux_integrated.

Some useful formula
===================

Some formula useful for Maxwellian are summarized here.

.. math::

I^{(n)}(v) \equiv \int x^n\exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv

.. math::

I^{(0)}(v) = \int \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv &=& \sigma\sqrt\frac{\pi}{2} \mathrm{erf}\Bigl(\frac{v}{\sqrt{2}\sigma}\Bigr) + C    \\
I^{(1)}(v) = \int v \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv &=& -\sigma^2\exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) + C \\
I^{(2)}(v) = \int v^2 \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv &=& \sqrt\frac{\pi}{2} \sigma^3 \mathrm{erf}\Bigl(\frac{v}{\sqrt{2}\sigma}\Bigr) - \sigma^2 v \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) + C  \\
I^{(3)}(v) = \int v^3 \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv &=& -\sigma^2(2\sigma^2+v^2)\exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) + C \\
I^{(4)}(v) = \int v^4 \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv &=& 3\sqrt\frac{\pi}{2} \sigma^5 \mathrm{erf}\Bigl(\frac{v}{\sqrt{2}\sigma}\Bigr) - \sigma^2 v (3\sigma^2+v^2) \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) + C  ~~~~~~ \textrm{To be checked}\\

Further expression:

.. math::

I^{(n+2)}(v) = \sigma^2(n+1)\cdot I^{(n)}(v) - v^{n+1}\sigma^2\exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr)

or

.. math::

\int v^{n+2} \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv = \sigma^2(n+1)\int v^n \exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr) dv - v^{n+1}\sigma^2\exp\Bigl(-\frac{v^2}{2\sigma^2}\Bigr)

Here :math:\mathrm{erf} is the error function, which can be defined as

.. math::

\mathrm{erf}(x)=\frac{2}{\sqrt\pi}\int_0^x \exp(-t^2) dt

By definition, :math:\mathrm{erf}(\pm\infty)=\pm 1, and :math:\mathrm{erf}(0)=0.
In python, one can get it by either :func:math.erf or scipy.special.erf() function.

Sometimes, one defines :math:\mathrm{erfc}(x) \equiv 1 - \mathrm{erf}(x), also
one can use :func:math.erfc or :func:scipy.special.erfc.

Definite integration:

===================== =============================== ====================================
Integral              :math:v \in [-\infty, \infty] :math:v\in [0, \infty]
--------------------- ------------------------------- ------------------------------------
:math:I^{(0)}       :math:\sqrt{2\pi}\sigma       :math:\sqrt\frac{\pi}{2}\sigma
:math:I^{(1)}       :math:0                       :math:\sigma^2
:math:I^{(2)}       :math:\sqrt{2\pi}\sigma^3     :math:\sqrt\frac{\pi}{2}\sigma^3
:math:I^{(3)}       :math:0                       :math:2\sigma^4
:math:I^{(4)} (TBC) :math:3\sqrt{2\pi}\sigma^5    :math:3\sqrt\frac{\pi}{2}\sigma^5
===================== =============================== ====================================

'''
import logging as _logging
import warnings as _warnings

import numpy as _np
from scipy import optimize
import scipy.special
from scipy import constants as _scicon
from scipy import constants as _c

from irfpy.util import unitq as u
from irfpy.util.unitq import k
from irfpy.util import physics as ph
from irfpy.util import exception as _ex

_logger = _logging.getLogger(__name__)

[docs]def mkfunc(n0, v0, vth, multiarray=False):
r''' Returns a function of the Maxwellian.

:param n0: Density (scalar)
:type n0: Float
:param v0: 3-D array (list, tuple, numpy.array, or whatever) of
initial velocity, (vx0, vy0, vz0)
:type v0: array-like
:param vth: Thermal velocity (scalar)
:type vth: Float
:returns: A Maxwell distribution function, which get 3-element array
(list, tuple, numpy.array) to return the distribution function.

.. math::

f([v_x,v_y,v_z]) = \frac{n_0}{(2\pi v_{th}^2)^{3/2}}
\exp -\frac{(v_x-v_{x0})^2 + (v_y-v_{y0})^2 +
(v_z-v_{z0})^2 }{2 v_{th}^2}

:rtype: A function.

For n=5 /cc, 400 km/s bulk velocity with 40 km/s thermal velocity
maxwellian is formed from the following.  Use MKSA system.

>>> fsw = mkfunc(5e6, [-400e3, 0, 0], 40e3)
>>> print('%.2e' % fsw([-400e3, 0, 0]))    # The peak
4.96e-09

You may give (3, N) shaped array or (3, N, M)... shaped array.

>>> fval = fsw([[-400e3, -400e3], [0, -40e3], [0, 0]])
>>> from irfpy.util.with_context import printoptions
>>> with printoptions(precision=3):
...     print(fval)       # doctest: +NORMALIZE_WHITESPACE
[4.960e-09   3.009e-09]

Note again that the expression of the Maxwellian does not include mass.
Thus, this function can be used regardless the particles.
'''
if _np.isscalar(v0) or len(v0) != 3:
raise ValueError('Velocity should be 3-element vector here (not %d).'
% len(v0))

def __func(varr):
''' Maxwell function.

:param varr: Velocity vector.
:type varr: np.array with the shape of (3, ...)
:returns: Velocity distribution function.
'''
v = _np.array(varr)
coeff = (2 * _np.pi * vth ** 2) ** (-1.5)
return n0 * coeff * _np.exp(-((v[0] - v0[0]) ** 2
+ (v[1] - v0[1]) ** 2 + (v[2] - v0[2]) ** 2) / (2 * vth ** 2))

return __func

[docs]class MaxwellDistributionTT:
r''' Create the function for the Maxwell distribution with temperature tensor as argument.

:param n0: Density in MKSA (/m^3)
:param v0: Bulk velocity vector in MKSA (m/s). (3,) shape numpy array.
:param temperature_tensor: Temperature tensor in Kelvin. (3, 3) shape array, but must be symmetric.
If it is not symmetric, :class:NotSymmetricException is raised.
:keyword tol: Tolerance if the temperature tensor is symmetric or not.  If the tensor is judged as
non-symmetric, the :class:NotSymmetricException is raised
:return: A function, fv(vx, vy, vz).  All the units are assumed MKSA.

.. math::

f(\vec{v}) = n_0 \left(\frac{m}{2\pi k}\right)^{3/2}\frac{1}{\sqrt{|T|}}\exp-\frac{m}{2k}(\vec{v}-\vec{v_0})^T\mathbf{T}^{-1}(\vec{v}-\vec{v_0}))

>>> f = MaxwellDistributionTT(1e6, [300e3, 0e3, 0e3], [[15000, 0, 0], [0, 5000, 0,], [0, 0, 30000]])
>>> print(f)
f(\vec{v}) = n_0 \left(\frac{m}{2\pi k}\right)^{3/2}\frac{1}{\sqrt{|T|}}\exp-\frac{m}{2k}(\vec{v}-\vec{v_0})^T\mathbf{T}^{-1}(\vec{v}-\vec{v_0}))
n_0 = 1000000.0
v_0 = 300000.0 0.0 0.0
T = 15000 0 0
0 5000 0
0 0 30000
m = 1.67262192369e-27

>>> print('{:.2e}'.format(f([300e3, 0, 0])))
5.64e-08
>>> print('{:.2e}'.format(f([350e3, 0, 0])))
2.33e-12
>>> print('{:.2e}'.format(f([300e3, 50e3, 0])))
3.96e-21
>>> print('{:.2e}'.format(f([300e3, 0, 50e3])))
3.63e-10
'''
_logger = _logging.getLogger(__name__)

def __init__(self, n0, v0, temperature_tensor, tol=1e-8, m=_scicon.proton_mass):
# Argument check.
vb = _np.array(v0)
if vb.shape != (3,):
raise ValueError('V0 should be (3,)-shape array: {}'.format(vb))
is_symmetric_tensor(temperature_tensor, tol=tol, raises=True)

self.n0 = n0
self.v0 = vb
self.tt = temperature_tensor
self._ttinv = _np.linalg.inv(self.tt)
self.m = m
self._coeff2 = self.m / 2 / _np.pi / _scicon.Boltzmann
self._coeff3 = 1 / _np.sqrt(_np.linalg.det(self.tt))
self._coeff4 = n0 * (self._coeff2 ** 1.5) * self._coeff3

_logger.debug('coeff2 = m/(2 pi k) = %e', self._coeff2)
_logger.debug('coeff3 = 1 / sqrt(det(T)) = %e', self._coeff3)
_logger.debug('coeff4 = n0 coeff2 ** 1.5 / coeff3 = %e', self._coeff4)

def __call__(self, vvec):
coeff2 = self._coeff2 * _np.pi
coeff4 = self._coeff4

v = (vvec - self.v0)    # (3,) shape
tv = self._ttinv.dot(v)  # (3) shape
tv = v.dot(tv)    # scalar
inexp = -coeff2 * tv

exp = _np.exp(inexp)

f = coeff4 * exp

return f

def __str__(self):
s = r'f(\vec{v}) = n_0 \left(\frac{m}{2\pi k}\right)^{3/2}\frac{1}{\sqrt{|T|}}\exp-\frac{m}{2k}(\vec{v}-\vec{v_0})^T\mathbf{T}^{-1}(\vec{v}-\vec{v_0}))'
s += '\n'
s += 'n_0 = {}'.format(self.n0)
s += '\n'
s += 'v_0 = {v[0]:} {v[1]:} {v[2]:}'.format(v=self.v0)
s += '\n'
s += 'T = {T[0]:} {T[1]:} {T[2]:}\n'.format(T=self.tt[0])
s += '    {T[0]:} {T[1]:} {T[2]:}\n'.format(T=self.tt[1])
s += '    {T[0]:} {T[1]:} {T[2]:}\n'.format(T=self.tt[2])
s += 'm = {}'.format(self.m)
return s

[docs]class LnMaxwellDistributionTT:
r''' Create the function for the log-natural Maxwell distribution with temperature tensor as argument.

:param n0: Density in MKSA (/m^3)
:param v0: Bulk velocity vector in MKSA (m/s). (3,) shape numpy array.
:param temperature_tensor: Temperature tensor in Kelvin. (3, 3) shape array, but must be symmetric.
If it is not symmetric, :class:NotSymmetricException is raised.
:keyword tol: Tolerance if the temperature tensor is symmetric or not.  If the tensor is judged as
non-symmetric, the :class:NotSymmetricException is raised
:return: A function, lnfv(vx, vy, vz); it returns the log of VDF.  All the units are assumed MKSA.

.. math::

\ln f(\vec{v}) &= \ln\left(n_0 \left(\frac{m}{2\pi k}\right)^{3/2}\frac{1}{\sqrt{|T|}}\exp-\frac{m}{2\pi k}(\vec{v}-\vec{v_0})^T\mathbf{T}^{-1}(\vec{v}-\vec{v_0}))\right)  \\
&= \ln n_0 + \frac{3}{2}(\ln m-\ln 2\pi k) - \frac{1}{2}\ln|T| - \frac{m}{2k}(\vec{v}-\vec{v_b})^T\mathbf{T}(\vec{v}-\vec{v_b})

.. note::

This is much faster than :class:MaxwellDistributionTT (~1000 times).
However, np.exp costs significant amount time.

>>> f = LnMaxwellDistributionTT(1e6, [300e3, 0e3, 0e3], [[15000, 0, 0], [0, 5000, 0,], [0, 0, 30000]])
>>> print(f)
\ln f(\vec{v}) = \ln\left(n_0 \left(\frac{m}{2\pi k}\right)^{3/2}\frac{1}{\sqrt{|T|}}\exp-\frac{m}{2k}(\vec{v}-\vec{v_0})^T\mathbf{T}^{-1}(\vec{v}-\vec{v_0}))\right)
n_0 = 1000000.0
v_0 = 300000.0 0.0 0.0
T = 15000 0 0
0 5000 0
0 0 30000
m = 1.67262192369e-27

>>> import numpy as np
>>> print('{:.2e}'.format(np.exp(f([300e3, 0, 0]))))
5.64e-08
>>> print('{:.2e}'.format(np.exp(f([350e3, 0, 0]))))
2.33e-12
>>> print('{:.2e}'.format(np.exp(f([300e3, 50e3, 0]))))
3.96e-21
>>> print('{:.2e}'.format(np.exp(f([300e3, 0, 50e3]))))
3.63e-10
'''
def __init__(self, n0, v0, temperature_tensor, tol=1e-8, m=_scicon.proton_mass):
# Argument check.
vb = _np.array(v0)
if vb.shape != (3,):
raise ValueError('V0 should be (3,)-shape array: {}'.format(vb))
is_symmetric_tensor(temperature_tensor, tol=tol, raises=True)

self.n0 = n0
self.v0 = vb
self.tt = temperature_tensor
self._ttinv = _np.linalg.inv(self.tt)
self.m = m

self._coeff1 = _np.log(n0)
self._coeff2 = (_np.log(m) - _np.log(2) - _np.log(_np.pi) - _np.log(_scicon.Boltzmann)) * 1.5
self._coeff3 = -0.5 * _np.log(_np.linalg.det(self.tt))
self._coeff4 = -m / (2 * _scicon.Boltzmann)

def __call__(self, vvec):
v = (vvec - self.v0)    # (3,) shape
tv = self._ttinv.dot(v)  # (3) shape
tv = v.dot(tv)    # (3) shape

lnf = self._coeff1 + self._coeff2 + self._coeff3 + self._coeff4 * tv

return lnf

def __str__(self):
s = r'\ln f(\vec{v}) = \ln\left(n_0 \left(\frac{m}{2\pi k}\right)^{3/2}\frac{1}{\sqrt{|T|}}\exp-\frac{m}{2k}(\vec{v}-\vec{v_0})^T\mathbf{T}^{-1}(\vec{v}-\vec{v_0}))\right)'
s += '\n'
s += 'n_0 = {}'.format(self.n0)
s += '\n'
s += 'v_0 = {v[0]:} {v[1]:} {v[2]:}'.format(v=self.v0)
s += '\n'
s += 'T = {T[0]:} {T[1]:} {T[2]:}\n'.format(T=self.tt[0])
s += '    {T[0]:} {T[1]:} {T[2]:}\n'.format(T=self.tt[1])
s += '    {T[0]:} {T[1]:} {T[2]:}\n'.format(T=self.tt[2])
s += 'm = {}'.format(self.m)
return s

[docs]class NotSymmetricException(_ex.IrfpyException):
""" Exception raised when the matrix/tensor is not symmetric.
"""
def __init__(self, value):
self.value = value

def __str__(self):
return repr(self.value)

[docs]def is_symmetric_tensor(tensor2d, tol=1e-8, raises=False):
r""" Check if the given tensor is symmetric or not.

:param tensor2d: 2-D tensor to be checked.
:keyword tol: Tolerance
:keyword raises: Instead of returning False, an exception (:class:NotSymmetricException) is raised.
:return: If all the elements in tensor2d is smaller than tol, the tensor2d is considered as symmetric and
return True. Otherwise, either False is returned, or :class:NotSymmetricException is raised,
depending on the raises keyword.

>>> t = [[1, 3, 2],
...      [3, 1, 4],
...      [2, 4, 3.0]]
>>> is_symmetric_tensor(t)
True

>>> t = [[1, 3, 2],
...      [3, 4, 5],
...      [1.999, 5, 3]]
>>> is_symmetric_tensor(t)
False

>>> is_symmetric_tensor(t, raises=True)    # doctests: +IGNORE_EXCEPTION_DETAIL
Traceback (most recent call last):
irfpy.util.maxwell.NotSymmetricException: 'Given tensor/matrix is not symmetric: [[1.    3.    2.   ]\n [3.    4.    5.   ]\n [1.999 5.    3.   ]]'

>>> is_symmetric_tensor(t, tol=0.01)
True
"""
t = _np.array(tensor2d)
diff = _np.abs(t - t.T)
diffmax = diff.max()

if diffmax < tol:
return True
else:
if raises:
raise NotSymmetricException('Given tensor/matrix is not symmetric: {}'.format(str(t)))
else:
return False

[docs]def mkfunc_u(n0_u, v0_u, vth_u):
''' Unitful version of :func:mkfunc.

>>> f_u = mkfunc_u(30 / u.cc, [250, -15, 0] * u.km / u.s, 35 * u.km / u.s)
>>> print('%.3e' % f_u([230, 0, 0] * u.km / u.s).nr(u.nit('s**3/cm**6')))
3.442e-20

If you use non-unitful version, the results should be the same.

>>> f = mkfunc(30e6, [250e3, -15e3, 0], 35e3)
>>> print('%.3e' % f([230e3, 0, 0]))
3.442e-08

'''

n0 = n0_u.nr(u.m ** -3)
v0 = v0_u.nr(u.m / u.s)
vth = vth_u.nr(u.m / u.s)

func = mkfunc(n0, v0, vth)

def ffunc(v_u):
'''Maxwell function with unit support

:param v_u: Velocity vector
:returns: Velocity distribution function
'''
v = v_u.nr(u.m / u.s)
f = func(v)
return f * u.s ** 3 / u.m ** 6

return ffunc

[docs]class MaxwellDistributionG:
''' Maxwell Distribution in Gridded space

This class derives the differential flux and the velocity distribution function,
and save these values on a gridded bin.

.. note::

This class / functionality has been replaced by
:func:irfpy.util.energyspectrum.maxwellDifferentialFluxGrid
and
:func:    irfpy.util.vdf.maxwellVelocityDistributionGrid
functions.

First, create object with a set of parameter.

>>> maxwellobj = MaxwellDistributionG(m=k.proton_mass, n=10 / u.cc,
...             vsw = [-400, 0, 50] * u.km / u.s,
...             vth = 150 * u.km / u.s)

As seen in the sample, parameters are unitted.
From profiler point of view, the overhead in the computational resource (time)
is very small by using unitq.

Second, set the grid if you want.  Only E-t-p contour grid is supported.
By default, grid is defined as 96 step for energy direction (100 - 20000 eV, logspace)
181 steps for theta, and 361 steps for phi.

>>> maxwellobj.set_etp_grid([100, 300, 1000, 3000, 10000],
...                 [-90, -60, -30, -15, 0, 15, 30, 60, 90],
...                 [-180, -120, -60, 0, 60, 120, 180])

Third, calculate.

>>> maxwellobj.do_calc()

Now you can refer to j and fval.

>>> print('%.3e (1/cm**2 /sr /eV / s)' % maxwellobj.j.max())
3.164e+05 (1/cm**2 /sr /eV / s)

>>> u.pprint(maxwellobj.fval.max(), fmt='%.3e')
1.724e-10 s**3/m**6

Other attributes that you may refer to is

- m_u: Mass specified
- n_u: Density specified
- vsw_u: Bulk velocity specified
- vth_u: Thermal velocity specified.
- dive: Number of energy bin
- divt: Number of theta (polar) bin
- divp: Number of phi (azimuth) bin
- j: Differential flux
- fval: VDF
'''
logger = _logging.getLogger(__name__ + '.MaxwellDistributionG')

def __init__(self, m=k.proton_mass, n=10 / u.cc,
vsw=[400, 50, -30] * u.km / u.s,
vth=60 * u.km / u.s):
''' Initialize by the solar wind parameter
'''
_warnings.warn(
"MaxwellDistributionG is deprecated. Use irfpy.util.energyspectrum.maxwellDiifferentialGrid "
"or irfpy.util.vdf.maxwellVelocityDistribution for future use.",
DeprecationWarning)

self.m_u = m
self.m = m.nr(u.kg)
self.n_u = n
self.n = n.nr(u.m**-3)
self.vsw_u = vsw
self.vsw = vsw.nr(u.m / u.s)
self.vth_u = vth
self.vth = vth.nr(u.m / u.s)

# Maxwell function, interface with MKSA.
maxwell_func = mkfunc(self.n, self.vsw, self.vth)
self.maxwell_func = maxwell_func

self.dive = -1
self.divt = -1
self.divp = -1

self.set_etp_grid()

self.j = None
self.fval = None

[docs]    def set_etp_grid(self, egrid1d=_np.logspace(_np.log10(100), _np.log10(20000), 96),
tgrid1d=_np.linspace(-90, 90, 181),
pgrid1d=_np.linspace(-180, 180, 361)):
''' Define the grid. Not the energy in eV, angles in degree.
'''
self.egrid1d = egrid1d   # eV
self.tgrid1d = tgrid1d   # deg
self.pgrid1d = pgrid1d   # deg
self.vgrid1d = ph.energy_2_velocity(self.m, egrid1d * u.k2.e)
self.logger.debug('Exam velocity from {0} to {1}.'.format(min(self.vgrid1d), max(self.vgrid1d)))

self.dive = len(egrid1d)
self.divt = len(tgrid1d)
self.divp = len(pgrid1d)

self.j = None
self.fval = None

[docs]    def get_vtp_grid3d(self):
''' Calculate the 3-D grid and return

:returns: (vgrid, tgrid, pgrid). Velocity, theta and phi (degress) on the grid.
Each is 3-D array with shape of (dive, divt, divp).
'''
etp_bin = _np.meshgrid(self.vgrid1d, self.tgrid1d, self.pgrid1d)
vgrid, tgrid, pgrid = etp_bin   # Value follows the order of argument.
# each shape follows t, v, p.  (order of 1, 0, 2)
vgrid = vgrid.transpose([1, 0, 2])
tgrid = tgrid.transpose([1, 0, 2])
pgrid = pgrid.transpose([1, 0, 2])

return vgrid, tgrid, pgrid

[docs]    def get_egrid3d(self):
''' Calculate 3-D grid of energy

:returns: Energy on the grid in eV.  Shape of (dive, divt, divp)
'''
vgrid, tgrid, pgrid = self.get_vtp_grid3d()
egrid = ph.kinetic_energy(self.m_u, vgrid) / u.k2.e
return egrid

[docs]    def get_vxyz_array(self):
''' Calculate corresponding velocity on the grid

:returns: (vxarr, vyarr, vzarr). Each data is 3-D array with shape of (dive, divt, divp).
'''
vgrid, tgrid, pgrid = self.get_vtp_grid3d()

vxarr = vgrid * _np.cos(tgrid) * _np.cos(pgrid)
vyarr = vgrid * _np.cos(tgrid) * _np.sin(pgrid)
vzarr = vgrid * _np.sin(tgrid)
self.logger.debug('Shape of array = %s' % str(vxarr.shape))

return vxarr, vyarr, vzarr

[docs]    def do_calc(self):
''' Run the calculation
'''

m, n, vsw, vth = self.m, self.n, self.vsw, self.vth

vsw_abs = _np.sqrt((vsw ** 2).sum())
esw = ph.kinetic_energy(m, vsw_abs)   # Now in J.
esw /= u.k2.e   # Now in eV.

vxarr, vyarr, vzarr = self.get_vxyz_array()
#        vxarr = vxarr.nr(u.m/u.s)
#        vyarr = vyarr.nr(u.m/u.s)
#        vzarr = vzarr.nr(u.m/u.s)
vvec = _np.array([vxarr, vyarr, vzarr])
self.fval = self.maxwell_func(vvec)   # Now in u.s ** 3 / u.m ** 6
self.fval = self.fval * u.s**3 / u.m**6

#        vxarr, vyarr, vzarr = self.get_vxyz_array()

varr = _np.sqrt(vxarr ** 2 + vyarr ** 2 + vzarr ** 2)   # in m/s

# Start calculating differential flux.
self.j = ph.vdf_2_dflux(self.fval.n, m, varr) * u.k2.e * 1e-4  # in cm-eV system, / u.cm ** 2 / u.sr / u.eV / u.s

def __str__(self):
from io import StringIO
s = StringIO()
s.write('Maxwell distribution, gridded to nE=%d nT=%d nP=%d\n' % (self.dive, self.divt, self.divp))
s.write('  n = %g / cm3\n' % self.n_u.nr('cm ** -3'))
s.write('  vsw = %g %g %g km/s\n' % tuple(self.vsw_u.nr('km/s').tolist()))
s.write('  vth = %g\n' % self.vth_u.nr('km/s'))

if self.fval is None or self.j is None:
s.write('----Not yet calculated. (do_calc)\n')
else:
s.write('----\n')
s.write('  peak vdf = %g s^3/m^6\n' % (self.fval.max()))

idx = self.fval.argmax()
vxarr, vyarr, vzarr = self.get_vxyz_array()
vxatmax = vxarr[_np.unravel_index(idx, self.fval.shape)] / 1000.
vyatmax = vyarr[_np.unravel_index(idx, self.fval.shape)] / 1000.
vzatmax = vzarr[_np.unravel_index(idx, self.fval.shape)] / 1000.
s.write('      at  %g  %g  %g km/s\n' % (vxatmax, vyatmax, vzatmax))
s.write('  peak j = %g /m2 sr J s\n' % (self.j.max()))
s.write('         = %g /cm2 sr eV s\n' % (self.j.max() / u.m ** 2 / u.J).rescale(1 / u.cm**2 / u.eV))

idx = self.j.argmax()
egrid3d = self.get_egrid3d()
vgrid3d, tgrid3d, pgrid3d = self.get_vtp_grid3d()
eatmax = egrid3d[_np.unravel_index(idx, self.j.shape)]
tatmax = tgrid3d[_np.unravel_index(idx, self.j.shape)]
patmax = pgrid3d[_np.unravel_index(idx, self.j.shape)]
s.write('      at  %g eV,  t= %g deg,  p= %g deg\n' % (eatmax, tatmax, patmax))

return s.getvalue()

[docs]def mk1d(A, B, C):
r''' Returns a callable of 1-D expression of Maxwellian distribution.

Returns a callable of 1-D expression of the Maxwellian distribution
(standard distribution).  The formulation is

.. math::

f(v) = A \exp -\frac{(v-B)^2}{2 C^2}

See :meth:mkfunc also.
'''
def __func(v):
if _np.isscalar(v):
return A * _np.exp(-(v - B) ** 2 / (2 * C ** 2))
else:
return A * _np.exp(-(_np.array(v) - B) ** 2 / (2 * C ** 2))
return __func

[docs]def mkslice(n0, v0, vth):
r''' Returns a callable of 1-D slice of the distribution function.

Maxwell distribution :math:f(v_x ,v_y ,v_z) is 3-D,
but one can slice it along the direction of the bulk velocity
:math:\mathbf{v_0} = (v_{x0} ,v_{y0} ,v_{z0}).
This method can be widely used for plasma data analysis
because one can pick energy spectra of the direction along the bulk velocity.

The formulation is

.. math::

f(v) = n_0 ( 2\pi v_{th}^2 ) ^{-3/2} \exp -\frac{(v-v_0)^2}{2 v_{th}^2}

Here, :math:v is the scalar value, which indeed expresses the vector
:math:\mathbf{v} = v \mathbf{v_0}/|\mathbf{v_0}|.
Never consider such like :math:v=|\mathbf{v}|.

:param n0: Density (scalar)
:param v0: Velocity (scalar)
:param vth: Thermal velocity (scalar)
:retturns: A Maxwell distribution function, which get scalar
velocity to return the distribution function.
'''
def __func(v):
coeff = (2 * _np.pi * vth ** 2) ** (-1.5)
_logger.debug('(2 pi vth^2)^-1.5 = %e' % coeff)
if _np.isscalar(v):
va = _np.array([v])
else:
va = _np.array(v)
fv = n0 * coeff * _np.exp(-(va - v0) ** 2 / (2 * vth ** 2))
if _np.isscalar(v):
return fv[0]
else:
return fv

return __func

[docs]def fit1d(x_array, y_array):
r''' Fitting by 1-dimensional Maxwell distiribution function.

:param x_array: A sequence of x values.
:type x_array: A sequence of flaot
:param y_array: A sequence of y values.
:type y_array: A sequence of flaot
:returns: ((A, B, C), success_code).  (A, B, C) is defined below.
success_code is the same as defied in :meth:optimize.leastsq.

Fitting by least square algorithm.
Initial values for iteration is estimated automatically.
In order to reduce confusion in the sqrt term in the coefficient,
the coefficient it iself is one parameter, A.
So users must normalize the coefficiency to convert it to density or
probablility functions.

The returned values is the best fit parameters of (A,B,C) which is
formulated as

.. math::

y_{array} = A \exp(-\frac{(x-B)^2}{2C^2})
'''
# Convert to numpy.array
x = _np.array(x_array).copy()
y = _np.array(y_array).copy()
y = _np.where(y > 0, y, 0)

# where is y max?
ix0 = y.argmax()
x0 = x[ix0]
y0 = y.max()
_logger.debug('Y0 = %g at x0=%g' % (y0, x0))

# Normalize.
z = y / y0
z0 = 1

# Neighbor
if ix0 == 0:
ix1 = 1
else:
ix1 = ix0 - 1

x1 = x[ix1]
y1 = y[ix1]
z1 = z[ix1]
_logger.debug('Y1 = %g (z1=%g) at x1=%g' % (y1, z1, x1))

# Estimated B is x0 obviously.
B = x0
_logger.debug('Estimated B = %g' % B)

# Estimated C is abs(x0-x1) / sqrt(2*log(z1))
C = abs(x0 - x1) / _np.sqrt(2 * abs(_np.log(z1)))
_logger.debug('Estimated C = %g' % C)

# Estimated A is y0
A = 1
_logger.debug('Estimated A = %g' % A)

#    print x.shape, z.shape
fitfunc = lambda p, _x: p[0] * _np.exp(-(p[1] - _x) ** 2 / (2 * p[2] ** 2))
errfunc = lambda p, _x, _y: fitfunc(p, _x) - _y
p1, success = optimize.leastsq(errfunc, (A, B, C), args=(x, z))

A, B, C = p1
A = A * y0

return (A, B, C), success

[docs]def getRandom(vel0, vth, n=10000):
''' Get a random distributed numpy.array of the 3-D Maxwellian.

vel0 should be 3-d array of the bulk velocity vector.
vth should be the scalar.
The returned value is a numpy.array with a shape of [3, n]
'''
vx = getRandom1D(vel0[0], vth, n=n)
vy = getRandom1D(vel0[1], vth, n=n)
vz = getRandom1D(vel0[2], vth, n=n)

parts = _np.array([vx, vy, vz])

return parts

[docs]def getRandom1D(v0, vth, n=10000):
''' Get a random distributed numpy.array of values for 1-D Maxwellian.

Get a random distributed numpy.array of the value following 1-D
Maxwellian using scipy.  Indeed this function is just a wrapper
to scipy.norm functionality.  The thermal velocity is same as the
scale parameter of the normal distribution,
and the bulk velocity is same as the location parameter.
'''
from scipy.stats import norm
return norm.rvs(size=n, scale=vth, loc=v0)

[docs]def flux_integrated(n0, v0, vth, vzmin=-_np.infty):
r""" Integrate the maxwellian, with a specific minimum range.

:param n0: Density, scalar.
:param v0: Velocity, vector. (vx0, vy0, vz0)
:param vth: Thermal velocity. Scalar.
:keyword: vzmin: If given, the lower end of integral (in z axis) is given.

From the Maxwellian velocity distribution,

.. math::

f(v_x,v_y,v_z) = n_0 ( 2\pi v_\mathrm{th}^2 ) ^{-\frac{3}{2}}
\exp -\frac{ (v_x-v_{x0})^2 + (v_y-v_{y0})^2 + (v_z-v_{z0})^2 }
{2 v_\mathrm{th}^2}

the following integral is returned.

.. math::

F = \int_{-\infty}^{\infty} dv_x
\int_{-\infty}^{\infty} dv_y
\int_{v_{zmin}}^{\infty} dv_z \bigl[v_z f(v_x, v_y, v_z)\bigr]

A mathematics says that the value is

.. math::

F = \frac{n_0v_{th}}{\sqrt{2\pi}}\exp-\frac{(v_{zmin}-v_{z0})^2}{2v_{th}^2} + \frac{n_0v_{z0}}{2}\mathrm{erfc}\frac{v_{zmin}-v_{z0}}{\sqrt{2}{v_{th}}}

This integral is useful for calculating the flux but only for those
with the given speed.

Example 1: For the default parameter, the returned value gives
just n0 v0z, since the integral is taken over the whole velocity space.

>>> n0 = 30
>>> v0 = (0, 10, 20)
>>> vth = 15
>>> print(flux_integrated(n0, v0, vth))
600.0

Example 2: If the lower edge of the vz integral range is higher than
the core component, the flux approaches to 0.

>>> n0 = 30
>>> v0 = (0, 15, 40)
>>> vth = 20
>>> vzmin = 140    # This is vz0 + 5 vth.
>>> f = flux_integrated(n0, v0, vth, vzmin=vzmin)
>>> print('{:.2e}'.format(f))
1.24e-03

Example 3: If no bulk velocity, and integral range from 0, then

>>> n0 = 1
>>> v0 = (0, 0, 0)
>>> vth = 5
>>> vzmin = 0
>>> f = flux_integrated(n0, v0, vth, vzmin=vzmin)
>>> print('{:.4f}'.format(f))
1.9947
"""
_logger = _logging.getLogger(__name__ + ".flux_integrated")
pi = _np.pi
_, _, vz0 = v0   # Only vz0 is used.

_logger.debug('n0={}'.format(n0))
_logger.debug('vz0={}'.format(vz0))
_logger.debug('vth={}'.format(vth))
_logger.debug('vzmin={}'.format(vzmin))
_logger.debug('---')

coeff1 = n0 * vth / _np.sqrt(2 * pi)
_logger.debug('coeff1={}'.format(coeff1))
inexp = ((vzmin - vz0) ** 2) / (2 * (vth ** 2))
_logger.debug('inexp={}'.format(inexp))

coeff2 = n0 * vz0 / 2
_logger.debug('coeff2={}'.format(coeff2))
inerfc = (vzmin - vz0) / (_np.sqrt(2) * vth)
_logger.debug('inerfc={}'.format(inerfc))

f = coeff1 * _np.exp(-inexp) + coeff2 * scipy.special.erfc(inerfc)
_logger.debug('f={}'.format(f))

return f

[docs]def flux_precip(n0, vz0, vth):
''' Return the precipitating flux along +z.

This function is to calculate the precipitation flux with a non-zero temperature.

This is equivalent to :func:flux_integrated(n0, (0, 0, vz0), vth, vzmin=0) <flux_integrated>.

.. todo::

More documentation

:param n0: Density
:param vz0: Velocity along the precipitating direction.
Usually, the direction of nadir or direction of the magnetic field
according to the problem of interest.
:param vth: Thermal velocity.

The calculation is based on the MKSA system.
If you need non-MKSA system calculation, unit conversions are necessary.

>>> print('{:.4f}'.format(flux_precip(1, 0, 5)))
1.9947
'''
n0 = float(n0)
vz0 = float(vz0)
vth = float(vth)
term1 = n0 * vth / _np.sqrt(2 * _np.pi)
inexp = -0.5 * ((vz0 / vth) ** 2)
term2 = n0 * vz0 / 2.0
inerfc = -(vz0 / vth) / _np.sqrt(2.)

return term1 * _np.exp(inexp) + term2 * scipy.special.erfc(inerfc)

[docs]def differential_flux_u(density, bulk_velocity, thermal_velocity, m=_c.proton_mass):
raise NotImplementedError('')

[docs]class LnDifferentialFlux:
""" Create a function to return the log natural of differential flux.
"""

e = _c.elementary_charge
log_e = _np.log(e)
log_2 = _np.log(2)
log_2pi = _np.log(2 * _np.pi)
log_1em4 = _np.log(1e-4)

def __init__(self, density_mksa, bulk_velocity_mksa, thermal_velocity_mksa, m=_c.proton_mass):
r"""

:param density_mksa: Density in m^-3
:param bulk_velocity_mksa: Bulk velocity vector.  m/s
:param thermal_velocity_mksa: Scalar thermal velocity. m/s
:keyword m: Mass. kg.

Return a function (a callable object) that calculate the log differential flux for a maxwellian.

If you want to calculate the differential flux corresponding to the density 10/cm^3,
Vx=150 km/s, Vy=Vz=1 m/s, and Vth=30 km/s (~10 eV for proton)

>>> dflux_func = LnDifferentialFlux(10e6, [150e3, 1, 1], 30e3)

The returned dflux_func is a function(-like) that can give you the differential flux.
The unit for the dflux_func is

- eV for energy (in/out)
- degrees for angles (in)
* polar angle, |theta|, from +z
* azimutha angle, |phi|, from +x
- cm for dimension (out)
- steradian for solid angle (out)
- /cm2 s sr eV for the differential flux (out)

Thus, if you want to know the differential flux for 120 eV from x-direction (|theta| =90, |phi| =0).

>>> lnj = dflux_func(120, 90, 0)    # 120 eV, theta=90, phi=0
>>> import numpy as np
>>> print('{:.2e}'.format(np.exp(lnj)))          # The returned flux is in /cm2 s sr eV
5.17e+06

"""
vth = thermal_velocity_mksa    # An alias

self.m = m
self.log_m = _np.log(m)
self.kT = m * (vth ** 2)
self.log_kT = _np.log(m) + 2 * _np.log(vth)
self.n = density_mksa
self.log_n = _np.log(self.n)
self.vb = bulk_velocity_mksa
self.coeff0 = self.log_2 + self.log_n - 2 * self.log_m + (self.log_m - self.log_2pi - self.log_kT) * 1.5
self.coeff1 = self.log_1em4 + self.log_e
self.coeff0 = self.coeff0 + self.coeff1
self.incoeff0 = self.m / (2 * self.kT)

def __call__(self, energy_eV, theta_deg, phi_deg):
ene = energy_eV
log_ene = _np.log(ene)
eneJ = ene * self.e
log_eneJ = log_ene + self.log_e
v_m_s = _np.sqrt(2 * eneJ / self.m)

st = _np.sin(theta)
vx = v_m_s * st * _np.cos(phi)
vy = v_m_s * st * _np.sin(phi)
vz = v_m_s * _np.cos(theta)

coeff0 = self.coeff0
coeff2 = log_eneJ
incoeff0 = self.incoeff0
incoeff1 = (vx - self.vb[0]) ** 2 + (vy - self.vb[1]) ** 2 + (vz - self.vb[2]) ** 2

jval = coeff0 + coeff2 - incoeff0 * incoeff1

return jval

[docs]class DifferentialFlux:
""" Return a function that calculate the differential flux.

:param density_mksa: Density in m^-3
:param bulk_velocity_mksa: Bulk velocity vector.  m/s
:param thermal_velocity_mksa: Scalar thermal velocity. m/s
:keyword m: Mass. kg
:return: The function to calculate the differentil flux, J.

J is the function to return the differential flux.  Note that the unit of the energy is eV and that of the area is cm^2.
- eV for energy (in/out)
- degrees for angles (in)
* polar angle, |theta|, from +z
* azimutha angle, |phi|, from +x
- cm for dimension (out)
- steradian for solid angle (out)
- /cm2 s sr eV for the differential flux (out)

>>> dflux_func = DifferentialFlux(10e6, [150e3, 1, 1], 30e3)

Thus, if you want to know the differential flux for 120 eV from x-direction (|theta| =90, |phi| =0).

>>> j = dflux_func(120, 90, 0)    # 120 eV, theta=90, phi=0
>>> print('{:.2e}'.format(j))          # The returned flux is in /cm2 s sr eV
5.17e+06
"""
def __init__(self, density_mksa, bulk_velocity_mksa, thermal_velocity_mksa, m=_c.proton_mass):
self.lnjfunc = LnDifferentialFlux(density_mksa, bulk_velocity_mksa, thermal_velocity_mksa, m=m)

def __call__(self, energy_eV, theta_deg, phi_deg):
lnvalue = self.lnjfunc(energy_eV, theta_deg, phi_deg)
return _np.exp(lnvalue)

[docs]class DifferentialFlux_u:
""" Return a function that calculate the differential flux, with unit support.

:param density:
:param bulk_velocity: Bulk velocity vector.
:param thermal_velocity: Scalar thermal velocity.
:keyword m: Mass
:return: The function to calculate the differentil flux, J. Unit support.

J is the function to return the differential flux.

>>> dflux_func = DifferentialFlux_u(10 / u.cc, [150 * u.km/u.s, 0.001 * u.km/u.s, 0.001 * u.km/u.s], 30 * u.km/u.s)

Thus, if you want to know the differential flux for 120 eV from x-direction (|theta| =90, |phi| =0).

>>> j = dflux_func(120 * u.eV, 90 * u.deg, 0 * u.deg)    # 120 eV, theta=90, phi=0
>>> print('{:.2e}'.format(j))          # The returned flux is in /cm2 s sr eV
5.17e+06 1/(cm**2*s*sr*eV)
"""
def __init__(self, density, bulk_velocity, thermal_velocity, m=k.proton_mass):
self.df_mksa = DifferentialFlux(density.nr(u.m ** -3),
[bulk_velocity[0].nr(u.m / u.s), bulk_velocity[1].nr(u.m / u.s), bulk_velocity[2].nr(u.m / u.s)],
thermal_velocity.nr(u.m / u.s),
m=m.nr(u.kg))

def __call__(self, energy, theta, phi):
df_value_mksa = self.df_mksa(energy.nr(u.eV), theta.nr(u.degree), phi.nr(u.degree))
return df_value_mksa / (u.cm ** 2) / u.s / u.steradian / u.eV