From a5c6479d829559b27c15399a503d7726ef9f54c6 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 12 Oct 2026 20:35:30 +0200 Subject: [PATCH 1/5] Particles on domains with a pole at eta1=0 (no hole around the axis) Markers can now move through the axis of a domain with a pole (e.g. IGAPolarCylinder, IGAPolarTorus, HollowCylinder/HollowTorus with a1=0), with polar splines for the FEEC parts. - update_position_stage (pusher_utilities_kernels): the stage update of the explicit Runge-Kutta pushers (PushEta, GC bxE* and B* explicit) is done in the coordinates xi = eta1 (cos 2 pi eta2, sin 2 pi eta2), in which the mapping and the marker velocity are regular at the pole. Straight lines through the axis of a cylinder are exact. - Axis map in apply_kinetic_bc_marker (Pyccel and CUDA) and Particles.apply_kinetic_bc: a marker with eta1 < 0 is mapped to (-eta1, eta2 + 1/2); eta1 is kept >= 1e-10, also at loading. The discrete-gradient GC pushers only get this map (mid-points in eta are not regular at the axis). - New field MarkerArguments.pole (and CudaMarkerArguments, pusher_args.cuh). Pole domains are not supported on the CuPy backend yet, nor with poc/pol_period != 1. - PushVxB: the projected equilibrium b2 was extracted twice (assertion with polar splines); only the b2 variable is extracted now. - PushGuidingCenterBxEstar/Parallel: phi, the electric field and b_tilde are passed to the kernels as tensor-product coefficients (PolarVector has no _data). - Reflecting boundary: only Cartesian velocities (vdim=3) are reflected; for Particles5D the columns 3:6 (v_parallel, mu, weight) were overwritten. - CanonicalMaxwellian2D: removed the unused read of domain.params["a1"] (KeyError on IGAPolarTorus). - AdhocTorus, AdhocTorusQPsi: psi derivatives are regular on the magnetic axis (were 0/0). - GuidingCenter: the magnetic energy en_fB summed over holes and ghosts instead of the valid markers (each removed marker added 1 to en_tot). - Tests in pic/tests/test_pole.py. Co-Authored-By: Claude Opus 5.5 --- src/struphy/fields_background/equils.py | 52 +-- src/struphy/kernel_arguments/pusher_args.cuh | 1 + .../kernel_arguments/pusher_args_cuda.py | 3 + .../kernel_arguments/pusher_args_kernels.py | 6 + src/struphy/kinetic_background/maxwellians.py | 1 - src/struphy/models/guiding_center.py | 4 +- src/struphy/pic/base.py | 32 ++ .../push_eta_stage/push_eta_stage_kernels.py | 12 +- ...sh_gc_Bstar_explicit_multistage_kernels.py | 13 +- ..._gc_bxEstar_explicit_multistage_kernels.py | 12 +- .../pic/pushing/pusher_utilities_kernels.cuh | 17 +- .../pic/pushing/pusher_utilities_kernels.py | 117 +++++- src/struphy/pic/tests/test_accum_gc_gradB.py | 2 +- src/struphy/pic/tests/test_kernel_backends.py | 4 +- src/struphy/pic/tests/test_kernel_setup.py | 8 +- src/struphy/pic/tests/test_pole.py | 369 ++++++++++++++++++ .../push_guiding_center_bx_estar.py | 50 ++- .../push_guiding_center_parallel.py | 50 ++- src/struphy/propagators/push_vxb.py | 19 +- 19 files changed, 659 insertions(+), 113 deletions(-) create mode 100644 src/struphy/pic/tests/test_pole.py diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index bac148c33..cde0e9ec7 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -1242,24 +1242,26 @@ def psi(self, R, Z, dR=0, dZ=0): if dR == 0 and dZ == 0: out = self.psi_r(r, der=0) else: - dr_dR = (R - self.params["R0"]) / r - dr_dZ = Z / r - - d2r_dR2 = (r - (R - self.params["R0"]) * dr_dR) / r**2 - d2r_dZ2 = (r - Z * dr_dZ) / r**2 - - d2r_dRdZ = -Z * (R - self.params["R0"]) / r**3 + # regular on the magnetic axis r=0: with g = psi'(r)/r -> psi''(0) and the direction cosines + # c_R = (R - R0)/r, c_Z = Z/r (set to 0 on the axis), the derivatives below have the correct limits + on_axis = r == 0.0 + r_safe = xp.where(on_axis, 1.0, r) + dpsi = self.psi_r(r, der=1) + d2psi = self.psi_r(r, der=2) + g = xp.where(on_axis, d2psi, dpsi / r_safe) + c_R = (R - self.params["R0"]) / r_safe + c_Z = Z / r_safe if dR == 1 and dZ == 0: - out = self.psi_r(r, der=1) * dr_dR + out = g * (R - self.params["R0"]) elif dR == 0 and dZ == 1: - out = self.psi_r(r, der=1) * dr_dZ + out = g * Z elif dR == 2 and dZ == 0: - out = self.psi_r(r, der=2) * dr_dR**2 + self.psi_r(r, der=1) * d2r_dR2 + out = d2psi * c_R**2 + g * (1.0 - c_R**2) elif dR == 0 and dZ == 2: - out = self.psi_r(r, der=2) * dr_dZ**2 + self.psi_r(r, der=1) * d2r_dZ2 + out = d2psi * c_Z**2 + g * (1.0 - c_Z**2) elif dR == 1 and dZ == 1: - out = self.psi_r(r, der=2) * dr_dR * dr_dZ + self.psi_r(r, der=1) * d2r_dRdZ + out = (d2psi - g) * c_R * c_Z else: raise NotImplementedError( "Only combinations (dR=0, dZ=0), (dR=1, dZ=0), (dR=0, dZ=1), (dR=2, dZ=0), (dR=0, dZ=2) and (dR=1, dZ=1) possible!", @@ -1588,24 +1590,26 @@ def psi(self, R, Z, dR=0, dZ=0): if dR == 0 and dZ == 0: out = self.psi_r(r, der=0) else: - dr_dR = (R - self.params["R0"]) / r - dr_dZ = Z / r - - d2r_dR2 = (r - (R - self.params["R0"]) * dr_dR) / r**2 - d2r_dZ2 = (r - Z * dr_dZ) / r**2 - - d2r_dRdZ = -Z * (R - self.params["R0"]) / r**3 + # regular on the magnetic axis r=0: with g = psi'(r)/r -> psi''(0) and the direction cosines + # c_R = (R - R0)/r, c_Z = Z/r (set to 0 on the axis), the derivatives below have the correct limits + on_axis = r == 0.0 + r_safe = xp.where(on_axis, 1.0, r) + dpsi = self.psi_r(r, der=1) + d2psi = self.psi_r(r, der=2) + g = xp.where(on_axis, d2psi, dpsi / r_safe) + c_R = (R - self.params["R0"]) / r_safe + c_Z = Z / r_safe if dR == 1 and dZ == 0: - out = self.psi_r(r, der=1) * dr_dR + out = g * (R - self.params["R0"]) elif dR == 0 and dZ == 1: - out = self.psi_r(r, der=1) * dr_dZ + out = g * Z elif dR == 2 and dZ == 0: - out = self.psi_r(r, der=2) * dr_dR**2 + self.psi_r(r, der=1) * d2r_dR2 + out = d2psi * c_R**2 + g * (1.0 - c_R**2) elif dR == 0 and dZ == 2: - out = self.psi_r(r, der=2) * dr_dZ**2 + self.psi_r(r, der=1) * d2r_dZ2 + out = d2psi * c_Z**2 + g * (1.0 - c_Z**2) elif dR == 1 and dZ == 1: - out = self.psi_r(r, der=2) * dr_dR * dr_dZ + self.psi_r(r, der=1) * d2r_dRdZ + out = (d2psi - g) * c_R * c_Z else: raise NotImplementedError( "Only combinations (dR=0, dZ=0), (dR=1, dZ=0), (dR=0, dZ=1), (dR=2, dZ=0), (dR=0, dZ=2) and (dR=1, dZ=1) possible!", diff --git a/src/struphy/kernel_arguments/pusher_args.cuh b/src/struphy/kernel_arguments/pusher_args.cuh index 59815fced..e3a547e1e 100644 --- a/src/struphy/kernel_arguments/pusher_args.cuh +++ b/src/struphy/kernel_arguments/pusher_args.cuh @@ -18,6 +18,7 @@ struct MarkerArgs { int first_free_idx; int mu_idx; long long* bc_type; + bool pole; }; struct DerhamArgs { diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index 08fefff46..feac42098 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -45,6 +45,7 @@ class CudaMarkerArguments(CudaStructArguments): ("first_free_idx", "int"), ("mu_idx", "int"), ("bc_type", "long long*"), + ("pole", "bool"), ) def __init__( @@ -61,6 +62,7 @@ def __init__( first_free_idx: int, mu_idx: int, bc_type, + pole: bool, ): self.markers = _device_array("markers", markers, np.float64) if markers.ndim != 2: @@ -77,6 +79,7 @@ def __init__( self.first_free_idx = first_free_idx self.mu_idx = mu_idx self.bc_type = _device_array("bc_type", bc_type, np.int64) + self.pole = pole self.pack() diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index 2e66dc58e..963bbbacb 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -45,6 +45,10 @@ class MarkerArguments: Kinetic boundary condition in each logical direction (0: periodic, 1: reflect, 2: remove, 3: refill (handled in Python)), see :func:`~struphy.pic.pushing.pusher_utilities_kernels.apply_kinetic_bc_marker`. + + pole : bool + Whether the domain has a pole at eta1=0 (markers then move through the axis), + see :func:`~struphy.pic.pushing.pusher_utilities_kernels.update_position_stage`. """ def __init__( @@ -61,6 +65,7 @@ def __init__( first_free_idx: int, mu_idx: int, bc_type: "int[:]", + pole: bool, ): self.markers = markers self.valid_mks = valid_mks @@ -77,6 +82,7 @@ def __init__( self.first_free_idx = first_free_idx # index after which auxiliary saving is possible self.mu_idx = mu_idx # particle magnetic moment self.bc_type = bc_type # kinetic boundary conditions per axis + self.pole = pole # pole at eta1=0 class DerhamArguments: diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index bc0e62c12..9293d2c79 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -647,7 +647,6 @@ def eval_psic(self, *coords, use_cbufs: bool = True): cbufs = self.cbufs if use_cbufs else None - a1 = self.equil.domain.params["a1"] B0 = self.equil.params["B0"] R0 = self.equil.params["R0"] diff --git a/src/struphy/models/guiding_center.py b/src/struphy/models/guiding_center.py index 0c9ac633d..2e067a5a1 100644 --- a/src/struphy/models/guiding_center.py +++ b/src/struphy/models/guiding_center.py @@ -117,8 +117,8 @@ def _compute_en_fB(self): particles = self.kinetic_ions.var.particles particles.save_magnetic_background_energy() # The weights (column 5) already carry the 1/N_p of the Monte-Carlo estimate, see KineticEnergyPIC. - energy = particles.markers[~particles.valid_mks, 5].dot( - particles.markers[~particles.valid_mks, 8], + energy = particles.markers[particles.valid_mks, 5].dot( + particles.markers[particles.valid_mks, 8], ) return energy diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 5e1d217b2..132c908ec 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -63,6 +63,9 @@ 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 + def _to_numpy_for_kernel(value): """Convert CuPy arrays to NumPy for compiled kernel calls.""" @@ -344,6 +347,18 @@ def __init__( 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) @@ -1399,6 +1414,11 @@ def draw_markers( counter += 1 + # keep markers off a pole at eta1=0, where the mapping is singular + if self._pole: + eta1 = self._markers[:n_mks_load_loc, 0] + eta1[eta1 < POLE_ETA1_MIN] = POLE_ETA1_MIN + # check if all particle positions are inside the unit cube [0, 1]^3 n_mks_load_loc = self.n_mks_load[self._mpi_rank] @@ -1826,6 +1846,14 @@ def apply_kinetic_bc(self, newton=False): 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: + crossed = self._find_outside_particles(0) + self.markers[crossed, 0] *= -1.0 + self.markers[crossed, 1] = (self.markers[crossed, 1] + 0.5) % 1.0 + near_pole = xp.logical_and(self.valid_mks, self.markers[:, 0] < POLE_ETA1_MIN) + self.markers[near_pole, 0] = POLE_ETA1_MIN + # apply boundary conditions if self._remove_axes: # extract the (n_rows, 3) logical coordinates once per contiguous @@ -1906,6 +1934,9 @@ def apply_kinetic_bc(self, newton=False): 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)) @@ -2576,6 +2607,7 @@ def _allocate_marker_array(self, dry_run: bool = False): self.first_free_idx, self.mu_idx, self._bc_type, + self._pole, ) def _initialize_sorting_boxes(self): diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py index e5f0245f1..56996db6f 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py @@ -44,7 +44,6 @@ def push_eta_stage( markers = args_markers.markers n_markers = args_markers.n_markers first_pusher_idx = args_markers.first_pusher_idx - first_free_idx = args_markers.first_free_idx if stage == n_stages - 1: last = 1.0 @@ -80,15 +79,8 @@ def push_eta_stage( # pull-back of velocity linalg_kernels.matrix_vector(dfinv, v, k) - # accumulation for last stage - markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k - - # update positions for intermediate stages or last stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] - + dt * a[stage] * k - + last * markers[ip, first_free_idx : first_free_idx + 3] - ) + # accumulation for last stage and update of positions (regular at a pole) + pusher_utilities_kernels.update_position_stage(ip, k, dt, a[stage], b[stage], last, args_markers) # apply kinetic boundary conditions pusher_utilities_kernels.apply_kinetic_bc_marker(ip, args_markers, args_domain, False) diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py index 03ac967b6..138e57fc6 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py @@ -201,16 +201,11 @@ def push_gc_Bstar_explicit_multistage( k_v = linalg_kernels.scalar_dot(b_star, e_star) k_v /= b_star_parallel * epsilon - # accumulation for last stage - markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k - markers[ip, first_free_idx + 3] += dt * b[stage] * k_v + # accumulation for last stage and update of positions (regular at a pole) + pusher_utilities_kernels.update_position_stage(ip, k, dt, a[stage], b[stage], last, args_markers) - # update positions for intermediate stages or last stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] - + dt * a[stage] * k - + last * markers[ip, first_free_idx : first_free_idx + 3] - ) + # same for v_parallel + markers[ip, first_free_idx + 3] += dt * b[stage] * k_v markers[ip, 3] = ( markers[ip, first_pusher_idx + 3] + dt * a[stage] * k_v + last * markers[ip, first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py index 905bfb8e0..07dedcdec 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py @@ -67,7 +67,6 @@ def push_gc_bxEstar_explicit_multistage( n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx first_pusher_idx = args_markers.first_pusher_idx - first_free_idx = args_markers.first_free_idx # get number of stages n_stages = shape(b)[0] @@ -166,15 +165,8 @@ def push_gc_bxEstar_explicit_multistage( k[:] = Exb / b_star_parallel - # accumulation for last stage - markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k - - # update positions for intermediate stages or last stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] - + dt * a[stage] * k - + last * markers[ip, first_free_idx : first_free_idx + 3] - ) + # accumulation for last stage and update of positions (regular at a pole) + pusher_utilities_kernels.update_position_stage(ip, k, dt, a[stage], b[stage], last, args_markers) # apply kinetic boundary conditions pusher_utilities_kernels.apply_kinetic_bc_marker(ip, args_markers, args_domain, False) diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index 48b06ef19..bde5e205b 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -37,8 +37,9 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, * @param args_domain Mapping arguments needed for velocity reflection. * @param newton Accumulate periodic shifts if true; overwrite them otherwise. * - * Call after updating the position. Boundary type 0 wraps periodically, - * type 1 mirrors position and logical velocity and sets first_pusher_idx to -1, + * Call after updating the position. With a pole at eta1=0, a marker with + * eta1 < 0 is mapped to (-eta1, eta2 + 1/2) first. Boundary type 0 wraps periodically, + * type 1 mirrors position and logical velocity (for vdim=3) and sets first_pusher_idx to -1, * type 2 clears all columns except the ID to -1, and type 3 is handled on * the host. The marker view carries shape and strides; j is a CUDA-only * column index replacing Pyccel array slices. @@ -49,6 +50,16 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma int first_pusher_idx = args_markers.first_pusher_idx; int first_shift_idx = args_markers.first_shift_idx; + // Pole at eta1=0: a marker with eta1 < 0 has crossed the axis. + if (args_markers.pole) { + if (args_markers.markers(ip, 0) < 0.) { + args_markers.markers(ip, 0) = -args_markers.markers(ip, 0); + double e2 = args_markers.markers(ip, 1) + 0.5; + args_markers.markers(ip, 1) = e2 - floor(e2); + } + if (args_markers.markers(ip, 0) < 1e-10) args_markers.markers(ip, 0) = 1e-10; + } + // Remove markers before applying periodic or reflecting boundaries. for (int axis = 0; axis < 3; ++axis) { if (bc_type[axis] == 2 && @@ -93,6 +104,8 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma } if (n_reflected == 0) return; args_markers.markers(ip, first_pusher_idx) = -1.; + // Only Cartesian velocities are reflected (columns 3:6 hold other data for vdim != 3). + if (args_markers.vdim != 3) return; for (int axis = 0; axis < 3; ++axis) { if (reflected[axis]) reflect_velocity(ip, args_markers.markers, args_domain, axis); } diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.py b/src/struphy/pic/pushing/pusher_utilities_kernels.py index b21c1bfb0..160123840 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.py +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.py @@ -1,4 +1,4 @@ -from numpy import empty, floor, mod, shape, sqrt, zeros +from numpy import arctan2, cos, empty, floor, mod, pi, shape, sin, sqrt, zeros from pyccel.decorators import pure, stack_array import struphy.bsplines.bsplines_kernels as bsplines_kernels @@ -22,12 +22,18 @@ def apply_kinetic_bc_marker( to be called at the end of a pusher kernel's particle loop, after the position update. Same result as :meth:`~struphy.pic.base.Particles.apply_kinetic_bc` for this marker. + If the domain has a pole at :math:`\eta_1=0` (``args_markers.pole``), a marker with :math:`\eta_1 < 0` + has crossed the axis and is mapped to :math:`(-\eta_1, \eta_2 + 1/2)`, with unchanged velocities; + :math:`\eta_1` is kept above ``1e-10``, where the mapping is regular. + The boundary condition of axis 0 then applies at :math:`\eta_1 = 1` only. + The boundary condition of each axis is read from ``args_markers.bc_type``: * 0 (periodic): wrap :math:`\eta_i` into :math:`[0, 1)` and store the shift :math:`\pm 1` at ``first_shift_idx + axis`` (set for ``newton=False``, accumulated for ``newton=True``). * 1 (reflect): mirror :math:`\eta_i` at the boundary, flip the logical velocity component - and mark the marker as done for multi-stage/iterative pushers (``first_pusher_idx = -1``). + (only for Cartesian velocities, ``vdim=3``) and mark the marker as done for + multi-stage/iterative pushers (``first_pusher_idx = -1``). * 2 (remove): turn the marker into a hole (all columns except the ID set to -1). * 3 (refill): skipped here, handled by :meth:`~struphy.pic.base.Particles.apply_kinetic_bc`. @@ -51,6 +57,14 @@ def apply_kinetic_bc_marker( first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx + # pole at eta1=0 + if args_markers.pole: + if markers[ip, 0] < 0.0: + markers[ip, 0] = -markers[ip, 0] + markers[ip, 1] = mod(markers[ip, 1] + 0.5, 1.0) + if markers[ip, 0] < 1e-10: + markers[ip, 0] = 1e-10 + # remove for axis in range(3): if bc_type[axis] == 2: @@ -96,6 +110,10 @@ def apply_kinetic_bc_marker( markers[ip, first_pusher_idx] = -1.0 + # only Cartesian velocities are reflected (columns 3:6 hold other data for vdim != 3) + if args_markers.vdim != 3: + return + # reflect velocities (Jacobian at the reflected position) dfm = empty((3, 3), dtype=float) dfinv = empty((3, 3), dtype=float) @@ -124,3 +142,98 @@ def apply_kinetic_bc_marker( linalg_kernels.matrix_vector(dfm, v_logical, v) markers[ip, 3:6] = v[:] + + +def update_position_stage( + ip: int, + k: "float[:]", + dt: float, + a_stage: float, + b_stage: float, + last: float, + args_markers: "MarkerArguments", +): + r"""Position update of one stage of an explicit Runge-Kutta pusher for the marker ``markers[ip]``, + to be called after the logical velocity :math:`\mathbf k = \dot{\boldsymbol \eta}` has been evaluated at + the current position ``markers[ip, 0:3]``. + + The stage contribution :math:`\Delta t\, b_\textrm{stage}\, \mathbf k` is accumulated at + ``first_free_idx``, and the new position is + + .. math:: + + \boldsymbol \eta = \boldsymbol \eta^n + \Delta t\, a_\textrm{stage}\, \mathbf k + \textrm{last} \sum_\textrm{stages} \Delta t\, b\, \mathbf k\,, + + with :math:`\boldsymbol \eta^n` at ``first_pusher_idx``. + + If the domain has a pole at :math:`\eta_1=0` (``args_markers.pole``), the update is done in the coordinates + + .. math:: + + (\xi_1, \xi_2, \eta_3) = \big(\eta_1 \cos(2\pi\eta_2),\, \eta_1 \sin(2\pi\eta_2),\, \eta_3\big)\,, + + in which the mapping (and thus the velocity :math:`\dot{\boldsymbol \xi} = \partial \boldsymbol \xi / \partial \boldsymbol \eta\, \mathbf k`) + is regular at the pole, and markers move through the axis. The accumulated stages are then :math:`\dot{\boldsymbol \xi}`, + and the position is transformed back with :math:`\eta_1 = |\boldsymbol \xi|` (kept above ``1e-10``) + and :math:`\eta_2 = \arctan(\xi_2 / \xi_1) / 2\pi \in [0, 1)`. + + Parameters + ---------- + ip : int + Row index of the marker. + + k : array[float] + Logical velocity at the current position. + + dt : float + Time step. + + a_stage : float + Butcher coefficient of the position update of this stage. + + b_stage : float + Butcher weight of this stage. + + last : float + 1.0 for the last stage (adds the accumulated stages), 0.0 otherwise. + + args_markers : MarkerArguments + Marker arguments. + """ + + markers = args_markers.markers + first_pusher_idx = args_markers.first_pusher_idx + first_free_idx = args_markers.first_free_idx + + if not args_markers.pole: + # accumulation for last stage + markers[ip, first_free_idx : first_free_idx + 3] += dt * b_stage * k + + # update positions for intermediate stages or last stage + markers[ip, 0:3] = ( + markers[ip, first_pusher_idx : first_pusher_idx + 3] + + dt * a_stage * k + + last * markers[ip, first_free_idx : first_free_idx + 3] + ) + return + + # velocity in the regular coordinates (xi1, xi2), at the current position + theta = 2.0 * pi * markers[ip, 1] + k_xi1 = cos(theta) * k[0] - 2.0 * pi * markers[ip, 0] * sin(theta) * k[1] + k_xi2 = sin(theta) * k[0] + 2.0 * pi * markers[ip, 0] * cos(theta) * k[1] + + # accumulation for last stage + markers[ip, first_free_idx] += dt * b_stage * k_xi1 + markers[ip, first_free_idx + 1] += dt * b_stage * k_xi2 + markers[ip, first_free_idx + 2] += dt * b_stage * k[2] + + # update positions for intermediate stages or last stage + theta = 2.0 * pi * markers[ip, first_pusher_idx + 1] + xi1 = markers[ip, first_pusher_idx] * cos(theta) + dt * a_stage * k_xi1 + last * markers[ip, first_free_idx] + xi2 = markers[ip, first_pusher_idx] * sin(theta) + dt * a_stage * k_xi2 + last * markers[ip, first_free_idx + 1] + + markers[ip, 0] = sqrt(xi1**2 + xi2**2) + if markers[ip, 0] < 1e-10: + markers[ip, 0] = 1e-10 + markers[ip, 1] = mod(arctan2(xi2, xi1) / (2.0 * pi), 1.0) + markers[ip, 2] = markers[ip, first_pusher_idx + 2] + dt * a_stage * k[2] + last * markers[ip, first_free_idx + 2] diff --git a/src/struphy/pic/tests/test_accum_gc_gradB.py b/src/struphy/pic/tests/test_accum_gc_gradB.py index aa61079ea..9950cae73 100644 --- a/src/struphy/pic/tests/test_accum_gc_gradB.py +++ b/src/struphy/pic/tests/test_accum_gc_gradB.py @@ -42,7 +42,7 @@ def rand(space): markers[:, 4] = rng.uniform(0.1, 1.0, n_mk) markers[:, 5] = rng.uniform(0.5, 1.5, n_mk) args_markers = MarkerArguments( - markers, np.ones(n_mk, dtype=bool), n_mk, 2, 5, 6, 10, 20, 25, 30, 4, np.zeros(3, dtype=int) + markers, np.ones(n_mk, dtype=bool), n_mk, 2, 5, 6, 10, 20, 25, 30, 4, np.zeros(3, dtype=int), False ) space = "H1vec" if basis_u == 0 else "Hdiv" diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 41e64dade..2f936bd87 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -135,7 +135,9 @@ def test_cuda_classes_reject_host_arrays(): with cunumpy.use_backend("numpy"): markers = np.zeros((10, N_COLS)) with pytest.raises(TypeError, match="CuPy array"): - CudaMarkerArguments(markers, np.ones(10, dtype=bool), 10, *MARKER_INDICES, np.zeros(3, dtype=np.int64)) + CudaMarkerArguments( + markers, np.ones(10, dtype=bool), 10, *MARKER_INDICES, np.zeros(3, dtype=np.int64), False + ) @requires_cupy diff --git a/src/struphy/pic/tests/test_kernel_setup.py b/src/struphy/pic/tests/test_kernel_setup.py index 2cf5ede92..00b9d127c 100644 --- a/src/struphy/pic/tests/test_kernel_setup.py +++ b/src/struphy/pic/tests/test_kernel_setup.py @@ -38,7 +38,9 @@ def marker_args(): markers[0, 8:14] = markers[0, :6] markers[0, 11] = 1.0 # Initial parallel velocity. markers[1, 0] = -1.0 # A hole must never be written to. - return MarkerArguments(markers, np.array([True, False]), 1, 3, 6, 7, 8, 14, 17, 18, 4, np.zeros(3, dtype=int)) + return MarkerArguments( + markers, np.array([True, False]), 1, 3, 6, 7, 8, 14, 17, 18, 4, np.zeros(3, dtype=int), False + ) @pytest.fixture @@ -163,7 +165,9 @@ def test_gc_discrete_gradient_residual_at_zero_velocity(derham_args, kernel, n_f markers[0, :3] = (0.4, 0.3, 0.2) markers[0, 8:11] = markers[0, :3] markers[0, 18] = markers[0, 21] = 1.0 # B*_parallel and B* = (0, 0, 1) at time n. - args_markers = MarkerArguments(markers, np.array([True]), 1, 2, 5, 6, 8, 12, 16, 17, 4, np.zeros(3, dtype=int)) + args_markers = MarkerArguments( + markers, np.array([True]), 1, 2, 5, 6, 8, 12, 16, 17, 4, np.zeros(3, dtype=int), False + ) kernel(0.1, 0, args_markers, Cuboid().args_domain, derham_args, 1.0, *(coefficients(1.0),) * n_fields, False) assert markers[0, 16] == 0.0 np.testing.assert_array_equal(markers[0, :4], (0.4, 0.3, 0.2, 0.0)) diff --git a/src/struphy/pic/tests/test_pole.py b/src/struphy/pic/tests/test_pole.py new file mode 100644 index 000000000..36eb46aab --- /dev/null +++ b/src/struphy/pic/tests/test_pole.py @@ -0,0 +1,369 @@ +"""Markers on domains with a pole at eta1=0 (no hole around the axis). + +Run the MPI tests with:: + + mpirun -n 4 pytest -p no:xdist --with-mpi src/struphy/pic/tests/test_pole.py +""" + +import numpy as np +import pytest +from maybempi import MPI + +from struphy.pic.pushing.pusher_utilities_kernels import apply_kinetic_bc_marker + +# column layout of the hand-made marker arrays below (as MarkerArguments indices) +N_COLS = 25 +FIRST_PUSHER_IDX = 8 +FIRST_SHIFT_IDX = 14 +RESIDUAL_IDX = 17 +FIRST_FREE_IDX = 18 + + +def _marker_args(markers, vdim, bc_type, pole): + from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments + + n_mks = markers.shape[0] + return MarkerArguments( + markers, + np.ones(n_mks, dtype=bool), + n_mks, + vdim, + 3 + vdim, + 7, + FIRST_PUSHER_IDX, + FIRST_SHIFT_IDX, + RESIDUAL_IDX, + FIRST_FREE_IDX, + 4, + np.array(bc_type, dtype=int), + pole, + ) + + +@pytest.mark.mpi_skip +def test_bc_kernel_axis_map(): + """A marker with eta1 < 0 is mapped through the axis; eta1 is kept off the pole.""" + from struphy import domains + + domain = domains.HollowCylinder(a1=0.0, a2=1.0, Lz=4.0) + + markers = np.zeros((3, N_COLS)) + markers[:, -1] = np.arange(3) + markers[0, :6] = (-0.2, 0.3, 0.5, 0.1, 0.2, 0.3) + markers[1, :6] = (-0.1, 0.7, 0.5, 0.1, 0.2, 0.3) + markers[2, :6] = (0.0, 0.7, 0.5, 0.1, 0.2, 0.3) + args_markers = _marker_args(markers, 3, (2, 0, 0), True) + + for ip in range(3): + apply_kinetic_bc_marker(ip, args_markers, domain.args_domain, False) + + assert np.allclose(markers[0, :3], (0.2, 0.8, 0.5)) + assert np.allclose(markers[1, :3], (0.1, 0.2, 0.5)) + assert markers[2, 0] == 1e-10 + # Cartesian velocities are unchanged, the marker is not removed + assert np.all(markers[:, 3:6] == (0.1, 0.2, 0.3)) + + # without a pole, the same marker is removed + markers[0, :6] = (-0.2, 0.3, 0.5, 0.1, 0.2, 0.3) + apply_kinetic_bc_marker(0, _marker_args(markers, 3, (2, 0, 0), False), domain.args_domain, False) + assert np.all(markers[0, :-1] == -1.0) + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("vdim", [2, 3]) +def test_bc_kernel_reflect_vdim(vdim): + """Reflection mirrors the position; only Cartesian velocities (vdim=3) are reflected.""" + from struphy import domains + + domain = domains.HollowCylinder(a1=0.2, a2=1.0, Lz=4.0) + + markers = np.zeros((1, N_COLS)) + markers[0, :6] = (1.1, 0.3, 0.5, 0.4, 0.2, 0.7) + apply_kinetic_bc_marker(0, _marker_args(markers, vdim, (1, 0, 0), False), domain.args_domain, False) + + assert np.isclose(markers[0, 0], 0.9) + assert markers[0, FIRST_PUSHER_IDX] == -1.0 + if vdim == 3: + # the radial velocity is reversed, the speed is kept (orthogonal mapping) + assert not np.allclose(markers[0, 3:6], (0.4, 0.2, 0.7)) + assert np.isclose(np.linalg.norm(markers[0, 3:6]), np.linalg.norm((0.4, 0.2, 0.7))) + else: + # v_parallel, mu and the weight are unchanged + assert np.all(markers[0, 3:6] == (0.4, 0.2, 0.7)) + + +def _particles_on_pole_domain(comm, domain, ppc=10, num_elements=(8, 9, 4), bufsize=0.25): + from struphy import BoundaryParameters, LoadingParameters + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.pic.particles import Particles6D + from struphy.topology.grids import TensorProductGrid + + derham = Derham(TensorProductGrid(num_elements=num_elements), DerhamOptions(degree=(2, 2, 1)), comm=comm) + domain_decomp = (derham.domain_array, derham.domain_decomposition.nprocs) + + particles = Particles6D( + comm_world=comm, + domain_decomp=domain_decomp, + loading_params=LoadingParameters(ppc=ppc, seed=1234, moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0)), + boundary_params=BoundaryParameters(bc=("remove", "periodic", "periodic")), + bufsize=bufsize, + domain=domain, + ) + particles.draw_markers() + particles.mpi_sort_markers() + return particles + + +def _by_unique_id(comm, ids, values): + """Dictionary marker ID -> values over all ranks, without IDs that occur more than once + (pseudo-random loading can give duplicate IDs at rank boundaries).""" + all_ids = np.concatenate(comm.allgather(ids)) + all_values = np.concatenate(comm.allgather(values)) + unique, counts = np.unique(all_ids, return_counts=True) + keep = set(unique[counts == 1]) + return {i: v for i, v in zip(all_ids, all_values) if i in keep} + + +@pytest.mark.mpi(min_size=1) +def test_python_bc_axis_map(): + """Particles.apply_kinetic_bc maps markers with eta1 < 0 through the axis.""" + from struphy import domains + + comm = MPI.COMM_WORLD + domain = domains.HollowCylinder(a1=0.0, a2=1.0, Lz=4.0) + particles = _particles_on_pole_domain(comm, domain) + + assert particles.args_markers.pole + assert np.all(particles.positions[:, 0] >= 1e-10) + + valid = np.nonzero(particles.valid_mks)[0][:3] + markers = particles.markers + markers[valid, 0] = -0.25 + markers[valid, 1] = 0.6 + n_valid = np.count_nonzero(particles.valid_mks) + + particles.apply_kinetic_bc() + + assert np.allclose(markers[valid, 0], 0.25) + assert np.allclose(markers[valid, 1], 0.1) + assert np.count_nonzero(particles.valid_mks) == n_valid + assert particles.n_lost_markers == 0 + + +@pytest.mark.mpi(min_size=1) +@pytest.mark.parametrize("algo", ["forward_euler", "heun2", "rk4"]) +def test_push_eta_through_pole(algo): + """Straight lines through the axis of a cylinder without hole are exact, since the mapping is linear in xi.""" + from struphy import domains + from struphy.ode.utils import ButcherTableau + from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage + from struphy.pic.pushing.pusher import Pusher + + comm = MPI.COMM_WORLD + a2 = 1.0 + domain = domains.HollowCylinder(a1=0.0, a2=a2, Lz=4.0) + # all markers cross the axis, i.e. most of them change the process + particles = _particles_on_pole_domain(comm, domain, bufsize=2.0) + markers = particles.markers + + # aim all markers at (a point next to) the axis, such that they cross it and end inside the disc + dt, n_steps = 0.1, 10 + T = dt * n_steps + valid = particles.valid_mks + x0, y0, z0 = domain(markers[valid, :3].copy()) + markers[valid, 3] = (-1.8 * x0 + 0.01) / T + markers[valid, 4] = (-1.8 * y0 - 0.02) / T + markers[valid, 5] = 0.3 + + # expected positions per marker ID (markers change rank while being pushed) + expected = np.stack((x0 + markers[valid, 3] * T, y0 + markers[valid, 4] * T, z0 + markers[valid, 5] * T), axis=1) + expected_by_id = _by_unique_id(comm, markers[valid, -1], expected) + n_valid_0 = comm.allreduce(np.count_nonzero(valid)) + + butcher = ButcherTableau(algo=algo) + pusher = Pusher( + particles, + push_eta_stage, + (butcher.a_stage, butcher.b, butcher.c, butcher.n_stages), + domain.args_domain, + alpha_in_kernel=1.0, + n_stages=butcher.n_stages, + pushes_eta=True, + local_eval_only=True, + ) + for _ in range(n_steps): + pusher(dt) + + valid = particles.valid_mks + assert comm.allreduce(np.count_nonzero(valid)) == n_valid_0 + assert particles.n_lost_markers == 0 + + ids = markers[valid, -1] + compare = np.array([i in expected_by_id for i in ids], dtype=bool) + assert comm.allreduce(np.count_nonzero(compare)) > 0.9 * n_valid_0 + x, y, z = domain(markers[valid, :3][compare].copy()) + x_expected, y_expected, z_expected = np.array([expected_by_id[i] for i in ids[compare]]).T + assert np.allclose(x, x_expected, atol=1e-10) + assert np.allclose(y, y_expected, atol=1e-10) + # z is periodic with period Lz + assert np.allclose(np.mod(z - z_expected + 2.0, 4.0) - 2.0, 0.0, atol=1e-10) + + +# grid of the simulation tests: the number of eta2 elements must be a multiple of 3 for polar splines, +# and the IGAPolar* mappings must have the same number of elements and degree in (eta1, eta2) +NUM_ELEMENTS = (8, 9, 4) +DEGREE = (3, 3, 1) + + +def _run_pole_simulation(tmp_folder, model, domain, equil, background, bc, dt, n_steps, comm): + from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + Time, + grids, + ) + + sim = Simulation( + model=model, + env=EnvironmentOptions(out_folders=tmp_folder, sim_folder="sim_pole"), + time_opts=Time(dt=dt, Tend=dt * n_steps), + domain=domain, + equil=equil, + grid=grids.TensorProductGrid(num_elements=NUM_ELEMENTS), + derham_opts=DerhamOptions(degree=DEGREE, bcs=(("free", "free"), None, None), polar_splines=True), + comm=comm, + ) + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=4, seed=4321), + boundary_params=BoundaryParameters(bc=bc), + # few markers per process: the discrete-gradient pushers (sorting at mid-points in each iteration) + # need a larger buffer, also on domains without a pole + bufsize=3.0, + ) + model.kinetic_ions.var.add_background(background) + return sim + + +def _read_scalars(sim, names): + import h5py + + with h5py.File(sim.data.file_path, "r") as file: + return {name: file["scalar/" + name][()] for name in names} + + +def _inner_density(*etas): + # called with (eta1, eta2, eta3) or with one (n, 3) array of positions + eta1 = etas[0] if len(etas) == 3 else etas[0][:, 0] + return np.where(eta1 < 0.6, 1.0, 0.0) + + +@pytest.mark.mpi(min_size=1) +def test_vlasov_through_pole(tmp_path_factory): + """Vlasov (PushVxB with polar splines and PushEta through the axis) conserves the kinetic energy.""" + from struphy import domains, equils, maxwellians + from struphy.models import Vlasov + + comm = MPI.COMM_WORLD + tmp_folder = comm.bcast(str(tmp_path_factory.mktemp("vlasov_pole")) if comm.Get_rank() == 0 else None) + + R0 = 5.0 + model = Vlasov() + sim = _run_pole_simulation( + tmp_folder, + model, + domains.IGAPolarCylinder(num_elements=NUM_ELEMENTS[:2], degree=DEGREE[:2], a=1.0, Lz=2 * np.pi * R0), + equils.ScrewPinch(a=1.0, R0=R0, B0=1.0), + # only markers that cannot reach the outer wall carry weight (reflection on a spline mapping changes |v|) + maxwellians.Maxwellian3D(n=(_inner_density, None), vth1=(0.1, None), vth2=(0.1, None), vth3=(0.1, None)), + ("reflect", "periodic", "periodic"), + dt=0.05, + n_steps=10, + comm=comm, + ) + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + model.propagators.push_eta.options = model.propagators.push_eta.Options() + sim.run() + + assert model.kinetic_ions.var.particles.n_lost_markers == 0 + if comm.Get_rank() == 0: + en = _read_scalars(sim, ["kinetic_energy"])["kinetic_energy"] + assert np.all(np.isfinite(en)) + assert np.allclose(en, en[0], rtol=1e-12) + + +def _gc_energies(sim, particles): + """Marker IDs, positions and guiding-center energies v_par^2/2 + mu |B0| of the valid markers, + with the projected (spline) |B0| that the pushers use.""" + particles.save_magnetic_background_energy() + markers = particles.markers + valid = particles.valid_mks + eta = markers[valid, :3].copy() + energy = 0.5 * markers[valid, 3] ** 2 + markers[valid, particles.first_diagnostics_idx] + return markers[valid, -1].copy(), eta, energy + + +@pytest.mark.mpi(min_size=1) +@pytest.mark.parametrize("algo", ["explicit", "discrete_gradient_1st_order"]) +def test_guiding_center_through_pole(tmp_path_factory, algo): + """GuidingCenter on a torus without hole, polar splines: markers drift through the axis. + + The marker energies are conserved (explicit pushers in the regular coordinates xi); + the discrete-gradient pushers only map markers through the axis and are less accurate near it.""" + from struphy import domains, equils, maxwellians + from struphy.models import GuidingCenter + + comm = MPI.COMM_WORLD + tmp_folder = comm.bcast(str(tmp_path_factory.mktemp(f"gc_pole_{algo}")) if comm.Get_rank() == 0 else None) + + # large epsilon: large drifts, such that markers cross the axis + model = GuidingCenter(epsilon=1.0) + sim = _run_pole_simulation( + tmp_folder, + model, + domains.IGAPolarTorus(num_elements=NUM_ELEMENTS[:2], degree=DEGREE[:2], a=1.0, R0=3.0, tor_period=1), + equils.AdhocTorus(a=1.0, R0=3.0, B0=1.0), + maxwellians.GyroMaxwellian2D(n=(1.0, None), vth_para=(0.5, None), vth_perp=(0.5, None)), + ("remove", "periodic", "periodic"), + dt=0.05, + n_steps=10, + comm=comm, + ) + model.propagators.push_bxe.options = model.propagators.push_bxe.Options(algo=algo) + model.propagators.push_parallel.options = model.propagators.push_parallel.Options(algo=algo) + sim.allocate() + + particles = model.kinetic_ions.var.particles + ids, eta0, energy0 = _gc_energies(sim, particles) + initial = _by_unique_id(comm, ids, np.concatenate((eta0, energy0[:, None]), axis=1)) + + dt = sim.time_opts.dt + for _ in range(10): + model.propagators.push_bxe(dt) + model.propagators.push_parallel(dt) + + ids, eta, energy = _gc_energies(sim, particles) + assert np.all(np.isfinite(eta)) and np.all(np.isfinite(energy)) + assert np.all(eta[:, 0] > 0.0) and np.all(eta[:, 0] <= 1.0) + + compare = np.array([i in initial for i in ids], dtype=bool) + start = np.array([initial[i] for i in ids[compare]]).reshape(-1, 4) + rel_err = np.abs(energy[compare] - start[:, 3]) / start[:, 3] + + # markers that went through the axis region: close to the axis and eta2 changed by about 1/2 + near_axis = (start[:, 0] < 0.1) & (eta[compare, 0] < 0.1) + jump = np.abs(np.mod(eta[compare, 1] - start[:, 1] + 0.5, 1.0) - 0.5) > 0.25 + crossed = near_axis & jump + assert comm.allreduce(np.count_nonzero(crossed)) > 0 + + max_err = comm.allreduce(np.max(rel_err, initial=0.0), op=MPI.MAX) + max_err_axis = comm.allreduce(np.max(rel_err[near_axis], initial=0.0), op=MPI.MAX) + if algo == "explicit": + assert max_err < 1e-3 + assert max_err_axis < 1e-4 + else: + # the discrete-gradient schemes evaluate at mid-points in eta, which are not regular at the axis + assert max_err < 1e-1 diff --git a/src/struphy/propagators/push_guiding_center_bx_estar.py b/src/struphy/propagators/push_guiding_center_bx_estar.py index 96c9aa89f..97234b63e 100644 --- a/src/struphy/propagators/push_guiding_center_bx_estar.py +++ b/src/struphy/propagators/push_guiding_center_bx_estar.py @@ -178,29 +178,28 @@ def allocate(self): self._absB0 = self.projected_equil.absB0 curl_unit_b_dot_b0 = self.projected_equil.curl_unit_b_dot_b0 + # the kernels need tensor-product coefficients: transposed extraction operator + # PolarVector --> StencilVector (identity map in case of no polar splines) + self._E0T = self.derham.extraction_ops["0"].transpose() + # magnetic perturbation if self.b_tilde is not None: self._B_dot_b = self.derham.V0.zeros() self._grad_b_full = self.derham.V1.zeros() self._PB = getattr(self.basis_ops, "PB") + self._B_dot_b_pol = self._PB.codomain.zeros() - B_dot_b = self._PB.dot(self.b_tilde.spline.vector, out=self._B_dot_b) - B_dot_b.update_ghost_regions() - - grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) - grad_b_full.update_ghost_regions() - - grad_b_full += self._gradB1 - B_dot_b += self._absB0 + self._update_b_tilde() else: self._grad_b_full = self._gradB1 self._B_dot_b = self._absB0 - # allocate electric field + # allocate electric potential (self._phi holds its tensor-product coefficients) if self._allocate_phi: self.phi.allocate(self.derham, self.domain) - self._phi = self.phi.spline.vector + self._phi = self.derham.V0.zeros() + self._update_phi() self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.V1.zeros() @@ -482,24 +481,37 @@ def allocate(self): pushes_eta=True, ) + def _update_phi(self): + """Extract the electric potential to the tensor-product space.""" + phi = self._E0T.dot(self.phi.spline.vector, out=self._phi) + phi.update_ghost_regions() + + def _update_b_tilde(self): + """Parallel magnetic perturbation and its gradient (plus equilibrium) in the tensor-product spaces.""" + B_dot_b_pol = self._PB.dot(self.b_tilde.spline.vector, out=self._B_dot_b_pol) + B_dot_b = self._E0T.dot(B_dot_b_pol, out=self._B_dot_b) + B_dot_b.update_ghost_regions() + + grad_b_full = self.derham.grad_tp.dot(B_dot_b, out=self._grad_b_full) + grad_b_full.update_ghost_regions() + + grad_b_full += self._gradB1 + B_dot_b += self._absB0 + @profile def __call__(self, dt): + # electric potential + self._update_phi() + # electric field # TODO: add out to __neg__ of StencilVector if self._evaluate_e_field: - e_field = self.derham.grad.dot(-self._phi, out=self._e_field) + e_field = self.derham.grad_tp.dot(-self._phi, out=self._e_field) e_field.update_ghost_regions() # magnetic perturbation if self.b_tilde is not None: - B_dot_b = self._PB.dot(self.b_tilde.spline.vector, out=self._B_dot_b) - B_dot_b.update_ghost_regions() - - grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) - grad_b_full.update_ghost_regions() - - grad_b_full += self._gradB1 - B_dot_b += self._absB0 + self._update_b_tilde() # call pusher self._pusher(dt) diff --git a/src/struphy/propagators/push_guiding_center_parallel.py b/src/struphy/propagators/push_guiding_center_parallel.py index cd90a538b..b388d7510 100644 --- a/src/struphy/propagators/push_guiding_center_parallel.py +++ b/src/struphy/propagators/push_guiding_center_parallel.py @@ -195,29 +195,28 @@ def allocate(self): self._absB0 = self.projected_equil.absB0 curl_unit_b_dot_b0 = self.projected_equil.curl_unit_b_dot_b0 + # the kernels need tensor-product coefficients: transposed extraction operator + # PolarVector --> StencilVector (identity map in case of no polar splines) + self._E0T = self.derham.extraction_ops["0"].transpose() + # magnetic perturbation if self.b_tilde is not None: self._B_dot_b = self.derham.V0.zeros() self._grad_b_full = self.derham.V1.zeros() self._PB = getattr(self.basis_ops, "PB") + self._B_dot_b_pol = self._PB.codomain.zeros() - B_dot_b = self._PB.dot(self.b_tilde.spline.vector, out=self._B_dot_b) - B_dot_b.update_ghost_regions() - - grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) - grad_b_full.update_ghost_regions() - - grad_b_full += self._gradB1 - B_dot_b += self._absB0 + self._update_b_tilde() else: self._grad_b_full = self._gradB1 self._B_dot_b = self._absB0 - # allocate electric field + # allocate electric potential (self._phi holds its tensor-product coefficients) if self._allocate_phi: self.phi.allocate(self.derham, domain=self.domain) - self._phi = self.phi.spline.vector + self._phi = self.derham.V0.zeros() + self._update_phi() self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.V1.zeros() @@ -508,24 +507,37 @@ def allocate(self): pushes_eta=True, ) + def _update_phi(self): + """Extract the electric potential to the tensor-product space.""" + phi = self._E0T.dot(self.phi.spline.vector, out=self._phi) + phi.update_ghost_regions() + + def _update_b_tilde(self): + """Parallel magnetic perturbation and its gradient (plus equilibrium) in the tensor-product spaces.""" + B_dot_b_pol = self._PB.dot(self.b_tilde.spline.vector, out=self._B_dot_b_pol) + B_dot_b = self._E0T.dot(B_dot_b_pol, out=self._B_dot_b) + B_dot_b.update_ghost_regions() + + grad_b_full = self.derham.grad_tp.dot(B_dot_b, out=self._grad_b_full) + grad_b_full.update_ghost_regions() + + grad_b_full += self._gradB1 + B_dot_b += self._absB0 + @profile def __call__(self, dt): + # electric potential + self._update_phi() + # electric field # TODO: add out to __neg__ of StencilVector if self._evaluate_e_field: - e_field = self.derham.grad.dot(-self._phi, out=self._e_field) + e_field = self.derham.grad_tp.dot(-self._phi, out=self._e_field) e_field.update_ghost_regions() # magnetic perturbation if self.b_tilde is not None: - B_dot_b = self._PB.dot(self.b_tilde.spline.vector, out=self._B_dot_b) - B_dot_b.update_ghost_regions() - - grad_b_full = self.derham.grad.dot(B_dot_b, out=self._grad_b_full) - grad_b_full.update_ghost_regions() - - grad_b_full += self._gradB1 - B_dot_b += self._absB0 + self._update_b_tilde() # call pusher self._pusher(dt) diff --git a/src/struphy/propagators/push_vxb.py b/src/struphy/propagators/push_vxb.py index 11c31476a..cfe13d330 100644 --- a/src/struphy/propagators/push_vxb.py +++ b/src/struphy/propagators/push_vxb.py @@ -108,7 +108,10 @@ def allocate(self): self._epsilon = self.variables.ions.species.equation_params.epsilon assert self.derham is not None, f"{self.__class__.__name__} needs a Derham object." - # TODO: treat PolarVector as well, but polar splines are being reworked at the moment + # transposed extraction operator PolarVector --> BlockVector (identity map in case of no polar splines) + self._E2T: LinearOperator = self.derham.extraction_ops["2"].transpose() + + # the projected equilibrium is already extracted to the tensor-product space if self.projected_equil is not None: self.b2_0 = self.projected_equil.b2 assert self.b2_0.space == self.derham.V2 @@ -116,7 +119,7 @@ def allocate(self): self.b2_0 = self.derham.V2.zeros() if self.b2_var is not None: - assert self.b2_var.spline.vector.space == self.derham.V2 + assert self.b2_var.spline.vector.space == self._E2T.domain # allocate dummy vectors to avoid temporary array allocations self._tmp = self.derham.V2.zeros() @@ -147,18 +150,12 @@ def allocate(self): pushes_eta=False, ) - # transposed extraction operator PolarVector --> BlockVector (identity map in case of no polar splines) - self._E2T: LinearOperator = self.derham.extraction_ops["2"].transpose() - @profile def __call__(self, dt): - # sum up total magnetic field - tmp = self.b2_0.copy(out=self._tmp) + # sum up total magnetic field in the tensor-product space (the variable is extracted first) + b_full: BlockVector = self.b2_0.copy(out=self._b_full) if self.b2_var is not None: - tmp += self.b2_var.spline.vector - - # extract coefficients to tensor product space - b_full: BlockVector = self._E2T.dot(tmp, out=self._b_full) + b_full += self._E2T.dot(self.b2_var.spline.vector, out=self._tmp) b_full.update_ghost_regions() b_full /= self._epsilon From 39fcf9366c319dc68dbbaa8da00582baa9ebcddf Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 12 Oct 2026 20:41:07 +0200 Subject: [PATCH 2/5] Pole support: review fixes - Particles.apply_kinetic_bc: the axis map took markers outside on both sides of eta1; a marker with eta1 > 1 was moved to the axis instead of getting the boundary condition at eta1=1. Only eta1 < 0 is mapped now (test added). - The eta1 clamp at the pole is applied after all loading branches (external, restart, ...). - GC propagators: the gradients are again taken with derham.grad (including the boundary operators) and then extracted to the tensor-product space, unchanged without polar splines. The GC pole test evaluates the electric field. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 13 ++++++------ src/struphy/pic/tests/test_pole.py | 21 +++++++++++-------- .../push_guiding_center_bx_estar.py | 13 ++++++++---- .../push_guiding_center_parallel.py | 13 ++++++++---- 4 files changed, 37 insertions(+), 23 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 132c908ec..9026258ba 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1414,17 +1414,17 @@ def draw_markers( counter += 1 - # keep markers off a pole at eta1=0, where the mapping is singular - if self._pole: - eta1 = self._markers[:n_mks_load_loc, 0] - eta1[eta1 < POLE_ETA1_MIN] = POLE_ETA1_MIN - # 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 + if self._initialized_sorting and sort: logger.info("\nSorting the markers after initial draw") if self.mpi_comm is not None: @@ -1848,7 +1848,8 @@ def apply_kinetic_bc(self, newton=False): # pole at eta1=0: markers with eta1 < 0 have crossed the axis (as in apply_kinetic_bc_marker) if self._pole: - crossed = self._find_outside_particles(0) + self._find_outside_particles(0) + crossed = xp.nonzero(self._is_outside_left)[0] self.markers[crossed, 0] *= -1.0 self.markers[crossed, 1] = (self.markers[crossed, 1] + 0.5) % 1.0 near_pole = xp.logical_and(self.valid_mks, self.markers[:, 0] < POLE_ETA1_MIN) diff --git a/src/struphy/pic/tests/test_pole.py b/src/struphy/pic/tests/test_pole.py index 36eb46aab..92b682e92 100644 --- a/src/struphy/pic/tests/test_pole.py +++ b/src/struphy/pic/tests/test_pole.py @@ -137,18 +137,21 @@ def test_python_bc_axis_map(): assert particles.args_markers.pole assert np.all(particles.positions[:, 0] >= 1e-10) - valid = np.nonzero(particles.valid_mks)[0][:3] + valid = np.nonzero(particles.valid_mks)[0][:4] markers = particles.markers - markers[valid, 0] = -0.25 - markers[valid, 1] = 0.6 + markers[valid[:3], 0] = -0.25 + markers[valid[:3], 1] = 0.6 + # outside at eta1=1: removed (boundary condition of axis 0), not mapped through the axis + markers[valid[3], 0] = 1.05 n_valid = np.count_nonzero(particles.valid_mks) particles.apply_kinetic_bc() - assert np.allclose(markers[valid, 0], 0.25) - assert np.allclose(markers[valid, 1], 0.1) - assert np.count_nonzero(particles.valid_mks) == n_valid - assert particles.n_lost_markers == 0 + assert np.allclose(markers[valid[:3], 0], 0.25) + assert np.allclose(markers[valid[:3], 1], 0.1) + assert np.all(markers[valid[3], :-1] == -1.0) + assert particles.n_lost_markers == 1 + assert np.count_nonzero(particles.valid_mks) == n_valid - 1 @pytest.mark.mpi(min_size=1) @@ -332,8 +335,8 @@ def test_guiding_center_through_pole(tmp_path_factory, algo): n_steps=10, comm=comm, ) - model.propagators.push_bxe.options = model.propagators.push_bxe.Options(algo=algo) - model.propagators.push_parallel.options = model.propagators.push_parallel.Options(algo=algo) + model.propagators.push_bxe.options = model.propagators.push_bxe.Options(algo=algo, evaluate_e_field=True) + model.propagators.push_parallel.options = model.propagators.push_parallel.Options(algo=algo, evaluate_e_field=True) sim.allocate() particles = model.kinetic_ions.var.particles diff --git a/src/struphy/propagators/push_guiding_center_bx_estar.py b/src/struphy/propagators/push_guiding_center_bx_estar.py index 97234b63e..f6fc1da6f 100644 --- a/src/struphy/propagators/push_guiding_center_bx_estar.py +++ b/src/struphy/propagators/push_guiding_center_bx_estar.py @@ -178,9 +178,10 @@ def allocate(self): self._absB0 = self.projected_equil.absB0 curl_unit_b_dot_b0 = self.projected_equil.curl_unit_b_dot_b0 - # the kernels need tensor-product coefficients: transposed extraction operator - # PolarVector --> StencilVector (identity map in case of no polar splines) + # the kernels need tensor-product coefficients: transposed extraction operators + # PolarVector --> Stencil-/BlockVector (identity maps in case of no polar splines) self._E0T = self.derham.extraction_ops["0"].transpose() + self._E1T = self.derham.extraction_ops["1"].transpose() # magnetic perturbation if self.b_tilde is not None: @@ -189,6 +190,7 @@ def allocate(self): self._PB = getattr(self.basis_ops, "PB") self._B_dot_b_pol = self._PB.codomain.zeros() + self._grad_b_pol = self.derham.grad.codomain.zeros() self._update_b_tilde() else: @@ -202,6 +204,7 @@ def allocate(self): self._update_phi() self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.V1.zeros() + self._e_field_pol = self.derham.grad.codomain.zeros() # choose method particles = self.variables.ions.particles @@ -492,7 +495,8 @@ def _update_b_tilde(self): B_dot_b = self._E0T.dot(B_dot_b_pol, out=self._B_dot_b) B_dot_b.update_ghost_regions() - grad_b_full = self.derham.grad_tp.dot(B_dot_b, out=self._grad_b_full) + grad_b_pol = self.derham.grad.dot(B_dot_b_pol, out=self._grad_b_pol) + grad_b_full = self._E1T.dot(grad_b_pol, out=self._grad_b_full) grad_b_full.update_ghost_regions() grad_b_full += self._gradB1 @@ -506,7 +510,8 @@ def __call__(self, dt): # electric field # TODO: add out to __neg__ of StencilVector if self._evaluate_e_field: - e_field = self.derham.grad_tp.dot(-self._phi, out=self._e_field) + e_field_pol = self.derham.grad.dot(-self.phi.spline.vector, out=self._e_field_pol) + e_field = self._E1T.dot(e_field_pol, out=self._e_field) e_field.update_ghost_regions() # magnetic perturbation diff --git a/src/struphy/propagators/push_guiding_center_parallel.py b/src/struphy/propagators/push_guiding_center_parallel.py index b388d7510..7aa53f662 100644 --- a/src/struphy/propagators/push_guiding_center_parallel.py +++ b/src/struphy/propagators/push_guiding_center_parallel.py @@ -195,9 +195,10 @@ def allocate(self): self._absB0 = self.projected_equil.absB0 curl_unit_b_dot_b0 = self.projected_equil.curl_unit_b_dot_b0 - # the kernels need tensor-product coefficients: transposed extraction operator - # PolarVector --> StencilVector (identity map in case of no polar splines) + # the kernels need tensor-product coefficients: transposed extraction operators + # PolarVector --> Stencil-/BlockVector (identity maps in case of no polar splines) self._E0T = self.derham.extraction_ops["0"].transpose() + self._E1T = self.derham.extraction_ops["1"].transpose() # magnetic perturbation if self.b_tilde is not None: @@ -206,6 +207,7 @@ def allocate(self): self._PB = getattr(self.basis_ops, "PB") self._B_dot_b_pol = self._PB.codomain.zeros() + self._grad_b_pol = self.derham.grad.codomain.zeros() self._update_b_tilde() else: @@ -219,6 +221,7 @@ def allocate(self): self._update_phi() self._evaluate_e_field = self.options.evaluate_e_field self._e_field = self.derham.V1.zeros() + self._e_field_pol = self.derham.grad.codomain.zeros() # choose method particles = self.variables.ions.particles @@ -518,7 +521,8 @@ def _update_b_tilde(self): B_dot_b = self._E0T.dot(B_dot_b_pol, out=self._B_dot_b) B_dot_b.update_ghost_regions() - grad_b_full = self.derham.grad_tp.dot(B_dot_b, out=self._grad_b_full) + grad_b_pol = self.derham.grad.dot(B_dot_b_pol, out=self._grad_b_pol) + grad_b_full = self._E1T.dot(grad_b_pol, out=self._grad_b_full) grad_b_full.update_ghost_regions() grad_b_full += self._gradB1 @@ -532,7 +536,8 @@ def __call__(self, dt): # electric field # TODO: add out to __neg__ of StencilVector if self._evaluate_e_field: - e_field = self.derham.grad_tp.dot(-self._phi, out=self._e_field) + e_field_pol = self.derham.grad.dot(-self.phi.spline.vector, out=self._e_field_pol) + e_field = self._E1T.dot(e_field_pol, out=self._e_field) e_field.update_ghost_regions() # magnetic perturbation From b3858dbac60869b6baaa0ccfce968db9a41945c1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 12 Oct 2026 21:37:59 +0200 Subject: [PATCH 3/5] Disc loading under MPI: markers on their own process, marker array large enough With spatial="disc" and more than one process, markers were lost or the run hung (not related to the pole): - pseudo_random loading kept the uniform draws on the domain of the process and applied eta1 -> sqrt(eta1) afterwards, which moved many markers to the domain of another process. Without sorting boxes they were not sent there, and the first push evaluated the fields outside the local domain (e.g. 15% of the GuidingCenter markers were removed in the first step on 4 processes). The transform is now applied before the domain check. - After loading, markers are sent to their process also without sorting boxes (e.g. Sobol loading draws on the whole domain on each process). - The number of markers per process at loading was estimated from the uniform measure; with the disc measure 2*eta1, processes at larger eta1 overflowed the marker array (ValueError on some processes, hang on the others). The estimate uses the disc measure now. - Test test_draw_disc_on_own_process (pseudo_random and sobol_standard). Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 25 +++++++++++--- src/struphy/pic/tests/test_draw_parallel.py | 38 +++++++++++++++++++++ 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 9026258ba..cd539effe 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1237,6 +1237,9 @@ def draw_markers( # Generate a chunk of random particles num_to_add_glob = min(chunk_size, int(self.Np) - num_loaded_particles_glob) temp = xp.random.rand(num_to_add_glob, 3 + self.vdim) + # 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=1))[0] @@ -1388,11 +1391,12 @@ def draw_markers( "Inverse transform sampling of given vdim is not implemented!", ) - # inversion method for drawing uniformly on the disc + # inversion method for drawing uniformly on the disc (done above for pseudo-random markers) if self.spatial == "disc": - self._markers[:n_mks_load_loc, 0] = xp.sqrt( - self._markers[:n_mks_load_loc, 0], - ) + if self.loading != "pseudo_random": + self._markers[:n_mks_load_loc, 0] = xp.sqrt( + self._markers[:n_mks_load_loc, 0], + ) else: assert self.spatial == "uniform", f'Spatial drawing must be "uniform" or "disc", is {self.spatial}.' @@ -1431,6 +1435,9 @@ def draw_markers( self.mpi_sort_markers() self.do_sort() logger.info("Done.") + elif self.mpi_comm is not None and self.mpi_size > 1: + # markers drawn on the domain of another process (e.g. Sobol loading) are sent to their process + self.mpi_sort_markers(apply_bc=False) def initialize_weights( self, @@ -2538,7 +2545,15 @@ def _n_mks_load_and_Np_per_clone(self): _n_mks_load_tot = self.Np _ppc = self.ppc - n_mks_load = self.gather_scalar_in_subcomm_array(int(_ppc * n_cells_loc)) + 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 = 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) diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index ff18a3321..90e0344fd 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -196,3 +196,41 @@ def test_marker_ids(loading, Np=1000): }, ], ) + + +@pytest.mark.parametrize("loading", ["pseudo_random", "sobol_standard"]) +def test_draw_disc_on_own_process(loading, ppc=50): + """Markers drawn uniformly on the disc are on their own process after `draw_markers` (without sorting boxes).""" + + import numpy as np + from maybempi import MPI + + from struphy import LoadingParameters, domains + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.pic.particles import Particles6D + from struphy.topology.grids import TensorProductGrid + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + domain = domains.HollowCylinder(a1=0.1, a2=1.0, Lz=4.0) + derham = Derham(TensorProductGrid(num_elements=(8, 9, 4)), DerhamOptions(degree=(1, 1, 1)), comm=comm) + domain_decomp = (derham.domain_array, derham.domain_decomposition.nprocs) + + particles = Particles6D( + comm_world=comm, + domain_decomp=domain_decomp, + loading_params=LoadingParameters(ppc=ppc, loading=loading, seed=1234, spatial="disc"), + domain=domain, + ) + particles.draw_markers() + + # all markers are on their own process and none is lost + positions = particles.positions + assert np.all(particles._is_on_domain_of_rank(positions, rank)) + assert comm.allreduce(positions.shape[0]) == particles.Np + + # uniform on the disc: eta1 has the density 2 eta1 with mean 2/3 and variance 1/18 + mean = comm.allreduce(np.sum(positions[:, 0])) / particles.Np + assert abs(mean - 2.0 / 3.0) < 5.0 * np.sqrt(1.0 / 18.0 / particles.Np) From e2c88ae3c320e132331d2583192dfe629e8703f7 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 12 Oct 2026 21:41:20 +0200 Subject: [PATCH 4/5] Disc loading: review fixes - The sort after loading only runs without sorting boxes and not for restart loading (the holes of the restart markers array are not updated at that point). - At least one marker per process in the disc estimate (avoids a zero-size estimate near the axis). - Test: reduce before asserting (no deadlock if one process fails); moved above the main block. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 12 ++++-- src/struphy/pic/tests/test_draw_parallel.py | 46 +++++++++++---------- 2 files changed, 33 insertions(+), 25 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index cd539effe..914d10498 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1435,8 +1435,14 @@ def draw_markers( self.mpi_sort_markers() self.do_sort() logger.info("Done.") - elif self.mpi_comm is not None and self.mpi_size > 1: - # markers drawn on the domain of another process (e.g. Sobol loading) are sent to their process + elif ( + not self._initialized_sorting + and self.loading != "restart" + and self.mpi_comm is not None + and self.mpi_size > 1 + ): + # without sorting boxes, markers drawn on the domain of another process (e.g. Sobol loading + # or the disc transform) are sent to their process self.mpi_sort_markers(apply_bc=False) def initialize_weights( @@ -2549,7 +2555,7 @@ def _n_mks_load_and_Np_per_clone(self): # 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 = int(_n_mks_load_tot * (d[1] ** 2 - d[0] ** 2) * (d[4] - d[3]) * (d[7] - d[6])) + 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) diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 90e0344fd..1bcd1fa88 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -177,27 +177,6 @@ def test_marker_ids(loading, Np=1000): assert n_at_origin == 0 -if __name__ == "__main__": - # test_draw([8, 9, 10], [2, 3, 4], [False, False, True], ['Cuboid', { - # 'l1': 1., 'r1': 2., 'l2': 10., 'r2': 20., 'l3': 100., 'r3': 200.}]) - test_draw( - [8, 9, 10], - [2, 3, 4], - [False, False, True], - [ - "Cuboid", - { - "l1": 0.0, - "r1": 1.0, - "l2": 0.0, - "r2": 1.0, - "l3": 0.0, - "r3": 1.0, - }, - ], - ) - - @pytest.mark.parametrize("loading", ["pseudo_random", "sobol_standard"]) def test_draw_disc_on_own_process(loading, ppc=50): """Markers drawn uniformly on the disc are on their own process after `draw_markers` (without sorting boxes).""" @@ -227,10 +206,33 @@ def test_draw_disc_on_own_process(loading, ppc=50): particles.draw_markers() # all markers are on their own process and none is lost + # (reduce before asserting, such that a failure on one process does not deadlock the others) positions = particles.positions - assert np.all(particles._is_on_domain_of_rank(positions, rank)) + on_own_process = comm.allreduce(bool(np.all(particles._is_on_domain_of_rank(positions, rank))), op=MPI.LAND) + assert on_own_process assert comm.allreduce(positions.shape[0]) == particles.Np # uniform on the disc: eta1 has the density 2 eta1 with mean 2/3 and variance 1/18 mean = comm.allreduce(np.sum(positions[:, 0])) / particles.Np assert abs(mean - 2.0 / 3.0) < 5.0 * np.sqrt(1.0 / 18.0 / particles.Np) + + +if __name__ == "__main__": + # test_draw([8, 9, 10], [2, 3, 4], [False, False, True], ['Cuboid', { + # 'l1': 1., 'r1': 2., 'l2': 10., 'r2': 20., 'l3': 100., 'r3': 200.}]) + test_draw( + [8, 9, 10], + [2, 3, 4], + [False, False, True], + [ + "Cuboid", + { + "l1": 0.0, + "r1": 1.0, + "l2": 0.0, + "r2": 1.0, + "l3": 0.0, + "r3": 1.0, + }, + ], + ) From 02a786be5bcc97cd15f26ea54254c863acb77388 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Tue, 13 Oct 2026 08:34:33 +0200 Subject: [PATCH 5/5] Pass pole to MarkerArguments in kernel test helpers The pole argument added to MarkerArguments/CudaMarkerArguments was missing in marker_arguments() (kernel_test_args.py) and make_arguments() (test_kernel_backends.py), which broke the CUDA emulation parity tests and the fake-CuPy accumulator test on CI. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 1 + src/struphy/pic/tests/test_kernel_backends.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index abe1ef445..d1deb316a 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -42,6 +42,7 @@ def marker_arguments(bc): N_MARKERS, *MARKER_INDICES, xp.asarray(bc, dtype=np.int64), + False, # Cuboid has no pole ) return args_markers, Cuboid().args_domain diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 2f936bd87..01cb000c3 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -48,7 +48,7 @@ def make_arguments(n_markers, seed=0): valid = cunumpy.asarray(rng.random(n_markers) > 0.1) bc = cunumpy.zeros(3, dtype=np.int64) args_class = CudaMarkerArguments if cunumpy.get_backend() == "cupy" else MarkerArguments - return args_class(markers, valid, n_markers, *MARKER_INDICES, bc), Cuboid().args_domain + return args_class(markers, valid, n_markers, *MARKER_INDICES, bc, False), Cuboid().args_domain def make_pusher(kernel):