import copy
import logging
import cunumpy as xp
from struphy.fields_background import equils
from struphy.fields_background.base import FluidEquilibrium, FluidEquilibriumWithB
from struphy.fields_background.projected_equils import ProjectedFluidEquilibriumWithB
from struphy.geometry.base import Domain
from struphy.geometry.utilities import TransformedPformComponent
from struphy.initial.base import Perturbation
from struphy.kinetic_background import maxwellians
from struphy.kinetic_background.base import Maxwellian, SumKineticBackground
from struphy.pic.base import ORBIT_POSITIONS, Particles
from struphy.pic.diagnostics.kernels.eval_canonical_toroidal_moment_5d import eval_canonical_toroidal_moment_5d
from struphy.pic.diagnostics.kernels.eval_canonical_toroidal_moment_6d import eval_canonical_toroidal_moment_6d
from struphy.pic.diagnostics.kernels.eval_energy_5d import eval_energy_5d
from struphy.pic.diagnostics.kernels.eval_guiding_center_from_6d import eval_guiding_center_from_6d
from struphy.pic.diagnostics.kernels.eval_magnetic_background_energy import eval_magnetic_background_energy
from struphy.pic.diagnostics.kernels.eval_magnetic_energy_PBb import eval_magnetic_energy_PBb
from struphy.pic.diagnostics.kernels.eval_magnetic_moment_5d import eval_magnetic_moment_5d
logger = logging.getLogger("struphy")
[docs]
class Particles6D(Particles):
"""
Particles in the full 6D phase space :math:`(\\boldsymbol \\eta, \\mathbf v) \\in [0, 1]^3 \\times \\mathbb R^3`,
as used e.g. in full-orbit (Vlasov) kinetic models.
Each marker carries a logical (curvilinear) position :math:`\\boldsymbol \\eta_p` together with a velocity
:math:`\\mathbf v_p` expressed in the *Cartesian* velocity space attached to that position
(i.e. velocities are not transformed by the curvilinear map, unlike positions).
See :class:`~struphy.pic.base.Particles` for the structure of the numpy marker array and the meaning of its columns.
"""
# Class properties
vdim = 3
"""Dimension of the (Cartesian) velocity space, here 3."""
coordinate_labels = ("$\\eta_1$", "$\\eta_2$", "$\\eta_3$", "$v_x$", "$v_y$", "$v_z$")
"""Labels for the coordinates in the phase space. Length is 6, with the first 3 being the spatial coordinates and the last 3 being the velocity coordinates."""
orbit_quantities = (
*ORBIT_POSITIONS,
(3, "v1", "$v_x$", "Cartesian velocity x"),
(4, "v2", "$v_y$", "Cartesian velocity y"),
(5, "v3", "$v_z$", "Cartesian velocity z"),
(6, "weight", "$w$", "marker weight"),
)
"""Marker columns saved as orbits, see :attr:`~struphy.pic.base.Particles.orbit_quantities`."""
default_background = maxwellians.Maxwellian3D()
"""Default kinetic background is a 3D Cartesian Maxwellian."""
default_n_cols = {"diagnostics": 0, "aux": 6}
"""Default number of buffer columns reserved for diagnostics and auxiliary (pusher/free) use (6 for the symmetric
3x3 particle-matrix fillings stored per marker by the matrix-free accumulation, see
:class:`~struphy.pic.accumulation.particles_to_grid.MarkerOperator`)."""
@property
def mu_idx(self):
return self.first_diagnostics_idx + 4
[docs]
def __post_init__(self):
"""If the background is a :class:`~struphy.kinetic_background.maxwellians.CanonicalMaxwellian`,
set up the discrete magnetic field (needed to evaluate canonical invariants) from the projected equilibrium."""
if isinstance(self.background, maxwellians.CanonicalMaxwellian2D):
assert isinstance(self.projected_equil, ProjectedFluidEquilibriumWithB), (
"CanonicalMaxwellian2D needs background with magnetic field."
)
self._absB0_h = self.projected_equil.absB0
self._b2_h = self.projected_equil.b2
self._derham = self.projected_equil.derham
@property
def sampling_density(self):
"""Sampling density function as volume form, used to draw markers via inverse transform/rejection
sampling and to compute their initial weights (see :meth:`~struphy.pic.base.Particles.draw_markers`).
This is a :class:`~struphy.kinetic_background.maxwellians.Maxwellian3D` in the Cartesian velocities
``vx, vy, vz``, parametrized by the mean velocities and thermal velocities in :attr:`loading_params`,
with density normalized to 1 (i.e. uniform in ``eta1, eta2, eta3``), further multiplied by the
Jacobian factor ``2 * eta1`` if :attr:`spatial` is ``"disc"`` (to sample uniformly in physical
space on a disc, where ``eta1`` plays the role of a normalized radius).
"""
if not hasattr(self, "_sampling_density"):
self._sampling_density = maxwellians.Maxwellian3D(
n=(1.0, None),
u1=(self.loading_params.moments[0], None),
u2=(self.loading_params.moments[1], None),
u3=(self.loading_params.moments[2], None),
vth1=(self.loading_params.moments[3], None),
vth2=(self.loading_params.moments[4], None),
vth3=(self.loading_params.moments[5], None),
uniform_on_disc=(self.spatial == "disc"),
)
return self._sampling_density
[docs]
def s0(self, eta1, eta2, eta3, vx, vy, vz, flat_eval=False, remove_holes=True):
"""Sampling density function as 0 form. This is the quantity stored in each
marker's ``s0`` column (see the class docstring) and used to compute initial weights
``w0 = f_init / s0 / Np``.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
vx, vy, vz : array_like
Cartesian velocity evaluation points.
flat_eval : bool
If true, perform flat (marker) evaluation (etas must be same size 1D).
remove_holes : bool
If True, holes are removed from the returned array. If False, holes are evaluated to -1.
Returns
-------
out : array-like
The 0-form sampling density.
-------
"""
assert self.domain, "self.domain must be set to call the sampling density 0-form."
return self.domain.transform(
self.sampling_density(eta1, eta2, eta3, vx, vy, vz),
eta1,
eta2,
eta3,
flat_eval=flat_eval,
kind="3_to_0",
remove_outside=remove_holes,
)
[docs]
def save_constants_of_motion(self):
"""
Calculate each marker's guiding-center constants of motion (only the equilibrium
magnetic field is considered) and assign them into the diagnostics columns of the marker array:
* ``0:3``: guiding-center position (logical :math:`\\boldsymbol \\eta`)
* ``3``: energy
* ``4``: magnetic moment
* ``5``: canonical toroidal momentum
* ``6``: parallel velocity
"""
assert isinstance(self.equil, FluidEquilibriumWithB), "Constants of motion need background with magnetic field."
# idx and slice
idx_gc_r = self.first_diagnostics_idx
slice_gc = slice(self.first_diagnostics_idx, self.first_diagnostics_idx + 3)
idx_energy = self.first_diagnostics_idx + 3
idx_can_momentum = self.first_diagnostics_idx + 5
# save cartesian positions
self.markers[slice_gc, ~self.holes] = self.domain(self.positions)
# eval guiding center phase space
eval_guiding_center_from_6d(
self.markers,
self._derham.args_derham,
self.domain.args_domain,
self.first_diagnostics_idx,
self.equation_params.epsilon,
self._b2_h[0]._data,
self._b2_h[1]._data,
self._b2_h[2]._data,
self._absB0_h._data,
)
# apply domain inverse map to get logical guiding center positions
# TODO: currently only possible with the geometry where its inverse map is defined.
assert hasattr(self.domain, "inverse_map")
self.markers[slice_gc, ~self.holes] = xp.stack(self.domain.inverse_map(*self.markers[slice_gc, ~self.holes]))
# eval energy
self.markers[idx_energy, ~self.holes] = (
self.markers[3, ~self.holes] ** 2 + self.markers[4, ~self.holes] ** 2 + self.markers[5, ~self.holes] ** 2
) / (2)
# eval psi at etas
a1 = self.equil.domain.params["a1"]
R0 = self.equil.params["R0"]
B0 = self.equil.params["B0"]
r = self.markers[idx_gc_r, ~self.holes] * (1 - a1) + a1
self.markers[idx_can_momentum, ~self.holes] = self.equil.psi_r(r)
# send particles to the guiding center positions
self.markers[self.first_pusher_idx : self.first_pusher_idx + 3, ~self.holes] = self.markers[
slice_gc, ~self.holes
]
if self.mpi_comm is not None:
self.mpi_sort_markers(alpha=1)
eval_canonical_toroidal_moment_6d(
self.markers,
self._derham.args_derham,
self.first_diagnostics_idx,
self.equation_params.epsilon,
B0,
R0,
self._absB0_h._data,
)
# send back and clear buffer
if self.mpi_comm is not None:
self.mpi_sort_markers()
self.markers[self.first_pusher_idx : self.first_pusher_idx + 3, ~self.holes] = 0
[docs]
class DeltaFParticles6D(Particles6D):
"""
A class for kinetic species in full 6D phase space that solve for delta_f = f - f0.
See :class:`~struphy.pic.particles.Particles6D` for more information.
"""
[docs]
def __post_init__(self):
"""Force the control-variate weight update off, since delta-f weights already evolve
the perturbation directly (there is no separate background contribution to subtract)."""
self.weights_params.control_variate = False
def _set_initial_condition(self):
"""Zero out the density of the (unperturbed) background before setting the initial
condition, so that only the perturbation :math:`\\delta f` is initialized on the markers."""
self.set_n_to_zero(self.initial_condition)
super()._set_initial_condition()
[docs]
def set_n_to_zero(self, background: Maxwellian | SumKineticBackground):
"""Recursively set the density moment ``n`` of ``background`` (and, if it is a
:class:`~struphy.kinetic_background.base.SumKineticBackground`, of both its summands) to zero,
keeping any perturbation attached to it.
Parameters
----------
background : Maxwellian | SumKineticBackground
The kinetic background whose density is to be zeroed.
"""
if isinstance(background, Maxwellian):
background.params["n"] = (0.0, background.params["n"][1])
else:
assert isinstance(background, SumKineticBackground)
self.set_n_to_zero(background._f1)
self.set_n_to_zero(background._f2)
[docs]
class PauliParticles6D(Particles6D):
r"""
Classical Pauli particles (Xiao and Qin, Comput. Phys. Commun. 265, 107981 (2021)): full-orbit (6D) markers
with an additional, constant magnetic moment :math:`\mu_p`, stored in the first diagnostics column
(:attr:`mu_idx`).
The markers are loaded on the slow manifold of the classical Pauli particle, whose dynamics is that of the
guiding center: after drawing the markers from the initial condition as for :class:`Particles6D`
(and computing their initial weights and sampling densities), each marker gets the magnetic moment of its
perpendicular velocity and its velocity is reduced to the parallel component,
.. math::
\mu_p = \frac{|\mathbf v_p \times \mathbf b_0(\mathbf x_p)|^2}{2 |\mathbf B_0(\mathbf x_p)|}\,,\qquad
\mathbf v_p \to (\mathbf v_p \cdot \mathbf b_0)\, \mathbf b_0(\mathbf x_p)\,,
see :meth:`reduce_to_slow_manifold`. The marker position is from then on the guiding-center position, and
the markers represent the gyro-average of the initial condition. The magnetic moment of the remaining
perpendicular velocity, :math:`\mu'_p` (:meth:`mu_prime`), stays close to zero as long as the dynamics remains
on the slow manifold.
With the control variate, the distribution functions (background and initial condition) are evaluated at the
gyrotropic velocity :math:`v_\parallel \mathbf b_0 + \sqrt{2 \mu_p |\mathbf B_0|}\, \mathbf e_\perp` with some unit
vector :math:`\mathbf e_\perp \perp \mathbf b_0` (see :attr:`f_coords`), i.e. the background is the guiding-center
distribution :math:`f_0(\mathbf x, v_\parallel, \mu)`. It must therefore be gyrotropic: Maxwellians with
isotropic thermal velocities and a mean velocity parallel to :math:`\mathbf B_0` (no diamagnetic drift).
The polarization drift on the slow manifold is compressible, so the phase-space density of the markers is not
constant along their orbits, and subtracting :math:`f_0` along the markers removes the polarization density from
the deposited charge: it must then be added in the field equation (as for guiding centers, see
:class:`~struphy.propagators.adiabatic_poisson_solve.AdiabaticPoissonSolve` with ``polarization=True``).
See :class:`~struphy.pic.base.Particles` for the structure of the numpy marker array and the meaning of its columns.
"""
orbit_quantities = Particles6D.orbit_quantities + ((9, "mu", "$\\mu$", "magnetic moment"),)
"""Marker columns saved as orbits: those of :class:`Particles6D` and the magnetic moment."""
default_n_cols = {"diagnostics": 1, "aux": 6}
"""One diagnostics column holds the magnetic moment."""
@property
def mu_idx(self):
"""Column of the magnetic moment (the first diagnostics column)."""
return self.first_diagnostics_idx
[docs]
def __post_init__(self):
"""Check that there is a magnetic equilibrium and, with the control variate, that the kinetic backgrounds
are gyrotropic Maxwellians (isotropic thermal velocities)."""
super().__post_init__()
assert isinstance(self.equil, FluidEquilibriumWithB), "Classical Pauli particles need an equilibrium with B."
if self.weights_params.control_variate:
for f in (self.background, self.initial_condition):
for maxw in _maxwellians_in(f):
assert isinstance(maxw, maxwellians.Maxwellian3D), (
f"Classical Pauli particles need Maxwellian3D backgrounds, not {type(maxw).__name__}."
)
assert maxw.params["vth1"] == maxw.params["vth2"] == maxw.params["vth3"], (
"Classical Pauli particles with the control variate need gyrotropic backgrounds: "
"set vth1 = vth2 = vth3."
)
# a mean velocity must be parallel to B0, which cannot be checked on the parameters
if any(maxw.params[u] != (0.0, None) for u in ("u1", "u2", "u3")):
logger.warning(
"PauliParticles6D with the control variate: the mean velocity of the background must be "
"parallel to B0 (gyrotropic background), otherwise the delta-f weights are wrong."
)
@property
def magnetic_moments(self):
"""Array of shape ``(n_mks_loc,)`` holding the magnetic moments of the markers."""
return self.markers[self.mu_idx, self.valid_mks]
@magnetic_moments.setter
def magnetic_moments(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (self.n_mks_loc,)
self._markers[self.mu_idx, self.valid_mks] = new
def _unit_b_and_absB(self, positions):
"""Cartesian unit vector and strength of the equilibrium magnetic field at the marker positions
(array of shape ``(3, n)``)."""
# no kernel call without markers (e.g. on a process without markers)
if positions.shape[1] == 0:
return xp.zeros((3, 0), dtype=float), xp.zeros(0, dtype=float)
# the evaluation kernels need C-ordered positions with standard strides (a fancy-indexed single marker
# has the strides (8, 24) and is still flagged C-contiguous)
b, _ = self.equil.b_cart(positions.copy(order="C"))
absB = xp.sqrt(b[0] ** 2 + b[1] ** 2 + b[2] ** 2)
return b / absB, absB
[docs]
def initialize_weights(self, **kwargs):
"""Compute the initial weights and sampling densities as for :class:`Particles6D` with the drawn
velocities, then reduce the markers to the slow manifold (:meth:`reduce_to_slow_manifold`).
Markers loaded from a file (``loading="restart"`` or ``"external"``) are already on the slow manifold
and carry their magnetic moments (column :attr:`mu_idx`), so they are not reduced again."""
self._on_slow_manifold = False
super().initialize_weights(**kwargs)
if self.loading not in ("restart", "external"):
self.reduce_to_slow_manifold()
self._on_slow_manifold = True
# the gyrotropic velocity of f_coords is the drawn one up to its direction, so the weights stay the same
if self.control_variate:
self.update_weights()
[docs]
def reduce_to_slow_manifold(self):
r"""Set the magnetic moments :math:`\mu_p = |\mathbf v_{\perp,p}|^2 / (2 |\mathbf B_0|)` and remove the
perpendicular velocities, :math:`\mathbf v_p \to (\mathbf v_p \cdot \mathbf b_0) \mathbf b_0`."""
if self.n_mks_loc > 0:
b, absB = self._unit_b_and_absB(self.positions)
v = self.velocities
v_par = xp.sum(v * b, axis=0)
v_perp = v - v_par * b
self.magnetic_moments = xp.sum(v_perp**2, axis=0) / (2.0 * absB)
self.velocities = v_par * b
[docs]
def mu_prime(self):
r"""Magnetic moment of the perpendicular marker velocities, :math:`\mu'_p = |\mathbf v_p \times \mathbf b_0|^2 / (2 |\mathbf B_0|)`
(Eq. (7) of Xiao and Qin), array of shape ``(n_mks_loc,)``. It is close to zero on the slow manifold."""
if self.n_mks_loc == 0:
return xp.zeros(0, dtype=float)
b, absB = self._unit_b_and_absB(self.positions)
v = self.velocities
v_perp = v - xp.sum(v * b, axis=0) * b
return xp.sum(v_perp**2, axis=0) / (2.0 * absB)
@property
def f_coords(self):
r"""Coordinates of the distribution function: the positions and, on the slow manifold, the gyrotropic
velocities :math:`v_\parallel \mathbf b_0 + \sqrt{2 \mu_p |\mathbf B_0|}\, \mathbf e_\perp`."""
coords = super().f_coords
if not getattr(self, "_on_slow_manifold", False) or coords.shape[1] == 0:
return coords
b, absB = self._unit_b_and_absB(coords[:3])
v_par = xp.sum(coords[3:6] * b, axis=0)
# e_perp = b x e_k / |b x e_k| with the Cartesian unit vector e_k least aligned with b
k = xp.argmin(xp.abs(b), axis=0)
e_k = xp.zeros_like(b)
e_k[k, xp.arange(b.shape[1])] = 1.0
e_perp = xp.cross(b, e_k, axis=0)
e_perp /= xp.sqrt(xp.sum(e_perp**2, axis=0))
v_perp = xp.sqrt(2.0 * self.markers[self.mu_idx, self.valid_mks] * absB)
coords[3:6] = v_par * b + v_perp * e_perp
return coords
@f_coords.setter
def f_coords(self, new):
Particles6D.f_coords.fset(self, new)
def _maxwellians_in(f):
"""The Maxwellians of a kinetic background (summands of a :class:`~struphy.kinetic_background.base.SumKineticBackground`)."""
if f is None:
return []
if isinstance(f, SumKineticBackground):
return _maxwellians_in(f._f1) + _maxwellians_in(f._f2)
return [f]
[docs]
class Particles5D(Particles):
"""
Particles in the 5D guiding-center, drift-kinetic or gyro-kinetic phase space
:math:`(\\boldsymbol \\eta, v_\\parallel, \\mu) \\in [0, 1]^3 \\times \\mathbb R \\times \\mathbb R_{\\geq 0}`.
Each marker carries a logical (curvilinear) position :math:`\\boldsymbol \\eta_p` together with the
velocity coordinates
.. math::
v_{\\parallel, p} = \\mathbf v_p \\cdot \\mathbf b_0(\\boldsymbol \\eta_p) \\,, \\qquad
\\mu_p = \\frac{1}{2 |\\mathbf B_0|} m |\\mathbf v_p|^2 - v_{\\parallel, p}^2 \\,,
defined with respect to the equilibrium magnetic field :math:`\\mathbf B_0` and its unit vector
:math:`\\mathbf b_0 = \\mathbf B_0 / |\\mathbf B_0|` (unlike :class:`Particles6D`, velocities are thus
not Cartesian but expressed in a field-aligned basis that itself depends on :math:`\\boldsymbol \\eta_p`).
By default, two diagnostics columns are reserved (``default_n_cols["diagnostics"] = 2``), holding
each marker's perpendicular energy and canonical toroidal momentum
(see :meth:`save_constants_of_motion`).
See :class:`~struphy.pic.base.Particles` for the structure of the numpy marker array and the meaning of its columns.
"""
# Class properties
vdim = 2
"""Dimension of the velocity space, here 2 (:math:`v_\\parallel, \\mu`)."""
coordinate_labels = ("$\\eta_1$", "$\\eta_2$", "$\\eta_3$", "$v_\\parallel$", "$\\mu$")
"""Labels for the coordinates in the phase space. Length is 5, with the first 3 being the spatial coordinates and the last 2 being the velocity coordinates."""
orbit_quantities = (
*ORBIT_POSITIONS,
(3, "v_par", "$v_\\parallel$", "parallel velocity"),
(4, "mu", "$\\mu$", "magnetic moment"),
(5, "weight", "$w$", "marker weight"),
(9, "p_phi", "$p_\\phi$", "canonical toroidal momentum (set by save_constants_of_motion)"),
)
"""Marker columns saved as orbits, see :attr:`~struphy.pic.base.Particles.orbit_quantities`."""
mu_idx = 4
"""Column index of particle magnetic moment."""
default_background = maxwellians.GyroMaxwellian2D()
"""Default kinetic background is a gyrotropic Maxwellian in :math:`(v_\\parallel, \\mu)`."""
default_n_cols = {"diagnostics": 2, "aux": 12}
"""Default number of buffer columns is 2 diagnostics (perpendicular energy, canonical toroidal
momentum, see :meth:`save_constants_of_motion`) and 12 auxiliary columns."""
[docs]
def __post_init__(self):
"""Retrieve the discrete equilibrium magnetic-field quantities (:math:`|B_0|`, unit 1-form
:math:`\\mathbf b_0`, Derham complex) needed to project marker velocities onto
:math:`v_\\parallel, \\mu` and to evaluate diagnostics, and allocate the temporary
FE coefficient vectors used for that."""
assert self.projected_equil is not None, "Particles5D needs a projected MHD equilibrium."
# magnetic background
if self.projected_equil is not None:
assert isinstance(self.projected_equil, ProjectedFluidEquilibriumWithB), (
"Particles5D needs background with magnetic field."
)
self._absB0_h = self.projected_equil.absB0
self._unit_b1_h = self.projected_equil.unit_b1
self._derham = self.projected_equil.derham
self._tmp0 = self.derham.V0.zeros()
self._tmp2 = self.derham.V2.zeros()
@property
def magn_bckgr(self):
"""Equilibrium fluid background carrying the magnetic field :math:`\\mathbf B_0` with respect to
which :math:`v_\\parallel, \\mu` are defined."""
return self.equil
@property
def absB0_h(self):
"""Discrete 0-form coefficients of :math:`|B_0|`."""
return self._absB0_h
@property
def unit_b1_h(self):
"""Discrete 1-form coefficients of the equilibrium field-aligned unit vector :math:`\\mathbf b_0 = \\mathbf B_0/|B_0|`."""
return self._unit_b1_h
@property
def epsilon(self):
"""Normalization parameter :math:`\\epsilon` (from :attr:`equation_params`) entering the
guiding-center equations of motion, e.g. the canonical toroidal momentum evaluation."""
return self.equation_params.epsilon
@property
def derham(self):
"""Discrete Derham complex of the projected equilibrium."""
return self._derham
@property
def sampling_density(self):
"""
Sampling density function as volume form, used to draw markers via inverse transform/rejection
sampling and to compute their initial weights (see :meth:`~struphy.pic.base.Particles.draw_markers`).
This is a :class:`~struphy.kinetic_background.maxwellians.GyroMaxwellian2D` in
:math:`(v_\\parallel, \\mu)`, parametrized by the mean/thermal parallel velocity
and by the equilibrium magnetic field in :attr:`loading_params`. It is normalized to
1 in logical space (i.e. uniform in ``eta1, eta2, eta3``) and already includes the polar-coordinate
Jacobian factor :math:`|\\mathbf B_0|` (``volume_form=True``), further multiplied by ``2 * eta1`` if
:attr:`spatial` is ``"disc"``.
"""
if not hasattr(self, "_sampling_density"):
self._sampling_density = maxwellians.GyroMaxwellian2D(
n=(1.0, None),
u_para=(self.loading_params.moments[0], None),
u_perp=(0.0, None),
vth_para=(self.loading_params.moments[2], None),
vth_perp=(self.loading_params.moments[3], None),
volume_form=True,
# equil=self.magn_bckgr,
B0=self.loading_params.B0,
uniform_on_disc=(self.spatial == "disc"),
)
return self._sampling_density
[docs]
def s3(self, eta1, eta2, eta3, v_para, mu):
"""
Sampling density function as 3-form, i.e. :meth:`sampling_density` with the velocity-space
(:math:`B_0`) Jacobian factor divided back out,
leaving a density that is a volume form in :math:`\\boldsymbol \\eta` only.
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 3-form sampling density.
-------
"""
return self.sampling_density(eta1, eta2, eta3, v_para, mu) / self.sampling_density.velocity_jacobian_det(
eta1, eta2, eta3, v_para, mu
)
[docs]
def s0(self, eta1, eta2, eta3, v_para, mu, flat_eval=False, remove_holes=True):
"""
Sampling density function as 0-form, i.e. :meth:`s3` pushed forward to a pointwise density by
dividing out the spatial metric Jacobian determinant. This is the quantity stored in each marker's
``s0`` column and used to compute initial weights ``w0 = f_init / s0 / Np``.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
v_para, mu : array_like
Parallel velocity and magnetic moment evaluation points.
flat_eval : bool
If true, perform flat (marker) evaluation (etas must be same size 1D).
remove_holes : bool
If True, holes are removed from the returned array. If False, holes are evaluated to -1.
Returns
-------
out : array-like
The 0-form sampling density.
-------
"""
return self.domain.transform(
self.s3(eta1, eta2, eta3, v_para, mu),
eta1,
eta2,
eta3,
flat_eval=flat_eval,
kind="3_to_0",
remove_outside=remove_holes,
)
[docs]
def save_constants_of_motion(self):
"""
Calculate each marker's guiding-center energy and canonical toroidal momentum (only the
equilibrium magnetic field is considered) and assign them into the diagnostics columns of
the marker array:
* ``first_diagnostics_idx + 0``: energy
* ``first_diagnostics_idx + 1``: canonical toroidal momentum
The magnetic moment itself is not a diagnostics column here (unlike in
:class:`Particles5Dvperp`) since it is already a phase-space coordinate, see :attr:`mu_idx`.
"""
assert isinstance(self.equil, FluidEquilibriumWithB), "Constants of motion need background with magnetic field."
# idx and slice
idx_can_momentum = self.first_diagnostics_idx + 1
eval_energy_5d(
self.markers,
self.derham.args_derham,
self.first_diagnostics_idx,
self.mu_idx,
self.absB0_h._data,
)
# eval psi at etas
a1 = self.equil.domain.params["a1"]
R0 = self.equil.params["R0"]
B0 = self.equil.params["B0"]
r = self.markers[0, ~self.holes] * (1 - a1) + a1
self.markers[idx_can_momentum, ~self.holes] = self.equil.psi_r(r)
eval_canonical_toroidal_moment_5d(
self.markers,
self.derham.args_derham,
self.first_diagnostics_idx,
self.mu_idx,
idx_can_momentum,
self.equation_params.epsilon,
B0,
R0,
self.absB0_h._data,
)
[docs]
def save_magnetic_energy(self, PBb):
r"""
Calculate the (time-dependent) magnetic field energy at each marker's position and assign it
into the energy diagnostics column (``self.first_diagnostics_idx``).
Parameters
----------
PBb : BlockVector
Finite element coefficients of the time-dependent magnetic field, projected onto V0.
"""
E0T = self.derham.extraction_ops["0"].transpose()
PBbt = E0T.dot(PBb, out=self._tmp0)
PBbt.update_ghost_regions()
eval_magnetic_energy_PBb(
self.markers,
self.derham.args_derham,
self.domain.args_domain,
self.first_diagnostics_idx,
self.mu_idx,
self.absB0_h._data,
PBbt._data,
)
[docs]
def save_magnetic_background_energy(self):
r"""
Evaluate the equilibrium magnetic-moment energy :math:`\mu_p |B_0(\boldsymbol \eta_p)|` for each marker.
The result is stored in the energy diagnostics column (``self.first_diagnostics_idx``).
"""
eval_magnetic_background_energy(
self.markers,
self.derham.args_derham,
self.domain.args_domain,
self.first_diagnostics_idx,
self.mu_idx,
self.absB0_h._data,
)
[docs]
class Particles5Dvperp(Particles):
"""
Particles in the 5D guiding-center, drift-kinetic or gyro-kinetic phase space
:math:`(\\boldsymbol \\eta, v_\\parallel, v_\\perp) \\in [0, 1]^3 \\times \\mathbb R \\times \\mathbb R_{\\geq 0}`.
Each marker carries a logical (curvilinear) position :math:`\\boldsymbol \\eta_p` together with the
parallel and perpendicular velocity coordinates
.. math::
v_{\\parallel, p} = \\mathbf v_p \\cdot \\mathbf b_0(\\boldsymbol \\eta_p) \\,, \\qquad
v_{\\perp, p} = \\left| \\mathbf v_p - v_{\\parallel, p} \\, \\mathbf b_0(\\boldsymbol \\eta_p) \\right| \\,,
defined with respect to the equilibrium magnetic field :math:`\\mathbf B_0` and its unit vector
:math:`\\mathbf b_0 = \\mathbf B_0 / |\\mathbf B_0|` (unlike :class:`Particles6D`, velocities are thus
not Cartesian but expressed in a field-aligned basis that itself depends on :math:`\\boldsymbol \\eta_p`).
By default, three diagnostics columns are reserved (``default_n_cols["diagnostics"] = 3``), holding
each marker's guiding-center energy, magnetic moment and canonical toroidal momentum
(see :meth:`save_constants_of_motion`).
See :class:`~struphy.pic.base.Particles` for the structure of the numpy marker array and the meaning of its columns.
"""
# Class properties
vdim = 2
"""Dimension of the velocity space, here 2 (:math:`v_\\parallel, v_\\perp`)."""
coordinate_labels = ("$\\eta_1$", "$\\eta_2$", "$\\eta_3$", "$v_\\parallel$", "$v_\\perp$")
"""Labels for the coordinates in the phase space. Length is 5, with the first 3 being the spatial coordinates and the last 2 being the velocity coordinates."""
orbit_quantities = (
*ORBIT_POSITIONS,
(3, "v_par", "$v_\\parallel$", "parallel velocity"),
(4, "v_perp", "$v_\\perp$", "perpendicular velocity"),
(5, "weight", "$w$", "marker weight"),
(10, "p_phi", "$p_\\phi$", "canonical toroidal momentum (set by save_constants_of_motion)"),
)
"""Marker columns saved as orbits, see :attr:`~struphy.pic.base.Particles.orbit_quantities`."""
default_background = maxwellians.GyroMaxwellian2Dvperp()
"""Default kinetic background is a gyrotropic Maxwellian in :math:`(v_\\parallel, v_\\perp)`."""
default_n_cols = {"diagnostics": 3, "aux": 12}
"""Default number of buffer columns is 3 diagnostics (energy, magnetic moment, canonical toroidal
momentum, see :meth:`save_constants_of_motion`) and 12 auxiliary columns."""
@property
def mu_idx(self):
return self.first_diagnostics_idx + 1
[docs]
def __post_init__(self):
"""Retrieve the discrete equilibrium magnetic-field quantities (:math:`|B_0|`, unit 1-form
:math:`\\mathbf b_0`, Derham complex) needed to project marker velocities onto
:math:`v_\\parallel, v_\\perp` and to evaluate diagnostics, and allocate the temporary
FE coefficient vectors used for that."""
assert self.projected_equil is not None, "Particles5Dvperp needs a projected MHD equilibrium."
# magnetic background
if self.projected_equil is not None:
assert isinstance(self.projected_equil, ProjectedFluidEquilibriumWithB), (
"Particles5Dvperp needs background with magnetic field."
)
self._absB0_h = self.projected_equil.absB0
self._unit_b1_h = self.projected_equil.unit_b1
self._derham = self.projected_equil.derham
self._tmp0 = self.derham.V0.zeros()
self._tmp2 = self.derham.V2.zeros()
@property
def magn_bckgr(self):
"""Equilibrium fluid background carrying the magnetic field :math:`\\mathbf B_0` with respect to
which :math:`v_\\parallel, v_\\perp` are defined."""
return self.equil
@property
def absB0_h(self):
"""Discrete 0-form coefficients of :math:`|B_0|`."""
return self._absB0_h
@property
def unit_b1_h(self):
"""Discrete 1-form coefficients of the equilibrium field-aligned unit vector :math:`\\mathbf b_0 = \\mathbf B_0/|B_0|`."""
return self._unit_b1_h
@property
def epsilon(self):
"""Normalization parameter :math:`\\epsilon` (from :attr:`equation_params`) entering the
guiding-center equations of motion, e.g. the canonical toroidal momentum evaluation."""
return self.equation_params.epsilon
@property
def derham(self):
"""Discrete Derham complex of the projected equilibrium."""
return self._derham
@property
def sampling_density(self):
"""
Sampling density function as volume form, used to draw markers via inverse transform/rejection
sampling and to compute their initial weights (see :meth:`~struphy.pic.base.Particles.draw_markers`).
This is a :class:`~struphy.kinetic_background.maxwellians.GyroMaxwellian2Dvperp` in
:math:`(v_\\parallel, v_\\perp)`, parametrized by the mean/thermal parallel and perpendicular velocities
in :attr:`loading_params`. It is normalized to
1 in logical space (i.e. uniform in ``eta1, eta2, eta3``) and already includes the polar-coordinate
Jacobian factor :math:`|v_\\perp|` (``volume_form=True``), further multiplied by ``2 * eta1`` if
:attr:`spatial` is ``"disc"``.
Returns
-------
out : array-like
The volume-form sampling density.
-------
"""
if not hasattr(self, "_sampling_density"):
self._sampling_density = maxwellians.GyroMaxwellian2Dvperp(
n=(1.0, None),
u_para=(self.loading_params.moments[0], None),
u_perp=(self.loading_params.moments[1], None),
vth_para=(self.loading_params.moments[2], None),
vth_perp=(self.loading_params.moments[3], None),
volume_form=True,
equil=self.magn_bckgr,
uniform_on_disc=(self.spatial == "disc"),
)
return self._sampling_density
[docs]
def s3(self, eta1, eta2, eta3, v_para, v_perp):
"""
Sampling density function as 3-form, i.e. :meth:`sampling_density` with the velocity-space
(:math:`|v_\\perp|`) Jacobian factor divided back out,
leaving a density that is a volume form in :math:`\\boldsymbol \\eta` only.
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 3-form sampling density.
-------
"""
return self.sampling_density(eta1, eta2, eta3, v_para, v_perp) / self.sampling_density.velocity_jacobian_det(
eta1, eta2, eta3, v_para, v_perp
)
[docs]
def s0(self, eta1, eta2, eta3, v_para, v_perp, flat_eval=False, remove_holes=True):
"""
Sampling density function as 0-form, i.e. :meth:`s3` pushed forward to a pointwise density by
dividing out the spatial metric Jacobian determinant. This is the quantity stored in each marker's
``s0`` column and used to compute initial weights ``w0 = f_init / s0 / Np``.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
v_para, v_perp : array_like
Parallel and perpendicular velocity evaluation points.
flat_eval : bool
If true, perform flat (marker) evaluation (etas must be same size 1D).
remove_holes : bool
If True, holes are removed from the returned array. If False, holes are evaluated to -1.
Returns
-------
out : array-like
The 0-form sampling density.
-------
"""
return self.domain.transform(
self.s3(eta1, eta2, eta3, v_para, v_perp),
eta1,
eta2,
eta3,
flat_eval=flat_eval,
kind="3_to_0",
remove_outside=remove_holes,
)
[docs]
def draw_markers(self, sort: bool = True):
super().draw_markers(sort=sort)
# magnetic moment is an adiabatic invariant: evaluate once at draw time (diagnostics column 1)
eval_magnetic_moment_5d(
self.markers,
self.derham.args_derham,
self.first_diagnostics_idx,
self._absB0_h._data,
)
[docs]
def save_constants_of_motion(self):
"""
Calculate each marker's guiding-center energy and canonical toroidal momentum (only the
equilibrium magnetic field is considered) and assign them into the diagnostics columns of
the marker array:
* ``first_diagnostics_idx + 0``: energy
* ``first_diagnostics_idx + 1``: magnetic moment (set once in :meth:`draw_markers`, unchanged here
since it is an adiabatic invariant)
* ``first_diagnostics_idx + 2``: canonical toroidal momentum
"""
assert isinstance(self.equil, FluidEquilibriumWithB), "Constants of motion need background with magnetic field."
# idx and slice
idx_can_momentum = self.first_diagnostics_idx + 2
eval_energy_5d(
self.markers,
self.derham.args_derham,
self.first_diagnostics_idx,
self.mu_idx,
self.absB0_h._data,
)
# eval psi at etas
a1 = self.equil.domain.params["a1"]
R0 = self.equil.params["R0"]
B0 = self.equil.params["B0"]
r = self.markers[0, ~self.holes] * (1 - a1) + a1
self.markers[idx_can_momentum, ~self.holes] = self.equil.psi_r(r)
eval_canonical_toroidal_moment_5d(
self.markers,
self.derham.args_derham,
self.first_diagnostics_idx,
self.mu_idx,
idx_can_momentum,
self.equation_params.epsilon,
B0,
R0,
self.absB0_h._data,
)
[docs]
def save_magnetic_energy(self, PBb):
r"""
Calculate the (time-dependent) magnetic field energy at each marker's position and assign it
into the energy diagnostics column (``self.first_diagnostics_idx``).
Parameters
----------
PBb : BlockVector
Finite element coefficients of the time-dependent magnetic field, projected onto V0.
"""
E0T = self.derham.extraction_ops["0"].transpose()
PBbt = E0T.dot(PBb, out=self._tmp0)
PBbt.update_ghost_regions()
eval_magnetic_energy_PBb(
self.markers,
self.derham.args_derham,
self.domain.args_domain,
self.first_diagnostics_idx,
self.mu_idx,
self.absB0_h._data,
PBbt._data,
)
[docs]
def save_magnetic_background_energy(self):
r"""
Evaluate the equilibrium magnetic-moment energy :math:`\mu_p |B_0(\boldsymbol \eta_p)|` for each marker.
The result is stored in the energy diagnostics column (``self.first_diagnostics_idx``).
"""
eval_magnetic_background_energy(
self.markers,
self.derham.args_derham,
self.domain.args_domain,
self.first_diagnostics_idx,
self.mu_idx,
self.absB0_h._data,
)
[docs]
def save_magnetic_moment(self):
r"""
Calculate the magnetic moment of each marker and assign it into the magnetic-moment
diagnostics column (``self.first_diagnostics_idx + 1``).
"""
eval_magnetic_moment_5d(
self.markers,
self.derham.args_derham,
self.first_diagnostics_idx,
self.absB0_h._data,
)
[docs]
class Particles3D(Particles):
"""
Particles in pure 3D configuration space :math:`\\boldsymbol \\eta \\in [0, 1]^3`, with no velocity
space attached (``vdim = 0``) — each marker only carries a logical (curvilinear) position, used e.g.
to represent a (massless) tracer or cold-plasma fluid density.
See :class:`~struphy.pic.base.Particles` for the structure of the numpy marker array and the meaning of its columns.
Parameters
----------
name : str
Name of particle species.
Np : int
Number of particles.
bc : list
Either 'remove', 'reflect', 'periodic' or 'refill' in each direction.
loading : str
Drawing of markers; either 'pseudo_random', 'sobol_standard',
'sobol_antithetic', 'external' or 'restart'.
**kwargs : dict
Parameters for markers, see :class:`~struphy.pic.base.Particles`.
"""
# Class properties
vdim = 0
"""Dimension of the velocity space, here 0 (no velocity coordinates)."""
coordinate_labels = ("$\\eta_1$", "$\\eta_2$", "$\\eta_3$")
"""Labels for the coordinates in the phase space. Length is 3, with all being spatial coordinates."""
orbit_quantities = (*ORBIT_POSITIONS, (3, "weight", "$w$", "marker weight"))
"""Marker columns saved as orbits, see :attr:`~struphy.pic.base.Particles.orbit_quantities`."""
default_background = maxwellians.ColdPlasma()
"""Default kinetic background is a cold-plasma (velocity-independent) density."""
default_n_cols = {"diagnostics": 0, "aux": 5}
"""Default number of buffer columns reserved for diagnostics and auxiliary (pusher/free) use."""
@property
def mu_idx(self):
return self.first_free_idx
[docs]
def __post_init__(self):
"""No additional setup is required for this class."""
@property
def sampling_density(self):
"""Sampling density function as volume form."""
if not hasattr(self, "_sampling_density"):
self._sampling_density = maxwellians.ColdPlasma(n=(1.0, None), uniform_on_disc=(self.spatial == "disc"))
return self._sampling_density
[docs]
def s0(self, eta1, eta2, eta3, flat_eval=False, remove_holes=True):
"""Sampling density function as 0 form, i.e. :meth:`sampling_density` pushed forward to a pointwise
(non-volume-form) density by dividing out the metric Jacobian determinant.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
flat_eval : bool
If true, perform flat (marker) evaluation (etas must be same size 1D).
remove_holes : bool
If True, holes are removed from the returned array. If False, holes are evaluated to -1.
Returns
-------
out : array-like
The 0-form sampling density.
-------
"""
return self.domain.transform(
self.sampling_density(eta1, eta2, eta3),
eta1,
eta2,
eta3,
flat_eval=flat_eval,
kind="3_to_0",
remove_outside=remove_holes,
)
[docs]
class ParticlesSPH(Particles):
"""
Particles for Smoothed Particle Hydrodynamics (SPH) models. The particle distribution itself lives
in pure 3D configuration space :math:`\\boldsymbol \\eta \\in [0, 1]^3`, exactly as for :class:`Particles3D`
(:meth:`sampling_density` and :meth:`s0` depend only on :math:`\\boldsymbol \\eta_p`).
Each marker additionally carries a Cartesian velocity :math:`\\mathbf v_p` in its marker-array columns,
but this is a per-particle *helper* quantity (e.g. the SPH velocity-field sample used by pushers and
kernel-based reconstructions) rather than a coordinate of a sampled phase-space density.
See :class:`~struphy.pic.base.Particles` for the structure of the numpy marker array and the meaning of its columns.
Parameters
----------
name : str
Name of the particle species.
**params : dict
Parameters for markers, see :class:`~struphy.pic.base.Particles`.
"""
# Class properties
vdim = 3
"""Dimension of the per-marker Cartesian velocity attribute, here 3 (not a sampled coordinate, see class docstring)."""
coordinate_labels = ("$\\eta_1$", "$\\eta_2$", "$\\eta_3$", "$v_x$", "$v_y$", "$v_z$")
"""Labels for the coordinates in the phase space. Length is 6, with the first 3 being the spatial coordinates and the last 3 being the velocity coordinates."""
orbit_quantities = (
*ORBIT_POSITIONS,
(3, "v1", "$v_x$", "Cartesian velocity x"),
(4, "v2", "$v_y$", "Cartesian velocity y"),
(5, "v3", "$v_z$", "Cartesian velocity z"),
(6, "weight", "$w$", "marker weight"),
)
"""Marker columns saved as orbits, see :attr:`~struphy.pic.base.Particles.orbit_quantities`."""
default_background = equils.ConstantVelocity()
"""Default fluid background is a spatially constant velocity field."""
default_n_cols = {"diagnostics": 0, "aux": 24}
"""Default number of buffer columns reserved for diagnostics and auxiliary (pusher/free) use."""
@property
def mu_idx(self):
return self.first_free_idx
[docs]
def __post_init__(self):
"""Attach the domain to the background (needed to evaluate it at marker positions).
SPH does not support clone-based (tile-copied) MPI parallelization."""
assert self.clone_config is None, "SPH can only be launched with --nclones 1"
self.background.domain = self.domain
@property
def sampling_density(self):
"""Sampling density function as volume form, used to draw markers via inverse transform/rejection
sampling and to compute their initial weights (see :meth:`~struphy.pic.base.Particles.draw_markers`).
This density is purely spatial: uniform (normalized to 1) if :attr:`spatial` is ``"uniform"``, or
multiplied by the Jacobian factor ``2 * eta1`` if :attr:`spatial` is ``"disc"``.
"""
if not hasattr(self, "_sampling_density"):
def func(eta1, eta2, eta3, *v):
if self.spatial == "uniform":
return 0 * eta1 + 1.0
elif self.spatial == "disc":
return 2 * eta1
else:
raise NotImplementedError(f'Spatial drawing must be "uniform" or "disc", is {self.spatial}.')
self._sampling_density = func
return self._sampling_density
[docs]
def s0(self, eta1, eta2, eta3, *v, flat_eval=False, remove_holes=True):
"""Sampling density function as 0 form, i.e. :meth:`sampling_density` pushed forward to a pointwise
(non-volume-form) density by dividing out the metric Jacobian determinant.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
*v : array_like
Accepted for a call signature compatible with generic phase-space evaluation, but unused
(see :meth:`sampling_density`).
flat_eval : bool
If true, perform flat (marker) evaluation (etas must be same size 1D).
remove_holes : bool
If True, holes are removed from the returned array. If False, holes are evaluated to -1.
Returns
-------
out : array-like
The 0-form sampling density.
-------
"""
return self.domain.transform(
self.sampling_density(eta1, eta2, eta3, *v),
eta1,
eta2,
eta3,
flat_eval=flat_eval,
kind="3_to_0",
remove_outside=remove_holes,
)
