"Base classes for particle deposition (accumulation) on the grid."
from dataclasses import dataclass
import cunumpy as xp
from cunumpy.kernels import Kernel, PyccelKernel
from feectools.linalg.basic import LinearOperator, Vector, VectorSpace
from feectools.linalg.block import BlockVector, BlockVectorSpace
from feectools.linalg.stencil import StencilMatrix, StencilVector
from maybempi import MPI
from scope_profiler import ProfileManager
from struphy.feec.mass import WeightedMassOperators
from struphy.feec.psydac_derham import Derham
from struphy.io.options import LiteralOptions
from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments
from struphy.models.variables import PICVariable, SPHVariable
from struphy.pic.accumulation.filter import AccumFilter, FilterParameters
from struphy.pic.base import Particles
from struphy.utils.cuda_arguments import check_mapping_on_device
from struphy.utils.utils import __dataclass_repr_no_defaults__, check_option
def _accumulation_kernel(kernel) -> Kernel:
"""The kernel as a :class:`cunumpy.kernels.Kernel` (pyccel or CUDA depending on the backend)."""
if not isinstance(kernel, Kernel):
kernel = Kernel(kernel)
return kernel
[docs]
class Accumulator:
r"""
Approximates integrals of the form
.. math::
I_A &= \int_\Omega \int_{\mathbb R^3} \Lambda^\mu_{ijk}(\boldsymbol \eta) \, A^{\mu, \nu}(\boldsymbol \eta, \mathbf v) \, \Lambda^\nu_{mno}(\boldsymbol \eta) \, f^{\textrm{vol}}(\boldsymbol \eta, \mathbf v)\,\mathrm d\mathbf v \textrm d \boldsymbol \eta\,,
\\[2mm]
I_B &= \int_\Omega \int_{\mathbb R^3} \Lambda^\mu_{ijk}(\boldsymbol \eta) \, B^\mu(\boldsymbol \eta, \mathbf v) \, f^{\textrm{vol}}(\boldsymbol \eta, \mathbf v)\,\mathrm d\mathbf v \textrm d \boldsymbol \eta\,,
for given weight functions :math:`A^{\mu,\nu}` and :math:`B^\mu` by Monte-Carlo quadrature through the particle distribution function :math:`f^{\textrm{vol}}`:
.. math::
f^{\textrm{vol}}(\boldsymbol \eta, \mathbf v) \approx \sum_{p=0}^{N-1} w_p \, \delta(\boldsymbol \eta - \boldsymbol \eta_p) \, \delta(\mathbf v - \mathbf v_p)\,.
This results in stencil (block) matrices and vectors
.. math::
M &= (M^{\mu,\nu})_{\mu,\nu}\,,\qquad && M^{\mu,\nu} \in \mathbb R^{\mathbb N^\alpha_\mu \times \mathbb N^\alpha_\nu}\,,
\\[2mm]
V &= (V^\mu)_\mu\,,\qquad &&V^\mu \in \mathbb R^{\mathbb N^\alpha_\mu}\,,
where :math:`N^\alpha_\mu` denotes the dimension of the :math:`\mu`-th component
of the :class:`~struphy.feec.psydac_derham.Derham` space
:math:`V_h^\alpha` (:math:`\mu,\nu = 1,2,3` for vector-valued spaces),
with entries obtained by summing over all particles :math:`p`,
.. math::
M^{\mu,\nu}_{ijk,mno} &= \sum_{p=0}^{N-1} w_p\, \Lambda^\mu_{ijk}(\boldsymbol \eta_p) \, A^{\mu,\nu}_p \, \Lambda^\nu_{mno}(\boldsymbol \eta_p) \,,
\\[2mm]
V^\mu_{ijk} &= \sum_{p=0}^{N-1} w_p\, \Lambda^\mu_{ijk}(\boldsymbol \eta_p) \, B^\mu_p \,.
Here, :math:`\Lambda^\mu_{ijk}(\boldsymbol \eta_p)` denotes the :math:`ijk`-th basis function
of the :math:`\mu`-th component of a Derham space.
Parameters
----------
particles : Particles
Particles object holding the markers to accumulate.
space_id : str
Space identifier for the matrix/vector (H1, Hcurl, Hdiv, L2 or H1vec) to be accumulated into.
kernel : Kernel | PyccelKernel
The accumulation kernel, e.g. ``charge_density_0form`` from ``struphy.pic.accumulation.kernels.charge_density_0form``.
On the CuPy backend, a kernel without CUDA version raises NotImplementedError.
derham : Derham
Discrete FE spaces object.
args_domain : DomainArguments
Mapping infos.
add_vector : bool
True if, additionally to a matrix, a vector in the same space is to be accumulated. Default=False.
symmetry : str
In case of space_id=Hcurl/Hdiv, the symmetry property of the block matrix: diag, asym, symm, pressure or None (=full matrix, default)
filter_params : dict
Params for the accumulation filter: use_filter(string, either `three_point or `fourier), repeat(int), alpha(float) and modes(list with int).
Note
----
Struphy accumulation kernels live in ``struphy/pic/accumulation/kernels/``, one folder per kernel
(see :mod:`struphy.pic.accumulation.kernels`).
"""
[docs]
def __init__(
self,
particles: Particles,
space_id: str,
kernel: Kernel | PyccelKernel,
mass_ops: WeightedMassOperators,
args_domain: DomainArguments,
*,
add_vector: bool = False,
symmetry: str = None,
filter_params: FilterParameters = None,
):
self._particles = particles
self._space_id = space_id
self._kernel = _accumulation_kernel(kernel)
self._derham = mass_ops.derham
self._args_domain = args_domain
# profiling region names (precomputed, they are looked up on every call)
self._region_name = "accum: " + kernel.name
self._comm_region_name = "accum comm: " + kernel.name
self._symmetry = symmetry
self._form = self.derham.space_to_form[space_id]
# initialize matrices (instances of WeightedMassOperator)
self._operators = []
# special treatment in model LinearMHDVlasovPC (symmetry=pressure, six symmetric BlockMatrices are needed)
if symmetry == "pressure":
for _ in range(6):
operator = mass_ops.create_weighted_mass(
space_id,
space_id,
weights="symm",
)
self._operators.append(operator)
# "normal" treatment (just one matrix)
else:
operator = mass_ops.create_weighted_mass(
space_id,
space_id,
weights=symmetry,
)
self._operators.append(operator)
# collect all _data attributes needed in accumulation kernel
self._args_data = ()
for op in self._operators:
if isinstance(op.matrix, StencilMatrix):
self._args_data += (op.matrix._data,)
else:
for a, row in enumerate(op.matrix.blocks):
for b, bl in enumerate(row):
if symmetry in ["pressure", "symm", "asym", "diag"]:
if b >= a and bl is not None:
self._args_data += (bl._data,)
else:
if bl is not None:
self._args_data += (bl._data,)
# initialize vectors
self._vectors = []
self._vectors_temp = []
self._vectors_out = []
if add_vector:
# special treatment in model LinearMHDVlasovPC (symmetry=pressure, three BlockVectors are needed)
if symmetry == "pressure":
for _ in range(3):
self._vectors += [BlockVector(self.derham.coeff_spaces[self.form])]
self._vectors_temp += [
BlockVector(self.derham.coeff_spaces[self.form]),
]
self._vectors_out += [
BlockVector(self.derham.coeff_spaces[self.form]),
]
# normal treatment (just one vector)
else:
for op in self._operators:
if isinstance(op.matrix, StencilMatrix):
self._vectors += [StencilVector(op.matrix.domain)]
self._vectors_temp += [StencilVector(op.matrix.domain)]
self._vectors_out += [StencilVector(op.matrix.domain)]
else:
self._vectors += [BlockVector(op.matrix.domain)]
self._vectors_temp += [BlockVector(op.matrix.domain)]
self._vectors_out += [BlockVector(op.matrix.domain)]
for vec in self._vectors:
if isinstance(vec, StencilVector):
self._args_data += (vec._data,)
else:
for bl in vec.blocks:
self._args_data += (bl._data,)
# initialize filter
self._accfilter = AccumFilter(filter_params, self._derham, self._space_id)
[docs]
def __call__(self, *optional_args):
"""
Performs the accumulation into the matrix/vector by calling the chosen accumulation kernel.
Parameters
----------
particles : Particles
Particles object holding the markers information in format particles.markers.shape == (:, n_markers).
optional_args : any
Additional arguments to be passed to the accumulator kernel, besides the mandatory arguments
which are prepared automatically (spline bases info, mapping info, data arrays).
Examples would be parameters for a background kinetic distribution or spline coefficients of a background magnetic field.
Entries must be pyccel-conform types.
"""
with ProfileManager.profile_region(self._region_name):
self._accumulate(*optional_args)
def _accumulate(self, *optional_args):
"""Body of :meth:`__call__`, see there."""
# flags for break
vec_finished = False
# reset data
for dat in self._args_data:
dat[:] = 0.0
# accumulate into matrix (and vector) with markers
with ProfileManager.profile_region("kernel: " + self.kernel.name):
self.kernel(
self.particles.args_markers,
self.derham.args_derham,
self.args_domain,
*self._args_data,
*optional_args,
n_threads=self.particles.args_markers.n_markers, # the markers (columns) that can hold markers
)
# apply filter
if self.accfilter.params.use_filter is not None:
for vec in self._vectors:
with ProfileManager.profile_region(self._comm_region_name):
vec.exchange_assembly_data()
vec.update_ghost_regions()
self.accfilter(vec)
vec_finished = True
if self.particles.clone_config is None:
num_clones = 1
else:
num_clones = self.particles.clone_config.num_clones
if num_clones > 1:
with ProfileManager.profile_region(self._comm_region_name):
for data_array in self._args_data:
self.particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
data_array,
op=MPI.SUM,
)
# finish vector: accumulate ghost regions and update ghost regions
if not vec_finished:
with ProfileManager.profile_region(self._comm_region_name):
for vec in self._vectors:
vec.exchange_assembly_data()
vec.update_ghost_regions()
# finish matrix: accumulate ghost regions, update ghost regions and copy data for symmetric/antisymmetric block matrices
with ProfileManager.profile_region(self._comm_region_name):
for op in self._operators:
op.matrix.exchange_assembly_data()
op.matrix.update_ghost_regions()
if self.symmetry == "symm":
self._operators[0].matrix[0, 1].transpose(
out=self._operators[0].matrix[1, 0],
)
self._operators[0].matrix[0, 2].transpose(
out=self._operators[0].matrix[2, 0],
)
self._operators[0].matrix[1, 2].transpose(
out=self._operators[0].matrix[2, 1],
)
elif self.symmetry == "asym":
self._operators[0].matrix[0, 1].transpose(
out=self._operators[0].matrix[1, 0],
)
self._operators[0].matrix[1, 0] *= -1
self._operators[0].matrix[0, 2].transpose(
out=self._operators[0].matrix[2, 0],
)
self._operators[0].matrix[2, 0] *= -1
self._operators[0].matrix[1, 2].transpose(
out=self._operators[0].matrix[2, 1],
)
self._operators[0].matrix[2, 1] *= -1
elif self.symmetry == "pressure":
for i in range(6):
self._operators[i].matrix[0, 1].transpose(
out=self._operators[i].matrix[1, 0],
)
self._operators[i].matrix[0, 2].transpose(
out=self._operators[i].matrix[2, 0],
)
self._operators[i].matrix[1, 2].transpose(
out=self._operators[i].matrix[2, 1],
)
@property
def particles(self):
"""Particle object."""
return self._particles
@property
def kernel(self) -> Kernel:
"""The accumulation kernel."""
return self._kernel
[docs]
def kernels(self) -> tuple[Kernel, ...]:
"""The kernels this accumulator calls (the accumulation kernel)."""
return (self._kernel,)
@property
def derham(self):
"""Discrete Derham complex on the logical unit cube."""
return self._derham
@property
def args_domain(self):
"""Mapping info for evaluating metric coefficients."""
return self._args_domain
@property
def space_id(self):
"""Space identifier for the matrix/vector (H1, Hcurl, Hdiv, L2 or H1vec) to be accumulated into."""
return self._space_id
@property
def form(self):
"""p-form ("0", "1", "2", "3" or "v") to be accumulated into."""
return self._form
@property
def symmetry(self):
"""Symmetry of the accumulation matrix (diagonal, symmetric, asymmetric, etc.)."""
return self._symmetry
@property
def operators(self):
"""List of WeightedMassOperators of the accumulator. Matrices can be accessed e.g. with operators[0].matrix."""
return self._operators
@property
def vectors(self):
"""List of Stencil-/Block-/PolarVectors of the accumulator."""
out = []
for vec, vec_temp, vec_out in zip(self._vectors, self._vectors_temp, self._vectors_out):
self._derham.extraction_ops[self.form].dot(vec, out=vec_temp)
self._derham.boundary_ops[self.form].dot(vec_temp, out=vec_out)
out += [vec_out]
return out
@property
def accfilter(self):
"""Callable filters"""
return self._accfilter
[docs]
def show_accumulated_spline_field(self, mass_ops: WeightedMassOperators, eta_direction=0, component=0):
r"""1D plot of the spline field corresponding to the accumulated vector.
The latter can be viewed as the rhs of an L2-projection:
.. math::
\mathbb M \mathbf a = \sum_p \boldsymbol \Lambda(\boldsymbol \eta_p) * B_p\,.
The FE coefficients :math:`\mathbf a` determine a FE :class:`~struphy.feec.psydac_derham.SplineFunction`.
"""
from matplotlib import pyplot as plt
from struphy.feec.mass import L2Projector
# L2 projection
proj = L2Projector(self.space_id, mass_ops)
a = proj.solve(self.vectors[0])
# create field and assign coeffs
field = self.derham.create_spline_function("accum_field", self.space_id)
field.vector = a
# plot field
eta = xp.linspace(0, 1, 100)
if eta_direction == 0:
args = (eta, 0.5, 0.5)
elif eta_direction == 1:
args = (0.5, eta, 0.5)
else:
args = (0.5, 0.5, eta)
vals = mass_ops.domain.push(field, *args, kind="1", squeeze_out=True)
plt.plot(eta, vals[component])
plt.title(
f'Spline field accumulated with the kernel "{self.kernel}"',
)
plt.xlabel(rf"$\eta_{eta_direction + 1}$")
plt.ylabel("field amplitude")
plt.show()
[docs]
class AccumulatorVector:
r"""
Approximates integrals of the form
.. math::
I_B = \int_\Omega \int_{\mathbb R^3} \Lambda^\mu_{ijk}(\boldsymbol \eta) \, B^\mu(\boldsymbol \eta, \mathbf v) \, f^{\textrm{vol}}(\boldsymbol \eta, \mathbf v)\,\mathrm d\mathbf v \textrm d \boldsymbol \eta\,,
for a given weight function and :math:`B^\mu` by Monte-Carlo quadrature through the particle distribution function :math:`f^{\textrm{vol}}`:
.. math::
f^{\textrm{vol}}(\boldsymbol \eta, \mathbf v) \approx \sum_{p=0}^{N-1} w_p \, \delta(\boldsymbol \eta - \boldsymbol \eta_p) \, \delta(\mathbf v - \mathbf v_p)\,.
This results in a stencil (block) vector
.. math::
V = (V^\mu)_\mu\,,\qquad V^\mu \in \mathbb R^{\mathbb N^\alpha_\mu}\,,
where :math:`N^\alpha_\mu` denotes the dimension of the :math:`\mu`-th component
of the :class:`~struphy.feec.psydac_derham.Derham` space
:math:`V_h^\alpha` (:math:`\mu,\nu = 1,2,3` for vector-valued spaces),
with entries obtained by summing over all particles :math:`p`,
.. math::
V^\mu_{ijk} = \sum_{p=0}^{N-1} w_p\, \Lambda^\mu_{ijk}(\boldsymbol \eta_p) \, B^\mu_p \,.
Here, :math:`\Lambda^\mu_{ijk}(\boldsymbol \eta_p)` denotes the :math:`ijk`-th basis function
of the :math:`\mu`-th component of a Derham space.
Similar to :class:`~struphy.pic.accumulation.particles_to_grid.Accumulator` but only for vectors :math:`V`.
Parameters
----------
particles : Particles
Particles object holding the markers to accumulate.
space_id : str
Space identifier for the matrix/vector (H1, Hcurl, Hdiv, L2 or H1vec) to be accumulated into.
kernel : Kernel | PyccelKernel
The accumulation kernel, e.g. ``charge_density_0form`` from ``struphy.pic.accumulation.kernels.charge_density_0form``.
On the CuPy backend, a kernel without CUDA version raises NotImplementedError.
derham : Derham
Discrete FE spaces object.
args_domain : DomainArguments
Mapping infos.
filter_params : FilterParameters
Params for the accumulation filter.
n_vectors : int
Number of vectors accumulated by one kernel call (in the order of the kernel arguments), e.g. 3 for the
pressure coupling. Default=1.
"""
[docs]
def __init__(
self,
particles: Particles,
space_id: str,
kernel: Kernel | PyccelKernel,
mass_ops: WeightedMassOperators,
args_domain: DomainArguments,
filter_params: FilterParameters = None,
n_vectors: int = 1,
):
self._particles = particles
self._space_id = space_id
self._kernel = _accumulation_kernel(kernel)
self._derham = mass_ops.derham
self._args_domain = args_domain
# profiling region names (precomputed, they are looked up on every call)
self._region_name = "accum: " + kernel.name
self._comm_region_name = "accum comm: " + kernel.name
self._form = self.derham.space_to_form[space_id]
# initialize vectors
self._vectors = []
self._vectors_temp = []
self._vectors_out = []
# collect all _data attributes needed in accumulation kernel
self._args_data = ()
if space_id in ("H1", "L2"):
vector_class = StencilVector
elif space_id in ("Hcurl", "Hdiv", "H1vec"):
vector_class = BlockVector
coeff_space = self.derham.fem_spaces[self.form].coeff_space
for _ in range(n_vectors):
self._vectors += [vector_class(coeff_space)]
self._vectors_temp += [vector_class(coeff_space)]
self._vectors_out += [vector_class(coeff_space)]
for vec in self._vectors:
if isinstance(vec, StencilVector):
self._args_data += (vec._data,)
else:
for bl in vec.blocks:
self._args_data += (bl._data,)
# initialize filter
self._accfilter = AccumFilter(filter_params, self._derham, self._space_id)
[docs]
def __call__(self, *optional_args):
"""
Performs the accumulation into the vector by calling the chosen accumulation kernel.
Parameters
----------
optional_args : any
Additional arguments to be passed to the accumulator kernel, besides the mandatory arguments
which are prepared automatically (spline bases info, mapping info, data arrays).
Examples would be parameters for a background kinetic distribution or spline coefficients of a background magnetic field.
Entries must be pyccel-conform types.
"""
with ProfileManager.profile_region(self._region_name):
self._accumulate(*optional_args)
def _accumulate(self, *optional_args):
"""Body of :meth:`__call__`, see there."""
# flags for break
vec_finished = False
# reset data
for dat in self._args_data:
dat[:] = 0.0
# accumulate into matrix (and vector) with markers
with ProfileManager.profile_region("kernel: " + self.kernel.name):
self.kernel(
self.particles.args_markers,
self.derham.args_derham,
self.args_domain,
*self._args_data,
*optional_args,
n_threads=self.particles.args_markers.n_markers, # the markers (columns) that can hold markers
)
# apply filter
if self.accfilter.params.use_filter is not None:
for vec in self._vectors:
with ProfileManager.profile_region(self._comm_region_name):
vec.exchange_assembly_data()
vec.update_ghost_regions()
self.accfilter(vec)
vec_finished = True
if self.particles.clone_config is None:
num_clones = 1
else:
num_clones = self.particles.clone_config.num_clones
if num_clones > 1:
with ProfileManager.profile_region(self._comm_region_name):
for data_array in self._args_data:
self.particles.clone_config.inter_comm.Allreduce(
MPI.IN_PLACE,
data_array,
op=MPI.SUM,
)
# finish vector: accumulate ghost regions and update ghost regions
if not vec_finished:
with ProfileManager.profile_region(self._comm_region_name):
for vec in self._vectors:
vec.exchange_assembly_data()
vec.update_ghost_regions()
@property
def particles(self):
"""Particle object."""
return self._particles
@property
def kernel(self) -> Kernel:
"""The accumulation kernel."""
return self._kernel
[docs]
def kernels(self) -> tuple[Kernel, ...]:
"""The kernels this accumulator calls (the accumulation kernel)."""
return (self._kernel,)
@property
def derham(self):
"""Discrete Derham complex on the logical unit cube."""
return self._derham
@property
def args_domain(self):
"""Mapping arguments."""
return self._args_domain
@property
def space_id(self):
"""Space identifier for the matrix/vector (H1, Hcurl, Hdiv, L2 or H1vec) to be accumulated into."""
return self._space_id
@property
def form(self):
"""p-form ("0", "1", "2", "3" or "v") to be accumulated into."""
return self._form
@property
def vectors(self):
"""List of Stencil-/Block-/PolarVectors of the accumulator."""
out = []
for vec, vec_temp, vec_out in zip(self._vectors, self._vectors_temp, self._vectors_out):
self._derham.extraction_ops[self.form].dot(vec, out=vec_temp)
self._derham.boundary_ops[self.form].dot(vec_temp, out=vec_out)
out += [vec_out]
return out
@property
def accfilter(self):
"""Callable filters"""
return self._accfilter
[docs]
def show_accumulated_spline_field(self, mass_ops, eta_direction=(True, False, False), save_L2=False):
r"""1 or 2D plot of the spline field corresponding to the accumulated vector.
The latter can be viewed as the rhs of an L2-projection:
.. math::
\mathbb M \mathbf a = \sum_p \boldsymbol \Lambda(\boldsymbol \eta_p) * B_p\,.
The FE coefficients :math:`\mathbf a` determine a FE :class:`~struphy.feec.psydac_derham.SplineFunction`.
:param eta_direction: axes of eta to show accumulation (eta1, eta2, eta3).
"""
assert sum(eta_direction) < 3, "Current implementation is only possible with 1 and 2D visualization"
from matplotlib import pyplot as plt
from struphy.feec.mass import L2Projector
# L2 projection
proj = L2Projector(self.space_id, mass_ops)
a = proj.solve(self.vectors[0])
if save_L2:
return a
# create field and assign coeffs
field = self.derham.create_spline_function("accum_field", self.space_id)
field.vector = a
# plot field
# initialize axis and slicing
eta = xp.linspace(0, 1, 100)
args = [0.5, 0.5, 0.5]
# fill slices to plot with eta
plt_axis = xp.flatnonzero(eta_direction)
for idx in plt_axis:
args[idx] = eta
args = tuple(args)
# field value at specified axes
field_value = field(*args, squeeze_out=True)
# One-dimensional case
if len(plt_axis) == 1:
plt.plot(eta, field_value)
plt.xlabel(rf"$\eta_{plt_axis[0] + 1}$")
plt.ylabel("field amplitude")
# Two-dimensional case
elif len(plt_axis) == 2:
Eta1, Eta2 = xp.meshgrid(eta, eta, indexing="ij")
pcm = plt.pcolor(Eta1, Eta2, field_value)
plt.colorbar(pcm, label="field amplitude")
plt.xlabel(rf"$\eta_{plt_axis[0] + 1}$")
plt.ylabel(rf"$\eta_{plt_axis[1] + 1}$")
plt.title(
f'Spline field accumulated with the kernel "{self.kernel}"',
)
plt.show()
[docs]
@dataclass
class ParticlesToGrid:
r"""Lightweight, serializable description of a particle-to-grid coupling
(for example charge- or current deposition) into FEEC degrees of freedom.
A ``ParticlesToGrid`` does not perform any accumulation itself: it simply bundles
the pieces needed to build an :class:`~struphy.pic.accumulation.particles_to_grid.AccumulatorVector`.
Parameters
----------
pic_variable : PICVariable | SPHVariable
The kinetic variable whose markers (``pic_variable.particles``) are deposited on the grid.
accum_space : {"H1", "Hcurl", "Hdiv", "L2", "H1vec"}
FEEC space identifier of the vector to accumulate into.
accum_kernel : Kernel | PyccelKernel
Accumulation kernel matching ``accum_space``, for example
``charge_density_0form`` from ``struphy.pic.accumulation.kernels.charge_density_0form``.
Examples
--------
>>> from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form
>>> from struphy.pic.accumulation.particles_to_grid import ParticlesToGrid
>>> from struphy.propagators.poisson_solve import PoissonSolve
>>> rho = ParticlesToGrid(
... kinetic_ions.var,
... "H1",
... charge_density_0form,
... )
>>> poisson = PoissonSolve(rho=rho, rho_coeffs=alpha**2 / epsilon)
"""
pic_variable: PICVariable | SPHVariable = None
accum_space: LiteralOptions.OptsFEECSpace = None
accum_kernel: Kernel | PyccelKernel = None
def __post_init__(self):
if self.accum_space is not None:
check_option(self.accum_space, LiteralOptions.OptsFEECSpace)
assert isinstance(self.accum_kernel, (Kernel, PyccelKernel)) or self.accum_kernel is None
def __repr_no_defaults__(self):
return __dataclass_repr_no_defaults__(self)
[docs]
class MarkerOperator(LinearOperator):
r"""Matrix-free particle matrix: the matrix-free counterpart of the matrices assembled by :class:`Accumulator`.
For a vector-valued space :math:`V^\alpha_h` (:math:`\alpha = 1, 2, v`) with the basis functions
:math:`\Lambda^\alpha_{\mu, ijk}`, and a 3x3 matrix :math:`F_p` per marker (e.g. of rank one,
:math:`F_p = c_p\, \mathbf e_p \mathbf e_p^\top`, or antisymmetric), the particle matrix is
.. math::
A_{(\mu, ijk), (\nu, mno)} = \sum_p \Lambda^\alpha_{\mu, ijk}(\boldsymbol \eta_p)\, F_{p, \mu\nu}\,
\Lambda^\alpha_{\nu, mno}(\boldsymbol \eta_p)\,.
Assembling :math:`A` costs, per marker, the products of all pairs of basis functions that do not vanish at
:math:`\boldsymbol \eta_p` in each block. This operator is never assembled; with the spline function
:math:`S_h` of the coefficients :math:`s` (see :class:`~struphy.feec.psydac_derham.SplineFunction`),
.. math::
\sum_{\nu, mno} F_{p, \mu\nu}\, \Lambda^\alpha_{\nu, mno}(\boldsymbol \eta_p)\, s_{\nu, mno}
= \big(F_p\, S_h(\boldsymbol \eta_p)\big)_\mu\,,
it is applied as
.. math::
(A s)_{\mu, ijk} = \sum_p \Lambda^\alpha_{\mu, ijk}(\boldsymbol \eta_p)\,
\big(F_p\, S_h(\boldsymbol \eta_p)\big)_\mu\,,
i.e. one evaluation of :math:`S_h` and one deposition per marker (by ``kernel``, called through an
:class:`AccumulatorVector` without filter). The data of :math:`F_p` must be stored in the marker array
beforehand, e.g. by the accumulation kernel of a vector. Like the assembled particle matrices, the operator acts
on the coefficients with boundary conditions and polar extraction: with the extraction and boundary operators
:math:`\mathbb E` and :math:`\mathbb B` of the space,
.. math::
s \mapsto \mathbb B \mathbb E\, A\, \mathbb E^\top \mathbb B^\top s\,.
Parameters
----------
particles : Particles
The markers.
space_id : str
The FE space of the coefficients (``"Hcurl"``, ``"Hdiv"`` or ``"H1vec"``).
kernel : Kernel
The kernel ``(args_markers, args_derham, args_domain, vec1, vec2, vec3, s1, s2, s3, *args)`` that deposits
:math:`A s` into ``vec``, where ``s1, s2, s3`` are the coefficient arrays of the three components of
:math:`s`.
mass_ops : WeightedMassOperators
Weighted mass operators (for the de Rham sequence).
args_domain : DomainArguments
Arguments of the mapping.
args : tuple
Further arguments of the kernel.
symmetry : str
``"symm"`` (:math:`F_p^\top = F_p`) or ``"asym"`` (:math:`F_p^\top = -F_p`), for the transpose.
n_blocks : int
Number of copies of the space in the domain (and codomain). For ``n_blocks > 1``, the operator is the
``n_blocks x n_blocks`` block matrix :math:`(A_{ab})_{a,b}` on the product space, e.g. the pressure coupling
with :math:`A_{ab} = \sum_p v_{p,a} v_{p,b}\, \vec \Lambda_p F_p \vec \Lambda_p^\top`; then ``kernel`` takes
the ``n_blocks`` output vectors followed by the ``n_blocks`` coefficient vectors (three arrays each), and
``symmetry`` refers to the whole block matrix.
"""
[docs]
def __init__(
self,
particles: Particles,
space_id: str,
kernel: Kernel | PyccelKernel,
mass_ops: WeightedMassOperators,
args_domain: DomainArguments,
args: tuple = (),
symmetry: str = "symm",
n_blocks: int = 1,
):
assert symmetry in ("symm", "asym")
assert n_blocks >= 1
derham = mass_ops.derham
form = derham.space_to_form[space_id]
self._E = derham.extraction_ops[form]
self._B = derham.boundary_ops[form]
# the transposes are new operators: built once (dot is called in every iteration of a solver)
self._ET = self._E.T
self._BT = self._B.T
self._acc = AccumulatorVector(particles, space_id, kernel, mass_ops, args_domain, n_vectors=n_blocks)
self._args = args
self._symmetry = symmetry
self._n_blocks = n_blocks
# one full coefficient vector per block (with ghost regions, evaluated at the markers)
self._x_full = [self._E.domain.zeros() for _ in range(n_blocks)]
self._x_tmp = self._B.domain.zeros()
if n_blocks == 1:
self._space = self._B.codomain
else:
self._space = BlockVectorSpace(*(self._B.codomain,) * n_blocks)
self._data = ()
for x_full in self._x_full:
blocks = x_full.blocks if hasattr(x_full, "blocks") else [x_full]
self._data += tuple(blk._data for blk in blocks)
@property
def domain(self) -> VectorSpace:
return self._space
@property
def codomain(self) -> VectorSpace:
return self._space
@property
def dtype(self):
return self._space.dtype
[docs]
def kernels(self) -> tuple[Kernel, ...]:
"""The kernels this operator calls (the kernel that applies the matrix)."""
return self._acc.kernels()
[docs]
def tosparse(self):
raise NotImplementedError("MarkerOperator is matrix-free.")
[docs]
def toarray(self):
raise NotImplementedError("MarkerOperator is matrix-free.")
[docs]
def transpose(self, conjugate=False):
return self if self._symmetry == "symm" else -1.0 * self
[docs]
def dot(self, v: Vector, out: Vector = None) -> Vector:
assert v.space == self._space
v_blocks = [v] if self._n_blocks == 1 else v.blocks
# coefficients of the full space, with ghost regions for the evaluation at the markers
for v_blk, x_full in zip(v_blocks, self._x_full):
self._BT.dot(v_blk, out=self._x_tmp)
self._ET.dot(self._x_tmp, out=x_full)
x_full.update_ghost_regions()
self._acc(*self._data, *self._args)
res = self._acc.vectors
if out is None:
out = self._space.zeros()
if self._n_blocks == 1:
res[0].copy(out=out)
else:
for res_blk, out_blk in zip(res, out.blocks):
res_blk.copy(out=out_blk)
return out
