"Only particle variables are updated."
import logging
from dataclasses import dataclass, replace
from typing import Literal
from line_profiler import profile
from struphy.io.options import LiteralOptions, OptionsBase
from struphy.models.variables import FEECVariable, PICVariable
from struphy.ode.utils import ButcherTableau
from struphy.pic.pushing.kernel_setup import KernelSetup
from struphy.pic.pushing.kernels.bstar_2form import bstar_2form
from struphy.pic.pushing.kernels.bstar_parallel_3form import bstar_parallel_3form
from struphy.pic.pushing.kernels.driftkinetic_hamiltonian import driftkinetic_hamiltonian
from struphy.pic.pushing.kernels.grad_driftkinetic_hamiltonian import grad_driftkinetic_hamiltonian
from struphy.pic.pushing.kernels.push_gc_Bstar_discrete_gradient_1st_order import (
push_gc_Bstar_discrete_gradient_1st_order,
)
from struphy.pic.pushing.kernels.push_gc_Bstar_discrete_gradient_1st_order_newton import (
push_gc_Bstar_discrete_gradient_1st_order_newton,
)
from struphy.pic.pushing.kernels.push_gc_Bstar_discrete_gradient_2nd_order import (
push_gc_Bstar_discrete_gradient_2nd_order,
)
from struphy.pic.pushing.kernels.push_gc_Bstar_explicit_multistage import push_gc_Bstar_explicit_multistage
from struphy.pic.pushing.pusher import Pusher
from struphy.propagators.base import Propagator
from struphy.utils.utils import check_option
logger = logging.getLogger("struphy")
[docs]
class PushGuidingCenterParallel(Propagator):
r"""For each marker :math:`p`, solves
.. math::
\left\{
\begin{aligned}
\frac{\textnormal d \mathbf X_p(t)}{\textnormal d t} &= v_{\parallel,p}(t) \frac{\mathbf B^*}{B^*_\parallel}(\mathbf X_p(t)) \,,
\\
\frac{\textnormal d v_{\parallel,p}(t)}{\textnormal d t} &= \frac{1}{\varepsilon} \frac{\mathbf B^*}{B^*_\parallel} \cdot \mathbf E^* (\mathbf X_p(t)) \,,
\end{aligned}
\right.
where
.. math::
\mathbf E^* = -\nabla \phi - \varepsilon \mu_p \nabla |\mathbf B|\,,\qquad \mathbf B^* = \mathbf B + \varepsilon v_\parallel \nabla \times \mathbf b_0\,,\qquad B^*_\parallel = \mathbf B^* \cdot \mathbf b_0\,,
where :math:`\mathbf B = \mathbf B_0 + \tilde{\mathbf B}` can be the full magnetic field (equilibrium + perturbation).
The electric potential ``phi`` and/or the magnetic perturbation ``b_tilde``
can be ignored by passing ``None``. Instead of the potential, an electric field :math:`\mathbf E \in H(\textnormal{curl})` can be passed (``e_field``, explicit algorithm); then :math:`\mathbf E^* = \mathbf E - \varepsilon \mu_p \nabla |\mathbf B|` (gauge-invariant electromagnetic formulation).
In logical space this is given by :math:`\mathbf X = F(\boldsymbol \eta)`:
.. math::
\left\{
\begin{aligned}
\frac{\textnormal d \boldsymbol \eta_p(t)}{\textnormal d t} &= v_{\parallel,p}(t) \frac{\hat{\mathbf B}^{*2}}{\hat B^{*3}_\parallel}(\boldsymbol \eta_p(t)) \,,
\\
\frac{\textnormal d v_{\parallel,p}(t)}{\textnormal d t} &= \frac{1}{\varepsilon} \frac{\hat{\mathbf B}^{*2}}{\hat B^{*3}_\parallel} \cdot \hat{\mathbf E}^{*1} (\boldsymbol \eta_p(t)) \,.
\end{aligned}
\right.
Available algorithms:
* Explicit from :class:`~struphy.ode.utils.ButcherTableau`
* :func:`~struphy.pic.pushing.kernels.push_gc_Bstar_discrete_gradient_1st_order.push_gc_Bstar_discrete_gradient_1st_order_kernels.push_gc_Bstar_discrete_gradient_1st_order`
* :func:`~struphy.pic.pushing.kernels.push_gc_Bstar_discrete_gradient_1st_order_newton.push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.push_gc_Bstar_discrete_gradient_1st_order_newton`
* :func:`~struphy.pic.pushing.kernels.push_gc_Bstar_discrete_gradient_2nd_order.push_gc_Bstar_discrete_gradient_2nd_order_kernels.push_gc_Bstar_discrete_gradient_2nd_order`
"""
[docs]
class Variables:
"""Container for variables advanced by :class:`PushGuidingCenterParallel`.
Attributes
----------
ions : PICVariable
Guiding-center particle variable in ``"Particles5D"`` space.
"""
def __init__(self):
self._ions: PICVariable = None
@property
def ions(self) -> PICVariable:
return self._ions
@ions.setter
def ions(self, new):
assert isinstance(new, PICVariable)
assert "Particles5D" in new.space
self._ions = new
def __init__(
self,
phi: FEECVariable = None,
b_tilde: FEECVariable = None,
e_field: FEECVariable = None,
):
r"""
Parameters
----------
phi : FEECVariable, default=None
Electric potential in ``"H1"`` space contributing to :math:`E^*`.
If ``None``, an empty ``FEECVariable(space="H1")`` is created internally.
b_tilde : FEECVariable, default=None
Magnetic perturbation in ``"Hdiv"`` space contributing to :math:`B^*`.
If ``None``, the perturbation is ignored.
e_field : FEECVariable, default=None
Electric field in ``"Hcurl"`` space, used in :math:`E^* = E - \varepsilon \mu \nabla |B|` instead of
:math:`-\nabla \phi` (gauge-invariant electromagnetic formulation). Cannot be combined with ``phi``. Only the
explicit algorithm is consistent with a general :math:`E` (the discrete gradients conserve a Hamiltonian
with a potential); a discrete-gradient ``algo`` is replaced by ``"explicit"``, with a warning.
"""
assert phi is None or e_field is None, "Pass either the potential phi or the electric field e_field, not both."
if e_field is not None:
assert isinstance(e_field, FEECVariable) and e_field.space == "Hcurl", "e_field must be in Hcurl."
self.e_field = e_field
self.variables = self.Variables()
self.phi = phi if phi is not None else FEECVariable(space="H1")
if phi is None:
self._allocate_phi = True
else:
self._allocate_phi = False
self.b_tilde = b_tilde
[docs]
@dataclass(repr=False)
class Options(OptionsBase):
"""Configuration options for :class:`PushGuidingCenterParallel`.
Parameters
----------
evaluate_e_field : bool, default=False
If ``True``, evaluate and include electric-field contributions in
drift-kinetic kernels.
algo : {"discrete_gradient_2nd_order", "discrete_gradient_1st_order", "discrete_gradient_1st_order_newton", "explicit"}, default="discrete_gradient_1st_order"
Guiding-center pushing algorithm.
butcher : ButcherTableau, default=None
Butcher tableau used in explicit mode.
If ``None`` and ``algo="explicit"``, defaults to
``ButcherTableau()``.
maxiter : int, default=20
Maximum number of fixed-point or Newton iterations in
discrete-gradient modes.
tol : float, default=1e-7
Convergence tolerance for iterative discrete-gradient updates.
mpi_sort : LiteralOptions.OptsMPIsort, default="each"
MPI sorting policy for particle exchange.
"""
# specific literals
OptsAlgo = Literal[
"discrete_gradient_2nd_order",
"discrete_gradient_1st_order",
"discrete_gradient_1st_order_newton",
"explicit",
]
# propagator options
evaluate_e_field: bool = False
algo: OptsAlgo = "discrete_gradient_1st_order"
butcher: ButcherTableau = None
maxiter: int = 20
tol: float = 1e-7
mpi_sort: LiteralOptions.OptsMPIsort = "each"
def __post_init__(self):
# checks
check_option(self.algo, self.OptsAlgo)
check_option(self.mpi_sort, LiteralOptions.OptsMPIsort)
# defaults
if self.algo == "explicit" and self.butcher is None:
self.butcher = ButcherTableau()
@property
def options(self) -> Options:
if not hasattr(self, "_options"):
self._options = self.Options()
return self._options
@options.setter
def options(self, new):
assert isinstance(new, self.Options)
self._options = new
logger.info(f"\nNew options for propagator '{self.__class__.__name__}':\n{self._options}")
@profile
def allocate(self):
# scaling factor
self._epsilon = self.variables.ions.species.equation_params.epsilon
# magnetic equilibrium field
self._gradB1 = self.projected_equil.gradB1
b2 = self.projected_equil.b2
curl_unit_b2 = self.projected_equil.curl_unit_b2
self._absB0 = self.projected_equil.absB0
curl_unit_b_dot_b0 = self.projected_equil.curl_unit_b_dot_b0
# the kernels need tensor-product coefficients: transposed extraction operators
# PolarVector --> Stencil-/BlockVector (identity maps in case of no polar splines)
self._E0T = self.derham.extraction_ops["0"].transpose()
self._E1T = self.derham.extraction_ops["1"].transpose()
self._E2T = self.derham.extraction_ops["2"].transpose()
# magnetic perturbation (in B* = B0 + b_tilde + epsilon v_par curl b0 and in B*_parallel and |B|)
if self.b_tilde is not None:
self._B_dot_b = self.derham.V0.zeros()
self._grad_b_full = self.derham.V1.zeros()
self._b2_eq = b2
self._b2_full = b2.space.zeros()
b2 = self._b2_full
self._PB = getattr(self.basis_ops, "PB")
self._B_dot_b_pol = self._PB.codomain.zeros()
self._grad_b_pol = self.derham.grad.codomain.zeros()
self._update_b_tilde()
else:
self._grad_b_full = self._gradB1
self._B_dot_b = self._absB0
# allocate electric potential (self._phi holds its tensor-product coefficients)
if self._allocate_phi:
self.phi.allocate(self.derham, domain=self.domain)
self._phi = self.derham.V0.zeros()
self._update_phi()
self._evaluate_e_field = self.options.evaluate_e_field
if self.e_field is not None:
self._evaluate_e_field = True
# the discrete gradients conserve the Hamiltonian, which contains no potential for a general E:
# they would remove the work of E along the step (inconsistent scheme), hence the explicit algorithm
if "discrete_gradient" in self.options.algo:
logger.warning(
f'{self.__class__.__name__} with e_field: algo="{self.options.algo}" needs a potential, '
'using algo="explicit" (default Butcher tableau).'
)
self._options = replace(self.options, algo="explicit", butcher=None)
self._e_field = self.derham.V1.zeros()
self._e_field_pol = self.derham.grad.codomain.zeros()
# choose method
particles = self.variables.ions.particles
if "discrete_gradient" in self.options.algo:
# place for storing data during iteration
first_free_idx = particles.args_markers.first_free_idx
if "1st_order" in self.options.algo:
# init kernels
self.add_init_kernel(
KernelSetup(
kernel=driftkinetic_hamiltonian,
output_indices=(first_free_idx,),
args=(
self.derham.args_derham,
self._epsilon,
self._B_dot_b._data,
self._phi._data,
self._evaluate_e_field,
),
),
)
self.add_init_kernel(
KernelSetup(
kernel=bstar_parallel_3form,
output_indices=(first_free_idx + 1,),
args=(
self.derham.args_derham,
self._epsilon,
self._B_dot_b._data,
curl_unit_b_dot_b0._data,
),
),
)
self.add_init_kernel(
KernelSetup(
kernel=bstar_2form,
output_indices=(first_free_idx + 2, first_free_idx + 3, first_free_idx + 4),
args=(
self.derham.args_derham,
self._epsilon,
b2[0]._data,
b2[1]._data,
b2[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
),
),
)
if "newton" in self.options.algo:
# eval kernels
self.add_eval_kernel(
KernelSetup(
kernel=driftkinetic_hamiltonian,
output_indices=(first_free_idx + 5,),
args=(
self.derham.args_derham,
self._epsilon,
self._B_dot_b._data,
self._phi._data,
self._evaluate_e_field,
),
alpha=(1.0, 0.0, 0.0, 0.0),
),
)
self.add_eval_kernel(
KernelSetup(
kernel=driftkinetic_hamiltonian,
output_indices=(first_free_idx + 6,),
args=(
self.derham.args_derham,
self._epsilon,
self._B_dot_b._data,
self._phi._data,
self._evaluate_e_field,
),
alpha=(1.0, 1.0, 0.0, 0.0),
),
)
self.add_eval_kernel(
KernelSetup(
kernel=grad_driftkinetic_hamiltonian,
output_indices=(first_free_idx + 7, None, None),
args=(
self.derham.args_derham,
self._epsilon,
self._grad_b_full[0]._data,
self._grad_b_full[1]._data,
self._grad_b_full[2]._data,
self._e_field[0]._data,
self._e_field[1]._data,
self._e_field[2]._data,
self._evaluate_e_field,
),
alpha=(1.0, 0.0, 0.0, 0.0),
),
)
self.add_eval_kernel(
KernelSetup(
kernel=grad_driftkinetic_hamiltonian,
output_indices=(first_free_idx + 8, first_free_idx + 9, None),
args=(
self.derham.args_derham,
self._epsilon,
self._grad_b_full[0]._data,
self._grad_b_full[1]._data,
self._grad_b_full[2]._data,
self._e_field[0]._data,
self._e_field[1]._data,
self._e_field[2]._data,
self._evaluate_e_field,
),
alpha=(1.0, 1.0, 0.0, 0.0),
),
)
# pusher kernel
kernel = push_gc_Bstar_discrete_gradient_1st_order_newton
alpha_in_kernel = 1.0 # evaluate at eta^{n+1,k} and save
args_kernel = (
self.derham.args_derham,
self._epsilon,
self._grad_b_full[0]._data,
self._grad_b_full[1]._data,
self._grad_b_full[2]._data,
self._B_dot_b._data,
self._e_field[0]._data,
self._e_field[1]._data,
self._e_field[2]._data,
self._phi._data,
self._evaluate_e_field,
)
else:
# eval kernels
self.add_eval_kernel(
KernelSetup(
kernel=driftkinetic_hamiltonian,
output_indices=(first_free_idx + 5,),
args=(
self.derham.args_derham,
self._epsilon,
self._B_dot_b._data,
self._phi._data,
self._evaluate_e_field,
),
alpha=1.0,
),
) # evaluate at Z^{n+1,k} and save
# pusher kernel
kernel = push_gc_Bstar_discrete_gradient_1st_order
alpha_in_kernel = 0.5 # evaluate at mid-point
args_kernel = (
self.derham.args_derham,
self._epsilon,
self._grad_b_full[0]._data,
self._grad_b_full[1]._data,
self._grad_b_full[2]._data,
self._e_field[0]._data,
self._e_field[1]._data,
self._e_field[2]._data,
self._evaluate_e_field,
)
elif "2nd_order" in self.options.algo:
# init kernels (evaluate at eta^n and save)
self.add_init_kernel(
KernelSetup(
kernel=driftkinetic_hamiltonian,
output_indices=(first_free_idx,),
args=(
self.derham.args_derham,
self._epsilon,
self._B_dot_b._data,
self._phi._data,
self._evaluate_e_field,
),
),
)
# eval kernels
self.add_eval_kernel(
KernelSetup(
kernel=driftkinetic_hamiltonian,
output_indices=(first_free_idx + 1,),
args=(
self.derham.args_derham,
self._epsilon,
self._B_dot_b._data,
self._phi._data,
self._evaluate_e_field,
),
alpha=1.0,
),
) # evaluate at Z^{n+1,k} and save
# pusher kernel
kernel = push_gc_Bstar_discrete_gradient_2nd_order
alpha_in_kernel = 0.5 # evaluate at mid-point
args_kernel = (
self.derham.args_derham,
self._epsilon,
self._grad_b_full[0]._data,
self._grad_b_full[1]._data,
self._grad_b_full[2]._data,
b2[0]._data,
b2[1]._data,
b2[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
self._B_dot_b._data,
curl_unit_b_dot_b0._data,
self._e_field[0]._data,
self._e_field[1]._data,
self._e_field[2]._data,
self._evaluate_e_field,
)
# Pusher instance
self._pusher = Pusher(
particles,
kernel,
args_kernel,
self.domain.args_domain,
alpha_in_kernel=alpha_in_kernel,
init_kernels=self.init_kernels,
eval_kernels=self.eval_kernels,
maxiter=self.options.maxiter,
tol=self.options.tol,
mpi_sort=self.options.mpi_sort,
pushes_eta=True,
)
else:
if self.options.butcher is None:
butcher = ButcherTableau()
else:
butcher = self.options.butcher
# temp fix due to refactoring of ButcherTableau:
kernel = push_gc_Bstar_explicit_multistage
args_kernel = (
self.derham.args_derham,
self._epsilon,
self._grad_b_full[0]._data,
self._grad_b_full[1]._data,
self._grad_b_full[2]._data,
b2[0]._data,
b2[1]._data,
b2[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
self._B_dot_b._data,
curl_unit_b_dot_b0._data,
self._e_field[0]._data,
self._e_field[1]._data,
self._e_field[2]._data,
self._evaluate_e_field,
butcher.a_stage,
butcher.b,
butcher.c,
)
self._pusher = Pusher(
particles,
kernel,
args_kernel,
self.domain.args_domain,
alpha_in_kernel=1.0,
n_stages=butcher.n_stages,
mpi_sort=self.options.mpi_sort,
pushes_eta=True,
)
def _update_phi(self):
"""Extract the electric potential to the tensor-product space."""
phi = self._E0T.dot(self.phi.spline.vector, out=self._phi)
phi.update_ghost_regions()
def _update_b_tilde(self):
"""Total magnetic field (2-form), parallel magnetic field and its gradient (equilibrium plus perturbation) in the
tensor-product spaces."""
B_dot_b_pol = self._PB.dot(self.b_tilde.spline.vector, out=self._B_dot_b_pol)
B_dot_b = self._E0T.dot(B_dot_b_pol, out=self._B_dot_b)
B_dot_b.update_ghost_regions()
grad_b_pol = self.derham.grad.dot(B_dot_b_pol, out=self._grad_b_pol)
grad_b_full = self._E1T.dot(grad_b_pol, out=self._grad_b_full)
grad_b_full.update_ghost_regions()
grad_b_full += self._gradB1
B_dot_b += self._absB0
# total magnetic field B0 + b_tilde (2-form) in B*
b2_full = self._E2T.dot(self.b_tilde.spline.vector, out=self._b2_full)
b2_full += self._b2_eq
b2_full.update_ghost_regions()
@profile
def __call__(self, dt):
# electric potential
self._update_phi()
# electric field
# TODO: add out to __neg__ of StencilVector
if self.e_field is not None:
e_field = self._E1T.dot(self.e_field.spline.vector, out=self._e_field)
e_field.update_ghost_regions()
elif self._evaluate_e_field:
e_field_pol = self.derham.grad.dot(-self.phi.spline.vector, out=self._e_field_pol)
e_field = self._E1T.dot(e_field_pol, out=self._e_field)
e_field.update_ghost_regions()
# magnetic perturbation
if self.b_tilde is not None:
self._update_b_tilde()
# call pusher
self._pusher(dt)
# update_weights
if self.variables.ions.species.weights_params.control_variate:
self.variables.ions.particles.update_weights()
