"Particle and FEEC variables are updated."
import logging
from dataclasses import dataclass
from typing import Literal
import cunumpy as xp
from feectools.linalg.solvers import inverse
from line_profiler import profile
from maybempi import MPI
from struphy.feec import preconditioner
from struphy.io.options import LiteralOptions, OptionsBase
from struphy.linear_algebra.solver import DiscreteGradientSolverParameters, SolverParameters
from struphy.models.variables import FEECVariable, PICVariable
from struphy.ode.utils import ButcherTableau
from struphy.pic.accumulation.filter import FilterParameters
from struphy.pic.accumulation.kernels.cc_lin_mhd_5d_gradB import cc_lin_mhd_5d_gradB
from struphy.pic.accumulation.kernels.cc_lin_mhd_5d_gradB_dg import cc_lin_mhd_5d_gradB_dg
from struphy.pic.accumulation.kernels.cc_lin_mhd_5d_gradB_dg_init import cc_lin_mhd_5d_gradB_dg_init
from struphy.pic.accumulation.particles_to_grid import AccumulatorVector
from struphy.pic.diagnostics.kernels.eval_gradB_ediff import eval_gradB_ediff
from struphy.pic.pushing.kernels.push_gc_cc_J2_dg_Hdiv import push_gc_cc_J2_dg_Hdiv
from struphy.pic.pushing.kernels.push_gc_cc_J2_dg_init_Hdiv import push_gc_cc_J2_dg_init_Hdiv
from struphy.pic.pushing.kernels.push_gc_cc_J2_stage_H1vec import push_gc_cc_J2_stage_H1vec
from struphy.pic.pushing.kernels.push_gc_cc_J2_stage_Hdiv import push_gc_cc_J2_stage_Hdiv
from struphy.propagators.base import Propagator
from struphy.utils.utils import check_option
logger = logging.getLogger("struphy")
[docs]
class CurrentCoupling5DGradB(Propagator):
r""":ref:`FEEC <gempic>` discretization of the following equations:
find :math:`\mathbf U \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}` and :math:`\mathbf B \in H(\textnormal{div})` such that
.. math::
\left\{
\begin{aligned}
\int \rho_0 &\frac{\partial \tilde{\mathbf U}}{\partial t} \cdot \mathbf V \, \textnormal{d} \mathbf x = - \frac{A_\textnormal{h}}{A_b} \iint \mu \frac{f^\text{vol}}{B^*_\parallel} (\mathbf b_0 \times \nabla B_\parallel) \times \mathbf B \cdot \mathbf V \,\textnormal{d} \mathbf x \textnormal{d} v_\parallel \textnormal{d} \mu \quad \forall \,\mathbf V \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}\,,
\\
&\frac{\partial f}{\partial t} = \frac{1}{B^*_\parallel} \left[ \mathbf b_0 \times (\tilde{\mathbf U} \times \mathbf B) \right] \cdot \nabla f\,.
\end{aligned}
\right.
:ref:`time_discret`: Explicit ('rk4', 'forward_euler', 'heun2', 'rk2', 'heun3').
.. math::
\begin{bmatrix}
\dot{\mathbf u}\\ \dot{\mathbf H}
\end{bmatrix}
=
\begin{bmatrix}
0 & (\mathbb{M}^{2,n})^{-1} \mathbb{L}² \frac{1}{\bar{\sqrt{g}}} \frac{1}{\bar B^{*0}_\parallel}\bar{B}^\times_f \bar{G}^{-1} \bar{b}^\times_0 \bar{G}^{-1}
\\
-\bar{G}^{-1} \bar{b}^\times_0 \bar{G}^{-1} \bar{B}^\times_f \frac{1}{\bar B^{*0}_\parallel} \frac{1}{\bar{\sqrt{g}}} (\mathbb{L}²)^\top (\mathbb{M}^{2,n})^{-1} & 0
\end{bmatrix}
\begin{bmatrix}
\mathbb M^{2,n} \mathbf u
\\
\frac{A_\textnormal{h}}{A_b} \bar M \bar W \overline{\nabla B}_\parallel
\end{bmatrix} \,.
For the detail explanation of the notations, see `2022_DriftKineticCurrentCoupling <https://gitlab.mpcdf.mpg.de/struphy/struphy-projects/-/blob/main/running-projects/2022_DriftKineticCurrentCoupling.md?ref_type=heads>`_.
"""
[docs]
class Variables:
"""Container for variables advanced by :class:`CurrentCoupling5DGradB`.
Attributes
----------
u : FEECVariable
Fluid-like FEEC variable in one of ``"Hcurl"``, ``"Hdiv"``, or
``"H1vec"``.
energetic_ions : PICVariable
Energetic-ion particle variable in ``"Particles5D"`` space.
"""
def __init__(self):
self._u: FEECVariable = None
self._energetic_ions: PICVariable = None
@property
def u(self) -> FEECVariable:
return self._u
@u.setter
def u(self, new):
assert isinstance(new, FEECVariable)
assert new.space in ("Hcurl", "Hdiv", "H1vec")
self._u = new
@property
def energetic_ions(self) -> PICVariable:
return self._energetic_ions
@energetic_ions.setter
def energetic_ions(self, new):
assert isinstance(new, PICVariable)
assert "Particles5D" in new.space
self._energetic_ions = new
def __init__(self, b_tilde: FEECVariable = None):
"""
Parameters
----------
b_tilde : FEECVariable, default=None
Magnetic perturbation 1-form (``"Hcurl"`` space) entering the grad-B
coupling term. If ``None``, only the equilibrium field is used.
"""
self.variables = self.Variables()
self.b_tilde = b_tilde
[docs]
@dataclass(repr=False)
class Options(OptionsBase):
r"""Configuration options for :class:`CurrentCoupling5DGradB`.
Parameters
----------
ep_scale : float, default=None
Factor of the energetic-particle contributions (accumulation): the mass ratio
:math:`A_\textnormal{h}/A_\textnormal{b}` of the energetic ions and the MHD bulk if ``None``.
algo : {"discrete_gradient", "explicit"}, default="explicit"
Time integration algorithm.
- ``"explicit"``: explicit Runge-Kutta update with ``butcher``.
- ``"discrete_gradient"``: nonlinear discrete-gradient update.
butcher : ButcherTableau, default=None
Butcher tableau used in explicit mode. If ``None`` and
``algo="explicit"``, defaults to ``ButcherTableau()``.
u_space : LiteralOptions.OptsVecSpace, default="Hdiv"
FEEC space used for the unknown ``u`` variable. ``"Hdiv"`` and
``"H1vec"`` are supported for ``algo="explicit"``, only ``"Hdiv"``
for ``algo="discrete_gradient"``.
solver : LiteralOptions.OptsSymmSolver, default="pcg"
Symmetric iterative solver used for mass-matrix inversions.
precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner"
Preconditioner used with the mass matrix.
solver_params : SolverParameters, default=None
Iterative-solver controls. If ``None``, defaults to
``SolverParameters()``.
filter_params : FilterParameters, default=None
Filter of the accumulated (particle-to-grid) vectors; defaults to ``FilterParameters()`` (no filter).
The three-point (binomial) filter ``FilterParameters(use_filter="three_point")`` reduces the noise,
but the propagator then no longer conserves the energy exactly (the filter is applied to the
accumulated vectors only, not to the fields seen by the particles).
dg_solver_params : DiscreteGradientSolverParameters, default=None
Nonlinear solver controls for discrete-gradient mode. If ``None``
and ``algo="discrete_gradient"``, defaults to
``DiscreteGradientSolverParameters()``.
"""
# specific literals
OptsAlgo = Literal[
"discrete_gradient",
"explicit",
]
# propagator options
ep_scale: float = None
algo: OptsAlgo = "explicit"
butcher: ButcherTableau = None
u_space: LiteralOptions.OptsVecSpace = "Hdiv"
solver: LiteralOptions.OptsSymmSolver = "pcg"
precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner"
solver_params: SolverParameters = None
filter_params: FilterParameters = None
dg_solver_params: DiscreteGradientSolverParameters = None
def __post_init__(self):
# checks
check_option(self.algo, self.OptsAlgo)
check_option(self.u_space, LiteralOptions.OptsVecSpace)
check_option(self.solver, LiteralOptions.OptsSymmSolver)
check_option(self.precond, LiteralOptions.OptsMassPrecond)
assert self.ep_scale is None or isinstance(self.ep_scale, float)
# supported u_space per algorithm (pusher kernels exist only for these)
if self.algo == "explicit":
supported_u_spaces = ("Hdiv", "H1vec")
else:
supported_u_spaces = ("Hdiv",)
if self.u_space not in supported_u_spaces:
raise ValueError(
f"{self.u_space = } is not supported by CurrentCoupling5DGradB with {self.algo = }, "
f"choose from {supported_u_spaces}.",
)
# defaults
if self.algo == "explicit" and self.butcher is None:
self.butcher = ButcherTableau()
if self.algo == "discrete_gradient" and self.dg_solver_params is None:
self.dg_solver_params = DiscreteGradientSolverParameters()
if self.solver_params is None:
self.solver_params = SolverParameters()
if self.filter_params is None:
self.filter_params = FilterParameters()
@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):
# factor of the energetic-particle contributions (default: mass ratio A_h/A_b)
self._ep_scale = self.options.ep_scale
if self._ep_scale is None:
self._ep_scale = self.variables.energetic_ions.species.mass_number / self.variables.u.species.mass_number
logger.info(f"{self.__class__.__name__}: ep_scale = A_h/A_b = {self._ep_scale} (mass ratio, default).")
if self.options.u_space == "H1vec":
self._u_form_int = 0
else:
self._u_form_int = int(self.derham.space_to_form[self.options.u_space])
# call operatros
id_M = "M" + self.derham.space_to_form[self.options.u_space] + "n"
self._A = getattr(self.mass_ops, id_M)
self._PB = getattr(self.basis_ops, "PB")
# Preconditioner
if self.options.precond is None:
pc = None
else:
pc_class = getattr(preconditioner, self.options.precond)
pc = pc_class(getattr(self.mass_ops, id_M))
# linear solver
self._A_inv = inverse(
self._A,
self.options.solver,
pc=pc,
tol=self.options.solver_params.tol,
maxiter=self.options.solver_params.maxiter,
verbose=self.options.solver_params.verbose,
recycle=self.options.solver_params.recycle,
)
# magnetic equilibrium field
unit_b1 = self.projected_equil.unit_b1
curl_unit_b2 = self.projected_equil.curl_unit_b2
self._b2 = self.projected_equil.b2
gradB1 = self.projected_equil.gradB1
absB0 = self.projected_equil.absB0
# scaling factor
epsilon = self.variables.energetic_ions.species.equation_params.epsilon
if self.options.algo == "explicit":
# temporary vectors to avoid memory allocation
self._b_full = self._b2.space.zeros()
self._u_new = self.variables.u.spline.vector.space.zeros()
self._u_temp = self.variables.u.spline.vector.space.zeros()
self._ku = self.variables.u.spline.vector.space.zeros()
self._PB_b = self._PB.codomain.zeros()
self._grad_PB_b = self.derham.grad.codomain.zeros()
# define Accumulator (vector only) and arguments
self._ACC = AccumulatorVector(
self.variables.energetic_ions.particles,
self.options.u_space,
cc_lin_mhd_5d_gradB,
self.mass_ops,
self.domain.args_domain,
filter_params=self.options.filter_params,
)
self._args_accum_kernel = (
epsilon,
self._ep_scale,
self._b_full[0]._data,
self._b_full[1]._data,
self._b_full[2]._data,
unit_b1[0]._data,
unit_b1[1]._data,
unit_b1[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
self._grad_PB_b[0]._data,
self._grad_PB_b[1]._data,
self._grad_PB_b[2]._data,
gradB1[0]._data,
gradB1[1]._data,
gradB1[2]._data,
self._u_form_int,
)
# define Pusher
if self.options.u_space == "Hdiv":
self._pusher_kernel = push_gc_cc_J2_stage_Hdiv
elif self.options.u_space == "H1vec":
self._pusher_kernel = push_gc_cc_J2_stage_H1vec
else:
raise ValueError(
f'{self.options.u_space =} not valid, choose from "Hdiv" or "H1vec".',
)
# temp fix due to refactoring of ButcherTableau:
self._args_pusher_kernel = (
self.domain.args_domain,
self.derham.args_derham,
epsilon,
self._b_full[0]._data,
self._b_full[1]._data,
self._b_full[2]._data,
unit_b1[0]._data,
unit_b1[1]._data,
unit_b1[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
self._u_temp[0]._data,
self._u_temp[1]._data,
self._u_temp[2]._data,
self.options.butcher.a_stage,
self.options.butcher.b,
self.options.butcher.c,
)
else:
# temporary vectors to avoid memory allocation
self._b_full = self._b2.space.zeros()
self._PB_b = self._PB.codomain.zeros()
self._grad_PB_b = self.derham.grad.codomain.zeros()
self._u_old = self.variables.u.spline.vector.space.zeros()
self._u_new = self.variables.u.spline.vector.space.zeros()
self._u_diff = self.variables.u.spline.vector.space.zeros()
self._u_mid = self.variables.u.spline.vector.space.zeros()
self._M2n_dot_u = self.variables.u.spline.vector.space.zeros()
self._ku = self.variables.u.spline.vector.space.zeros()
self._u_temp = self.variables.u.spline.vector.space.zeros()
# magnetic perturbation (zero if not given)
if self.b_tilde is None:
b_tilde = self._b2.space.zeros()
else:
b_tilde = self.b_tilde.spline.vector
# Call the accumulation and Pusher class
accum_kernel_init = cc_lin_mhd_5d_gradB_dg_init
accum_kernel = cc_lin_mhd_5d_gradB_dg
self._accum_kernel_en_fB_mid = eval_gradB_ediff
self._args_accum_kernel = (
epsilon,
self._ep_scale,
b_tilde[0]._data,
b_tilde[1]._data,
b_tilde[2]._data,
self._b2[0]._data,
self._b2[1]._data,
self._b2[2]._data,
unit_b1[0]._data,
unit_b1[1]._data,
unit_b1[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
self._grad_PB_b[0]._data,
self._grad_PB_b[1]._data,
self._grad_PB_b[2]._data,
gradB1[0]._data,
gradB1[1]._data,
gradB1[2]._data,
self._u_form_int,
)
self._args_accum_kernel_en_fB_mid = (
self.domain.args_domain,
self.derham.args_derham,
gradB1[0]._data,
gradB1[1]._data,
gradB1[2]._data,
self._grad_PB_b[0]._data,
self._grad_PB_b[1]._data,
self._grad_PB_b[2]._data,
)
self._ACC_init = AccumulatorVector(
self.variables.energetic_ions.particles,
self.options.u_space,
accum_kernel_init,
self.mass_ops,
self.domain.args_domain,
filter_params=self.options.filter_params,
)
self._ACC = AccumulatorVector(
self.variables.energetic_ions.particles,
self.options.u_space,
accum_kernel,
self.mass_ops,
self.domain.args_domain,
filter_params=self.options.filter_params,
)
self._args_pusher_kernel_init = (
self.domain.args_domain,
self.derham.args_derham,
epsilon,
self._b_full[0]._data,
self._b_full[1]._data,
self._b_full[2]._data,
unit_b1[0]._data,
unit_b1[1]._data,
unit_b1[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
self.variables.u.spline.vector[0]._data,
self.variables.u.spline.vector[1]._data,
self.variables.u.spline.vector[2]._data,
)
self._args_pusher_kernel = (
self.domain.args_domain,
self.derham.args_derham,
epsilon,
self._b_full[0]._data,
self._b_full[1]._data,
self._b_full[2]._data,
unit_b1[0]._data,
unit_b1[1]._data,
unit_b1[2]._data,
curl_unit_b2[0]._data,
curl_unit_b2[1]._data,
curl_unit_b2[2]._data,
self._u_mid[0]._data,
self._u_mid[1]._data,
self._u_mid[2]._data,
self._u_temp[0]._data,
self._u_temp[1]._data,
self._u_temp[2]._data,
)
self._pusher_kernel_init = push_gc_cc_J2_dg_init_Hdiv
self._pusher_kernel = push_gc_cc_J2_dg_Hdiv
def __call__(self, dt):
# current FE coeffs
un = self.variables.u.spline.vector
# particle markers and idx
particles = self.variables.energetic_ions.particles
holes = particles.holes
args_markers = particles.args_markers
markers = args_markers.markers
first_pusher_idx = args_markers.first_pusher_idx
first_free_idx = args_markers.first_free_idx
# clear buffer
markers[first_pusher_idx:-2] = 0.0
# save old marker positions
markers[first_pusher_idx : first_pusher_idx + 3] = markers[:3]
# sum up total magnetic field b_full1 = b_eq + b_tilde (in-place)
b_full = self._b2.copy(out=self._b_full)
if self.b_tilde is not None:
b_full += self.b_tilde.spline.vector
b_full.update_ghost_regions()
if self.options.algo == "explicit":
if self.b_tilde is not None:
PB_b = self._PB.dot(self.b_tilde.spline.vector, out=self._PB_b)
grad_PB_b = self.derham.grad.dot(PB_b, out=self._grad_PB_b)
grad_PB_b.update_ghost_regions()
# save old u
u_new = un.copy(out=self._u_new)
u_temp = un.copy(out=self._u_temp)
for stage in range(self.options.butcher.n_stages):
# accumulate
self._ACC(
*self._args_accum_kernel,
)
# push particles
self._pusher_kernel(
dt,
stage,
args_markers,
*self._args_pusher_kernel,
)
# kinetic boundary conditions are applied per marker inside the kernel
particles.finish_kernel_bc()
if particles.mpi_comm is not None:
particles.mpi_sort_markers(apply_bc=False)
# solve linear system for updating u coefficients
ku = self._A_inv.dot(self._ACC.vectors[0], out=self._ku)
info = self._A_inv._info
# calculate u^{n+1}_k
u_temp = un.copy(out=self._u_temp)
u_temp += ku * dt * self.options.butcher.a_stage[stage]
u_temp.update_ghost_regions()
# calculate u^{n+1}
u_new += ku * dt * self.options.butcher.b[stage]
u_new.update_ghost_regions()
if self.options.solver_params.info and MPI.COMM_WORLD.Get_rank() == 0:
logger.info(f"Stage: {stage}")
logger.info(f"Status for CurrentCoupling5DGradB: {info['success']}")
logger.info(f"Iterations for CurrentCoupling5DGradB: {info['niter']}")
logger.info("")
# update u coefficients
diffs = self.update_feec_variables(u=u_new)
# clear the buffer
markers[first_pusher_idx:-2] = 0.0
# update_weights
if self.variables.energetic_ions.species.weights_params.control_variate:
particles.update_weights()
if self.options.solver_params.info and MPI.COMM_WORLD.Get_rank() == 0:
logger.info(f"Maxdiff up for CurrentCoupling5DGradB: {diffs['u']}")
logger.info("")
else:
# total number of markers
n_mks_tot = particles.Np
# relaxation factor
alpha = self.options.dg_solver_params.relaxation_factor
# eval parallel tilde b and its gradient
PB_b = self._PB_b
if self.b_tilde is not None:
PB_b = self._PB.dot(self.b_tilde.spline.vector, out=self._PB_b)
PB_b.update_ghost_regions()
grad_PB_b = self.derham.grad.dot(PB_b, out=self._grad_PB_b)
grad_PB_b.update_ghost_regions()
# save old u
u_old = un.copy(out=self._u_old)
u_new = un.copy(out=self._u_new)
# save en_U_old
self._A.dot(un, out=self._M2n_dot_u)
en_U_old = un.inner(self._M2n_dot_u) / 2.0
# save en_fB_old
particles.save_magnetic_energy(PB_b)
en_fB_old = xp.sum(markers[8, ~holes].dot(markers[5, ~holes])) * self._ep_scale
en_fB_old /= n_mks_tot
buffer_array = xp.array([en_fB_old])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
if particles.clone_config is not None:
particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
en_fB_old = buffer_array[0]
en_tot_old = en_U_old + en_fB_old
# initial guess
self._ACC_init(*self._args_accum_kernel)
ku = self._A_inv.dot(self._ACC_init.vectors[0], out=self._ku)
u_new += ku * dt
u_new.update_ghost_regions()
# save en_U_new
self._A.dot(u_new, out=self._M2n_dot_u)
en_U_new = u_new.inner(self._M2n_dot_u) / 2.0
# push eta
self._pusher_kernel_init(
dt,
args_markers,
*self._args_pusher_kernel_init,
)
if particles.mpi_comm is not None:
particles.mpi_sort_markers(apply_bc=False)
# save en_fB_new
particles.save_magnetic_energy(PB_b)
en_fB_new = xp.sum(markers[8, ~holes].dot(markers[5, ~holes])) * self._ep_scale
en_fB_new /= n_mks_tot
buffer_array = xp.array([en_fB_new])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
if particles.clone_config is not None:
particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
en_fB_new = buffer_array[0]
# fixed-point iterations
iter_num = 0
while True:
iter_num += 1
if self.options.dg_solver_params.verbose and MPI.COMM_WORLD.Get_rank() == 0:
logger.info(f"# of iteration: {iter_num}")
# calculate discrete gradient
# save u^{n+1, k}
u_old = u_new.copy(out=self._u_old)
u_diff = u_old.copy(out=self._u_diff)
u_diff -= un
u_diff.update_ghost_regions()
u_mid = u_old.copy(out=self._u_mid)
u_mid += un
u_mid /= 2.0
u_mid.update_ghost_regions()
# save H^{n+1, k}
markers[first_free_idx : first_free_idx + 3, ~holes] = markers[0:3, ~holes]
# calculate denominator ||z^{n+1, k} - z^n||^2
sum_u_diff_loc = xp.sum((u_diff.toarray() ** 2))
sum_H_diff_loc = xp.sum(
(markers[:3, ~holes] - markers[first_pusher_idx : first_pusher_idx + 3, ~holes]) ** 2,
)
buffer_array = xp.array([sum_u_diff_loc])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
denominator = buffer_array[0]
buffer_array = xp.array([sum_H_diff_loc])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
if particles.clone_config is not None:
particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
denominator += buffer_array[0]
# sorting markers at mid-point
if particles.mpi_comm is not None:
particles.mpi_sort_markers(apply_bc=False, alpha=0.5)
self._accum_kernel_en_fB_mid(
args_markers,
*self._args_accum_kernel_en_fB_mid,
first_free_idx + 3,
)
en_fB_mid = xp.sum(markers[first_free_idx + 3, ~holes].dot(markers[5, ~holes])) * self._ep_scale
en_fB_mid /= n_mks_tot
buffer_array = xp.array([en_fB_mid])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
if particles.clone_config is not None:
particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
en_fB_mid = buffer_array[0]
if denominator == 0.0:
const = 0.0
else:
const = (en_fB_new - en_fB_old - en_fB_mid) / denominator
# update u^{n+1, k}
self._ACC(*self._args_accum_kernel, const)
ku = self._A_inv.dot(self._ACC.vectors[0], out=self._ku)
u_new = un.copy(out=self._u_new)
u_new += ku * dt
u_new *= alpha
u_new += u_old * (1.0 - alpha)
u_new.update_ghost_regions()
# update en_U_new
self._A.dot(u_new, out=self._M2n_dot_u)
en_U_new = u_new.inner(self._M2n_dot_u) / 2.0
# update H^{n+1, k}
self._pusher_kernel(
dt,
args_markers,
*self._args_pusher_kernel,
const,
alpha,
)
sum_H_diff_loc = xp.sum(
xp.abs(markers[0:3, ~holes] - markers[first_free_idx : first_free_idx + 3, ~holes]),
)
if particles.mpi_comm is not None:
particles.mpi_sort_markers(apply_bc=False)
# update en_fB_new
particles.save_magnetic_energy(PB_b)
en_fB_new = xp.sum(markers[8, ~holes].dot(markers[5, ~holes])) * self._ep_scale
en_fB_new /= n_mks_tot
buffer_array = xp.array([en_fB_new])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
if particles.clone_config is not None:
particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
en_fB_new = buffer_array[0]
# calculate total energy difference
e_diff = xp.abs(en_U_new + en_fB_new - en_tot_old)
# calculate ||z^{n+1, k} - z^{n+1, k-1||
sum_u_diff_loc = xp.sum(xp.abs(u_new.toarray() - u_old.toarray()))
buffer_array = xp.array([sum_u_diff_loc])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
diff = buffer_array[0]
buffer_array = xp.array([sum_H_diff_loc])
if particles.mpi_comm is not None:
particles.mpi_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
if particles.clone_config is not None:
particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
buffer_array,
op=MPI.SUM,
)
diff += buffer_array[0]
# check convergence
if diff < self.options.dg_solver_params.tol:
if self.options.dg_solver_params.verbose and MPI.COMM_WORLD.Get_rank() == 0:
logger.info(f"converged diff: {diff}")
logger.info(f"converged e_diff: {e_diff}")
if particles.mpi_comm is not None:
particles.mpi_comm.Barrier()
break
else:
if self.options.dg_solver_params.verbose and MPI.COMM_WORLD.Get_rank() == 0:
logger.info(f"not converged diff: {diff}")
logger.info(f"not converged e_diff: {e_diff}")
if iter_num == self.options.dg_solver_params.maxiter:
if self.options.dg_solver_params.info and MPI.COMM_WORLD.Get_rank() == 0:
logger.info(
f"{iter_num =}, maxiter={self.options.dg_solver_params.maxiter} reached! diff: {diff}, e_diff: {e_diff}",
)
if particles.mpi_comm is not None:
particles.mpi_comm.Barrier()
break
# sorting markers
if particles.mpi_comm is not None:
particles.mpi_sort_markers()
else:
particles.apply_kinetic_bc()
# update u coefficients
diffs = self.update_feec_variables(u=u_new)
# clear the buffer
markers[first_pusher_idx:-2] = 0.0
# update_weights
if self.variables.energetic_ions.species.weights_params.control_variate:
particles.update_weights()
if self.options.dg_solver_params.info and MPI.COMM_WORLD.Get_rank() == 0:
logger.info(f"Maxdiff up for CurrentCoupling5DGradB: {diffs['u']}")
logger.info("")
