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 fd7724cc4..6c7f1fccb 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) @@ -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 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 45f3e37e2..259091753 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -36,8 +36,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. @@ -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); } 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/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index 1d3d1566e..02f331bec 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -44,6 +44,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_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_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 09020c588..1de1a0f4b 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -227,6 +227,46 @@ def test_draw_without_boxes_on_rank_domain(loading, sort, Np=1000): assert n_markers == Np +@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 + # (reduce before asserting, such that a failure on one process does not deadlock the others) + positions = particles.positions + 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.}]) diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index a22d9e05f..3fab53103 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -57,7 +57,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): @@ -151,7 +151,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..92b682e92 --- /dev/null +++ b/src/struphy/pic/tests/test_pole.py @@ -0,0 +1,372 @@ +"""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][:4] + markers = particles.markers + 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[: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) +@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, 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 + 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..f6fc1da6f 100644 --- a/src/struphy/propagators/push_guiding_center_bx_estar.py +++ b/src/struphy/propagators/push_guiding_center_bx_estar.py @@ -178,31 +178,33 @@ 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 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: 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() + self._grad_b_pol = self.derham.grad.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() + self._e_field_pol = self.derham.grad.codomain.zeros() # choose method particles = self.variables.ions.particles @@ -482,24 +484,39 @@ 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_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 + 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_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 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..7aa53f662 100644 --- a/src/struphy/propagators/push_guiding_center_parallel.py +++ b/src/struphy/propagators/push_guiding_center_parallel.py @@ -195,31 +195,33 @@ 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 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: 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() + self._grad_b_pol = self.derham.grad.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() + self._e_field_pol = self.derham.grad.codomain.zeros() # choose method particles = self.variables.ions.particles @@ -508,24 +510,39 @@ 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_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 + 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_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 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