Particles on domains with a pole (markers through the axis, polar splines in particle propagators)

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!",

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 {

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()
 
 

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:

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"]
 

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
 

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)
@@ -1224,6 +1239,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]
@@ -1380,11 +1398,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}.'
 
@@ -1412,10 +1431,16 @@ def draw_markers(
             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
+
         # 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)
+        # 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 (self.mpi_size > 1 or sort_boxes):
+        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()
 
@@ -1852,6 +1877,15 @@ 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:
+            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)
+            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
@@ -1932,6 +1966,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))
@@ -2532,7 +2569,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 = 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)
@@ -2602,6 +2647,7 @@ def _allocate_marker_array(self, dry_run: bool = False):
             self.first_free_idx,
             self.mu_idx,
             self._bc_type,
+            self._pole,
         )
 
         # on the CuPy backend with several ranks, find out now (collectively, in the setup) whether MPI takes device

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)

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]
         )

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)

src/struphy/pic/pushing/pusher_utilities_kernels.cuh

@@ -36,8 +36,9 @@ __device__ inline void reflect_velocity(long long ip, Array2D<double> 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.
@@ -48,6 +49,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 &&
@@ -92,6 +103,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);
     }