"Maxwellian (Gaussian) distributions in velocity space."
import copy
import logging
import warnings
from inspect import signature
from typing import Callable, Literal
import cunumpy as xp
from struphy.fields_background.base import AxisymmMHDequilibrium, FluidEquilibriumWithB
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, check_B0_option, eval_B0_option
logger = logging.getLogger("struphy")
[docs]
class Maxwellian3D(Maxwellian):
r"""A :class:`~struphy.kinetic_background.base.Maxwellian` depending :math:`(\eta_1, \eta_2, \eta_3)`
and 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).
uniform_on_disc : bool
Whether the density n is uniform on the disc.
"""
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),
uniform_on_disc: bool = False,
):
# 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 vdim(self):
"""Dimension of the velocity space."""
return 3
@property
def velocity_coords(self) -> LiteralOptions.VelocityCoordinates:
"""Velocity coordinates of the background."""
return "cartesian"
[docs]
def velocity_jacobian_det(self, eta1, eta2, eta3, vx, vy, vz):
"""Jacobian determinant is 1 (Cartesian velocity coordinates).
Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords)
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
vx, vy, vz : array_like
Velocity evaluation points.
Returns
-------
out : array-like
The Jacobian determinant evaluated at given logical coordinates.
-------
"""
assert eta1.ndim == eta2.ndim == eta3.ndim == 1
assert vx.ndim == vy.ndim == vz.ndim == 1
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, new_factors: dict):
for kw, arg in new_factors.items():
assert kw in self._moment_factors, f"Unknown moment factor {kw}."
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, \mu)`, :math:`n=2`,
where :math:`v_\parallel = \mathbf v \cdot \mathbf b_0` and
:math:`\mu = v_\perp^2/(2B_0)` is the magnetic moment, with :math:`B_0` the background magnetic field strength.
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).
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).
B0: float | Callable | "equil"
Background magnetic field strength :math:`B_0` in :math:`\mu = v_\perp^2/(2B_0)`, see :meth:`eval_B0`:
a constant, a callable of the stacked logical positions ``etas`` of shape ``(3, N)`` returning shape ``(N,)``,
or ``"equil"`` (default) for the local field strength :math:`|\mathbf B_0|(\boldsymbol \eta)` of the
equilibrium attached with :meth:`attach_equil` (done automatically for the background and initial
condition of a :class:`~struphy.pic.particles.Particles5D` species).
uniform_on_disc : bool
Whether the density n is uniform on the disc.
"""
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),
volume_form: bool = True,
B0: float | Callable | Literal["equil"] = "equil",
uniform_on_disc: bool = False,
):
# use setter to store input parameters
self.params = copy.deepcopy(locals())
self.check_maxw_params()
check_B0_option(B0)
# equilibrium providing |B_0| for B0="equil" (see attach_equil)
self._B0_equil = None
# last (B0 option, equilibrium, positions, values) of eval_B0: f0 and its velocity Jacobian are evaluated
# at the same markers
self._B0_cache = None
# volume form represenation
self._volume_form = volume_form
# 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 vdim(self):
"""Dimension of the velocity space."""
return 2
@property
def velocity_coords(self) -> LiteralOptions.VelocityCoordinates:
"""Velocity coordinates of the background."""
return "vpara_mu"
[docs]
def velocity_jacobian_det(self, eta1, eta2, eta3, v_para, mu):
r"""Jacobian determinant of the velocity coordinate transformation to :math:`(v_\parallel, mu)`, is :math:`B_0`.
Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords)
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
v_para, mu : array_like
Parallel velocity and magnetic moment evaluation points.
Returns
-------
out : array-like
The Jacobian determinant evaluated at given logical coordinates.
-------
"""
assert eta1.ndim == eta2.ndim == eta3.ndim == 1
assert v_para.ndim == mu.ndim == 1
return self.eval_B0(xp.stack((eta1, eta2, eta3), axis=0))
[docs]
def attach_equil(self, equil: FluidEquilibriumWithB):
"""Attach the equilibrium whose field strength :math:`|\\mathbf B_0|` is used for ``B0="equil"``
(replacing an equilibrium attached before). Does nothing for a constant or callable ``B0``.
The background object itself is changed: if it is shared by particles with different equilibria (e.g.
two simulations set up in one script before running), the last equilibrium attached is used by all of
them, and a warning is logged."""
if isinstance(self.params["B0"], str):
if equil is None:
return
assert isinstance(equil, FluidEquilibriumWithB), (
f'GyroMaxwellian2D with B0="equil" needs an equilibrium with magnetic field, got {type(equil).__name__}.'
)
if self._B0_equil is not None and self._B0_equil is not equil:
logger.warning(
'GyroMaxwellian2D with B0="equil" gets a different equilibrium attached: the background object '
"is shared, and all particles using it take |B_0| from the new equilibrium from now on. This is "
"fine if the simulation set up with the previous equilibrium is not used any more (e.g. a scan); "
"otherwise use a separate background object for each simulation."
)
self._B0_equil = equil
self._B0_cache = None
[docs]
def eval_B0(self, etas):
"""Evaluate the magnetic field strength :math:`B_0` in :math:`\\mu = v_\\perp^2/(2B_0)` at flat positions.
Parameters
----------
etas : array_like
Stacked logical positions of shape ``(3, N)``.
Returns
-------
out : array-like
:math:`B_0` at the positions, shape ``(N,)``.
"""
B0 = self.params["B0"]
if not (isinstance(B0, str) or callable(B0)):
return eval_B0_option(B0, etas)
# a position-dependent B0 is often evaluated twice at the same markers (f0 and velocity_jacobian_det in
# the control variate); the comparison is much cheaper than an evaluation of |B_0|
cache = self._B0_cache
if (
cache is not None
and cache[0] is B0
and cache[1] is self._B0_equil
and cache[2].shape == etas.shape
and xp.array_equal(cache[2], etas)
):
return cache[3]
out = eval_B0_option(B0, etas, self._B0_equil, owner="GyroMaxwellian2D")
self._B0_cache = (B0, self._B0_equil, xp.array(etas, copy=True), out)
return out
[docs]
def eval_B0_on_args(self, etas):
""":meth:`eval_B0` at the positions ``eta1, eta2, eta3`` of a Maxwellian evaluation, flat (1d arrays)
or on a phase-space meshgrid (``3 + vdim`` dimensions, positions on the first 3 axes): there it is
evaluated on the spatial points only and broadcast to the velocity axes. A constant ``B0`` is returned
as float."""
B0 = self.params["B0"]
if not (isinstance(B0, str) or callable(B0)):
return float(B0)
if xp.ndim(etas[0]) == 1:
return self.eval_B0(xp.stack(etas, axis=0))
# spatial part of the meshgrid (velocity axes at their first point)
spatial = [a[(slice(None),) * 3 + (0,) * (xp.ndim(a) - 3)] for a in etas]
out = self.eval_B0(xp.stack([xp.ravel(a) for a in spatial], axis=0))
out = xp.reshape(out, xp.shape(spatial[0]) + (1,) * (xp.ndim(etas[0]) - 3))
return xp.broadcast_to(out, xp.shape(etas[0]))
@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 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, new_factors: dict):
for kw, arg in new_factors.items():
assert kw in self._moment_factors, f"Unknown moment factor {kw}."
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]
class GyroMaxwellian2Dvperp(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 = \mathbf 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).
uniform_on_disc : bool
Whether the density n is uniform on the disc (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,
uniform_on_disc: bool = False,
):
# 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 vdim(self):
"""Dimension of the velocity space."""
return 2
@property
def velocity_coords(self) -> LiteralOptions.VelocityCoordinates:
"""Velocity coordinates of the background."""
return "vpara_vperp"
[docs]
def velocity_jacobian_det(self, eta1, eta2, eta3, v_para, v_perp):
r"""Jacobian determinant of the velocity coordinate transformation to :math:`(v_\parallel, v_\perp)`, is :math:`v_\perp`.
Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords)
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
v_para, v_perp : array_like
Parallel and perpendicular velocity evaluation points.
Returns
-------
out : array-like
The Jacobian determinant evaluated at given logical coordinates.
-------
"""
assert eta1.ndim == eta2.ndim == eta3.ndim == 1
assert v_para.ndim == v_perp.ndim == 1
return v_perp
@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, new_factors: dict):
for kw, arg in new_factors.items():
assert kw in self._moment_factors, f"Unknown moment factor {kw}."
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]
class CanonicalMaxwellian2D(GyroMaxwellian2D):
r"""Canonical Maxwellian distribution function in
:math:`(\eta_1, \eta_2, \eta_3, v_\parallel, \mu)` coordinates.
Uses caching for evaluation of the canonical toroidal momentum in these coordinates.
The distribution is parameterized by the density and thermal speed as functions of the
(shifted) canonical toroidal momentum :math:`\psi_c`:
.. math::
\psi_c = P_\phi - \text{sign}(v_\parallel)\sqrt{2(\epsilon - \mu B_\text{ref})}\frac{m_sF}{q_sB_\text{ref}}
\mathcal{H}(\epsilon - \mu B_\text{ref})\,,\qquad P_\phi = \psi + \frac{m_s F}{q_s B}v_\parallel\,,
where :math:`P_\phi` is the canonical toroidal momentum and the reference field :math:`B_\text{ref}` is set by
``shift`` (see below). The shift (Angelino et al., Phys. Plasmas 13, 052304 (2006)) makes :math:`\psi_c = \psi`
where :math:`|\mathbf B| = B_\text{ref}` on the orbit of a passing particle, such that the profiles in
:math:`\psi_c` are close to the profiles in the flux label of the position.
- 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, vth : tuple
Moments of the canonical Maxwellian as tuples. The first entry defines the background
(float for constant background or callable), the second entry defines a Perturbation (can be None).
maxw_params : dict
Parameters for the kinetic background.
vth : tuple[float | Callable, Perturbation]
Thermal-speed background and optional perturbation.
equil : AxisymmMHDequilibrium, 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.
cache_size : int, optional
Number of rows in the cache buffer for :math:`\psi_c` evaluation. If ``None``, no caching is used.
Must be able to accomodate all markers on the current process.
shift : str | None, default="Bmax"
Reference field :math:`B_\text{ref}` of the shift of the canonical toroidal momentum:
* ``"Bmax"``: the maximum of :math:`|\mathbf B|` on the flux surface :math:`\psi = P_\phi`, see
:meth:`~struphy.fields_background.base.AxisymmMHDequilibrium.absB_max`. A particle can only bounce
(:math:`v_\parallel = 0`) on this surface, hence the shift vanishes for all particles that can bounce,
:math:`\psi_c` is a function of the constants of motion :math:`P_\phi, \epsilon, \mu` and the
distribution is an exact equilibrium of the guiding-center motion (with the constant poloidal current
function :math:`F = -B_0 R_0` assumed in :math:`P_\phi`, as in
:class:`~struphy.fields_background.equils.AdhocTorus`). It has a net parallel flow.
* ``"B0"``: the on-axis field (Angelino et al.). The shift jumps at the bounce points of the trapped
particles with :math:`\epsilon > \mu B_0` (the sign of :math:`v_\parallel` changes there), so the
distribution is not an exact equilibrium (default up to Struphy 3.6).
* ``None``: no shift, :math:`\psi_c = P_\phi` (exact equilibrium, but the profiles are shifted by the
drift-orbit width with respect to the position).
shifted : bool, optional
Deprecated, use ``shift``: ``True`` is ``shift="B0"``, ``False`` is ``shift=None``.
"""
shift_modes = {None: 0, "B0": 1, "Bmax": 2}
"""Integer codes of the values of ``shift``, as passed to the compiled kernels."""
def __init__(
self,
n: tuple[float | Callable, Perturbation] = (1.0, None),
vth: tuple[float | Callable, Perturbation] = (1.0, None),
volume_form: bool = True,
uniform_on_disc: bool = False,
equil: AxisymmMHDequilibrium = None,
epsilon: float = 1.0,
cache_size: int | None = None,
shift: str | None = "Bmax",
shifted: bool | None = None,
):
assert isinstance(equil, AxisymmMHDequilibrium)
if shifted is not None:
warnings.warn(
'CanonicalMaxwellian2D(shifted=...) is deprecated, use shift="B0" (shifted=True) or shift=None '
"(shifted=False).",
DeprecationWarning,
stacklevel=2,
)
shift = "B0" if shifted else None
assert shift in self.shift_modes, f"shift must be one of {list(self.shift_modes)}, but is {shift!r}."
super().__init__(
n=n,
u_para=(0.0, None),
u_perp=(0.0, None),
vth_para=vth,
vth_perp=vth,
volume_form=volume_form,
B0=equil.absB0,
uniform_on_disc=uniform_on_disc,
)
# store additional parameters
self._equil = equil
self._epsilon = epsilon
self._shift = shift
self.params["vth"] = vth
self.params["equil"] = equil
self.params["epsilon"] = epsilon
self.params["cache_size"] = cache_size
self.params["shift"] = shift
# factors multiplied onto the defined moments n and vth (can be set via setter)
self._moment_factors = {
"n": 1.0,
"vth": 1.0,
}
# create cache for psi_c evaluation
if cache_size is not None:
self.cbufs = {}
self.cbufs["absB0"] = xp.empty(cache_size, dtype=float)
self.cbufs["x"] = xp.empty(cache_size, dtype=float)
self.cbufs["y"] = xp.empty(cache_size, dtype=float)
self.cbufs["z"] = xp.empty(cache_size, dtype=float)
self.cbufs["R"] = xp.empty(cache_size, dtype=float)
self.cbufs["P"] = xp.empty(cache_size, dtype=float)
self.cbufs["Z"] = xp.empty(cache_size, dtype=float)
self.cbufs["psi"] = xp.empty(cache_size, dtype=float)
self.cbufs["energy"] = xp.empty(cache_size, dtype=float)
self.cbufs["psic"] = xp.empty(cache_size, dtype=float)
self.cbufs["positive_mask"] = xp.empty(cache_size, dtype=bool)
self.cbufs["correction"] = xp.empty(cache_size, dtype=float)
self.cbufs["Bref"] = xp.empty(cache_size, dtype=float)
logger.debug(f"Created {len(self.cbufs)} cache buffers for psi_c evaluation, each with size {cache_size}.")
else:
self.cbufs = None
# psi_c of the current __call__, shared by the moments n and vth of that call (None outside of __call__)
self._psi_c_call = None
def __call__(self, *phase_space_coords):
r"""Evaluates the canonical Maxwellian, see :meth:`~struphy.kinetic_background.base.Maxwellian.__call__`.
:math:`\psi_c` is evaluated once from the given coordinates and shared by the moments :math:`n` and
:math:`v_\text{th}` of this call. It is not kept between calls: markers move and change slots (MPI sorting,
refilling) between calls, so a stored :math:`\psi_c` can belong to other phase-space points."""
if any(callable(self.params[name][0]) for name in ("n", "vth")):
self._psi_c_call = self.eval_psic(*self._psic_coords(*phase_space_coords))
try:
return super().__call__(*phase_space_coords)
finally:
self._psi_c_call = None
@staticmethod
def _psic_coords(eta1, eta2, eta3, v_parallel, mu):
"""Arrange phase-space points as the input of :meth:`eval_psic`: one (5, n_markers) array for flat
(marker) evaluation, or the five 5d meshgrid arrays."""
if eta1.ndim == 1:
return [xp.stack((eta1, eta2, eta3, v_parallel, mu), axis=0)]
elif eta1.ndim == 5:
return (eta1, eta2, eta3, v_parallel, mu)
else:
raise ValueError(f"Input arrays must be 1d or 5d (from meshgrid), but got {eta1.ndim}d.")
@property
def equil(self) -> AxisymmMHDequilibrium:
"""One of :mod:`~struphy.fields_background.equils`
in case that moments are to be set in that way, None otherwise.
"""
return self._equil
@property
def epsilon(self) -> float:
"""Epsilon parameter in the canonical toroidal momentum."""
return self._epsilon
@property
def shift(self) -> str | None:
"""Reference field of the shift of the canonical toroidal momentum (``"Bmax"``, ``"B0"`` or ``None``)."""
return self._shift
@GyroMaxwellian2D.moment_factors.setter
def moment_factors(self, new_factors: dict):
# n and vth are scalar factors here (no u)
for kw, arg in new_factors.items():
assert kw in self._moment_factors, f"Unknown moment factor {kw}."
self._moment_factors[kw] = arg
def _evaluate_moment(
self,
eta1,
eta2,
eta3,
v_parallel,
mu,
*,
name: str = "n",
add_perturbation: bool = None,
):
"""Scalar moment evaluation as background + perturbation.
Incontrast to standard Maxwellians, here the moments are evaluated
at the phase space coordinates.
Parameters
----------
eta1, eta2, eta3, v_parallel, mu : numpy.arrays
Phase space evaluation points. All arrays must be of same shape (can be 1d for flat evaluation).
name : str
Which moment to evaluate (see varaible "dct" below).
add_perturbation : bool | None
Whether to add the perturbation defined in params. If None, is taken from self.add_perturbation.
Returns
-------
A float (background value) or a numpy.array of the evaluated scalar moment.
"""
# collect arguments
assert isinstance(eta1, xp.ndarray)
assert isinstance(eta2, xp.ndarray)
assert isinstance(eta3, xp.ndarray)
assert isinstance(v_parallel, xp.ndarray)
assert isinstance(mu, xp.ndarray)
assert eta1.shape == eta2.shape == eta3.shape == v_parallel.shape == mu.shape
params = self.params[name]
assert isinstance(params, tuple)
assert len(params) == 2
if eta1.ndim not in (1, 5):
raise ValueError(f"Input arrays must be 1d or 5d (from meshgrid), but got {eta1.ndim}d.")
# initialize output
out = 0.0 * eta1
# evaluate background
background = params[0]
if isinstance(background, (float, int)):
out += background
else:
assert callable(background)
sig = signature(background)
assert len(sig.parameters) == 1, (
f"Background function {background} must take one argument (psi_c), but takes {len(sig.parameters)}."
)
# inside __call__, psi_c has been evaluated at the same points; otherwise (direct call of n or vth)
# evaluate it here
psi_c = self._psi_c_call
if psi_c is None:
psi_c = self.eval_psic(*self._psic_coords(eta1, eta2, eta3, v_parallel, mu))
out += background(psi_c)
# add perturbation
if add_perturbation is None:
add_perturbation = self.add_perturbation
perturbation = params[1]
if perturbation is not None and add_perturbation:
assert isinstance(perturbation, Perturbation)
if eta1.ndim == 1:
out += perturbation(eta1, eta2, eta3)
else:
raise NotImplementedError("Perturbation evaluation for meshgrid input not implemented yet.")
# uniform density on disc (n=2 eta_1)
if name == "n" and self.params.get("uniform_on_disc", False):
out *= 2.0 * eta1
return out
[docs]
def eval_psic(self, *coords, use_cbufs: bool = True):
r"""Shifted canonical toroidal momentum evaluated at given particle positions and velocities.
If ``use_cbufs`` is ``False``, the cache buffers are not used (and not overwritten) for marker evaluation.
Otherwise the result for markers is a view of the buffer ``cbufs["psic"]``, which the next evaluation
overwrites: copy it to keep it."""
cbufs = self.cbufs if use_cbufs else None
B0 = self.equil.params["B0"]
R0 = self.equil.params["R0"]
if len(coords) == 1:
if use_cbufs and self.cbufs is None:
logger.warning(
f"Initialize {self.__class__.__name__} with `cache_size` for faster psi_c evaluation for markers!"
)
etas = coords[0][:3] # these are views (no mem allocation)
vparallel = coords[0][3]
mu = coords[0][4]
n_markers = etas.shape[1]
if cbufs is None:
absB0 = self.equil.absB0(etas)
x, y, z = self.equil.domain(etas)
else:
absB0 = cbufs["absB0"][:n_markers]
x = cbufs["x"][:n_markers]
y = cbufs["y"][:n_markers]
z = cbufs["z"][:n_markers]
absB0[:] = self.equil.absB0(etas)
x[:], y[:], z[:] = self.equil.domain(etas)
else:
assert len(coords) == 5
assert coords[0].ndim == coords[1].ndim == coords[2].ndim == coords[3].ndim == coords[4].ndim == 5
eta1 = coords[0][:, :, :, 0, 0]
eta2 = coords[1][:, :, :, 0, 0]
eta3 = coords[2][:, :, :, 0, 0]
etas = (eta1, eta2, eta3)
absB0 = self.equil.absB0(*etas)[:, :, :, None, None]
x, y, z = self.equil.domain(*etas)
vparallel = coords[3]
mu = coords[4]
if cbufs is None or len(coords) != 1:
R, P, Z = self.equil.inverse_map(x, y, z)
psi = self.equil.psi(R, Z)
if len(coords) != 1:
psi = psi[:, :, :, None, None]
energy = 1 / 2 * vparallel**2 + mu * absB0
psi_c = psi - self._epsilon * B0 * R0 / absB0 * vparallel
if self._shift is not None:
self._add_shift(
psi_c,
energy,
vparallel,
mu,
positive_mask=xp.empty(psi_c.shape, dtype=bool),
correction=xp.empty(psi_c.shape, dtype=float),
Bref=xp.empty(psi_c.shape, dtype=float),
)
else:
R = cbufs["R"][:n_markers]
P = cbufs["P"][:n_markers]
Z = cbufs["Z"][:n_markers]
psi = cbufs["psi"][:n_markers]
energy = cbufs["energy"][:n_markers]
psi_c = cbufs["psic"][:n_markers]
positive_mask = cbufs["positive_mask"][:n_markers]
correction = cbufs["correction"][:n_markers]
Bref = cbufs["Bref"][:n_markers]
R[:], P[:], Z[:] = self.equil.inverse_map(x, y, z)
psi[:] = self.equil.psi(R, Z)
energy[:] = 1 / 2 * vparallel**2 + mu * absB0
psi_c[:] = psi - self._epsilon * B0 * R0 / absB0 * vparallel
if self._shift is not None:
self._add_shift(
psi_c, energy, vparallel, mu, positive_mask=positive_mask, correction=correction, Bref=Bref
)
return psi_c
def _add_shift(self, psi_c, energy, vparallel, mu, *, positive_mask, correction, Bref):
r"""Add the shift :math:`-\text{sign}(v_\parallel)\sqrt{2(\epsilon - \mu B_\text{ref})}
\frac{m_sF}{q_sB_\text{ref}}\mathcal{H}(\epsilon - \mu B_\text{ref})` in place to the canonical toroidal
momentum ``psi_c`` (:math:`P_\phi` on input). ``positive_mask``, ``correction`` and ``Bref`` are work arrays
of the shape of ``psi_c``."""
B0 = self.equil.params["B0"]
R0 = self.equil.params["R0"]
if self._shift == "B0":
positive_mask[...] = (energy - mu * B0) > 0
correction[...] = 0.0
correction[positive_mask] = (
self._epsilon
* xp.sign(vparallel[positive_mask])
* xp.sqrt(2 * (energy[positive_mask] - mu[positive_mask] * B0))
* R0
)
elif self._shift == "Bmax":
# maximum of |B| on the flux surface psi = P_phi, the only surface where the particle can bounce
Bref[...] = self.equil.absB_max(psi_c)
positive_mask[...] = (energy - mu * Bref) > 0
correction[...] = 0.0
correction[positive_mask] = (
self._epsilon
* xp.sign(vparallel[positive_mask])
* xp.sqrt(2 * (energy[positive_mask] - mu[positive_mask] * Bref[positive_mask]))
* R0
* (B0 / Bref[positive_mask])
)
else:
raise ValueError(f"Unknown shift {self._shift!r}.")
psi_c[...] += correction
[docs]
def shift_kernel_args(self):
"""Arguments ``(shift_mode, psi_axis, psi_edge, bmax_table)`` of the compiled kernels that evaluate the
shifted canonical toroidal momentum of markers, see :attr:`shift_modes` and
:attr:`~struphy.fields_background.base.AxisymmMHDequilibrium.absB_max_table`."""
return self.shift_kernel_args_for(self.equil, self.shift)
[docs]
@classmethod
def shift_kernel_args_for(cls, equil: AxisymmMHDequilibrium, shift: str | None):
"""Kernel arguments of :meth:`shift_kernel_args` for a given equilibrium and ``shift``."""
assert shift in cls.shift_modes, f"shift must be one of {list(cls.shift_modes)}, but is {shift!r}."
psi_axis, psi_edge = equil.psi_range
if shift == "Bmax":
bmax_table = xp.ascontiguousarray(equil.absB_max_table, dtype=float)
else:
bmax_table = xp.zeros(2, dtype=float)
return cls.shift_modes[shift], float(psi_axis), float(psi_edge), bmax_table
[docs]
def eval_rc(self, eta1, eta2, eta3, vparallel, mu):
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.
"""
# calculate psic
psic = self.eval_psic(eta1, eta2, eta3, vparallel, mu)
# 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, eta1, eta2, eta3, vparallel, mu):
"""Zero-th moment (density)."""
out = self._evaluate_moment(eta1, eta2, eta3, vparallel, mu, name="n")
return out * self.moment_factors["n"]
[docs]
def u(self, eta1, eta2, eta3, vparallel, mu):
"""Mean velocities (zero for the canonical Maxwellian)."""
return [0.0 * eta1, 0.0 * eta1]
[docs]
def vth(self, eta1, eta2, eta3, vparallel, mu):
"""Thermal velocities."""
out = self._evaluate_moment(eta1, eta2, eta3, vparallel, mu, name="vth")
out = out * self.moment_factors["vth"]
return [out, out]
[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,
uniform_on_disc: bool = False,
):
# 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 vdim(self):
"""Dimension of the velocity space (vdim = 0)."""
return 0
@property
def velocity_coords(self) -> LiteralOptions.VelocityCoordinates:
"""Velocity coordinates of the background."""
return None
@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 + 0.0 * eta1
[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)
