import logging
import math
import os
import warnings
from abc import ABCMeta, abstractmethod
from typing import TYPE_CHECKING
import cunumpy as xp
import h5py
import numpy as np
from feectools.ddm.partition import factorint
from line_profiler import profile
from maybempi import MPI, SerialComm
from scope_profiler import ProfileManager
from struphy.bsplines.bsplines import quadrature_grid
from struphy.fields_background import equils
from struphy.fields_background.base import (
FluidEquilibrium,
FluidEquilibriumWithB,
NumericalFluidEquilibrium,
)
from struphy.fields_background.equils import set_defaults
from struphy.fields_background.projected_equils import ProjectedFluidEquilibrium
from struphy.geometry.base import Domain
from struphy.geometry.utilities import TransformedPformComponent
from struphy.initial.base import Perturbation
from struphy.io.options import LiteralOptions
from struphy.io.output_handling import DataContainer
from struphy.kernel_arguments.pusher_args_cuda import CudaMarkerArguments
from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments
from struphy.kinetic_background.base import KineticBackground, Maxwellian
from struphy.kinetic_background.maxwellians import CanonicalMaxwellian2D
from struphy.particles.parameters import (
BoundaryParameters,
LoadingParameters,
SavingParameters,
SortingParameters,
WeightsParameters,
)
from struphy.pic import sampling_kernels, sobol_seq
from struphy.pic.kernels.box_and_sort_particles import box_and_sort_particles
from struphy.pic.kernels.box_particles import box_particles
from struphy.pic.kernels.fill_boxes import fill_boxes
from struphy.pic.kernels.sort_and_fill_boxes import sort_and_fill_boxes
from struphy.pic.marker_layout import marker_array, padded_n_markers
from struphy.pic.pushing.kernels.reflect import reflect
from struphy.pic.pushing.kernels.sph_mean_velocity_coeffs import sph_mean_velocity_coeffs
from struphy.pic.pushing.kernels.sph_viscosity_tensor import sph_viscosity_tensor
from struphy.pic.sorting import SortingBoxes
from struphy.pic.sph.kernels.box_based_evaluation_flat import box_based_evaluation_flat
from struphy.pic.sph.kernels.box_based_evaluation_meshgrid import box_based_evaluation_meshgrid
from struphy.pic.sph.kernels.naive_evaluation_flat import naive_evaluation_flat
from struphy.pic.sph.kernels.naive_evaluation_meshgrid import naive_evaluation_meshgrid
from struphy.pic.sph_eval_kernels import (
distance,
)
from struphy.utils import utils
from struphy.utils.clone_config import CloneConfig
from struphy.utils.cuda_arguments import check_mapping_on_device
if TYPE_CHECKING: # importing mpi4py.MPI initializes MPI, which is slow; only needed for annotations
from mpi4py.MPI import Intracomm
logger = logging.getLogger("struphy")
# smallest eta1 of a marker on a domain with a pole at eta1=0 (as in pusher_utilities_kernels)
POLE_ETA1_MIN = 1e-10
ORBIT_POSITIONS = (
(0, "x", "$x$", "physical position x"),
(1, "y", "$y$", "physical position y"),
(2, "z", "$z$", "physical position z"),
)
"""Orbit quantities of the position columns, which post-processing maps to physical coordinates."""
[docs]
class Particles(metaclass=ABCMeta):
r"""
Base class for particle species.
The marker information is stored in a 2D array of shape ``(n_cols, n_markers_max)``, one column per marker
(see :mod:`struphy.pic.marker_layout`).
In ``markers[j, ip]`` the index ``ip`` refers to a specific particle,
the index ``j`` to its attributes; ``n_cols`` is the number of attributes and ``n_markers_max`` the number of markers
(including holes), and the attribute indices are called "columns" below and in the rest of the code.
They are as follows:
* ``0:3``: position in the logical unit cube (:math:`\boldsymbol \eta_p \in [0, 1]^3`)
* ``3:3 + vdim``: velocities
* ``3 + vdim``: (time-dependent) weight :math:`w_k(t)`
* ``4 + vdim``: PDF :math:`s^0 = s^3/\sqrt g` at particle position
* ``5 + vdim``: initial weight :math:`w_0`
* ``6 + vdim <= j < -2``: buffer columns, laid out in consecutive blocks (each block's starting
column and width are given by a pair of attributes/properties):
* :attr:`first_diagnostics_idx` (width :attr:`n_cols_diagnostics`): free columns for
model-specific diagnostics (e.g. canonical momentum, magnetic moment, ...).
* :attr:`first_pusher_idx` (width :attr:`n_cols_pusher` :math:`= 3 + \mathrm{vdim}`): scratch
space for a :class:`~struphy.pic.pushing.pusher.Pusher` call, used to hold the phase space
coordinates at the start of a push (or of a sub-stage, for multi-stage pushers).
* :attr:`first_shift_idx` (width :attr:`n_cols_shift` :math:`= 3`): accumulated shifts in
:math:`\eta`-space due to boundary conditions (e.g. periodic wrap-around), added back onto
the pusher's initial positions when reconstructing a marker's unwrapped trajectory.
* :attr:`residual_idx` (width 1): residual of the current iteration, for pushers that solve a
nonlinear/implicit equation iteratively.
* :attr:`first_free_idx` (width :attr:`n_cols_aux`): general-purpose scratch columns available
to any routine that needs temporary per-marker storage (e.g. field evaluations).
The total number of columns is given by :attr:`n_cols`, i.e. ``first_free_idx + n_cols_aux + 2``
(the ``+ 2`` accounts for the last two columns below).
* ``-2``: number of the sorting box the particle is in
* ``-1``: particle ID
Direct indexing into ``markers`` is rarely needed outside of the ``Particles`` class itself.
Instead, the most commonly used columns are exposed as convenience properties, each returning
(or setting) a 2D array of shape ``(n_mks_loc, ...)`` restricted to the valid markers on the
current process (i.e. ``markers[..., self.valid_mks]``):
* :attr:`positions` (columns ``0:3``): marker positions :math:`\boldsymbol\eta_p`.
* :attr:`velocities` (columns ``3:3 + vdim``): marker velocities.
* :attr:`phasespace_coords` (columns ``0:3 + vdim``): positions and velocities combined.
* :attr:`weights` (column ``3 + vdim``): current weights :math:`w_k(t)`.
* :attr:`sampling_density_values` (column ``4 + vdim``): PDF :math:`s^0` at the particle position.
* :attr:`weights0` (column ``5 + vdim``): initial weights :math:`w_0`.
* :attr:`marker_ids` (column ``-1``): unique particle IDs.
Each of these properties has a matching setter (e.g. ``self.positions = new_positions``) that
validates the shape of the assigned array and writes it back into ``self._markers`` at the
corresponding columns, for the valid markers only.
Parameters
----------
comm_world : Intracomm
World MPI communicator.
clone_config : CloneConfig
Manages the configuration for clone-based (copied grids) parallel processing using MPI.
domain_decomp : tuple
The first entry is a domain_array (see :attr:`~struphy.feec.psydac_derham.Derham.domain_array`) and
the second entry is the number of MPI processes in each direction.
loading_params : LoadingParameters
Parameterts for particle loading.
weights_params : WeightsParameters
Parameters for particle weights.
boundary_params : BoundaryParameters
Parameters for particle boundary conditions.
sorting_params : SortingParameters
Parameters for particle sorting.
saving_params : SavingParameters
Parameters for particle saving.
bufsize : float
Size of buffer (as multiple of total size, default=.25) in markers array.
domain : Domain
Struphy domain object.
equil : FluidEquilibrium
Struphy fluid equilibrium object.
projected_equil : ProjectedFluidEquilibrium
Struphy fluid equilibrium projected into a discrete Derham complex.
background : KineticBackground
Kinetic background.
initial_condition : KineticBackground
Kinetic initial condition.
n_as_volume_form: bool
Whether the number density n is given as a volume form or scalar function (=default).
perturbations : Perturbation | list
Kinetic perturbation parameters.
equation_params : dict
Normalization parameters (epsilon, alpha, ...)
dry_run : bool
If True, only compute the sizing of the marker array (:attr:`n_markers_max`, :attr:`n_cols`, ...)
and return early, without allocating any of the (potentially large) marker/sorting/buffer
arrays. Used by :attr:`nbytes_local` to estimate the memory footprint before actually
allocating the particles, see :meth:`~struphy.models.variables.PICVariable.estimate_mem`.
"""
def __init__(
self,
comm_world: "Intracomm" = None,
clone_config: CloneConfig = None,
domain_decomp: tuple = None,
# mpi_dims_mask: tuple | list = None,
# boxes_per_dim: tuple | list = None,
# box_bufsize: float = 5.0,
n_cols_diagnostics: int = None,
n_cols_aux: int = None,
loading_params: LoadingParameters = None,
weights_params: WeightsParameters = None,
boundary_params: BoundaryParameters = None,
sorting_params: SortingParameters = None,
saving_params: SavingParameters = None,
bufsize: float = 0.25,
domain: Domain = None,
equil: FluidEquilibrium = None,
projected_equil: ProjectedFluidEquilibrium = None,
background: KineticBackground | FluidEquilibrium = None,
initial_condition: KineticBackground = None,
perturbations: dict[str, Perturbation] = None,
n_as_volume_form: bool = False,
equation_params: dict = None,
dry_run: bool = False,
):
self._args_backend = xp.get_backend()
self._clone_config = clone_config
if self.clone_config is None:
self._mpi_comm = comm_world
self._num_clones = 1
self._clone_id = 0
else:
self._mpi_comm = self.clone_config.sub_comm
self._num_clones = self.clone_config.num_clones
self._clone_id = self.clone_config.clone_id
# defaults
if n_cols_diagnostics is None:
n_cols_diagnostics = self.default_n_cols["diagnostics"]
if n_cols_aux is None:
n_cols_aux = self.default_n_cols["aux"]
if loading_params is None:
loading_params = LoadingParameters()
if weights_params is None:
weights_params = WeightsParameters()
if boundary_params is None:
boundary_params = BoundaryParameters()
if sorting_params is None:
sorting_params = SortingParameters()
if saving_params is None:
saving_params = SavingParameters()
# other parameters
self._loading_params = loading_params
self._weights_params = weights_params
self._boundary_params = boundary_params
self._sorting_params = sorting_params
self._saving_params = saving_params
self._domain = domain
self._equil = equil
self._projected_equil = projected_equil
self._equation_params = equation_params
self._n_cols_diagnostics = n_cols_diagnostics
self._n_cols_aux = n_cols_aux
# check for mpi communicator (i.e. sub_comm of clone)
if self.mpi_comm is None:
self._mpi_size = 1
self._mpi_rank = 0
self._Barrier = lambda: None
else:
self._mpi_size = self.mpi_comm.Get_size()
self._mpi_rank = self.mpi_comm.Get_rank()
self._Barrier = self.mpi_comm.Barrier
# domain decomposition (MPI) and cell information
self._boxes_per_dim = self.sorting_params.boxes_per_dim
self._box_bufsize = self.sorting_params.box_bufsize
if domain_decomp is None:
self._domain_array, self._nprocs = self._get_domain_decomp(self.sorting_params.dims_mask)
else:
self._domain_array = domain_decomp[0]
self._nprocs = domain_decomp[1]
# domain_array does not change after setup: one host copy for the host-side bookkeeping
# (Python numbers) and the host kernels, also on the CuPy backend
self._domain_array_host = xp.to_numpy(self._domain_array)
# total number of cells (equal to mpi_size if no grid)
n_cells = int(np.sum(np.prod(self._domain_array_host[:, 2::3], axis=1, dtype=int))) * self.num_clones
# total number of boxes
if self.boxes_per_dim is None:
n_boxes = self.mpi_size * self.num_clones
else:
assert all([nboxes >= nproc for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]), (
f"There must be at least one box {self.boxes_per_dim =} on each process {self.nprocs =} in each direction."
)
assert all([nboxes % nproc == 0 for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]), (
f"Number of boxes {self.boxes_per_dim =} must be divisible by number of processes {self.nprocs =} in each direction."
)
n_boxes = math.prod(self.boxes_per_dim) * self.num_clones
# total number of markers (Np) and particles per cell (ppc)
Np = self.loading_params.Np
ppc = self.loading_params.ppc
ppb = self.loading_params.ppb
if ppb is not None:
self._ppb = ppb
self._Np = int(self.ppb * n_boxes)
self._ppc = self.Np / n_cells
elif ppc is not None:
self._ppc = ppc
self._Np = int(self.ppc * n_cells)
self._ppb = self.Np / n_boxes
elif Np is not None:
self._Np = int(Np)
self._ppc = self.Np / n_cells
self._ppb = self.Np / n_boxes
assert self.Np >= self.mpi_size
# boundary conditions
bc = boundary_params.bc
bc_refill = boundary_params.bc_refill
if bc is None:
bc = ["periodic", "periodic", "periodic"]
for bci in bc:
assert bci in ("remove", "reflect", "periodic", "refill")
if bci == "reflect":
assert domain is not None, "Reflecting boundary conditions require a domain."
if bc_refill is not None:
for bc_refilli in bc_refill:
assert bc_refilli in ("outer", "inner")
self._bc = bc
self._periodic_axes = [axis for axis, b_c in enumerate(bc) if b_c == "periodic"]
self._reflect_axes = [axis for axis, b_c in enumerate(bc) if b_c == "reflect"]
self._remove_axes = [axis for axis, b_c in enumerate(bc) if b_c == "remove"]
# boundary condition type per axis for the per-marker kernel apply_kinetic_bc_marker
# (0: periodic, 1: reflect, 2: remove, 3: handled in Python by apply_kinetic_bc)
self._bc_in_python = bc_refill is not None
bc_codes = {"periodic": 0, "reflect": 1, "remove": 3 if self._bc_in_python else 2}
self._bc_type = xp.array([bc_codes.get(b_c, 3) for b_c in bc], dtype=int)
# pole at eta1=0 (no hole around the axis): markers move through the axis,
# the boundary condition of axis 0 then applies at eta1=1 only
# (see apply_kinetic_bc_marker and update_position_stage in pusher_utilities_kernels)
self._pole = domain is not None and domain.pole
if self._pole:
if domain.params.get("poc", 1) != 1 or domain.params.get("pol_period", 1) != 1:
raise NotImplementedError(
"Markers on a domain with a pole need the full poloidal angle (poc=1 or pol_period=1)."
)
if self._args_backend == "cupy":
raise NotImplementedError("Markers on a domain with a pole are not supported on the CuPy backend yet.")
# create marker array
self._bufsize = bufsize
self._allocate_marker_array(dry_run=dry_run)
if dry_run:
return
bc_sph = boundary_params.bc_sph
if bc_sph is None:
bc_sph = [bci if bci == "periodic" else "mirror" for bci in self.bc]
self._mean_velocity_index = None
for bci in bc_sph:
assert bci in ("periodic", "mirror", "fixed", "noslip")
if bci == "noslip":
if boundary_params.mean_velocity_index is None:
self._mean_velocity_index = (
self.first_free_idx
) # index in marker array where mean velocity for noslip BC is stored
else:
self._mean_velocity_index = boundary_params.mean_velocity_index
self._bc_sph = bc_sph
# initialize sorting boxes
self._initialize_sorting_boxes()
# background
if background is None:
self._background = self.default_background
logger.info(f"Background set to default {self.background = }.")
else:
self._background = background
# background p-form description in [eta, v] (False means 0-form, True means volume form -> divide by det)
if isinstance(self.background, FluidEquilibrium):
self._is_volume_form = (n_as_volume_form, False)
else:
self._is_volume_form = (
n_as_volume_form,
self.background.volume_form,
)
# set background function
self._set_background_function()
self._set_background_coordinates()
# perturbation parameters (needed for fluid background)
self._perturbations = perturbations
# initial condition
if initial_condition is None:
self._initial_condition = self.background
else:
self._initial_condition = initial_condition
# for loading
# if self.loading_params["moments"] is None and not isinstance(self, ParticlesSPH) and isinstance(self.bckgr_params, dict):
self._generate_sampling_moments()
# create buffers for mpi_sort_markers
self._sorting_etas = xp.zeros((3, self.n_markers_max), dtype=float)
self._is_on_proc_domain = xp.zeros((3, self.n_markers_max), dtype=bool)
self._can_stay = xp.zeros(self.n_markers_max, dtype=bool)
self._reqs = [None] * self.mpi_size
self._recvbufs = [None] * self.mpi_size
self._send_to_i = [None] * self.mpi_size
self._send_list = [None] * self.mpi_size
# Domain decomposition is static for the lifetime of this object, so the
# neighbour/non-neighbour split used by _sendrecv_get_destinations only needs
# computing once here rather than on every mpi_sort_markers call.
self._neighbor_ranks, self._non_neighbor_ranks = self._compute_neighbor_ranks()
# post init
self.__post_init__()
# ----------------
# Abstract methods
# ----------------
@property
@abstractmethod
def vdim(self):
"""Dimension of the velocity space."""
pass
@property
@abstractmethod
def coordinate_labels(self) -> tuple[str]:
"""Labels for the coordinates in the phase space.
Length must be 3 + vdim, where the first 3 are the spatial coordinates and the last vdim are the velocity coordinates."""
pass
orbit_quantities: tuple[tuple[int, str, str, str], ...] = ()
"""Marker columns saved in the ``orbits`` post-processing product, as
``(column, name, long_name, description)``. The first three are :data:`ORBIT_POSITIONS`;
the marker index becomes the ``marker`` coordinate of the product."""
@property
@abstractmethod
def mu_idx(self):
"""Index of the column in the marker array where the magnetic moment is stored."""
pass
@property
@abstractmethod
def default_background(cls):
"""The default background (of type Maxwellian)."""
pass
@property
@abstractmethod
def default_n_cols(self):
"Dictionary of the form {'diagnostics': 3, 'aux': 12} for default number of columns."
pass
[docs]
@abstractmethod
def __post_init__(self):
"""Can be used for checks on the constructor arguments and for setting additional attributes in subclasses."""
pass
@property
@abstractmethod
def sampling_density(self):
"""Marker sampling density function :math:`s^\textrm{vol}` as a volume form, see :ref:`monte_carlo`.
Must be normalized to 1. Its coordinates are the coordinates used in Monte-Carlo Integrals
approximated by the particles."""
pass
[docs]
@abstractmethod
def s0(self, eta1, eta2, eta3, *v, flat_eval=False, remove_holes=True):
r"""0-form corresponding to :ref:`~struphy.pic.base.Particles.sampling_density`.
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``."""
pass
# ------------------------------------
# Columns and indexing of marker array
# ------------------------------------
@property
def n_cols_diagnostics(self):
"""Number of columns for storing diagnostics for each marker."""
return self._n_cols_diagnostics
@property
def n_cols_aux(self):
"""Number of auxiliary columns for each marker (e.g. for storing evaluation data)."""
return self._n_cols_aux
@property
def first_diagnostics_idx(self):
"""Starting index for diagnostics columns:
after 3 positions, vdim velocities, weight, s0 and w0."""
return 3 + self.vdim + 3
@property
def first_pusher_idx(self):
"""Starting index for storing initial conditions for a Pusher call."""
return self.first_diagnostics_idx + self.n_cols_diagnostics
@property
def n_cols_pusher(self):
"""Dimension of the phase space (for storing initial conditions for a Pusher call)."""
return 3 + self.vdim
@property
def first_shift_idx(self):
"""First index for storing shifts due to boundary conditions in eta-space."""
return self.first_pusher_idx + self.n_cols_pusher
@property
def n_cols_shift(self):
"""Number of columns for storing shifts due to boundary conditions in eta-space."""
return 3
@property
def residual_idx(self):
"""Column for storing the residual in iterative pushers."""
return self.first_shift_idx + self.n_cols_shift
@property
def first_free_idx(self):
"""First index for storing auxiliary quantities for each particle."""
return self.residual_idx + 1
@property
def n_cols(self):
"""Number of quantities ("columns") per marker, the length of the first axis of the markers array.
The last 2 columns refer to box number and particle ID, respectively."""
return self.first_free_idx + self.n_cols_aux + 2
@property
def n_markers_max(self):
"""Number of marker slots (columns of :attr:`markers`), holes included: the most markers this process can
hold."""
if not hasattr(self, "_n_markers_max"):
input("\nWarning: marker array not yet created, creating now ...")
self._allocate_marker_array()
return self._n_markers_max
@property
def mean_velocity_index(self):
"""Index in marker array where mean velocity for noslip BC is stored."""
return self._mean_velocity_index
@property
def nbytes_local(self) -> int:
"""Estimated local (per-MPI-rank) memory footprint, in bytes, of all marker-related arrays
(markers, sorting buffers, lost-marker container). Only depends on :attr:`n_markers_max` and
:attr:`n_cols`, so it is valid whether or not the arrays were actually allocated
(see the ``dry_run`` argument of :meth:`__init__`)."""
float_size = 8 # dtype=float
bool_size = 1 # dtype=bool
n_markers_max = self.n_markers_max
n_cols = self.n_cols
nbytes = 0
nbytes += n_markers_max * n_cols * float_size # markers
nbytes += n_markers_max * 3 * float_size # sorting_etas (mpi_sort_markers buffer)
nbytes += n_markers_max * 3 * bool_size # is_on_proc_domain
nbytes += n_markers_max * bool_size # can_stay
# holes, ghost_particles, valid_mks, is_outside_right, is_outside_left, is_outside
nbytes += n_markers_max * bool_size * 6
nbytes += int(n_markers_max * 0.5) * 10 * float_size # lost_markers
return int(nbytes)
@property
def kinds(self):
"""Name of the class."""
return self.__class__.__name__
@property
def index(self):
"""Dict holding the column indices referring to specific marker parameters (coordinates)."""
out = {}
out["pos"] = slice(0, 3) # positions
out["vel"] = slice(3, 3 + self.vdim) # velocities
out["coords"] = slice(0, 3 + self.vdim) # phasespace_coords
out["weights"] = 3 + self.vdim # weights
out["s0"] = 4 + self.vdim # sampling density at t=0
out["w0"] = 5 + self.vdim # weights at t=0
out["box"] = -2 # sorting box index
out["ids"] = -1 # marker_inds
out["mu"] = self.mu_idx # magnetic moment
return out
@property
def f_coords_index(self):
"""Dict holding the column indices referring to coords of the distribution fuction."""
return self._f_coords_index
@property
def f_jacobian_coords_index(self):
"""Dict holding the column indices referring to coords of the velocity jacobian determinant of the distribution fuction."""
return self._f_jacobian_coords_index
# ------------------------
# Marker parameter classes
# ------------------------
@property
def loading_params(self) -> LoadingParameters:
"""Parameters for particle loading."""
return self._loading_params
@property
def weights_params(self) -> WeightsParameters:
"""Parameters for particle weights."""
return self._weights_params
@property
def boundary_params(self) -> BoundaryParameters:
"""Parameters for marker loading."""
return self._boundary_params
@property
def sorting_params(self) -> SortingParameters:
"""Parameters for marker sorting."""
return self._sorting_params
@property
def saving_params(self) -> SavingParameters:
"""Parameters for marker/distribution function saving."""
return self._saving_params
@property
def loading(self) -> LiteralOptions.OptsLoading:
"""Type of particle loading."""
return self.loading_params.loading
@property
def spatial(self):
"""Drawing particles uniformly on the unit cube('uniform') or on the disc('disc')"""
return self.loading_params.spatial
@property
def bc(self):
"""List of particle boundary conditions in each direction."""
return self._bc
@property
def bc_refill(self):
"""How to re-enter particles if bc is 'refill'."""
return self.boundary_params.bc_refill
@property
def bc_sph(self):
"""List of boundary conditions for sph evaluation in each direction."""
return self._bc_sph
# -----------------------------------------
# MPI decomposition and clone configuration
# -----------------------------------------
@property
def mpi_comm(self):
"""MPI communicator."""
return self._mpi_comm
@property
def mpi_size(self):
"""Number of MPI processes."""
return self._mpi_size
@property
def mpi_rank(self):
"""Rank of current process."""
return self._mpi_rank
@property
def clone_config(self):
"""Manages the configuration for clone-based (copied grids) parallel processing using MPI."""
return self._clone_config
@property
def num_clones(self):
"""Total number of clones."""
return self._num_clones
@property
def clone_id(self):
"""Clone id of current process."""
return self._clone_id
@property
def domain_array(self):
"""
A 2d array[float] of shape (comm.Get_size(), 9). The row index denotes the process number and
for n=0,1,2:
* domain_array[i, 3*n + 0] holds the LEFT domain boundary of process i in direction eta_(n+1).
* domain_array[i, 3*n + 1] holds the RIGHT domain boundary of process i in direction eta_(n+1).
* domain_array[i, 3*n + 2] holds the number of cells of process i in direction eta_(n+1).
"""
return self._domain_array
@property
def mpi_dims_mask(self):
"""3-list | tuple; True if the dimension is to be used in the domain decomposition (=default for each dimension).
If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed."""
return self.sorting_params.dims_mask
@property
def nprocs(self):
"""Number of MPI processes in each dimension."""
return self._nprocs
# --------------------------------
# Markers and their data structure
# --------------------------------
@property
def Np(self):
"""Total number of markers/particles, from user input."""
return self._Np
@property
def Np_per_clone(self):
"""Array where i-th entry corresponds to the number of loaded particles on clone i.
(This is not necessarily the number of valid markers per clone, see self.n_mks_on_each_clone)."""
return self._Np_per_clone
@property
def ppc(self):
"""Particles per cell (=Np if no grid is present)."""
return self._ppc
@property
def ppb(self):
"""Particles per sorting box."""
return self._ppb
@property
def bufsize(self):
"""Relative size of buffer in markers array."""
return self._bufsize
@property
def n_mks_load(self):
"""Array of number of markers on each process at loading stage"""
return self._n_mks_load
@property
def markers(self):
"""2D array of shape ``(n_cols, n_markers_max)`` holding the marker information, including holes.
The i-th column holds the i-th marker info (see :mod:`struphy.pic.marker_layout`).
===== ============== ======================= ======= ====== ====== ========== === === ===
index | 0 | 1 | 2 | | 3 | ... | 3+(vdim-1)| 3+vdim 4+vdim 5+vdim >=6+vdim ... -2 -1
===== ============== ======================= ======= ====== ====== ========== === === ===
value position (eta) velocities weight s0 w0 other ... box ID
===== ============== ======================= ======= ====== ====== ========== === === ===
The row indices referring to different attributes can be obtained from
:attr:`~struphy.pic.base.Particles.index`.
"""
return self._markers
@property
def holes(self):
"""Array of booleans stating if an entry in the markers array is a hole."""
if not hasattr(self, "_holes"):
self._holes = self.markers[0] == -1.0
return self._holes
@property
def ghost_particles(self):
"""Array of booleans stating if an entry in the markers array is a ghost particle."""
if not hasattr(self, "_ghost_particles"):
self._ghost_particles = self.markers[-1] == -2.0
return self._ghost_particles
@property
def markers_wo_holes(self):
"""Array holding the marker information, excluding holes. The i-th column holds the i-th marker info."""
return self.markers[:, ~self.holes]
@property
def markers_wo_holes_and_ghost(self):
"""Array holding the marker information, excluding holes and ghosts (only valid markers). The i-th column holds the i-th marker info."""
return self.markers[:, self.valid_mks]
@property
def lost_markers(self):
"""Array containing the last infos of removed markers"""
return self._lost_markers
@property
def n_lost_markers(self):
"""Number of removed particles."""
return self._n_lost_markers
@property
def valid_mks(self):
"""Array of booleans stating if an entry in the markers array is a true local particle (not a hole or ghost)."""
if not hasattr(self, "_valid_mks"):
self._valid_mks = ~xp.logical_or(self.holes, self.ghost_particles)
return self._valid_mks
@property
def n_mks_loc(self):
"""Number of valid markers on process (without holes and ghosts)."""
return xp.count_nonzero(self.valid_mks)
@property
def n_mks_on_each_proc(self):
"""Array where i-th entry corresponds to the number of valid markers on i-th process (without holes and ghosts)."""
return self.gather_scalar_in_subcomm_array(self.n_mks_loc)
@property
def n_mks_on_clone(self):
"""Number of valid markers on current clone (without holes and ghosts)."""
return xp.sum(self.n_mks_on_each_proc)
@property
def n_mks_on_each_clone(self):
"""Number of valid markers on current clone (without holes and ghosts)."""
return self.gather_scalar_in_intercomm_array(self.n_mks_on_clone)
@property
def n_mks_global(self):
"""Number of valid markers on current clone (without holes and ghosts)."""
return xp.sum(self.n_mks_on_each_clone)
@property
def positions(self):
"""Array of shape ``(3, n_mks_loc)`` holding the marker positions in logical space. The i-th column holds the i-th marker info."""
return self.markers[self.index["pos"], self.valid_mks]
@positions.setter
def positions(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (3, self.n_mks_loc)
self._markers[self.index["pos"], self.valid_mks] = new
@property
def velocities(self):
"""Array of shape ``(vdim, n_mks_loc)`` holding the marker velocities in logical space. The i-th column holds the i-th marker info."""
return self.markers[self.index["vel"], self.valid_mks]
@velocities.setter
def velocities(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (self.vdim, self.n_mks_loc), f"{self.n_mks_loc =} and {self.vdim =} but {new.shape =}"
self._markers[self.index["vel"], self.valid_mks] = new
@property
def phasespace_coords(self):
"""Array of shape ``(3 + vdim, n_mks_loc)`` holding the marker positions and velocities in logical space. The i-th column holds the i-th marker info."""
return self.markers[self.index["coords"], self.valid_mks]
@phasespace_coords.setter
def phasespace_coords(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (3 + self.vdim, self.n_mks_loc)
self._markers[self.index["coords"], self.valid_mks] = new
@property
def weights(self):
"""Array holding the current marker weights. The i-th entry holds the i-th marker info."""
return self.markers[self.index["weights"], self.valid_mks]
@weights.setter
def weights(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (self.n_mks_loc,)
self._markers[self.index["weights"], self.valid_mks] = new
@property
def sampling_density_values(self):
"""Array holding the current marker 0form sampling density s0. The i-th entry holds the i-th marker info."""
return self.markers[self.index["s0"], self.valid_mks]
@sampling_density_values.setter
def sampling_density_values(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (self.n_mks_loc,)
self._markers[self.index["s0"], self.valid_mks] = new
@property
def weights0(self):
"""Array holding the initial marker weights. The i-th entry holds the i-th marker info."""
return self.markers[self.index["w0"], self.valid_mks]
@weights0.setter
def weights0(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (self.n_mks_loc,)
self._markers[self.index["w0"], self.valid_mks] = new
@property
def marker_ids(self):
"""Array holding the marker id's on the current process."""
return self.markers[self.index["ids"], self.valid_mks]
@marker_ids.setter
def marker_ids(self, new):
assert isinstance(new, xp.ndarray)
assert new.shape == (self.n_mks_loc,)
self._markers[self.index["ids"], self.valid_mks] = new
@property
def f_coords(self):
"""Coordinates of the distribution function."""
return self.markers[self.f_coords_index, self.valid_mks]
@f_coords.setter
def f_coords(self, new):
assert isinstance(new, xp.ndarray)
self.markers[self.f_coords_index, self.valid_mks] = new
@property
def f_jacobian_coords(self):
"""Coordinates of the velocity jacobian determinant of the distribution fuction."""
if isinstance(self.f_jacobian_coords_index, list):
return self.markers[xp.ix_(self.f_jacobian_coords_index, ~self.holes)]
else:
return self.markers[self.f_jacobian_coords_index, ~self.holes]
@f_jacobian_coords.setter
def f_jacobian_coords(self, new):
assert isinstance(new, xp.ndarray)
if isinstance(self.f_jacobian_coords_index, list):
self.markers[
xp.ix_(
self.f_jacobian_coords_index,
~self.holes,
)
] = new
else:
self.markers[self.f_jacobian_coords_index, ~self.holes] = new
@property
def args_markers(self) -> MarkerArguments | CudaMarkerArguments:
"""Arguments for marker kernels, selected when this particle object is created."""
return self._args_markers
[docs]
def kernels(self) -> tuple:
"""The kernels these particles call when they apply their boundary conditions and sort into boxes.
:func:`~struphy.pic.pushing.kernels.reflect.reflect` if an axis has reflecting boundary
conditions, and the box kernels of :mod:`struphy.pic.kernels` if there are sorting boxes. Used by
:meth:`~struphy.simulation.sim.Simulation.compile_cuda_kernels`.
"""
kernels = (reflect,) if self._reflect_axes else ()
if self._sorting_boxes is not None:
kernels += (box_particles, fill_boxes, sort_and_fill_boxes, box_and_sort_particles)
return kernels
# -------------------------------------------
# Initial condition and background -> weights
# -------------------------------------------
@property
def background(self) -> KineticBackground:
"""Kinetic background."""
return self._background
@property
def perturbations(self) -> dict[str, Perturbation]:
"""Kinetic perturbations, keys are the names of moments of the distribution function ("n", "u1", etc.)."""
return self._perturbations
@property
def reject_weights(self):
"""Whether to reect weights below threshold."""
return self.weights_params.reject_weights
@property
def threshold(self):
"""Threshold for rejecting weights."""
return self.weights_params.threshold
@property
def initial_condition(self) -> KineticBackground:
"""Kinetic initial condition"""
return self._initial_condition
@property
def f_init(self):
"""Callable initial condition (background + perturbation).
For kinetic models this is a Maxwellian.
For SPH models this is a :class:`~struphy.fields_background.base.FluidEquilibrium`."""
assert hasattr(self, "_f_init"), AttributeError(
'The method "_set_initial_condition" has not yet been called.',
)
return self._f_init
@property
def u_init(self):
"""Callable initial condition (background + perturbation) for the Cartesian velocity
in SPH models."""
assert hasattr(self, "_u_init"), AttributeError(
'The method "_set_initial_condition" has not yet been called.',
)
return self._u_init
@property
def f0(self) -> Maxwellian:
"""Callable background distribution function, used as the control variate in
:meth:`~struphy.pic.base.Particles.update_weights`."""
assert hasattr(self, "_f0"), AttributeError(
"No background distribution available, please run self._set_background_function()",
)
return self._f0
@property
def is_volume_form(self):
"""Tuple of size 2 for (position, velocity), defining the p-form representation of f_init: True means volume-form, False means 0-form."""
return self._is_volume_form
@property
def control_variate(self):
"""Boolean for whether to use the :ref:`control_var` during time stepping."""
return self.weights_params.control_variate
# ----------------
# Particle sorting
# ----------------
@property
def boxes_per_dim(self):
"""Tuple, number of sorting boxes per dimension."""
return self._boxes_per_dim
@property
def sorting_boxes(self):
"""The :class:`~struphy.pic.sorting.SortingBoxes` instance holding the
sorting-box data structure used by :meth:`~struphy.pic.base.Particles.put_particles_in_boxes`
and :meth:`~struphy.pic.base.Particles.do_sort`."""
if not hasattr(self, "_sorting_boxes"):
self._initialize_sorting_boxes()
return self._sorting_boxes
@property
def tesselation(self):
"""Tesselation of the current process domain."""
return self._tesselation
# -------
# Physics
# -------
@property
def equation_params(self):
"""Parameters appearing in model equation due to Struphy normalization."""
return self._equation_params
@property
def domain(self) -> Domain:
"""From :mod:`struphy.geometry.domains`."""
return self._domain
@property
def equil(self) -> FluidEquilibrium:
"""From :mod:`struphy.fields_background.equils`."""
return self._equil
@property
def projected_equil(self) -> ProjectedFluidEquilibrium:
"""MHD equilibrium projected on 3d Derham sequence with commuting projectors."""
return self._projected_equil
# =================
# =================
# Public methods
# =================
# =================
[docs]
def draw_markers(
self,
sort: bool = True,
):
r"""
Drawing markers
* for PIC: according to the volume density :math:`s^\textrm{vol}_{\textnormal{in}}`
* for SPH: from unity/disc in space and according to the vector-field representation of the fluid velocity
In Struphy, the initial marker distribution :math:`s^\textrm{vol}_{\textnormal{in}}` is always of the form
.. math::
s^\textrm{vol}_{\textnormal{in}}(\eta,v) = n^3(\eta)\, \mathcal M(v)\,,
with :math:`\mathcal M(v)` a multi-variate Gaussian:
.. math::
\mathcal M(v) = \prod_{i=1}^{d_v} \frac{1}{\sqrt{2\pi}\,v_{\mathrm{th},i}}
\exp\left[-\frac{(v_i-u_i)^2}{2 v_{\mathrm{th},i}^2}\right]\,,
where :math:`d_v` stands for the dimension in velocity space, :math:`u_i` are velocity constant shifts
and :math:`v_{\mathrm{th},i}` are constant thermal velocities (standard deviations).
The function :math:`n^3:(0,1)^3 \to \mathbb R^+` is a normalized 3-form on the unit cube,
.. math::
\int_{(0,1)^3} n^3(\eta)\,\textnormal d \eta = 1\,.
The following choices are available in Struphy:
1. Uniform distribution on the unit cube: :math:`n^3(\eta) = 1`
2. Uniform distribution on the disc: :math:`n^3(\eta) = 2\eta_1` (radial coordinate = volume element of square-to-disc mapping)
Velocities are sampled via inverse transform sampling.
In case of Particles6D, velocities are sampled as a Maxwellian in each 3 directions,
.. math::
r_i = \int^{v_i}_{-\infty} \mathcal M(v^\prime_i) \textnormal{d} v^\prime_i = \frac{1}{2}\left[ 1 + \text{erf}\left(\frac{v_i - u_i}{\sqrt{2}v_{\mathrm{th},i}}\right)\right] \,,
where :math:`r_i \in \mathcal R(0,1)` is a uniformly drawn random number in the unit interval. So then
.. math::
v_i = \text{erfinv}(2r_i - 1)\sqrt{2}v_{\mathrm{th},i} + u_i \,.
In case of Particles5Dvperp, parallel velocity is sampled as a Maxwellian and perpendicular particle speed :math:`v_\perp = \sqrt{v_1^2 + v_2^2}`
is sampled as a 2D Maxwellian in polar coordinates,
.. math::
\mathcal{M}(v_1, v_2) \, \textnormal{d} v_1 \textnormal{d} v_2 &= \prod_{i=1}^{2} \frac{1}{\sqrt{2\pi}}\frac{1}{v_{\mathrm{th},i}}
\exp\left[-\frac{(v_i-u_i)^2}{2 v_{\mathrm{th},i}^2}\right] \textnormal{d} v_i\,,
\\
&= \frac{1}{v_\mathrm{th}^2}v_\perp \exp\left[-\frac{(v_\perp-u)^2}{2 v_\mathrm{th}^2}\right] \textnormal{d} v_\perp\,,
\\
&= \mathcal{M}^{\textnormal{pol}}(v_\perp) \, \textnormal{d} v_\perp \,.
Then,
.. math::
r = \int^{v_\perp}_0 \mathcal{M}^{\textnormal{pol}} \textnormal{d} v_\perp = 1 - \exp\left[-\frac{(v_\perp-u)^2}{2 v_\mathrm{th}^2}\right] \,.
So then,
.. math::
v_\perp = \sqrt{- \ln(1-r)}\sqrt{2}v_\mathrm{th} + u \,.
All needed parameters can be set in the parameter file, see :ref:`params_yml`.
On several processes, the markers are always sent to the process of their domain after drawing
(:meth:`mpi_sort_markers`). They are also sorted into boxes if sort is True (default) and sorting
boxes were set up.
Parameters
----------
sort : Bool
Whether to sort the markers into boxes after initial drawing (only if sorting boxes were set up).
"""
import scipy.special as sp
from struphy.pic.particles import Particles5D, Particles5Dvperp, Particles6D, ParticlesSPH
# number of markers on the local process at loading stage
n_mks_load_loc = self.n_mks_load[self.mpi_rank]
# Np_per_clone_loc = self.Np_per_clone[self.clone_id]
# fill holes in markers array with -1 (all holes are at end of array at loading stage); every column can hold
# a loaded marker until they are counted below (restart files can have markers in any column)
self._markers[:, n_mks_load_loc:] = -1.0
self._set_n_markers_used(self.n_markers_max)
# number of holes and markers on process
self.update_holes()
self._update_ghost_particles()
# cumulative sum of number of markers on each process at loading stage.
n_mks_load_cum_sum = xp.cumsum(self.n_mks_load)
Np_per_clone_cum_sum = xp.cumsum(self.Np_per_clone)
_first_marker_id = (Np_per_clone_cum_sum - self.Np_per_clone)[self.clone_id] + (
n_mks_load_cum_sum - self.n_mks_load
)[self._mpi_rank]
logger.debug("\nMARKERS:")
logger.debug(f"{'Np:':<25}{self.Np}")
logger.debug(f"{'ppc:':<25}{self.ppc}")
logger.debug(f"{'ppb:':<25}{self.ppb}")
logger.debug(f"{'bc:':<25}{self.bc}")
logger.debug(f"{'bc_refill:':<25}{self.bc_refill}")
logger.debug(f"{'loading:':<25}{self.loading}")
logger.debug(f"{'type:':<25}{type(self).__name__}")
logger.debug(f"{'control_variate:':<25}{self.control_variate}")
logger.debug(f"{'domain_array[0]:':<25}{self.domain_array[0]}")
logger.debug(f"{'boxes_per_dim:':<25}{self.boxes_per_dim}")
logger.debug(f"{'mpi_dims_mask:':<25}{self.mpi_dims_mask}")
if self.loading == "external":
self._load_external(n_mks_load_loc, n_mks_load_cum_sum)
elif self.loading == "restart":
self._load_restart()
elif self.loading == "tesselation":
self._load_tesselation()
if isinstance(self, ParticlesSPH):
self._set_initial_condition()
self.velocities = xp.array(self.u_init(self.positions))
# set markers ID in last column
self.marker_ids = _first_marker_id + xp.arange(n_mks_load_loc, dtype=float)
else:
logger.debug("\nLoading fresh markers:")
for key, val in self.loading_params.__dict__.items():
logger.debug(f"{key + ' :':<25}{val}")
# 1. standard random number generator (pseudo-random)
if self.loading == "pseudo_random":
# set seed
_seed = self.loading_params.seed
if _seed is not None:
xp.random.seed(_seed)
# counting integers
num_loaded_particles_loc = 0 # number of particles alreday loaded (local)
num_loaded_particles_glob = 0 # number of particles already loaded (each clone)
chunk_size = 10000 # TODO: number of particle chunk
# Total number of markers to draw (sum over all clones)
while num_loaded_particles_glob < int(self.Np):
# Generate a chunk of random particles
num_to_add_glob = min(chunk_size, int(self.Np) - num_loaded_particles_glob)
# one column per particle; drawn row by row, so that the random numbers of a particle are
# consecutive in the stream
temp = xp.random.rand(num_to_add_glob, 3 + self.vdim).T
# inversion method for drawing uniformly on the disc, before the domain check
if self.spatial == "disc":
temp[0] = xp.sqrt(temp[0])
# check which particles are on the current process domain
is_on_proc_domain = self._is_on_domain_of_rank(temp[:3], self.mpi_rank)
valid_idx = xp.nonzero(xp.all(is_on_proc_domain, axis=0))[0]
valid_particles = temp[:, valid_idx]
valid_particles = xp.array_split(valid_particles, self.num_clones, axis=1)[self.clone_id]
num_valid = valid_particles.shape[1]
# Add the valid particles to the phasespace_coords array
self._markers[
: 3 + self.vdim,
num_loaded_particles_loc : num_loaded_particles_loc + num_valid,
] = valid_particles
num_loaded_particles_glob += num_to_add_glob
num_loaded_particles_loc += num_valid
# make sure all particles are loaded
assert self.Np == int(num_loaded_particles_glob), f"{self.Np =}, {int(num_loaded_particles_glob) =}"
# set new n_mks_load and Np_per_clone
self.gather_scalar_in_subcomm_array(num_loaded_particles_loc, out=self.n_mks_load)
self.gather_scalar_in_intercomm_array(int(xp.sum(self.n_mks_load)), out=self.Np_per_clone)
n_mks_load_loc = self.n_mks_load[self.mpi_rank]
n_mks_load_cum_sum = xp.cumsum(self.n_mks_load)
# recompute first marker ID from the actual number of loaded markers
Np_per_clone_cum_sum = xp.cumsum(self.Np_per_clone)
_first_marker_id = (Np_per_clone_cum_sum - self.Np_per_clone)[self.clone_id] + (
n_mks_load_cum_sum - self.n_mks_load
)[self._mpi_rank]
# set new holes in markers array to -1
self._markers[:, num_loaded_particles_loc:] = -1.0
self.update_holes()
# 2. plain sobol numbers with skip of first 1000 numbers
elif self.loading == "sobol_standard":
# generated on the host, then copied to the device (setup only)
self.phasespace_coords = xp.to_cunumpy(
sobol_seq.i4_sobol_generate(
3 + self.vdim,
int(n_mks_load_loc),
1000 + int((n_mks_load_cum_sum - self.n_mks_load)[self._mpi_rank]),
).T
)
# 3. symmetric sobol numbers in all 6 dimensions with skip of first 1000 numbers
elif self.loading == "sobol_antithetic":
assert self.vdim == 3, NotImplementedError(
'"sobol_antithetic" requires vdim=3 at the moment.',
)
# each sobol point yields 64 symmetric markers; round up and truncate the last group
n_sobol = -(-self.n_mks_load // 64)
temp_markers = xp.to_cunumpy(
sobol_seq.i4_sobol_generate(
3 + self.vdim,
int(n_sobol[self._mpi_rank]),
1000 + int((xp.cumsum(n_sobol) - n_sobol)[self._mpi_rank]),
)
)
temp_symmetric = xp.zeros((3 + self.vdim, 64 * n_sobol[self._mpi_rank]), dtype=float)
sampling_kernels.set_particles_symmetric_3d_3v(
temp_markers,
temp_symmetric,
)
self._markers[: 3 + self.vdim, :n_mks_load_loc] = temp_symmetric[:, :n_mks_load_loc]
# 4. Wrong specification
else:
raise ValueError(
"Specified particle loading method does not exist!",
)
# initial velocities - SPH case: v(0) = u(x(0)) for given velocity u(x)
if isinstance(self, ParticlesSPH):
self._set_initial_condition()
self.velocities = xp.array(self.u_init(self.positions))
else:
# inverse transform sampling in velocity space
# Avoid exact 0 or 1 from low-discrepancy sequences: erfinv(±1)
# and log(0) produce infinities or invalid polar velocities.
eps = xp.finfo(float).eps
self._markers[3 : 3 + self.vdim, :n_mks_load_loc] = xp.clip(
self._markers[3 : 3 + self.vdim, :n_mks_load_loc],
eps,
1.0 - eps,
)
u_mean = xp.array(self.loading_params.moments[: self.vdim])
v_th = xp.array(self.loading_params.moments[self.vdim :])
B0 = self.loading_params.B0
# Particles6D: (1d Maxwellian, 1d Maxwellian, 1d Maxwellian)
if isinstance(self, Particles6D):
self.velocities = (
sp.erfinv(
2 * self.velocities - 1,
)
* xp.sqrt(2)
* v_th[:, None]
+ u_mean[:, None]
)
# Particles5D: (1d Maxwellian, muB0-Maxwellian as volume-form)
elif isinstance(self, Particles5D):
self._markers[3, :n_mks_load_loc] = (
sp.erfinv(
2 * self.velocities[0] - 1,
)
* xp.sqrt(2)
* v_th[0]
+ u_mean[0]
)
self._markers[4, :n_mks_load_loc] = -xp.log(1.0 - self.velocities[1]) * v_th[1] ** 2 / B0
# mu is a magnetic moment and must be >= 0.
# A mean shift in this coordinate is not physically consistent.
if abs(float(u_mean[1])) > 0.0:
raise ValueError(
"For Particles5D, the second velocity coordinate is polar "
"(v_perp), so loading_params.moments[1] must be 0.0."
)
# Particles5Dvperp: (1d Maxwellian, polar Maxwellian as volume-form)
elif isinstance(self, Particles5Dvperp):
self._markers[3, :n_mks_load_loc] = (
sp.erfinv(
2 * self.velocities[0] - 1,
)
* xp.sqrt(2)
* v_th[0]
+ u_mean[0]
)
self._markers[4, :n_mks_load_loc] = (
xp.sqrt(
-xp.log(1.0 - self.velocities[1]),
)
* xp.sqrt(2)
* v_th[1]
)
# v_perp is a polar velocity coordinate and must be >= 0.
# A mean shift in this coordinate is not physically consistent.
if abs(float(u_mean[1])) > 0.0:
raise ValueError(
"For Particles5Dvperp, the second velocity coordinate is polar "
"(v_perp), so loading_params.moments[1] must be 0.0."
)
elif self.vdim == 0:
pass
else:
raise NotImplementedError(
"Inverse transform sampling of given vdim is not implemented!",
)
# inversion method for drawing uniformly on the disc (done above for pseudo-random markers)
if self.spatial == "disc":
if self.loading != "pseudo_random":
self._markers[0, :n_mks_load_loc] = xp.sqrt(
self._markers[0, :n_mks_load_loc],
)
else:
assert self.spatial == "uniform", f'Spatial drawing must be "uniform" or "disc", is {self.spatial}.'
self.marker_ids = _first_marker_id + xp.arange(n_mks_load_loc, dtype=float)
# set specific initial condition for some particles
if self.loading_params.specific_markers is not None:
specific_markers = self.loading_params.specific_markers
counter = 0
for i in range(len(specific_markers)):
if i == int(self.markers[-1, counter]):
for j in range(3 + self.vdim):
if specific_markers[i][j] is not None:
self._markers[j, counter] = specific_markers[i][j]
counter += 1
# check if all particle positions are inside the unit cube [0, 1]^3
n_mks_load_loc = self.n_mks_load[self._mpi_rank]
assert xp.all(~self.holes[:n_mks_load_loc])
assert xp.all(self.holes[n_mks_load_loc:])
# keep markers off a pole at eta1=0, where the mapping is singular (holes have eta1=-1)
if self._pole:
eta1 = self._markers[0]
eta1[xp.logical_and(eta1 >= 0.0, eta1 < POLE_ETA1_MIN)] = POLE_ETA1_MIN
# the loaded markers occupy the first columns (fresh loading) or any columns (restart)
self.count_markers_used()
# send the markers to the process of their domain: needed on several processes whether or not there are
# sorting boxes, since accumulations before the first push (e.g. an initial Poisson solve) assume it (#764);
# not for restart loading without boxes (the holes of the restart markers array are not updated yet)
sort_boxes = self._initialized_sorting and sort
if self.mpi_comm is not None and (sort_boxes or (self.mpi_size > 1 and self.loading != "restart")):
logger.info("\nSending the markers to their processes after initial draw")
self.mpi_sort_markers()
if sort_boxes:
logger.info("\nSorting the markers into boxes after initial draw")
self.do_sort()
logger.info("Done.")
elif self.mpi_comm is not None and self.loading in ("sobol_standard", "sobol_antithetic"):
# Sobol markers are drawn on the whole domain on every process (pseudo-random ones only on the
# process domain): send them to their process, otherwise the first accumulation (e.g. the initial
# Poisson solve) writes outside the local arrays and corrupts the heap.
self.mpi_sort_markers()
[docs]
def initialize_weights(
self,
*,
bckgr_params: dict = None,
pert_params: dict = None,
# reject_weights: bool = False,
# threshold: float = 1e-8,
):
r"""
Computes the initial weights
.. math::
w_{k0} := \frac{f^0(t, q_k(t)) }{s^0(t, q_k(t)) } = \frac{f^0(0, q_k(0)) }{s^0(0, q_k(0)) } = \frac{f^0_{\textnormal{in}}(q_{k0}) }{s^0_{\textnormal{in}}(q_{k0}) }
from the initial distribution function :math:`f^0_{\textnormal{in}}` specified in the parmeter file
and from the initial volume density :math:`s^n_{\textnormal{vol}}` specified in :meth:`~struphy.pic.base.Particles.draw_markers`.
Moreover, it sets the corresponding columns for "w0", "s0" and "weights" in the markers array.
If :attr:`~struphy.pic.base.Particles.control_variate` is True, the background :attr:`~struphy.pic.base.Particles.f0` is subtracted.
Parameters
----------
bckgr_params : dict
Kinetic background parameters.
pert_params : dict
Kinetic perturbation parameters for initial condition.
"""
from struphy.pic.particles import ParticlesSPH
if self.loading == "tesselation":
if not self.is_volume_form[0]:
fvol = TransformedPformComponent([self.f_init], "0", "3", domain=self.domain)
else:
fvol = self.f_init
cell_avg = self.tesselation.cell_averages(fvol, n_quad=self.loading_params.n_quad)
self.weights0 = cell_avg.flatten()
else:
assert self.domain is not None, "A domain is needed to initialize weights."
if self.n_mks_loc == 0:
return
# set initial condition
if bckgr_params is not None:
self._bckgr_params = bckgr_params
if pert_params is not None:
self._pert_params = pert_params
if not isinstance(self, ParticlesSPH):
self._set_initial_condition()
# if isinstance(self.f_init, CanonicalMaxwellian):
# self.save_constants_of_motion()
# evaluate initial distribution function
if isinstance(self, ParticlesSPH):
f_init = self.f_init(self.positions)
else:
f_init = self.f_init(*self.f_coords)
# if f_init is vol-form, transform to 0-form
if self.is_volume_form[0]:
f_init /= self.domain.jacobian_det(self.positions)
if self.is_volume_form[1]:
f_init /= self.f_init.velocity_jacobian_det(
*self.f_jacobian_coords,
)
# compute s0 and save at vdim + 4
self.sampling_density_values = self.s0(*self.phasespace_coords, flat_eval=True)
# compute w0 and save at vdim + 5
self.weights0 = f_init / self.sampling_density_values / self.Np
if self.reject_weights:
reject = self.markers[self.index["w0"]] < self.threshold
self._markers[:, reject] = -1.0
self.update_holes()
self._reset_marker_ids()
logger.info(
f"\nWeights < {self.threshold} have been rejected, number of valid markers on process {self.mpi_rank} is {self.n_mks_loc}.",
)
# compute (time-dependent) weights at vdim + 3
if self.control_variate:
self.update_weights()
else:
self.weights = self.weights0
[docs]
@profile
def update_weights(self):
"""
Applies the control variate method, i.e. updates the time-dependent marker weights
according to the algorithm in :ref:`control_var`.
The background :attr:`~struphy.pic.base.Particles.f0` is used for this.
"""
from struphy.pic.particles import ParticlesSPH
if self.n_mks_loc == 0:
return
if isinstance(self, ParticlesSPH):
f0 = self.f0.n0(self.positions)
else:
# in case of CanonicalMaxwellian, evaluate constants_of_motion
# if isinstance(self.f0, CanonicalMaxwellian):
# self.save_constants_of_motion()
f0 = self.f0(*self.f_coords)
# if f_init is vol-form, transform to 0-form
if self.is_volume_form[0]:
f0 /= self.domain.jacobian_det(self.positions)
if self.is_volume_form[1]:
f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords)
self.weights = self.weights0 - f0 / self.sampling_density_values / self.Np
[docs]
@profile
def binning(
self,
components: tuple[bool],
bin_edges: tuple[xp.ndarray],
output_quantity: LiteralOptions.BinningQuantity = "density",
divide_by_jac: bool = True,
):
r"""Computes full-f and delta-f distribution functions via marker binning in logical space.
``xp.histogramdd`` (NumPy or CuPy, on the device of the markers) is used, following the algorithm outlined in :ref:`binning`.
Parameters
----------
components : tuple[bool]
List of length 3 + vdim; an entry is True if the direction in phase space is to be binned.
bin_edges : tuple[array]
List of bin edges (resolution) having the length of True entries in components.
output_quantity : BinningOutput
String literal used to determine weights in binning and the type of output
divide_by_jac : bool
Whether to divide the weights by the Jacobian determinant for binning (default: True).
Returns
-------
f_slice : array-like
The reconstructed full-f distribution function.
df_slice : array-like
The reconstructed delta-f distribution function.
"""
# components is a host list: counted with Python, not with xp (CuPy rejects lists)
assert sum(bool(c) for c in components) == len(bin_edges)
# volume of a bin
bin_vol = 1.0
for be in bin_edges:
bin_vol *= be[1] - be[0]
# extend components list to number of columns of markers array
_n = len(components)
slicing = list(components) + [False] * (self.n_cols - _n)
# determine type of output quantity
# Note: "density" Literal does not have "_"
quantity, *v_axis = output_quantity.rsplit(sep="_", maxsplit=1)
v_axis = [int(char) - 1 for char in "".join(v_axis)] # convert dimension axis to index
# the indices of the valid markers, found once; below, only the rows (quantities) that are used are gathered
# for them, not all quantities of the markers. Everything stays on the device of the markers.
valid = xp.nonzero(self.valid_mks)[0]
if quantity != "density":
velocities = self.markers[self.index["vel"], valid]
# determine histogram weights multiplier
if quantity == "density":
multiplier = 1
elif quantity == "current":
multiplier = velocities[v_axis[0]]
elif quantity == "energy_tensor":
multiplier = velocities[v_axis[0]] * velocities[v_axis[1]]
elif quantity == "heat_flux":
velocity_norm2 = xp.linalg.norm(velocities, axis=0) ** 2
multiplier = velocity_norm2 * velocities[v_axis[0]]
# compute weights of histogram:
_weights0 = self.markers[self.index["w0"], valid] * self.Np * multiplier
_weights = self.markers[self.index["weights"], valid] * self.Np * multiplier
if divide_by_jac:
jacobian_det = self.domain.jacobian_det(self.markers[self.index["pos"], valid], remove_outside=False)
_weights /= jacobian_det
# _weights /= self.velocity_jacobian_det(*self.phasespace_coords)
_weights0 /= jacobian_det
# _weights0 /= self.velocity_jacobian_det(*self.phasespace_coords)
# the binned quantities, one by one (an index list would have to be uploaded to the device first)
sample = xp.stack([self.markers[i, valid] for i, binned in enumerate(slicing) if binned], axis=1)
f_slice = xp.histogramdd(
sample,
bins=bin_edges,
weights=_weights0,
)[0]
df_slice = xp.histogramdd(
sample,
bins=bin_edges,
weights=_weights,
)[0]
f_slice /= self.Np * bin_vol
df_slice /= self.Np * bin_vol
return f_slice, df_slice
[docs]
def show_distribution_function(self, components: list[bool], bin_edges: list[np.ndarray], do_plot=False):
"""
1D and 2D plots of slices of the distribution function via marker binning.
This routine is mainly for de-bugging.
Parameters
----------
components : list[bool]
List of length 3+vdim giving the directions in phase space in which to bin.
Up to two entries can be True, the rest must be False. The True entries correspond to the axes of the binning.
bin_edges : list[np.ndarray]
List of bin edges (resolution) having the length of True entries in components.
do_plot : bool
Whether to show the plot (default: False).
Returns
-------
err : float
Maximum relative error between the binned distribution function and the analytic initial condition.
"""
import matplotlib.pyplot as plt
assert len(components) == 3 + self.vdim, f"components must be of length {3 + self.vdim}, is {len(components)}."
n_dim = np.count_nonzero(components)
assert n_dim == 1 or n_dim == 2, f"Distribution function can only be shown in 1D or 2D slices, not {n_dim}."
f_slice, df_slice = self.binning(components, bin_edges)
bin_centers = [bi[:-1] + (bi[1] - bi[0]) / 2 for bi in bin_edges]
indices = np.nonzero(components)[0]
if n_dim == 1:
plt.plot(bin_centers[0], f_slice, linewidth=2, label="binned f")
i = int(indices[0])
resol = bin_centers[0]
integrate_resol = [0.5, 0.5, 0.5] + [32] * self.vdim
integrate_resol[i] = None
if i < 3:
v_lim = 5
else:
v_lim = bin_edges[0][-1]
f_init, pts, _, _ = self.f_init.reduced_eval(
dim_1=i, v_lim=v_lim, resol=resol, integrate_resol=integrate_resol
)
if do_plot:
plt.plot(pts, f_init, "r--", label="analytic initial condition")
plt.xlabel(self.coordinate_labels[i])
plt.ylabel("f")
plt.legend()
else:
i = int(indices[0])
j = int(indices[1])
if do_plot:
plt.subplot(1, 2, 1)
plt.contourf(bin_centers[0], bin_centers[1], f_slice.T, levels=20)
plt.colorbar()
plt.xlabel(self.coordinate_labels[i])
plt.ylabel(self.coordinate_labels[j])
plt.title("Binned f")
resol = tuple(bin_centers)
integrate_resol = [0.5, 0.5, 0.5] + [100] * self.vdim
integrate_resol[i] = None
integrate_resol[j] = None
v_lim = [5, 5]
if i > 2:
v_lim[0] = bin_edges[0][-1]
if j > 2:
v_lim[1] = bin_edges[1][-1]
v_lim = tuple(v_lim)
f_init, pts1, pts2, _ = self.f_init.reduced_eval(
dim_1=i,
dim_2=j,
v_lim=v_lim,
resol=resol,
integrate_resol=integrate_resol,
)
if do_plot:
plt.subplot(1, 2, 2)
plt.contourf(pts1, pts2, f_init.T, levels=20)
plt.colorbar()
plt.xlabel(self.coordinate_labels[i])
plt.ylabel(self.coordinate_labels[j])
plt.title("Analytic initial condition")
if do_plot:
plt.show()
err = np.max(np.abs(f_init - f_slice)) / np.max(f_init)
return err
[docs]
@profile
@ProfileManager.profile("mpi_sort_markers")
def mpi_sort_markers(
self,
apply_bc: bool = True,
alpha: tuple | list | int | float = 1.0,
do_test: bool = False,
remove_ghost: bool = True,
):
"""
Sorts markers according to MPI domain decomposition.
Markers are sent to the process corresponding to the alpha-weighted position
alpha*markers[0:3] + (1 - alpha)*markers[first_pusher_idx:first_pusher_idx + 3].
Periodic boundary conditions are taken into account
when computing the alpha-weighted position.
Parameters
----------
apply_bc : bool
Whether to apply kinetic boundary conditions before sorting.
alpha : tuple | list | int | float
For i=1,2,3 the sorting is according to alpha[i]*markers[i] + (1 - alpha[i])*markers[first_pusher_idx + i].
If int or float then alpha = (alpha, alpha, alpha). alpha must be between 0 and 1.
do_test : bool
Check if all markers are on the right process after sorting.
remove_ghost : bool
Remove ghost particles before send.
Notes
-----
On the CuPy backend, the markers are exchanged as device buffers with CUDA-aware MPI
(staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The
per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`).
"""
if remove_ghost:
self._remove_ghost_particles()
self._Barrier()
# before sorting, apply kinetic bc
if apply_bc:
self.apply_kinetic_bc()
if isinstance(alpha, int) or isinstance(alpha, float):
alpha = (alpha, alpha, alpha)
# create new markers_to_be_sent array and make corresponding holes in markers array
hole_inds_after_send, send_inds = self._sendrecv_determine_mtbs(alpha=alpha)
# determine where to send markers_to_be_sent
send_info = self._sendrecv_get_destinations(send_inds)
# set new holes in markers array to -1
self._markers[:, send_inds] = -1.0
# transpose send_info
recv_info = self._sendrecv_all_to_all(send_info)
# send and receive markers
self._sendrecv_markers(recv_info, hole_inds_after_send)
# the sent markers are holes now, the received ones filled the lowest holes (host counts, no device read)
n_received = int(np.sum(recv_info)) - int(recv_info[self.mpi_rank])
self._track_markers(n_emptied=int(send_inds.size), n_filled=n_received)
# new holes and new number of holes and markers on process
self.update_holes()
# refresh ghost mask: received markers may land in columns that previously held
# ghost particles. update_holes alone recomputes valid_mks from a stale
# _ghost_particles mask, which would wrongly exclude these incoming real markers.
self._update_ghost_particles()
# check if all markers are on the right process after sorting
if do_test:
# eta = 1 is sorted like eta = 0 in periodic directions (see _sendrecv_determine_mtbs)
etas = self.positions
for axis in self._periodic_axes:
etas[axis] %= 1.0
all_on_right_proc = xp.all(self._is_on_domain_of_rank(etas, self.mpi_rank))
assert all_on_right_proc
# assert self.phasespace_coords.size > 0, f'No particles on process {self.mpi_rank}, please rebalance, aborting ...'
self._Barrier()
[docs]
@profile
@ProfileManager.profile("apply_kinetic_bc")
def apply_kinetic_bc(self, newton=False):
"""
Apply boundary conditions to markers that are outside of the logical unit cube.
Parameters
----------
newton : bool
Whether the shift due to boundary conditions should be computed
for a Newton step or for a strandard (explicit or Picard) step.
"""
# pole at eta1=0: markers with eta1 < 0 have crossed the axis (as in apply_kinetic_bc_marker)
if self._pole:
self._find_outside_particles(0)
crossed = xp.nonzero(self._is_outside_left)[0]
self.markers[0, crossed] *= -1.0
self.markers[1, crossed] = (self.markers[1, crossed] + 0.5) % 1.0
near_pole = xp.logical_and(self.valid_mks, self.markers[0] < POLE_ETA1_MIN)
self.markers[0, near_pole] = POLE_ETA1_MIN
# apply boundary conditions
if self._remove_axes:
# snapshot of the (3, n_markers_max) logical coordinates for the loop over the axes
# (see _find_outside_particles)
self._eta_bc_buf[:, : self.n_markers_used] = self.markers[:3, : self.n_markers_used]
for axis in self._remove_axes:
outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf)
if len(outside_inds) == 0:
continue
if self.bc_refill is not None:
self._particle_refilling()
self._markers[:-1, self._is_outside] = -1.0
n_removed = len(xp.nonzero(self._is_outside)[0])
self._n_lost_markers += n_removed
self._track_markers(n_emptied=n_removed)
# removed markers are holes now and refilled markers have moved: refresh both,
# such that a marker outside on several axes is only counted (and refilled) once
self.update_holes()
self._eta_bc_buf[:, outside_inds] = self.markers[:3, outside_inds]
if self._periodic_axes:
self._eta_bc_buf[:, : self.n_markers_used] = self.markers[:3, : self.n_markers_used]
for axis in self._periodic_axes:
outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf)
if len(outside_inds) == 0:
continue
self.markers[axis, outside_inds] = self.markers[axis, outside_inds] % 1.0
# set shift for alpha-weighted mid-point computation
outside_right_inds = xp.nonzero(self._is_outside_right)[0]
outside_left_inds = xp.nonzero(self._is_outside_left)[0]
if newton:
self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_right_inds] += 1.0
self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_left_inds] += -1.0
else:
self.markers[self.first_pusher_idx + 3 + self.vdim + axis, : self.n_markers_used] = 0.0
self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_right_inds] = 1.0
self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_left_inds] = -1.0
# put all coordinate inside the unit cube (avoid wrong Jacobian evaluations)
outside_inds_per_axis = {}
if self._reflect_axes:
self._eta_bc_buf[:, : self.n_markers_used] = self.markers[:3, : self.n_markers_used]
for axis in self._reflect_axes:
outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf)
self.markers[axis, self._is_outside_left] *= -1.0
self.markers[axis, self._is_outside_right] *= -1.0
self.markers[axis, self._is_outside_right] += 2.0
self.markers[self.first_pusher_idx, self._is_outside] = -1.0
outside_inds_per_axis[axis] = outside_inds
for axis in self._reflect_axes:
if len(outside_inds_per_axis[axis]) == 0:
continue
# only Cartesian velocities are reflected (columns 3:6 hold other data for vdim != 3)
if self.vdim != 3:
continue
# flip velocity
outside_inds = outside_inds_per_axis[axis]
reflect(self.markers, self.domain.args_domain, outside_inds, axis, n_threads=len(outside_inds))
[docs]
def finish_kernel_bc(self, newton=False):
"""Bookkeeping after a pusher kernel that applied the kinetic boundary conditions per marker
(see :func:`~struphy.pic.pushing.pusher_utilities_kernels.apply_kinetic_bc_marker`). Refilling is not done
in the kernel and is applied here by :meth:`apply_kinetic_bc`. Markers removed in the
kernel have become holes; they are counted as lost markers.
Parameters
----------
newton : bool
Whether the shift due to boundary conditions should be computed
for a Newton step or for a standard (explicit or Picard) step.
"""
if self._bc_in_python:
self.apply_kinetic_bc(newton=newton)
self.update_holes()
elif self._remove_axes:
n_holes_before = xp.count_nonzero(self.holes)
self.update_holes()
n_removed = int(xp.count_nonzero(self.holes) - n_holes_before)
self._n_lost_markers += n_removed
self._track_markers(n_emptied=n_removed)
[docs]
def update_holes(self):
"""Recompute the :attr:`~struphy.pic.base.Particles.holes` mask (markers with ``markers[0] == -1``)
and, from it, refresh :attr:`~struphy.pic.base.Particles.valid_mks`.
Must be called after any operation that creates, removes or moves markers
(e.g. sorting, boundary conditions, refilling), since holes are tracked per marker index.
The markers (columns) from :attr:`n_markers_used` on are holes and are not read."""
n = self._n_markers_used
self._holes[:n] = self.markers[0, :n] == -1.0
self._holes[n:] = True
self._update_valid_mks()
@property
def n_markers_used(self) -> int:
"""Number of marker slots (columns of :attr:`markers`, out of :attr:`n_markers_max`) that can hold markers:
every slot from this one on is a hole.
Kernels loop over these markers only (``args_markers.n_markers``), and so do the row operations of the
pushers and the boundary conditions (``markers[j, :n_markers_used]``). It is an upper bound kept on the host,
so that it never waits for the device: markers always fill the lowest holes, so ``k`` incoming markers end
below (occupied slots) + ``k`` (see :meth:`_track_markers`). Loading and box sorting count it exactly
(:meth:`count_markers_used`), and so must code that writes markers into holes itself, outside the methods of
this class.
"""
return self._n_markers_used
def _set_n_markers_used(self, n_markers_used: int):
"""Set :attr:`n_markers_used`, and the number of markers the kernels loop over."""
self._n_markers_used = min(int(n_markers_used), self.n_markers_max)
self._args_markers.n_markers = self._n_markers_used
[docs]
def count_markers_used(self):
"""Set :attr:`n_markers_used` to one past the last column that is not a hole, and count the occupied columns.
Called after loading and after box sorting; call it, before :meth:`update_holes`, after writing markers
into holes outside the methods of this class. Reads two numbers from the device on the CuPy backend."""
occupied = xp.nonzero(self.markers[0] != -1.0)[0]
n_occupied = int(occupied.size)
self._n_occupied_bound = n_occupied
self._set_n_markers_used(int(occupied[-1]) + 1 if n_occupied > 0 else 0)
def _track_markers(self, n_emptied: int = 0, n_filled: int = 0):
"""Update :attr:`n_markers_used` on the host after ``n_emptied`` occupied columns became holes and then
``n_filled`` markers were written into the lowest holes.
At most :attr:`_n_occupied_bound` columns are occupied, so the ``n_filled`` lowest holes lie below
(occupied columns) + ``n_filled``. Holes created without a count (e.g. removed ghost particles) leave the
bound too high, which is safe.
"""
occupied = max(0, self._n_occupied_bound - int(n_emptied))
if n_filled > 0:
self._set_n_markers_used(max(self._n_markers_used, occupied + int(n_filled)))
self._n_occupied_bound = min(self.n_markers_max, occupied + int(n_filled))
[docs]
def set_velocities_comp(self, velocity, comp):
"""Set one or several velocity components to the same constant value, for all valid markers.
Parameters
----------
velocity : float
The constant value to assign to the selected velocity components.
comp : iterable[int]
Velocity components to set (0-based, e.g. 0 for v1, 1 for v2, ...).
"""
for c in comp:
self._markers[3 + c, self.valid_mks] = velocity
[docs]
@profile
@ProfileManager.profile("put_particles_in_boxes")
def put_particles_in_boxes(self):
"""Assign the right box to the particles and the list of the particles to each box.
If sorting_boxes was instantiated with an MPI comm, then the particles in the
neighbouring boxes of neighbouring processes are also communicated (as ghost particles)."""
self._remove_ghost_particles()
max_in_box = box_particles(
self.markers,
self.holes,
self._sorting_boxes.nx,
self._sorting_boxes.ny,
self._sorting_boxes.nz,
self._domain_array_host[self.mpi_rank],
self._sorting_boxes._boxes,
self._sorting_boxes._next_index,
self._sorting_boxes.box_index,
)
self._check_box_capacity(max_in_box)
if self.sorting_boxes.communicate:
self._communicate_boxes()
self._fill_boxes()
self._update_ghost_particles()
# if self.verbose:
# valid_box_ids = xp.nonzero(self._sorting_boxes._boxes[:, 0] != -1)[0]
# logger.info(f"Boxes holding at least one particle: {valid_box_ids}")
# for i in valid_box_ids:
# n_mks_box = xp.count_nonzero(self._sorting_boxes._boxes[i] != -1)
# logger.info(f"Number of markers in box {i} is {n_mks_box}")
[docs]
@profile
@ProfileManager.profile("do_sort")
def do_sort(self, use_numpy_argsort=False):
"""Assign the particles to their sorting boxes and reorder the markers array accordingly,
so that markers in the same box occupy contiguous columns.
The box array :attr:`~struphy.pic.sorting.SortingBoxes.boxes` is rebuilt for the sorted markers,
so it stays valid after the sort.
Parameters
----------
use_numpy_argsort : bool
If True, sort via :func:`numpy.argsort` on the box row; if False (default),
use the kernel :func:`~struphy.pic.kernels.sort_and_fill_boxes.sort_and_fill_boxes`.
"""
boxes = self._sorting_boxes
if use_numpy_argsort:
self.put_particles_in_boxes()
self._sort_boxed_particles_numpy()
self.update_holes()
self._fill_boxes()
elif boxes.communicate:
# SPH: the ghost particles are exchanged between assigning the boxes and sorting
self.put_particles_in_boxes()
sort_and_fill_boxes(
self._markers,
self._holes,
boxes._boxes,
boxes._next_index,
boxes._cumul_next_index,
boxes._swap_line_1,
boxes._swap_line_2,
boxes.box_index,
)
else:
# assign the boxes, sort and rebuild the box array in one kernel call
# (on the CuPy backend one host round trip of the markers, #694)
self._remove_ghost_particles()
max_in_box = box_and_sort_particles(
self._markers,
self._holes,
boxes.nx,
boxes.ny,
boxes.nz,
self._domain_array_host[self.mpi_rank],
boxes._boxes,
boxes._next_index,
boxes._cumul_next_index,
boxes._swap_line_1,
boxes._swap_line_2,
boxes.box_index,
)
self._check_box_capacity(max_in_box)
# The markers have just been reordered. The masks are per marker index,
# so they must be rebuilt before any later use of valid_mks/f_coords;
# the holes may have moved anywhere, so the used columns are counted again.
self.count_markers_used()
self.update_holes()
self._update_ghost_particles()
self._update_valid_mks()
[docs]
def eval_density(
self,
eta1,
eta2,
eta3,
h1,
h2,
h3,
kernel_type="gaussian_1d",
derivative=0,
fast=True,
):
"""Evaluate particle number density (0-form) using an SPH smoothing kernel.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points. Inputs may be 1-D arrays (flat evaluation) or
broadcastable meshgrid arrays; the output will match the shape of `eta1`.
h1, h2, h3 : float
Support radius of the smoothing kernel in each logical dimension.
kernel_type : str, optional
Name of the smoothing kernel (must be a key in `self.ker_dct()`).
derivative : int, optional
Selects whether to evaluate the kernel derivative along a coordinate
direction: 0 (default) returns the scalar density, 1/2/3 returns the
corresponding component of the density gradient with respect to
logical coordinates.
fast : bool, optional
If True, use the box-based neighbor search (faster for many particles);
if False, use the naive all-pairs evaluation (simpler, slower).
Returns
-------
out : xp.ndarray
Estimated number density (or requested derivative component) at the
provided evaluation points. The array uses the same shape as `eta1`
and is returned as a `cunumpy` (`xp`) array.
Notes
-----
This method is a thin wrapper around :meth:`eval_sph` and internally
evaluates the column given by `self.index['weights']` (particle weights).
"""
return self._eval_sph(
eta1,
eta2,
eta3,
self.index["weights"],
kernel_type=kernel_type,
derivative=derivative,
h1=h1,
h2=h2,
h3=h3,
fast=fast,
)
[docs]
def eval_velocity(
self,
eta1,
eta2,
eta3,
h1,
h2,
h3,
kernel_type="gaussian_1d",
derivative=0,
fast=True,
) -> tuple:
"""Estimate mean velocity components using SPH smoothing.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points. May be 1-D arrays or broadcastable meshgrid
arrays; the returned component arrays match the shape of `eta1`.
h1, h2, h3 : float
Support radius of the smoothing kernel in each logical dimension.
kernel_type : str, optional
Name of the smoothing kernel (must be a key in `self.ker_dct()`).
derivative : int, optional
If 0 (default) evaluate the mean velocity; if 1/2/3 return the
corresponding component of the spatial derivative of the velocity.
fast : bool, optional
If True use the box-based neighbor search (faster for many particles);
if False use the naive all-pairs evaluation.
Returns
-------
(v1, v2, v3) : tuple of xp.ndarray
Three arrays containing the estimated velocity components at the
provided evaluation points. Each array has the same shape as `eta1`.
Notes
-----
This method first computes SPH coefficients by calling
the Pyccel kernel `sph_mean_velocity_coeffs` to
assemble mean-velocity coefficients into the markers array, then calls
:meth:`eval_sph` for each velocity component.
"""
first_free_idx = self.args_markers.first_free_idx
self.put_particles_in_boxes()
sph_mean_velocity_coeffs(
xp.array((0.0, 0.0, 0.0)),
xp.array((first_free_idx, first_free_idx + 1, first_free_idx + 2), dtype=int),
self.args_markers,
self.domain.args_domain,
self.sorting_boxes.boxes,
self.sorting_boxes.neighbours,
self.holes,
self.boundary_params.bc_sph[0] == "periodic",
self.boundary_params.bc_sph[1] == "periodic",
self.boundary_params.bc_sph[2] == "periodic",
self.ker_dct()[kernel_type],
h1,
h2,
h3,
)
v1 = self._eval_sph(
eta1,
eta2,
eta3,
first_free_idx,
kernel_type=kernel_type,
derivative=derivative,
h1=h1,
h2=h2,
h3=h3,
fast=fast,
)
v2 = self._eval_sph(
eta1,
eta2,
eta3,
first_free_idx + 1,
kernel_type=kernel_type,
derivative=derivative,
h1=h1,
h2=h2,
h3=h3,
fast=fast,
)
v3 = self._eval_sph(
eta1,
eta2,
eta3,
first_free_idx + 2,
kernel_type=kernel_type,
derivative=derivative,
h1=h1,
h2=h2,
h3=h3,
fast=fast,
)
return v1, v2, v3
[docs]
def eval_div_viscosity(
self,
eta1,
eta2,
eta3,
h1,
h2,
h3,
kernel_type="gaussian_1d",
mu: float = 1.0,
fast=True,
) -> tuple:
"""Compute divergence of the viscous stress (mu * viscosity tensor).
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points where the divergence is evaluated.
h1, h2, h3 : float
Support radius of the smoothing kernel in each logical dimension.
kernel_type : str, optional
Name of the smoothing kernel (must be a key in `self.ker_dct()`).
mu : float, optional
Dynamic viscosity coefficient used in the viscosity kernel.
fast : bool, optional
If True use the box-based neighbor search; if False use naive
evaluation.
Returns
-------
(gamma_x, gamma_y, gamma_z) : tuple of xp.ndarray
Components of the divergence of the viscous stress evaluated at the
provided points. Each array matches the shape of `eta1`.
Notes
-----
The routine populates intermediate marker columns using two Pyccel
kernels: `sph_mean_velocity_coeffs` (mean velocity) and
`sph_viscosity_tensor` (viscosity tensor components in Piola form). It then
evaluates the necessary logical derivatives via :meth:`eval_sph`, sums
contributions and divides by the Jacobian determinant to produce the three
Cartesian divergence components.
"""
first_free_idx = self.args_markers.first_free_idx
self.put_particles_in_boxes()
# 1st kernel
sph_mean_velocity_coeffs(
xp.array((0.0, 0.0, 0.0)),
xp.array((first_free_idx, first_free_idx + 1, first_free_idx + 2), dtype=int),
self.args_markers,
self.domain.args_domain,
self.sorting_boxes.boxes,
self.sorting_boxes.neighbours,
self.holes,
self.boundary_params.bc_sph[0] == "periodic",
self.boundary_params.bc_sph[1] == "periodic",
self.boundary_params.bc_sph[2] == "periodic",
self.ker_dct()[kernel_type],
h1,
h2,
h3,
)
# 2nd kernel
sph_viscosity_tensor(
xp.array((0.0, 0.0, 0.0)),
xp.arange(first_free_idx + 3, first_free_idx + 12, dtype=int),
self.args_markers,
self.domain.args_domain,
self.sorting_boxes.boxes,
self.sorting_boxes.neighbours,
self.holes,
self.boundary_params.bc_sph[0] == "periodic",
self.boundary_params.bc_sph[1] == "periodic",
self.boundary_params.bc_sph[2] == "periodic",
self.ker_dct()[kernel_type],
h1,
h2,
h3,
mu,
)
# grid evaluation
gamma = []
for j in range(3):
gamma += [[]]
for k in range(3):
gamma[-1] += [
self._eval_sph(
eta1,
eta2,
eta3,
first_free_idx + 3 * (j + 1) + k,
kernel_type=kernel_type,
derivative=k + 1,
h1=h1,
h2=h2,
h3=h3,
fast=fast,
)
]
# the stored tensor is in Piola form, divide by the Jacobian determinant at the evaluation points
etas = xp.stack([xp.ravel(eta1), xp.ravel(eta2), xp.ravel(eta3)], axis=0)
det_df = self.domain.jacobian_det(etas, remove_outside=False).reshape(xp.shape(eta1))
gamma_x = (gamma[0][0] + gamma[0][1] + gamma[0][2]) / det_df
gamma_y = (gamma[1][0] + gamma[1][1] + gamma[1][2]) / det_df
gamma_z = (gamma[2][0] + gamma[2][1] + gamma[2][2]) / det_df
return gamma_x, gamma_y, gamma_z
[docs]
@classmethod
def ker_dct(self):
"""Dict mapping the name of each available SPH smoothing kernel (e.g. ``"gaussian_1d"``)
to its integer kernel ID, used in the Pyccel kernels.
Kernel IDs must have three digits, see :meth:`~struphy.pic.smoothing_kernels.smoothing_kernel`."""
return {
"trigonometric_1d": 100,
"gaussian_1d": 110,
"linear_1d": 120,
"trigonometric_2d": 340,
"gaussian_2d": 350,
"linear_2d": 360,
"trigonometric_3d": 670,
"gaussian_3d": 680,
"linear_isotropic_3d": 690,
"linear_3d": 700,
}
[docs]
def gather_scalar_in_subcomm_array(self, scalar: int, out: xp.ndarray = None):
"""Return an array of length sub_comm.size, where the i-th entry corresponds to the value
of the scalar on process i.
Parameters
----------
scalar : int
The scalar value on each process.
out : xp.ndarray
The returned array (optional).
"""
if out is None:
_tmp = xp.zeros(self.mpi_size, dtype=int)
else:
assert out.size == self.mpi_size
_tmp = out
_tmp[self.mpi_rank] = scalar
if self.mpi_comm is not None:
# own entry set beforehand, so the result does not depend on how the communicator fills the buffer
gathered = np.zeros(self.mpi_size, dtype=int)
gathered[self.mpi_rank] = scalar
self.mpi_comm.Allgather(
np.array([scalar], dtype=int),
gathered,
)
_tmp[:] = xp.asarray(gathered)
return _tmp
[docs]
def gather_scalar_in_intercomm_array(self, scalar: int, out: xp.ndarray = None):
"""Return an array of length inter_comm.size, where the i-th entry corresponds to the value
of the scalar on clone i.
Parameters
----------
scalar : int
The scalar value on each clone.
out : xp.ndarray
The returned array (optional).
"""
if out is None:
_tmp = xp.zeros(self.num_clones, dtype=int)
else:
assert out.size == self.num_clones
_tmp = out
_tmp[self.clone_id] = scalar
if self.clone_config is not None:
# own entry set beforehand, so the result does not depend on how the communicator fills the buffer
gathered = np.zeros(self.num_clones, dtype=int)
gathered[self.clone_id] = scalar
self.clone_config.inter_comm.Allgather(
np.array([scalar], dtype=int),
gathered,
)
_tmp[:] = xp.asarray(gathered)
return _tmp
# =================
# =================
# Private methods
# =================
# =================
def _update_valid_mks(self):
"""Refresh :attr:`~struphy.pic.base.Particles.valid_mks`: a row is a valid marker
if and only if it is neither a hole nor a ghost particle."""
self._valid_mks[:] = ~xp.logical_or(self.holes, self.ghost_particles)
def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None):
"""
Compute domain decomposition for mesh-less methods (no Derham object).
Parameters
----------
mpi_dims_mask: list | tuple of bool
True if the dimension is to be used in the domain decomposition (=default for each dimension).
If mpi_dims_mask[i]=False, the i-th dimension will not be decomposed.
Returns
-------
dom_arr : xp.ndarray
A 2d array of shape (#MPI processes, 9). The row index denotes the process rank. The columns are for n=0,1,2:
- arr[3*n + 0, i] holds the LEFT domain boundary of process i in direction eta_(n+1).
- arr[3*n + 1, i] holds the RIGHT domain boundary of process i in direction eta_(n+1).
- arr[3*n + 2, i] holds the number of cells of process i in direction eta_(n+1).
nprocs : tuple
The number of processes in each direction.
"""
if mpi_dims_mask is None:
mpi_dims_mask = [True, True, True]
dom_arr = xp.zeros((self.mpi_size, 9), dtype=float)
# factorize mpi size
factors = factorint(self.mpi_size)
factors_vec = []
for fac, multiplicity in factors.items():
for m in range(multiplicity):
factors_vec += [fac]
# processes in each direction
skip_dims = False
boxes_per_dim = (1, 1, 1)
if self.boxes_per_dim is not None:
boxes_per_dim = self.boxes_per_dim
if not all([bpd == 1 for bpd in self.boxes_per_dim]):
skip_dims = True
nprocs = [1, 1, 1]
for m, fac in enumerate(factors_vec):
mm = m % 3
while (boxes_per_dim[mm] == 1 and skip_dims) or not mpi_dims_mask[mm]:
mm = (mm + 1) % 3
nprocs[mm] *= fac
assert xp.prod(xp.array(nprocs)) == self.mpi_size
# domain decomposition
breaks = [xp.linspace(0.0, 1.0, nproc + 1) for nproc in nprocs]
# fill domain array
for n in range(self.mpi_size):
# determine (ijk box index) corresponding to n (inverse flattening)
i = n // (nprocs[1] * nprocs[2])
nn = n % (nprocs[1] * nprocs[2])
j = nn // nprocs[2]
k = nn % nprocs[2]
dom_arr[n, 0] = breaks[0][i]
dom_arr[n, 1] = breaks[0][i + 1]
dom_arr[n, 2] = 1
dom_arr[n, 3] = breaks[1][j]
dom_arr[n, 4] = breaks[1][j + 1]
dom_arr[n, 5] = 1
dom_arr[n, 6] = breaks[2][k]
dom_arr[n, 7] = breaks[2][k + 1]
dom_arr[n, 8] = 1
return dom_arr, tuple(nprocs)
def _set_background_function(self):
"""Set :attr:`~struphy.pic.base.Particles.f0` to :attr:`~struphy.pic.base.Particles.background`."""
self._f0 = self.background
def _set_background_coordinates(self):
"""Set the default marker-array column indices at which :attr:`f0` and its
Jacobian determinant are evaluated (both default to ``self.index["coords"]``,
the phase-space coordinate columns)."""
self._f_coords_index = self.index["coords"]
self._f_jacobian_coords_index = self.index["coords"]
def _n_mks_load_and_Np_per_clone(self):
"""Return two arrays: 1) an array of sub_comm.size where the i-th entry corresponds to the number of markers drawn on process i,
and 2) an array of size num_clones where the i-th entry corresponds to the number of markers on clone i."""
# number of cells on current process
n_cells_loc = int(np.prod(self._domain_array_host[self.mpi_rank, 2::3], dtype=int))
# array of number of markers on each process at loading stage
if self.clone_config is not None:
_n_cells_clone = int(np.sum(np.prod(self._domain_array_host[:, 2::3], axis=1, dtype=int)))
_n_mks_load_tot = self.clone_config.get_Np_clone(self.Np)
_ppc = _n_mks_load_tot / _n_cells_clone
else:
_n_mks_load_tot = self.Np
_ppc = self.ppc
if self.spatial == "disc":
# markers drawn uniformly on the disc have the density 2*eta1 in eta1
# (more markers on processes at larger eta1)
d = self.domain_array[self.mpi_rank]
n_mks_loc = max(1, int(_n_mks_load_tot * (d[1] ** 2 - d[0] ** 2) * (d[4] - d[3]) * (d[7] - d[6])))
else:
n_mks_loc = int(_ppc * n_cells_loc)
n_mks_load = self.gather_scalar_in_subcomm_array(n_mks_loc)
# add deviation from Np to rank 0
n_mks_load[0] += _n_mks_load_tot - xp.sum(n_mks_load)
# check if all markers are there
assert xp.sum(n_mks_load) == _n_mks_load_tot
# Np on each clone
Np_per_clone = self.gather_scalar_in_intercomm_array(_n_mks_load_tot)
assert xp.sum(Np_per_clone) == self.Np
return n_mks_load, Np_per_clone
def _allocate_marker_array(self, dry_run: bool = False):
"""Create marker array :attr:`~struphy.pic.base.Particles.markers`.
If dry_run is True, only :attr:`n_markers_max` (and :attr:`n_cols`) are computed and no array
is actually allocated; see :attr:`nbytes_local`."""
if not hasattr(self, "_n_mks_load"):
self._n_mks_load, self._Np_per_clone = self._n_mks_load_and_Np_per_clone()
# number of markers on the local process at loading stage
n_mks_load_loc = self.n_mks_load[self._mpi_rank]
bufsize = self.bufsize + 1.0 / xp.sqrt(n_mks_load_loc)
# allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer;
# one column per marker, every row starting on a 256-byte boundary (see struphy.pic.marker_layout)
self._n_markers_max = padded_n_markers(round(float(n_mks_load_loc * (1 + bufsize))))
# Have at least 3 spare places in markers array
assert self.first_free_idx + 2 < self.n_cols - 2, (
f"{self.first_free_idx + 2} is not smaller than {self.n_cols - 2 =}; not enough columns in marker array !!"
)
if dry_run:
return
self._markers = marker_array(self.n_cols, self.n_markers_max)
# allocate auxiliary arrays
self._holes = xp.zeros(self.n_markers_max, dtype=bool)
self._ghost_particles = xp.zeros(self.n_markers_max, dtype=bool)
self._valid_mks = xp.zeros(self.n_markers_max, dtype=bool)
self._is_outside_right = xp.zeros(self.n_markers_max, dtype=bool)
self._is_outside_left = xp.zeros(self.n_markers_max, dtype=bool)
self._is_outside = xp.zeros(self.n_markers_max, dtype=bool)
# scratch copy of markers[:3], refreshed once per apply_kinetic_bc boundary-condition-type
# loop (see there)
self._eta_bc_buf = marker_array(3, self.n_markers_max)
# create array container (3 x positions, vdim x velocities, weight, s0, w0, ID) for removed markers
self._n_lost_markers = 0
self._lost_markers = xp.zeros((10, int(self.n_markers_max * 0.5)), dtype=float)
# arguments for kernels: the pyccel class on the NumPy backend, its CUDA version on the CuPy backend
args_class = CudaMarkerArguments if self._args_backend == "cupy" else MarkerArguments
self._args_markers = args_class(
self.markers,
self.valid_mks,
self.Np,
self.vdim,
self.index["weights"],
self.first_diagnostics_idx,
self.first_pusher_idx,
self.first_shift_idx,
self.residual_idx,
self.first_free_idx,
self.mu_idx,
self._bc_type,
self._pole,
)
# every column can hold a marker until the markers are loaded (see n_markers_used)
self._n_markers_used = self.n_markers_max
self._n_occupied_bound = self.n_markers_max
# on the CuPy backend with several ranks, find out now (collectively, in the setup) whether MPI takes device
# buffers; otherwise the probe runs on the first marker exchange, inside the first time step.
# The probe is collective: all ranks of mpi_comm must allocate their marker arrays together, never one alone.
if xp.get_backend() == "cupy" and self.mpi_size > 1:
self._mpi_cuda_aware()
def _initialize_sorting_boxes(self):
"""Initializes the sorting boxes.
Each MPI process has exactly the same box structure and numbering.
For instance, if boxes_per_dim = (16, 1, 1) and there are 2 MPI processes,
each process would get 8 boxes in the first direction.
Hence boxes_per_dim has to be divisible by the number of ranks in each direction.
"""
from struphy.pic.particles import ParticlesSPH
self._initialized_sorting = False
if self.boxes_per_dim is not None:
# split boxes across MPI processes
nboxes = [nboxes // nproc for nboxes, nproc in zip(self.boxes_per_dim, self.nprocs)]
# check whether this process touches the domain boundary (Python bools, also on the CuPy backend)
x_l, x_r, _, y_l, y_r, _, z_l, z_r, _ = self._domain_array_host[self.mpi_rank].tolist()
is_domain_boundary = {}
is_domain_boundary["x_m"] = x_l == 0.0
is_domain_boundary["x_p"] = x_r == 1.0
is_domain_boundary["y_m"] = y_l == 0.0
is_domain_boundary["y_p"] = y_r == 1.0
is_domain_boundary["z_m"] = z_l == 0.0
is_domain_boundary["z_p"] = z_r == 1.0
self._sorting_boxes = SortingBoxes(
self.markers.shape,
isinstance(self, ParticlesSPH),
nx=nboxes[0],
ny=nboxes[1],
nz=nboxes[2],
bc_sph=self.bc_sph,
is_domain_boundary=is_domain_boundary,
comm=self.mpi_comm,
box_bufsize=self._box_bufsize,
)
if self.sorting_boxes.communicate:
self._get_neighbouring_proc()
self._initialized_sorting = True
else:
self._sorting_boxes = None
def _generate_sampling_moments(self):
"""Automatically determine moments for sampling distribution (Gaussian) from the given background."""
if self.loading_params.moments is None:
self.loading_params.moments = tuple([0.0] * self.vdim + [1.0] * self.vdim)
# TODO: reformulate this function with KineticBackground methods
# ns = []
# us = []
# vths = []
# for fi, params in self.bckgr_params.items():
# if fi[-2] == "_":
# fi_type = fi[:-2]
# else:
# fi_type = fi
# us.append([])
# vths.append([])
# bckgr = getattr(maxwellians, fi_type)
# for key in default_maxw_params:
# if key[0] == "n":
# if key in params:
# ns += [params[key]]
# else:
# ns += [1.0]
# elif key[0] == "u":
# if key in params:
# us[-1] += [params[key]]
# else:
# us[-1] += [0.0]
# elif key[0] == "v":
# if key in params:
# vths[-1] += [params[key]]
# else:
# vths[-1] += [1.0]
# assert len(ns) == len(us) == len(vths)
# ns = xp.array(ns)
# us = xp.array(us)
# vths = xp.array(vths)
# Use the mean of shifts and thermal velocity such that outermost shift+thermal is
# new shift + new thermal
# mean_us = xp.mean(us, axis=0)
# us_ext = us + vths * xp.where(us >= 0, 1, -1)
# us_ext_dist = us_ext - mean_us[None, :]
# new_vths = xp.max(xp.abs(us_ext_dist), axis=0)
# new_moments = []
# new_moments += [*mean_us]
# new_moments += [*new_vths]
# new_moments = [float(moment) for moment in new_moments]
# self.loading_params["moments"] = new_moments
def _set_initial_condition(self):
"""Build :attr:`~struphy.pic.base.Particles.f_init` (and, for SPH particles,
:attr:`~struphy.pic.base.Particles.u_init`).
For non-SPH particles, :attr:`f_init` is simply
:attr:`~struphy.pic.base.Particles.initial_condition`. For SPH particles,
:attr:`f_init`/:attr:`u_init` are built from the fluid equilibrium density/velocity
(:attr:`~struphy.pic.base.Particles.f0`) plus, if given, a single density ('n') or
velocity ('u1') :class:`~struphy.initial.base.Perturbation` from
:attr:`~struphy.pic.base.Particles.perturbations`, transformed to the matching
(0-form / vector-field) representation."""
from struphy.pic.particles import ParticlesSPH
if not isinstance(self, ParticlesSPH):
self._f_init = self.initial_condition
self._u_init = None
else:
assert isinstance(self.f0, FluidEquilibrium)
_density = None
_u1 = None
if self.perturbations is not None:
for (
moment,
pert,
) in self.perturbations.items(): # only one perturbation is taken into account at the moment
assert isinstance(moment, str)
if pert is None:
continue
assert isinstance(pert, Perturbation)
if moment == "n":
if pert.given_in_basis is None:
pert.given_in_basis = "0"
_density = TransformedPformComponent(
pert,
pert.given_in_basis,
"0",
comp=pert.comp,
domain=self.domain,
)
elif moment == "u1":
if pert.given_in_basis is None:
pert.given_in_basis = "v"
_u1 = TransformedPformComponent(
pert,
pert.given_in_basis,
"v",
comp=pert.comp,
domain=self.domain,
)
# self._u_init = lambda e1, e2, e3: self.f0.uv(e1, e2, e3) + _u1(e1, e2, e3)
# TODO: add other velocity components
def _f_init(*etas, flat_eval=False):
if len(etas) == 1:
if _density is None:
out = self.f0.n0(etas[0])
else:
out = self.f0.n0(etas[0]) + _density(*etas[0])
else:
assert len(etas) == 3
E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts(
etas[0],
etas[1],
etas[2],
flat_eval=flat_eval,
)
out = self.f0.n0(E1, E2, E3)
if _density is not None:
out1 = _density(E1, E2, E3)
assert out.shape == out1.shape
out += out1
if flat_eval:
out = xp.squeeze(out)
return out
def _u_init(*etas, flat_eval=False):
if len(etas) == 1:
out = self.f0.uv(etas[0])
if _u1 is not None:
out[0] += _u1(*etas[0])
else:
assert len(etas) == 3
E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts(
etas[0],
etas[1],
etas[2],
flat_eval=flat_eval,
)
out = self.f0.uv(E1, E2, E3)
if _u1 is not None:
out1 = _u1(E1, E2, E3)
assert out[0].shape == out1.shape
out[0] += out1
if flat_eval:
out = xp.squeeze(out)
return out
self._f_init = _f_init
self._u_init = _u_init
def _load_external(
self,
n_mks_load_loc: int,
n_mks_load_cum_sum: xp.ndarray,
):
"""Load markers from external .hdf5 file.
The dataset ``markers`` of the file has the layout of the marker array, shape ``(n_cols, n_markers)``
(one column per marker).
Parameters
----------
n_mks_load_loc: int
Number of markers on the local process at loading stage.
n_mks_load_cum_sum: xp.ndarray
Cumulative sum of number of markers on each process at loading stage.
"""
if self.mpi_rank == 0:
with h5py.File(self.loading_params.dir_external, "r") as file:
logger.info(f"\nLoading markers from file: {file}")
self._markers[:, : n_mks_load_cum_sum[0]] = file["markers"][:, : n_mks_load_cum_sum[0]]
for i in range(1, self._mpi_size):
self._mpi_comm.Send(
file["markers"][:, n_mks_load_cum_sum[i - 1] : n_mks_load_cum_sum[i]],
dest=i,
tag=123,
)
else:
recvbuf = xp.zeros(
(self.n_cols, n_mks_load_loc),
dtype=float,
)
self._mpi_comm.Recv(recvbuf, source=0, tag=123)
self._markers[:, :n_mks_load_loc] = recvbuf
def _load_restart(self):
"""Load markers from restart .hdf5 file."""
# Read struphy state file
state = utils.read_state()
o_path = state["o_path"]
if self.loading_params.dir_particles_abs is None:
data_path = os.path.join(
o_path,
self.loading_params.dir_particles,
)
else:
data_path = self.loading_params.dir_particles_abs
data = DataContainer(data_path, comm=self.mpi_comm)
with h5py.File(data.file_path, "a") as file:
saved = file["restart/" + self.loading_params.restart_key][-1, :, :]
# the number of marker slots n_markers_max is padded (see struphy.pic.marker_layout), so a file written with
# another padding has a different number of slots: copy the common ones; the slots after them are holes, in the
# buffer (set before loading) and in the file
n_common = min(saved.shape[1], self.n_markers_max)
assert np.all(saved[0, n_common:] == -1.0), (
f"the restart file has {saved.shape[1]} markers on rank {self.mpi_rank}, but there are only "
f"{self.n_markers_max} here, and the extra ones are not holes"
)
self._markers[:, :n_common] = saved[:, :n_common]
def _load_tesselation(self, n_quad: int = 1):
"""
Load markers on a grid defined by the center-of-mass points of a tesselation.
Parameters
----------
n_quad: int
Number of quadrature points for the Gauss-Legendre quadrature for cell averages.
"""
self._tesselation = Tesselation(
self.ppb,
comm=self.mpi_comm,
domain_array=self.domain_array,
sorting_boxes=self.sorting_boxes,
)
eta1, eta2, eta3 = self.tesselation.draw_markers()
self._markers[0, : eta1.size] = eta1
self._markers[1, : eta2.size] = eta2
self._markers[2, : eta3.size] = eta3
self._update_valid_mks()
def _reset_marker_ids(self):
"""Reset the marker ids (last column in marker array) according to the current distribution of particles.
The first marker on rank 0 gets the id '0', the last marker on the last rank gets the id 'n_mks_global - 1'."""
n_mks_proc_cumsum = xp.cumsum(self.n_mks_on_each_proc)
n_mks_clone_cumsum = xp.cumsum(self.n_mks_on_each_clone)
first_marker_id = (n_mks_clone_cumsum - self.n_mks_on_each_clone)[self.clone_id] + (
n_mks_proc_cumsum - self.n_mks_on_each_proc
)[self.mpi_rank]
self.marker_ids = first_marker_id + xp.arange(self.n_mks_loc, dtype=int)
def _find_outside_particles(self, axis, eta=None):
"""Find markers whose ``axis``-th logical coordinate lies outside ``[0, 1]``
(holes and ghost particles are excluded), updating
:attr:`_is_outside_left`/:attr:`_is_outside_right`/:attr:`_is_outside` accordingly.
Parameters
----------
axis : int
Row of the markers array (0, 1 or 2) holding the logical coordinate to check.
eta : xp.ndarray[float], optional
Snapshot ``(3, n_markers_max)`` of ``markers[:3]`` (see :meth:`apply_kinetic_bc`).
If ``None``, reads straight from ``markers``.
Returns
-------
outside_inds : xp.ndarray[int]
Indices of the markers that are outside the logical unit cube.
"""
# only the markers (columns) that can hold markers (the others are holes)
n = self.n_markers_used
col = self.markers[axis, :n] if eta is None else eta[axis, :n]
# determine particles outside of the logical unit cube
self._is_outside_right[:n] = col > 1.0
self._is_outside_right[n:] = False
self._is_outside_left[:n] = col < 0.0
self._is_outside_left[n:] = False
self._is_outside_right[self.holes] = False
self._is_outside_right[self.ghost_particles] = False
self._is_outside_left[self.holes] = False
self._is_outside_left[self.ghost_particles] = False
self._is_outside[:] = xp.logical_or(
self._is_outside_right,
self._is_outside_left,
)
# indices or particles that are outside of the logical unit cube
outside_inds = xp.nonzero(self._is_outside)[0]
return outside_inds
def _particle_refilling(self):
r"""
When particles move outside of the domain, refills them.
TODO: Currently only valid for HollowTorus geometry with AdhocTorus equilibrium.
In case of guiding-center orbit, refills particles at the opposite poloidal angle of the same magnetic flux surface.
.. math::
\theta_\text{refill} &= - \theta_\text{loss}
\\
\phi_\text{refill} &= -2 q(r_\text{loss}) \theta_\text{loss}
In case of full orbit, refills particles at the same gyro orbit until their guiding-centers are also outside of the domain.
When their guiding-centers also reach at the boundary, refills them as we did with guiding-center orbit.
"""
for kind in self.bc_refill:
# sorting out particles which are out of the domain
if kind == "inner":
outside_inds = xp.nonzero(self._is_outside_left)[0]
self.markers[0, outside_inds] = 1e-4
r_loss = self.domain.params["a1"]
else:
outside_inds = xp.nonzero(self._is_outside_right)[0]
self.markers[0, outside_inds] = 1 - 1e-4
r_loss = 1.0
if len(outside_inds) == 0:
continue
# in case of Particles6D, do gyro boundary transfer
if self.vdim == 3:
gyro_inside_inds = self._gyro_transfer(outside_inds)
# mark the particle as done for multiple step pushers
self.markers[self.first_pusher_idx, outside_inds[gyro_inside_inds]] = -1.0
self._is_outside[outside_inds[gyro_inside_inds]] = False
# exclude particles whose guiding center positions are still inside.
if len(gyro_inside_inds) > 0:
outside_inds = outside_inds[~gyro_inside_inds]
# do phi boundary transfer = phi_loss - 2*q(r_loss)*theta_loss
self.markers[2, outside_inds] -= 2 * self.equil.q_r(r_loss) * self.markers[1, outside_inds]
# theta_boudary_transfer = - theta_loss
self.markers[1, outside_inds] = 1.0 - self.markers[1, outside_inds]
# mark the particle as done for multiple step pushers
self.markers[self.first_pusher_idx, outside_inds] = -1.0
self._is_outside[outside_inds] = False
def _gyro_transfer(self, outside_inds):
r"""Refills particles at the same gyro orbit.
Their perpendicular velocity directions are also changed accordingly:
First, refills the particles at the other side of the cross point (between gyro circle and domain boundary),
.. math::
\theta_\text{refill} = \theta_\text{gc} - \left(\theta_\text{loss} - \theta_\text{gc} \right) \,.
Then changes the direction of the perpendicular velocity,
.. math::
\vec{v}_{\perp, \text{refill}} = \frac{\vec{\rho}_g}{|\vec{\rho}_g|} \times \vec{b}_0 |\vec{v}_{\perp, \text{loss}}| \,,
where :math:`\vec{\rho}_g = \vec{x}_\text{refill} - \vec{X}_\text{gc}` is the cartesian radial vector.
Parameters
----------
outside_inds : xp.array (int)
An array of indices of particles which are outside of the domain.
Returns
-------
out : xp.array (bool)
An array of indices of particles where its guiding centers are outside of the domain.
"""
# incoming markers must be "Particles6D".
assert self.vdim == 3
# TODO: currently assumes periodic boundary condition along poloidal and toroidal angle
self.markers[1:3, outside_inds] = self.markers[1:3, outside_inds] % 1
v = self.markers[3:6, outside_inds]
# eval cartesian equilibrium magnetic field at the marker positions
assert isinstance(self.equil, FluidEquilibriumWithB), "Gyro transfer function needs a magnetic background."
b_cart, xyz = self.equil.b_cart(self.markers[:, outside_inds])
# calculate magnetic field amplitude and normalized magnetic field
absB0 = xp.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2)
norm_b_cart = b_cart / absB0
# calculate parallel and perpendicular velocities
v_parallel = xp.einsum("ij,ij->j", v, norm_b_cart)
v_perp = xp.cross(norm_b_cart, xp.cross(v, norm_b_cart, axis=0), axis=0)
v_perp_square = xp.sqrt(v_perp[0] ** 2 + v_perp[1] ** 2 + v_perp[2] ** 2)
assert xp.all(xp.isclose(v_perp, v - norm_b_cart * v_parallel))
# calculate Larmor radius
Larmor_r = xp.cross(norm_b_cart, v_perp, axis=0) / absB0 * self.equation_params.epsilon
# transform cartesian coordinates to logical coordinates
# TODO: currently only possible with the geomoetry where its inverse map is defined.
assert hasattr(self.domain, "inverse_map")
xyz -= Larmor_r
gc_etas = self.domain.inverse_map(*xyz, bounded=False)
# gyro transfer
self.markers[1, outside_inds] = (gc_etas[1] - (self.markers[1, outside_inds] - gc_etas[1]) % 1) % 1
new_xyz = self.domain(self.markers[:, outside_inds])
# eval cartesian equilibrium magnetic field at the marker positions
b_cart = self.equil.b_cart(self.markers[:, outside_inds])[0]
# calculate magnetic field amplitude and normalized magnetic field
absB0 = xp.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2)
norm_b_cart = b_cart / absB0
Larmor_r = new_xyz - xyz
Larmor_r /= xp.sqrt(Larmor_r[0] ** 2 + Larmor_r[1] ** 2 + Larmor_r[2] ** 2)
new_v_perp = xp.cross(Larmor_r, norm_b_cart, axis=0) * v_perp_square
self.markers[3:6, outside_inds] = norm_b_cart * v_parallel + new_v_perp
return xp.logical_and(1.0 > gc_etas[0], gc_etas[0] > 0.0)
def _sort_boxed_particles_numpy(self):
"""Sort the particles by box using numpy.argsort."""
sorting_axis = self._sorting_boxes.box_index
if not hasattr(self, "_argsort_array"):
self._argsort_array = xp.zeros(self.n_markers_max, dtype=int)
self._argsort_array[:] = self._markers[sorting_axis].argsort()
self._markers[:] = self._markers[:, self._argsort_array]
def _fill_boxes(self):
"""Assign the particles to boxes from the box row of the markers, after checking that the box
array has enough columns (see :meth:`_check_box_capacity`)."""
max_in_box = fill_boxes(
self.markers,
self.holes,
self._sorting_boxes._boxes,
self._sorting_boxes._next_index,
self._sorting_boxes.box_index,
)
self._check_box_capacity(max_in_box)
def _check_box_capacity(self, max_in_box: int):
"""Raise (or abort under MPI) if the largest number of markers in a box, as returned by the box
kernels (counted on the host, a Python int), does not fit into the box array; the kernels then
leave the box array untouched. This detects load imbalance with respect to the sorting boxes."""
# the last column must stay -1 (terminator for the box readers), hence >=
if max_in_box >= self._sorting_boxes.boxes.shape[1]:
msg = f'Strong load imbalance detected in sorting boxes: \
max number of markers in a box ({max_in_box}) on rank {self.mpi_rank} \
does not fit into the column-size of the box array ({self._sorting_boxes.boxes.shape[1]}, last column is reserved). \
Increase the value of "box_bufsize" in the markers parameters for the next run.'
if self.mpi_comm is None:
raise RuntimeError(msg)
warnings.warn(msg)
self.mpi_comm.Abort()
def _update_ghost_particles(self):
"""Refresh :attr:`~struphy.pic.base.Particles.ghost_particles`: a marker is flagged
as a ghost particle when its ID column (last column) equals -2, the marker set by
:meth:`_prepare_ghost_particles`/:meth:`_sendrecv_markers_boxes` for SPH ghost-box
particles received from a neighbouring process."""
self._ghost_particles[:] = self.markers[-1] == -2.0
self._update_valid_mks()
def _remove_ghost_particles(self):
"""Discard all current ghost particles: turn their marker-array columns into new
holes (so the space can be reused before the next SPH ghost-box update)."""
self._update_ghost_particles()
new_holes = xp.nonzero(self.ghost_particles)[0]
self._markers[:, new_holes] = -1.0
self.update_holes()
def _prepare_ghost_particles(self):
"""Markers for boundary conditions and MPI communication.
Does the following:
1. determine which markers belong to boxes that are at the boundary and put these markers in a new array (e.g. markers_x_m)
2. set their last index to -2 to indicate that they will be "ghost particles" after sending
3. set their new box number (boundary conditions enter here)
4. optional: mirror position for boundary conditions
"""
shifts = self.sorting_boxes.bc_sph_index_shifts
## Faces
# ghost marker arrays
self._markers_x_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m)
self._markers_x_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p)
self._markers_y_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_m)
self._markers_y_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_p)
self._markers_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_z_m)
self._markers_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_z_p)
# Put last index to -2 to indicate that they are ghosts on the new process
self._markers_x_m[-1] = -2.0
self._markers_x_p[-1] = -2.0
self._markers_y_m[-1] = -2.0
self._markers_y_p[-1] = -2.0
self._markers_z_m[-1] = -2.0
self._markers_z_p[-1] = -2.0
# Adjust box number
self._markers_x_m[self._sorting_boxes.box_index] += shifts["x_m"]
self._markers_x_p[self._sorting_boxes.box_index] -= shifts["x_p"]
self._markers_y_m[self._sorting_boxes.box_index] += shifts["y_m"]
self._markers_y_p[self._sorting_boxes.box_index] -= shifts["y_p"]
self._markers_z_m[self._sorting_boxes.box_index] += shifts["z_m"]
self._markers_z_p[self._sorting_boxes.box_index] -= shifts["z_p"]
# Mirror position for boundary condition
if self.bc_sph[0] in ("mirror", "fixed", "noslip"):
self._mirror_particles(
"_markers_x_m",
"_markers_x_p",
is_domain_boundary=self.sorting_boxes.is_domain_boundary,
mean_velocity_index=self.mean_velocity_index,
)
if self.bc_sph[1] in ("mirror", "fixed", "noslip"):
self._mirror_particles(
"_markers_y_m",
"_markers_y_p",
is_domain_boundary=self.sorting_boxes.is_domain_boundary,
mean_velocity_index=self.mean_velocity_index,
)
if self.bc_sph[2] in ("mirror", "fixed", "noslip"):
self._mirror_particles(
"_markers_z_m",
"_markers_z_p",
is_domain_boundary=self.sorting_boxes.is_domain_boundary,
mean_velocity_index=self.mean_velocity_index,
)
## Edges x-y
# ghost marker arrays
self._markers_x_m_y_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m)
self._markers_x_m_y_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p)
self._markers_x_p_y_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m)
self._markers_x_p_y_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p)
# Put last index to -2 to indicate that they are ghosts on the new process
self._markers_x_m_y_m[-1] = -2.0
self._markers_x_m_y_p[-1] = -2.0
self._markers_x_p_y_m[-1] = -2.0
self._markers_x_p_y_p[-1] = -2.0
# Adjust box number
self._markers_x_m_y_m[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"]
self._markers_x_m_y_p[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"]
self._markers_x_p_y_m[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"]
self._markers_x_p_y_p[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"]
# Mirror position for boundary condition
if self.bc_sph[0] in ("mirror", "fixed", "noslip") or self.bc_sph[1] in ("mirror", "fixed", "noslip"):
self._mirror_particles(
"_markers_x_m_y_m",
"_markers_x_m_y_p",
"_markers_x_p_y_m",
"_markers_x_p_y_p",
is_domain_boundary=self.sorting_boxes.is_domain_boundary,
mean_velocity_index=self.mean_velocity_index,
)
## Edges x-z
# ghost marker arrays
self._markers_x_m_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_z_m)
self._markers_x_m_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_z_p)
self._markers_x_p_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_z_m)
self._markers_x_p_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_z_p)
# Put last index to -2 to indicate that they are ghosts on the new process
self._markers_x_m_z_m[-1] = -2.0
self._markers_x_m_z_p[-1] = -2.0
self._markers_x_p_z_m[-1] = -2.0
self._markers_x_p_z_p[-1] = -2.0
# Adjust box number
self._markers_x_m_z_m[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["z_m"]
self._markers_x_m_z_p[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["z_p"]
self._markers_x_p_z_m[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["z_m"]
self._markers_x_p_z_p[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["z_p"]
# Mirror position for boundary condition
if self.bc_sph[0] in ("mirror", "fixed", "noslip") or self.bc_sph[2] in ("mirror", "fixed", "noslip"):
self._mirror_particles(
"_markers_x_m_z_m",
"_markers_x_m_z_p",
"_markers_x_p_z_m",
"_markers_x_p_z_p",
is_domain_boundary=self.sorting_boxes.is_domain_boundary,
mean_velocity_index=self.mean_velocity_index,
)
## Edges y-z
# ghost marker arrays
self._markers_y_m_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_m_z_m)
self._markers_y_m_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_m_z_p)
self._markers_y_p_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_p_z_m)
self._markers_y_p_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_p_z_p)
# Put last index to -2 to indicate that they are ghosts on the new process
self._markers_y_m_z_m[-1] = -2.0
self._markers_y_m_z_p[-1] = -2.0
self._markers_y_p_z_m[-1] = -2.0
self._markers_y_p_z_p[-1] = -2.0
# Adjust box number
self._markers_y_m_z_m[self._sorting_boxes.box_index] += shifts["y_m"] + shifts["z_m"]
self._markers_y_m_z_p[self._sorting_boxes.box_index] += shifts["y_m"] - shifts["z_p"]
self._markers_y_p_z_m[self._sorting_boxes.box_index] += -shifts["y_p"] + shifts["z_m"]
self._markers_y_p_z_p[self._sorting_boxes.box_index] += -shifts["y_p"] - shifts["z_p"]
# Mirror position for boundary condition
if self.bc_sph[1] in ("mirror", "fixed", "noslip") or self.bc_sph[2] in ("mirror", "fixed", "noslip"):
self._mirror_particles(
"_markers_y_m_z_m",
"_markers_y_m_z_p",
"_markers_y_p_z_m",
"_markers_y_p_z_p",
is_domain_boundary=self.sorting_boxes.is_domain_boundary,
mean_velocity_index=self.mean_velocity_index,
)
## Corners
# ghost marker arrays
self._markers_x_m_y_m_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m_z_m)
self._markers_x_m_y_m_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_m_z_p)
self._markers_x_m_y_p_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p_z_m)
self._markers_x_m_y_p_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_m_y_p_z_p)
self._markers_x_p_y_m_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m_z_m)
self._markers_x_p_y_m_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_m_z_p)
self._markers_x_p_y_p_z_m = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p_z_m)
self._markers_x_p_y_p_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p_z_p)
# Put last index to -2 to indicate that they are ghosts on the new process
self._markers_x_m_y_m_z_m[-1] = -2.0
self._markers_x_m_y_m_z_p[-1] = -2.0
self._markers_x_m_y_p_z_m[-1] = -2.0
self._markers_x_m_y_p_z_p[-1] = -2.0
self._markers_x_p_y_m_z_m[-1] = -2.0
self._markers_x_p_y_m_z_p[-1] = -2.0
self._markers_x_p_y_p_z_m[-1] = -2.0
self._markers_x_p_y_p_z_p[-1] = -2.0
# Adjust box number
self._markers_x_m_y_m_z_m[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] + shifts["z_m"]
self._markers_x_m_y_m_z_p[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] - shifts["z_p"]
self._markers_x_m_y_p_z_m[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] + shifts["z_m"]
self._markers_x_m_y_p_z_p[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] - shifts["z_p"]
self._markers_x_p_y_m_z_m[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] + shifts["z_m"]
self._markers_x_p_y_m_z_p[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] - shifts["z_p"]
self._markers_x_p_y_p_z_m[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] + shifts["z_m"]
self._markers_x_p_y_p_z_p[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] - shifts["z_p"]
# Mirror position for boundary condition
if any([bci in ("mirror", "fixed", "noslip") for bci in self.bc_sph]):
self._mirror_particles(
"_markers_x_m_y_m_z_m",
"_markers_x_m_y_m_z_p",
"_markers_x_m_y_p_z_m",
"_markers_x_m_y_p_z_p",
"_markers_x_p_y_m_z_m",
"_markers_x_p_y_m_z_p",
"_markers_x_p_y_p_z_m",
"_markers_x_p_y_p_z_p",
is_domain_boundary=self.sorting_boxes.is_domain_boundary,
mean_velocity_index=self.mean_velocity_index,
)
def _mirror_particles(
self, *marker_array_names, is_domain_boundary: dict | None = None, mean_velocity_index: int | None = None
):
"""
Mirror the positions and velocities of the particles in the ghost marker arrays for the boundary conditions.
For "mirror" boundary condition, the positions are mirrored and the velocities are unchanged.
For "fixed" boundary condition, the positions are mirrored and the velocities are set to zero (or to the value of f_init if provided).
For "noslip" boundary condition, the positions are mirrored and the velocities are inverted to have zero velocity at the boundary.
Parameters
----------
marker_array_names : str
The names of the marker arrays to be mirrored (e.g. "_markers_x_m", "_markers_x_p", etc.).
is_domain_boundary : dict
A dictionary indicating whether the boundary condition is applied at the domain boundary (e.g. {"x_m": True, "x_p": True, "y_m": True, "y_p": True, "z_m": True, "z_p": True}).
mean_velocity_index : int, optional
The index of the mean velocity in the marker array (if applicable), by default None.
"""
self._fixed_markers_set = {}
for arr_name in marker_array_names:
assert isinstance(arr_name, str)
arr = getattr(self, arr_name)
if arr.size == 0:
continue
# x-direction
if self.bc_sph[0] in ("mirror", "fixed", "noslip"):
if "x_m" in arr_name and is_domain_boundary["x_m"]:
arr[0] *= -1.0
if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set:
boundary_values = self.f_init(
*arr[:3],
flat_eval=True,
) # evaluation outside of the unit cube - maybe not working for all f_init!
arr[self.index["weights"]] = (
-boundary_values
/ self.s0(
*arr[:3],
flat_eval=True,
remove_holes=False,
)
/ self.Np
) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right
self._fixed_markers_set[arr_name] = True
elif self.bc_sph[0] == "noslip":
# invert the velocities to have zero velocity at the boundary
arr[3] *= -1.0
arr[4] *= -1.0
arr[5] *= -1.0
if mean_velocity_index is not None:
arr[mean_velocity_index] *= -1.0
arr[mean_velocity_index + 1] *= -1.0
arr[mean_velocity_index + 2] *= -1.0
elif "x_p" in arr_name and is_domain_boundary["x_p"]:
arr[0] = 2.0 - arr[0]
if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set:
boundary_values = self.f_init(
*arr[:3],
flat_eval=True,
) # evaluation outside of the unit cube - maybe not working for all f_init!
arr[self.index["weights"]] = (
-boundary_values
/ self.s0(
*arr[:3],
flat_eval=True,
remove_holes=False,
)
/ self.Np
) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right
self._fixed_markers_set[arr_name] = True
elif self.bc_sph[0] == "noslip":
# invert the velocities to have zero velocity at the boundary
arr[3] *= -1.0
arr[4] *= -1.0
arr[5] *= -1.0
if mean_velocity_index is not None:
arr[mean_velocity_index] *= -1.0
arr[mean_velocity_index + 1] *= -1.0
arr[mean_velocity_index + 2] *= -1.0
# y-direction
if self.bc_sph[1] in ("mirror", "fixed", "noslip"):
if "y_m" in arr_name and is_domain_boundary["y_m"]:
arr[1] *= -1.0
if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set:
boundary_values = self.f_init(
*arr[:3],
flat_eval=True,
) # evaluation outside of the unit cube - maybe not working for all f_init!
arr[self.index["weights"]] = (
-boundary_values
/ self.s0(
*arr[:3],
flat_eval=True,
remove_holes=False,
)
/ self.Np
) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right
self._fixed_markers_set[arr_name] = True
elif self.bc_sph[1] == "noslip":
# invert the velocities to have zero velocity at the boundary
arr[3] *= -1.0
arr[4] *= -1.0
arr[5] *= -1.0
if mean_velocity_index is not None:
arr[mean_velocity_index] *= -1.0
arr[mean_velocity_index + 1] *= -1.0
arr[mean_velocity_index + 2] *= -1.0
elif "y_p" in arr_name and is_domain_boundary["y_p"]:
arr[1] = 2.0 - arr[1]
if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set:
boundary_values = self.f_init(
*arr[:3],
flat_eval=True,
) # evaluation outside of the unit cube - maybe not working for all f_init!
arr[self.index["weights"]] = (
-boundary_values
/ self.s0(
*arr[:3],
flat_eval=True,
remove_holes=False,
)
/ self.Np
) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right
self._fixed_markers_set[arr_name] = True
elif self.bc_sph[1] == "noslip":
# invert the velocities to have zero velocity at the boundary
arr[3] *= -1.0
arr[4] *= -1.0
arr[5] *= -1.0
if mean_velocity_index is not None:
arr[mean_velocity_index] *= -1.0
arr[mean_velocity_index + 1] *= -1.0
arr[mean_velocity_index + 2] *= -1.0
# z-direction
if self.bc_sph[2] in ("mirror", "fixed", "noslip"):
if "z_m" in arr_name and is_domain_boundary["z_m"]:
arr[2] *= -1.0
if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set:
boundary_values = self.f_init(
*arr[:3],
flat_eval=True,
) # evaluation outside of the unit cube - maybe not working for all f_init!
arr[self.index["weights"]] = (
-boundary_values
/ self.s0(
*arr[:3],
flat_eval=True,
remove_holes=False,
)
/ self.Np
) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right
self._fixed_markers_set[arr_name] = True
elif self.bc_sph[2] == "noslip":
# invert the velocities to have zero velocity at the boundary
arr[3] *= -1.0
arr[4] *= -1.0
arr[5] *= -1.0
if mean_velocity_index is not None:
arr[mean_velocity_index] *= -1.0
arr[mean_velocity_index + 1] *= -1.0
arr[mean_velocity_index + 2] *= -1.0
elif "z_p" in arr_name and is_domain_boundary["z_p"]:
arr[2] = 2.0 - arr[2]
if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set:
boundary_values = self.f_init(
*arr[:3],
flat_eval=True,
) # evaluation outside of the unit cube - maybe not working for all f_init!
arr[self.index["weights"]] = (
-boundary_values
/ self.s0(
*arr[:3],
flat_eval=True,
remove_holes=False,
)
/ self.Np
) # clarify in case of tesselation: multiple by tile volume (=1/Np) to get the integral value right
self._fixed_markers_set[arr_name] = True
elif self.bc_sph[2] == "noslip":
# invert the velocities to have zero velocity at the boundary
arr[3] *= -1.0
arr[4] *= -1.0
arr[5] *= -1.0
if mean_velocity_index is not None:
arr[mean_velocity_index] *= -1.0
arr[mean_velocity_index + 1] *= -1.0
arr[mean_velocity_index + 2] *= -1.0
def _determine_markers_in_box(self, list_boxes):
"""Gather the markers currently sorted into any of the given boxes into a new array
(used to collect the particles on a domain/process boundary before turning them
into ghost particles).
Parameters
----------
list_boxes : list[int]
Flat box indices (as computed by
:func:`~struphy.pic.sorting_kernels.flatten_index`) whose particles are collected.
Returns
-------
markers_in_box : xp.ndarray
Copy of the marker-array columns (markers) belonging to any of ``list_boxes``.
"""
indices = []
for i in list_boxes:
indices += list(self._sorting_boxes._boxes[i][self._sorting_boxes._boxes[i] != -1])
indices = xp.array(indices, dtype=int)
# C-contiguous, as the receive buffers (n_cols, n) of _sendrecv_markers_boxes: selecting columns gives a
# Fortran-ordered copy, which MPI would send marker by marker
markers_in_box = xp.ascontiguousarray(self.markers[:, indices])
return markers_in_box
def _get_destinations_box(self):
"""Route the ghost markers prepared by :meth:`_prepare_ghost_particles` (one array
per face/edge/corner) to the neighbouring process on that side (found earlier by
:meth:`_get_neighbouring_proc`), accumulating, per destination rank, the number of
markers to send (:attr:`_send_info_box`) and the markers themselves
(:attr:`_send_list_box`, used by :meth:`_self_communication_boxes` and
:meth:`_sendrecv_markers_boxes`)."""
self._send_info_box = xp.zeros(self.mpi_size, dtype=int)
self._send_list_box = [xp.zeros((self.n_cols, 0))] * self.mpi_size
# Faces
# if self._x_m_proc is not None:
self._send_info_box[self._x_m_proc] += self._markers_x_m.shape[1]
self._send_list_box[self._x_m_proc] = xp.concatenate(
(self._send_list_box[self._x_m_proc], self._markers_x_m), axis=1
)
# if self._x_p_proc is not None:
self._send_info_box[self._x_p_proc] += self._markers_x_p.shape[1]
self._send_list_box[self._x_p_proc] = xp.concatenate(
(self._send_list_box[self._x_p_proc], self._markers_x_p), axis=1
)
# if self._y_m_proc is not None:
self._send_info_box[self._y_m_proc] += self._markers_y_m.shape[1]
self._send_list_box[self._y_m_proc] = xp.concatenate(
(self._send_list_box[self._y_m_proc], self._markers_y_m), axis=1
)
# if self._y_p_proc is not None:
self._send_info_box[self._y_p_proc] += self._markers_y_p.shape[1]
self._send_list_box[self._y_p_proc] = xp.concatenate(
(self._send_list_box[self._y_p_proc], self._markers_y_p), axis=1
)
# if self._z_m_proc is not None:
self._send_info_box[self._z_m_proc] += self._markers_z_m.shape[1]
self._send_list_box[self._z_m_proc] = xp.concatenate(
(self._send_list_box[self._z_m_proc], self._markers_z_m), axis=1
)
# if self._z_p_proc is not None:
self._send_info_box[self._z_p_proc] += self._markers_z_p.shape[1]
self._send_list_box[self._z_p_proc] = xp.concatenate(
(self._send_list_box[self._z_p_proc], self._markers_z_p), axis=1
)
# x-y edges
# if self._x_m_y_m_proc is not None:
self._send_info_box[self._x_m_y_m_proc] += self._markers_x_m_y_m.shape[1]
self._send_list_box[self._x_m_y_m_proc] = xp.concatenate(
(self._send_list_box[self._x_m_y_m_proc], self._markers_x_m_y_m),
axis=1,
)
# if self._x_m_y_p_proc is not None:
self._send_info_box[self._x_m_y_p_proc] += self._markers_x_m_y_p.shape[1]
self._send_list_box[self._x_m_y_p_proc] = xp.concatenate(
(self._send_list_box[self._x_m_y_p_proc], self._markers_x_m_y_p),
axis=1,
)
# if self._x_p_y_m_proc is not None:
self._send_info_box[self._x_p_y_m_proc] += self._markers_x_p_y_m.shape[1]
self._send_list_box[self._x_p_y_m_proc] = xp.concatenate(
(self._send_list_box[self._x_p_y_m_proc], self._markers_x_p_y_m),
axis=1,
)
# if self._x_p_y_p_proc is not None:
self._send_info_box[self._x_p_y_p_proc] += self._markers_x_p_y_p.shape[1]
self._send_list_box[self._x_p_y_p_proc] = xp.concatenate(
(self._send_list_box[self._x_p_y_p_proc], self._markers_x_p_y_p),
axis=1,
)
# x-z edges
# if self._x_m_z_m_proc is not None:
self._send_info_box[self._x_m_z_m_proc] += self._markers_x_m_z_m.shape[1]
self._send_list_box[self._x_m_z_m_proc] = xp.concatenate(
(self._send_list_box[self._x_m_z_m_proc], self._markers_x_m_z_m),
axis=1,
)
# if self._x_m_z_p_proc is not None:
self._send_info_box[self._x_m_z_p_proc] += self._markers_x_m_z_p.shape[1]
self._send_list_box[self._x_m_z_p_proc] = xp.concatenate(
(self._send_list_box[self._x_m_z_p_proc], self._markers_x_m_z_p),
axis=1,
)
# if self._x_p_z_m_proc is not None:
self._send_info_box[self._x_p_z_m_proc] += self._markers_x_p_z_m.shape[1]
self._send_list_box[self._x_p_z_m_proc] = xp.concatenate(
(self._send_list_box[self._x_p_z_m_proc], self._markers_x_p_z_m),
axis=1,
)
# if self._x_p_z_p_proc is not None:
self._send_info_box[self._x_p_z_p_proc] += self._markers_x_p_z_p.shape[1]
self._send_list_box[self._x_p_z_p_proc] = xp.concatenate(
(self._send_list_box[self._x_p_z_p_proc], self._markers_x_p_z_p),
axis=1,
)
# y-z edges
# if self._y_m_z_m_proc is not None:
self._send_info_box[self._y_m_z_m_proc] += self._markers_y_m_z_m.shape[1]
self._send_list_box[self._y_m_z_m_proc] = xp.concatenate(
(self._send_list_box[self._y_m_z_m_proc], self._markers_y_m_z_m),
axis=1,
)
# if self._y_m_z_p_proc is not None:
self._send_info_box[self._y_m_z_p_proc] += self._markers_y_m_z_p.shape[1]
self._send_list_box[self._y_m_z_p_proc] = xp.concatenate(
(self._send_list_box[self._y_m_z_p_proc], self._markers_y_m_z_p),
axis=1,
)
# if self._y_p_z_m_proc is not None:
self._send_info_box[self._y_p_z_m_proc] += self._markers_y_p_z_m.shape[1]
self._send_list_box[self._y_p_z_m_proc] = xp.concatenate(
(self._send_list_box[self._y_p_z_m_proc], self._markers_y_p_z_m),
axis=1,
)
# if self._y_p_z_p_proc is not None:
self._send_info_box[self._y_p_z_p_proc] += self._markers_y_p_z_p.shape[1]
self._send_list_box[self._y_p_z_p_proc] = xp.concatenate(
(self._send_list_box[self._y_p_z_p_proc], self._markers_y_p_z_p),
axis=1,
)
# corners
# if self._x_m_y_m_z_m_proc is not None:
self._send_info_box[self._x_m_y_m_z_m_proc] += self._markers_x_m_y_m_z_m.shape[1]
self._send_list_box[self._x_m_y_m_z_m_proc] = xp.concatenate(
(self._send_list_box[self._x_m_y_m_z_m_proc], self._markers_x_m_y_m_z_m),
axis=1,
)
# if self._x_m_y_m_z_p_proc is not None:
self._send_info_box[self._x_m_y_m_z_p_proc] += self._markers_x_m_y_m_z_p.shape[1]
self._send_list_box[self._x_m_y_m_z_p_proc] = xp.concatenate(
(self._send_list_box[self._x_m_y_m_z_p_proc], self._markers_x_m_y_m_z_p),
axis=1,
)
# if self._x_m_y_p_z_m_proc is not None:
self._send_info_box[self._x_m_y_p_z_m_proc] += self._markers_x_m_y_p_z_m.shape[1]
self._send_list_box[self._x_m_y_p_z_m_proc] = xp.concatenate(
(self._send_list_box[self._x_m_y_p_z_m_proc], self._markers_x_m_y_p_z_m),
axis=1,
)
# if self._x_m_y_p_z_p_proc is not None:
self._send_info_box[self._x_m_y_p_z_p_proc] += self._markers_x_m_y_p_z_p.shape[1]
self._send_list_box[self._x_m_y_p_z_p_proc] = xp.concatenate(
(self._send_list_box[self._x_m_y_p_z_p_proc], self._markers_x_m_y_p_z_p),
axis=1,
)
# if self._x_p_y_m_z_m_proc is not None:
self._send_info_box[self._x_p_y_m_z_m_proc] += self._markers_x_p_y_m_z_m.shape[1]
self._send_list_box[self._x_p_y_m_z_m_proc] = xp.concatenate(
(self._send_list_box[self._x_p_y_m_z_m_proc], self._markers_x_p_y_m_z_m),
axis=1,
)
# if self._x_p_y_m_z_p_proc is not None:
self._send_info_box[self._x_p_y_m_z_p_proc] += self._markers_x_p_y_m_z_p.shape[1]
self._send_list_box[self._x_p_y_m_z_p_proc] = xp.concatenate(
(self._send_list_box[self._x_p_y_m_z_p_proc], self._markers_x_p_y_m_z_p),
axis=1,
)
# if self._x_p_y_p_z_m_proc is not None:
self._send_info_box[self._x_p_y_p_z_m_proc] += self._markers_x_p_y_p_z_m.shape[1]
self._send_list_box[self._x_p_y_p_z_m_proc] = xp.concatenate(
(self._send_list_box[self._x_p_y_p_z_m_proc], self._markers_x_p_y_p_z_m),
axis=1,
)
# if self._x_p_y_p_z_p_proc is not None:
self._send_info_box[self._x_p_y_p_z_p_proc] += self._markers_x_p_y_p_z_p.shape[1]
self._send_list_box[self._x_p_y_p_z_p_proc] = xp.concatenate(
(self._send_list_box[self._x_p_y_p_z_p_proc], self._markers_x_p_y_p_z_p),
axis=1,
)
def _self_communication_boxes(self):
"""Communicate the particles in case a process is it's own neighbour
(in case of periodicity with low number of procs/boxes)"""
if self._send_info_box[self.mpi_rank] > 0:
self.update_holes()
holes_inds = xp.nonzero(self.holes)[0]
if holes_inds.size < self._send_info_box[self.mpi_rank]:
warnings.warn(
f'Strong load imbalance detected: \
number of holes ({holes_inds.size}) on rank {self.mpi_rank} \
is smaller than number of incoming particles ({self._send_info_box[self.mpi_rank]}). \
Increasing the value of "bufsize" in the markers parameters for the next run.',
)
self.mpi_comm.Abort()
# _tmp = self.markers.copy()
# _n_rows_old = _tmp.shape[0]
# logger.info(f"old: {self.markers.shape = }")
# self._bufsize *= 2.0
# self._allocate_marker_array()
# logger.info(f"new: {self.markers.shape = }\n")
# self.markers[:, :] = -1.0
# self.markers[:, :_n_rows_old] = _tmp
# self.update_holes()
# self._update_ghost_particles()
# self._update_valid_mks()
# holes_inds = xp.nonzero(self.holes)[0]
self.markers[:, holes_inds[xp.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[
self.mpi_rank
]
self._track_markers(n_filled=int(self._send_info_box[self.mpi_rank]))
def _sendrecv_all_to_all_boxes(self):
"""
Distribute info on how many markers will be sent/received to/from each process via all-to-all
for the communication of particles in boundary boxes.
"""
self._recv_info_box = xp.zeros(self.mpi_comm.Get_size(), dtype=int)
self.mpi_comm.Alltoall(self._send_info_box, self._recv_info_box)
def _sendrecv_markers_boxes(self):
"""
Use non-blocking communication. In-place modification of markers
for the communication of particles in boundary boxes.
"""
# i-th entry holds the number (not the index) of the first hole to be filled by data from process i
first_hole = xp.cumsum(self._recv_info_box) - self._recv_info_box
hole_inds = xp.nonzero(self._holes)[0]
# Initialize send and receive commands
reqs = []
recvbufs = []
send_reqs = []
for i, (data, N_recv) in enumerate(zip(self._send_list_box, list(self._recv_info_box))):
if i == self.mpi_comm.Get_rank():
reqs += [None]
recvbufs += [None]
else:
send_reqs += [self.mpi_comm.Isend(data, dest=i, tag=self.mpi_comm.Get_rank())]
recvbufs += [xp.zeros((self.n_cols, N_recv), dtype=float)]
reqs += [self.mpi_comm.Irecv(recvbufs[-1], source=i, tag=i)]
# Wait for buffer, then put markers into holes
test_reqs = [False] * (self._recv_info_box.size - 1)
while len(test_reqs) > 0:
# loop over all receive requests
for i, req in enumerate(reqs):
if req is None:
continue
else:
# check if data has been received
if req.Test():
if hole_inds.size < first_hole[i] + self._recv_info_box[i]:
warnings.warn(
f'Strong load imbalance detected: \
number of holes ({hole_inds.size}) on rank {self.mpi_rank} \
is smaller than number of incoming particles ({first_hole[i] + self._recv_info_box[i]}). \
Increasing the value of "bufsize" in the markers parameters for the next run.',
)
self.mpi_comm.Abort()
# exit()
self._markers[:, hole_inds[first_hole[i] + xp.arange(self._recv_info_box[i])]] = recvbufs[i]
test_reqs.pop()
reqs[i] = None
# the send buffers must not be reused before the sends have completed
for req in send_reqs:
req.Wait()
# the received ghost particles filled the lowest holes (counted with the own slot of first_hole)
self._track_markers(n_filled=int(xp.sum(self._recv_info_box)))
self._Barrier()
def _get_neighbouring_proc(self):
"""Find the neighbouring processes for the sending of boxes.
The left (right) neighbour in direction 1 is called x_m_proc (x_p_proc), etc.
By default every process is its own neighbour.
"""
# Faces
self._x_m_proc = None
self._x_p_proc = None
self._y_m_proc = None
self._y_p_proc = None
self._z_m_proc = None
self._z_p_proc = None
# Edges
self._x_m_y_m_proc = None
self._x_m_y_p_proc = None
self._x_p_y_m_proc = None
self._x_p_y_p_proc = None
self._x_m_z_m_proc = None
self._x_m_z_p_proc = None
self._x_p_z_m_proc = None
self._x_p_z_p_proc = None
self._y_m_z_m_proc = None
self._y_m_z_p_proc = None
self._y_p_z_m_proc = None
self._y_p_z_p_proc = None
# Corners
self._x_m_y_m_z_m_proc = None
self._x_m_y_m_z_p_proc = None
self._x_m_y_p_z_m_proc = None
self._x_p_y_m_z_m_proc = None
self._x_m_y_p_z_p_proc = None
self._x_p_y_m_z_p_proc = None
self._x_p_y_p_z_m_proc = None
self._x_p_y_p_z_p_proc = None
# periodicitiy for distance computation
periodic1 = self.bc_sph[0] == "periodic"
periodic2 = self.bc_sph[1] == "periodic"
periodic3 = self.bc_sph[2] == "periodic"
# Determine which proc are on which side (on the host copy, also on the CuPy backend)
dd = self._domain_array_host
rank = self.mpi_rank
x_l = dd[rank][0]
x_r = dd[rank][1]
y_l = dd[rank][3]
y_r = dd[rank][4]
z_l = dd[rank][6]
z_r = dd[rank][7]
for i in range(self.mpi_size):
xl_i = dd[i][0]
xr_i = dd[i][1]
yl_i = dd[i][3]
yr_i = dd[i][4]
zl_i = dd[i][6]
zr_i = dd[i][7]
is_same_x_l = abs(distance(xl_i, x_l, periodic1)) < 1e-5
is_same_x_r = abs(distance(xr_i, x_r, periodic1)) < 1e-5
is_same_y_l = abs(distance(yl_i, y_l, periodic2)) < 1e-5
is_same_y_r = abs(distance(yr_i, y_r, periodic2)) < 1e-5
is_same_z_l = abs(distance(zl_i, z_l, periodic3)) < 1e-5
is_same_z_r = abs(distance(zr_i, z_r, periodic3)) < 1e-5
is_neigh_x_l = abs(distance(xr_i, x_l, periodic1)) < 1e-5
is_neigh_x_r = abs(distance(xl_i, x_r, periodic1)) < 1e-5
is_neigh_y_l = abs(distance(yr_i, y_l, periodic2)) < 1e-5
is_neigh_y_r = abs(distance(yl_i, y_r, periodic2)) < 1e-5
is_neigh_z_l = abs(distance(zr_i, z_l, periodic3)) < 1e-5
is_neigh_z_r = abs(distance(zl_i, z_r, periodic3)) < 1e-5
# Faces
# Process on the left (minus axis) in the x direction
if is_same_y_l and is_same_y_r and is_same_z_l and is_same_z_r and is_neigh_x_l:
self._x_m_proc = i
# Process on the right (plus axis) in the x direction
if is_same_y_l and is_same_y_r and is_same_z_l and is_same_z_r and is_neigh_x_r:
self._x_p_proc = i
# Process on the left (minus axis) in the y direction
if is_same_x_l and is_same_x_r and is_same_z_l and is_same_z_r and is_neigh_y_l:
self._y_m_proc = i
# Process on the right (plus axis) in the y direction
if is_same_x_l and is_same_x_r and is_same_z_l and is_same_z_r and is_neigh_y_r:
self._y_p_proc = i
# Process on the left (minus axis) in the z direction
if is_same_x_l and is_same_x_r and is_same_y_l and is_same_y_r and is_neigh_z_l:
self._z_m_proc = i
# Process on the right (plus axis) in the z direction
if is_same_x_l and is_same_x_r and is_same_y_l and is_same_y_r and is_neigh_z_r:
self._z_p_proc = i
# Edges
# Process on the left in x and left in y axis
if is_same_z_l and is_same_z_r and is_neigh_x_l and is_neigh_y_l:
self._x_m_y_m_proc = i
# Process on the left in x and right in y axis
if is_same_z_l and is_same_z_r and is_neigh_x_l and is_neigh_y_r:
self._x_m_y_p_proc = i
# Process on the right in x and left in y axis
if is_same_z_l and is_same_z_r and is_neigh_x_r and is_neigh_y_l:
self._x_p_y_m_proc = i
# Process on the right in x and right in y axis
if is_same_z_l and is_same_z_r and is_neigh_x_r and is_neigh_y_r:
self._x_p_y_p_proc = i
# Process on the left in x and left in z axis
if is_same_y_l and is_same_y_r and is_neigh_x_l and is_neigh_z_l:
self._x_m_z_m_proc = i
# Process on the left in x and right in z axis
if is_same_y_l and is_same_y_r and is_neigh_x_l and is_neigh_z_r:
self._x_m_z_p_proc = i
# Process on the right in x and left in z axis
if is_same_y_l and is_same_y_r and is_neigh_x_r and is_neigh_z_l:
self._x_p_z_m_proc = i
# Process on the right in x and right in z axis
if is_same_y_l and is_same_y_r and is_neigh_x_r and is_neigh_z_r:
self._x_p_z_p_proc = i
# Process on the left in y and left in z axis
if is_same_x_l and is_same_x_r and is_neigh_y_l and is_neigh_z_l:
self._y_m_z_m_proc = i
# Process on the left in y and right in z axis
if is_same_x_l and is_same_x_r and is_neigh_y_l and is_neigh_z_r:
self._y_m_z_p_proc = i
# Process on the right in y and left in z axis
if is_same_x_l and is_same_x_r and is_neigh_y_r and is_neigh_z_l:
self._y_p_z_m_proc = i
# Process on the right in y and right in z axis
if is_same_x_l and is_same_x_r and is_neigh_y_r and is_neigh_z_r:
self._y_p_z_p_proc = i
# Corners
# Process on the left in x, left in y and left in z axis
if is_neigh_x_l and is_neigh_y_l and is_neigh_z_l:
self._x_m_y_m_z_m_proc = i
# Process on the left in x, left in y and right in z axis
if is_neigh_x_l and is_neigh_y_l and is_neigh_z_r:
self._x_m_y_m_z_p_proc = i
# Process on the left in x, right in y and left in z axis
if is_neigh_x_l and is_neigh_y_r and is_neigh_z_l:
self._x_m_y_p_z_m_proc = i
# Process on the left in x, right in y and right in z axis
if is_neigh_x_l and is_neigh_y_r and is_neigh_z_r:
self._x_m_y_p_z_p_proc = i
# Process on the right in x, left in y and left in z axis
if is_neigh_x_r and is_neigh_y_l and is_neigh_z_l:
self._x_p_y_m_z_m_proc = i
# Process on the right in x, left in y and right in z axis
if is_neigh_x_r and is_neigh_y_l and is_neigh_z_r:
self._x_p_y_m_z_p_proc = i
# Process on the right in x, right in y and left in z axis
if is_neigh_x_r and is_neigh_y_r and is_neigh_z_l:
self._x_p_y_p_z_m_proc = i
# Process on the right in x, right in y and right in z axis
if is_neigh_x_r and is_neigh_y_r and is_neigh_z_r:
self._x_p_y_p_z_p_proc = i
# set empty faces in x
if self._x_m_proc is None:
self._x_m_proc = rank
if self._x_p_proc is None:
self._x_p_proc = rank
# set empty faces in y
if self._y_m_proc is None:
self._y_m_proc = rank
if self._y_p_proc is None:
self._y_p_proc = rank
# set empty faces in z
if self._z_m_proc is None:
self._z_m_proc = rank
if self._z_p_proc is None:
self._z_p_proc = rank
# set empty edges in xy
if self._x_m_y_m_proc is None:
if self._x_m_proc == rank:
self._x_m_y_m_proc = self._y_m_proc
elif self._y_m_proc == rank:
self._x_m_y_m_proc = self._x_m_proc
if self._x_m_y_p_proc is None:
if self._x_m_proc == rank:
self._x_m_y_p_proc = self._y_p_proc
elif self._y_p_proc == rank:
self._x_m_y_p_proc = self._x_m_proc
if self._x_p_y_m_proc is None:
if self._x_p_proc == rank:
self._x_p_y_m_proc = self._y_m_proc
elif self._y_m_proc == rank:
self._x_p_y_m_proc = self._x_p_proc
if self._x_p_y_p_proc is None:
if self._x_p_proc == rank:
self._x_p_y_p_proc = self._y_p_proc
elif self._y_p_proc == rank:
self._x_p_y_p_proc = self._x_p_proc
# set empty edges in xz
if self._x_m_z_m_proc is None:
if self._x_m_proc == rank:
self._x_m_z_m_proc = self._z_m_proc
elif self._z_m_proc == rank:
self._x_m_z_m_proc = self._x_m_proc
if self._x_m_z_p_proc is None:
if self._x_m_proc == rank:
self._x_m_z_p_proc = self._z_p_proc
elif self._z_p_proc == rank:
self._x_m_z_p_proc = self._x_m_proc
if self._x_p_z_m_proc is None:
if self._x_p_proc == rank:
self._x_p_z_m_proc = self._z_m_proc
elif self._z_m_proc == rank:
self._x_p_z_m_proc = self._x_p_proc
if self._x_p_z_p_proc is None:
if self._x_p_proc == rank:
self._x_p_z_p_proc = self._z_p_proc
elif self._z_p_proc == rank:
self._x_p_z_p_proc = self._x_p_proc
# set empty edges in yz
if self._y_m_z_m_proc is None:
if self._y_m_proc == rank:
self._y_m_z_m_proc = self._z_m_proc
elif self._z_m_proc == rank:
self._y_m_z_m_proc = self._y_m_proc
if self._y_m_z_p_proc is None:
if self._y_m_proc == rank:
self._y_m_z_p_proc = self._z_p_proc
elif self._z_p_proc == rank:
self._y_m_z_p_proc = self._y_m_proc
if self._y_p_z_m_proc is None:
if self._y_p_proc == rank:
self._y_p_z_m_proc = self._z_m_proc
elif self._z_m_proc == rank:
self._y_p_z_m_proc = self._y_p_proc
if self._y_p_z_p_proc is None:
if self._y_p_proc == rank:
self._y_p_z_p_proc = self._z_p_proc
elif self._z_p_proc == rank:
self._y_p_z_p_proc = self._y_p_proc
# set empty corners: across a non-periodic domain boundary in one direction
# (face neighbour == rank) the corner neighbour is the edge neighbour in the
# two remaining directions (the edges have been completed above)
if self._x_m_y_m_z_m_proc is None:
if self._x_m_proc == rank:
self._x_m_y_m_z_m_proc = self._y_m_z_m_proc
elif self._y_m_proc == rank:
self._x_m_y_m_z_m_proc = self._x_m_z_m_proc
elif self._z_m_proc == rank:
self._x_m_y_m_z_m_proc = self._x_m_y_m_proc
if self._x_m_y_m_z_p_proc is None:
if self._x_m_proc == rank:
self._x_m_y_m_z_p_proc = self._y_m_z_p_proc
elif self._y_m_proc == rank:
self._x_m_y_m_z_p_proc = self._x_m_z_p_proc
elif self._z_p_proc == rank:
self._x_m_y_m_z_p_proc = self._x_m_y_m_proc
if self._x_m_y_p_z_m_proc is None:
if self._x_m_proc == rank:
self._x_m_y_p_z_m_proc = self._y_p_z_m_proc
elif self._y_p_proc == rank:
self._x_m_y_p_z_m_proc = self._x_m_z_m_proc
elif self._z_m_proc == rank:
self._x_m_y_p_z_m_proc = self._x_m_y_p_proc
if self._x_m_y_p_z_p_proc is None:
if self._x_m_proc == rank:
self._x_m_y_p_z_p_proc = self._y_p_z_p_proc
elif self._y_p_proc == rank:
self._x_m_y_p_z_p_proc = self._x_m_z_p_proc
elif self._z_p_proc == rank:
self._x_m_y_p_z_p_proc = self._x_m_y_p_proc
if self._x_p_y_m_z_m_proc is None:
if self._x_p_proc == rank:
self._x_p_y_m_z_m_proc = self._y_m_z_m_proc
elif self._y_m_proc == rank:
self._x_p_y_m_z_m_proc = self._x_p_z_m_proc
elif self._z_m_proc == rank:
self._x_p_y_m_z_m_proc = self._x_p_y_m_proc
if self._x_p_y_m_z_p_proc is None:
if self._x_p_proc == rank:
self._x_p_y_m_z_p_proc = self._y_m_z_p_proc
elif self._y_m_proc == rank:
self._x_p_y_m_z_p_proc = self._x_p_z_p_proc
elif self._z_p_proc == rank:
self._x_p_y_m_z_p_proc = self._x_p_y_m_proc
if self._x_p_y_p_z_m_proc is None:
if self._x_p_proc == rank:
self._x_p_y_p_z_m_proc = self._y_p_z_m_proc
elif self._y_p_proc == rank:
self._x_p_y_p_z_m_proc = self._x_p_z_m_proc
elif self._z_m_proc == rank:
self._x_p_y_p_z_m_proc = self._x_p_y_p_proc
if self._x_p_y_p_z_p_proc is None:
if self._x_p_proc == rank:
self._x_p_y_p_z_p_proc = self._y_p_z_p_proc
elif self._y_p_proc == rank:
self._x_p_y_p_z_p_proc = self._x_p_z_p_proc
elif self._z_p_proc == rank:
self._x_p_y_p_z_p_proc = self._x_p_y_p_proc
@profile
def _communicate_boxes(self):
"""Refresh the SPH ghost-box layer: build the outgoing ghost markers
(:meth:`_prepare_ghost_particles`), route them to the neighbouring processes
(:meth:`_get_destinations_box`), deliver the ones staying on this process
(:meth:`_self_communication_boxes`), and, if running under MPI, exchange the rest
with neighbouring processes (:meth:`_sendrecv_all_to_all_boxes`,
:meth:`_sendrecv_markers_boxes`) before marking the received markers as ghost
particles (:meth:`_update_ghost_particles`)."""
# if verbose:
# n_valid = xp.count_nonzero(self.valid_mks)
# n_holes = xp.count_nonzero(self.holes)
# n_ghosts = xp.count_nonzero(self.ghost_particles)
# logger.info(f"before communicate_boxes: {self.mpi_rank = }, {n_valid = } {n_holes = }, {n_ghosts = }")
self._prepare_ghost_particles()
self._get_destinations_box()
self._self_communication_boxes()
self.update_holes()
if self.mpi_comm is not None:
self._Barrier()
self._sendrecv_all_to_all_boxes()
self._sendrecv_markers_boxes()
self.update_holes()
self._update_ghost_particles()
# if verbose:
# n_valid = xp.count_nonzero(self.valid_mks)
# n_holes = xp.count_nonzero(self.holes)
# n_ghosts = xp.count_nonzero(self.ghost_particles)
# logger.info(f"after communicate_boxes: {self.mpi_rank = }, {n_valid = }, {n_holes = }, {n_ghosts = }")
def _eval_sph(
self,
eta1: xp.ndarray,
eta2: xp.ndarray,
eta3: xp.ndarray,
index: int,
out: xp.ndarray = None,
fast: bool = True,
kernel_type: str = "gaussian_1d",
derivative: int = 0,
h1: float = 0.1,
h2: float = 0.1,
h3: float = 0.1,
):
r"""Perform a (meshgrid) SPH evaluation of a function :math:`\rho: [0, 1]^3 \to \mathbb R` in the following sense:
.. math::
\rho(\boldsymbol \eta_i) = \sum_{j=0}^{N-1} \rho_j\, W_h(\boldsymbol \eta_i - \boldsymbol \eta_j)\,.
The coefficients :math:`\rho_j` must be available in the marker array, stored at some index ``self.markers[index, j]``.
In case that `derivative=k` where `k` is not zero, the `k`-th component of the gradient of :math:`\rho` is computed:
.. math::
\textrm{derivative}=k:\qquad [\nabla \rho(\boldsymbol \eta_i)]_k = \sum_{j=0}^{N-1} \rho_j \frac{\partial W_h}{\partial \eta_k}(\boldsymbol \eta_i - \boldsymbol \eta_j)\,.
The possible choices for :math:`W_h` are listed in :ref:`smoothing_kernels`
and in :meth:`~struphy.pic.base.Particles.ker_dct`.
Parameters
----------
eta1, eta2, eta3 : array_like
Logical evaluation points.
index : int
At which index of the markers array are located the coefficients :math:`\rho_j`.
out : array_like
Output will be store in this array. A new array is created if not provided.
fast : bool
If true, uses an optimized evaluation algorithm taking advantage of the box structure.
This assume that the boxes are bigger then the radius used for the smoothing kernel.
kernel_type : str
Name of the smoothing kernel, see :mod:`~struphy.pic.sph_smoothing_kernels`
and :meth:`~struphy.pic.base.Particles.ker_dct`.
derivative: int
0: no kernel derivative
1: first component of grad
2: second component of grad
3: third component of grad
h1, h2, h3 : float
Radius of the smoothing kernel in each dimension.
"""
_shp = xp.shape(eta1)
assert _shp == xp.shape(eta2) == xp.shape(eta3)
if out is not None:
assert _shp == xp.shape(out)
else:
out = xp.zeros_like(eta1)
assert derivative in {0, 1, 2, 3}, f"derivative must be 0, 1, 2 or 3, but is {derivative}."
ker_id = self.ker_dct()[kernel_type]
ker_id += derivative
# for the moment we always assume periodicity for the evaluation near the boundary, TODO: fill ghost boxes with suitable markers for other bcs?
periodic1, periodic2, periodic3 = [True] * 3 # [bci == "periodic" for bci in self.bc]
self.put_particles_in_boxes()
if fast:
if len(_shp) == 1:
func = box_based_evaluation_flat
elif len(_shp) == 3:
if _shp[0] > 1:
assert eta1[0, 0, 0] != eta1[1, 0, 0], "Meshgrids must be obtained with indexing='ij'!"
if _shp[1] > 1:
assert eta2[0, 0, 0] != eta2[0, 1, 0], "Meshgrids must be obtained with indexing='ij'!"
if _shp[2] > 1:
assert eta3[0, 0, 0] != eta3[0, 0, 1], "Meshgrids must be obtained with indexing='ij'!"
func = box_based_evaluation_meshgrid
func(
self.args_markers,
eta1,
eta2,
eta3,
self.sorting_boxes.nx,
self.sorting_boxes.ny,
self.sorting_boxes.nz,
self.domain_array[self.mpi_rank],
self.sorting_boxes.boxes,
self.sorting_boxes.neighbours,
self.holes,
periodic1,
periodic2,
periodic3,
index,
ker_id,
h1,
h2,
h3,
out,
)
else:
if len(_shp) == 1:
func = naive_evaluation_flat
elif len(_shp) == 3:
func = naive_evaluation_meshgrid
func(
self.args_markers,
eta1,
eta2,
eta3,
self.holes,
periodic1,
periodic2,
periodic3,
index,
ker_id,
h1,
h2,
h3,
out,
)
return out
### MPI comm for domain decomposition ###
def _is_on_domain_of_rank(self, etas, rank):
"""
Check which logical positions lie on the domain of a given process.
Uses the half-open interval [left, right) in each direction, so that a position
exactly on a boundary between two processes belongs to exactly one of them.
The right end eta = 1 of the unit cube is included in the last process.
Parameters
----------
etas : array[float]
Logical positions of shape (3, n), one column per position.
rank : int
Process rank (row index of domain_array).
Returns
-------
is_on_domain : array[bool]
Array of shape (3, n), True where the position lies in the process domain in that direction.
"""
left = self.domain_array[rank, 0::3, None]
right = self.domain_array[rank, 1::3, None]
# include eta = 1 in the last process: shift its right boundary to the next float
right = xp.where(right == 1.0, xp.nextafter(right, 2.0), right)
return xp.logical_and(etas >= left, etas < right)
def _sendrecv_determine_mtbs(
self,
alpha: list | tuple | xp.ndarray = (1.0, 1.0, 1.0),
):
"""
Determine which markers have to be sent from current process and put them in a new array.
Corresponding columns in markers array become holes and are therefore set to -1.
This can be done purely with numpy functions (fast, vectorized).
Parameters
----------
alpha : list | tuple
For i=1,2,3 the sorting is according to alpha[i]*markers[i] + (1 - alpha[i])*markers[first_pusher_idx + i].
alpha[i] must be between 0 and 1.
Returns
-------
hole_inds_after_send : array[int]
Indices of empty columns in markers after send.
sorting_etas : array[float]
Eta-values of shape (3, n_markers_max) according to which the sorting is performed.
"""
# position that determines the sorting (including periodic shift of boundary conditions)
if not isinstance(alpha, xp.ndarray):
# check a tuple/list on the host: on CuPy, asserting on a device comparison reads it back
alpha_host = np.array(alpha, dtype=float)
assert alpha_host.size == 3
assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0)
if alpha_host[0] == alpha_host[1] == alpha_host[2]:
# the common case (alpha = 1 or alpha = 0.5 in every direction): a Python scalar,
# which needs no host-to-device copy on CuPy and gives the same values on NumPy
alpha = float(alpha_host[0])
else:
alpha = xp.asarray(alpha_host)[:, None]
else:
assert alpha.size == 3
assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0)
alpha = alpha.reshape(3, 1)
# only the markers (columns) that can hold markers; the columns after them are holes, which stay
n = self.n_markers_used
bi = self.first_pusher_idx
sorting_etas = self._sorting_etas[:, :n]
sorting_etas[:] = (
alpha * (self.markers[:3, :n] + self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3, :n])
+ (1.0 - alpha) * self.markers[bi : bi + 3, :n]
)
# for non-periodic axes, eta = 1 is the right end of the domain and must not be wrapped to eta = 0
non_periodic_axes = [axis for axis in range(3) if axis not in self._periodic_axes]
at_right_end = {axis: sorting_etas[axis] == 1.0 for axis in non_periodic_axes}
xp.mod(sorting_etas, 1.0, out=sorting_etas)
for axis, mask in at_right_end.items():
sorting_etas[axis, mask] = 1.0
# check which particles are on the current process domain (into the first columns of the full-size buffer)
on_domain = self._is_on_proc_domain[:, :n]
on_domain[:] = self._is_on_domain_of_rank(sorting_etas, self.mpi_rank)
# to stay on the current process, all three directions must be True.
# Combining the three rows is faster than reducing over the axis of size 3
self._can_stay[:n] = on_domain[0] & on_domain[1] & on_domain[2]
self._can_stay[n:] = True
# holes and ghosts can stay, too
self._can_stay[self.holes] = True
self._can_stay[self.ghost_particles] = True
# True values can stay on the process, False must be sent, already empty markers (-1) cannot be sent
send_inds = xp.nonzero(~self._can_stay)[0]
hole_inds_after_send = xp.nonzero(xp.logical_or(~self._can_stay, self.holes))[0]
return hole_inds_after_send, send_inds
def _compute_neighbor_ranks(self) -> tuple[list[int], list[int]]:
"""Split every other rank into geometric neighbours of this rank's sub-domain
box and everyone else, for :meth:`_sendrecv_get_destinations`.
A rank is a neighbour (Moore neighbourhood: touching along every axis, sharing
at least a corner) if, for each of the 3 axes, its box interval touches or
overlaps this rank's -- including wrap-around for axes with a periodic boundary
condition (:attr:`_periodic_axes`), where the box touching ``eta=1`` is also
adjacent to the box touching ``eta=0``.
Domain decomposition is a static, non-overlapping Cartesian tiling of the unit
cube (see :meth:`_get_domain_decomp`) computed once at construction, so this is
safe to compute once here rather than re-derived every call.
Returns
-------
neighbor_ranks : list[int]
Ranks (excluding this one) whose box touches or overlaps this rank's.
non_neighbor_ranks : list[int]
Every other rank (excluding this one and ``neighbor_ranks``).
"""
# Tiny (mpi_size, 9) array -- the host copy, so the O(mpi_size)
# comparisons below don't pay a device sync per rank checked.
domain_array_host = self._domain_array_host
own = domain_array_host[self.mpi_rank]
tol = 1e-12
neighbor_ranks = []
non_neighbor_ranks = []
for j in range(self.mpi_size):
if j == self.mpi_rank:
continue
other = domain_array_host[j]
touching = True
for axis in range(3):
own_l, own_r = own[3 * axis], own[3 * axis + 1]
other_l, other_r = other[3 * axis], other[3 * axis + 1]
axis_touches = other_r >= own_l - tol and other_l <= own_r + tol
if not axis_touches and axis in self._periodic_axes:
axis_touches = (abs(own_r - 1.0) < tol and abs(other_l) < tol) or (
abs(own_l) < tol and abs(other_r - 1.0) < tol
)
if not axis_touches:
touching = False
break
(neighbor_ranks if touching else non_neighbor_ranks).append(j)
return neighbor_ranks, non_neighbor_ranks
def _sendrecv_get_destinations(self, send_inds):
"""
Determine to which process particles have to be sent.
Parameters
----------
send_inds : array[int]
Indices of particles which will be sent.
Returns
-------
send_info : array[int]
Amount of particles sent to i-th process.
"""
# One entry for each process. The counts are host integers on every backend: the
# size of a nonzero() result is known on the host (CuPy synchronizes to size it), so
# a small NumPy array holds them without any host/device copy, and the Alltoall of
# the counts (_sendrecv_all_to_all) passes host buffers to MPI.
send_info = np.zeros(self.mpi_size, dtype=int)
# Gathered once and reused for every rank below, instead of re-gathering
# self.markers[:, send_inds] and self._sorting_etas[:, send_inds] fresh on every
# iteration of the rank loop (as the previous version did).
candidates = self.markers[:, send_inds]
etas_to_send = self._sorting_etas[:, send_inds]
# Reset every rank's send buffer to empty first. The neighbour/non-neighbour
# loop below can `break` before visiting every rank once all candidates are
# matched, but _sendrecv_markers unconditionally Isends self._send_list[i] for
# every i != mpi_rank -- a rank skipped this call must not keep a stale,
# non-empty buffer from a previous call (send/recv size would then disagree
# with send_info, which is always correct since it defaults to 0 above).
empty_local = xp.empty(0, dtype=int)
empty_markers = candidates[:, :0]
for i in range(self.mpi_size):
self._send_to_i[i] = empty_local
self._send_list[i] = empty_markers
# A marker leaving this rank's domain is overwhelmingly likely to land in a
# geometrically adjacent sub-domain -- one push (sub-)step is small compared to
# a sub-domain -- so check neighbour ranks (_compute_neighbor_ranks) first
# instead of every rank. `remaining` tracks positions into send_inds/
# etas_to_send not yet matched to a destination; if neighbours don't cover
# everyone (a marker moved further than one sub-domain this step), the leftover
# few are checked against every other rank in the second pass, so this changes
# only how many ranks get checked in the common case, not correctness.
remaining = xp.arange(send_inds.shape[0])
for rank_group in (self._neighbor_ranks, self._non_neighbor_ranks):
if remaining.size == 0:
break
etas_remaining = etas_to_send[:, remaining]
still_remaining = xp.ones(remaining.shape[0], dtype=bool)
for i in rank_group:
conds = self._is_on_domain_of_rank(etas_remaining, i)
matched_local = xp.nonzero(xp.all(conds, axis=0))[0]
matched = remaining[matched_local]
self._send_to_i[i] = matched
send_info[i] = matched.size
# C-contiguous, as the receive buffers (n_cols, n): selecting columns gives a Fortran-ordered copy,
# which MPI would send marker by marker
self._send_list[i] = xp.ascontiguousarray(candidates[:, matched])
still_remaining[matched_local] = False
remaining = remaining[still_remaining]
return send_info
def _sendrecv_all_to_all(self, send_info):
"""
Distribute info on how many markers will be sent/received to/from each process via all-to-all.
Parameters
----------
send_info : array[int]
Amount of markers to be sent to i-th process.
Returns
-------
recv_info : array[int]
Amount of marticles to be received from i-th process.
Notes
-----
``send_info`` and ``recv_info`` are small NumPy arrays (one integer per rank) on every
backend: the counts size the receive buffers and the holes to fill, which Python needs
on the host anyway. On CuPy, only the markers themselves are device buffers
(:meth:`_sendrecv_markers`).
"""
recv_info = np.zeros(self.mpi_size, dtype=int)
self.mpi_comm.Alltoall(send_info, recv_info)
return recv_info
def _sendrecv_markers(self, recv_info, hole_inds_after_send):
"""
Use non-blocking communication. In-place modification of markers
Parameters
----------
recv_info : array[int]
Amount of markers to be received from i-th process.
hole_inds_after_send : array[int]
Indices of empty columns (holes) in markers after send.
"""
if xp.get_array_backend(self._markers) == "cupy":
self._sendrecv_markers_device(recv_info, hole_inds_after_send)
return
# i-th entry holds the number (not the index) of the first hole to be filled by data from process i
first_hole = xp.cumsum(recv_info) - recv_info
# Initialize send and receive commands
send_reqs = []
for i, (data, N_recv) in enumerate(zip(self._send_list, list(recv_info))):
if i == self.mpi_rank:
self._reqs[i] = None
self._recvbufs[i] = None
else:
send_reqs += [self.mpi_comm.Isend(data, dest=i, tag=self.mpi_rank)]
self._recvbufs[i] = xp.zeros((self.n_cols, N_recv), dtype=float)
self._reqs[i] = self.mpi_comm.Irecv(self._recvbufs[i], source=i, tag=i)
# Wait for buffer, then put markers into holes
test_reqs = [False] * (recv_info.size - 1)
while len(test_reqs) > 0:
# loop over all receive requests
for i, req in enumerate(self._reqs):
if req is None:
continue
else:
# check if data has been received
if req.Test():
if hole_inds_after_send.size < first_hole[i] + recv_info[i]:
warnings.warn(
f'Strong load imbalance detected: \
number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \
is smaller than number of incoming particles ({first_hole[i] + recv_info[i]}). \
Increasing the value of "bufsize" in the markers parameters for the next run.',
)
self.mpi_comm.Abort()
self.markers[:, hole_inds_after_send[first_hole[i] + xp.arange(recv_info[i])]] = self._recvbufs[
i
]
test_reqs.pop()
self._reqs[i] = None
# the send buffers must not be reused before the sends have completed
for req in send_reqs:
req.Wait()
def _mpi_cuda_aware(self) -> bool:
"""Whether MPI takes device (CuPy) buffers, for :meth:`_sendrecv_markers_device`.
Uses the answer recorded by :func:`cunumpy.mpi.mpi_is_cuda_aware` or
:func:`cunumpy.mpi.set_mpi_cuda_aware` (e.g. by feectools or at startup). If nothing
was recorded, probes :attr:`mpi_comm` once; the probe is collective, which is safe
here because :meth:`mpi_sort_markers` is called by every rank of :attr:`mpi_comm`.
"""
cuda_aware = xp.mpi.get_mpi_cuda_aware()
if cuda_aware is None:
cuda_aware = xp.mpi.mpi_is_cuda_aware(self.mpi_comm)
if not cuda_aware and not Particles._warned_mpi_staging:
Particles._warned_mpi_staging = True
logger.warning(
"MPI is not CUDA-aware: markers exchanged between ranks are staged through host memory.",
)
return cuda_aware
_warned_mpi_staging = False
def _sendrecv_markers_device(self, recv_info, hole_inds_after_send):
"""Exchange markers and scatter the incoming ones into empty slots (holes).
:func:`cunumpy.mpi.exchange` handles GPU synchronization, optional host
staging, and completion of all MPI transfers.
Parameters
----------
recv_info : array[int]
Number of markers to receive from each process (NumPy array).
hole_inds_after_send : array[int]
Indices of empty columns (holes) in markers after send (device array).
"""
recv_info = np.asarray(recv_info)
n_recv = int(recv_info.sum())
if hole_inds_after_send.size < n_recv:
warnings.warn(
f"Strong load imbalance detected: "
f"number of holes ({hole_inds_after_send.size}) "
f"on rank {self.mpi_rank} is smaller than number of "
f"incoming particles ({n_recv}). "
'Increase "bufsize" in the markers parameters for the next run.',
)
self.mpi_comm.Abort()
# one receive buffer (n_cols, n) per source rank: a slice of a common (n_cols, n_recv) buffer would not be
# contiguous
first_hole = np.cumsum(recv_info) - recv_info
sends = []
receives = []
for rank in range(self.mpi_size):
if rank == self.mpi_rank:
continue
if recv_info[rank] > 0:
recvbuf = xp.empty((self.n_cols, int(recv_info[rank])), dtype=self.markers.dtype)
receives.append((recvbuf, rank, rank))
data = self._send_list[rank]
if data.shape[1] > 0:
sends.append((data, rank, self.mpi_rank))
xp.mpi.exchange(
self.mpi_comm,
sends=sends,
receives=receives,
cuda_aware=self._mpi_cuda_aware(),
)
for recvbuf, rank, _ in receives:
start = int(first_hole[rank])
stop = start + int(recv_info[rank])
self._markers[:, hole_inds_after_send[start:stop]] = recvbuf
class Tesselation:
"""
Subdivide each :class:`~struphy.pic.sorting.SortingBoxes` box on the current MPI
process into ``tiles_pb`` equally sized, axis-aligned tiles (1 sorting box, hence 1
set of tiles, if ``sorting_boxes=None``).
The tesselation gives a deterministic, evenly spaced set of points inside the process
domain: :meth:`draw_markers` places one marker at the midpoint of each tile (used for
``loading="tesselation"``, e.g. deterministic SPH marker loading with ``ppb`` markers per
box), and :meth:`cell_averages` integrates a given function over each tile via
Gauss-Legendre quadrature (used to initialize marker weights from cell averages of the
background distribution).
The number of tiles per direction is chosen by factorizing ``tiles_pb`` and greedily
assigning factors to whichever direction currently has the fewest tiles (see
:meth:`get_tiles`), so ``tiles_pb`` should factorize into small primes for a
close-to-cubic tile shape.
Parameters
----------
tiles_pb : int
Number of equally sized tiles per box defined in sorting boxes (there is 1 box if sorting_boxes=None).
This is equal to particels per box (ppb) when used for SPH markers.
comm : Intracomm
MPI communicator.
domain_array : xp.ndarray
A 2d array[float] of shape (comm.Get_size(), 9) holding info on the domain decomposition.
sorting_boxes : SortingBoxes
Box info for SPH evaluations.
"""
def __init__(
self,
tiles_pb: int | float,
*,
comm: "Intracomm" = None,
domain_array: xp.ndarray = None,
sorting_boxes: SortingBoxes = None,
):
if isinstance(tiles_pb, int):
self._tiles_pb = tiles_pb
else:
if tiles_pb == int(tiles_pb):
self._tiles_pb = int(tiles_pb)
else:
self._tiles_pb = int(tiles_pb + 1)
if comm is None:
self._rank = 0
else:
self._rank = comm.Get_rank()
assert domain_array is not None
if domain_array is None:
self._starts = xp.zeros(3)
self._ends = xp.ones(3)
else:
self._starts = domain_array[self.rank, 0::3]
self._ends = domain_array[self.rank, 1::3]
if sorting_boxes is None:
self._boxes_per_dim = [1, 1, 1]
else:
self._boxes_per_dim = [
sorting_boxes.nx,
sorting_boxes.ny,
sorting_boxes.nz,
]
self._box_widths = [(ri - le) / nb for ri, le, nb in zip(self._ends, self._starts, self.boxes_per_dim)]
n_boxes = 1
for nbi in self.boxes_per_dim:
n_boxes *= nbi
if n_boxes == 1:
self._dims_mask = [True] * 3
else:
self._dims_mask = xp.array(self.boxes_per_dim) > 1
min_tiles = 2 ** xp.count_nonzero(self.dims_mask)
assert self.tiles_pb >= min_tiles, (
f"At least {min_tiles} tiles per sorting box is enforced, but you have {self.tiles_pb}!"
)
self._n_tiles = n_boxes * self.tiles_pb
self.get_tiles()
def get_tiles(self):
"""Compute the tesselation of a single sorting box: split ``tiles_pb`` into prime
factors and distribute them, one at a time, over the directions marked True in
:attr:`dims_mask` (each factor going to the direction with the currently fewest
tiles), yielding :attr:`nt_per_dim`. From this, compute the tile breakpoints
(:attr:`tile_breaks`), midpoints (:attr:`tile_midpoints`) and volume
(:attr:`tile_volume`) within a single sorting box, all expressed relative to the
box's own origin (i.e. on ``[0, box_width]`` in each direction)."""
# factorize tiles per box
factors = factorint(self.tiles_pb)
factors_vec = []
for fac, multiplicity in factors.items():
for m in range(multiplicity):
factors_vec += [fac]
# logger.info(f'{self.tiles_pb = }')
# logger.info(f'{factors_vec = }')
# logger.info(f'{self.dims_mask = }')
# tiles in one sorting box
self._nt_per_dim = xp.array([1, 1, 1])
_ids = xp.nonzero(self._dims_mask)[0]
for fac in factors_vec:
_nt = self.nt_per_dim[self._dims_mask]
d = _ids[xp.argmin(_nt)]
self._nt_per_dim[d] *= fac
# logger.info(f'{_nt = }, {d = }, {self.nt_per_dim = }')
assert xp.prod(self.nt_per_dim) == self.tiles_pb
# tiles between [0, box_width] in each direction
self._tile_breaks = [xp.linspace(0.0, bw, nt + 1) for bw, nt in zip(self.box_widths, self.nt_per_dim)]
self._tile_midpoints = [(xp.roll(tbs, -1)[:-1] + tbs[:-1]) / 2 for tbs in self.tile_breaks]
self._tile_volume = 1.0
for tb in self.tile_breaks:
self._tile_volume *= tb[1]
def draw_markers(self):
"""Place one marker at the midpoint of every tile of every sorting box on the
current process, i.e. a deterministic alternative to random marker loading.
Returns
-------
eta1, eta2, eta3 : xp.ndarray
1d arrays of logical-space marker coordinates, one entry per tile
(length :attr:`n_tiles`)."""
_, eta1 = self._tile_output_arrays()
eta2 = xp.zeros_like(eta1)
eta3 = xp.zeros_like(eta1)
nt_x, nt_y, nt_z = self.nt_per_dim
for i in range(self.boxes_per_dim[0]):
x_midpoints = self._get_midpoints(i, 0)
for j in range(self.boxes_per_dim[1]):
y_midpoints = self._get_midpoints(j, 1)
for k in range(self.boxes_per_dim[2]):
z_midpoints = self._get_midpoints(k, 2)
xx, yy, zz = xp.meshgrid(
x_midpoints,
y_midpoints,
z_midpoints,
indexing="ij",
)
eta1[
i * nt_x : (i + 1) * nt_x,
j * nt_y : (j + 1) * nt_y,
k * nt_z : (k + 1) * nt_z,
] = xx
eta2[
i * nt_x : (i + 1) * nt_x,
j * nt_y : (j + 1) * nt_y,
k * nt_z : (k + 1) * nt_z,
] = yy
eta3[
i * nt_x : (i + 1) * nt_x,
j * nt_y : (j + 1) * nt_y,
k * nt_z : (k + 1) * nt_z,
] = zz
return eta1.flatten(), eta2.flatten(), eta3.flatten()
def _get_quad_pts(self, n_quad=None):
"""Compute the Gauss-Legendre quadrature points and weights on a single tile
(:attr:`tile_quad_pts`, :attr:`tile_quad_wts`), rescaled from ``[-1, 1]`` to the
tile's own extent, i.e. the first tile of :attr:`tile_breaks` in each direction."""
if n_quad is None:
n_quad = [1, 1, 1]
elif isinstance(n_quad, int):
n_quad = [n_quad] * 3
self._tile_quad_pts = []
self._tile_quad_wts = []
for nq, tb in zip(n_quad, self.tile_breaks):
pts_loc, wts_loc = xp.polynomial.legendre.leggauss(nq)
pts, wts = quadrature_grid(tb[:2], pts_loc, wts_loc)
self._tile_quad_pts += [pts[0]]
self._tile_quad_wts += [wts[0]]
def cell_averages(self, fun, n_quad=None):
"""Compute the cell average of ``fun`` over every tile on the current process,
via Gauss-Legendre quadrature of the given order in each direction.
Parameters
----------
fun : callable
Function of (eta1, eta2, eta3) to be averaged; called once per sorting box
with meshgrid-shaped arrays of quadrature points.
n_quad : int | list[int]
Number of Gauss-Legendre quadrature points per tile, either the same in all
three directions (int) or one value per direction (list of 3 ints).
Defaults to 1 point per direction.
Returns
-------
out : xp.ndarray
3d array of shape ``n_tiles_per_dim * boxes_per_dim`` (one entry per tile on
the current process) holding the cell average of ``fun`` over each tile.
"""
self._get_quad_pts(n_quad=n_quad)
# logger.info(f'{self.tile_quad_pts = }')
single_box_out, out = self._tile_output_arrays()
nt_x, nt_y, nt_z = self.nt_per_dim
for i in range(self.boxes_per_dim[0]):
x_pts = self._get_box_quad_pts(i, 0)
for j in range(self.boxes_per_dim[1]):
y_pts = self._get_box_quad_pts(j, 1)
for k in range(self.boxes_per_dim[2]):
z_pts = self._get_box_quad_pts(k, 2)
xx, yy, zz = xp.meshgrid(
x_pts.flatten(),
y_pts.flatten(),
z_pts.flatten(),
indexing="ij",
)
fun_vals = fun(xx, yy, zz)
sampling_kernels.tile_int_kernel(
fun_vals,
*self.tile_quad_wts,
single_box_out,
)
# single_box_out /= self.tile_volume
out[
i * nt_x : (i + 1) * nt_x,
j * nt_y : (j + 1) * nt_y,
k * nt_z : (k + 1) * nt_z,
] = single_box_out
return out
def _tile_output_arrays(self):
"""Allocate two 3d arrays filled with zeros, to be filled tile-by-tile and
box-by-box in :meth:`draw_markers` and :meth:`cell_averages`:
* the first, of shape :attr:`nt_per_dim`, holds one entry per tile within a
single sorting box;
* the second, of shape ``nt_per_dim * boxes_per_dim``, holds one entry per tile
on the current process (i.e. the first array tiled over all sorting boxes).
"""
# self._quad_pts = [xp.zeros((nt, nq)).flatten() for nt, nq in zip(self.nt_per_dim, self.tile_quad_pts)]
single_box_out = xp.zeros(self.nt_per_dim)
out = xp.tile(single_box_out, self.boxes_per_dim)
return single_box_out, out
def _get_midpoints(self, i: int, dim: int):
"""Compute the tile midpoints within the i-th sorting box in direction dim,
i.e. :attr:`tile_midpoints` shifted from the single-box-local frame ``[0,
box_width]`` to the process-global frame.
Parameters
----------
i : int
Index of the sorting box in direction dim, starting at 0.
dim : int
Direction, either 0, 1, or 2.
Returns
-------
xp.ndarray
1d array of tile midpoints, length ``nt_per_dim[dim]``.
"""
xl = self.starts[dim] + i * self.box_widths[dim]
return xl + self.tile_midpoints[dim]
def _get_box_quad_pts(self, i: int, dim: int):
"""Compute all quadrature points for cell averages within the i-th sorting box in direction dim.
Parameters
----------
i : int
Index of the box, starting at 0.
dim : int
Direction, either 0, 1, or 2.
Returns
-------
x_pts : xp.array
2d array of shape (n_tiles_pb, n_tile_quad_pts)
"""
xl = self.starts[dim] + i * self.box_widths[dim]
x_tile_breaks = xl + self.tile_breaks[dim][:-1]
x_tile_pts = self.tile_quad_pts[dim]
x_pts = xp.tile(x_tile_breaks, (x_tile_pts.size, 1)).T + x_tile_pts
return x_pts
@property
def tiles_pb(self):
"""Number of equally sized tiles per sorting box."""
return self._tiles_pb
@property
def n_tiles(self):
"""Total number of tiles on current process."""
return self._n_tiles
@property
def nt_per_dim(self):
"""3-list of number of equally sized tiles per sorting box per direction."""
return self._nt_per_dim
@property
def starts(self):
"""3-list of domain starts (left boundaries) on current process."""
return self._starts
@property
def ends(self):
"""3-list of domain ends (right boundaries) on current process."""
return self._ends
@property
def tile_breaks(self):
"""3-list of tile break points within the single sorting box [0.0, sorting_box_width], in each direction."""
return self._tile_breaks
@property
def tile_midpoints(self):
"""3-list of tile midpoints within the single sorting box [0.0, sorting_box_width], in each direction."""
return self._tile_midpoints
@property
def tile_volume(self):
"""Volume of a single tile."""
return self._tile_volume
@property
def tile_quad_pts(self):
"""3-list of quadrature points (GL) within a single tile, in each direction."""
return self._tile_quad_pts
@property
def tile_quad_wts(self):
"""3-list of quadrature weights (GL) within a single tile, in each direction."""
return self._tile_quad_wts
@property
def rank(self):
"""Current process rank."""
return self._rank
@property
def boxes_per_dim(self):
"""Sorting boxes per direction."""
return self._boxes_per_dim
@property
def box_widths(self):
"""3-list of sorting box widths in each direction."""
return self._box_widths
@property
def dims_mask(self):
"""Boolean array of size 3; entry is True if direction participates in tesselation."""
return self._dims_mask
