"Maxwellian (Gaussian) distributions in velocity space."
import copy
from typing import Callable
import cunumpy as xp
from struphy.fields_background.base import FluidEquilibriumWithB
from struphy.fields_background.equils import set_defaults
from struphy.geometry.base import Domain
from struphy.initial.base import Perturbation
from struphy.io.options import LiteralOptions
from struphy.kinetic_background.base import Maxwellian
[docs]
class Maxwellian3D(Maxwellian):
r"""A :class:`~struphy.kinetic_background.base.Maxwellian` depending on three (:math:`n=3`) Cartesian velocities.
Parameters
----------
n, ui, vthi : tuple
Moments of the Maxwellian as tuples. The first entry defines the background
(float for constant background or callable), the second entry defines a Perturbation (can be None).
"""
def __init__(
self,
n: tuple[float | Callable, Perturbation] = (1.0, None),
u1: tuple[float | Callable, Perturbation] = (0.0, None),
u2: tuple[float | Callable, Perturbation] = (0.0, None),
u3: tuple[float | Callable, Perturbation] = (0.0, None),
vth1: tuple[float | Callable, Perturbation] = (1.0, None),
vth2: tuple[float | Callable, Perturbation] = (1.0, None),
vth3: tuple[float | Callable, Perturbation] = (1.0, None),
):
# use setter to store input parameters
self.params = copy.deepcopy(locals())
self.check_maxw_params()
# factors multiplied onto the defined moments n, u and vth (can be set via setter)
self._moment_factors = {
"n": 1.0,
"u": [1.0, 1.0, 1.0],
"vth": [1.0, 1.0, 1.0],
}
@property
def coords(self):
"""Coordinates of the Maxwellian6D, :math:`(v_1, v_2, v_3)`."""
return "cartesian"
@property
def vdim(self):
"""Dimension of the velocity space."""
return 3
@property
def is_polar(self):
"""List of booleans of length vdim. True for a velocity coordinate that is a radial polar coordinate (v_perp)."""
return [False, False, False]
[docs]
def velocity_jacobian_det(self, eta1, eta2, eta3, *v):
"""Jacobian determinant of the velocity coordinate transformation from Maxwellian6D('cartesian') to Particles6D('cartesian').
Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords.T)
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
*v : array_like
Velocity evaluation points.
Returns
-------
out : array-like
The Jacobian determinant evaluated at given logical coordinates.
-------
"""
assert eta1.ndim == 1
assert eta2.ndim == 1
assert eta3.ndim == 1
assert len(v) == 3
return 1.0 + 0 * eta1
@property
def volume_form(self):
"""Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian)."""
return False
@property
def moment_factors(self):
"""Collection of factors multiplied onto the defined moments n, u, and vth."""
return self._moment_factors
@moment_factors.setter
def moment_factors(self, **kwargs):
for kw, arg in kwargs:
if kw in {"u", "vth"}:
assert len(arg) == 3
self._moment_factors[kw] = arg
[docs]
def n(self, eta1, eta2, eta3):
"""Zero-th moment (density)."""
out = self._evaluate_moment(eta1, eta2, eta3, name="n")
return out * self.moment_factors["n"]
[docs]
def u(self, eta1, eta2, eta3):
"""Mean velocities."""
out = []
out += [self._evaluate_moment(eta1, eta2, eta3, name="u1")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="u2")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="u3")]
return [ou * mom_fac for ou, mom_fac in zip(out, self.moment_factors["u"])]
[docs]
def vth(self, eta1, eta2, eta3):
"""Thermal velocities."""
out = []
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth1")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth2")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth3")]
return [ou * mom_fac for ou, mom_fac in zip(out, self.moment_factors["vth"])]
[docs]
class GyroMaxwellian2D(Maxwellian):
r"""A gyrotropic :class:`~struphy.kinetic_background.base.Maxwellian` depending on
two velocities :math:`(v_\parallel, v_\perp)`, :math:`n=2`,
where :math:`v_\parallel = \matbf v \cdot \mathbf b_0` and :math:`v_\perp`
is the radial component of a polar coordinate system perpendicular
to the magentic direction :math:`\mathbf b_0`.
Parameters
----------
n, u_para, u_perp, vth_para, vth_perp : tuple
Moments of the Maxwellian as tuples. The first entry defines the background
(float for constant background or callable), the second entry defines a Perturbation (can be None).
equil : FluidEquilibriumWithB
Fluid background.
volume_form : bool
Whether to represent the Maxwellian as a volume form;
if True it is multiplied by the Jacobian determinant |v_perp|
of the polar coordinate transofrmation (default = False).
"""
def __init__(
self,
n: tuple[float | Callable, Perturbation] = (1.0, None),
u_para: tuple[float | Callable, Perturbation] = (0.0, None),
u_perp: tuple[float | Callable, Perturbation] = (0.0, None),
vth_para: tuple[float | Callable, Perturbation] = (1.0, None),
vth_perp: tuple[float | Callable, Perturbation] = (1.0, None),
equil: FluidEquilibriumWithB = None,
volume_form: bool = True,
):
# use setter to store input parameters
self.params = copy.deepcopy(locals())
self.check_maxw_params()
# volume form represenation
self._volume_form = volume_form
self._equil = equil
# factors multiplied onto the defined moments n, u and vth (can be set via setter)
self._moment_factors = {
"n": 1.0,
"u": [1.0, 1.0],
"vth": [1.0, 1.0],
}
@property
def coords(self):
r"""Coordinates of the Maxwellian5D, :math:`(v_\parallel, v_\perp)`."""
return "vpara_vperp"
@property
def vdim(self):
"""Dimension of the velocity space."""
return 2
@property
def is_polar(self):
"""List of booleans of length vdim. True for a velocity coordinate that is a radial polar coordinate (v_perp)."""
return [False, True]
[docs]
def velocity_jacobian_det(self, eta1, eta2, eta3, *v):
r"""Jacobian determinant of the velocity coordinate transformation from Maxwellian5D('vpara_vperp') to Particles5D('vpara_mu').
.. math::
\begin{aligned}
F &: (v_\parallel, v_\perp) \to (v_\parallel, \mu) \,,
\\[3mm]
DF &= \begin{bmatrix} \frac{\partial v_\parallel}{\partial v_\parallel} & \frac{\partial v_\parallel}{\partial v_\perp} \\
\frac{\partial \mu}{\partial v_\parallel} & \frac{\partial \mu}{\partial v_\perp} \end{bmatrix} =
\begin{bmatrix} 1 & 0 \\
0 & \frac{v_\perp}{B} \end{bmatrix} \,,
\\[3mm]
J_F &= \frac{v_\perp}{B} \,,
\end{aligned}
where :math:`\mu = \frac{v_\perp^2}{2B}`.
Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords.T)
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
*v : array_like
Velocity evaluation points.
Returns
-------
out : array-like
The Jacobian determinant evaluated at given logical coordinates.
-------
"""
assert eta1.ndim == 1
assert eta2.ndim == 1
assert eta3.ndim == 1
assert len(v) == 2
# call equilibrium
etas = (xp.vstack((eta1, eta2, eta3)).T).copy()
absB0 = self.equil.absB0(etas)
# J = v_perp/B
jacobian_det = v[1] / absB0
return jacobian_det
@property
def volume_form(self) -> bool:
"""Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian |v_perp|)."""
return self._volume_form
@property
def equil(self) -> FluidEquilibriumWithB:
"""Fluid background with B-field."""
return self._equil
@property
def moment_factors(self):
"""Collection of factors multiplied onto the defined moments n, u, and vth."""
return self._moment_factors
@moment_factors.setter
def moment_factors(self, **kwargs):
for kw, arg in kwargs:
if kw in {"u", "vth"}:
assert len(arg) == 2
self._moment_factors[kw] = arg
[docs]
def n(self, eta1, eta2, eta3):
"""Zero-th moment (density)."""
out = self._evaluate_moment(eta1, eta2, eta3, name="n")
return out * self.moment_factors["n"]
[docs]
def u(self, eta1, eta2, eta3):
"""Mean velocities."""
out = []
out += [self._evaluate_moment(eta1, eta2, eta3, name="u_para")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="u_perp")]
return [ou * mom_fac for ou, mom_fac in zip(out, self.moment_factors["u"])]
[docs]
def vth(self, eta1, eta2, eta3):
"""Thermal velocities."""
out = []
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth_para")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth_perp")]
return [ou * mom_fac for ou, mom_fac in zip(out, self.moment_factors["vth"])]
[docs]
def plot_density_profile(
self,
dim_1: LiteralOptions.KineticDimensionsToPlot = "e1",
dim_2: LiteralOptions.KineticDimensionsToPlot | None = None,
v_lim: float = 5.0,
resol: int = 100,
integrate_resol: int = 10,
logical_coord: tuple[float] = (0.5, 0.5, 0.5),
in_physical: bool = False,
domain: Domain | None = None,
proj_axis: tuple[float,] = (0, 1),
plot_3D: bool = False,
title: str | None = None,
use_mu: bool = False,
equil: FluidEquilibriumWithB | None = None,
):
if equil is None:
equil = self.equil
super().plot_density_profile(
dim_1,
dim_2,
v_lim,
resol,
integrate_resol,
logical_coord,
in_physical,
domain,
proj_axis,
plot_3D,
title,
use_mu=use_mu,
equil=equil,
)
[docs]
class CanonicalMaxwellian:
r"""Canonical Maxwellian distribution function in constants-of-motion coordinates.
The distribution is parameterized by the density and thermal speed as functions of the
canonical toroidal momentum :math:`\psi_c`:
.. math::
\psi_c = \psi + \frac{m_s F}{q_s B}v_\parallel - \text{sign}(v_\parallel)\sqrt{2(\epsilon - \mu B)}\frac{m_sF}{q_sB} \mathcal{H}(\epsilon - \mu B),
- Energy
.. math::
\epsilon = \frac{1}{2}m_sv_\parallel² + \mu B,
- Magnetic moment
.. math::
\mu = \frac{m_s v_\perp²}{2B},
where :math:`\psi` is the poloidal magnetic flux function, :math:`F=F(\psi)` is the poloidal current function and :math:`\mathcal{H}` is the Heaviside function.
With the three constants of motion, a canonical Maxwellian distribution function is defined as
.. math::
F(\psi_c, \epsilon, \mu) = \frac{n(\psi_c)}{(2\pi)^{3/2}v_\text{th}³(\psi_c)} \text{exp}\left[ - \frac{\epsilon}{v_\text{th}²(\psi_c)}\right].
Parameters
----------
n : tuple[float | Callable, Perturbation]
Density background and optional perturbation.
vth : tuple[float | Callable, Perturbation]
Thermal-speed background and optional perturbation.
equil : FluidEquilibriumWithB, optional
Fluid equilibrium used to evaluate background profiles in the magnetic geometry.
volume_form : bool, default=True
If ``True``, represent the distribution as a volume form and include the appropriate
velocity-space Jacobian when evaluating it.
"""
def __init__(
self,
n: tuple[float | Callable, Perturbation] = (1.0, None),
vth: tuple[float | Callable, Perturbation] = (1.0, None),
equil: FluidEquilibriumWithB = None,
volume_form: bool = True,
):
# use setter to store input parameters
self.params = copy.deepcopy(locals())
self.check_maxw_params()
# volume form represenation
self._volume_form = volume_form
self._equil = equil
# factors multiplied onto the defined moments n and vth (can be set via setter)
self._moment_factors = {
"n": 1.0,
"vth": 1.0,
}
@property
def coords(self):
r"""Coordinates of the CanonicalMaxwellian, :math:`(\epsilon, \mu, \psi_c)`."""
return "constants_of_motion"
@property
def equil(self) -> FluidEquilibriumWithB:
"""One of :mod:`~struphy.fields_background.equils`
in case that moments are to be set in that way, None otherwise.
"""
return self._equil
[docs]
def check_maxw_params(self):
for k, v in self.params.items():
assert isinstance(k, str)
if isinstance(v, tuple):
assert len(v) == 2
assert isinstance(v[0], (float, int, Callable))
assert isinstance(v[1], Perturbation) or v[1] is None
[docs]
def velocity_jacobian_det(self, eta1, eta2, eta3, energy):
r"""TODO"""
assert eta1.ndim == 1
assert eta2.ndim == 1
assert eta3.ndim == 1
if self.params["type"] == "Particles6D":
return xp.sqrt(2.0 * energy) * 4.0 * xp.pi
else:
# call equilibrium
etas = (xp.vstack((eta1, eta2, eta3)).T).copy()
absB0 = self.equil.absB0(etas)
return xp.sqrt(energy) * 2.0 * xp.sqrt(2.0) / absB0
[docs]
def gaussian(self, e, vth=1.0):
"""3-dim. normal distribution, to which array-valued thermal velocities can be passed.
Parameters
----------
e : float | array-like
Energy.
vth : float | array-like
Thermal velocity evaluated at psic.
Returns
-------
An array of size(e).
"""
if isinstance(vth, xp.ndarray):
assert e.shape == vth.shape, f"{e.shape =} but {vth.shape =}"
return 2.0 * xp.sqrt(e / xp.pi) / vth**3 * xp.exp(-e / vth**2)
def __call__(self, *args):
"""Evaluates the canonical Maxwellian distribution function.
There are two use-cases for this function in the code:
1. Evaluating for particles ("flat evaluation", inputs are all 1D of length N_p)
2. Evaluating the function on a meshgrid (in phase space).
Hence all arguments must always have
1. the same shape
2. either ndim = 1 or ndim = 3.
Parameters
----------
*args : array_like
Position-velocity arguments in the order energy, magnetic moment, canonical toroidal momentum.
Returns
-------
f : xp.ndarray
The evaluated Maxwellian.
"""
# Check that all args have the same shape
shape0 = xp.shape(args[0])
for i, arg in enumerate(args):
assert xp.shape(arg) == shape0, f"Argument {i} has {xp.shape(arg) =}, but must be {shape0 =}."
assert xp.ndim(arg) == 1 or xp.ndim(arg) == 3, (
f"{xp.ndim(arg) =} not allowed for canonical Maxwellian evaluation."
) # flat or meshgrid evaluation
# Get result evaluated with each particles' psic
res = self.n(args[2])
vths = self.vth(args[2])
# take care of correct broadcasting, assuming args come from phase space meshgrid
if xp.ndim(args[0]) == 3:
# move eta axes to the back
arg_t = xp.moveaxis(args[0], 0, -1)
arg_t = xp.moveaxis(arg_t, 0, -1)
arg_t = xp.moveaxis(arg_t, 0, -1)
# broadcast
res_broad = res + 0.0 * arg_t
# move eta axes to the front
res = xp.moveaxis(res_broad, -1, 0)
res = xp.moveaxis(res, -1, 0)
res = xp.moveaxis(res, -1, 0)
# Multiply result with gaussian in energy
if xp.ndim(args[0]) == 3:
vth_broad = vths + 0.0 * arg_t
vth = xp.moveaxis(vth_broad, -1, 0)
vth = xp.moveaxis(vth, -1, 0)
vth = xp.moveaxis(vth, -1, 0)
else:
vth = vths
res *= self.gaussian(args[0], vth=vth)
return res
@property
def volume_form(self):
"""Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian |v_perp|)."""
return self._volume_form
@property
def moment_factors(self):
"""Collection of factors multiplied onto the defined moments n, u, and vth."""
return self._moment_factors
@moment_factors.setter
def moment_factors(self, **kwargs):
for kw, arg in kwargs:
self._moment_factors[kw] = arg
[docs]
def rc(self, psic):
r""" Square root of radially normalized canonical toroidal momentum.
.. math::
\begin{aligned}
r_c^2 &= \frac{\psi_c - \psi_\text{axis}}{\psi_\text{edge} - \psi_\text{axis}} \,,
\\[3mm]
r_c &= \begin{cases}
\sqrt{\frac{\psi_c - \psi_\text{axis}}{\psi_\text{edge} - \psi_\text{axis}}} & \text{if} \quad \frac{\psi_c - \psi_\text{axis}}{\psi_\text{edge} - \psi_\text{axis}} \geq 0 \,, \\
-\sqrt{\frac{\psi_c - \psi_\text{axis}}{\psi_\text{edge} - \psi_\text{axis}}} & \text{if} \quad \frac{\psi_c - \psi_\text{axis}}{\psi_\text{edge} - \psi_\text{axis}} < 0 \,,
\end{cases}
\end{aligned}
where :math:`\psi_\text{axis}` and :math:`\psi_\text{edge}` are poloidal magnetic flux function at the center and edge of poloidal plane respectively.
Parameters
----------
psic : numpy.arrays
Evaluation points. All arrays must be of same shape (can be 1d for flat evaluation).
Returns
-------
A numpy.array of the evaluated :math:`r_c`.
"""
# calculate rc²
rc_squared = (psic - self.equil.psi_range[0]) / (self.equil.psi_range[1] - self.equil.psi_range[0])
# sorting out indices of negative rc²
neg_index = xp.logical_not(rc_squared >= 0)
# make them positive
rc_squared[neg_index] *= -1
# calculate rc
rc = xp.sqrt(rc_squared)
rc[neg_index] *= -1
return rc
[docs]
def n(self, psic, add_perturbation: bool = None):
"""Density as background + perturbation.
Parameters
----------
psic : numpy.array
Evaluation points. All arrays must be of same shape (can be 1d for flat evaluation).
Returns
-------
A float (background value) or a numpy.array of the evaluated density.
"""
# collect arguments
assert isinstance(psic, xp.ndarray)
# assuming that input comes from meshgrid.
if psic.ndim == 3:
psic = psic[0, 0, :]
# set background density
if isinstance(self.params["n"][0], (float, int)):
res = self.params["n"][0] + 0.0 * psic
else:
nfun = self.params["n"][1]
# for typ, params in mom_funcs.items():
# nfun = getattr(moment_functions, typ)(**params)
res = nfun(eta1=self.rc(psic))
# add perturbation
if add_perturbation is None:
add_perturbation = self.add_perturbation
perturbation = self.params["n"][1]
if perturbation is not None and add_perturbation:
assert isinstance(perturbation, Perturbation)
res = perturbation(eta1=self.rc(psic))
# if eta1.ndim == 1:
# out += perturbation(eta1, eta2, eta3)
# else:
# out += perturbation(*etas)
return res * self.moment_factors["n"]
[docs]
def vth(self, psic):
"""Thermal velocities as background + perturbation.
Parameters
----------
psic : numpy.arrays
Evaluation points. All arrays must be of same shape (can be 1d for flat evaluation).
Returns
-------
A list[float] (background value) or a list[numpy.array] of the evaluated thermal velocities.
"""
# collect arguments
assert isinstance(psic, xp.ndarray)
# assuming that input comes from meshgrid.
if psic.ndim == 3:
psic = psic[0, 0, :]
res = self.params["vth"][0] + 0.0 * psic
# TODO: add perturbation
return res * self.moment_factors["vth"]
@property
def add_perturbation(self) -> bool:
if not hasattr(self, "_add_perturbation"):
self._add_perturbation = True
return self._add_perturbation
@add_perturbation.setter
def add_perturbation(self, new):
assert isinstance(new, bool)
self._add_perturbation = new
[docs]
class ColdPlasma(Maxwellian):
r"""Base class for a distribution as a Dirac-delta in velocity (vth = 0).
The __call__ method returns the density evaluation."""
[docs]
@classmethod
def default_maxw_params(cls):
"""Default parameters dictionary defining the constant value of the constant background."""
return {
"n": 5.0,
"u1": 0.0,
"u2": 0.0,
"u3": 0.0,
"vth1": 0.0,
"vth2": 0.0,
"vth3": 0.0,
}
def __init__(
self,
n: tuple[float | Callable, Perturbation] = (1.0, None),
u1: tuple[float | Callable, Perturbation] = (0.0, None),
u2: tuple[float | Callable, Perturbation] = (0.0, None),
u3: tuple[float | Callable, Perturbation] = (0.0, None),
equil: FluidEquilibriumWithB = None,
):
# use setter to store input parameters
self.params = copy.deepcopy(locals())
self._params["vth1"] = (0.0, None)
self._params["vth2"] = (0.0, None)
self._params["vth3"] = (0.0, None)
self.check_maxw_params()
self._equil = equil
@property
def coords(self):
"""Coordinates of the constant background."""
return None
@property
def vdim(self):
"""Dimension of the velocity space (vdim = 0)."""
return 0
@property
def is_polar(self):
"""List of booleans of length vdim. True for a velocity coordinate that is a radial polar coordinate (v_perp)."""
return []
@property
def volume_form(self):
"""Boolean. True if the background is represented as a volume form (thus including the velocity Jacobian)."""
return False
@property
def equil(self) -> FluidEquilibriumWithB:
"""Fluid background with B-field."""
return self._equil
[docs]
def velocity_jacobian_det(self, eta1, eta2, eta3, *v):
"""Jacobian determinant of the velocity coordinate transformation."""
return 1.0
[docs]
def n(self, eta1, eta2, eta3):
"""Zero-th moment (density)."""
out = self._evaluate_moment(eta1, eta2, eta3, name="n")
return out
[docs]
def u(self, eta1, eta2, eta3):
"""Mean velocities."""
out = []
out += [self._evaluate_moment(eta1, eta2, eta3, name="u1")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="u2")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="u3")]
return out
[docs]
def vth(self, eta1, eta2, eta3):
"""Thermal velocities (are zero here, see __init__)."""
out = []
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth1")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth2")]
out += [self._evaluate_moment(eta1, eta2, eta3, name="vth3")]
return out
def __call__(self, eta1, eta2, eta3):
return self.n(eta1, eta2, eta3)
