From 131737f130c140a27bbe679565f00f95e6bb7df3 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 14 Oct 2026 19:21:03 +0200 Subject: [PATCH 01/12] Store the marker buffer column-major (struct-of-arrays) The marker buffer Particles.markers keeps its shape (n_rows, n_cols) and its indexing, but is now allocated in Fortran order, so every column is one contiguous block. Column operations (the pusher preamble, hole detection, boundary conditions, masks in the diagnostics) touch only the columns involved, and on the GPU neighbouring threads read neighbouring addresses. - struphy.pic.marker_layout: MARKER_ORDER, marker_array, as_marker_array, check_marker_layout; n_rows is padded to a multiple of 32 so that every column starts on a 256-byte boundary. - pyccel: MarkerArguments.markers and the 16 kernels that take the marker buffer declare float[:, :](order=F); pyccel rejects a C-ordered array. - CUDA: CudaMarkerArguments takes the F-contiguous buffer; prepare_push takes an Array2D view (a raw pointer argument is made C-contiguous by the launch) and numbers its threads column by column. - Geometry and spline evaluation at markers convert caller arrays (row gathers are C-ordered) to the marker layout. - Binning, the box count of the sorting boxes, the DKEA magnetic energy and two emptiness checks gather only the columns they use instead of whole rows. - Restart: a file with another number of rows (padding) is loaded row by row; extra rows must be holes. The HDF5 layout on disk is unchanged. - VlasovAmpereCoupling: the info output used a non-existent particles.bufferindex; it reads the saved velocities now. - Tests: test_marker_layout.py (NumPy, fake CuPy, 1 and 2 ranks), an emulation test of prepare_push, column-major markers in the test helpers; bench_marker_layout.py times the layout-dependent operations. Part 1 of #707 (part 2, tracking the range of used rows, is a separate PR). Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 2 +- doc/sections/subsections/dev-pic.rst | 7 + .../eval_spline_mpi_markers_kernels.py | 2 +- src/struphy/feec/psydac_derham.py | 6 +- src/struphy/geometry/base.py | 6 +- .../kernel_evaluate_pic_kernels.py | 2 +- .../kernel_pullpush_pic_kernels.py | 4 +- src/struphy/kernel_arguments/array_view.cuh | 7 +- .../kernel_arguments/pusher_args_cuda.py | 17 +- .../kernel_arguments/pusher_args_kernels.py | 4 +- .../drift_kinetic_electrostatic_adiabatic.py | 5 +- .../charge_density_0form_cuda.cu | 2 +- .../linear_vlasov_ampere_cuda.cu | 2 +- .../vlasov_maxwell/vlasov_maxwell_cuda.cu | 2 +- src/struphy/pic/base.py | 63 +++--- .../canonical_kinetic_particles_kernels.py | 2 +- ...al_canonical_toroidal_moment_5d_kernels.py | 2 +- ...al_canonical_toroidal_moment_6d_kernels.py | 2 +- .../eval_energy_5d/eval_energy_5d_kernels.py | 2 +- .../eval_guiding_center_from_6d_kernels.py | 2 +- ...eval_magnetic_background_energy_kernels.py | 2 +- .../eval_magnetic_energy_PBb_kernels.py | 2 +- .../eval_magnetic_moment_5d_kernels.py | 2 +- src/struphy/pic/marker_layout.py | 59 ++++++ .../bstar_parallel_3form_cuda.cu | 2 +- .../push_eta_stage/push_eta_stage_cuda.cu | 2 +- .../push_v_with_efield_cuda.cu | 2 +- .../push_vxb_analytic_cuda.cu | 2 +- .../push_vxb_implicit_cuda.cu | 2 +- .../push_weights_with_efield_lin_va_cuda.cu | 2 +- .../kernels/reflect/reflect_kernels.py | 2 +- src/struphy/pic/pushing/prepare_push_cuda.cu | 26 +-- src/struphy/pic/pushing/pusher.py | 4 +- src/struphy/pic/sorting_kernels.py | 6 +- src/struphy/pic/tests/bench_marker_layout.py | 191 ++++++++++++++++++ src/struphy/pic/tests/cuda_parity_cases.py | 4 +- src/struphy/pic/tests/kernel_test_args.py | 8 +- src/struphy/pic/tests/test_accum_gc_gradB.py | 2 +- src/struphy/pic/tests/test_cuda_emulation.py | 20 ++ src/struphy/pic/tests/test_kernel_backends.py | 13 +- src/struphy/pic/tests/test_kernel_setup.py | 4 +- src/struphy/pic/tests/test_marker_layout.py | 157 ++++++++++++++ src/struphy/pic/tests/test_pole.py | 4 +- src/struphy/pic/tests/test_reflect.py | 3 +- src/struphy/pic/utilities_kernels.py | 2 +- .../propagators/vlasov_ampere_coupling.py | 21 +- 46 files changed, 569 insertions(+), 116 deletions(-) create mode 100644 src/struphy/pic/marker_layout.py create mode 100644 src/struphy/pic/tests/bench_marker_layout.py create mode 100644 src/struphy/pic/tests/test_marker_layout.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index c2c09d408..6af88b79e 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -378,7 +378,7 @@ stencil matrix data) are in cunumpy since 0.6.1 (`Array6D`); `linear_vlasov_ampe - **Scalar types.** Members of the argument structs are checked when they are packed (PR 8). Scalars passed directly to a kernel (`dt`, `stage`) are not checked against the kernel signature: a Python `float` for an `int` parameter arrives wrong without an error. Either `CudaKernel` parses the `extern "C"` signature once and casts/checks (what cunumpy's `CudaKernel` does), or we accept it until the move to cunumpy. Until then: kernels take `double dt, int stage` in exactly this order, and the pyccel signature is the reference. - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. -- **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. The array view represents strides explicitly on the CUDA side. +- **Marker layout.** Decided (#707): the markers array (`n_markers × n_cols`) is column-major (struct-of-arrays, `order="F"`) on both backends, see `struphy.pic.marker_layout`. Neighbouring threads (markers) read neighbouring addresses, and column operations (the pusher preamble, hole detection, masks in the diagnostics) touch only the columns involved. Indexing is unchanged (`markers[ip, j]`); the pyccel kernels declare `float[:, :](order=F)`, and the CUDA kernels read the buffer through `Array2D`, which carries the strides. `n_rows` is a multiple of 32, so every column starts on a 256-byte boundary. - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet. - **Single-source alternatives.** Decided (#690): no code generation; every CUDA kernel is written by hand, see [Hand-written CUDA](#hand-written-cuda-decision-for-690). - **Polar splines and MHD equilibria on the GPU** (left after PR 19): both run on CuPy now, polar splines since #695 diff --git a/doc/sections/subsections/dev-pic.rst b/doc/sections/subsections/dev-pic.rst index 6ce11453e..413d922a9 100644 --- a/doc/sections/subsections/dev-pic.rst +++ b/doc/sections/subsections/dev-pic.rst @@ -11,6 +11,13 @@ Particle base class :exclude-members: __init__ +Marker array layout +^^^^^^^^^^^^^^^^^^^ + +.. automodule:: struphy.pic.marker_layout + :members: + + Particle subclasses ^^^^^^^^^^^^^^^^^^^ diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py index 9b166e165..8a63df173 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py @@ -8,7 +8,7 @@ def eval_spline_mpi_markers( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 _data: "float[:,:,:]", args_spline: "SplineArguments", values: "float[:]", diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 31fca09c3..5ae7d73f3 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -49,6 +49,7 @@ from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.kernel_arguments.spline_args_kernels import SplineArguments +from struphy.pic.marker_layout import MARKER_ORDER from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator @@ -2929,8 +2930,9 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): marker_evaluation = True is_sparse_meshgrid = False assert etas[0].ndim == 2 - # copy positions, such that flagging does not modify the caller's array - markers = xp.array(etas[0][:, :3], dtype=float) + # copy positions, such that flagging does not modify the caller's array; in the column-major marker + # layout that the kernel takes + markers = xp.array(etas[0][:, :3], dtype=float, order=MARKER_ORDER) self._flag_pts_not_on_proc(markers) tmp_shape = (markers.shape[0],) # 3D meshgrid evaluation diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 44c48ccd1..67149166a 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -19,6 +19,7 @@ from struphy.kernel_arguments.pusher_args_cuda import CudaDomainArguments from struphy.kernel_arguments.pusher_args_kernels import DomainArguments from struphy.linear_algebra import linalg_kron +from struphy.pic.marker_layout import as_marker_array from struphy.utils.cuda_arguments import check_mapping_on_device from struphy.utils.docstring_converter import rst_to_html, rst_to_latex, rst_to_markdown from struphy.utils.ipython_compat import HTML, display @@ -1054,7 +1055,8 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): # markers evaluation if len(etas) == 1: - markers = etas[0] + # the kernel takes markers in the column-major marker layout (no copy for the marker buffer) + markers = as_marker_array(etas[0]) # to keep C-ordering the (3, 3)-part is in the last indices out = xp.empty((markers.shape[0], 3, 3), dtype=float) @@ -1220,6 +1222,8 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa markers = xp.stack(etas, axis=1) else: markers = etas[0] + # the kernel takes markers in the column-major marker layout (no copy for the marker buffer) + markers = as_marker_array(markers) # coordinates (:, 3) and argument evaluation (without holes) if callable(a): diff --git a/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py b/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py index 94cb2494c..74bff0066 100644 --- a/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py +++ b/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py @@ -10,7 +10,7 @@ @stack_array("tmp0", "tmp1", "tmp2", "tmp3", "out") def kernel_evaluate_pic( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 kind_coeff: int, args: "DomainArguments", mat_f: "float[:,:,:]", diff --git a/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_kernels.py b/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_kernels.py index 6c15edf1d..d453ea4aa 100644 --- a/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_kernels.py +++ b/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_kernels.py @@ -11,7 +11,7 @@ @stack_array("tmp1", "tmp2") def kernel_pullpush_pic( a: "float[:,:]", - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 kind_transform: int, kind_fun: int, args_domain: "DomainArguments", @@ -26,7 +26,7 @@ def kernel_pullpush_pic( a : float[:,:] Values of scalar function a[0, ip] or values of components of a vector valued function (a[0, ip], a[1, ip], a[2, ip]). - markers : float[:,:] + markers : float[:,:](order=F) Evaluation points in marker format (eta1 = markers[:, 0], eta2 = markers[:, 1], eta3 = markers[:, 2]). kind_transform : int diff --git a/src/struphy/kernel_arguments/array_view.cuh b/src/struphy/kernel_arguments/array_view.cuh index 12d8afa7c..85ca437fb 100644 --- a/src/struphy/kernel_arguments/array_view.cuh +++ b/src/struphy/kernel_arguments/array_view.cuh @@ -8,9 +8,10 @@ // for kernel arguments and in Argument._pack (struphy/utils/cuda_arguments.py) for struct fields. // // The layout is NumPy's: for a C-contiguous (row-major) array of shape (n0, n1, n2), -// strides = (n1 * n2, n2, 1), so the last index runs fastest; for shape (n_markers, n_cols), -// strides = (n_cols, 1) and markers(ip, j) is data[ip * n_cols + j]. A sliced or transposed -// CuPy array keeps its own strides, so it is indexed correctly without a copy. +// strides = (n1 * n2, n2, 1), so the last index runs fastest. The marker buffer of shape +// (n_markers, n_cols) is F-contiguous (column-major, struphy.pic.marker_layout): strides = (1, n_markers) +// and markers(ip, j) is data[ip + j * n_markers]. A sliced or transposed CuPy array keeps its own +// strides, so it is indexed correctly without a copy. template struct Array1D { T* data; long long shape[1], strides[1]; diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index feac42098..dbdeed54e 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -11,17 +11,21 @@ import numpy as np from cunumpy.arguments import CudaStructArguments +from struphy.pic.marker_layout import MARKER_ORDER -def _device_array(name, arr, dtype, ndim=None): - """Check that ``arr`` is a C-contiguous CuPy array of ``dtype`` (and ``ndim``) on the current device (never copies).""" + +def _device_array(name, arr, dtype, ndim=None, order="C"): + """Check that ``arr`` is a CuPy array of ``dtype`` (and ``ndim``), contiguous in ``order`` ("C" or "F"), + on the current device (never copies).""" import cupy as cp if not isinstance(arr, cp.ndarray): raise TypeError(f"{name} must be a CuPy array, got {type(arr)}") if arr.device.id != cp.cuda.runtime.getDevice(): raise ValueError(f"{name} must be on the current CUDA device") - if arr.dtype != dtype or not arr.flags.c_contiguous: - raise TypeError(f"{name} must be a C-contiguous array of dtype {np.dtype(dtype)}") + contiguous = arr.flags.f_contiguous if order == "F" else arr.flags.c_contiguous + if arr.dtype != dtype or not contiguous: + raise TypeError(f"{name} must be a {order}-contiguous array of dtype {np.dtype(dtype)}") if ndim is not None and arr.ndim != ndim: raise TypeError(f"{name} must be a {ndim}-dimensional array, got {arr.ndim} dimensions") return arr @@ -64,9 +68,8 @@ def __init__( bc_type, pole: bool, ): - self.markers = _device_array("markers", markers, np.float64) - if markers.ndim != 2: - raise TypeError("markers must be a two-dimensional array.") + # column-major like the marker buffer (struphy.pic.marker_layout); Array2D carries the strides + self.markers = _device_array("markers", markers, np.float64, ndim=2, order=MARKER_ORDER) self.valid_mks = _device_array("valid_mks", valid_mks, np.bool_) self.Np = Np self.vdim = vdim diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index 963bbbacb..d697fe708 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -11,7 +11,7 @@ class MarkerArguments: Paramaters ---------- markers : array[float] - Markers array. + Markers array, column-major (see :mod:`struphy.pic.marker_layout`). Np : int Total number of particles. @@ -53,7 +53,7 @@ class MarkerArguments: def __init__( self, - markers: "float[:, :]", + markers: "float[:, :](order=F)", # noqa: F821 valid_mks: "bool[:]", Np: int, vdim: int, diff --git a/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py b/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py index 83d3f647e..8bd3c7779 100644 --- a/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py +++ b/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py @@ -203,8 +203,9 @@ def _compute_en_particle_magnetic(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. - markers = particles.markers_wo_holes_and_ghost - return markers[:, 5].dot(markers[:, 8]) + # only the two columns needed, not the whole rows of the valid markers + valid_mks = particles.valid_mks + return particles.markers[valid_mks, 5].dot(particles.markers[valid_mks, 8]) ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): diff --git a/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu b/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu index b03038867..4d3695561 100644 --- a/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu +++ b/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu @@ -7,7 +7,7 @@ using namespace struphy_cuda; * * B_p = w_p, deposited with the N-splines of V0; one thread per marker row. * - * @param args_markers Marker buffer (n_markers x n_cols, row-major) and the weight column weight_idx. + * @param args_markers Marker buffer (n_markers x n_cols, column-major) and the weight column weight_idx. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param args_domain Mapping arguments (unused; part of the common accumulation signature). * @param vec 0-form stencil vector data; written with atomic additions. diff --git a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu index e4ce83e46..407c211a8 100644 --- a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu +++ b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu @@ -13,7 +13,7 @@ using namespace struphy_cuda; * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; * one thread per marker row. * - * @param args_markers Marker buffer (n_markers x n_cols, row-major) and the total marker number Np; the weight is in + * @param args_markers Marker buffer (n_markers x n_cols, column-major) and the total marker number Np; the weight is in * column 6 and s_0 in column 7, as in pyccel. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param args_domain Mapping arguments; every spline and analytic mapping. diff --git a/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu index 880538d1e..ee2259a3f 100644 --- a/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu +++ b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu @@ -13,7 +13,7 @@ using namespace struphy_cuda; * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; * one thread per marker row. * - * @param args_markers Marker buffer (n_markers x n_cols, row-major); the weight is in column 6, as in pyccel. + * @param args_markers Marker buffer (n_markers x n_cols, column-major); the weight is in column 6, as in pyccel. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 6c7f1fccb..64f72911a 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -38,6 +38,7 @@ WeightsParameters, ) from struphy.pic import sampling_kernels, sobol_seq +from struphy.pic.marker_layout import marker_array, padded_n_rows from struphy.pic.pushing.kernels.reflect import reflect from struphy.pic.pushing.kernels.sph_mean_velocity_coeffs import sph_mean_velocity_coeffs from struphy.pic.pushing.kernels.sph_viscosity_tensor import sph_viscosity_tensor @@ -1494,7 +1495,7 @@ def initialize_weights( else: assert self.domain is not None, "A domain is needed to initialize weights." - if xp.size(self.markers_wo_holes_and_ghost) == 0: + if self.n_mks_loc == 0: return # set initial condition @@ -1555,7 +1556,7 @@ def update_weights(self): """ from struphy.pic.particles import ParticlesSPH - if xp.size(self.markers_wo_holes_and_ghost) == 0: + if self.n_mks_loc == 0: return if isinstance(self, ParticlesSPH): @@ -1626,9 +1627,12 @@ def binning( quantity, *v_axis = output_quantity.rsplit(sep="_", maxsplit=1) v_axis = [int(char) - 1 for char in "".join(v_axis)] # convert dimension axis to index - # the valid markers, gathered once; everything below stays on the device of the markers - markers = self.markers_wo_holes_and_ghost - velocities = markers[:, self.index["vel"]] + # the rows of the valid markers, found once; below, only the columns that are used are gathered from them + # (each one contiguous in the column-major marker buffer), not whole rows. Everything stays on the device + # of the markers. + rows = xp.nonzero(self.valid_mks)[0] + if quantity != "density": + velocities = self.markers[rows, self.index["vel"]] # determine histogram weights multiplier if quantity == "density": @@ -1642,11 +1646,11 @@ def binning( multiplier = velocity_norm2 * velocities[:, v_axis[0]] # compute weights of histogram: - _weights0 = markers[:, self.index["w0"]] * self.Np * multiplier - _weights = markers[:, self.index["weights"]] * self.Np * multiplier + _weights0 = self.markers[rows, self.index["w0"]] * self.Np * multiplier + _weights = self.markers[rows, self.index["weights"]] * self.Np * multiplier if divide_by_jac: - jacobian_det = self.domain.jacobian_det(markers[:, self.index["pos"]], remove_outside=False) + jacobian_det = self.domain.jacobian_det(self.markers[rows, self.index["pos"]], remove_outside=False) _weights /= jacobian_det # _weights /= self.velocity_jacobian_det(*self.phasespace_coords.T) @@ -1654,7 +1658,7 @@ def binning( # _weights0 /= self.velocity_jacobian_det(*self.phasespace_coords.T) # the binned columns, one by one (an index list would have to be uploaded to the device first) - sample = xp.stack([markers[:, i] for i, binned in enumerate(slicing) if binned], axis=1) + sample = xp.stack([self.markers[rows, i] for i, binned in enumerate(slicing) if binned], axis=1) f_slice = xp.histogramdd( sample, @@ -1888,9 +1892,8 @@ def apply_kinetic_bc(self, newton=False): # apply boundary conditions if self._remove_axes: - # extract the (n_rows, 3) logical coordinates once per contiguous - # cache line instead of re-striding into the full row-major - # markers array once per axis (see _find_outside_particles) + # snapshot of the (n_rows, 3) logical coordinates for the loop over the axes + # (see _find_outside_particles) self._eta_bc_buf[:] = self.markers[:, :3] for axis in self._remove_axes: @@ -2603,8 +2606,9 @@ def _allocate_marker_array(self, dry_run: bool = False): n_mks_load_loc = self.n_mks_load[self._mpi_rank] bufsize = self.bufsize + 1.0 / xp.sqrt(n_mks_load_loc) - # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer - self._n_rows = round(float(n_mks_load_loc * (1 + bufsize))) + # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer; + # column-major, every column starting on a 256-byte boundary (see struphy.pic.marker_layout) + self._n_rows = padded_n_rows(round(float(n_mks_load_loc * (1 + bufsize)))) # Have at least 3 spare places in markers array assert self.first_free_idx + 2 < self.n_cols - 2, ( @@ -2614,7 +2618,7 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return - self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float) + self._markers = marker_array(self.n_rows, self.n_cols) # allocate auxiliary arrays self._holes = xp.zeros(self.n_rows, dtype=bool) @@ -2623,10 +2627,9 @@ def _allocate_marker_array(self, dry_run: bool = False): self._is_outside_right = xp.zeros(self.n_rows, dtype=bool) self._is_outside_left = xp.zeros(self.n_rows, dtype=bool) self._is_outside = xp.zeros(self.n_rows, dtype=bool) - # contiguous scratch copy of markers[:, :3], refreshed once per apply_kinetic_bc - # boundary-condition-type loop (see there) instead of re-striding into the full - # (n_rows, n_cols) row-major marker array once per axis. - self._eta_bc_buf = xp.zeros((self.n_rows, 3), dtype=float) + # scratch copy of markers[:, :3], refreshed once per apply_kinetic_bc boundary-condition-type + # loop (see there); column-major like the markers, so the copy is three contiguous columns + self._eta_bc_buf = marker_array(self.n_rows, 3) # create array container (3 x positions, vdim x velocities, weight, s0, w0, ID) for removed markers self._n_lost_markers = 0 @@ -2935,7 +2938,17 @@ def _load_restart(self): data = DataContainer(data_path, comm=self.mpi_comm) with h5py.File(data.file_path, "a") as file: - self._markers[:, :] = file["restart/" + self.loading_params.restart_key][-1, :, :] + saved = file["restart/" + self.loading_params.restart_key][-1, :, :] + + # the number of rows is padded (see struphy.pic.marker_layout), so a file written with another padding has a + # different number of rows: copy the common rows; the rows after them are holes, in the buffer (set before + # loading) and in the file + n_rows = min(saved.shape[0], self.n_rows) + assert np.all(saved[n_rows:, 0] == -1.0), ( + f"the restart file has {saved.shape[0]} marker rows on rank {self.mpi_rank}, but there are only " + f"{self.n_rows} rows here, and the extra ones are not holes" + ) + self._markers[:n_rows, :] = saved[:n_rows] def _load_tesselation(self, n_quad: int = 1): """ @@ -2979,10 +2992,8 @@ def _find_outside_particles(self, axis, eta=None): Column of the markers array (0, 1 or 2) holding the logical coordinate to check. eta : xp.ndarray[float], optional - Pre-extracted, contiguous ``(n_rows, 3)`` copy of ``markers[:, :3]`` (see - :meth:`apply_kinetic_bc`). If ``None``, reads straight from ``markers`` -- - correct but slower, since a single-column slice of the row-major ``markers`` - array is strided (see the comment in :meth:`apply_kinetic_bc`). + Snapshot ``(n_rows, 3)`` of ``markers[:, :3]`` (see :meth:`apply_kinetic_bc`). + If ``None``, reads straight from ``markers``. Returns ------- @@ -3165,11 +3176,11 @@ def _check_and_assign_particles_to_boxes(self): from cunumpy.xp import array_backend if array_backend.backend == "numpy": - bcount = xp.bincount(xp.int64(self.markers_wo_holes[:, -2])) + bcount = xp.bincount(xp.int64(self.markers[~self.holes, -2])) else: import cupy as cp - indices = self.markers_wo_holes[:, -2] + indices = self.markers[~self.holes, -2] indices = indices.astype(cp.int64) bcount = cp.bincount(indices) diff --git a/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py b/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py index 32d55dff2..07e8df407 100644 --- a/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py @@ -13,7 +13,7 @@ @stack_array("dfm", "dfinv", "dfinv_t", "v", "a_form", "dfta_form") def canonical_kinetic_particles( res: "float[:]", - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", args_domain: "DomainArguments", a1_1: "float[:,:,:]", diff --git a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py index 350bcc803..23f3f9a6f 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py @@ -8,7 +8,7 @@ def eval_canonical_toroidal_moment_5d( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", first_diagnostics_idx: int, mu_idx: int, diff --git a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py index f13d23a25..78f5e1b98 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py @@ -8,7 +8,7 @@ def eval_canonical_toroidal_moment_6d( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", first_diagnostics_idx: int, epsilon: float, diff --git a/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py index 67721d35e..810e7c438 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py @@ -8,7 +8,7 @@ def eval_energy_5d( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", first_diagnostics_idx: int, mu_idx: int, diff --git a/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py index c5c08d3d4..a95427436 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py @@ -12,7 +12,7 @@ @stack_array("v", "dfm", "b2", "norm_b_cart", "temp", "v_perp", "Larmor_r") def eval_guiding_center_from_6d( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, diff --git a/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py index 0f1a8201a..61a973d52 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py @@ -10,7 +10,7 @@ @stack_array("dfm", "norm_b1", "b") def eval_magnetic_background_energy( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, diff --git a/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py index d079bb11b..43e2f4b39 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py @@ -11,7 +11,7 @@ @stack_array("dfm", "eta") def eval_magnetic_energy_PBb( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, diff --git a/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py index 68dc460f3..14eb8dcef 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py @@ -8,7 +8,7 @@ def eval_magnetic_moment_5d( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", first_diagnostics_idx: int, absB: "float[:,:,:]", diff --git a/src/struphy/pic/marker_layout.py b/src/struphy/pic/marker_layout.py new file mode 100644 index 000000000..f2f89e61d --- /dev/null +++ b/src/struphy/pic/marker_layout.py @@ -0,0 +1,59 @@ +"""Memory layout of the marker buffer. + +:attr:`Particles.markers ` has shape ``(n_rows, n_cols)``, one row per marker, +and is stored column-major (Fortran order, struct-of-arrays): every column is one contiguous block of ``n_rows`` +values. Column operations (``markers[:, j]``, ``markers[:, a:b] = ...``, a mask on one column) then touch only +the columns involved, and on the GPU neighbouring threads (markers) read neighbouring addresses. Indexing is +unchanged: ``markers[ip, j]`` is still the ``j``-th entry of marker ``ip``. + +Kernels see the buffer as follows: + +* pyccel kernels declare it as ``float[:, :](order=F)`` (in :class:`~struphy.kernel_arguments.pusher_args_kernels.MarkerArguments` + and in the kernels that take ``markers`` directly); pyccel checks the order at every call, so a C-ordered array + raises a ``TypeError`` instead of being read wrongly. +* CUDA kernels read it through ``Array2D`` views, which carry the strides of the array. + +Copies made by fancy or boolean row indexing (``markers[valid_mks]``), ``.copy()`` and ``xp.concatenate`` are +C-ordered. They are fine to read, or to write back into the buffer with ``markers[rows] = ...``; to pass such an +array to a kernel, convert it with :func:`as_marker_array` first. + +See struphy-hub/struphy#707. +""" + +import cunumpy as xp + +MARKER_ORDER = "F" +"""Memory order of the marker buffer and of every marker array passed to a kernel.""" + +ROW_ALIGNMENT = 32 +"""The number of rows is a multiple of this, so that every column starts on a 256-byte boundary (32 doubles).""" + + +def padded_n_rows(n_rows: int) -> int: + """Smallest multiple of :data:`ROW_ALIGNMENT` that is at least ``n_rows``.""" + return -(-int(n_rows) // ROW_ALIGNMENT) * ROW_ALIGNMENT + + +def marker_array(n_rows: int, n_cols: int): + """Zero marker array of shape ``(n_rows, n_cols)`` in :data:`MARKER_ORDER` on the active backend.""" + return xp.zeros((n_rows, n_cols), dtype=float, order=MARKER_ORDER) + + +def as_marker_array(markers): + """``markers`` in :data:`MARKER_ORDER`, for passing to a kernel; no copy if it already has that order.""" + return xp.asfortranarray(markers) + + +def has_marker_layout(markers) -> bool: + """Whether ``markers`` is a 2D array in :data:`MARKER_ORDER` (as the kernels expect).""" + return markers.ndim == 2 and bool(markers.flags.f_contiguous) + + +def check_marker_layout(markers, name: str = "markers"): + """Raise a ``TypeError`` if ``markers`` is not a 2D array in :data:`MARKER_ORDER`.""" + if not has_marker_layout(markers): + raise TypeError( + f"{name} must be a 2D Fortran-ordered (column-major) array, got shape {markers.shape} with " + f"C_CONTIGUOUS={markers.flags.c_contiguous}, F_CONTIGUOUS={markers.flags.f_contiguous}; " + "see struphy.pic.marker_layout" + ) diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu index 7a62cb4c8..c0c889853 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu @@ -12,7 +12,7 @@ using namespace struphy_cuda; * * @param alpha Weights of the average (at least four entries): alpha[0:3] for eta, alpha[3] for v_parallel. * @param output_indices Marker column of the result in output_indices[0]; -1 skips the output. - * @param args_markers Marker buffer (n_markers x n_cols, row-major) and column indices. + * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param epsilon Scaling parameter epsilon in front of v_parallel. diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu index 9a6c24d52..252fd6d40 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu @@ -9,7 +9,7 @@ using namespace struphy_cuda; * * @param dt Time step. * @param stage Index of the current stage, 0 <= stage < n_stages. - * @param args_markers Marker buffer (n_markers x n_cols, row-major) and column indices. + * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param a Butcher coefficients a_stage (n_stages entries). * @param b Butcher weights b (n_stages entries). diff --git a/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu b/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu index e5ede6bf2..a5828b27b 100644 --- a/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu @@ -11,7 +11,7 @@ using namespace struphy_cuda; * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, row-major), validity mask and column indices. + * @param args_markers Marker buffer (n_markers x n_cols, column-major), validity mask and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param e1_1 Coefficients of the first component of the electric 1-form. diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu index 7d2a270ec..c8131f4b1 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu @@ -13,7 +13,7 @@ using namespace struphy_cuda; * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, row-major) and column indices. + * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param b2_1 Coefficients of the first component of the magnetic 2-form. diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu index 71ef3a461..687328fe5 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu @@ -12,7 +12,7 @@ using namespace struphy_cuda; * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, row-major) and column indices. + * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param b2_1 Coefficients of the first component of the magnetic 2-form. diff --git a/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu b/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu index f3dfb6e00..1b6fd9f87 100644 --- a/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu @@ -13,7 +13,7 @@ using namespace struphy_cuda; * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, row-major), total marker number Np and column indices. + * @param args_markers Marker buffer (n_markers x n_cols, column-major), total marker number Np and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param e1_1 Coefficients of the first component of the electric 1-form. diff --git a/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py b/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py index b32974a0b..88c790688 100644 --- a/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py +++ b/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py @@ -11,7 +11,7 @@ @stack_array("dfm", "dfinv", "eta", "v", "v_logical") def reflect( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_domain: "DomainArguments", outside_inds: "int[:]", axis: "int", diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index 71c02556d..37a853e15 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -1,3 +1,5 @@ +#include "cunumpy/array_view.cuh" + /** * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: * @@ -5,24 +7,22 @@ * markers[:, first_shift_idx:residual_idx] = 0.0 * markers[:, residual_idx:-2] = 0.0 * - * One thread per entry of the contiguous column range [first_pusher_idx, n_cols - 2) of every row, so that - * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written - * in a single pass (instead of one strided pass over the buffer per slice). + * One thread per entry of the column range [first_pusher_idx, n_cols - 2) of every row, numbered column by + * column. In the column-major marker buffer (struphy.pic.marker_layout) neighbouring threads then read and write + * neighbouring rows of the same column, and the whole range is written in a single launch (instead of one per + * slice). The buffer is an Array2D view: a raw pointer argument would be made C-contiguous (copied) by the launch. * - * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. - * @param n_rows Number of rows of the buffer. - * @param n_cols Number of columns of the buffer. + * @param markers Marker buffer (n_rows x n_cols, column-major), every row including holes. * @param first_pusher_idx First column of the saved initial phase space coordinates. * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ -extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_pusher_idx, - int n_init) { - long long width = n_cols - 2 - first_pusher_idx; +extern "C" __global__ void prepare_push(Array2D markers, int first_pusher_idx, int n_init) { + long long n_rows = markers.shape[0]; + long long width = markers.shape[1] - 2 - first_pusher_idx; long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; if (t >= n_rows * width) return; - long long ip = t / width; - int col = first_pusher_idx + (int)(t - ip * width); - double* row = markers + ip * n_cols; - row[col] = col < first_pusher_idx + n_init ? row[col - first_pusher_idx] : 0.0; + long long k = t / n_rows; + long long ip = t - k * n_rows; + markers(ip, first_pusher_idx + k) = k < n_init ? markers(ip, k) : 0.0; } diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index c9271883f..7fa7069de 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -218,12 +218,10 @@ def _push(self, dt: float): logger.debug(f"{self.particles.n_cols =}") if self._cuda: - # the three assignments below in one pass over the buffer (each slice is a strided pass on its own) + # the three assignments below in one launch n_rows, n_cols = markers.shape _prepare_push_cuda()( markers, - n_rows, - n_cols, first_pusher_idx, 3 + vdim, n_threads=n_rows * (n_cols - 2 - first_pusher_idx), diff --git a/src/struphy/pic/sorting_kernels.py b/src/struphy/pic/sorting_kernels.py index 46d84007f..8e55838b0 100644 --- a/src/struphy/pic/sorting_kernels.py +++ b/src/struphy/pic/sorting_kernels.py @@ -154,7 +154,7 @@ def get_next_index(box_nb: "float", next_index: "int[:]", cumul_next_index: "int def sort_boxed_particles( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 swap_line_1: "float[:]", swap_line_2: "float[:]", n_boxes: "int", @@ -198,7 +198,7 @@ def sort_boxed_particles( def assign_box_to_each_particle( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 holes: "bool[:]", nx: "int", ny: "int", @@ -230,7 +230,7 @@ def assign_box_to_each_particle( def assign_particles_to_boxes( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 holes: "bool[:]", boxes: "int[:,:]", next_index: "int[:]", diff --git a/src/struphy/pic/tests/bench_marker_layout.py b/src/struphy/pic/tests/bench_marker_layout.py new file mode 100644 index 000000000..db60bdff9 --- /dev/null +++ b/src/struphy/pic/tests/bench_marker_layout.py @@ -0,0 +1,191 @@ +""" +Standalone benchmark (not a pytest test) of the operations on the marker buffer that depend on its memory layout +(see :mod:`struphy.pic.marker_layout` and struphy-hub/struphy#707): the column assignments at the start of every +push, two pushers, one accumulation kernel, hole detection, boundary conditions, the MPI sort, binning and the +gathers of valid markers used by the diagnostics. + +Runs on the active cunumpy backend (``CUNUMPY_BACKEND=numpy`` or ``cupy``); on the GPU every timing ends with a +device synchronization. Prints the median wall time of each operation in milliseconds, and writes them to a JSON +file with ``--json``, so that two runs (e.g. before and after a layout change) can be compared with ``--compare``. + +Usage:: + + mpirun -n 1 python src/struphy/pic/tests/bench_marker_layout.py [--Np 1000000] [--repeats 5] [--json out.json] + python src/struphy/pic/tests/bench_marker_layout.py --compare before.json after.json +""" + +import argparse +import json +import statistics +import time + + +def _time(fun, repeats: int) -> float: + """Median wall time of ``fun()`` in milliseconds, after one warm-up call.""" + import cunumpy as xp + + fun() + xp.synchronize() + times = [] + for _ in range(repeats): + t0 = time.perf_counter() + fun() + xp.synchronize() + times.append(time.perf_counter() - t0) + return 1e3 * statistics.median(times) + + +def run(Np: int, repeats: int, num_elements: int, degree: int, save_markers: str = None) -> dict: + import cunumpy as xp + import numpy as np + from maybempi import MPI + + from struphy import BoundaryParameters, LoadingParameters, WeightsParameters, domains + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.ode.utils import ButcherTableau + from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form + from struphy.pic.particles import Particles6D + from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage + from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic + from struphy.pic.pushing.pusher import Pusher, _prepare_push_cuda + from struphy.topology.grids import TensorProductGrid + + comm = MPI.COMM_WORLD + domain = domains.Cuboid(r1=2.0, r2=3.0, r3=4.0) + grid = TensorProductGrid(num_elements=[num_elements] * 3) + derham = Derham(grid, DerhamOptions(degree=[degree] * 3), comm=comm) + + particles = Particles6D( + comm_world=comm, + domain_decomp=(derham.domain_array, derham.domain_decomposition.nprocs), + loading_params=LoadingParameters(Np=Np, seed=1234, spatial="uniform"), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(bc=("periodic", "periodic", "periodic")), + domain=domain, + ) + particles.draw_markers() + particles.mpi_sort_markers() + particles.initialize_weights() + + markers = particles.markers + first_pusher_idx = particles.first_pusher_idx + first_shift_idx = particles.first_shift_idx + residual_idx = particles.residual_idx + n_init = 3 + particles.vdim + + def preamble(): + """The column assignments at the start of Pusher._push.""" + if xp.get_backend() == "cupy": + n_rows, n_cols = markers.shape + _prepare_push_cuda()( + markers, + first_pusher_idx, + n_init, + n_threads=n_rows * (n_cols - 2 - first_pusher_idx), + ) + else: + markers[:, first_pusher_idx:first_shift_idx] = markers[:, :n_init] + markers[:, first_shift_idx:residual_idx] = 0.0 + markers[:, residual_idx:-2] = 0.0 + + butcher = ButcherTableau("rk4") + push_eta = Pusher( + particles, + push_eta_stage, + (xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(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, + ) + + _, b2 = create_equal_random_arrays(derham.V2fem, seed=3456, flattened=True) + push_vxb = Pusher( + particles, + push_vxb_analytic, + (derham.args_derham, b2[0]._data, b2[1]._data, b2[2]._data), + domain.args_domain, + alpha_in_kernel=1.0, + pushes_eta=False, + ) + + rho = derham.coeff_spaces["H1"].zeros()._data + + def accumulate(): + rho[:] = 0.0 + charge_density_0form(particles.args_markers, derham.args_derham, domain.args_domain, rho) + + bin_edges = (xp.linspace(0.0, 1.0, 65), xp.linspace(-5.0, 5.0, 65)) + components = (True, False, False, True, False, False) + + ops = { + "preamble (column assignments)": preamble, + "push_eta_stage (RK4 pusher)": lambda: push_eta(1e-3), + "push_vxb_analytic (pusher)": lambda: push_vxb(1e-3), + "charge_density_0form (accumulation)": accumulate, + "update_holes": particles.update_holes, + "apply_kinetic_bc (periodic)": particles.apply_kinetic_bc, + "mpi_sort_markers": particles.mpi_sort_markers, + "binning e1_v1": lambda: particles.binning(components, bin_edges), + "gather valid rows (markers_wo_holes_and_ghost)": lambda: particles.markers_wo_holes_and_ghost, + "gather weights of valid markers": lambda: particles.weights, + } + + info = { + "backend": xp.get_backend(), + "n_rows": int(markers.shape[0]), + "n_cols": int(markers.shape[1]), + "Np": Np, + "f_contiguous": bool(markers.flags.f_contiguous), + } + result = {"info": info, "ms": {name: _time(fun, repeats) for name, fun in ops.items()}} + + if save_markers is not None: + # the valid markers after all timed calls, sorted by ID, for a bit-for-bit comparison of two runs + valid = xp.to_numpy(particles.markers_wo_holes_and_ghost) + np.save(save_markers, valid[np.argsort(valid[:, -1])]) + return result + + +def _print(result: dict): + info = result["info"] + order = "F (column-major)" if info["f_contiguous"] else "C (row-major)" + print(f"backend {info['backend']}, markers {info['n_rows']} x {info['n_cols']} in {order}, Np = {info['Np']}") + width = max(len(name) for name in result["ms"]) + for name, ms in result["ms"].items(): + print(f" {name:<{width}} {ms:10.2f} ms") + + +def _compare(before: dict, after: dict): + width = max(len(name) for name in before["ms"]) + print(f" {'operation':<{width}} {'before':>10} {'after':>10} {'speedup':>8}") + for name, ms_before in before["ms"].items(): + ms_after = after["ms"].get(name) + if ms_after is None: + continue + print(f" {name:<{width}} {ms_before:8.2f}ms {ms_after:8.2f}ms {ms_before / ms_after:7.2f}x") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--Np", type=int, default=1_000_000, help="number of markers") + parser.add_argument("--repeats", type=int, default=5, help="timed calls per operation (median is reported)") + parser.add_argument("--num-elements", type=int, default=8, help="cells per direction") + parser.add_argument("--degree", type=int, default=3, help="spline degree per direction") + parser.add_argument("--json", help="write the timings to this file") + parser.add_argument("--save-markers", help="save the valid markers after the run (sorted by ID) to this .npy file") + parser.add_argument("--compare", nargs=2, metavar=("BEFORE", "AFTER"), help="compare two JSON files and exit") + args = parser.parse_args() + + if args.compare: + with open(args.compare[0]) as f, open(args.compare[1]) as g: + _compare(json.load(f), json.load(g)) + else: + result = run(args.Np, args.repeats, args.num_elements, args.degree, args.save_markers) + _print(result) + if args.json: + with open(args.json, "w") as f: + json.dump(result, f, indent=2) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index 7e4bb00e2..1f28821e5 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -103,7 +103,7 @@ def push_weights_with_efield_lin_va_args(bc): def reflect_args(axis): """Reverse one logical velocity component of 10 markers in a Cuboid.""" - markers = np.random.default_rng(32).random((140, 8)) + markers = np.asfortranarray(np.random.default_rng(32).random((140, 8))) outside_inds = np.arange(139, 129, -1, dtype=np.int64) return (xp.asarray(markers), Cuboid(r1=2.0, r2=3.0, r3=4.0).args_domain, xp.asarray(outside_inds), axis) @@ -167,7 +167,7 @@ def vlasov_maxwell_args(domain): def eval_spline_mpi_markers_args(kind): """Spline values at 129 markers, one of them flagged.""" - markers = np.random.default_rng(5).random((129, 3)) + markers = np.asfortranarray(np.random.default_rng(5).random((129, 3))) markers[0, 0] = -1.0 # not on the process domain: skipped return (xp.asarray(markers), *spline_evaluation_arguments(kind), xp.zeros(129)) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index 02f331bec..ae30c272f 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -17,6 +17,7 @@ from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.ode.utils import ButcherTableau +from struphy.pic.marker_layout import as_marker_array N_MARKERS = 129 # not a multiple of the block size N_COLS = 25 @@ -39,7 +40,7 @@ def marker_arguments(bc): valid[:2] = False args_class = CudaMarkerArguments if xp.get_backend() == "cupy" else MarkerArguments args_markers = args_class( - xp.asarray(markers), + as_marker_array(xp.asarray(markers)), xp.asarray(valid), N_MARKERS, *MARKER_INDICES, @@ -213,7 +214,8 @@ def geometry_domain(index): def logical_markers(n=N_MARKERS, seed=5): - """Logical points in marker format, shape (n, 7), away from the poles (eta1 >= 0.1). + """Logical points in marker format, shape (n, 7), column-major like the marker buffer, away from the poles + (eta1 >= 0.1). Row 0 is a hole (-1), row 1 is outside in eta1 and row 2 in eta3, so ``remove_outside`` changes the result. """ @@ -223,4 +225,4 @@ def logical_markers(n=N_MARKERS, seed=5): markers[0, :] = -1.0 markers[1, 0] = 1.2 markers[2, 2] = -0.3 - return markers + return np.asfortranarray(markers) diff --git a/src/struphy/pic/tests/test_accum_gc_gradB.py b/src/struphy/pic/tests/test_accum_gc_gradB.py index 9950cae73..0fce51a4a 100644 --- a/src/struphy/pic/tests/test_accum_gc_gradB.py +++ b/src/struphy/pic/tests/test_accum_gc_gradB.py @@ -36,7 +36,7 @@ def rand(space): zero = derham.coeff_spaces["Hcurl"].zeros() n_mk = 5 - markers = np.zeros((n_mk, 40)) + markers = np.zeros((n_mk, 40), order="F") markers[:, :3] = rng.uniform(0.05, 0.95, (n_mk, 3)) markers[:, 3] = rng.standard_normal(n_mk) markers[:, 4] = rng.uniform(0.1, 1.0, n_mk) diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index 0e69c148d..a2548df03 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -111,3 +111,23 @@ def test_emulated_v1_symm_filler(): for result, reference in zip(data, expected): assert np.any(reference != 0.0) np.testing.assert_allclose(result, reference, rtol=1e-12, atol=1e-13) + + +@requires_compiler +@pytest.mark.parametrize("n_rows", [96, 129]) +def test_emulated_prepare_push(n_rows): + """The fused preamble of Pusher._push (prepare_push_cuda.cu) writes the column-major buffer like the three column + assignments, and leaves the columns before first_pusher_idx and the last two columns alone.""" + from struphy.pic.marker_layout import marker_array + from struphy.pic.pushing.pusher import _prepare_push_cuda + + n_cols, first_pusher_idx, n_init = 26, 8, 6 + markers = marker_array(n_rows, n_cols) + markers[:] = np.random.default_rng(11).random((n_rows, n_cols)) + expected = markers.copy() + expected[:, first_pusher_idx : first_pusher_idx + n_init] = expected[:, :n_init] + expected[:, first_pusher_idx + n_init : -2] = 0.0 + + width = n_cols - 2 - first_pusher_idx + emulate_struct_kernel(_prepare_push_cuda(), markers, first_pusher_idx, n_init, n_threads=n_rows * width) + np.testing.assert_array_equal(markers, expected) diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 3fab53103..cccb16898 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -20,6 +20,7 @@ from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments, MarkerArguments from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.kernel_arguments.spline_args_kernels import SplineArguments +from struphy.pic.marker_layout import as_marker_array from struphy.pic.tests.kernel_test_args import N_GEOMETRY_DOMAINS from struphy.utils.cuda_arguments import ( CUDA_OPTIONS, @@ -53,7 +54,7 @@ def make_arguments(n_markers, seed=0): rng = np.random.default_rng(seed) - markers = cunumpy.asarray(rng.random((n_markers, N_COLS))) + markers = as_marker_array(cunumpy.asarray(rng.random((n_markers, N_COLS)))) 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 @@ -107,13 +108,17 @@ def test_generated_spline_header(tmp_path): @pytest.mark.parametrize("cuda_class, source, class_name, cuda_only, header", ARGUMENT_PAIRS) -def test_argument_classes_correspond(cuda_class, source, class_name, cuda_only, header): +def test_argument_classes_correspond(cuda_class, source, class_name, cuda_only, header, tmp_path): """Each CUDA argument class mirrors its pyccel class: same constructor, same attributes in the same order.""" assert cuda_class.__name__ == "Cuda" + class_name cuda_parameters = list(inspect.signature(cuda_class.__init__).parameters)[1:] assert cuda_parameters == pyccel_init_parameters(source, class_name) - # attributes set from constructor arguments; derived ones (e.g. n_markers, bn1) are skipped by the parser - pyccel_fields = [field.name for field in CudaStruct.from_pyccel_class(ARGS_DIR / source, class_name).fields] + # attributes set from constructor arguments; derived ones (e.g. n_markers, bn1) are skipped by the parser. + # cunumpy's parser does not read pyccel's memory-order suffix ("float[:, :](order=F)" of the column-major + # markers), so it is removed in a copy of the source; only the names are compared here. + pyccel_source = tmp_path / source + pyccel_source.write_text((ARGS_DIR / source).read_text().replace("(order=F)", "")) + pyccel_fields = [field.name for field in CudaStruct.from_pyccel_class(pyccel_source, class_name).fields] cuda_fields = [field.name for field in cuda_class.struct.fields] assert [name for name in cuda_fields if name not in cuda_only] == pyccel_fields diff --git a/src/struphy/pic/tests/test_kernel_setup.py b/src/struphy/pic/tests/test_kernel_setup.py index 00b9d127c..ae3de77f1 100644 --- a/src/struphy/pic/tests/test_kernel_setup.py +++ b/src/struphy/pic/tests/test_kernel_setup.py @@ -33,7 +33,7 @@ @pytest.fixture def marker_args(): - markers = np.zeros((2, 40)) + markers = np.zeros((2, 40), order="F") markers[0, :7] = (0.4, 0.3, 0.2, 2.0, 4.0, 6.0, 2.0) markers[0, 8:14] = markers[0, :6] markers[0, 11] = 1.0 # Initial parallel velocity. @@ -161,7 +161,7 @@ def test_sph_kernel_gradients_vanish_at_zero(kernel_type, point): ) def test_gc_discrete_gradient_residual_at_zero_velocity(derham_args, kernel, n_fields): # A marker at rest (v = mu = 0, no E-field) stays put; the residual must be 0, not NaN (issue #589). - markers = np.zeros((1, 40)) + markers = np.zeros((1, 40), order="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. diff --git a/src/struphy/pic/tests/test_marker_layout.py b/src/struphy/pic/tests/test_marker_layout.py new file mode 100644 index 000000000..7ccc9374f --- /dev/null +++ b/src/struphy/pic/tests/test_marker_layout.py @@ -0,0 +1,157 @@ +"""The marker buffer is column-major (struct-of-arrays), see :mod:`struphy.pic.marker_layout`. + +The CuPy cases run on a GPU, or without one on cunumpy's fake CuPy with every CUDA launch emulated on the CPU:: + + mpirun -n 1 env CUNUMPY_FAKE_CUPY=1 pytest src/struphy/pic/tests/test_marker_layout.py + mpirun -n 2 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi src/struphy/pic/tests/test_marker_layout.py +""" + +import contextlib + +import cunumpy as xp +import numpy as np +import pytest + +from struphy.models.tests.gpu_e2e import device_context, requires_cupy_backend +from struphy.pic.marker_layout import ( + ROW_ALIGNMENT, + as_marker_array, + check_marker_layout, + has_marker_layout, + marker_array, + padded_n_rows, +) + +# CuPy on a GPU, or cunumpy's fake CuPy +BACKENDS = ["numpy", pytest.param("cupy", marks=requires_cupy_backend)] + + +def test_padded_n_rows(): + assert [padded_n_rows(n) for n in (0, 1, 31, 32, 33, 64, 65)] == [0, 32, 32, 32, 64, 64, 96] + assert ROW_ALIGNMENT == 32 + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_layout_helpers(backend): + with xp.use_backend(backend): + markers = marker_array(64, 7) + assert markers.shape == (64, 7) and has_marker_layout(markers) + check_marker_layout(markers) + # no copy for an array in the marker layout, nor for a block of its columns + assert as_marker_array(markers) is markers + assert has_marker_layout(markers[:, 2:5]) + + # row gathers and .copy() are C-ordered: converted (with a copy) before they go to a kernel + rows = markers[xp.arange(0, 64, 2)] + assert not has_marker_layout(rows) + with pytest.raises(TypeError, match="Fortran-ordered"): + check_marker_layout(rows) + converted = as_marker_array(rows) + assert has_marker_layout(converted) + np.testing.assert_array_equal(xp.to_numpy(converted), xp.to_numpy(rows)) + + +def test_pyccel_rejects_row_major_markers(): + """pyccel checks the memory order of the marker arrays at every call, so a missed conversion fails loudly.""" + from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments + + args = (np.ones(4, dtype=bool), 4, 3, 6, 7, 8, 14, 17, 18, 4, np.zeros(3, dtype=np.int64), False) + MarkerArguments(np.zeros((4, 25), order="F"), *args) + with pytest.raises(TypeError, match="ordering"): + MarkerArguments(np.zeros((4, 25)), *args) + + +def _device_context(backend): + """Emulated CUDA launches on the fake CuPy (nothing on a GPU or on NumPy).""" + return device_context() if backend == "cupy" else contextlib.nullcontext() + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_markers_stay_column_major(backend): + """Loading, pushing, boundary conditions, sorting between ranks, evaluation at markers and binning keep the + buffer column-major, and evaluation at row-major copies of the markers gives the same values.""" + from maybempi import MPI + + from struphy import BoundaryParameters, LoadingParameters, WeightsParameters, domains, perturbations + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.ode.utils import ButcherTableau + from struphy.pic.particles import Particles6D + from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage + from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic + from struphy.pic.pushing.pusher import Pusher + from struphy.topology.grids import TensorProductGrid + + comm = MPI.COMM_WORLD + with xp.use_backend(backend), _device_context(backend): + domain = domains.Colella(Lx=2.0, Ly=3.0, alpha=0.1, Lz=4.0) + derham = Derham(TensorProductGrid(num_elements=(8, 6, 4)), DerhamOptions(degree=(2, 3, 1)), comm=comm) + particles = Particles6D( + comm_world=comm, + domain_decomp=(derham.domain_array, derham.domain_decomposition.nprocs), + loading_params=LoadingParameters(Np=1000, seed=1234, spatial="uniform"), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(bc=("periodic", "reflect", "periodic")), + domain=domain, + ) + particles.draw_markers() + particles.mpi_sort_markers() + particles.initialize_weights() + + buffer = particles.markers + check_marker_layout(buffer) + assert buffer.shape[0] % ROW_ALIGNMENT == 0 + + butcher = ButcherTableau("rk4") + push_eta = Pusher( + particles, + push_eta_stage, + (xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(butcher.c), butcher.n_stages), + domain.args_domain, + alpha_in_kernel=1.0, + n_stages=butcher.n_stages, + pushes_eta=True, + ) + _, b2 = create_equal_random_arrays(derham.V2fem, seed=3456, flattened=True) + push_vxb = Pusher( + particles, + push_vxb_analytic, + (derham.args_derham, b2[0]._data, b2[1]._data, b2[2]._data), + domain.args_domain, + alpha_in_kernel=1.0, + pushes_eta=False, + ) + for _ in range(3): + push_eta(0.1) + push_vxb(0.1) + particles.mpi_sort_markers() + + # the buffer is never replaced, only written to + assert particles.markers is buffer + check_marker_layout(particles.markers) + + # evaluation at row-major copies of the markers equals evaluation at the buffer (converted at the entry) + rows = particles.markers_wo_holes_and_ghost + assert rows.flags.c_contiguous + n_valid = rows.shape[0] + det_buffer = domain.jacobian_det(particles.markers, remove_outside=False)[particles.valid_mks] + np.testing.assert_array_equal(xp.to_numpy(domain.jacobian_det(rows)), xp.to_numpy(det_buffer)) + np.testing.assert_array_equal( + xp.to_numpy(domain.pull(lambda x, y, z: x * y + z, rows, kind="3")), + xp.to_numpy(domain.pull(lambda x, y, z: x * y + z, as_marker_array(rows), kind="3")), + ) + + # spline evaluation at row-major markers (local: every rank has its own markers) + field = derham.create_spline_function("f", "H1") + field.initialize_coeffs(perturbations=perturbations.ModesCos(ls=(1,), ms=(1,), ns=(1,), amps=(1.0,))) + values = field(rows, local=True) + assert values.shape == (n_valid,) + np.testing.assert_array_equal(xp.to_numpy(values), xp.to_numpy(field(as_marker_array(rows), local=True))) + + # binning reads single columns of the valid markers + f_slice, _ = particles.binning( + (True, False, False, True, False, False), (xp.linspace(0.0, 1.0, 9), xp.linspace(-4.0, 4.0, 9)) + ) + assert f_slice.shape == (8, 8) + check_marker_layout(particles.markers) diff --git a/src/struphy/pic/tests/test_pole.py b/src/struphy/pic/tests/test_pole.py index 92b682e92..31a2b1874 100644 --- a/src/struphy/pic/tests/test_pole.py +++ b/src/struphy/pic/tests/test_pole.py @@ -47,7 +47,7 @@ def test_bc_kernel_axis_map(): domain = domains.HollowCylinder(a1=0.0, a2=1.0, Lz=4.0) - markers = np.zeros((3, N_COLS)) + markers = np.zeros((3, N_COLS), order="F") 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) @@ -77,7 +77,7 @@ def test_bc_kernel_reflect_vdim(vdim): domain = domains.HollowCylinder(a1=0.2, a2=1.0, Lz=4.0) - markers = np.zeros((1, N_COLS)) + markers = np.zeros((1, N_COLS), order="F") 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) diff --git a/src/struphy/pic/tests/test_reflect.py b/src/struphy/pic/tests/test_reflect.py index 8b2e69de7..5830346e6 100644 --- a/src/struphy/pic/tests/test_reflect.py +++ b/src/struphy/pic/tests/test_reflect.py @@ -27,7 +27,8 @@ def test_reflect_in_place(backend, axis, count): expected = original.copy() expected[indices, 2 * (3 + axis)] *= -1 with xp.use_backend(backend): - storage = xp.asarray(original.copy()) + # every other column of a column-major buffer: strided, in the marker layout + storage = xp.asfortranarray(xp.asarray(original)) markers = storage[:, ::2] outside_inds = xp.asarray(np.repeat(indices, 2))[::2] domain = Cuboid(r1=2.0, r2=3.0, r3=4.0) diff --git a/src/struphy/pic/utilities_kernels.py b/src/struphy/pic/utilities_kernels.py index 73ec15bdb..d781ac4ca 100644 --- a/src/struphy/pic/utilities_kernels.py +++ b/src/struphy/pic/utilities_kernels.py @@ -19,7 +19,7 @@ @stack_array("dfm", "norm_b1", "b") def eval_magnetic_energy( - markers: "float[:,:]", + markers: "float[:,:](order=F)", # noqa: F821 args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, diff --git a/src/struphy/propagators/vlasov_ampere_coupling.py b/src/struphy/propagators/vlasov_ampere_coupling.py index 77a129ae0..18dbe2884 100644 --- a/src/struphy/propagators/vlasov_ampere_coupling.py +++ b/src/struphy/propagators/vlasov_ampere_coupling.py @@ -233,20 +233,11 @@ def __call__(self, dt): logger.info(f"Iterations for VlasovMaxwell: {info['niter']}") logger.info(f"Maxdiff e1 for VlasovMaxwell: {max_de}") particles = self.variables.ions.particles - buffer_idx = particles.bufferindex - max_diff = xp.max( - xp.abs( - xp.sqrt( - particles.markers_wo_holes[:, 3] ** 2 - + particles.markers_wo_holes[:, 4] ** 2 - + particles.markers_wo_holes[:, 5] ** 2, - ) - - xp.sqrt( - particles.markers_wo_holes[:, buffer_idx + 3] ** 2 - + particles.markers_wo_holes[:, buffer_idx + 4] ** 2 - + particles.markers_wo_holes[:, buffer_idx + 5] ** 2, - ), - ), - ) + # the velocities before the push are saved at first_pusher_idx + 3 (see Pusher._push); + # gather only the six columns needed, not the whole rows + not_holes = ~particles.holes + v_new = particles.markers[not_holes, 3:6] + v_old = particles.markers[not_holes, particles.first_pusher_idx + 3 : particles.first_pusher_idx + 6] + max_diff = xp.max(xp.abs(xp.linalg.norm(v_new, axis=1) - xp.linalg.norm(v_old, axis=1))) logger.info(f"Maxdiff |v| for VlasovMaxwell: {max_diff}") logger.info("") From e41792274838f984c8b47045e63b5af36b2acd90 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 14 Oct 2026 21:39:29 +0200 Subject: [PATCH 02/12] Marker layout: accept single rows as pyccel does, C-ordered diffusion noise pyccel checks order=F by the strides (strides[0] <= strides[1]), not by the F-contiguous flag: a single row in C order is F-contiguous for NumPy but has decreasing strides, and was passed through unchanged by as_marker_array (domain evaluation at one point, SPH boundary mirroring). has_marker_layout now checks the strides too, and as_marker_array copies in that case. The noise of PushRandomDiffusion was a copy of markers[:, :3] and kept the column-major order, but the kernel declares it C-ordered. get_kinetic_energy_particles converts the markers it is given. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/marker_layout.py | 14 ++++++++++---- src/struphy/pic/tests/test_marker_layout.py | 11 +++++++++++ src/struphy/pic/utilities.py | 3 ++- src/struphy/propagators/push_random_diffusion.py | 3 ++- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/src/struphy/pic/marker_layout.py b/src/struphy/pic/marker_layout.py index f2f89e61d..e1736190c 100644 --- a/src/struphy/pic/marker_layout.py +++ b/src/struphy/pic/marker_layout.py @@ -40,13 +40,19 @@ def marker_array(n_rows: int, n_cols: int): def as_marker_array(markers): - """``markers`` in :data:`MARKER_ORDER`, for passing to a kernel; no copy if it already has that order.""" - return xp.asfortranarray(markers) + """``markers`` in :data:`MARKER_ORDER`, for passing to a kernel; no copy if it already has that layout.""" + if has_marker_layout(markers): + return markers + return xp.array(markers, order=MARKER_ORDER, copy=True) def has_marker_layout(markers) -> bool: - """Whether ``markers`` is a 2D array in :data:`MARKER_ORDER` (as the kernels expect).""" - return markers.ndim == 2 and bool(markers.flags.f_contiguous) + """Whether ``markers`` is a 2D array in :data:`MARKER_ORDER`, as the kernels expect. + + Besides the F-contiguity, the strides must not decrease along the axes, which is what pyccel checks for + ``order=F``: a single row in C order is F-contiguous for NumPy, but its strides are ``(8 * n_cols, 8)``. + """ + return markers.ndim == 2 and bool(markers.flags.f_contiguous) and markers.strides[0] <= markers.strides[1] def check_marker_layout(markers, name: str = "markers"): diff --git a/src/struphy/pic/tests/test_marker_layout.py b/src/struphy/pic/tests/test_marker_layout.py index 7ccc9374f..92372672d 100644 --- a/src/struphy/pic/tests/test_marker_layout.py +++ b/src/struphy/pic/tests/test_marker_layout.py @@ -50,6 +50,11 @@ def test_layout_helpers(backend): assert has_marker_layout(converted) np.testing.assert_array_equal(xp.to_numpy(converted), xp.to_numpy(rows)) + # a single row in C order is F-contiguous for NumPy, but its strides decrease, which pyccel rejects + row = xp.asarray([[0.2, 0.4, 0.6]]) + assert row.flags.f_contiguous and not has_marker_layout(row) + assert has_marker_layout(as_marker_array(row)) + def test_pyccel_rejects_row_major_markers(): """pyccel checks the memory order of the marker arrays at every call, so a missed conversion fails loudly.""" @@ -60,6 +65,12 @@ def test_pyccel_rejects_row_major_markers(): with pytest.raises(TypeError, match="ordering"): MarkerArguments(np.zeros((4, 25)), *args) + # a single row: as_marker_array gives it the strides pyccel accepts + one = (np.ones(1, dtype=bool), 1, *args[2:]) + with pytest.raises(TypeError, match="ordering"): + MarkerArguments(np.zeros((1, 25)), *one) + MarkerArguments(as_marker_array(np.zeros((1, 25))), *one) + def _device_context(backend): """Emulated CUDA launches on the fake CuPy (nothing on a GPU or on NumPy).""" diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index fc3e4725b..229361d2e 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -3,6 +3,7 @@ from struphy.feec.utilities import get_quad_grids from struphy.pic.diagnostics.kernels.canonical_kinetic_particles import canonical_kinetic_particles from struphy.pic.diagnostics.kernels.thermal_energy import thermal_energy +from struphy.pic.marker_layout import as_marker_array def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): @@ -24,7 +25,7 @@ def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): res = xp.empty(1, dtype=float) canonical_kinetic_particles( res, - particles.markers, + as_marker_array(particles.markers), derham.args_derham, domain.args_domain, fe_coeffs.blocks[0]._data, diff --git a/src/struphy/propagators/push_random_diffusion.py b/src/struphy/propagators/push_random_diffusion.py index 2101ef0f4..dce8153e7 100644 --- a/src/struphy/propagators/push_random_diffusion.py +++ b/src/struphy/propagators/push_random_diffusion.py @@ -110,7 +110,8 @@ def allocate(self): particles = self.variables.var.particles - self._noise = array(particles.markers[:, :3]) + # C-ordered, as the kernel declares it (a copy of the column-major markers would keep their order) + self._noise = array(particles.markers[:, :3], order="C") self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: From db4a1d956ab8447451f4a1a0af6e574ac1b7311a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 14 Oct 2026 22:24:29 +0200 Subject: [PATCH 03/12] LinearVlasovAmpere verification: check the net charge over the ranks of the clone test_delta_f_initial_condition checked the mean of the delta-f weights per rank, assuming every rank holds a random subset of the markers. The markers are sorted into the domain of their rank, so on more than two ranks each rank holds a part of the cos(kx) perturbation with a net charge of its own (mean/rms = (2/pi)/(1/sqrt(2)) = 0.9 on four ranks). The sums are now reduced over the clone's communicator; each clone holds a random subset of all markers. The test failed on four ranks before this branch as well; testmon had not selected it on recent PRs. Co-Authored-By: Claude Opus 5.5 --- .../test_verif_LinearVlasovAmpereOneSpecies.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py b/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py index 832cd1a57..ecdcc3a9f 100644 --- a/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/verification/test_verif_LinearVlasovAmpereOneSpecies.py @@ -4,6 +4,7 @@ import cunumpy as xp import h5py +import numpy as np from maybempi import MPI from struphy import ( @@ -100,11 +101,17 @@ def test_delta_f_initial_condition(exit_before_run: bool = False): # 1. no net charge: the weights of a cosine perturbation cancel over the box, so their # mean is Monte-Carlo noise, small against their own spread. The control variate applied # by mistake made every weight equal, i.e. |mean| == max|w|. - # Checked per rank: each one holds a random subset of the markers, so its own mean is - # as good an estimator of the net charge as the global one. - weights = xp.asarray(model.kinetic_ions.var.particles.weights) - mean = float(weights.mean()) - rms = float(xp.sqrt((weights**2).mean())) + # Summed over the ranks of the clone: the markers are sorted into the domain of their rank, + # so on more than two ranks each one holds a part of the cosine with a net charge of its + # own. Each clone holds a random subset of all markers, so its mean is as good an estimator + # of the net charge as the global one. + particles = model.kinetic_ions.var.particles + weights = xp.to_numpy(particles.weights) + sums = np.array([weights.sum(), (weights**2).sum(), weights.size], dtype=float) + if particles.mpi_comm is not None: + particles.mpi_comm.Allreduce(MPI.IN_PLACE, sums, op=MPI.SUM) + mean = sums[0] / sums[2] + rms = np.sqrt(sums[1] / sums[2]) assert abs(mean) < 0.1 * rms, ( f"Assertion for delta-f weights failed: net charge {mean =} is not small against {rms =}; " "the markers do not carry the density perturbation alone." From 301223cb6f5472481dfc6846dbf7cfd2593a7184 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 15 Oct 2026 17:18:31 +0200 Subject: [PATCH 04/12] Marker buffer as (n_cols, n_rows) in C order: layout module and kernels Instead of storing the (n_rows, n_cols) marker buffer in Fortran order, the buffer is now a plain C-ordered array of shape (n_cols, n_rows), one column per marker, indexed markers[j, ip]. The memory layout (struct-of-arrays) is the same, but no array needs an order flag: the pyccel kernels declare a plain float[:,:] again and index markers[j, ip], with the marker count shape(markers)[1]; the CUDA kernels index markers(j, ip). marker_layout keeps ROW_ALIGNMENT/padded_n_rows and marker_array(n_cols, n_rows); MARKER_ORDER, as_marker_array, has_marker_layout and check_marker_layout are removed. MarkerArguments/CudaMarkerArguments take n_markers = markers.shape[1] (C-contiguous again on the CUDA side). cunumpy infers the CUDA launch size from the first axis of the first array, which is now n_cols: pushers and accumulators pass n_threads=args_markers.n_markers, the geometry and spline marker kernels markers.shape[1]. The pusher preamble and prepare_push work on rows. The random-diffusion noise is (3, n_rows); the symmetric sobol sampling kernels fill (3 + vdim, n) arrays. Orbit post-processing keeps its (n, q) arrays. Co-Authored-By: Claude Opus 5.5 --- .../eval_spline_mpi_markers_cuda.cu | 16 +-- .../eval_spline_mpi_markers_kernels.py | 18 ++-- .../kernel_evaluate_pic_cuda.cu | 22 ++--- .../kernel_evaluate_pic_kernels.py | 10 +- .../kernel_pullpush_pic_cuda.cu | 22 ++--- .../kernel_pullpush_pic_kernels.py | 14 +-- src/struphy/kernel_arguments/array_view.cuh | 8 +- .../kernel_arguments/pusher_args_cuda.py | 18 ++-- .../kernel_arguments/pusher_args_kernels.py | 6 +- .../cc_lin_mhd_5d_D_kernels.py | 14 +-- .../cc_lin_mhd_5d_M_kernels.py | 14 +-- .../cc_lin_mhd_5d_curlb_kernels.py | 14 +-- .../cc_lin_mhd_5d_gradB_kernels.py | 16 +-- .../cc_lin_mhd_5d_gradB_dg_kernels.py | 14 +-- .../cc_lin_mhd_5d_gradB_dg_init_kernels.py | 16 +-- .../cc_lin_mhd_6d_1_kernels.py | 14 +-- .../cc_lin_mhd_6d_2_kernels.py | 16 +-- .../charge_density_0form_cuda.cu | 16 +-- .../charge_density_0form_kernels.py | 12 +-- .../div_u_weak_1form_kernels.py | 18 ++-- .../gc_density_0form_kernels.py | 12 +-- .../gc_mag_density_0form_kernels.py | 14 +-- .../linear_vlasov_ampere_cuda.cu | 26 ++--- .../linear_vlasov_ampere_kernels.py | 20 ++-- .../pc_lin_mhd_6d/pc_lin_mhd_6d_kernels.py | 14 +-- .../pc_lin_mhd_6d_full_kernels.py | 14 +-- .../vlasov_maxwell/vlasov_maxwell_cuda.cu | 24 ++--- .../vlasov_maxwell/vlasov_maxwell_kernels.py | 20 ++-- .../pic/accumulation/particles_to_grid.py | 4 +- .../canonical_kinetic_particles_kernels.py | 16 +-- ...al_canonical_toroidal_moment_5d_kernels.py | 28 +++--- ...al_canonical_toroidal_moment_6d_kernels.py | 28 +++--- .../eval_energy_5d/eval_energy_5d_kernels.py | 22 ++--- .../eval_gradB_ediff_kernels.py | 14 +-- .../eval_guiding_center_from_6d_kernels.py | 36 +++---- ...eval_magnetic_background_energy_kernels.py | 18 ++-- .../eval_magnetic_energy_PBb_kernels.py | 18 ++-- .../eval_magnetic_moment_5d_kernels.py | 18 ++-- src/struphy/pic/marker_layout.py | 64 ++++-------- .../bstar_2form/bstar_2form_kernels.py | 12 +-- .../bstar_parallel_3form_cuda.cu | 18 ++-- .../bstar_parallel_3form_kernels.py | 12 +-- .../driftkinetic_hamiltonian_kernels.py | 14 +-- .../grad_driftkinetic_hamiltonian_kernels.py | 10 +- .../push_bxu_H1vec/push_bxu_H1vec_kernels.py | 12 +-- .../push_bxu_Hcurl/push_bxu_Hcurl_kernels.py | 12 +-- .../push_bxu_Hdiv/push_bxu_Hdiv_kernels.py | 12 +-- .../push_bxu_Hdiv_pauli_kernels.py | 10 +- ...h_deterministic_diffusion_stage_kernels.py | 16 +-- .../push_eta_stage/push_eta_stage_cuda.cu | 22 ++--- .../push_eta_stage/push_eta_stage_kernels.py | 10 +- ...tar_discrete_gradient_1st_order_kernels.py | 34 +++---- ...crete_gradient_1st_order_newton_kernels.py | 34 +++---- ...tar_discrete_gradient_2nd_order_kernels.py | 30 +++--- ...sh_gc_Bstar_explicit_multistage_kernels.py | 20 ++-- ..._gc_bxEstar_dg_1st_order_newton_kernels.py | 30 +++--- ...tar_discrete_gradient_1st_order_kernels.py | 22 ++--- ...tar_discrete_gradient_2nd_order_kernels.py | 20 ++-- ..._gc_bxEstar_explicit_multistage_kernels.py | 12 +-- .../push_gc_cc_J1_H1vec_kernels.py | 12 +-- .../push_gc_cc_J1_Hcurl_kernels.py | 12 +-- .../push_gc_cc_J1_Hdiv_kernels.py | 12 +-- .../push_gc_cc_J2_dg_Hdiv_kernels.py | 14 +-- .../push_gc_cc_J2_dg_init_Hdiv_kernels.py | 12 +-- .../push_gc_cc_J2_stage_H1vec_kernels.py | 18 ++-- .../push_gc_cc_J2_stage_Hdiv_kernels.py | 20 ++-- .../push_pc_GXu/push_pc_GXu_kernels.py | 12 +-- .../push_pc_GXu_full_kernels.py | 12 +-- .../push_pc_eta_stage_H1vec_kernels.py | 18 ++-- .../push_pc_eta_stage_Hcurl_kernels.py | 18 ++-- .../push_pc_eta_stage_Hdiv_kernels.py | 18 ++-- .../push_random_diffusion_stage_kernels.py | 4 +- .../push_v_sph_pressure_kernels.py | 20 ++-- .../push_v_sph_pressure_ideal_gas_kernels.py | 20 ++-- .../push_v_viscosity_kernels.py | 18 ++-- .../push_v_with_efield_cuda.cu | 14 +-- .../push_v_with_efield_kernels.py | 8 +- .../push_vxb_analytic_cuda.cu | 18 ++-- .../push_vxb_analytic_kernels.py | 12 +-- .../push_vxb_implicit_cuda.cu | 18 ++-- .../push_vxb_implicit_kernels.py | 12 +-- .../push_weights_with_efield_lin_va_cuda.cu | 22 ++--- ...push_weights_with_efield_lin_va_kernels.py | 18 ++-- .../pushing/kernels/reflect/reflect_cuda.cu | 2 +- .../kernels/reflect/reflect_kernels.py | 8 +- .../sph_isotherm_kappa_kernels.py | 4 +- .../sph_mean_velocity_coeffs_kernels.py | 22 ++--- .../sph_pressure_coeffs_kernels.py | 24 ++--- .../sph_viscosity_tensor_kernels.py | 18 ++-- .../unit_b_1form/unit_b_1form_kernels.py | 8 +- src/struphy/pic/pushing/prepare_push_cuda.cu | 26 ++--- src/struphy/pic/pushing/pusher.py | 21 ++-- .../pic/pushing/pusher_utilities_kernels.cuh | 52 +++++----- .../pic/pushing/pusher_utilities_kernels.py | 98 +++++++++---------- src/struphy/pic/sampling_kernels.py | 8 +- src/struphy/pic/sorting_kernels.py | 34 +++---- src/struphy/pic/sph_eval_kernels.py | 30 +++--- src/struphy/pic/utilities_kernels.py | 18 ++-- 98 files changed, 888 insertions(+), 917 deletions(-) diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu index 08314d224..99ef61252 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_cuda.cu @@ -9,25 +9,25 @@ using namespace struphy_cuda; /** * Flat (marker) evaluation of a distributed spline, as in evaluation_kernels_3d.eval_spline_mpi_markers. * - * One thread per marker row. + * One thread per marker. * - * @param markers Marker coordinates (Np x 3 or wider, any strides); rows flagged with -1 in the first - * column are not on the process domain and are skipped. + * @param markers Marker coordinates (3 or more x Np, one column per marker, any strides); columns flagged with -1 + * in the first row are not on the process domain and are skipped. * @param _data Spline coefficients of the current process, any strides. * @param args_spline Kind of 1d basis (0 = N-spline, 1 = D-spline), spline degrees (1 to 8) and knot vectors of V0 * in each direction, and start indices of the splines on the current process. - * @param values Output values S_p = S(*markers[p, :]), one per row; skipped rows are left unchanged. + * @param values Output values S_p = S(*markers[:, p]), one per marker; skipped markers are left unchanged. */ extern "C" __global__ void eval_spline_mpi_markers(Array2D markers, Array3D _data, SplineArgs args_spline, Array1D values) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable long long ip = (long long)blockDim.x * blockIdx.x + threadIdx.x; - long long Np = markers.shape[0]; + long long Np = markers.shape[1]; if (ip >= Np) return; // point not in process domain - if (markers(ip, 0) == -1.) return; + if (markers(0, ip) == -1.) return; - values(ip) = evaluation_kernels_3d::eval_spline_mpi(markers(ip, 0), markers(ip, 1), markers(ip, 2), _data, + values(ip) = evaluation_kernels_3d::eval_spline_mpi(markers(0, ip), markers(1, ip), markers(2, ip), _data, args_spline); } diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py index 8a63df173..0ebbd256e 100644 --- a/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_markers/eval_spline_mpi_markers_kernels.py @@ -8,7 +8,7 @@ def eval_spline_mpi_markers( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", _data: "float[:,:,:]", args_spline: "SplineArguments", values: "float[:]", @@ -19,8 +19,8 @@ def eval_spline_mpi_markers( Parameters ---------- markers : array[float] - Marker coordinates in format [Np, 3]; markers not on local process domain must be flagged as -1. - Spline values are obtained as S_p = S(*markers[p, :]). + Marker coordinates in format [3, Np], one column per marker; markers not on local process domain must be flagged as -1. + Spline values are obtained as S_p = S(*markers[:, p]). _data : array[float] The spline coefficients c_ijk. @@ -30,19 +30,19 @@ def eval_spline_mpi_markers( current process. values : array[float] - Return 1D array for spline values S_p = S(*markers[p, :]). + Return 1D array for spline values S_p = S(*markers[:, p]). """ - Np = shape(markers)[0] + Np = shape(markers)[1] for ip in range(Np): - if markers[ip, 0] == -1: + if markers[0, ip] == -1: continue # point not in process domain values[ip] = eval_spline_mpi( - markers[ip, 0], - markers[ip, 1], - markers[ip, 2], + markers[0, ip], + markers[1, ip], + markers[2, ip], _data, args_spline, ) diff --git a/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_cuda.cu b/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_cuda.cu index 9644a2fca..3c6f4cae2 100644 --- a/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_cuda.cu +++ b/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_cuda.cu @@ -9,27 +9,27 @@ using namespace struphy_cuda; * * CUDA-only helper for the condition `e1 < 0.0 or e1 > 1.0 or ...` (a NaN coordinate counts as inside, as in pyccel). * - * @param markers Marker array view; the logical position is in columns 0-2. - * @param i Marker row. + * @param markers Marker array view; the logical position is in rows 0-2 (one column per marker). + * @param i Marker index. * @return True if the marker is outside (a hole has coordinates -1). */ __device__ inline bool marker_is_outside(const Array2D& markers, long long i) { - double e1 = markers(i, 0); - double e2 = markers(i, 1); - double e3 = markers(i, 2); + double e1 = markers(0, i); + double e2 = markers(1, i); + double e3 = markers(2, i); return e1 < 0.0 || e1 > 1.0 || e2 < 0.0 || e2 > 1.0 || e3 < 0.0 || e3 > 1.0; } /** * Evaluation of metric coefficients for given markers, as in kernel_evaluate_pic_kernels.kernel_evaluate_pic. * - * One thread per marker row i. The pyccel kernel writes the result of the i-th marker to row `counter` of mat_f + * One thread per marker i. The pyccel kernel writes the result of the i-th marker to row `counter` of mat_f * and returns the final counter. A CUDA kernel returns nothing, so the counter is not returned; struphy's caller * (Domain._evaluate_metric_coefficient) passes remove_outside=false, for which counter == i, and removes the rows of * outside markers itself, on both backends. With remove_outside=true each thread counts the inside markers before * row i (O(i) work per thread, O(N^2) in total): correct, so that both versions fill mat_f identically, but slow. * - * @param markers Evaluation points in marker format (eta1, eta2, eta3 = markers[:, 0:3]), any strides. + * @param markers Evaluation points in marker format (eta1, eta2, eta3 = markers[0:3], one column per marker), any strides. * @param kind_coeff Which coefficient: -1 identity, 0 mapping F, 1 DF, 2 det(DF), 3 DF^(-1), 4 G, 5 G^(-1). * @param args Mapping arguments; every spline and analytic mapping. * @param mat_f Output of shape (N, 3, 3), any strides. Rows of outside markers are set to -1 (identity: their @@ -40,13 +40,13 @@ __device__ inline bool marker_is_outside(const Array2D& markers, long lo extern "C" __global__ void kernel_evaluate_pic(Array2D markers, int kind_coeff, DomainArgs args, Array3D mat_f, bool remove_outside, bool avoid_round_off) { long long i = (long long)blockDim.x * blockIdx.x + threadIdx.x; - if (i >= markers.shape[0]) return; + if (i >= markers.shape[1]) return; double tmp0[3], tmp1[9], tmp2[9], tmp3[9], out[9]; - double e1 = markers(i, 0); - double e2 = markers(i, 1); - double e3 = markers(i, 2); + double e1 = markers(0, i); + double e2 = markers(1, i); + double e3 = markers(2, i); bool outside = marker_is_outside(markers, i); if (outside && remove_outside) return; diff --git a/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py b/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py index 74bff0066..ef735312b 100644 --- a/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py +++ b/src/struphy/geometry/kernels/kernel_evaluate_pic/kernel_evaluate_pic_kernels.py @@ -10,7 +10,7 @@ @stack_array("tmp0", "tmp1", "tmp2", "tmp3", "out") def kernel_evaluate_pic( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", kind_coeff: int, args: "DomainArguments", mat_f: "float[:,:,:]", @@ -36,13 +36,13 @@ def kernel_evaluate_pic( tmp3 = zeros((3, 3), dtype=float) out = zeros((3, 3), dtype=float) - np = shape(markers)[0] + np = shape(markers)[1] counter = 0 for i in range(np): - e1 = markers[i, 0] - e2 = markers[i, 1] - e3 = markers[i, 2] + e1 = markers[0, i] + e2 = markers[1, i] + e3 = markers[2, i] if e1 < 0.0 or e1 > 1.0 or e2 < 0.0 or e2 > 1.0 or e3 < 0.0 or e3 > 1.0: if remove_outside: diff --git a/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_cuda.cu b/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_cuda.cu index 466ce645f..934a12c1f 100644 --- a/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_cuda.cu +++ b/src/struphy/geometry/kernels/kernel_pullpush_pic/kernel_pullpush_pic_cuda.cu @@ -9,14 +9,14 @@ using namespace struphy_cuda; * * CUDA-only helper for the condition `e1 < 0.0 or e1 > 1.0 or ...` (a NaN coordinate counts as inside, as in pyccel). * - * @param markers Marker array view; the logical position is in columns 0-2. - * @param i Marker row. + * @param markers Marker array view; the logical position is in rows 0-2 (one column per marker). + * @param i Marker index. * @return True if the marker is outside (a hole has coordinates -1). */ __device__ inline bool marker_is_outside(const Array2D& markers, long long i) { - double e1 = markers(i, 0); - double e2 = markers(i, 1); - double e3 = markers(i, 2); + double e1 = markers(0, i); + double e2 = markers(1, i); + double e3 = markers(2, i); return e1 < 0.0 || e1 > 1.0 || e2 < 0.0 || e2 > 1.0 || e3 < 0.0 || e3 > 1.0; } @@ -24,7 +24,7 @@ __device__ inline bool marker_is_outside(const Array2D& markers, long lo * Pull-backs, push-forwards and transformations for given markers, as in * kernel_pullpush_pic_kernels.kernel_pullpush_pic. * - * One thread per marker row i (launched with n_threads = markers.shape[0], since the first array is a). The pyccel + * One thread per marker i (launched with n_threads = markers.shape[1]). The pyccel * kernel reads row `counter_a` of a and writes row `counter_o` of out, and returns the final counter_o. A CUDA kernel * returns nothing, so the counter is not returned; struphy's caller (Domain._pull_push_transform) passes * remove_outside=false and an `a` with one row per marker, for which counter_a == counter_o == i, and removes the @@ -33,7 +33,7 @@ __device__ inline bool marker_is_outside(const Array2D& markers, long lo * * @param a Values to transform, shape (N, n_comp) (with holes: one row per marker) or (n_inside, n_comp) (without * holes), n_comp 1 (scalar kinds) or 3 (vector kinds), any strides. - * @param markers Evaluation points in marker format (eta1, eta2, eta3 = markers[:, 0:3]), any strides. + * @param markers Evaluation points in marker format (eta1, eta2, eta3 = markers[0:3], one column per marker), any strides. * @param kind_transform Which general transformation: 0 pull, 1 push, otherwise tran. * @param kind_fun Which detailed transformation (see pull, push, tran in transform_kernels.cuh). * @param args_domain Mapping arguments; every spline and analytic mapping. @@ -45,7 +45,7 @@ extern "C" __global__ void kernel_pullpush_pic(Array2D a, Array2D out, bool remove_outside) { long long i = (long long)blockDim.x * blockIdx.x + threadIdx.x; - long long np = markers.shape[0]; + long long np = markers.shape[1]; if (i >= np) return; // CUDA-only: per-thread copies of the pyccel stack arrays tmp1 (length a.shape[1]) and tmp2 (out.shape[1]), @@ -57,9 +57,9 @@ extern "C" __global__ void kernel_pullpush_pic(Array2D a, Array2D 1.0 or e2 < 0.0 or e2 > 1.0 or e3 < 0.0 or e3 > 1.0: diff --git a/src/struphy/kernel_arguments/array_view.cuh b/src/struphy/kernel_arguments/array_view.cuh index 85ca437fb..9f07d50b0 100644 --- a/src/struphy/kernel_arguments/array_view.cuh +++ b/src/struphy/kernel_arguments/array_view.cuh @@ -8,10 +8,10 @@ // for kernel arguments and in Argument._pack (struphy/utils/cuda_arguments.py) for struct fields. // // The layout is NumPy's: for a C-contiguous (row-major) array of shape (n0, n1, n2), -// strides = (n1 * n2, n2, 1), so the last index runs fastest. The marker buffer of shape -// (n_markers, n_cols) is F-contiguous (column-major, struphy.pic.marker_layout): strides = (1, n_markers) -// and markers(ip, j) is data[ip + j * n_markers]. A sliced or transposed CuPy array keeps its own -// strides, so it is indexed correctly without a copy. +// strides = (n1 * n2, n2, 1), so the last index runs fastest; for the marker buffer of shape +// (n_cols, n_markers) (struphy.pic.marker_layout), strides = (n_markers, 1) and markers(j, ip) is +// data[j * n_markers + ip]. A sliced or transposed CuPy array keeps its own strides, so it is indexed +// correctly without a copy. template struct Array1D { T* data; long long shape[1], strides[1]; diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index dbdeed54e..16e4361a7 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -11,21 +11,17 @@ import numpy as np from cunumpy.arguments import CudaStructArguments -from struphy.pic.marker_layout import MARKER_ORDER - -def _device_array(name, arr, dtype, ndim=None, order="C"): - """Check that ``arr`` is a CuPy array of ``dtype`` (and ``ndim``), contiguous in ``order`` ("C" or "F"), - on the current device (never copies).""" +def _device_array(name, arr, dtype, ndim=None): + """Check that ``arr`` is a C-contiguous CuPy array of ``dtype`` (and ``ndim``) on the current device (never copies).""" import cupy as cp if not isinstance(arr, cp.ndarray): raise TypeError(f"{name} must be a CuPy array, got {type(arr)}") if arr.device.id != cp.cuda.runtime.getDevice(): raise ValueError(f"{name} must be on the current CUDA device") - contiguous = arr.flags.f_contiguous if order == "F" else arr.flags.c_contiguous - if arr.dtype != dtype or not contiguous: - raise TypeError(f"{name} must be a {order}-contiguous array of dtype {np.dtype(dtype)}") + if arr.dtype != dtype or not arr.flags.c_contiguous: + raise TypeError(f"{name} must be a C-contiguous array of dtype {np.dtype(dtype)}") if ndim is not None and arr.ndim != ndim: raise TypeError(f"{name} must be a {ndim}-dimensional array, got {arr.ndim} dimensions") return arr @@ -68,13 +64,13 @@ def __init__( bc_type, pole: bool, ): - # column-major like the marker buffer (struphy.pic.marker_layout); Array2D carries the strides - self.markers = _device_array("markers", markers, np.float64, ndim=2, order=MARKER_ORDER) + # shape (n_cols, n_rows), one column per marker (struphy.pic.marker_layout) + self.markers = _device_array("markers", markers, np.float64, ndim=2) self.valid_mks = _device_array("valid_mks", valid_mks, np.bool_) self.Np = Np self.vdim = vdim self.weight_idx = weight_idx - self.n_markers = markers.shape[0] + self.n_markers = markers.shape[1] self.first_diagnostics_idx = first_diagnostics_idx self.first_pusher_idx = first_pusher_idx self.first_shift_idx = first_shift_idx diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index d697fe708..25b9cf7b7 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -11,7 +11,7 @@ class MarkerArguments: Paramaters ---------- markers : array[float] - Markers array, column-major (see :mod:`struphy.pic.marker_layout`). + Markers array of shape (n_cols, n_rows), one column per marker (see :mod:`struphy.pic.marker_layout`). Np : int Total number of particles. @@ -53,7 +53,7 @@ class MarkerArguments: def __init__( self, - markers: "float[:, :](order=F)", # noqa: F821 + markers: "float[:, :]", valid_mks: "bool[:]", Np: int, vdim: int, @@ -72,7 +72,7 @@ def __init__( self.Np = Np self.vdim = vdim self.weight_idx = weight_idx - self.n_markers = markers.shape[0] + self.n_markers = markers.shape[1] # useful indices self.first_diagnostics_idx = first_diagnostics_idx diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_D/cc_lin_mhd_5d_D_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_D/cc_lin_mhd_5d_D_kernels.py index 1b53798a4..fdb65fe58 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_D/cc_lin_mhd_5d_D_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_D/cc_lin_mhd_5d_D_kernels.py @@ -85,7 +85,7 @@ def cc_lin_mhd_5d_D( tmp2 = empty((3, 3), dtype=float) # get local number of markers - n_markers_loc = shape(markers)[0] + n_markers_loc = shape(markers)[1] # fmt: off #$ omp parallel firstprivate(b_prod) private(b, b_para, b_star, b_star_para, curl_norm_b, density_const, det_df, df_inv, df_inv_t, dfm, eta1, eta2, eta3, filling_m12, filling_m13, filling_m23, g_inv, ip, norm_b1, span1, span2, span3, thread_derham, tmp1, tmp2, v, weight) @@ -99,15 +99,15 @@ def cc_lin_mhd_5d_D( # fmt: on for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] - v = markers[ip, 3] + v = markers[3, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, thread_derham) @@ -145,7 +145,7 @@ def cc_lin_mhd_5d_D( density_const = 1 - b_para / b_star_para # marker weight - weight = markers[ip, 5] + weight = markers[5, ip] if basis_u == 0: # filling functions diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_M/cc_lin_mhd_5d_M_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_M/cc_lin_mhd_5d_M_kernels.py index 6ba8aed0f..a55a23941 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_M/cc_lin_mhd_5d_M_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_M/cc_lin_mhd_5d_M_kernels.py @@ -65,7 +65,7 @@ def cc_lin_mhd_5d_M( filling_v = empty(3, dtype=float) # get number of markers - n_markers_loc = shape(markers)[0] + n_markers_loc = shape(markers)[1] # fmt: off #$ omp parallel private(det_df, dfm, eta1, eta2, eta3, filling_v, ip, mu, norm_b1, span1, span2, span3, thread_derham, weight) @@ -80,17 +80,17 @@ def cc_lin_mhd_5d_M( # fmt: on for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # marker weight and velocity - weight = markers[ip, 5] - mu = markers[ip, mu_idx] + weight = markers[5, ip] + mu = markers[mu_idx, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, thread_derham) diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_curlb/cc_lin_mhd_5d_curlb_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_curlb/cc_lin_mhd_5d_curlb_kernels.py index e3b3c25c4..a963a8b55 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_curlb/cc_lin_mhd_5d_curlb_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_curlb/cc_lin_mhd_5d_curlb_kernels.py @@ -99,21 +99,21 @@ def cc_lin_mhd_5d_curlb( tmp_v = empty(3, dtype=float) # get number of markers - n_markers_loc = shape(markers)[0] + n_markers_loc = shape(markers)[1] for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # marker weight and velocity - weight = markers[ip, 5] - v = markers[ip, 3] + weight = markers[5, ip] + v = markers[3, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py index e50a5e7a4..140aa7954 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py @@ -126,22 +126,22 @@ def cc_lin_mhd_5d_gradB( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # if particle is refilled - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # marker weight and velocity - weight = markers[ip, 5] - v = markers[ip, 3] - mu = markers[ip, mu_idx] + weight = markers[5, ip] + v = markers[3, ip] + mu = markers[mu_idx, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py index 015b4506a..712d346fa 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py @@ -100,23 +100,23 @@ def cc_lin_mhd_5d_gradB_dg( tmp_v = empty(3, dtype=float) # get number of markers - n_markers_loc = shape(markers)[0] + n_markers_loc = shape(markers)[1] for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 + eta_mid[:] = (markers[0:3, ip] + markers[first_pusher_idx : first_pusher_idx + 3, ip]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff[:] = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_diff[:] = markers[0:3, ip] - markers[first_pusher_idx : first_pusher_idx + 3, ip] # marker weight and velocity - weight = markers[ip, 5] - v = markers[ip, 3] - mu = markers[ip, mu_idx] + weight = markers[5, ip] + v = markers[3, ip] + mu = markers[mu_idx, ip] # b-field evaluation span1, span2, span3 = get_spans(eta_mid[0], eta_mid[1], eta_mid[2], args_derham) diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg_init/cc_lin_mhd_5d_gradB_dg_init_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg_init/cc_lin_mhd_5d_gradB_dg_init_kernels.py index b35097b7b..a56fe96f5 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg_init/cc_lin_mhd_5d_gradB_dg_init_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg_init/cc_lin_mhd_5d_gradB_dg_init_kernels.py @@ -93,22 +93,22 @@ def cc_lin_mhd_5d_gradB_dg_init( tmp_v = empty(3, dtype=float) # get number of markers - n_markers_loc = shape(markers)[0] + n_markers_loc = shape(markers)[1] for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # marker weight and velocity - weight = markers[ip, 5] - v = markers[ip, 3] - mu = markers[ip, mu_idx] + weight = markers[5, ip] + v = markers[3, ip] + mu = markers[mu_idx, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_1/cc_lin_mhd_6d_1_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_1/cc_lin_mhd_6d_1_kernels.py index 41d8bdce7..807fb8553 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_1/cc_lin_mhd_6d_1_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_1/cc_lin_mhd_6d_1_kernels.py @@ -64,7 +64,7 @@ def cc_lin_mhd_6d_1( tmp2 = empty((3, 3), dtype=float) # get local number of markers - n_markers_loc = shape(markers)[0] + n_markers_loc = shape(markers)[1] # fmt: off #$ omp parallel firstprivate(b_prod) private(b, det_df, df_inv, df_inv_t, dfm, eta1, eta2, eta3, filling_m12, filling_m13, filling_m23, g_inv, ip, span1, span2, span3, thread_derham, tmp1, tmp2, weight) @@ -78,17 +78,17 @@ def cc_lin_mhd_6d_1( # fmt: on for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: + if markers[0, ip] < boundary_cut or markers[0, ip] > 1.0 - boundary_cut: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, thread_derham) @@ -124,7 +124,7 @@ def cc_lin_mhd_6d_1( det_df = linalg_kernels.det(dfm) # marker weight - weight = markers[ip, 6] + weight = markers[6, ip] if basis_u == 0: # filling functions diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_2/cc_lin_mhd_6d_2_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_2/cc_lin_mhd_6d_2_kernels.py index a1315ee4c..cf9dc02f6 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_2/cc_lin_mhd_6d_2_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_6d_2/cc_lin_mhd_6d_2_kernels.py @@ -97,7 +97,7 @@ def cc_lin_mhd_6d_2( tmp_v = empty(3, dtype=float) # get number of markers - n_markers_loc = shape(markers)[0] + n_markers_loc = shape(markers)[1] # fmt: off #$ omp parallel firstprivate(b_prod) private(b, det_df, df_inv, df_inv_t, dfm, eta1, eta2, eta3, filling_m, filling_v, g_inv, ip, span1, span2, span3, thread_derham, tmp1, tmp2, tmp_m, tmp_t, tmp_v, v, weight) @@ -111,17 +111,17 @@ def cc_lin_mhd_6d_2( # fmt: on for ip in range(n_markers_loc): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: + if markers[0, ip] < boundary_cut or markers[0, ip] > 1.0 - boundary_cut: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, thread_derham) @@ -157,8 +157,8 @@ def cc_lin_mhd_6d_2( det_df = linalg_kernels.det(dfm) # marker weight and velocity - weight = markers[ip, 6] - v[:] = markers[ip, 3:6] + weight = markers[6, ip] + v[:] = markers[3:6, ip] if basis_u == 0: # needed metric coefficients diff --git a/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu b/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu index 4d3695561..465c17eb9 100644 --- a/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu +++ b/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_cuda.cu @@ -5,9 +5,9 @@ using namespace struphy_cuda; /** * Accumulate the marker weights into a 0-form vector, as in charge_density_0form_kernels.charge_density_0form. * - * B_p = w_p, deposited with the N-splines of V0; one thread per marker row. + * B_p = w_p, deposited with the N-splines of V0; one thread per marker. * - * @param args_markers Marker buffer (n_markers x n_cols, column-major) and the weight column weight_idx. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker) and the weight column weight_idx. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param args_domain Mapping arguments (unused; part of the common accumulation signature). * @param vec 0-form stencil vector data; written with atomic additions. @@ -17,22 +17,22 @@ using namespace struphy_cuda; */ extern "C" __global__ void charge_density_0form(MarkerArgs args_markers, DerhamArgs args_derham, DomainArgs args_domain, Array3D vec) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; if (ip >= args_markers.n_markers) return; int weight_idx = args_markers.weight_idx; // only do something if particle is a "true" particle (i.e. not a hole) - if (args_markers.markers(ip, 0) == -1.) return; + if (args_markers.markers(0, ip) == -1.) return; // marker positions - double eta1 = args_markers.markers(ip, 0); - double eta2 = args_markers.markers(ip, 1); - double eta3 = args_markers.markers(ip, 2); + double eta1 = args_markers.markers(0, ip); + double eta2 = args_markers.markers(1, ip); + double eta3 = args_markers.markers(2, ip); // filling is just the weights - double filling = args_markers.markers(ip, weight_idx); + double filling = args_markers.markers(weight_idx, ip); particle_to_mat_kernels::vec_fill_b_v0(args_derham, eta1, eta2, eta3, vec, filling); } diff --git a/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_kernels.py b/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_kernels.py index ccebefc23..5186fa136 100644 --- a/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_kernels.py +++ b/src/struphy/pic/accumulation/kernels/charge_density_0form/charge_density_0form_kernels.py @@ -36,18 +36,18 @@ def charge_density_0form( # fmt: off #$ omp for reduction ( + :vec) # fmt: on - for ip in range(shape(markers)[0]): + for ip in range(shape(markers)[1]): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # filling is just the weights - filling = markers[ip, weight_idx] + filling = markers[weight_idx, ip] particle_to_mat_kernels.vec_fill_b_v0( thread_derham, diff --git a/src/struphy/pic/accumulation/kernels/div_u_weak_1form/div_u_weak_1form_kernels.py b/src/struphy/pic/accumulation/kernels/div_u_weak_1form/div_u_weak_1form_kernels.py index e6d9c866a..8430e0c99 100644 --- a/src/struphy/pic/accumulation/kernels/div_u_weak_1form/div_u_weak_1form_kernels.py +++ b/src/struphy/pic/accumulation/kernels/div_u_weak_1form/div_u_weak_1form_kernels.py @@ -42,22 +42,22 @@ def div_u_weak_1form( # fmt: off #$ omp for reduction(+: vec1, vec2, vec3) # fmt: on - for ip in range(shape(markers)[0]): + for ip in range(shape(markers)[1]): # only do something if particle is a "true" particle (i.e. not a hole) if not valid_mks[ip]: continue # marker positions and velocites - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v1 = markers[ip, 3] - v2 = markers[ip, 4] - v3 = markers[ip, 5] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v1 = markers[3, ip] + v2 = markers[4, ip] + v3 = markers[5, ip] # weight and density - weight = markers[ip, weight_idx] - density = markers[ip, density_idx] + weight = markers[weight_idx, ip] + density = markers[density_idx, ip] # filling is just the weights fill1 = weight * v1 / density diff --git a/src/struphy/pic/accumulation/kernels/gc_density_0form/gc_density_0form_kernels.py b/src/struphy/pic/accumulation/kernels/gc_density_0form/gc_density_0form_kernels.py index 4f1b644cf..c88f15fb4 100644 --- a/src/struphy/pic/accumulation/kernels/gc_density_0form/gc_density_0form_kernels.py +++ b/src/struphy/pic/accumulation/kernels/gc_density_0form/gc_density_0form_kernels.py @@ -34,18 +34,18 @@ def gc_density_0form( # fmt: off #$ omp for reduction ( + :vec) # fmt: on - for ip in range(shape(markers)[0]): + for ip in range(shape(markers)[1]): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # filling = w_p/N - filling = markers[ip, weight_idx] + filling = markers[weight_idx, ip] particle_to_mat_kernels.vec_fill_b_v0( thread_derham, diff --git a/src/struphy/pic/accumulation/kernels/gc_mag_density_0form/gc_mag_density_0form_kernels.py b/src/struphy/pic/accumulation/kernels/gc_mag_density_0form/gc_mag_density_0form_kernels.py index ba59681c2..37247b347 100644 --- a/src/struphy/pic/accumulation/kernels/gc_mag_density_0form/gc_mag_density_0form_kernels.py +++ b/src/struphy/pic/accumulation/kernels/gc_mag_density_0form/gc_mag_density_0form_kernels.py @@ -35,19 +35,19 @@ def gc_mag_density_0form( # fmt: off #$ omp for reduction ( + :vec) # fmt: on - for ip in range(shape(markers)[0]): + for ip in range(shape(markers)[1]): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # marker weight and magnetic moment - weight = markers[ip, 5] - mu = markers[ip, mu_idx] + weight = markers[5, ip] + mu = markers[mu_idx, ip] # filling =mu*w_p/N filling = mu * weight * scale diff --git a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu index 407c211a8..1e4496ccd 100644 --- a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu +++ b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu @@ -11,9 +11,9 @@ using namespace struphy_cuda; * Filling functions (the factors alpha^2 kappa^2 / v_th^2 and alpha^2 kappa are applied by the propagator): * A_p^{mu, nu} = f_0(eta_p, v_p) / (N s_0) [DF^{-1}(eta_p) v_p]_mu [DF^{-1}(eta_p) v_p]_nu, * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; - * one thread per marker row. + * one thread per marker. * - * @param args_markers Marker buffer (n_markers x n_cols, column-major) and the total marker number Np; the weight is in + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker) and the total marker number Np; the weight is in * column 6 and s_0 in column 7, as in pyccel. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param args_domain Mapping arguments; every spline and analytic mapping. @@ -27,7 +27,7 @@ using namespace struphy_cuda; * @param vec1 First component of the V1 vector data; written atomically. * @param vec2 Second component of the V1 vector data; written atomically. * @param vec3 Third component of the V1 vector data; written atomically. - * @param f0_values Value of f0 for each marker row (n_markers entries). + * @param f0_values Value of f0 for each marker (n_markers entries). * * Holes (markers[ip, 0] == -1) and boundary particles (markers[ip, -1] == -2) are skipped. The pyccel loop runs over * all rows (shape(markers)[0]), which is n_markers here. Many threads add to the same entries, so the summation @@ -38,7 +38,7 @@ extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamA Array6D mat22, Array6D mat23, Array6D mat33, Array3D vec1, Array3D vec2, Array3D vec3, const double* f0_values) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; // get number of markers @@ -49,7 +49,7 @@ extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamA int n_markers_tot = args_markers.Np; // only do something if particle is a "true" particle (i.e. not a hole) - if (args_markers.markers(ip, 0) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(0, ip) == -1. || args_markers.markers(args_markers.markers.shape[0] - 1, ip) == -2.) return; // allocate for metric coeffs (3 x 3 matrices as nine row-major entries) @@ -59,14 +59,14 @@ extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamA double v[3], df_inv_v[3], filling_m[9], filling_v[3]; // marker positions - double eta1 = args_markers.markers(ip, 0); - double eta2 = args_markers.markers(ip, 1); - double eta3 = args_markers.markers(ip, 2); + double eta1 = args_markers.markers(0, ip); + double eta2 = args_markers.markers(1, ip); + double eta3 = args_markers.markers(2, ip); // get velocity - v[0] = args_markers.markers(ip, 3); - v[1] = args_markers.markers(ip, 4); - v[2] = args_markers.markers(ip, 5); + v[0] = args_markers.markers(3, ip); + v[1] = args_markers.markers(4, ip); + v[2] = args_markers.markers(5, ip); // evaluate Jacobian, result in dfm evaluation_kernels::df(eta1, eta2, eta3, args_domain, dfm); @@ -79,11 +79,11 @@ extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamA // filling_m = alpha^2 * kappa^2 * f0 / (N * s_0 * v_th^2) * (DF^{-1} v_p)_mu * (DF^{-1} v_p)_nu linalg_kernels::outer(df_inv_v, df_inv_v, filling_m); - double factor = f0_values[ip] / (args_markers.markers(ip, 7) * n_markers_tot); + double factor = f0_values[ip] / (args_markers.markers(7, ip) * n_markers_tot); for (int k = 0; k < 9; ++k) filling_m[k] *= factor; // filling_v = alpha^2 * kappa * w_p * DL^{-1} * v_p - for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_v[k]; + for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(6, ip) * df_inv_v[k]; // call the appropriate matvec filler particle_to_mat_kernels::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, diff --git a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_kernels.py b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_kernels.py index 10c3b8363..3b078af6b 100644 --- a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_kernels.py +++ b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_kernels.py @@ -58,7 +58,7 @@ def linear_vlasov_ampere( filling_v = empty(3, dtype=float) # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] # total number of markers (weights are w_p = delta f_p / (N * s_0)) n_markers_tot = args_markers.Np @@ -75,18 +75,18 @@ def linear_vlasov_ampere( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0 or markers[ip, -1] == -2.0: + if markers[0, ip] == -1.0 or markers[-1, ip] == -2.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # get velocity - v[0] = markers[ip, 3] - v[1] = markers[ip, 4] - v[2] = markers[ip, 5] + v[0] = markers[3, ip] + v[1] = markers[4, ip] + v[2] = markers[5, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -105,10 +105,10 @@ def linear_vlasov_ampere( # filling_m = alpha^2 * kappa^2 * f0 / (N * s_0 * v_th^2) * (DF^{-1} v_p)_mu * (DF^{-1} v_p)_nu linalg_kernels.outer(df_inv_v, df_inv_v, filling_m) - filling_m[:, :] *= f0_values[ip] / (markers[ip, 7] * n_markers_tot) + filling_m[:, :] *= f0_values[ip] / (markers[7, ip] * n_markers_tot) # filling_v = alpha^2 * kappa * w_p * DL^{-1} * v_p - filling_v[:] = markers[ip, 6] * df_inv_v + filling_v[:] = markers[6, ip] * df_inv_v # call the appropriate matvec filler particle_to_mat_kernels.m_v_fill_b_v1_symm( diff --git a/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d/pc_lin_mhd_6d_kernels.py b/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d/pc_lin_mhd_6d_kernels.py index 18abe86bd..f128ded90 100644 --- a/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d/pc_lin_mhd_6d_kernels.py +++ b/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d/pc_lin_mhd_6d_kernels.py @@ -98,21 +98,21 @@ def pc_lin_mhd_6d( tmp_v = empty(3, dtype=float) # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # marker weight and velocity - weight = markers[ip, 6] - v[:] = markers[ip, 3:6] + weight = markers[6, ip] + v[:] = markers[3:6, ip] # evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) diff --git a/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d_full/pc_lin_mhd_6d_full_kernels.py b/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d_full/pc_lin_mhd_6d_full_kernels.py index 29ef0ae03..05d0eb683 100644 --- a/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d_full/pc_lin_mhd_6d_full_kernels.py +++ b/src/struphy/pic/accumulation/kernels/pc_lin_mhd_6d_full/pc_lin_mhd_6d_full_kernels.py @@ -99,17 +99,17 @@ def pc_lin_mhd_6d_full( tmp_v = empty(3, dtype=float) # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # b-field evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -129,12 +129,12 @@ def pc_lin_mhd_6d_full( linalg_kernels.transpose(df_inv, df_inv_t) # filling functions - v[:] = markers[ip, 3:6] + v[:] = markers[3:6, ip] linalg_kernels.matrix_matrix(df_inv, df_inv_t, tmp1) linalg_kernels.matrix_vector(df_inv, v, tmp_v) - weight = markers[ip, 6] + weight = markers[6, ip] filling_m[:, :] = weight * tmp1 * ep_scale filling_v[:] = weight * tmp_v * ep_scale diff --git a/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu index ee2259a3f..a519f9bec 100644 --- a/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu +++ b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu @@ -11,9 +11,9 @@ using namespace struphy_cuda; * Filling functions: * A_p^{mu, nu} = w_p G^{-1}_{mu, nu}(eta_p), G^{-1} = DF^{-1} DF^{-T}, * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; - * one thread per marker row. + * one thread per marker. * - * @param args_markers Marker buffer (n_markers x n_cols, column-major); the weight is in column 6, as in pyccel. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker); the weight is in column 6, as in pyccel. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any @@ -35,7 +35,7 @@ extern "C" __global__ void vlasov_maxwell(MarkerArgs args_markers, DerhamArgs ar Array6D mat11, Array6D mat12, Array6D mat13, Array6D mat22, Array6D mat23, Array6D mat33, Array3D vec1, Array3D vec2, Array3D vec3) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; if (ip >= args_markers.n_markers) return; @@ -46,20 +46,20 @@ extern "C" __global__ void vlasov_maxwell(MarkerArgs args_markers, DerhamArgs ar double v[3], df_inv_times_v[3], filling_m[9], filling_v[3]; // only do something if particle is a "true" particle (i.e. not a hole) - if (args_markers.markers(ip, 0) == -1.) return; + if (args_markers.markers(0, ip) == -1.) return; // marker positions - double eta1 = args_markers.markers(ip, 0); - double eta2 = args_markers.markers(ip, 1); - double eta3 = args_markers.markers(ip, 2); + double eta1 = args_markers.markers(0, ip); + double eta2 = args_markers.markers(1, ip); + double eta3 = args_markers.markers(2, ip); // evaluate Jacobian, result in dfm evaluation_kernels::df(eta1, eta2, eta3, args_domain, dfm); // compute shifted and stretched velocity - v[0] = args_markers.markers(ip, 3); - v[1] = args_markers.markers(ip, 4); - v[2] = args_markers.markers(ip, 5); + v[0] = args_markers.markers(3, ip); + v[1] = args_markers.markers(4, ip); + v[2] = args_markers.markers(5, ip); // filling functions linalg_kernels::matrix_inv(dfm, df_inv); @@ -68,10 +68,10 @@ extern "C" __global__ void vlasov_maxwell(MarkerArgs args_markers, DerhamArgs ar linalg_kernels::matrix_vector(df_inv, v, df_inv_times_v); // filling_m = w_p * DF^{-1} * DF^{-T} - for (int k = 0; k < 9; ++k) filling_m[k] = args_markers.markers(ip, 6) * g_inv[k]; + for (int k = 0; k < 9; ++k) filling_m[k] = args_markers.markers(6, ip) * g_inv[k]; // filling_v = w_p * DF^{-1} * \V - for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_times_v[k]; + for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(6, ip) * df_inv_times_v[k]; // call the appropriate matvec filler particle_to_mat_kernels::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, diff --git a/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_kernels.py b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_kernels.py index 7dcf59ad7..762aea429 100644 --- a/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_kernels.py +++ b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_kernels.py @@ -66,15 +66,15 @@ def vlasov_maxwell( # fmt: off #$ omp for reduction ( + : mat11, mat12, mat13, mat22, mat23, mat33, vec1, vec2, vec3) # fmt: on - for ip in range(shape(markers)[0]): + for ip in range(shape(markers)[1]): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -86,9 +86,9 @@ def vlasov_maxwell( ) # compute shifted and stretched velocity - v[0] = markers[ip, 3] - v[1] = markers[ip, 4] - v[2] = markers[ip, 5] + v[0] = markers[3, ip] + v[1] = markers[4, ip] + v[2] = markers[5, ip] # filling functions linalg_kernels.matrix_inv(dfm, df_inv) @@ -97,10 +97,10 @@ def vlasov_maxwell( linalg_kernels.matrix_vector(df_inv, v, df_inv_times_v) # filling_m = w_p * DF^{-1} * DF^{-T} - filling_m[:, :] = markers[ip, 6] * g_inv + filling_m[:, :] = markers[6, ip] * g_inv # filling_v = w_p * DF^{-1} * \V - filling_v[:] = markers[ip, 6] * df_inv_times_v + filling_v[:] = markers[6, ip] * df_inv_times_v # call the appropriate matvec filler particle_to_mat_kernels.m_v_fill_b_v1_symm( diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index 560ce65db..67642b634 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -208,7 +208,7 @@ def __call__(self, *optional_args): Parameters ---------- particles : Particles - Particles object holding the markers information in format particles.markers.shape == (n_markers, :). + Particles object holding the markers information in format particles.markers.shape == (:, n_markers). optional_args : any Additional arguments to be passed to the accumulator kernel, besides the mandatory arguments @@ -237,6 +237,7 @@ def _accumulate(self, *optional_args): self.args_domain, *self._args_data, *optional_args, + n_threads=self.particles.args_markers.n_markers, # CUDA: one thread per marker ) # apply filter @@ -569,6 +570,7 @@ def _accumulate(self, *optional_args): self.args_domain, *self._args_data, *optional_args, + n_threads=self.particles.args_markers.n_markers, # CUDA: one thread per marker ) # apply filter diff --git a/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py b/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py index 07e8df407..f55431879 100644 --- a/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/canonical_kinetic_particles/canonical_kinetic_particles_kernels.py @@ -13,7 +13,7 @@ @stack_array("dfm", "dfinv", "dfinv_t", "v", "a_form", "dfta_form") def canonical_kinetic_particles( res: "float[:]", - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", args_domain: "DomainArguments", a1_1: "float[:,:,:]", @@ -41,7 +41,7 @@ def canonical_kinetic_particles( v = empty(3, dtype=float) # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] # fmt: off #$ omp parallel private(a_form, dfinv, dfinv_t, dfm, dfta_form, eta1, eta2, eta3, ip, span1, span2, span3, thread_derham, v, w) # fmt: on @@ -54,14 +54,14 @@ def canonical_kinetic_particles( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v[:] = markers[ip, 3:6] - w = markers[ip, 6] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v[:] = markers[3:6, ip] + w = markers[6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( eta1, diff --git a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py index 23f3f9a6f..52048269d 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_5d/eval_canonical_toroidal_moment_5d_kernels.py @@ -1,4 +1,4 @@ -"""Evaluate canonical toroidal momentum of each particles and assign it into markers[ip,idx_can_momentum].""" +"""Evaluate canonical toroidal momentum of each particles and assign it into markers[idx_can_momentum, ip].""" from numpy import shape, sign, sqrt @@ -8,7 +8,7 @@ def eval_canonical_toroidal_moment_5d( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", first_diagnostics_idx: int, mu_idx: int, @@ -19,25 +19,25 @@ def eval_canonical_toroidal_moment_5d( absB: "float[:,:,:]", ): """ - Evaluate canonical toroidal momentum of each particles and assign it into markers[ip,idx_can_momentum]. + Evaluate canonical toroidal momentum of each particles and assign it into markers[idx_can_momentum, ip]. """ # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] - v_para = markers[ip, 3] - mu = markers[ip, mu_idx] - energy = markers[ip, first_diagnostics_idx] - psi = markers[ip, idx_can_momentum] + v_para = markers[3, ip] + mu = markers[mu_idx, ip] + energy = markers[first_diagnostics_idx, ip] + psi = markers[idx_can_momentum, ip] # spline evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -51,7 +51,7 @@ def eval_canonical_toroidal_moment_5d( ) # shifted canonical toroidal momentum - markers[ip, idx_can_momentum] = psi - epsilon * B0 * R0 / abs_B * v_para + markers[idx_can_momentum, ip] = psi - epsilon * B0 * R0 / abs_B * v_para if energy - mu * B0 > 0: - markers[ip, idx_can_momentum] += epsilon * sign(v_para) * sqrt(2 * (energy - mu * B0)) * R0 + markers[idx_can_momentum, ip] += epsilon * sign(v_para) * sqrt(2 * (energy - mu * B0)) * R0 diff --git a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py index 78f5e1b98..f7f90d528 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_canonical_toroidal_moment_6d/eval_canonical_toroidal_moment_6d_kernels.py @@ -1,4 +1,4 @@ -"""Evaluate canonical toroidal momentum of each particles and assign it into markers[ip,first_diagnostics_idx+5].""" +"""Evaluate canonical toroidal momentum of each particles and assign it into markers[first_diagnostics_idx+5, ip].""" from numpy import shape, sign, sqrt @@ -8,7 +8,7 @@ def eval_canonical_toroidal_moment_6d( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", first_diagnostics_idx: int, epsilon: float, @@ -17,25 +17,25 @@ def eval_canonical_toroidal_moment_6d( absB: "float[:,:,:]", ): """ - Evaluate canonical toroidal momentum of each particles and assign it into markers[ip,first_diagnostics_idx+5]. + Evaluate canonical toroidal momentum of each particles and assign it into markers[first_diagnostics_idx+5, ip]. """ # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] - energy = markers[ip, first_diagnostics_idx + 3] - mu = markers[ip, first_diagnostics_idx + 4] - psi = markers[ip, first_diagnostics_idx + 5] - v_para = markers[ip, first_diagnostics_idx + 6] + energy = markers[first_diagnostics_idx + 3, ip] + mu = markers[first_diagnostics_idx + 4, ip] + psi = markers[first_diagnostics_idx + 5, ip] + v_para = markers[first_diagnostics_idx + 6, ip] # spline evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -49,7 +49,7 @@ def eval_canonical_toroidal_moment_6d( ) # shifted canonical toroidal momentum - markers[ip, first_diagnostics_idx + 5] = psi - epsilon * B0 * R0 / abs_B * v_para + markers[first_diagnostics_idx + 5, ip] = psi - epsilon * B0 * R0 / abs_B * v_para if energy - mu * B0 > 0: - markers[ip, first_diagnostics_idx + 5] += epsilon * sign(v_para) * sqrt(2 * (energy - mu * B0)) * R0 + markers[first_diagnostics_idx + 5, ip] += epsilon * sign(v_para) * sqrt(2 * (energy - mu * B0)) * R0 diff --git a/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py index 810e7c438..8538e7b2b 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_energy_5d/eval_energy_5d_kernels.py @@ -1,4 +1,4 @@ -"""Evaluate total energy of each particles and assign it into markers[ip,first_diagnostics_idx].""" +"""Evaluate total energy of each particles and assign it into markers[first_diagnostics_idx, ip].""" from numpy import shape @@ -8,30 +8,30 @@ def eval_energy_5d( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", first_diagnostics_idx: int, mu_idx: int, absB: "float[:,:,:]", ): """ - Evaluate total energy of each particles and assign it into markers[ip,first_diagnostics_idx]. + Evaluate total energy of each particles and assign it into markers[first_diagnostics_idx, ip]. """ # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] - v_parallel = markers[ip, 3] - mu = markers[ip, mu_idx] + v_parallel = markers[3, ip] + mu = markers[mu_idx, ip] # spline evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -45,4 +45,4 @@ def eval_energy_5d( ) # total energy - markers[ip, first_diagnostics_idx] = 1 / 2 * v_parallel**2 + mu * abs_B + markers[first_diagnostics_idx, ip] = 1 / 2 * v_parallel**2 + mu * abs_B diff --git a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py index e33183e72..615f7ea80 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py @@ -47,18 +47,18 @@ def eval_gradB_ediff( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 + eta_mid[:] = (markers[0:3, ip] + markers[first_pusher_idx : first_pusher_idx + 3, ip]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_diff = markers[0:3, ip] - markers[first_pusher_idx : first_pusher_idx + 3, ip] # marker weight and velocity - weight = markers[ip, 5] - mu = markers[ip, mu_idx] + weight = markers[5, ip] + mu = markers[mu_idx, ip] # b-field evaluation span1, span2, span3 = get_spans(eta_mid[0], eta_mid[1], eta_mid[2], args_derham) @@ -103,5 +103,5 @@ def eval_gradB_ediff( tmp = gradB + grad_PB_b - markers[ip, idx] = linalg_kernels.scalar_dot(eta_diff, tmp) - markers[ip, idx] *= mu + markers[idx, ip] = linalg_kernels.scalar_dot(eta_diff, tmp) + markers[idx, ip] *= mu diff --git a/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py index a95427436..95693e62a 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_guiding_center_from_6d/eval_guiding_center_from_6d_kernels.py @@ -12,7 +12,7 @@ @stack_array("v", "dfm", "b2", "norm_b_cart", "temp", "v_perp", "Larmor_r") def eval_guiding_center_from_6d( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, @@ -24,9 +24,9 @@ def eval_guiding_center_from_6d( ): r""" Evaluate guiding center phase space of each particles: - markers[ip, first_diagnostics_idx: first_diagnostics_idx+3] : logical guiding center positions - markers[ip, first_diagnostics_idx + 4] : magnetic moment - markers[ip, first_diagnostics_idx + 6] : parallel velocity + markers[first_diagnostics_idx: first_diagnostics_idx+3, ip] : logical guiding center positions + markers[first_diagnostics_idx + 4, ip] : magnetic moment + markers[first_diagnostics_idx + 6, ip] : parallel velocity """ v = empty(3, dtype=float) @@ -38,20 +38,20 @@ def eval_guiding_center_from_6d( Larmor_r = empty(3, dtype=float) # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - x = markers[ip, first_diagnostics_idx] - y = markers[ip, first_diagnostics_idx + 1] - z = markers[ip, first_diagnostics_idx + 2] - v[:] = markers[ip, 3:6] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + x = markers[first_diagnostics_idx, ip] + y = markers[first_diagnostics_idx + 1, ip] + z = markers[first_diagnostics_idx + 2, ip] + v[:] = markers[3:6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -104,10 +104,10 @@ def eval_guiding_center_from_6d( v_perp_square = v_perp[0] ** 2 + v_perp[1] ** 2 + v_perp[2] ** 2 # parallel velocity - markers[ip, first_diagnostics_idx + 6] = v_parallel + markers[first_diagnostics_idx + 6, ip] = v_parallel # magnetic moment - markers[ip, first_diagnostics_idx + 4] = 1 / 2 * v_perp_square / abs_B + markers[first_diagnostics_idx + 4, ip] = 1 / 2 * v_perp_square / abs_B # calculate Larmor radius vector linalg_kernels.cross(norm_b_cart, v_perp, Larmor_r) @@ -115,6 +115,6 @@ def eval_guiding_center_from_6d( Larmor_r *= epsilon # calculate cartesian guiding center positions - markers[ip, first_diagnostics_idx] = x - Larmor_r[0] - markers[ip, first_diagnostics_idx + 1] = y - Larmor_r[1] - markers[ip, first_diagnostics_idx + 2] = z - Larmor_r[2] + markers[first_diagnostics_idx, ip] = x - Larmor_r[0] + markers[first_diagnostics_idx + 1, ip] = y - Larmor_r[1] + markers[first_diagnostics_idx + 2, ip] = z - Larmor_r[2] diff --git a/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py index 61a973d52..e9c578127 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_magnetic_background_energy/eval_magnetic_background_energy_kernels.py @@ -10,7 +10,7 @@ @stack_array("dfm", "norm_b1", "b") def eval_magnetic_background_energy( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, @@ -19,11 +19,11 @@ def eval_magnetic_background_energy( ): r""" Evaluate :math:`mu_p |B_0(\boldsymbol \eta_p)|` for each marker. - The result is stored at markers[:, first_diagnostics_idx]. + The result is stored at markers[first_diagnostics_idx, :]. """ # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] # fmt: off #$ omp parallel private(abs_B, eta1, eta2, eta3, ip, mu, span1, span2, span3, thread_derham) @@ -37,14 +37,14 @@ def eval_magnetic_background_energy( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # spline evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, thread_derham) @@ -58,7 +58,7 @@ def eval_magnetic_background_energy( abs_B0, ) - markers[ip, first_diagnostics_idx] = mu * abs_B + markers[first_diagnostics_idx, ip] = mu * abs_B # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py index 43e2f4b39..93a04e296 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_magnetic_energy_PBb/eval_magnetic_energy_PBb_kernels.py @@ -11,7 +11,7 @@ @stack_array("dfm", "eta") def eval_magnetic_energy_PBb( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, @@ -21,26 +21,26 @@ def eval_magnetic_energy_PBb( ): r""" Evaluate :math:`mu_p |B(\boldsymbol \eta_p)_\parallel|` for each marker. - The result is stored at markers[:, first_diagnostics_idx]. + The result is stored at markers[first_diagnostics_idx, :]. """ eta = empty(3, dtype=float) dfm = empty((3, 3), dtype=float) # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta[:] = mod(markers[ip, 0:3], 1.0) + eta[:] = mod(markers[0:3, ip], 1.0) - weight = markers[ip, 7] - dweight = markers[ip, 5] + weight = markers[7, ip] + dweight = markers[5, ip] - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # spline evaluation span1, span2, span3 = get_spans(eta[0], eta[1], eta[2], args_derham) @@ -72,4 +72,4 @@ def eval_magnetic_energy_PBb( PBb, ) - markers[ip, first_diagnostics_idx] = mu * (abs_B + PB_b) + markers[first_diagnostics_idx, ip] = mu * (abs_B + PB_b) diff --git a/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py index 14eb8dcef..827e21e97 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_magnetic_moment_5d/eval_magnetic_moment_5d_kernels.py @@ -8,29 +8,29 @@ def eval_magnetic_moment_5d( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", first_diagnostics_idx: int, absB: "float[:,:,:]", ): """ Evaluate parallel velocity and magnetic moment of each particles - and assign it into markers[ip,first_diagnostics_idx+1]. + and assign it into markers[first_diagnostics_idx+1, ip]. """ # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] - v_perp = markers[ip, 4] + v_perp = markers[4, ip] # spline evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, args_derham) @@ -44,4 +44,4 @@ def eval_magnetic_moment_5d( ) # magnetic moment - markers[ip, first_diagnostics_idx + 1] = 1 / 2 * v_perp**2 / abs_B + markers[first_diagnostics_idx + 1, ip] = 1 / 2 * v_perp**2 / abs_B diff --git a/src/struphy/pic/marker_layout.py b/src/struphy/pic/marker_layout.py index e1736190c..d12d5ada2 100644 --- a/src/struphy/pic/marker_layout.py +++ b/src/struphy/pic/marker_layout.py @@ -1,32 +1,28 @@ """Memory layout of the marker buffer. -:attr:`Particles.markers ` has shape ``(n_rows, n_cols)``, one row per marker, -and is stored column-major (Fortran order, struct-of-arrays): every column is one contiguous block of ``n_rows`` -values. Column operations (``markers[:, j]``, ``markers[:, a:b] = ...``, a mask on one column) then touch only -the columns involved, and on the GPU neighbouring threads (markers) read neighbouring addresses. Indexing is -unchanged: ``markers[ip, j]`` is still the ``j``-th entry of marker ``ip``. +:attr:`Particles.markers ` has shape ``(n_cols, n_rows)``, one column per marker, +and is stored in C order (struct-of-arrays): every quantity (row ``j``) is one contiguous block of ``n_rows`` values. +``markers[j, ip]`` is the ``j``-th entry of marker ``ip``; ``markers[j]`` is quantity ``j`` of all markers, and +``markers[:, ip]`` all quantities of marker ``ip``. Operations on quantities (``markers[j]``, ``markers[a:b] = ...``, +a mask on one row) then touch only the rows involved, and on the GPU neighbouring threads (markers) read neighbouring +addresses. -Kernels see the buffer as follows: +The same convention, one column per point, holds for every array of markers or points passed to an evaluation +(:class:`~struphy.geometry.base.Domain`, :class:`~struphy.feec.psydac_derham.SplineFunction`, the equilibria): +``(k, n)`` with the logical coordinates in rows 0 to 2. -* pyccel kernels declare it as ``float[:, :](order=F)`` (in :class:`~struphy.kernel_arguments.pusher_args_kernels.MarkerArguments` - and in the kernels that take ``markers`` directly); pyccel checks the order at every call, so a C-ordered array - raises a ``TypeError`` instead of being read wrongly. -* CUDA kernels read it through ``Array2D`` views, which carry the strides of the array. - -Copies made by fancy or boolean row indexing (``markers[valid_mks]``), ``.copy()`` and ``xp.concatenate`` are -C-ordered. They are fine to read, or to write back into the buffer with ``markers[rows] = ...``; to pass such an -array to a kernel, convert it with :func:`as_marker_array` first. +The kernels take C-ordered arrays of this shape. Selecting markers (``markers[:, valid_mks]``) gives a Fortran-ordered +copy (NumPy keeps the selected axis innermost): fine to read or to write back with ``markers[:, idx] = ...``, but convert +it with ``xp.ascontiguousarray`` before passing it to a kernel or to MPI. See struphy-hub/struphy#707. """ import cunumpy as xp -MARKER_ORDER = "F" -"""Memory order of the marker buffer and of every marker array passed to a kernel.""" - ROW_ALIGNMENT = 32 -"""The number of rows is a multiple of this, so that every column starts on a 256-byte boundary (32 doubles).""" +"""The number of markers ``n_rows`` is a multiple of this, so that every row of the buffer starts on a 256-byte +boundary (32 doubles).""" def padded_n_rows(n_rows: int) -> int: @@ -34,32 +30,6 @@ def padded_n_rows(n_rows: int) -> int: return -(-int(n_rows) // ROW_ALIGNMENT) * ROW_ALIGNMENT -def marker_array(n_rows: int, n_cols: int): - """Zero marker array of shape ``(n_rows, n_cols)`` in :data:`MARKER_ORDER` on the active backend.""" - return xp.zeros((n_rows, n_cols), dtype=float, order=MARKER_ORDER) - - -def as_marker_array(markers): - """``markers`` in :data:`MARKER_ORDER`, for passing to a kernel; no copy if it already has that layout.""" - if has_marker_layout(markers): - return markers - return xp.array(markers, order=MARKER_ORDER, copy=True) - - -def has_marker_layout(markers) -> bool: - """Whether ``markers`` is a 2D array in :data:`MARKER_ORDER`, as the kernels expect. - - Besides the F-contiguity, the strides must not decrease along the axes, which is what pyccel checks for - ``order=F``: a single row in C order is F-contiguous for NumPy, but its strides are ``(8 * n_cols, 8)``. - """ - return markers.ndim == 2 and bool(markers.flags.f_contiguous) and markers.strides[0] <= markers.strides[1] - - -def check_marker_layout(markers, name: str = "markers"): - """Raise a ``TypeError`` if ``markers`` is not a 2D array in :data:`MARKER_ORDER`.""" - if not has_marker_layout(markers): - raise TypeError( - f"{name} must be a 2D Fortran-ordered (column-major) array, got shape {markers.shape} with " - f"C_CONTIGUOUS={markers.flags.c_contiguous}, F_CONTIGUOUS={markers.flags.f_contiguous}; " - "see struphy.pic.marker_layout" - ) +def marker_array(n_cols: int, n_rows: int): + """Zero marker array of shape ``(n_cols, n_rows)`` on the active backend.""" + return xp.zeros((n_cols, n_rows), dtype=float) diff --git a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py index 42bed58ef..5115d1edc 100644 --- a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py @@ -55,17 +55,17 @@ def bstar_2form( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) - v_k = markers[ip, 3] - v_n = markers[ip, first_pusher_idx + 3] + v_k = markers[3, ip] + v_n = markers[first_pusher_idx + 3, ip] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # spline evaluation @@ -100,4 +100,4 @@ def bstar_2form( # save for j in range(3): if output_indices[j] >= 0: - markers[ip, output_indices[j]] = b_star[j] + markers[output_indices[j], ip] = b_star[j] diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu index c0c889853..d49aa4064 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu @@ -8,11 +8,11 @@ using namespace struphy_cuda; * Evaluate B*_parallel as a 3-form, as in bstar_parallel_3form_kernels.bstar_parallel_3form. * * B^{*3}_parallel(Z_p) = sqrt(g) (B . b + epsilon v_parallel,p ((curl b_0) . b_0)(eta_p)) at the weighted average - * Z_{p,i} = alpha_i Z^{n+1,k}_{p,i} + (1 - alpha_i) Z^n_{p,i} (i = 1, ..., 4); one thread per marker row. + * Z_{p,i} = alpha_i Z^{n+1,k}_{p,i} + (1 - alpha_i) Z^n_{p,i} (i = 1, ..., 4); one thread per marker. * * @param alpha Weights of the average (at least four entries): alpha[0:3] for eta, alpha[3] for v_parallel. * @param output_indices Marker column of the result in output_indices[0]; -1 skips the output. - * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param epsilon Scaling parameter epsilon in front of v_parallel. @@ -25,7 +25,7 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long MarkerArgs args_markers, DomainArgs args_domain, DerhamArgs args_derham, double epsilon, Array3D B_dot_b_coeffs, Array3D curl_unit_b_dot_b0) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; int n_markers = args_markers.n_markers; if (ip >= n_markers) return; @@ -38,19 +38,19 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long int first_shift_idx = args_markers.first_shift_idx; // only do something if particle is a "true" particle (i.e. not a hole) - if (args_markers.markers(ip, 0) == -1.0) return; + if (args_markers.markers(0, ip) == -1.0) return; for (int j = 0; j < 3; ++j) { - eta_k[j] = args_markers.markers(ip, j) + args_markers.markers(ip, first_shift_idx + j); - eta_n[j] = args_markers.markers(ip, first_pusher_idx + j); + eta_k[j] = args_markers.markers(j, ip) + args_markers.markers(first_shift_idx + j, ip); + eta_n[j] = args_markers.markers(first_pusher_idx + j, ip); eta[j] = alpha[j] * eta_k[j] + (1.0 - alpha[j]) * eta_n[j]; // numpy.mod(eta, 1.0) (Fortran MODULO): the result has the sign of the divisor eta[j] = eta[j] - floor(eta[j]); } - double v_k = args_markers.markers(ip, 3); - double v_n = args_markers.markers(ip, first_pusher_idx + 3); + double v_k = args_markers.markers(3, ip); + double v_n = args_markers.markers(first_pusher_idx + 3, ip); double v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n; // evaluate Jacobian, result in dfm @@ -74,5 +74,5 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long b_star_parallel += B_dot_b; b_star_parallel *= det_df; - if (output_indices[0] >= 0) args_markers.markers(ip, output_indices[0]) = b_star_parallel; + if (output_indices[0] >= 0) args_markers.markers(output_indices[0], ip) = b_star_parallel; } diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py index b6356a1b0..d16897a16 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py @@ -50,17 +50,17 @@ def bstar_parallel_3form( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) - v_k = markers[ip, 3] - v_n = markers[ip, first_pusher_idx + 3] + v_k = markers[3, ip] + v_n = markers[first_pusher_idx + 3, ip] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # evaluate Jacobian, result in dfm @@ -99,4 +99,4 @@ def bstar_parallel_3form( b_star_parallel *= det_df if output_indices[0] >= 0: - markers[ip, output_indices[0]] = b_star_parallel + markers[output_indices[0], ip] = b_star_parallel diff --git a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py index d622133b3..7cbf655d4 100644 --- a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py @@ -49,20 +49,20 @@ def driftkinetic_hamiltonian( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) - v_k = markers[ip, 3] - v_n = markers[ip, first_pusher_idx + 3] + v_k = markers[3, ip] + v_n = markers[first_pusher_idx + 3, ip] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # spline evaluation span1, span2, span3 = get_spans(eta[0], eta[1], eta[2], args_derham) @@ -88,4 +88,4 @@ def driftkinetic_hamiltonian( # save if output_indices[0] >= 0: - markers[ip, output_indices[0]] = epsilon * v**2 / 2.0 + epsilon * mu * B_dot_b + phi + markers[output_indices[0], ip] = epsilon * v**2 / 2.0 + epsilon * mu * B_dot_b + phi diff --git a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py index 417d384e7..fa5ac6783 100644 --- a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py @@ -63,16 +63,16 @@ def grad_driftkinetic_hamiltonian( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # spline evaluation span1, span2, span3 = get_spans(eta[0], eta[1], eta[2], args_derham) @@ -108,4 +108,4 @@ def grad_driftkinetic_hamiltonian( # save for j in range(3): if output_indices[j] >= 0: - markers[ip, output_indices[j]] = grad_H[j] + markers[output_indices[j], ip] = grad_H[j] diff --git a/src/struphy/pic/pushing/kernels/push_bxu_H1vec/push_bxu_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_bxu_H1vec/push_bxu_H1vec_kernels.py index 1a19a25d5..4a11bcc46 100644 --- a/src/struphy/pic/pushing/kernels/push_bxu_H1vec/push_bxu_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_bxu_H1vec/push_bxu_H1vec_kernels.py @@ -70,17 +70,17 @@ def push_bxu_H1vec( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: + if markers[0, ip] < boundary_cut or markers[0, ip] > 1.0 - boundary_cut: continue # marker data - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -132,7 +132,7 @@ def push_bxu_H1vec( linalg_kernels.cross(b_cart, u_cart, e_cart) # update velocities - markers[ip, 3:6] += dt * e_cart + markers[3:6, ip] += dt * e_cart # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_bxu_Hcurl/push_bxu_Hcurl_kernels.py b/src/struphy/pic/pushing/kernels/push_bxu_Hcurl/push_bxu_Hcurl_kernels.py index e06179368..c04b46a0c 100644 --- a/src/struphy/pic/pushing/kernels/push_bxu_Hcurl/push_bxu_Hcurl_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_bxu_Hcurl/push_bxu_Hcurl_kernels.py @@ -72,17 +72,17 @@ def push_bxu_Hcurl( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: + if markers[0, ip] < boundary_cut or markers[0, ip] > 1.0 - boundary_cut: continue # marker data - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -136,7 +136,7 @@ def push_bxu_Hcurl( linalg_kernels.cross(b_cart, u_cart, e_cart) # update velocities - markers[ip, 3:6] += dt * e_cart + markers[3:6, ip] += dt * e_cart # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_bxu_Hdiv/push_bxu_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_bxu_Hdiv/push_bxu_Hdiv_kernels.py index 2c141d8c9..1d1223463 100644 --- a/src/struphy/pic/pushing/kernels/push_bxu_Hdiv/push_bxu_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_bxu_Hdiv/push_bxu_Hdiv_kernels.py @@ -70,17 +70,17 @@ def push_bxu_Hdiv( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # boundary cut - if markers[ip, 0] < boundary_cut or markers[ip, 0] > 1.0 - boundary_cut: + if markers[0, ip] < boundary_cut or markers[0, ip] > 1.0 - boundary_cut: continue # marker data - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -132,7 +132,7 @@ def push_bxu_Hdiv( linalg_kernels.cross(b_cart, u_cart, e_cart) # update velocities - markers[ip, 3:6] += dt * e_cart + markers[3:6, ip] += dt * e_cart # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_bxu_Hdiv_pauli/push_bxu_Hdiv_pauli_kernels.py b/src/struphy/pic/pushing/kernels/push_bxu_Hdiv_pauli/push_bxu_Hdiv_pauli_kernels.py index ba6e7c640..a9f8b1400 100644 --- a/src/struphy/pic/pushing/kernels/push_bxu_Hdiv_pauli/push_bxu_Hdiv_pauli_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_bxu_Hdiv_pauli/push_bxu_Hdiv_pauli_kernels.py @@ -105,13 +105,13 @@ def push_bxu_Hdiv_pauli( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker data - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -222,7 +222,7 @@ def push_bxu_Hdiv_pauli( e_cart[:] = e_cart - mu[ip] * b_grad # update velocities - markers[ip, 3:6] += dt * e_cart + markers[3:6, ip] += dt * e_cart # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py index 4a3882d55..c7211cd33 100644 --- a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py @@ -74,12 +74,12 @@ def push_deterministic_diffusion_stage( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] + e1 = markers[0, ip] + e2 = markers[1, ip] + e3 = markers[2, ip] # spline evaluation span1, span2, span3 = get_spans(e1, e2, e3, thread_derham) @@ -123,13 +123,13 @@ def push_deterministic_diffusion_stage( linalg_kernels.matrix_vector(ginv, tmp, k) # accumulation for last stage - markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k + markers[first_free_idx : first_free_idx + 3, ip] += 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] + markers[0:3, ip] = ( + markers[first_pusher_idx : first_pusher_idx + 3, ip] + dt * a[stage] * k - + last * markers[ip, first_free_idx : first_free_idx + 3] + + last * markers[first_free_idx : first_free_idx + 3, ip] ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu index 252fd6d40..caefe86d6 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu @@ -5,11 +5,11 @@ using namespace struphy_cuda; /** * One stage of an explicit Runge-Kutta position push, as in push_eta_stage_kernels.push_eta_stage. * - * Solves d eta_p / dt = DF^{-1}(eta_p) v_p with constant velocity, one thread per marker row. + * Solves d eta_p / dt = DF^{-1}(eta_p) v_p with constant velocity, one thread per marker. * * @param dt Time step. * @param stage Index of the current stage, 0 <= stage < n_stages. - * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param a Butcher coefficients a_stage (n_stages entries). * @param b Butcher weights b (n_stages entries). @@ -21,7 +21,7 @@ using namespace struphy_cuda; */ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, const double* a, const double* b, const double* c, int n_stages) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; int n_markers = args_markers.n_markers; if (ip >= n_markers) return; @@ -31,14 +31,14 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ double last = stage == n_stages - 1 ? 1. : 0.; // check if marker is a hole or a boundary particle - if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(first_pusher_idx, ip) == -1. || args_markers.markers(args_markers.markers.shape[0] - 1, ip) == -2.) return; - double e1 = args_markers.markers(ip, 0); - double e2 = args_markers.markers(ip, 1); - double e3 = args_markers.markers(ip, 2); + double e1 = args_markers.markers(0, ip); + double e2 = args_markers.markers(1, ip); + double e3 = args_markers.markers(2, ip); double v[3], dfm[9], dfinv[9], k[3]; - for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(ip, 3 + j); + for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(3 + j, ip); // evaluate Jacobian, result in dfm evaluation_kernels::df(e1, e2, e3, args_domain, dfm); @@ -51,10 +51,10 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ for (int j = 0; j < 3; ++j) { // accumulation for last stage - args_markers.markers(ip, first_free_idx + j) += dt * b[stage] * k[j]; + args_markers.markers(first_free_idx + j, ip) += dt * b[stage] * k[j]; // update positions for intermediate stages or last stage - args_markers.markers(ip, j) = args_markers.markers(ip, first_pusher_idx + j) + dt * a[stage] * k[j] + - last * args_markers.markers(ip, first_free_idx + j); + args_markers.markers(j, ip) = args_markers.markers(first_pusher_idx + j, ip) + dt * a[stage] * k[j] + + last * args_markers.markers(first_free_idx + j, ip); } // apply kinetic boundary conditions 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 3899a2d66..6cc20e660 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 @@ -56,13 +56,13 @@ def push_eta_stage( # fmt: on for ip in range(n_markers): # check if marker is a hole or a boundary particle - if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[first_pusher_idx, ip] == -1.0 or markers[-1, ip] == -2.0: continue - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + e1 = markers[0, ip] + e2 = markers[1, ip] + e3 = markers[2, ip] + v[:] = markers[3:6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py index 4f0a4ac73..fdf098388 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py @@ -89,32 +89,32 @@ def push_gc_Bstar_discrete_gradient_1st_order( for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n - v_k = markers[ip, 3] - v_n = markers[ip, first_pusher_idx + 3] + v_k = markers[3, ip] + v_n = markers[first_pusher_idx + 3, ip] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # Hamiltonian at n (from init_kernel) - H_n = markers[ip, first_free_idx] + H_n = markers[first_free_idx, ip] # Poisson matrix at n (from init_kernel) - b_star_parallel = epsilon * markers[ip, first_free_idx + 1] - b_star[:] = markers[ip, first_free_idx + 2 : first_free_idx + 5] + b_star_parallel = epsilon * markers[first_free_idx + 1, ip] + b_star[:] = markers[first_free_idx + 2 : first_free_idx + 5, ip] # Hamiltonian at (n+1, k) (from eval_kernel) - H_k = markers[ip, first_free_idx + 5] + H_k = markers[first_free_idx + 5, ip] # mid-point spline evaluation span1, span2, span3 = get_spans( @@ -174,19 +174,19 @@ def push_gc_Bstar_discrete_gradient_1st_order( k_v /= -b_star_parallel # compute values at (n+1, k+1) - markers[ip, 0:3] = eta_n + dt * k - markers[ip, 3] = v_n + dt * k_v + markers[0:3, ip] = eta_n + dt * k + markers[3, ip] = v_n + dt * k_v # residual (relative in v, absolute if v_k == 0) v_scale = abs(v_k) if v_scale == 0.0: v_scale = 1.0 - markers[ip, residual_idx] = sqrt( - (markers[ip, 0] - eta_k[0]) ** 2 - + (markers[ip, 1] - eta_k[1]) ** 2 - + (markers[ip, 2] - eta_k[2]) ** 2 - + ((markers[ip, 3] - v_k) / v_scale) ** 2, + markers[residual_idx, ip] = sqrt( + (markers[0, ip] - eta_k[0]) ** 2 + + (markers[1, ip] - eta_k[1]) ** 2 + + (markers[2, ip] - eta_k[2]) ** 2 + + ((markers[3, ip] - v_k) / v_scale) ** 2, ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py index 8f296d666..8b233e904 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py @@ -150,38 +150,38 @@ def push_gc_Bstar_discrete_gradient_1st_order_newton( for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] - eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] + eta_k[:] = markers[0:3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] + eta_k_shifted[:] = eta_k + markers[first_shift_idx : first_shift_idx + 3, ip] eta_diff[:] = eta_k_shifted - eta_n - v_k = markers[ip, 3] - v_n = markers[ip, first_pusher_idx + 3] + v_k = markers[3, ip] + v_n = markers[first_pusher_idx + 3, ip] v_diff = v_k - v_n - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # Hamiltonian at n - H_n = markers[ip, first_free_idx] + H_n = markers[first_free_idx, ip] # Poisson matrix at n - b_star_parallel = epsilon * markers[ip, first_free_idx + 1] - b_star[:] = markers[ip, first_free_idx + 2 : first_free_idx + 5] + b_star_parallel = epsilon * markers[first_free_idx + 1, ip] + b_star[:] = markers[first_free_idx + 2 : first_free_idx + 5, ip] # Hamiltonian at eta_1^(n+1, k) - H_k1 = markers[ip, first_free_idx + 5] + H_k1 = markers[first_free_idx + 5, ip] # Hamiltonian at eta_1^(n+1, k), eta_2^(n+1, k) - H_k12 = markers[ip, first_free_idx + 6] + H_k12 = markers[first_free_idx + 6, ip] # 1st comp of gradient of Hamiltonian at eta_1^(n+1, k) - grad_H_1 = markers[ip, first_free_idx + 7] + grad_H_1 = markers[first_free_idx + 7, ip] # 1st and 2nd comps of gradient of Hamiltonian at eta_1^(n+1, k), eta_2^(n+1, k) - grad_H_12[:] = markers[ip, first_free_idx + 8 : first_free_idx + 10] + grad_H_12[:] = markers[first_free_idx + 8 : first_free_idx + 10, ip] # evaluate H at (n+1, k) span1, span2, span3 = get_spans(eta_k[0], eta_k[1], eta_k[2], args_derham) @@ -326,15 +326,15 @@ def push_gc_Bstar_discrete_gradient_1st_order_newton( k_v = linalg_kernels.scalar_dot(C, func) k_v += func_v / schur_comp - markers[ip, 0:3] -= k - markers[ip, 3] -= k_v + markers[0:3, ip] -= k + markers[3, ip] -= k_v # residual (relative in v, absolute if v_k == 0) v_scale = abs(v_k) if v_scale == 0.0: v_scale = 1.0 - markers[ip, residual_idx] = sqrt(k[0] ** 2 + k[1] ** 2 + k[2] ** 2 + (k_v / v_scale) ** 2) + markers[residual_idx, ip] = sqrt(k[0] ** 2 + k[1] ** 2 + k[2] ** 2 + (k_v / v_scale) ** 2) # apply kinetic boundary conditions pusher_utilities_kernels.apply_kinetic_bc_marker(ip, args_markers, args_domain, True) diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py index 3c6b8d82e..a62707ceb 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py @@ -115,28 +115,28 @@ def push_gc_Bstar_discrete_gradient_2nd_order( for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n - v_k = markers[ip, 3] - v_n = markers[ip, first_pusher_idx + 3] + v_k = markers[3, ip] + v_n = markers[first_pusher_idx + 3, ip] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # Hamiltonian at n (from init_kernel) - H_n = markers[ip, first_free_idx] + H_n = markers[first_free_idx, ip] # Hamiltonian at (n+1, k) (from eval_kernel) - H_k = markers[ip, first_free_idx + 1] + H_k = markers[first_free_idx + 1, ip] # mid-point evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -254,19 +254,19 @@ def push_gc_Bstar_discrete_gradient_2nd_order( k_v /= -b_star_parallel # compute values at (n+1, k+1) - markers[ip, 0:3] = eta_n + dt * k - markers[ip, 3] = v_n + dt * k_v + markers[0:3, ip] = eta_n + dt * k + markers[3, ip] = v_n + dt * k_v # residual (relative in v, absolute if v_k == 0) v_scale = abs(v_k) if v_scale == 0.0: v_scale = 1.0 - markers[ip, residual_idx] = sqrt( - (markers[ip, 0] - eta_k[0]) ** 2 - + (markers[ip, 1] - eta_k[1]) ** 2 - + (markers[ip, 2] - eta_k[2]) ** 2 - + ((markers[ip, 3] - v_k) / v_scale) ** 2, + markers[residual_idx, ip] = sqrt( + (markers[0, ip] - eta_k[0]) ** 2 + + (markers[1, ip] - eta_k[1]) ** 2 + + (markers[2, ip] - eta_k[2]) ** 2 + + ((markers[3, ip] - v_k) / v_scale) ** 2, ) # apply kinetic boundary conditions 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 138e57fc6..819805431 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 @@ -93,18 +93,18 @@ def push_gc_Bstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue # if stage == 0.: # # save initial parallel velocity - # markers[ip, 14] = markers[ip, 3] + # markers[14, ip] = markers[3, ip] - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] - mu = markers[ip, mu_idx] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] + mu = markers[mu_idx, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -205,9 +205,9 @@ def push_gc_Bstar_explicit_multistage( pusher_utilities_kernels.update_position_stage(ip, k, dt, a[stage], b[stage], last, args_markers) # 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] + markers[first_free_idx + 3, ip] += dt * b[stage] * k_v + markers[3, ip] = ( + markers[first_pusher_idx + 3, ip] + dt * a[stage] * k_v + last * markers[first_free_idx + 3, ip] ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py index 3d9c1a3e2..d21f134dc 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py @@ -136,35 +136,35 @@ def push_gc_bxEstar_dg_1st_order_newton( for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] - eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] + eta_k[:] = markers[0:3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] + eta_k_shifted[:] = eta_k + markers[first_shift_idx : first_shift_idx + 3, ip] eta_diff[:] = eta_k_shifted - eta_n - v = markers[ip, 3] - mu = markers[ip, mu_idx] + v = markers[3, ip] + mu = markers[mu_idx, ip] # Hamiltonian at n - H_n = markers[ip, first_free_idx] + H_n = markers[first_free_idx, ip] # Poisson matrix at n - b_star_parallel = markers[ip, first_free_idx + 1] - unit_b1[:] = markers[ip, first_free_idx + 2 : first_free_idx + 5] + b_star_parallel = markers[first_free_idx + 1, ip] + unit_b1[:] = markers[first_free_idx + 2 : first_free_idx + 5, ip] # Hamiltonian at eta_1^(n+1, k) - H_k1 = markers[ip, first_free_idx + 5] + H_k1 = markers[first_free_idx + 5, ip] # Hamiltonian at eta_1^(n+1, k), eta_2^(n+1, k) - H_k12 = markers[ip, first_free_idx + 6] + H_k12 = markers[first_free_idx + 6, ip] # 1st comp of gradient of Hamiltonian at eta_1^(n+1, k) - grad_H_1 = markers[ip, first_free_idx + 7] + grad_H_1 = markers[first_free_idx + 7, ip] # 1st and 2nd comps of gradient of Hamiltonian at eta_1^(n+1, k), eta_2^(n+1, k) - grad_H_12[:] = markers[ip, first_free_idx + 8 : first_free_idx + 10] + grad_H_12[:] = markers[first_free_idx + 8 : first_free_idx + 10, ip] # evaluate H at (n+1, k) span1, span2, span3 = get_spans(eta_k[0], eta_k[1], eta_k[2], args_derham) @@ -282,10 +282,10 @@ def push_gc_bxEstar_dg_1st_order_newton( linalg_kernels.matrix_inv(Dfunc, Dfunc_inv) linalg_kernels.matrix_vector(Dfunc_inv, func, k) - markers[ip, 0:3] -= k + markers[0:3, ip] -= k # residual - markers[ip, residual_idx] = sqrt(k[0] ** 2 + k[1] ** 2 + k[2] ** 2) + markers[residual_idx, ip] = sqrt(k[0] ** 2 + k[1] ** 2 + k[2] ** 2) # apply kinetic boundary conditions pusher_utilities_kernels.apply_kinetic_bc_marker(ip, args_markers, args_domain, True) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py index 7c26a86dd..0a7ff66ae 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py @@ -82,27 +82,27 @@ def push_gc_bxEstar_discrete_gradient_1st_order( for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n - mu = markers[ip, mu_idx] + mu = markers[mu_idx, ip] # Hamiltonian at n (from init_kernel) - H_n = markers[ip, first_free_idx] + H_n = markers[first_free_idx, ip] # Poisson matrix at n (from init_kernel) - b_star_parallel = markers[ip, first_free_idx + 1] - unit_b1[:] = markers[ip, first_free_idx + 2 : first_free_idx + 5] + b_star_parallel = markers[first_free_idx + 1, ip] + unit_b1[:] = markers[first_free_idx + 2 : first_free_idx + 5, ip] # Hamiltonian at (n+1, k) (from eval_kernel) - H_k = markers[ip, first_free_idx + 5] + H_k = markers[first_free_idx + 5, ip] # mid-point spline evaluation span1, span2, span3 = get_spans( @@ -156,11 +156,11 @@ def push_gc_bxEstar_discrete_gradient_1st_order( k[:] = Exb / b_star_parallel # accumulation for last stage - markers[ip, 0:3] = eta_n + dt * k + markers[0:3, ip] = eta_n + dt * k # residual - markers[ip, residual_idx] = sqrt( - (markers[ip, 0] - eta_k[0]) ** 2 + (markers[ip, 1] - eta_k[1]) ** 2 + (markers[ip, 2] - eta_k[2]) ** 2, + markers[residual_idx, ip] = sqrt( + (markers[0, ip] - eta_k[0]) ** 2 + (markers[1, ip] - eta_k[1]) ** 2 + (markers[2, ip] - eta_k[2]) ** 2, ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py index bd0457365..cec53fd92 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py @@ -98,24 +98,24 @@ def push_gc_bxEstar_discrete_gradient_2nd_order( for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n - v = markers[ip, 3] - mu = markers[ip, mu_idx] + v = markers[3, ip] + mu = markers[mu_idx, ip] # Hamiltonian at n (from init_kernel) - H_n = markers[ip, first_free_idx] + H_n = markers[first_free_idx, ip] # Hamiltonian at (n+1, k) (from eval_kernel) - H_k = markers[ip, first_free_idx + 1] + H_k = markers[first_free_idx + 1, ip] # mid-point evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -212,11 +212,11 @@ def push_gc_bxEstar_discrete_gradient_2nd_order( k[:] = Exb / b_star_parallel # accumulation for last stage - markers[ip, 0:3] = eta_n + dt * k + markers[0:3, ip] = eta_n + dt * k # residual - markers[ip, residual_idx] = sqrt( - (markers[ip, 0] - eta_k[0]) ** 2 + (markers[ip, 1] - eta_k[1]) ** 2 + (markers[ip, 2] - eta_k[2]) ** 2, + markers[residual_idx, ip] = sqrt( + (markers[0, ip] - eta_k[0]) ** 2 + (markers[1, ip] - eta_k[1]) ** 2 + (markers[2, ip] - eta_k[2]) ** 2, ) # apply kinetic boundary conditions 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 07dedcdec..dcde056b4 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 @@ -78,14 +78,14 @@ def push_gc_bxEstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] - mu = markers[ip, mu_idx] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] + mu = markers[mu_idx, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J1_H1vec/push_gc_cc_J1_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J1_H1vec/push_gc_cc_J1_H1vec_kernels.py index 6185bd414..1a7a3c627 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J1_H1vec/push_gc_cc_J1_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J1_H1vec/push_gc_cc_J1_H1vec_kernels.py @@ -64,13 +64,13 @@ def push_gc_cc_J1_H1vec( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -147,4 +147,4 @@ def push_gc_cc_J1_H1vec( # curl_norm_b dot electric field temp = linalg_kernels.scalar_dot(e, curl_norm_b) - markers[ip, 3] += temp / abs_b_star_para * v * dt + markers[3, ip] += temp / abs_b_star_para * v * dt diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hcurl/push_gc_cc_J1_Hcurl_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hcurl/push_gc_cc_J1_Hcurl_kernels.py index b69d3bdb0..8a1aa2700 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hcurl/push_gc_cc_J1_Hcurl_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hcurl/push_gc_cc_J1_Hcurl_kernels.py @@ -63,13 +63,13 @@ def push_gc_cc_J1_Hcurl( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -154,4 +154,4 @@ def push_gc_cc_J1_Hcurl( # curl_norm_b dot electric field temp = linalg_kernels.scalar_dot(e, curl_norm_b) / det_df - markers[ip, 3] += temp / abs_b_star_para * v * dt + markers[3, ip] += temp / abs_b_star_para * v * dt diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hdiv/push_gc_cc_J1_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hdiv/push_gc_cc_J1_Hdiv_kernels.py index 2dc94c446..3ea8766c0 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hdiv/push_gc_cc_J1_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J1_Hdiv/push_gc_cc_J1_Hdiv_kernels.py @@ -59,13 +59,13 @@ def push_gc_cc_J1_Hdiv( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -145,4 +145,4 @@ def push_gc_cc_J1_Hdiv( # curl_norm_b dot electric field temp = linalg_kernels.scalar_dot(e, curl_norm_b) - markers[ip, 3] += temp / abs_b_star_para * v * dt + markers[3, ip] += temp / abs_b_star_para * v * dt diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py index c10e36e0a..b29f733f7 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py @@ -85,15 +85,15 @@ def push_gc_cc_J2_dg_Hdiv( for ip in range(n_markers): # check if marker is a hole - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # marker positions, mid point - eta_old[:] = markers[ip, 0:3] - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 + eta_old[:] = markers[0:3, ip] + eta_mid[:] = (markers[0:3, ip] + markers[first_pusher_idx : first_pusher_idx + 3, ip]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - v = markers[ip, 3] + v = markers[3, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -204,6 +204,6 @@ def push_gc_cc_J2_dg_Hdiv( e /= abs_b_star_para e /= det_df - markers[ip, 0:3] = markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * e - markers[ip, 0:3] *= alpha - markers[ip, 0:3] += eta_old * (1.0 - alpha) + markers[0:3, ip] = markers[first_pusher_idx : first_pusher_idx + 3, ip] - dt * e + markers[0:3, ip] *= alpha + markers[0:3, ip] += eta_old * (1.0 - alpha) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py index 388c069cc..5ac673ebc 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py @@ -72,13 +72,13 @@ def push_gc_cc_J2_dg_init_Hdiv( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -173,4 +173,4 @@ def push_gc_cc_J2_dg_init_Hdiv( e /= abs_b_star_para e /= det_df - markers[ip, 0:3] -= dt * e + markers[0:3, ip] -= dt * e diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py index d40c24fc7..24c8ca3a8 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py @@ -98,13 +98,13 @@ def push_gc_cc_J2_stage_H1vec( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -199,13 +199,13 @@ def push_gc_cc_J2_stage_H1vec( e /= abs_b_star_para # accumulation for last stage - markers[ip, first_free_idx : first_free_idx + 3] -= dt * b[stage] * e + markers[first_free_idx : first_free_idx + 3, ip] -= dt * b[stage] * e # update positions for intermediate stages or last stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] + markers[0:3, ip] = ( + markers[first_pusher_idx : first_pusher_idx + 3, ip] - dt * a[stage] * e - + last * markers[ip, first_free_idx : first_free_idx + 3] + + last * markers[first_free_idx : first_free_idx + 3, ip] ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py index c2a82a5f5..ca14ea4c2 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py @@ -94,17 +94,17 @@ def push_gc_cc_J2_stage_Hdiv( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue # if particle is refilled - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v = markers[ip, 3] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v = markers[3, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -200,13 +200,13 @@ def push_gc_cc_J2_stage_Hdiv( e /= det_df # accumulation for last stage - markers[ip, first_free_idx : first_free_idx + 3] -= dt * b[stage] * e + markers[first_free_idx : first_free_idx + 3, ip] -= dt * b[stage] * e # update positions for intermediate stages or last stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] + markers[0:3, ip] = ( + markers[first_pusher_idx : first_pusher_idx + 3, ip] - dt * a[stage] * e - + last * markers[ip, first_free_idx : first_free_idx + 3] + + last * markers[first_free_idx : first_free_idx + 3, ip] ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_pc_GXu/push_pc_GXu_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_GXu/push_pc_GXu_kernels.py index 05bde9799..a5732f72a 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_GXu/push_pc_GXu_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_GXu/push_pc_GXu_kernels.py @@ -71,13 +71,13 @@ def push_pc_GXu( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v[:] = markers[3:6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -125,4 +125,4 @@ def push_pc_GXu( linalg_kernels.matrix_vector(dfinv_t, e, e_cart) # update velocities - markers[ip, 3:6] -= dt * e_cart / 2.0 + markers[3:6, ip] -= dt * e_cart / 2.0 diff --git a/src/struphy/pic/pushing/kernels/push_pc_GXu_full/push_pc_GXu_full_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_GXu_full/push_pc_GXu_full_kernels.py index 7ba7eab1a..abc17253f 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_GXu_full/push_pc_GXu_full_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_GXu_full/push_pc_GXu_full_kernels.py @@ -71,13 +71,13 @@ def push_pc_GXu_full( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + v[:] = markers[3:6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -136,4 +136,4 @@ def push_pc_GXu_full( linalg_kernels.matrix_vector(dfinv_t, e, e_cart) # update velocities - markers[ip, 3:6] -= dt * e_cart / 2.0 + markers[3:6, ip] -= dt * e_cart / 2.0 diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py index e9fe1c0c5..0716c2dab 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py @@ -79,13 +79,13 @@ def push_pc_eta_stage_H1vec( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + e1 = markers[0, ip] + e2 = markers[1, ip] + e3 = markers[2, ip] + v[:] = markers[3:6, ip] # ----------------- stage n in Runge-Kutta method ------------------- # evaluate Jacobian, result in dfm @@ -127,13 +127,13 @@ def push_pc_eta_stage_H1vec( k[:] = k_v + u # accum k - markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k + markers[first_free_idx : first_free_idx + 3, ip] += dt * b[stage] * k # update markers for the next stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] + markers[0:3, ip] = ( + markers[first_pusher_idx : first_pusher_idx + 3, ip] + dt * k * a[stage] - + last * markers[ip, first_free_idx : first_free_idx + 3] + + last * markers[first_free_idx : first_free_idx + 3, ip] ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py index 55e02b6ff..1a8b89ec4 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py @@ -72,13 +72,13 @@ def push_pc_eta_stage_Hcurl( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + e1 = markers[0, ip] + e2 = markers[1, ip] + e3 = markers[2, ip] + v[:] = markers[3:6, ip] # ----------------- stage n in Runge-Kutta method ------------------- # evaluate Jacobian, result in dfm @@ -123,13 +123,13 @@ def push_pc_eta_stage_Hcurl( k[:] = k_v + k_u # accum k - markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k + markers[first_free_idx : first_free_idx + 3, ip] += dt * b[stage] * k # update markers for the next stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] + markers[0:3, ip] = ( + markers[first_pusher_idx : first_pusher_idx + 3, ip] + dt * k * a[stage] - + last * markers[ip, first_free_idx : first_free_idx + 3] + + last * markers[first_free_idx : first_free_idx + 3, ip] ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py index 0cae6016c..c2984395e 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py @@ -72,13 +72,13 @@ def push_pc_eta_stage_Hdiv( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + e1 = markers[0, ip] + e2 = markers[1, ip] + e3 = markers[2, ip] + v[:] = markers[3:6, ip] # ----------------- stage n in Runge-Kutta method ------------------- # evaluate Jacobian, result in dfm @@ -124,13 +124,13 @@ def push_pc_eta_stage_Hdiv( k[:] = k_v + k_u # accum k - markers[ip, first_free_idx : first_free_idx + 3] += dt * b[stage] * k + markers[first_free_idx : first_free_idx + 3, ip] += dt * b[stage] * k # update markers for the next stage - markers[ip, 0:3] = ( - markers[ip, first_pusher_idx : first_pusher_idx + 3] + markers[0:3, ip] = ( + markers[first_pusher_idx : first_pusher_idx + 3, ip] + dt * k * a[stage] - + last * markers[ip, first_free_idx : first_free_idx + 3] + + last * markers[first_free_idx : first_free_idx + 3, ip] ) # apply kinetic boundary conditions diff --git a/src/struphy/pic/pushing/kernels/push_random_diffusion_stage/push_random_diffusion_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_random_diffusion_stage/push_random_diffusion_stage_kernels.py index 8741c72d9..58b80181c 100644 --- a/src/struphy/pic/pushing/kernels/push_random_diffusion_stage/push_random_diffusion_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_random_diffusion_stage/push_random_diffusion_stage_kernels.py @@ -45,10 +45,10 @@ def push_random_diffusion_stage( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - markers[ip, 0:3] += scale * noise[ip, :] + markers[0:3, ip] += scale * noise[:, ip] # 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_v_sph_pressure/push_v_sph_pressure_kernels.py b/src/struphy/pic/pushing/kernels/push_v_sph_pressure/push_v_sph_pressure_kernels.py index d9de5beda..8066d6bad 100644 --- a/src/struphy/pic/pushing/kernels/push_v_sph_pressure/push_v_sph_pressure_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_v_sph_pressure/push_v_sph_pressure_kernels.py @@ -48,8 +48,8 @@ def push_v_sph_pressure( This kernel requires: - * The density :math:`\rho^{N,h}(\boldsymbol \eta_p)` to be pre-computed for each particle and stored at ``markers[:, first_free_idx]``) - * The coefficient :math:`w_i/\rho^{N,h}(\boldsymbol \eta_i)` to be pre-computed for each particle and stored at ``markers[:, first_free_idx + 1]``) + * The density :math:`\rho^{N,h}(\boldsymbol \eta_p)` to be pre-computed for each particle and stored at ``markers[first_free_idx, :]``) + * The coefficient :math:`w_i/\rho^{N,h}(\boldsymbol \eta_i)` to be pre-computed for each particle and stored at ``markers[first_free_idx + 1, :]``) This is accomplished by the kernel :func:`~struphy.pic.pushing.eval_kernels_sph.sph_pressure_coeffs`, which needs to be passed as an ``init_kernel`` to the :class:`~struphy.pic.pushing.pusher.Pusher`. @@ -63,7 +63,7 @@ def push_v_sph_pressure( Array containing the 27 neighbouring boxes of each box. holes : bool - 1D array of length markers.shape[0]. True if markers[i] is a hole. + 1D array of length markers.shape[1]. True if markers[:, i] is a hole. periodic1, periodic2, periodic3 : bool True if periodic in that dimension. @@ -93,7 +93,7 @@ def push_v_sph_pressure( weight_idx = args_markers.weight_idx first_free_idx = args_markers.first_free_idx valid_mks = args_markers.valid_mks - n_cols = shape(markers)[1] + n_cols = shape(markers)[0] # fmt: off #$ omp parallel firstprivate(grad_u) private(dfinv, dfinvT, eta1, eta2, eta3, grad_u_cart, ip, loc_box, n_at_eta, sum2, sum4, sum6, tmp1) @@ -103,11 +103,11 @@ def push_v_sph_pressure( if not valid_mks[ip]: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - n_at_eta = markers[ip, first_free_idx] - loc_box = int(markers[ip, n_cols - 2]) + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + n_at_eta = markers[first_free_idx, ip] + loc_box = int(markers[n_cols - 2, ip]) # first component grad_u[0] = sph_eval_kernels.box_based_kernel( @@ -252,7 +252,7 @@ def push_v_sph_pressure( linalg_kernels.matrix_vector(dfinvT, grad_u, grad_u_cart) # update velocities - markers[ip, 3:6] -= dt * (grad_u_cart - gravity) + markers[3:6, ip] -= dt * (grad_u_cart - gravity) # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_v_sph_pressure_ideal_gas/push_v_sph_pressure_ideal_gas_kernels.py b/src/struphy/pic/pushing/kernels/push_v_sph_pressure_ideal_gas/push_v_sph_pressure_ideal_gas_kernels.py index 68d68a166..ab6075404 100644 --- a/src/struphy/pic/pushing/kernels/push_v_sph_pressure_ideal_gas/push_v_sph_pressure_ideal_gas_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_v_sph_pressure_ideal_gas/push_v_sph_pressure_ideal_gas_kernels.py @@ -49,8 +49,8 @@ def push_v_sph_pressure_ideal_gas( This kernel requires: - * The density :math:`\rho^{N,h}(\boldsymbol \eta_p)` to be pre-computed for each particle and stored at ``markers[:, first_free_idx]``) - * The coefficient :math:`w_i (\rho^{N,h}(\boldsymbol \eta_i))^{\gamma - 2}` to be pre-computed for each particle and stored at ``markers[:, first_free_idx + 2]``) + * The density :math:`\rho^{N,h}(\boldsymbol \eta_p)` to be pre-computed for each particle and stored at ``markers[first_free_idx, :]``) + * The coefficient :math:`w_i (\rho^{N,h}(\boldsymbol \eta_i))^{\gamma - 2}` to be pre-computed for each particle and stored at ``markers[first_free_idx + 2, :]``) This is accomplished by the kernel :func:`~struphy.pic.pushing.eval_kernels_sph.sph_pressure_coeffs`, which needs to be passed as an ``init_kernel`` to the :class:`~struphy.pic.pushing.pusher.Pusher`. @@ -64,7 +64,7 @@ def push_v_sph_pressure_ideal_gas( Array containing the 27 neighbouring boxes of each box. holes : bool - 1D array of length markers.shape[0]. True if markers[i] is a hole. + 1D array of length markers.shape[1]. True if markers[:, i] is a hole. periodic1, periodic2, periodic3 : bool True if periodic in that dimension. @@ -94,7 +94,7 @@ def push_v_sph_pressure_ideal_gas( weight_idx = args_markers.weight_idx first_free_idx = args_markers.first_free_idx valid_mks = args_markers.valid_mks - n_cols = shape(markers)[1] + n_cols = shape(markers)[0] gamma = 5 / 3 @@ -106,11 +106,11 @@ def push_v_sph_pressure_ideal_gas( if not valid_mks[ip]: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - n_at_eta = markers[ip, first_free_idx] - loc_box = int(markers[ip, n_cols - 2]) + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + n_at_eta = markers[first_free_idx, ip] + loc_box = int(markers[n_cols - 2, ip]) # first component grad_u[0] = sph_eval_kernels.box_based_kernel( @@ -254,7 +254,7 @@ def push_v_sph_pressure_ideal_gas( linalg_kernels.matrix_vector(dfinvT, grad_u, grad_u_cart) # update velocities - markers[ip, 3:6] -= dt * (grad_u_cart - gravity) + markers[3:6, ip] -= dt * (grad_u_cart - gravity) # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_v_viscosity/push_v_viscosity_kernels.py b/src/struphy/pic/pushing/kernels/push_v_viscosity/push_v_viscosity_kernels.py index c39a55acf..5d66aa5ef 100644 --- a/src/struphy/pic/pushing/kernels/push_v_viscosity/push_v_viscosity_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_v_viscosity/push_v_viscosity_kernels.py @@ -46,14 +46,14 @@ def push_v_viscosity( * :math:`- w_i \, \sqrt g(\boldsymbol \eta_i) \, \sigma_{jm}(\boldsymbol \eta_i) \, (DF^{-1})_{km}(\boldsymbol \eta_i) / \rho^{N,h}(\boldsymbol \eta_i)` to be - pre-computed for each particle and stored at ``markers[:, first_free_idx + 3*(j+1) + k]`` + pre-computed for each particle and stored at ``markers[first_free_idx + 3*(j+1) + k, :]`` for :math:`j, k = 0, 1, 2` This is accomplished by the kernel :func:`~struphy.pic.pushing.eval_kernels_sph.sph_viscosity_tensor`, which itself requires the mean velocity coefficients :math:`w_i v_{k,i} / \rho^{N,h}(\boldsymbol \eta_i)` to be stored at - ``markers[:, first_free_idx:first_free_idx + 3]`` via + ``markers[first_free_idx:first_free_idx + 3, :]`` via :func:`~struphy.pic.pushing.eval_kernels_sph.sph_mean_velocity_coeffs`. Both kernels must be passed as ``init_kernel`` entries to the :class:`~struphy.pic.pushing.pusher.Pusher`. @@ -67,7 +67,7 @@ def push_v_viscosity( Array containing the 27 neighbouring boxes of each box. holes : bool - 1D array of length markers.shape[0]. True if markers[i] is a hole. + 1D array of length markers.shape[1]. True if markers[:, i] is a hole. periodic1, periodic2, periodic3 : bool True if periodic in that dimension. @@ -86,7 +86,7 @@ def push_v_viscosity( n_markers = args_markers.n_markers first_free_idx = args_markers.first_free_idx valid_mks = args_markers.valid_mks - n_cols = shape(markers)[1] + n_cols = shape(markers)[0] f_visc = zeros(3, dtype=float) # fmt: off @@ -97,10 +97,10 @@ def push_v_viscosity( if not valid_mks[ip]: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - loc_box = int(markers[ip, n_cols - 2]) + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + loc_box = int(markers[n_cols - 2, ip]) f_visc[:] = 0.0 for j in range(3): # row of viscosity tensor @@ -149,7 +149,7 @@ def push_v_viscosity( ) # update velocities - markers[ip, 3:6] -= dt * f_visc / detdf + markers[3:6, ip] -= dt * f_visc / detdf # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu b/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu index a5828b27b..a402de954 100644 --- a/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_cuda.cu @@ -7,11 +7,11 @@ using namespace struphy_cuda; /** * Velocity update in an electric field, as in push_v_with_efield_kernels.push_v_with_efield. * - * v_p += dt * const_factor * DF^{-T}(eta_p) e(eta_p), with e the 1-form spline field; one thread per marker row. + * v_p += dt * const_factor * DF^{-T}(eta_p) e(eta_p), with e the 1-form spline field; one thread per marker. * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, column-major), validity mask and column indices. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker), validity mask and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param e1_1 Coefficients of the first component of the electric 1-form. @@ -24,7 +24,7 @@ using namespace struphy_cuda; extern "C" __global__ void push_v_with_efield(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, DerhamArgs args_derham, Array3D e1_1, Array3D e1_2, Array3D e1_3, double const_factor) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; int n_markers = args_markers.n_markers; if (ip >= n_markers) return; @@ -32,9 +32,9 @@ extern "C" __global__ void push_v_with_efield(double dt, int stage, MarkerArgs a // only do something if particle is valid (i.e. not a hole or ghost) if (!args_markers.valid_mks[ip]) return; - double eta1 = args_markers.markers(ip, 0); - double eta2 = args_markers.markers(ip, 1); - double eta3 = args_markers.markers(ip, 2); + double eta1 = args_markers.markers(0, ip); + double eta2 = args_markers.markers(1, ip); + double eta3 = args_markers.markers(2, ip); double dfm[9], dfinv[9], dfinvt[9], e_form[3], e_cart[3]; // evaluate Jacobian, result in dfm @@ -56,5 +56,5 @@ extern "C" __global__ void push_v_with_efield(double dt, int stage, MarkerArgs a linalg_kernels::matrix_vector(dfinvt, e_form, e_cart); // update velocities - for (int j = 0; j < 3; ++j) args_markers.markers(ip, 3 + j) += dt * const_factor * e_cart[j]; + for (int j = 0; j < 3; ++j) args_markers.markers(3 + j, ip) += dt * const_factor * e_cart[j]; } diff --git a/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_kernels.py b/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_kernels.py index 429d22416..3c0d55786 100644 --- a/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_v_with_efield/push_v_with_efield_kernels.py @@ -69,9 +69,9 @@ def push_v_with_efield( if not valid_mks[ip]: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -105,7 +105,7 @@ def push_v_with_efield( linalg_kernels.matrix_vector(dfinvt, e_form, e_cart) # update velocities - markers[ip, 3:6] += dt * const_factor * e_cart + markers[3:6, ip] += dt * const_factor * e_cart # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu index c8131f4b1..8eb12c785 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu @@ -9,11 +9,11 @@ using namespace struphy_cuda; * Exact rotation of the velocity about the magnetic field, as in push_vxb_analytic_kernels.push_vxb_analytic. * * Solves d v_p / dt = v_p x DF B^2 / sqrt(g) with the field frozen at the marker position, one thread - * per marker row. + * per marker. * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param b2_1 Coefficients of the first component of the magnetic 2-form. @@ -25,7 +25,7 @@ using namespace struphy_cuda; extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, DerhamArgs args_derham, Array3D b2_1, Array3D b2_2, Array3D b2_3) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; int n_markers = args_markers.n_markers; if (ip >= n_markers) return; @@ -33,14 +33,14 @@ extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs ar int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(first_pusher_idx, ip) == -1. || args_markers.markers(args_markers.markers.shape[0] - 1, ip) == -2.) return; - double e1 = args_markers.markers(ip, 0); - double e2 = args_markers.markers(ip, 1); - double e3 = args_markers.markers(ip, 2); + double e1 = args_markers.markers(0, ip); + double e2 = args_markers.markers(1, ip); + double e3 = args_markers.markers(2, ip); double v[3], dfm[9], b_form[3], b_cart[3], b_norm[3], vperp[3], vxb_norm[3], b_normxvperp[3]; - for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(ip, 3 + j); + for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(3 + j, ip); // evaluate Jacobian, result in dfm evaluation_kernels::df(e1, e2, e3, args_domain, dfm); @@ -81,5 +81,5 @@ extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs ar // analytic rotation for (int j = 0; j < 3; ++j) - args_markers.markers(ip, 3 + j) = vpar * b_norm[j] + cos(b_abs * dt) * vperp[j] - sin(b_abs * dt) * b_normxvperp[j]; + args_markers.markers(3 + j, ip) = vpar * b_norm[j] + cos(b_abs * dt) * vperp[j] - sin(b_abs * dt) * b_normxvperp[j]; } diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py index 42c75ea3f..07fe0b7b3 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py @@ -68,13 +68,13 @@ def push_vxb_analytic( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[first_pusher_idx, ip] == -1.0 or markers[-1, ip] == -2.0: continue - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + e1 = markers[0, ip] + e2 = markers[1, ip] + e3 = markers[2, ip] + v[:] = markers[3:6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -126,7 +126,7 @@ def push_vxb_analytic( linalg_kernels.cross(b_norm, vperp, b_normxvperp) # analytic rotation - markers[ip, 3:6] = vpar * b_norm + cos(b_abs * dt) * vperp - sin(b_abs * dt) * b_normxvperp + markers[3:6, ip] = vpar * b_norm + cos(b_abs * dt) * vperp - sin(b_abs * dt) * b_normxvperp # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu index 687328fe5..51cb9d6f8 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu @@ -8,11 +8,11 @@ using namespace struphy_cuda; * Implicit (Crank-Nicolson) rotation of the velocity, as in push_vxb_implicit_kernels.push_vxb_implicit. * * Solves (I - dt/2 B_x) v^{n+1} = (I + dt/2 B_x) v^n, where B_x is the cross-product matrix of - * DF B^2 / sqrt(g) at the marker position, one thread per marker row. + * DF B^2 / sqrt(g) at the marker position, one thread per marker. * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, column-major) and column indices. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker) and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param b2_1 Coefficients of the first component of the magnetic 2-form. @@ -24,7 +24,7 @@ using namespace struphy_cuda; extern "C" __global__ void push_vxb_implicit(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, DerhamArgs args_derham, Array3D b2_1, Array3D b2_2, Array3D b2_3) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; int n_markers = args_markers.n_markers; if (ip >= n_markers) return; @@ -32,16 +32,16 @@ extern "C" __global__ void push_vxb_implicit(double dt, int stage, MarkerArgs ar int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_pusher_idx) == -1.) return; + if (args_markers.markers(first_pusher_idx, ip) == -1.) return; - double e1 = args_markers.markers(ip, 0); - double e2 = args_markers.markers(ip, 1); - double e3 = args_markers.markers(ip, 2); + double e1 = args_markers.markers(0, ip); + double e2 = args_markers.markers(1, ip); + double e3 = args_markers.markers(2, ip); double v[3], dfm[9], b_form[3], b_cart[3], vec[3], res[3]; double identity[9] = {1., 0., 0., 0., 1., 0., 0., 0., 1.}; double b_prod[9] = {0., 0., 0., 0., 0., 0., 0., 0., 0.}; double rhs[9], lhs[9], lhs_inv[9]; - for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(ip, 3 + j); + for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(3 + j, ip); // evaluate Jacobian, result in dfm evaluation_kernels::df(e1, e2, e3, args_domain, dfm); @@ -80,5 +80,5 @@ extern "C" __global__ void push_vxb_implicit(double dt, int stage, MarkerArgs ar linalg_kernels::matrix_vector(rhs, v, vec); linalg_kernels::matrix_vector(lhs_inv, vec, res); - for (int j = 0; j < 3; ++j) args_markers.markers(ip, 3 + j) = res[j]; + for (int j = 0; j < 3; ++j) args_markers.markers(3 + j, ip) = res[j]; } diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py index b30a8bc17..5327a1040 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py @@ -79,13 +79,13 @@ def push_vxb_implicit( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_pusher_idx] == -1.0: + if markers[first_pusher_idx, ip] == -1.0: continue - e1 = markers[ip, 0] - e2 = markers[ip, 1] - e3 = markers[ip, 2] - v[:] = markers[ip, 3:6] + e1 = markers[0, ip] + e2 = markers[1, ip] + e3 = markers[2, ip] + v[:] = markers[3:6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -137,7 +137,7 @@ def push_vxb_implicit( linalg_kernels.matrix_vector(rhs, v, vec) linalg_kernels.matrix_vector(lhs_inv, vec, res) - markers[ip, 3:6] = res + markers[3:6, ip] = res # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu b/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu index 1b6fd9f87..bbaf15846 100644 --- a/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_cuda.cu @@ -9,17 +9,17 @@ using namespace struphy_cuda; * push_weights_with_efield_lin_va_kernels.push_weights_with_efield_lin_va. * * w_p += kappa * dt / (2 * N * s_0 * v_th^2) * f_0 * (DF^{-1} v_p) . e(eta_p), with e the 1-form spline field - * (the sum of the old and new electric field); one thread per marker row. + * (the sum of the old and new electric field); one thread per marker. * * @param dt Time step. * @param stage Stage index (unused; part of the common pusher signature). - * @param args_markers Marker buffer (n_markers x n_cols, column-major), total marker number Np and column indices. + * @param args_markers Marker buffer (n_cols x n_markers, one column per marker), total marker number Np and column indices. * @param args_domain Mapping arguments; every spline and analytic mapping. * @param args_derham Spline degrees (1 to 8), knots and start indices. * @param e1_1 Coefficients of the first component of the electric 1-form. * @param e1_2 Coefficients of the second component of the electric 1-form. * @param e1_3 Coefficients of the third component of the electric 1-form. - * @param f0_values Value of f0 for each marker row (n_markers entries). + * @param f0_values Value of f0 for each marker (n_markers entries). * @param kappa Coupling strength between particles and fields. * @param vth Thermal velocity. * @@ -30,7 +30,7 @@ extern "C" __global__ void push_weights_with_efield_lin_va(double dt, int stage, Array3D e1_1, Array3D e1_2, Array3D e1_3, const double* f0_values, double kappa, double vth) { - // CUDA-only: the marker row of this thread replaces the pyccel loop variable + // CUDA-only: the marker of this thread replaces the pyccel loop variable int ip = blockDim.x * blockIdx.x + threadIdx.x; int n_markers = args_markers.n_markers; if (ip >= n_markers) return; @@ -38,16 +38,16 @@ extern "C" __global__ void push_weights_with_efield_lin_va(double dt, int stage, // total number of markers (weights are w_p = delta f_p / (N * s_0)) int n_markers_tot = args_markers.Np; - if (args_markers.markers(ip, 0) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) return; + if (args_markers.markers(0, ip) == -1. || args_markers.markers(args_markers.markers.shape[0] - 1, ip) == -2.) return; // position - double eta1 = args_markers.markers(ip, 0); - double eta2 = args_markers.markers(ip, 1); - double eta3 = args_markers.markers(ip, 2); + double eta1 = args_markers.markers(0, ip); + double eta2 = args_markers.markers(1, ip); + double eta3 = args_markers.markers(2, ip); double v[3], dfm[9], df_inv[9], df_inv_v[3], e_vec[3]; // get velocity - for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(ip, 3 + j); + for (int j = 0; j < 3; ++j) v[j] = args_markers.markers(3 + j, ip); // spline evaluation; CUDA-only scratch holds the spline values pyccel keeps in args_derham evaluation_kernels_3d::SplineScratch scratch; @@ -68,6 +68,6 @@ extern "C" __global__ void push_weights_with_efield_lin_va(double dt, int stage, // w_{n+1} = w_n + kappa * dt / (2 * N * s_0 * v_th^2) * f_0 * ( DF^{-1} v_p ) \cdot ( e_{n+1} + e_n ) double update = (df_inv_v[0] * e_vec[0] + df_inv_v[1] * e_vec[1] + df_inv_v[2] * e_vec[2]) * f0_values[ip] * - kappa * dt / (2 * n_markers_tot * args_markers.markers(ip, 7) * (vth * vth)); - args_markers.markers(ip, 6) += update; + kappa * dt / (2 * n_markers_tot * args_markers.markers(7, ip) * (vth * vth)); + args_markers.markers(6, ip) += update; } diff --git a/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_kernels.py b/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_kernels.py index 42fac988c..fb5c1d44b 100644 --- a/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_weights_with_efield_lin_va/push_weights_with_efield_lin_va_kernels.py @@ -66,18 +66,18 @@ def push_weights_with_efield_lin_va( #$ omp for # fmt: on for ip in range(n_markers): - if markers[ip, 0] == -1.0 or markers[ip, -1] == -2.0: + if markers[0, ip] == -1.0 or markers[-1, ip] == -2.0: continue # position - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] # get velocity - v[0] = markers[ip, 3] - v[1] = markers[ip, 4] - v[2] = markers[ip, 5] + v[0] = markers[3, ip] + v[1] = markers[4, ip] + v[2] = markers[5, ip] # spline evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, thread_derham) @@ -115,9 +115,9 @@ def push_weights_with_efield_lin_va( * f0_values[ip] * kappa * dt - / (2 * n_markers_tot * markers[ip, 7] * vth**2) + / (2 * n_markers_tot * markers[7, ip] * vth**2) ) - markers[ip, 6] += update + markers[6, ip] += update # fmt: off #$ omp end parallel diff --git a/src/struphy/pic/pushing/kernels/reflect/reflect_cuda.cu b/src/struphy/pic/pushing/kernels/reflect/reflect_cuda.cu index 341a3197b..964d2c34a 100644 --- a/src/struphy/pic/pushing/kernels/reflect/reflect_cuda.cu +++ b/src/struphy/pic/pushing/kernels/reflect/reflect_cuda.cu @@ -10,7 +10,7 @@ using namespace struphy_cuda; * * @param markers Marker array view; positions are already inside the logical cube. * @param args_domain Mapping arguments; every spline and analytic mapping. - * @param outside_inds View of the marker row indices to reflect. + * @param outside_inds View of the marker indices to reflect. * @param axis Logical velocity component to reverse (0, 1 or 2). */ extern "C" __global__ void reflect(Array2D markers, DomainArgs args_domain, diff --git a/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py b/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py index 88c790688..3624c0c29 100644 --- a/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py +++ b/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py @@ -11,7 +11,7 @@ @stack_array("dfm", "dfinv", "eta", "v", "v_logical") def reflect( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_domain: "DomainArguments", outside_inds: "int[:]", axis: "int", @@ -50,8 +50,8 @@ def reflect( v_logical = empty(3, dtype=float) for ip in outside_inds: - eta[:] = markers[ip, 0:3] - v[:] = markers[ip, 3:6] + eta[:] = markers[0:3, ip] + v[:] = markers[3:6, ip] # evaluate Jacobian, result in dfm evaluation_kernels.df( @@ -74,4 +74,4 @@ def reflect( linalg_kernels.matrix_vector(dfm, v_logical, v) # update the particle velocities - markers[ip, 3:6] = v[:] + markers[3:6, ip] = v[:] diff --git a/src/struphy/pic/pushing/kernels/sph_isotherm_kappa/sph_isotherm_kappa_kernels.py b/src/struphy/pic/pushing/kernels/sph_isotherm_kappa/sph_isotherm_kappa_kernels.py index 5849449be..3d1773236 100644 --- a/src/struphy/pic/pushing/kernels/sph_isotherm_kappa/sph_isotherm_kappa_kernels.py +++ b/src/struphy/pic/pushing/kernels/sph_isotherm_kappa/sph_isotherm_kappa_kernels.py @@ -21,8 +21,8 @@ def sph_isotherm_kappa( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue if output_indices[0] >= 0: - markers[ip, output_indices[0]] = 1.0 + markers[output_indices[0], ip] = 1.0 diff --git a/src/struphy/pic/pushing/kernels/sph_mean_velocity_coeffs/sph_mean_velocity_coeffs_kernels.py b/src/struphy/pic/pushing/kernels/sph_mean_velocity_coeffs/sph_mean_velocity_coeffs_kernels.py index 2ac0e7eeb..078cb111d 100644 --- a/src/struphy/pic/pushing/kernels/sph_mean_velocity_coeffs/sph_mean_velocity_coeffs_kernels.py +++ b/src/struphy/pic/pushing/kernels/sph_mean_velocity_coeffs/sph_mean_velocity_coeffs_kernels.py @@ -28,7 +28,7 @@ def sph_mean_velocity_coeffs( r"""For each particle, evaluate the smoothed SPH density :math:`\rho^{N,h}(\boldsymbol \eta_i)` and store the coefficient - * :math:`w_i v_{k,i} / \rho^{N,h}(\boldsymbol \eta_i)` at ``markers[:, output_indices[k]]`` for :math:`k = 0, 1, 2` + * :math:`w_i v_{k,i} / \rho^{N,h}(\boldsymbol \eta_i)` at ``markers[output_indices[k], :]`` for :math:`k = 0, 1, 2` where the smoothed SPH density is given by @@ -49,7 +49,7 @@ def sph_mean_velocity_coeffs( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - n_cols = shape(markers)[1] + n_cols = shape(markers)[0] weight_idx = args_markers.weight_idx valid_mks = args_markers.valid_mks @@ -62,10 +62,10 @@ def sph_mean_velocity_coeffs( # if holes[ip]: # continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - loc_box = int(markers[ip, n_cols - 2]) + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + loc_box = int(markers[n_cols - 2, ip]) n_at_eta = sph_eval_kernels.box_based_kernel( args_markers, eta1, @@ -84,12 +84,12 @@ def sph_mean_velocity_coeffs( h2, h3, ) - weight = markers[ip, weight_idx] - velocities = markers[ip, 3:6] + weight = markers[weight_idx, ip] + velocities = markers[3:6, ip] # save if output_indices[0] >= 0: - markers[ip, output_indices[0]] = weight / n_at_eta * velocities[0] + markers[output_indices[0], ip] = weight / n_at_eta * velocities[0] if output_indices[1] >= 0: - markers[ip, output_indices[1]] = weight / n_at_eta * velocities[1] + markers[output_indices[1], ip] = weight / n_at_eta * velocities[1] if output_indices[2] >= 0: - markers[ip, output_indices[2]] = weight / n_at_eta * velocities[2] + markers[output_indices[2], ip] = weight / n_at_eta * velocities[2] diff --git a/src/struphy/pic/pushing/kernels/sph_pressure_coeffs/sph_pressure_coeffs_kernels.py b/src/struphy/pic/pushing/kernels/sph_pressure_coeffs/sph_pressure_coeffs_kernels.py index f9ecdbc2a..14d51900c 100644 --- a/src/struphy/pic/pushing/kernels/sph_pressure_coeffs/sph_pressure_coeffs_kernels.py +++ b/src/struphy/pic/pushing/kernels/sph_pressure_coeffs/sph_pressure_coeffs_kernels.py @@ -27,9 +27,9 @@ def sph_pressure_coeffs( ): r"""For each particle, evaluate - * the density :math:`\rho^{N,h}(\boldsymbol \eta_i)` and store it at ``markers[:, output_indices[0]]``) - * the coefficient :math:`w_i/\rho^{N,h}(\boldsymbol \eta_i)` and store it at ``markers[:, output_indices[1]]``) - * the coefficient :math:`w_i (\rho^{N,h}(\boldsymbol \eta_i))^{\gamma - 2}` and store it at ``markers[:, output_indices[2]]``) + * the density :math:`\rho^{N,h}(\boldsymbol \eta_i)` and store it at ``markers[output_indices[0], :]``) + * the coefficient :math:`w_i/\rho^{N,h}(\boldsymbol \eta_i)` and store it at ``markers[output_indices[1], :]``) + * the coefficient :math:`w_i (\rho^{N,h}(\boldsymbol \eta_i))^{\gamma - 2}` and store it at ``markers[output_indices[2], :]``) where the smoothed SPH density is given by @@ -43,7 +43,7 @@ def sph_pressure_coeffs( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - n_cols = shape(markers)[1] + n_cols = shape(markers)[0] Np = args_markers.Np weight_idx = args_markers.weight_idx valid_mks = args_markers.valid_mks @@ -53,10 +53,10 @@ def sph_pressure_coeffs( if not valid_mks[ip]: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - loc_box = int(markers[ip, n_cols - 2]) + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + loc_box = int(markers[n_cols - 2, ip]) n_at_eta = sph_eval_kernels.box_based_kernel( args_markers, eta1, @@ -75,11 +75,11 @@ def sph_pressure_coeffs( h2, h3, ) - weight = markers[ip, weight_idx] + weight = markers[weight_idx, ip] # save if output_indices[0] >= 0: - markers[ip, output_indices[0]] = n_at_eta + markers[output_indices[0], ip] = n_at_eta if output_indices[1] >= 0: - markers[ip, output_indices[1]] = weight / n_at_eta + markers[output_indices[1], ip] = weight / n_at_eta if output_indices[2] >= 0: - markers[ip, output_indices[2]] = weight * n_at_eta ** (gamma - 2) + markers[output_indices[2], ip] = weight * n_at_eta ** (gamma - 2) diff --git a/src/struphy/pic/pushing/kernels/sph_viscosity_tensor/sph_viscosity_tensor_kernels.py b/src/struphy/pic/pushing/kernels/sph_viscosity_tensor/sph_viscosity_tensor_kernels.py index ab8cfdc95..e49676030 100644 --- a/src/struphy/pic/pushing/kernels/sph_viscosity_tensor/sph_viscosity_tensor_kernels.py +++ b/src/struphy/pic/pushing/kernels/sph_viscosity_tensor/sph_viscosity_tensor_kernels.py @@ -32,7 +32,7 @@ def sph_viscosity_tensor( deviatoric strain rate, and store the 9 coefficients * :math:`- w_i \, \sqrt g(\boldsymbol \eta_i) \, \sigma_{jm}(\boldsymbol \eta_i) \, (DF^{-1})_{km}(\boldsymbol \eta_i) - / \rho^{N,h}(\boldsymbol \eta_i)` at ``markers[:, output_indices[3*j + k]]`` for :math:`j, k = 0, 1, 2` + / \rho^{N,h}(\boldsymbol \eta_i)` at ``markers[output_indices[3*j + k], :]`` for :math:`j, k = 0, 1, 2` where :math:`\sqrt g = \det DF`, the smoothed SPH density is given by @@ -62,14 +62,14 @@ def sph_viscosity_tensor( (\nabla W_h)_k(\boldsymbol \eta_i - \boldsymbol \eta_l)\,. This kernel requires the coefficients of the mean velocity :math:`v_k^{N,h}` - for each particle to be pre-evaluated and stored at ``markers[:, first_free_idx:first_free_idx + 3]``, + for each particle to be pre-evaluated and stored at ``markers[first_free_idx:first_free_idx + 3, :]``, which can be achieved by the kernel :func:`~struphy.pic.pushing.eval_kernels_sph.sph_mean_velocity_coeffs`. """ # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - n_cols = shape(markers)[1] + n_cols = shape(markers)[0] weight_idx = args_markers.weight_idx first_free_idx = args_markers.first_free_idx valid_mks = args_markers.valid_mks @@ -87,10 +87,10 @@ def sph_viscosity_tensor( if not valid_mks[ip]: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] - loc_box = int(markers[ip, n_cols - 2]) + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] + loc_box = int(markers[n_cols - 2, ip]) n_at_eta = sph_eval_kernels.box_based_kernel( args_markers, eta1, @@ -109,7 +109,7 @@ def sph_viscosity_tensor( h2, h3, ) - weight = markers[ip, weight_idx] + weight = markers[weight_idx, ip] for j in range(3): for k in range(3): grad_v_at_eta[j, k] = sph_eval_kernels.box_based_kernel( @@ -162,4 +162,4 @@ def sph_viscosity_tensor( for j in range(3): for k in range(3): if output_indices[3 * j + k] >= 0: - markers[ip, output_indices[3 * j + k]] = d_piola[j, k] + markers[output_indices[3 * j + k], ip] = d_piola[j, k] diff --git a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py index 77cd4aa0a..f817053b7 100644 --- a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py +++ b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py @@ -44,11 +44,11 @@ def unit_b_1form( for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] + eta_k[:] = markers[0:3, ip] + markers[first_shift_idx : first_shift_idx + 3, ip] + eta_n[:] = markers[first_pusher_idx : first_pusher_idx + 3, ip] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) @@ -70,4 +70,4 @@ def unit_b_1form( # save for j in range(3): if output_indices[j] >= 0: - markers[ip, output_indices[j]] = unit_b1[j] + markers[output_indices[j], ip] = unit_b1[j] diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index 37a853e15..c19e15919 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -1,28 +1,28 @@ #include "cunumpy/array_view.cuh" /** - * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: + * Prepare the marker buffer for a push, as the three row-slice assignments at the start of Pusher._push: * - * markers[:, first_pusher_idx:first_shift_idx] = markers[:, :n_init] - * markers[:, first_shift_idx:residual_idx] = 0.0 - * markers[:, residual_idx:-2] = 0.0 + * markers[first_pusher_idx:first_shift_idx] = markers[:n_init] + * markers[first_shift_idx:residual_idx] = 0.0 + * markers[residual_idx:-2] = 0.0 * - * One thread per entry of the column range [first_pusher_idx, n_cols - 2) of every row, numbered column by - * column. In the column-major marker buffer (struphy.pic.marker_layout) neighbouring threads then read and write - * neighbouring rows of the same column, and the whole range is written in a single launch (instead of one per - * slice). The buffer is an Array2D view: a raw pointer argument would be made C-contiguous (copied) by the launch. + * One thread per entry of the row range [first_pusher_idx, n_cols - 2) of every marker, numbered row by row. In the + * marker buffer (n_cols x n_rows, one column per marker, struphy.pic.marker_layout) neighbouring threads then read + * and write neighbouring markers of the same row, and the whole range is written in a single launch (instead of one + * per slice). * - * @param markers Marker buffer (n_rows x n_cols, column-major), every row including holes. - * @param first_pusher_idx First column of the saved initial phase space coordinates. + * @param markers Marker buffer (n_cols x n_rows), every marker including holes. + * @param first_pusher_idx First row of the saved initial phase space coordinates. * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ extern "C" __global__ void prepare_push(Array2D markers, int first_pusher_idx, int n_init) { - long long n_rows = markers.shape[0]; - long long width = markers.shape[1] - 2 - first_pusher_idx; + long long n_rows = markers.shape[1]; + long long width = markers.shape[0] - 2 - first_pusher_idx; long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; if (t >= n_rows * width) return; long long k = t / n_rows; long long ip = t - k * n_rows; - markers(ip, first_pusher_idx + k) = k < n_init ? markers(ip, k) : 0.0; + markers(first_pusher_idx + k, ip) = k < n_init ? markers(k, ip) : 0.0; } diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 7fa7069de..0b8ae5e6a 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -82,11 +82,11 @@ class Pusher: alpha_in_kernel: float | int | tuple | list For i=0,1,2, the spline/geometry evaluations in kernel are at - alpha[i]*markers[:, i] + (1 - alpha[i])*markers[:, buffer_idx + i]. + alpha[i]*markers[i] + (1 - alpha[i])*markers[buffer_idx + i]. If float or int or then alpha = (alpha, alpha, alpha). alpha must be between 0 and 1. alpha[i]=0 means that evaluation is at the initial positions (time n), - stored at markers[:, buffer_idx + i]. + stored at markers[buffer_idx + i]. init_kernels : tuple[KernelSetup, ...] Evaluations at the initial state, executed once per push in tuple order. @@ -177,7 +177,7 @@ def __init__( # profiling region names (cached, they are looked up on every call) self._region_name = "pusher: " + self.kernel.name - self._residuals = xp.zeros(self.particles.markers.shape[0]) + self._residuals = xp.zeros(self.particles.n_rows) self._converged_loc = self._residuals == 1.0 self._not_converged_loc = self._residuals == 0.0 @@ -219,7 +219,7 @@ def _push(self, dt: float): if self._cuda: # the three assignments below in one launch - n_rows, n_cols = markers.shape + n_cols, n_rows = markers.shape _prepare_push_cuda()( markers, first_pusher_idx, @@ -231,13 +231,13 @@ def _push(self, dt: float): shift_slice = slice(first_shift_idx, residual_idx) # save initial phase space coordinates - markers[:, init_slice] = markers[:, : 3 + vdim] + markers[init_slice] = markers[: 3 + vdim] # set boundary shifts to zero - markers[:, shift_slice] = 0.0 + markers[shift_slice] = 0.0 - # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box - markers[:, residual_idx:-2] = 0.0 + # clear buffer rows starting from residual index, dont clear ID (last row) and loc_box + markers[residual_idx:-2] = 0.0 rank = self.particles.mpi_rank logger.debug(f"rank {rank}: starting {self.kernel} ...") @@ -285,6 +285,7 @@ def _push(self, dt: float): self.particles.args_markers, self._args_domain, *self._args_kernel, + n_threads=self.particles.args_markers.n_markers, # CUDA: one thread per marker ) # markers have moved @@ -300,7 +301,7 @@ def _push(self, dt: float): # compute number of non-converged particles (maxiter=1 for explicit schemes) if self.maxiter > 1: - self._residuals[:] = markers[:, residual_idx] + self._residuals[:] = markers[residual_idx] max_res = xp.max(self._residuals) if max_res < 0.0: max_res = None @@ -322,7 +323,7 @@ def _push(self, dt: float): ) # take converged markers out of the loop - markers[self._converged_loc, first_pusher_idx] = -1.0 + markers[first_pusher_idx, self._converged_loc] = -1.0 # maxiter=1 for explicit schemes if k == self.maxiter: diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index 259091753..91f2feaf5 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -6,7 +6,7 @@ namespace struphy_cuda::pusher_utilities_kernels { /** * Reflect one marker's velocity using the Pyccel boundary helper's calculation. * - * @param ip Marker row index. + * @param ip Marker index. * @param markers Marker array view. * @param args_domain Mapping arguments for the Jacobian at the reflected position. * @param axis Logical velocity component to reverse (0, 1 or 2). @@ -19,13 +19,13 @@ namespace struphy_cuda::pusher_utilities_kernels { __device__ inline void reflect_velocity(long long ip, Array2D markers, const DomainArgs& args_domain, int axis) { double dfm[9], dfinv[9], v[3], v_logical[3]; - evaluation_kernels::df(markers(ip, 0), markers(ip, 1), markers(ip, 2), args_domain, dfm); + evaluation_kernels::df(markers(0, ip), markers(1, ip), markers(2, ip), args_domain, dfm); linalg_kernels::matrix_inv(dfm, dfinv); - for (int j = 0; j < 3; ++j) v[j] = markers(ip, 3 + j); + for (int j = 0; j < 3; ++j) v[j] = markers(3 + j, ip); linalg_kernels::matrix_vector(dfinv, v, v_logical); v_logical[axis] *= -1.; linalg_kernels::matrix_vector(dfm, v_logical, v); - for (int j = 0; j < 3; ++j) markers(ip, 3 + j) = v[j]; + for (int j = 0; j < 3; ++j) markers(3 + j, ip) = v[j]; } /** @@ -51,20 +51,20 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma // 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(0, ip) < 0.) { + args_markers.markers(0, ip) = -args_markers.markers(0, ip); + double e2 = args_markers.markers(1, ip) + 0.5; + args_markers.markers(1, ip) = e2 - floor(e2); } - if (args_markers.markers(ip, 0) < 1e-10) args_markers.markers(ip, 0) = 1e-10; + if (args_markers.markers(0, ip) < 1e-10) args_markers.markers(0, ip) = 1e-10; } // Remove markers before applying periodic or reflecting boundaries. for (int axis = 0; axis < 3; ++axis) { if (bc_type[axis] == 2 && - (args_markers.markers(ip, axis) > 1. || args_markers.markers(ip, axis) < 0.)) { - int n_cols = args_markers.markers.shape[1]; - for (int j = 0; j < n_cols - 1; ++j) args_markers.markers(ip, j) = -1.; + (args_markers.markers(axis, ip) > 1. || args_markers.markers(axis, ip) < 0.)) { + int n_cols = args_markers.markers.shape[0]; + for (int j = 0; j < n_cols - 1; ++j) args_markers.markers(j, ip) = -1.; return; } } @@ -72,16 +72,16 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma // Wrap positions and update the shift exactly as in the Pyccel helper. for (int axis = 0; axis < 3; ++axis) { if (bc_type[axis] == 0) { - if (args_markers.markers(ip, axis) > 1.) { - args_markers.markers(ip, axis) -= floor(args_markers.markers(ip, axis)); - if (newton) args_markers.markers(ip, first_shift_idx + axis) += 1.; - else args_markers.markers(ip, first_shift_idx + axis) = 1.; - } else if (args_markers.markers(ip, axis) < 0.) { - args_markers.markers(ip, axis) -= floor(args_markers.markers(ip, axis)); - if (newton) args_markers.markers(ip, first_shift_idx + axis) += -1.; - else args_markers.markers(ip, first_shift_idx + axis) = -1.; + if (args_markers.markers(axis, ip) > 1.) { + args_markers.markers(axis, ip) -= floor(args_markers.markers(axis, ip)); + if (newton) args_markers.markers(first_shift_idx + axis, ip) += 1.; + else args_markers.markers(first_shift_idx + axis, ip) = 1.; + } else if (args_markers.markers(axis, ip) < 0.) { + args_markers.markers(axis, ip) -= floor(args_markers.markers(axis, ip)); + if (newton) args_markers.markers(first_shift_idx + axis, ip) += -1.; + else args_markers.markers(first_shift_idx + axis, ip) = -1.; } else if (!newton) { - args_markers.markers(ip, first_shift_idx + axis) = 0.; + args_markers.markers(first_shift_idx + axis, ip) = 0.; } } } @@ -90,19 +90,19 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma int n_reflected = 0; for (int axis = 0; axis < 3; ++axis) { if (bc_type[axis] == 1) { - if (args_markers.markers(ip, axis) > 1.) { - args_markers.markers(ip, axis) = 2. - args_markers.markers(ip, axis); + if (args_markers.markers(axis, ip) > 1.) { + args_markers.markers(axis, ip) = 2. - args_markers.markers(axis, ip); reflected[axis] = true; ++n_reflected; - } else if (args_markers.markers(ip, axis) < 0.) { - args_markers.markers(ip, axis) = -args_markers.markers(ip, axis); + } else if (args_markers.markers(axis, ip) < 0.) { + args_markers.markers(axis, ip) = -args_markers.markers(axis, ip); reflected[axis] = true; ++n_reflected; } } } if (n_reflected == 0) return; - args_markers.markers(ip, first_pusher_idx) = -1.; + args_markers.markers(first_pusher_idx, ip) = -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) { diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.py b/src/struphy/pic/pushing/pusher_utilities_kernels.py index 160123840..3d161dc9f 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.py +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.py @@ -18,7 +18,7 @@ def apply_kinetic_bc_marker( args_domain: "DomainArguments", newton: bool, ): - r"""Apply the kinetic boundary conditions to the single marker ``markers[ip]``, + r"""Apply the kinetic boundary conditions to the single marker ``markers[:, ip]``, 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. @@ -59,56 +59,56 @@ def apply_kinetic_bc_marker( # 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 + if markers[0, ip] < 0.0: + markers[0, ip] = -markers[0, ip] + markers[1, ip] = mod(markers[1, ip] + 0.5, 1.0) + if markers[0, ip] < 1e-10: + markers[0, ip] = 1e-10 # remove for axis in range(3): if bc_type[axis] == 2: - if markers[ip, axis] > 1.0 or markers[ip, axis] < 0.0: - n_cols = shape(markers)[1] - markers[ip, 0 : n_cols - 1] = -1.0 + if markers[axis, ip] > 1.0 or markers[axis, ip] < 0.0: + n_cols = shape(markers)[0] + markers[0 : n_cols - 1, ip] = -1.0 return # periodic for axis in range(3): if bc_type[axis] == 0: - if markers[ip, axis] > 1.0: - markers[ip, axis] = mod(markers[ip, axis], 1.0) + if markers[axis, ip] > 1.0: + markers[axis, ip] = mod(markers[axis, ip], 1.0) if newton: - markers[ip, first_shift_idx + axis] += 1.0 + markers[first_shift_idx + axis, ip] += 1.0 else: - markers[ip, first_shift_idx + axis] = 1.0 - elif markers[ip, axis] < 0.0: - markers[ip, axis] = mod(markers[ip, axis], 1.0) + markers[first_shift_idx + axis, ip] = 1.0 + elif markers[axis, ip] < 0.0: + markers[axis, ip] = mod(markers[axis, ip], 1.0) if newton: - markers[ip, first_shift_idx + axis] += -1.0 + markers[first_shift_idx + axis, ip] += -1.0 else: - markers[ip, first_shift_idx + axis] = -1.0 + markers[first_shift_idx + axis, ip] = -1.0 elif not newton: - markers[ip, first_shift_idx + axis] = 0.0 + markers[first_shift_idx + axis, ip] = 0.0 # reflect positions reflected = zeros(3, dtype=bool) n_reflected = 0 for axis in range(3): if bc_type[axis] == 1: - if markers[ip, axis] > 1.0: - markers[ip, axis] = 2.0 - markers[ip, axis] + if markers[axis, ip] > 1.0: + markers[axis, ip] = 2.0 - markers[axis, ip] reflected[axis] = True n_reflected += 1 - elif markers[ip, axis] < 0.0: - markers[ip, axis] = -markers[ip, axis] + elif markers[axis, ip] < 0.0: + markers[axis, ip] = -markers[axis, ip] reflected[axis] = True n_reflected += 1 if n_reflected == 0: return - markers[ip, first_pusher_idx] = -1.0 + markers[first_pusher_idx, ip] = -1.0 # only Cartesian velocities are reflected (columns 3:6 hold other data for vdim != 3) if args_markers.vdim != 3: @@ -124,12 +124,12 @@ def apply_kinetic_bc_marker( if not reflected[axis]: continue - v[:] = markers[ip, 3:6] + v[:] = markers[3:6, ip] evaluation_kernels.df( - markers[ip, 0], - markers[ip, 1], - markers[ip, 2], + markers[0, ip], + markers[1, ip], + markers[2, ip], args_domain, dfm, ) @@ -141,7 +141,7 @@ def apply_kinetic_bc_marker( v_logical[axis] *= -1.0 linalg_kernels.matrix_vector(dfm, v_logical, v) - markers[ip, 3:6] = v[:] + markers[3:6, ip] = v[:] def update_position_stage( @@ -153,9 +153,9 @@ def update_position_stage( last: float, args_markers: "MarkerArguments", ): - r"""Position update of one stage of an explicit Runge-Kutta pusher for the marker ``markers[ip]``, + 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 current position ``markers[0:3, ip]``. The stage contribution :math:`\Delta t\, b_\textrm{stage}\, \mathbf k` is accumulated at ``first_free_idx``, and the new position is @@ -207,33 +207,33 @@ def update_position_stage( if not args_markers.pole: # accumulation for last stage - markers[ip, first_free_idx : first_free_idx + 3] += dt * b_stage * k + markers[first_free_idx : first_free_idx + 3, ip] += 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] + markers[0:3, ip] = ( + markers[first_pusher_idx : first_pusher_idx + 3, ip] + dt * a_stage * k - + last * markers[ip, first_free_idx : first_free_idx + 3] + + last * markers[first_free_idx : first_free_idx + 3, ip] ) 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] + theta = 2.0 * pi * markers[1, ip] + k_xi1 = cos(theta) * k[0] - 2.0 * pi * markers[0, ip] * sin(theta) * k[1] + k_xi2 = sin(theta) * k[0] + 2.0 * pi * markers[0, ip] * 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] + markers[first_free_idx, ip] += dt * b_stage * k_xi1 + markers[first_free_idx + 1, ip] += dt * b_stage * k_xi2 + markers[first_free_idx + 2, ip] += 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] + theta = 2.0 * pi * markers[first_pusher_idx + 1, ip] + xi1 = markers[first_pusher_idx, ip] * cos(theta) + dt * a_stage * k_xi1 + last * markers[first_free_idx, ip] + xi2 = markers[first_pusher_idx, ip] * sin(theta) + dt * a_stage * k_xi2 + last * markers[first_free_idx + 1, ip] + + markers[0, ip] = sqrt(xi1**2 + xi2**2) + if markers[0, ip] < 1e-10: + markers[0, ip] = 1e-10 + markers[1, ip] = mod(arctan2(xi2, xi1) / (2.0 * pi), 1.0) + markers[2, ip] = markers[first_pusher_idx + 2, ip] + dt * a_stage * k[2] + last * markers[first_free_idx + 2, ip] diff --git a/src/struphy/pic/sampling_kernels.py b/src/struphy/pic/sampling_kernels.py index ce68d5aff..7aeed243b 100644 --- a/src/struphy/pic/sampling_kernels.py +++ b/src/struphy/pic/sampling_kernels.py @@ -45,8 +45,8 @@ def set_particles_symmetric_3d_3v(numbers: "float[:,:]", markers: "float[:,:]"): else: e[0] = 1 - e[0] - markers[i_part, 0:3] = e - markers[i_part, 3:6] = v + markers[0:3, i_part] = e + markers[3:6, i_part] = v @stack_array("e", "v") @@ -78,8 +78,8 @@ def set_particles_symmetric_2d_3v(numbers: "float[:,:]", markers: "float[:,:]"): else: e[0] = 1 - e[0] - markers[i_part, 1:3] = e - markers[i_part, 3:6] = v + markers[1:3, i_part] = e + markers[3:6, i_part] = v def tile_int_kernel( diff --git a/src/struphy/pic/sorting_kernels.py b/src/struphy/pic/sorting_kernels.py index 8e55838b0..99d4ec4f0 100644 --- a/src/struphy/pic/sorting_kernels.py +++ b/src/struphy/pic/sorting_kernels.py @@ -154,7 +154,7 @@ def get_next_index(box_nb: "float", next_index: "int[:]", cumul_next_index: "int def sort_boxed_particles( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", swap_line_1: "float[:]", swap_line_2: "float[:]", n_boxes: "int", @@ -178,7 +178,7 @@ def sort_boxed_particles( loc_i += next_index[box] # loop through the rest of the box while loc_i < cumul_next_index[box + 1]: - box_i = int(markers[loc_i, l + box_index]) + box_i = int(markers[l + box_index, loc_i]) if box_i == box: # is it in good position? if yes go to the next one # and say that this position is no longer available @@ -187,18 +187,18 @@ def sort_boxed_particles( else: # swap elements until comming back swap_i = -1 - swap_line_1[:] = markers[loc_i, :] + swap_line_1[:] = markers[:, loc_i] while swap_i != loc_i: # where should the actual particule go, put it and save data from the particle that should now move swap_i = get_next_index(swap_line_1[l + box_index], next_index, cumul_next_index) - swap_line_2[:] = markers[swap_i, :] - markers[swap_i, :] = swap_line_1[:] + swap_line_2[:] = markers[:, swap_i] + markers[:, swap_i] = swap_line_1[:] swap_line_1[:] = swap_line_2[:] loc_i += 1 def assign_box_to_each_particle( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", holes: "bool[:]", nx: "int", ny: "int", @@ -206,16 +206,16 @@ def assign_box_to_each_particle( domain_array: "float[:]", box_index: "int" = -2, ): - """Assign the right box to each particle, written into column -2 or marker array.""" - l = markers.shape[1] - for p in range(markers.shape[0]): + """Assign the right box to each particle, written into row -2 of the marker array.""" + l = markers.shape[0] + for p in range(markers.shape[1]): if holes[p]: n_box = (nx + 2) * (ny + 2) * (nz + 2) else: a = find_box( - markers[p, 0], - markers[p, 1], - markers[p, 2], + markers[0, p], + markers[1, p], + markers[2, p], nx, ny, nz, @@ -226,11 +226,11 @@ def assign_box_to_each_particle( else: n_box = a b = float(n_box) - markers[p, l + box_index] = b + markers[l + box_index, p] = b def assign_particles_to_boxes( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", holes: "bool[:]", boxes: "int[:,:]", next_index: "int[:]", @@ -239,11 +239,11 @@ def assign_particles_to_boxes( """Write particle indices into box array.""" boxes[:, :] = -1 next_index[:] = 0 - l = markers.shape[1] - for p in range(markers.shape[0]): + l = markers.shape[0] + for p in range(markers.shape[1]): if holes[p]: continue else: - a = int(markers[p, l + box_index]) + a = int(markers[l + box_index, p]) boxes[a, next_index[a]] = p next_index[a] += 1 diff --git a/src/struphy/pic/sph_eval_kernels.py b/src/struphy/pic/sph_eval_kernels.py index 09c0f1fd1..e3e97545f 100644 --- a/src/struphy/pic/sph_eval_kernels.py +++ b/src/struphy/pic/sph_eval_kernels.py @@ -1,4 +1,6 @@ # do not remove; needed to identify dependencies +from numpy import shape + import struphy.kernel_arguments.pusher_args_kernels as pusher_args_kernels import struphy.pic.sorting_kernels as sorting_kernels import struphy.pic.sph_smoothing_kernels as sph_smoothing_kernels @@ -57,7 +59,7 @@ def naive_evaluation_kernel( \rho(\boldsymbol \eta_i) = \sum_{j=0}^{N-1} \rho_j\, W_h(\boldsymbol \eta_i - \boldsymbol \eta_j)\,. - The coefficients :math:`\rho_j` must be available in the marker array, stored at some index ``self.markers[j, index]``. + The coefficients :math:`\rho_j` must be available in the marker array, stored at some index ``self.markers[index, j]``. In case that `derivative=k` where `k` is not zero, the `k`-th component of the gradient of :math:`\rho` is computed: .. math:: @@ -80,7 +82,7 @@ def naive_evaluation_kernel( Evaluation point in logical space. holes : bool[:] - 1D array of length ``markers.shape[0]``. ``True`` if particle ``i`` is a hole (inactive). + 1D array of length ``markers.shape[1]``. ``True`` if particle ``i`` is a hole (inactive). periodic1, periodic2, periodic3 : bool ``True`` if the domain is periodic in that dimension. @@ -102,15 +104,15 @@ def naive_evaluation_kernel( markers = args_markers.markers - n_particles = len(markers) + n_particles = shape(markers)[1] out = 0.0 for p in range(n_particles): # skip holes and ghost particles (copies of markers already counted, id = -2) - if not holes[p] and markers[p, -1] != -2.0: - r1 = distance(eta1, markers[p, 0], periodic1) - r2 = distance(eta2, markers[p, 1], periodic2) - r3 = distance(eta3, markers[p, 2], periodic3) - out += markers[p, index] * sph_smoothing_kernels.smoothing_kernel(kernel_type, r1, r2, r3, h1, h2, h3) + if not holes[p] and markers[-1, p] != -2.0: + r1 = distance(eta1, markers[0, p], periodic1) + r2 = distance(eta2, markers[1, p], periodic2) + r3 = distance(eta3, markers[2, p], periodic3) + out += markers[index, p] * sph_smoothing_kernels.smoothing_kernel(kernel_type, r1, r2, r3, h1, h2, h3) return out @@ -138,7 +140,7 @@ def box_based_kernel( \rho(\boldsymbol \eta_i) = \sum_{j=0}^{N-1} \rho_j\, W_h(\boldsymbol \eta_i - \boldsymbol \eta_j)\,. - The coefficients :math:`\rho_j` must be available in the marker array, stored at some index ``self.markers[j, index]``. + The coefficients :math:`\rho_j` must be available in the marker array, stored at some index ``self.markers[index, j]``. In case that `derivative=k` where `k` is not zero, the `k`-th component of the gradient of :math:`\rho` is computed: .. math:: @@ -173,7 +175,7 @@ def box_based_kernel( Total number of particles. holes : bool - 1D array of length markers.shape[0]. True if markers[i] is a hole. + 1D array of length markers.shape[1]. True if markers[:, i] is a hole. periodic1, periodic2, periodic3 : bool True if periodic in that dimension. @@ -200,10 +202,10 @@ def box_based_kernel( p = boxes[box_to_search, c] c += 1 if not holes[p]: - r1 = distance(eta1, markers[p, 0], periodic1) - r2 = distance(eta2, markers[p, 1], periodic2) - r3 = distance(eta3, markers[p, 2], periodic3) - out += markers[p, index] * sph_smoothing_kernels.smoothing_kernel(kernel_type, r1, r2, r3, h1, h2, h3) + r1 = distance(eta1, markers[0, p], periodic1) + r2 = distance(eta2, markers[1, p], periodic2) + r3 = distance(eta3, markers[2, p], periodic3) + out += markers[index, p] * sph_smoothing_kernels.smoothing_kernel(kernel_type, r1, r2, r3, h1, h2, h3) return out diff --git a/src/struphy/pic/utilities_kernels.py b/src/struphy/pic/utilities_kernels.py index d781ac4ca..c56cb329f 100644 --- a/src/struphy/pic/utilities_kernels.py +++ b/src/struphy/pic/utilities_kernels.py @@ -19,7 +19,7 @@ @stack_array("dfm", "norm_b1", "b") def eval_magnetic_energy( - markers: "float[:,:](order=F)", # noqa: F821 + markers: "float[:,:]", args_derham: "DerhamArguments", args_domain: "DomainArguments", first_diagnostics_idx: int, @@ -33,7 +33,7 @@ def eval_magnetic_energy( ): r""" Evaluate :math:`mu_p |B(\boldsymbol \eta_p)_\parallel|` for each marker. - The result is stored at markers[:, first_diagnostics_idx]. + The result is stored at markers[first_diagnostics_idx, :]. """ norm_b1 = empty(3, dtype=float) b = empty(3, dtype=float) @@ -41,7 +41,7 @@ def eval_magnetic_energy( dfm = empty((3, 3), dtype=float) # get number of markers - n_markers = shape(markers)[0] + n_markers = shape(markers)[1] # fmt: off #$ omp parallel private(abs_B, b, b_para, det_df, dfm, eta1, eta2, eta3, ip, mu, norm_b1, span1, span2, span3, thread_derham) @@ -55,14 +55,14 @@ def eval_magnetic_energy( # fmt: on for ip in range(n_markers): # only do something if particle is a "true" particle (i.e. not a hole) - if markers[ip, 0] == -1.0: + if markers[0, ip] == -1.0: continue - eta1 = markers[ip, 0] - eta2 = markers[ip, 1] - eta3 = markers[ip, 2] + eta1 = markers[0, ip] + eta2 = markers[1, ip] + eta3 = markers[2, ip] - mu = markers[ip, first_diagnostics_idx + 1] + mu = markers[first_diagnostics_idx + 1, ip] # spline evaluation span1, span2, span3 = get_spans(eta1, eta2, eta3, thread_derham) @@ -114,7 +114,7 @@ def eval_magnetic_energy( b_para = linalg_kernels.scalar_dot(norm_b1, b) b_para /= det_df - markers[ip, first_diagnostics_idx] = mu * (abs_B + b_para) + markers[first_diagnostics_idx, ip] = mu * (abs_B + b_para) # fmt: off #$ omp end parallel From 1f83eafc0793fea444dcf8c50b21b2d6a69b3653 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 15 Oct 2026 17:18:38 +0200 Subject: [PATCH 05/12] Particles: one column per marker Particles.markers is allocated as (n_cols, n_rows) and indexed markers[j, ip]. The properties follow: markers_wo_holes(_and_ghost) are (n_cols, n), positions (3, n), velocities (vdim, n), phasespace_coords, f_coords and f_jacobian_coords (k, n); the setters check these shapes. Hole/ghost masks read rows 0 and -1, marker selections are markers[:, sel], loading writes columns (pseudo-random markers are still drawn particle by particle, so the random stream is unchanged), and _is_on_domain_of_rank, the sorting etas and the bc buffer are (3, n). MPI: the send lists, received buffers and the SPH ghost-box arrays are (n_cols, n). Selecting columns gives a Fortran-ordered copy, which MPI would send marker by marker into C-ordered receive buffers, so the send lists and the ghost-box markers are made C-contiguous. The device exchange receives into one buffer per source rank (a slice of a common buffer would not be contiguous). External marker files and restart data have the layout of the buffer, (n_cols, n). Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 836 ++++++++++++++++++----------------- src/struphy/pic/particles.py | 33 +- src/struphy/pic/sorting.py | 6 +- src/struphy/pic/utilities.py | 3 +- 4 files changed, 447 insertions(+), 431 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 64f72911a..15d1da872 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -88,10 +88,12 @@ class Particles(metaclass=ABCMeta): r""" Base class for particle species. - The marker information is stored in a 2D numpy array. - In ``markers[ip, j]`` The row index ``ip`` refers to a specific particle, - the column index ``j`` to its attributes. - The columns are indexed as follows: + The marker information is stored in a 2D array of shape ``(n_cols, n_rows)``, one column per marker + (see :mod:`struphy.pic.marker_layout`). + In ``markers[j, ip]`` the index ``ip`` refers to a specific particle, + the index ``j`` to its attributes; ``n_cols`` is the number of attributes and ``n_rows`` the number of markers + (including holes), and the attribute indices are called "columns" below and in the rest of the code. + They are as follows: * ``0:3``: position in the logical unit cube (:math:`\boldsymbol \eta_p \in [0, 1]^3`) * ``3:3 + vdim``: velocities @@ -122,7 +124,7 @@ class Particles(metaclass=ABCMeta): Direct indexing into ``markers`` is rarely needed outside of the ``Particles`` class itself. Instead, the most commonly used columns are exposed as convenience properties, each returning (or setting) a 2D array of shape ``(n_mks_loc, ...)`` restricted to the valid markers on the - current process (i.e. ``markers[self.valid_mks, ...]``): + current process (i.e. ``markers[..., self.valid_mks]``): * :attr:`positions` (columns ``0:3``): marker positions :math:`\boldsymbol\eta_p`. * :attr:`velocities` (columns ``3:3 + vdim``): marker velocities. @@ -420,9 +422,9 @@ def __init__( self._generate_sampling_moments() # create buffers for mpi_sort_markers - self._sorting_etas = xp.zeros((self.markers.shape[0], 3), dtype=float) - self._is_on_proc_domain = xp.zeros((self.markers.shape[0], 3), dtype=bool) - self._can_stay = xp.zeros(self.markers.shape[0], dtype=bool) + self._sorting_etas = xp.zeros((3, self.n_rows), dtype=float) + self._is_on_proc_domain = xp.zeros((3, self.n_rows), dtype=bool) + self._can_stay = xp.zeros(self.n_rows, dtype=bool) self._reqs = [None] * self.mpi_size self._recvbufs = [None] * self.mpi_size self._send_to_i = [None] * self.mpi_size @@ -548,13 +550,13 @@ def first_free_idx(self): @property def n_cols(self): - """Total umber of columns in markers array. + """Number of quantities ("columns") per marker, the length of the first axis of the markers array. The last 2 columns refer to box number and particle ID, respectively.""" return self.first_free_idx + self.n_cols_aux + 2 @property def n_rows(self): - """Total number of rows in markers array.""" + """Total number of markers (columns) in markers array, including holes.""" if not hasattr(self, "_n_rows"): input("\nWarning: marker array not yet created, creating now ...") self._allocate_marker_array() @@ -764,8 +766,8 @@ def n_mks_load(self): @property def markers(self): - """2D numpy array holding the marker information, including holes. - The i-th row holds the i-th marker info. + """2D array of shape ``(n_cols, n_rows)`` holding the marker information, including holes. + The i-th column holds the i-th marker info (see :mod:`struphy.pic.marker_layout`). ===== ============== ======================= ======= ====== ====== ========== === === === index | 0 | 1 | 2 | | 3 | ... | 3+(vdim-1)| 3+vdim 4+vdim 5+vdim >=6+vdim ... -2 -1 @@ -773,7 +775,7 @@ def markers(self): value position (eta) velocities weight s0 w0 other ... box ID ===== ============== ======================= ======= ====== ====== ========== === === === - The column indices referring to different attributes can be obtained from + The row indices referring to different attributes can be obtained from :attr:`~struphy.pic.base.Particles.index`. """ return self._markers @@ -782,25 +784,25 @@ def markers(self): def holes(self): """Array of booleans stating if an entry in the markers array is a hole.""" if not hasattr(self, "_holes"): - self._holes = self.markers[:, 0] == -1.0 + self._holes = self.markers[0] == -1.0 return self._holes @property def ghost_particles(self): """Array of booleans stating if an entry in the markers array is a ghost particle.""" if not hasattr(self, "_ghost_particles"): - self._ghost_particles = self.markers[:, -1] == -2.0 + self._ghost_particles = self.markers[-1] == -2.0 return self._ghost_particles @property def markers_wo_holes(self): - """Array holding the marker information, excluding holes. The i-th row holds the i-th marker info.""" - return self.markers[~self.holes] + """Array holding the marker information, excluding holes. The i-th column holds the i-th marker info.""" + return self.markers[:, ~self.holes] @property def markers_wo_holes_and_ghost(self): - """Array holding the marker information, excluding holes and ghosts (only valid markers). The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks] + """Array holding the marker information, excluding holes and ghosts (only valid markers). The i-th column holds the i-th marker info.""" + return self.markers[:, self.valid_mks] @property def lost_markers(self): @@ -846,98 +848,98 @@ def n_mks_global(self): @property def positions(self): - """Array holding the marker positions in logical space. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["pos"]] + """Array of shape ``(3, n_mks_loc)`` holding the marker positions in logical space. The i-th column holds the i-th marker info.""" + return self.markers[self.index["pos"], self.valid_mks] @positions.setter def positions(self, new): assert isinstance(new, xp.ndarray) - assert new.shape == (self.n_mks_loc, 3) - self._markers[self.valid_mks, self.index["pos"]] = new + assert new.shape == (3, self.n_mks_loc) + self._markers[self.index["pos"], self.valid_mks] = new @property def velocities(self): - """Array holding the marker velocities in logical space. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["vel"]] + """Array of shape ``(vdim, n_mks_loc)`` holding the marker velocities in logical space. The i-th column holds the i-th marker info.""" + return self.markers[self.index["vel"], self.valid_mks] @velocities.setter def velocities(self, new): assert isinstance(new, xp.ndarray) - assert new.shape == (self.n_mks_loc, self.vdim), f"{self.n_mks_loc =} and {self.vdim =} but {new.shape =}" - self._markers[self.valid_mks, self.index["vel"]] = new + assert new.shape == (self.vdim, self.n_mks_loc), f"{self.n_mks_loc =} and {self.vdim =} but {new.shape =}" + self._markers[self.index["vel"], self.valid_mks] = new @property def phasespace_coords(self): - """Array holding the marker positions and velocities in logical space. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["coords"]] + """Array of shape ``(3 + vdim, n_mks_loc)`` holding the marker positions and velocities in logical space. The i-th column holds the i-th marker info.""" + return self.markers[self.index["coords"], self.valid_mks] @phasespace_coords.setter def phasespace_coords(self, new): assert isinstance(new, xp.ndarray) - assert new.shape == (self.n_mks_loc, 3 + self.vdim) - self._markers[self.valid_mks, self.index["coords"]] = new + assert new.shape == (3 + self.vdim, self.n_mks_loc) + self._markers[self.index["coords"], self.valid_mks] = new @property def weights(self): - """Array holding the current marker weights. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["weights"]] + """Array holding the current marker weights. The i-th entry holds the i-th marker info.""" + return self.markers[self.index["weights"], self.valid_mks] @weights.setter def weights(self, new): assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) - self._markers[self.valid_mks, self.index["weights"]] = new + self._markers[self.index["weights"], self.valid_mks] = new @property def sampling_density_values(self): - """Array holding the current marker 0form sampling density s0. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["s0"]] + """Array holding the current marker 0form sampling density s0. The i-th entry holds the i-th marker info.""" + return self.markers[self.index["s0"], self.valid_mks] @sampling_density_values.setter def sampling_density_values(self, new): assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) - self._markers[self.valid_mks, self.index["s0"]] = new + self._markers[self.index["s0"], self.valid_mks] = new @property def weights0(self): - """Array holding the initial marker weights. The i-th row holds the i-th marker info.""" - return self.markers[self.valid_mks, self.index["w0"]] + """Array holding the initial marker weights. The i-th entry holds the i-th marker info.""" + return self.markers[self.index["w0"], self.valid_mks] @weights0.setter def weights0(self, new): assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) - self._markers[self.valid_mks, self.index["w0"]] = new + self._markers[self.index["w0"], self.valid_mks] = new @property def marker_ids(self): """Array holding the marker id's on the current process.""" - return self.markers[self.valid_mks, self.index["ids"]] + return self.markers[self.index["ids"], self.valid_mks] @marker_ids.setter def marker_ids(self, new): assert isinstance(new, xp.ndarray) assert new.shape == (self.n_mks_loc,) - self._markers[self.valid_mks, self.index["ids"]] = new + self._markers[self.index["ids"], self.valid_mks] = new @property def f_coords(self): """Coordinates of the distribution function.""" - return self.markers[self.valid_mks, self.f_coords_index] + return self.markers[self.f_coords_index, self.valid_mks] @f_coords.setter def f_coords(self, new): assert isinstance(new, xp.ndarray) - self.markers[self.valid_mks, self.f_coords_index] = new + self.markers[self.f_coords_index, self.valid_mks] = new @property def f_jacobian_coords(self): """Coordinates of the velocity jacobian determinant of the distribution fuction.""" if isinstance(self.f_jacobian_coords_index, list): - return self.markers[xp.ix_(~self.holes, self.f_jacobian_coords_index)] + return self.markers[xp.ix_(self.f_jacobian_coords_index, ~self.holes)] else: - return self.markers[~self.holes, self.f_jacobian_coords_index] + return self.markers[self.f_jacobian_coords_index, ~self.holes] @f_jacobian_coords.setter def f_jacobian_coords(self, new): @@ -945,12 +947,12 @@ def f_jacobian_coords(self, new): if isinstance(self.f_jacobian_coords_index, list): self.markers[ xp.ix_( - ~self.holes, self.f_jacobian_coords_index, + ~self.holes, ) ] = new else: - self.markers[~self.holes, self.f_jacobian_coords_index] = new + self.markers[self.f_jacobian_coords_index, ~self.holes] = new @property def args_markers(self) -> MarkerArguments | CudaMarkerArguments: @@ -1181,7 +1183,7 @@ def draw_markers( # Np_per_clone_loc = self.Np_per_clone[self.clone_id] # fill holes in markers array with -1 (all holes are at end of array at loading stage) - self._markers[n_mks_load_loc:] = -1.0 + self._markers[:, n_mks_load_loc:] = -1.0 # number of holes and markers on process self.update_holes() @@ -1215,7 +1217,7 @@ def draw_markers( self._load_tesselation() if isinstance(self, ParticlesSPH): self._set_initial_condition() - self.velocities = xp.array(self.u_init(self.positions)).T + self.velocities = xp.array(self.u_init(self.positions)) # set markers ID in last column self.marker_ids = _first_marker_id + xp.arange(n_mks_load_loc, dtype=float) else: @@ -1239,21 +1241,23 @@ def draw_markers( while num_loaded_particles_glob < int(self.Np): # Generate a chunk of random particles num_to_add_glob = min(chunk_size, int(self.Np) - num_loaded_particles_glob) - temp = xp.random.rand(num_to_add_glob, 3 + self.vdim) + # one column per particle; drawn row by row, so that the random numbers of a particle are + # consecutive in the stream + temp = xp.random.rand(num_to_add_glob, 3 + self.vdim).T # inversion method for drawing uniformly on the disc, before the domain check if self.spatial == "disc": - temp[:, 0] = xp.sqrt(temp[:, 0]) + 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] - valid_particles = temp[valid_idx] - valid_particles = xp.array_split(valid_particles, self.num_clones)[self.clone_id] - num_valid = valid_particles.shape[0] + is_on_proc_domain = self._is_on_domain_of_rank(temp[:3], self.mpi_rank) + valid_idx = xp.nonzero(xp.all(is_on_proc_domain, axis=0))[0] + valid_particles = temp[:, valid_idx] + valid_particles = xp.array_split(valid_particles, self.num_clones, axis=1)[self.clone_id] + num_valid = valid_particles.shape[1] # Add the valid particles to the phasespace_coords array self._markers[ - num_loaded_particles_loc : num_loaded_particles_loc + num_valid, : 3 + self.vdim, + num_loaded_particles_loc : num_loaded_particles_loc + num_valid, ] = valid_particles num_loaded_particles_glob += num_to_add_glob num_loaded_particles_loc += num_valid @@ -1274,7 +1278,7 @@ def draw_markers( )[self._mpi_rank] # set new holes in markers array to -1 - self._markers[num_loaded_particles_loc:] = -1.0 + self._markers[:, num_loaded_particles_loc:] = -1.0 self.update_holes() # 2. plain sobol numbers with skip of first 1000 numbers @@ -1285,7 +1289,7 @@ def draw_markers( 3 + self.vdim, int(n_mks_load_loc), 1000 + int((n_mks_load_cum_sum - self.n_mks_load)[self._mpi_rank]), - ) + ).T ) # 3. symmetric sobol numbers in all 6 dimensions with skip of first 1000 numbers @@ -1304,12 +1308,12 @@ def draw_markers( ) ) - temp_symmetric = xp.zeros((64 * n_sobol[self._mpi_rank], 3 + self.vdim), dtype=float) + temp_symmetric = xp.zeros((3 + self.vdim, 64 * n_sobol[self._mpi_rank]), dtype=float) sampling_kernels.set_particles_symmetric_3d_3v( temp_markers, temp_symmetric, ) - self._markers[:n_mks_load_loc, : 3 + self.vdim] = temp_symmetric[:n_mks_load_loc] + self._markers[: 3 + self.vdim, :n_mks_load_loc] = temp_symmetric[:, :n_mks_load_loc] # 4. Wrong specification else: @@ -1320,14 +1324,14 @@ def draw_markers( # initial velocities - SPH case: v(0) = u(x(0)) for given velocity u(x) if isinstance(self, ParticlesSPH): self._set_initial_condition() - self.velocities = xp.array(self.u_init(self.positions)).T + self.velocities = xp.array(self.u_init(self.positions)) else: # inverse transform sampling in velocity space # Avoid exact 0 or 1 from low-discrepancy sequences: erfinv(±1) # and log(0) produce infinities or invalid polar velocities. eps = xp.finfo(float).eps - self._markers[:n_mks_load_loc, 3 : 3 + self.vdim] = xp.clip( - self._markers[:n_mks_load_loc, 3 : 3 + self.vdim], + self._markers[3 : 3 + self.vdim, :n_mks_load_loc] = xp.clip( + self._markers[3 : 3 + self.vdim, :n_mks_load_loc], eps, 1.0 - eps, ) @@ -1343,21 +1347,21 @@ def draw_markers( 2 * self.velocities - 1, ) * xp.sqrt(2) - * v_th - + u_mean + * v_th[:, None] + + u_mean[:, None] ) # Particles5D: (1d Maxwellian, muB0-Maxwellian as volume-form) elif isinstance(self, Particles5D): - self._markers[:n_mks_load_loc, 3] = ( + self._markers[3, :n_mks_load_loc] = ( sp.erfinv( - 2 * self.velocities[:, 0] - 1, + 2 * self.velocities[0] - 1, ) * xp.sqrt(2) * v_th[0] + u_mean[0] ) - self._markers[:n_mks_load_loc, 4] = -xp.log(1.0 - self.velocities[:, 1]) * v_th[1] ** 2 / B0 + self._markers[4, :n_mks_load_loc] = -xp.log(1.0 - self.velocities[1]) * v_th[1] ** 2 / B0 # mu is a magnetic moment and must be >= 0. # A mean shift in this coordinate is not physically consistent. @@ -1368,18 +1372,18 @@ def draw_markers( ) # Particles5Dvperp: (1d Maxwellian, polar Maxwellian as volume-form) elif isinstance(self, Particles5Dvperp): - self._markers[:n_mks_load_loc, 3] = ( + self._markers[3, :n_mks_load_loc] = ( sp.erfinv( - 2 * self.velocities[:, 0] - 1, + 2 * self.velocities[0] - 1, ) * xp.sqrt(2) * v_th[0] + u_mean[0] ) - self._markers[:n_mks_load_loc, 4] = ( + self._markers[4, :n_mks_load_loc] = ( xp.sqrt( - -xp.log(1.0 - self.velocities[:, 1]), + -xp.log(1.0 - self.velocities[1]), ) * xp.sqrt(2) * v_th[1] @@ -1402,8 +1406,8 @@ def draw_markers( # inversion method for drawing uniformly on the disc (done above for pseudo-random markers) if self.spatial == "disc": if self.loading != "pseudo_random": - self._markers[:n_mks_load_loc, 0] = xp.sqrt( - self._markers[:n_mks_load_loc, 0], + self._markers[0, :n_mks_load_loc] = xp.sqrt( + self._markers[0, :n_mks_load_loc], ) else: assert self.spatial == "uniform", f'Spatial drawing must be "uniform" or "disc", is {self.spatial}.' @@ -1416,13 +1420,10 @@ def draw_markers( counter = 0 for i in range(len(specific_markers)): - if i == int(self.markers[counter, -1]): + if i == int(self.markers[-1, counter]): for j in range(3 + self.vdim): if specific_markers[i][j] is not None: - self._markers[ - counter, - j, - ] = specific_markers[i][j] + self._markers[j, counter] = specific_markers[i][j] counter += 1 @@ -1434,7 +1435,7 @@ def draw_markers( # 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 = 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 @@ -1515,7 +1516,7 @@ def initialize_weights( if isinstance(self, ParticlesSPH): f_init = self.f_init(self.positions) else: - f_init = self.f_init(*self.f_coords.T) + f_init = self.f_init(*self.f_coords) # if f_init is vol-form, transform to 0-form if self.is_volume_form[0]: @@ -1523,18 +1524,18 @@ def initialize_weights( if self.is_volume_form[1]: f_init /= self.f_init.velocity_jacobian_det( - *self.f_jacobian_coords.T, + *self.f_jacobian_coords, ) # compute s0 and save at vdim + 4 - self.sampling_density_values = self.s0(*self.phasespace_coords.T, flat_eval=True) + self.sampling_density_values = self.s0(*self.phasespace_coords, flat_eval=True) # compute w0 and save at vdim + 5 self.weights0 = f_init / self.sampling_density_values / self.Np if self.reject_weights: - reject = self.markers[:, self.index["w0"]] < self.threshold - self._markers[reject] = -1.0 + reject = self.markers[self.index["w0"]] < self.threshold + self._markers[:, reject] = -1.0 self.update_holes() self._reset_marker_ids() logger.info( @@ -1565,14 +1566,14 @@ def update_weights(self): # in case of CanonicalMaxwellian, evaluate constants_of_motion # if isinstance(self.f0, CanonicalMaxwellian): # self.save_constants_of_motion() - f0 = self.f0(*self.f_coords.T) + f0 = self.f0(*self.f_coords) # if f_init is vol-form, transform to 0-form if self.is_volume_form[0]: f0 /= self.domain.jacobian_det(self.positions) if self.is_volume_form[1]: - f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords.T) + f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords) self.weights = self.weights0 - f0 / self.sampling_density_values / self.Np @@ -1620,45 +1621,44 @@ def binning( # extend components list to number of columns of markers array _n = len(components) - slicing = list(components) + [False] * (self.markers.shape[1] - _n) + slicing = list(components) + [False] * (self.n_cols - _n) # determine type of output quantity # Note: "density" Literal does not have "_" quantity, *v_axis = output_quantity.rsplit(sep="_", maxsplit=1) v_axis = [int(char) - 1 for char in "".join(v_axis)] # convert dimension axis to index - # the rows of the valid markers, found once; below, only the columns that are used are gathered from them - # (each one contiguous in the column-major marker buffer), not whole rows. Everything stays on the device - # of the markers. - rows = xp.nonzero(self.valid_mks)[0] + # the indices of the valid markers, found once; below, only the rows (quantities) that are used are gathered + # for them, not all quantities of the markers. Everything stays on the device of the markers. + valid = xp.nonzero(self.valid_mks)[0] if quantity != "density": - velocities = self.markers[rows, self.index["vel"]] + velocities = self.markers[self.index["vel"], valid] # determine histogram weights multiplier if quantity == "density": multiplier = 1 elif quantity == "current": - multiplier = velocities[:, v_axis[0]] + multiplier = velocities[v_axis[0]] elif quantity == "energy_tensor": - multiplier = velocities[:, v_axis[0]] * velocities[:, v_axis[1]] + multiplier = velocities[v_axis[0]] * velocities[v_axis[1]] elif quantity == "heat_flux": - velocity_norm2 = xp.linalg.norm(velocities, axis=1) ** 2 - multiplier = velocity_norm2 * velocities[:, v_axis[0]] + velocity_norm2 = xp.linalg.norm(velocities, axis=0) ** 2 + multiplier = velocity_norm2 * velocities[v_axis[0]] # compute weights of histogram: - _weights0 = self.markers[rows, self.index["w0"]] * self.Np * multiplier - _weights = self.markers[rows, self.index["weights"]] * self.Np * multiplier + _weights0 = self.markers[self.index["w0"], valid] * self.Np * multiplier + _weights = self.markers[self.index["weights"], valid] * self.Np * multiplier if divide_by_jac: - jacobian_det = self.domain.jacobian_det(self.markers[rows, self.index["pos"]], remove_outside=False) + jacobian_det = self.domain.jacobian_det(self.markers[self.index["pos"], valid], remove_outside=False) _weights /= jacobian_det - # _weights /= self.velocity_jacobian_det(*self.phasespace_coords.T) + # _weights /= self.velocity_jacobian_det(*self.phasespace_coords) _weights0 /= jacobian_det - # _weights0 /= self.velocity_jacobian_det(*self.phasespace_coords.T) + # _weights0 /= self.velocity_jacobian_det(*self.phasespace_coords) - # the binned columns, one by one (an index list would have to be uploaded to the device first) - sample = xp.stack([self.markers[rows, i] for i, binned in enumerate(slicing) if binned], axis=1) + # the binned quantities, one by one (an index list would have to be uploaded to the device first) + sample = xp.stack([self.markers[i, valid] for i, binned in enumerate(slicing) if binned], axis=1) f_slice = xp.histogramdd( sample, @@ -1793,7 +1793,7 @@ def mpi_sort_markers( Sorts markers according to MPI domain decomposition. Markers are sent to the process corresponding to the alpha-weighted position - alpha*markers[:, 0:3] + (1 - alpha)*markers[:, first_pusher_idx:first_pusher_idx + 3]. + alpha*markers[0:3] + (1 - alpha)*markers[first_pusher_idx:first_pusher_idx + 3]. Periodic boundary conditions are taken into account when computing the alpha-weighted position. @@ -1804,7 +1804,7 @@ def mpi_sort_markers( Whether to apply kinetic boundary conditions before sorting. alpha : tuple | list | int | float - For i=1,2,3 the sorting is according to alpha[i]*markers[:, i] + (1 - alpha[i])*markers[:, first_pusher_idx + i]. + For i=1,2,3 the sorting is according to alpha[i]*markers[i] + (1 - alpha[i])*markers[first_pusher_idx + i]. If int or float then alpha = (alpha, alpha, alpha). alpha must be between 0 and 1. do_test : bool @@ -1815,7 +1815,7 @@ def mpi_sort_markers( Notes ----- - On the CuPy backend, the marker rows are exchanged as device buffers with CUDA-aware MPI + On the CuPy backend, the markers are exchanged as device buffers with CUDA-aware MPI (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ @@ -1838,7 +1838,7 @@ def mpi_sort_markers( send_info = self._sendrecv_get_destinations(send_inds) # set new holes in markers array to -1 - self._markers[send_inds] = -1.0 + self._markers[:, send_inds] = -1.0 # transpose send_info recv_info = self._sendrecv_all_to_all(send_info) @@ -1849,7 +1849,7 @@ def mpi_sort_markers( # new holes and new number of holes and markers on process self.update_holes() - # refresh ghost mask: received markers may land in rows that previously held + # refresh ghost mask: received markers may land in columns that previously held # ghost particles. update_holes alone recomputes valid_mks from a stale # _ghost_particles mask, which would wrongly exclude these incoming real markers. self._update_ghost_particles() @@ -1859,7 +1859,7 @@ def mpi_sort_markers( # eta = 1 is sorted like eta = 0 in periodic directions (see _sendrecv_determine_mtbs) etas = self.positions for axis in self._periodic_axes: - etas[:, axis] %= 1.0 + etas[axis] %= 1.0 all_on_right_proc = xp.all(self._is_on_domain_of_rank(etas, self.mpi_rank)) @@ -1885,16 +1885,16 @@ def apply_kinetic_bc(self, newton=False): 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 + self.markers[0, crossed] *= -1.0 + self.markers[1, crossed] = (self.markers[1, crossed] + 0.5) % 1.0 + near_pole = xp.logical_and(self.valid_mks, self.markers[0] < POLE_ETA1_MIN) + self.markers[0, near_pole] = POLE_ETA1_MIN # apply boundary conditions if self._remove_axes: - # snapshot of the (n_rows, 3) logical coordinates for the loop over the axes + # snapshot of the (3, n_rows) logical coordinates for the loop over the axes # (see _find_outside_particles) - self._eta_bc_buf[:] = self.markers[:, :3] + self._eta_bc_buf[:] = self.markers[:3] for axis in self._remove_axes: outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf) @@ -1905,16 +1905,16 @@ def apply_kinetic_bc(self, newton=False): if self.bc_refill is not None: self._particle_refilling() - self._markers[self._is_outside, :-1] = -1.0 + self._markers[:-1, self._is_outside] = -1.0 self._n_lost_markers += len(xp.nonzero(self._is_outside)[0]) # removed markers are holes now and refilled markers have moved: refresh both, # such that a marker outside on several axes is only counted (and refilled) once self.update_holes() - self._eta_bc_buf[outside_inds] = self.markers[outside_inds, :3] + self._eta_bc_buf[:, outside_inds] = self.markers[:3, outside_inds] if self._periodic_axes: - self._eta_bc_buf[:] = self.markers[:, :3] + self._eta_bc_buf[:] = self.markers[:3] for axis in self._periodic_axes: outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf) @@ -1922,47 +1922,32 @@ def apply_kinetic_bc(self, newton=False): if len(outside_inds) == 0: continue - self.markers[outside_inds, axis] = self.markers[outside_inds, axis] % 1.0 + self.markers[axis, outside_inds] = self.markers[axis, outside_inds] % 1.0 # set shift for alpha-weighted mid-point computation outside_right_inds = xp.nonzero(self._is_outside_right)[0] outside_left_inds = xp.nonzero(self._is_outside_left)[0] if newton: - self.markers[ - outside_right_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] += 1.0 - self.markers[ - outside_left_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] += -1.0 + self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_right_inds] += 1.0 + self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_left_inds] += -1.0 else: - self.markers[ - :, - self.first_pusher_idx + 3 + self.vdim + axis, - ] = 0.0 - self.markers[ - outside_right_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] = 1.0 - self.markers[ - outside_left_inds, - self.first_pusher_idx + 3 + self.vdim + axis, - ] = -1.0 + self.markers[self.first_pusher_idx + 3 + self.vdim + axis] = 0.0 + self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_right_inds] = 1.0 + self.markers[self.first_pusher_idx + 3 + self.vdim + axis, outside_left_inds] = -1.0 # put all coordinate inside the unit cube (avoid wrong Jacobian evaluations) outside_inds_per_axis = {} if self._reflect_axes: - self._eta_bc_buf[:] = self.markers[:, :3] + self._eta_bc_buf[:] = self.markers[:3] for axis in self._reflect_axes: outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf) - self.markers[self._is_outside_left, axis] *= -1.0 - self.markers[self._is_outside_right, axis] *= -1.0 - self.markers[self._is_outside_right, axis] += 2.0 + self.markers[axis, self._is_outside_left] *= -1.0 + self.markers[axis, self._is_outside_right] *= -1.0 + self.markers[axis, self._is_outside_right] += 2.0 - self.markers[self._is_outside, self.first_pusher_idx] = -1.0 + self.markers[self.first_pusher_idx, self._is_outside] = -1.0 outside_inds_per_axis[axis] = outside_inds @@ -1997,11 +1982,11 @@ def finish_kernel_bc(self, newton=False): self._n_lost_markers += int(xp.count_nonzero(self.holes) - n_holes_before) def update_holes(self): - """Recompute the :attr:`~struphy.pic.base.Particles.holes` mask (rows with ``markers[:, 0] == -1``) + """Recompute the :attr:`~struphy.pic.base.Particles.holes` mask (markers with ``markers[0] == -1``) and, from it, refresh :attr:`~struphy.pic.base.Particles.valid_mks`. Must be called after any operation that creates, removes or moves markers - (e.g. sorting, boundary conditions, refilling), since holes are tracked per row index.""" - self._holes[:] = self.markers[:, 0] == -1.0 + (e.g. sorting, boundary conditions, refilling), since holes are tracked per marker index.""" + self._holes[:] = self.markers[0] == -1.0 self._update_valid_mks() def set_velocities_comp(self, velocity, comp): @@ -2015,10 +2000,8 @@ def set_velocities_comp(self, velocity, comp): comp : iterable[int] Velocity components to set (0-based, e.g. 0 for v1, 1 for v2, ...). """ - new = xp.ones(shape=(self.velocities.shape[0], 1)) * velocity - for c in comp: - self._markers[self.valid_mks, slice(3 + c, 3 + c + 1)] = new + self._markers[3 + c, self.valid_mks] = velocity @profile @ProfileManager.profile("put_particles_in_boxes") @@ -2055,7 +2038,7 @@ def put_particles_in_boxes(self): @ProfileManager.profile("do_sort") def do_sort(self, use_numpy_argsort=False): """Assign the particles to their sorting boxes and reorder the markers array accordingly, - so that markers in the same box occupy contiguous rows. + so that markers in the same box occupy contiguous columns. Parameters ---------- @@ -2082,7 +2065,7 @@ def do_sort(self, use_numpy_argsort=False): self._sorting_boxes._cumul_next_index, ) - # The marker rows have just been reordered. The masks are row-based, + # The markers have just been reordered. The masks are per marker index, # so they must be rebuilt before any later use of valid_mks/f_coords. self.update_holes() self._update_ghost_particles() @@ -2368,7 +2351,7 @@ def eval_div_viscosity( ] # the stored tensor is in Piola form, divide by the Jacobian determinant at the evaluation points - etas = xp.column_stack([xp.ravel(eta1), xp.ravel(eta2), xp.ravel(eta3)]) + etas = xp.stack([xp.ravel(eta1), xp.ravel(eta2), xp.ravel(eta3)], axis=0) det_df = self.domain.jacobian_det(etas, remove_outside=False).reshape(xp.shape(eta1)) gamma_x = (gamma[0][0] + gamma[0][1] + gamma[0][2]) / det_df @@ -2484,9 +2467,9 @@ def _get_domain_decomp(self, mpi_dims_mask: tuple | list = None): ------- dom_arr : xp.ndarray A 2d array of shape (#MPI processes, 9). The row index denotes the process rank. The columns are for n=0,1,2: - - arr[i, 3*n + 0] holds the LEFT domain boundary of process i in direction eta_(n+1). - - arr[i, 3*n + 1] holds the RIGHT domain boundary of process i in direction eta_(n+1). - - arr[i, 3*n + 2] holds the number of cells of process i in direction eta_(n+1). + - arr[3*n + 0, i] holds the LEFT domain boundary of process i in direction eta_(n+1). + - arr[3*n + 1, i] holds the RIGHT domain boundary of process i in direction eta_(n+1). + - arr[3*n + 2, i] holds the number of cells of process i in direction eta_(n+1). nprocs : tuple The number of processes in each direction. @@ -2607,7 +2590,7 @@ def _allocate_marker_array(self, dry_run: bool = False): bufsize = self.bufsize + 1.0 / xp.sqrt(n_mks_load_loc) # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer; - # column-major, every column starting on a 256-byte boundary (see struphy.pic.marker_layout) + # one column per marker, every row starting on a 256-byte boundary (see struphy.pic.marker_layout) self._n_rows = padded_n_rows(round(float(n_mks_load_loc * (1 + bufsize)))) # Have at least 3 spare places in markers array @@ -2618,7 +2601,7 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return - self._markers = marker_array(self.n_rows, self.n_cols) + self._markers = marker_array(self.n_cols, self.n_rows) # allocate auxiliary arrays self._holes = xp.zeros(self.n_rows, dtype=bool) @@ -2627,13 +2610,13 @@ def _allocate_marker_array(self, dry_run: bool = False): self._is_outside_right = xp.zeros(self.n_rows, dtype=bool) self._is_outside_left = xp.zeros(self.n_rows, dtype=bool) self._is_outside = xp.zeros(self.n_rows, dtype=bool) - # scratch copy of markers[:, :3], refreshed once per apply_kinetic_bc boundary-condition-type - # loop (see there); column-major like the markers, so the copy is three contiguous columns - self._eta_bc_buf = marker_array(self.n_rows, 3) + # scratch copy of markers[:3], refreshed once per apply_kinetic_bc boundary-condition-type + # loop (see there) + self._eta_bc_buf = marker_array(3, self.n_rows) # create array container (3 x positions, vdim x velocities, weight, s0, w0, ID) for removed markers self._n_lost_markers = 0 - self._lost_markers = xp.zeros((int(self.n_rows * 0.5), 10), dtype=float) + self._lost_markers = xp.zeros((10, int(self.n_rows * 0.5)), dtype=float) # arguments for kernels: the pyccel class on the NumPy backend, its CUDA version on the CuPy backend args_class = CudaMarkerArguments if self._args_backend == "cupy" else MarkerArguments @@ -2834,7 +2817,7 @@ def _f_init(*etas, flat_eval=False): if _density is None: out = self.f0.n0(etas[0]) else: - out = self.f0.n0(etas[0]) + _density(*etas[0].T) + out = self.f0.n0(etas[0]) + _density(*etas[0]) else: assert len(etas) == 3 E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts( @@ -2859,7 +2842,7 @@ def _u_init(*etas, flat_eval=False): if len(etas) == 1: out = self.f0.uv(etas[0]) if _u1 is not None: - out[0] += _u1(*etas[0].T) + out[0] += _u1(*etas[0]) else: assert len(etas) == 3 E1, E2, E3, is_sparse_meshgrid = Domain.prepare_eval_pts( @@ -2890,6 +2873,9 @@ def _load_external( ): """Load markers from external .hdf5 file. + The dataset ``markers`` of the file has the layout of the marker array, shape ``(n_cols, n_markers)`` + (one column per marker). + Parameters ---------- n_mks_load_loc: int @@ -2902,24 +2888,21 @@ def _load_external( with h5py.File(self.loading_params.dir_external, "r") as file: logger.info(f"\nLoading markers from file: {file}") - self._markers[ - : n_mks_load_cum_sum[0], - :, - ] = file["markers"][: n_mks_load_cum_sum[0], :] + self._markers[:, : n_mks_load_cum_sum[0]] = file["markers"][:, : n_mks_load_cum_sum[0]] for i in range(1, self._mpi_size): self._mpi_comm.Send( - file["markers"][n_mks_load_cum_sum[i - 1] : n_mks_load_cum_sum[i], :], + file["markers"][:, n_mks_load_cum_sum[i - 1] : n_mks_load_cum_sum[i]], dest=i, tag=123, ) else: recvbuf = xp.zeros( - (n_mks_load_loc, self.markers.shape[1]), + (self.n_cols, n_mks_load_loc), dtype=float, ) self._mpi_comm.Recv(recvbuf, source=0, tag=123) - self._markers[:n_mks_load_loc, :] = recvbuf + self._markers[:, :n_mks_load_loc] = recvbuf def _load_restart(self): """Load markers from restart .hdf5 file.""" @@ -2940,15 +2923,15 @@ def _load_restart(self): with h5py.File(data.file_path, "a") as file: saved = file["restart/" + self.loading_params.restart_key][-1, :, :] - # the number of rows is padded (see struphy.pic.marker_layout), so a file written with another padding has a - # different number of rows: copy the common rows; the rows after them are holes, in the buffer (set before - # loading) and in the file - n_rows = min(saved.shape[0], self.n_rows) - assert np.all(saved[n_rows:, 0] == -1.0), ( - f"the restart file has {saved.shape[0]} marker rows on rank {self.mpi_rank}, but there are only " - f"{self.n_rows} rows here, and the extra ones are not holes" + # the number of markers n_rows is padded (see struphy.pic.marker_layout), so a file written with another + # padding has a different number of markers: copy the common ones; the markers after them are holes, in the + # buffer (set before loading) and in the file + n_rows = min(saved.shape[1], self.n_rows) + assert np.all(saved[0, n_rows:] == -1.0), ( + f"the restart file has {saved.shape[1]} markers on rank {self.mpi_rank}, but there are only " + f"{self.n_rows} here, and the extra ones are not holes" ) - self._markers[:n_rows, :] = saved[:n_rows] + self._markers[:, :n_rows] = saved[:, :n_rows] def _load_tesselation(self, n_quad: int = 1): """ @@ -2966,9 +2949,9 @@ def _load_tesselation(self, n_quad: int = 1): sorting_boxes=self.sorting_boxes, ) eta1, eta2, eta3 = self.tesselation.draw_markers() - self._markers[: eta1.size, 0] = eta1 - self._markers[: eta2.size, 1] = eta2 - self._markers[: eta3.size, 2] = eta3 + self._markers[0, : eta1.size] = eta1 + self._markers[1, : eta2.size] = eta2 + self._markers[2, : eta3.size] = eta3 self._update_valid_mks() def _reset_marker_ids(self): @@ -2989,18 +2972,18 @@ def _find_outside_particles(self, axis, eta=None): Parameters ---------- axis : int - Column of the markers array (0, 1 or 2) holding the logical coordinate to check. + Row of the markers array (0, 1 or 2) holding the logical coordinate to check. eta : xp.ndarray[float], optional - Snapshot ``(n_rows, 3)`` of ``markers[:, :3]`` (see :meth:`apply_kinetic_bc`). + Snapshot ``(3, n_rows)`` of ``markers[:3]`` (see :meth:`apply_kinetic_bc`). If ``None``, reads straight from ``markers``. Returns ------- outside_inds : xp.ndarray[int] - Row indices of the markers that are outside the logical unit cube. + Indices of the markers that are outside the logical unit cube. """ - col = self.markers[:, axis] if eta is None else eta[:, axis] + col = self.markers[axis] if eta is None else eta[axis] # determine particles outside of the logical unit cube self._is_outside_right[:] = col > 1.0 @@ -3042,12 +3025,12 @@ def _particle_refilling(self): # sorting out particles which are out of the domain if kind == "inner": outside_inds = xp.nonzero(self._is_outside_left)[0] - self.markers[outside_inds, 0] = 1e-4 + self.markers[0, outside_inds] = 1e-4 r_loss = self.domain.params["a1"] else: outside_inds = xp.nonzero(self._is_outside_right)[0] - self.markers[outside_inds, 0] = 1 - 1e-4 + self.markers[0, outside_inds] = 1 - 1e-4 r_loss = 1.0 if len(outside_inds) == 0: @@ -3058,7 +3041,7 @@ def _particle_refilling(self): gyro_inside_inds = self._gyro_transfer(outside_inds) # mark the particle as done for multiple step pushers - self.markers[outside_inds[gyro_inside_inds], self.first_pusher_idx] = -1.0 + self.markers[self.first_pusher_idx, outside_inds[gyro_inside_inds]] = -1.0 self._is_outside[outside_inds[gyro_inside_inds]] = False # exclude particles whose guiding center positions are still inside. @@ -3066,13 +3049,13 @@ def _particle_refilling(self): outside_inds = outside_inds[~gyro_inside_inds] # do phi boundary transfer = phi_loss - 2*q(r_loss)*theta_loss - self.markers[outside_inds, 2] -= 2 * self.equil.q_r(r_loss) * self.markers[outside_inds, 1] + self.markers[2, outside_inds] -= 2 * self.equil.q_r(r_loss) * self.markers[1, outside_inds] # theta_boudary_transfer = - theta_loss - self.markers[outside_inds, 1] = 1.0 - self.markers[outside_inds, 1] + self.markers[1, outside_inds] = 1.0 - self.markers[1, outside_inds] # mark the particle as done for multiple step pushers - self.markers[outside_inds, self.first_pusher_idx] = -1.0 + self.markers[self.first_pusher_idx, outside_inds] = -1.0 self._is_outside[outside_inds] = False def _gyro_transfer(self, outside_inds): @@ -3108,13 +3091,13 @@ def _gyro_transfer(self, outside_inds): assert self.vdim == 3 # TODO: currently assumes periodic boundary condition along poloidal and toroidal angle - self.markers[outside_inds, 1:3] = self.markers[outside_inds, 1:3] % 1 + self.markers[1:3, outside_inds] = self.markers[1:3, outside_inds] % 1 - v = self.markers[outside_inds, 3:6].T + v = self.markers[3:6, outside_inds] # eval cartesian equilibrium magnetic field at the marker positions assert isinstance(self.equil, FluidEquilibriumWithB), "Gyro transfer function needs a magnetic background." - b_cart, xyz = self.equil.b_cart(self.markers[outside_inds, :]) + b_cart, xyz = self.equil.b_cart(self.markers[:, outside_inds]) # calculate magnetic field amplitude and normalized magnetic field absB0 = xp.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2) @@ -3139,12 +3122,12 @@ def _gyro_transfer(self, outside_inds): gc_etas = self.domain.inverse_map(*xyz, bounded=False) # gyro transfer - self.markers[outside_inds, 1] = (gc_etas[1] - (self.markers[outside_inds, 1] - gc_etas[1]) % 1) % 1 + self.markers[1, outside_inds] = (gc_etas[1] - (self.markers[1, outside_inds] - gc_etas[1]) % 1) % 1 - new_xyz = self.domain(self.markers[outside_inds, :]) + new_xyz = self.domain(self.markers[:, outside_inds]) # eval cartesian equilibrium magnetic field at the marker positions - b_cart = self.equil.b_cart(self.markers[outside_inds, :])[0] + b_cart = self.equil.b_cart(self.markers[:, outside_inds])[0] # calculate magnetic field amplitude and normalized magnetic field absB0 = xp.sqrt(b_cart[0] ** 2 + b_cart[1] ** 2 + b_cart[2] ** 2) @@ -3155,7 +3138,7 @@ def _gyro_transfer(self, outside_inds): new_v_perp = xp.cross(Larmor_r, norm_b_cart, axis=0) * v_perp_square - self.markers[outside_inds, 3:6] = (norm_b_cart * v_parallel).T + new_v_perp.T + self.markers[3:6, outside_inds] = norm_b_cart * v_parallel + new_v_perp return xp.logical_and(1.0 > gc_etas[0], gc_etas[0] > 0.0) @@ -3164,10 +3147,10 @@ def _sort_boxed_particles_numpy(self): sorting_axis = self._sorting_boxes.box_index if not hasattr(self, "_argsort_array"): - self._argsort_array = xp.zeros(self.markers.shape[0], dtype=int) - self._argsort_array[:] = self._markers[:, sorting_axis].argsort() + self._argsort_array = xp.zeros(self.n_rows, dtype=int) + self._argsort_array[:] = self._markers[sorting_axis].argsort() - self._markers[:, :] = self._markers[self._argsort_array] + self._markers[:] = self._markers[:, self._argsort_array] def _check_and_assign_particles_to_boxes(self): """Check whether the box array has enough columns (detect load imbalance wrt to sorting boxes), @@ -3176,11 +3159,11 @@ def _check_and_assign_particles_to_boxes(self): from cunumpy.xp import array_backend if array_backend.backend == "numpy": - bcount = xp.bincount(xp.int64(self.markers[~self.holes, -2])) + bcount = xp.bincount(xp.int64(self.markers[-2, ~self.holes])) else: import cupy as cp - indices = self.markers[~self.holes, -2] + indices = self.markers[-2, ~self.holes] indices = indices.astype(cp.int64) bcount = cp.bincount(indices) @@ -3208,15 +3191,15 @@ def _update_ghost_particles(self): as a ghost particle when its ID column (last column) equals -2, the marker set by :meth:`_prepare_ghost_particles`/:meth:`_sendrecv_markers_boxes` for SPH ghost-box particles received from a neighbouring process.""" - self._ghost_particles[:] = self.markers[:, -1] == -2.0 + self._ghost_particles[:] = self.markers[-1] == -2.0 self._update_valid_mks() def _remove_ghost_particles(self): - """Discard all current ghost particles: turn their marker-array rows into new + """Discard all current ghost particles: turn their marker-array columns into new holes (so the space can be reused before the next SPH ghost-box update).""" self._update_ghost_particles() - new_holes = xp.nonzero(self.ghost_particles) - self._markers[new_holes] = -1.0 + new_holes = xp.nonzero(self.ghost_particles)[0] + self._markers[:, new_holes] = -1.0 self.update_holes() def _prepare_ghost_particles(self): @@ -3241,20 +3224,20 @@ def _prepare_ghost_particles(self): self._markers_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_z_p) # Put last index to -2 to indicate that they are ghosts on the new process - self._markers_x_m[:, -1] = -2.0 - self._markers_x_p[:, -1] = -2.0 - self._markers_y_m[:, -1] = -2.0 - self._markers_y_p[:, -1] = -2.0 - self._markers_z_m[:, -1] = -2.0 - self._markers_z_p[:, -1] = -2.0 + self._markers_x_m[-1] = -2.0 + self._markers_x_p[-1] = -2.0 + self._markers_y_m[-1] = -2.0 + self._markers_y_p[-1] = -2.0 + self._markers_z_m[-1] = -2.0 + self._markers_z_p[-1] = -2.0 # Adjust box number - self._markers_x_m[:, self._sorting_boxes.box_index] += shifts["x_m"] - self._markers_x_p[:, self._sorting_boxes.box_index] -= shifts["x_p"] - self._markers_y_m[:, self._sorting_boxes.box_index] += shifts["y_m"] - self._markers_y_p[:, self._sorting_boxes.box_index] -= shifts["y_p"] - self._markers_z_m[:, self._sorting_boxes.box_index] += shifts["z_m"] - self._markers_z_p[:, self._sorting_boxes.box_index] -= shifts["z_p"] + self._markers_x_m[self._sorting_boxes.box_index] += shifts["x_m"] + self._markers_x_p[self._sorting_boxes.box_index] -= shifts["x_p"] + self._markers_y_m[self._sorting_boxes.box_index] += shifts["y_m"] + self._markers_y_p[self._sorting_boxes.box_index] -= shifts["y_p"] + self._markers_z_m[self._sorting_boxes.box_index] += shifts["z_m"] + self._markers_z_p[self._sorting_boxes.box_index] -= shifts["z_p"] # Mirror position for boundary condition if self.bc_sph[0] in ("mirror", "fixed", "noslip"): @@ -3290,16 +3273,16 @@ def _prepare_ghost_particles(self): self._markers_x_p_y_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p) # Put last index to -2 to indicate that they are ghosts on the new process - self._markers_x_m_y_m[:, -1] = -2.0 - self._markers_x_m_y_p[:, -1] = -2.0 - self._markers_x_p_y_m[:, -1] = -2.0 - self._markers_x_p_y_p[:, -1] = -2.0 + self._markers_x_m_y_m[-1] = -2.0 + self._markers_x_m_y_p[-1] = -2.0 + self._markers_x_p_y_m[-1] = -2.0 + self._markers_x_p_y_p[-1] = -2.0 # Adjust box number - self._markers_x_m_y_m[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] - self._markers_x_m_y_p[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] - self._markers_x_p_y_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] - self._markers_x_p_y_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] + self._markers_x_m_y_m[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] + self._markers_x_m_y_p[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] + self._markers_x_p_y_m[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] + self._markers_x_p_y_p[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] # Mirror position for boundary condition if self.bc_sph[0] in ("mirror", "fixed", "noslip") or self.bc_sph[1] in ("mirror", "fixed", "noslip"): @@ -3321,16 +3304,16 @@ def _prepare_ghost_particles(self): self._markers_x_p_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_z_p) # Put last index to -2 to indicate that they are ghosts on the new process - self._markers_x_m_z_m[:, -1] = -2.0 - self._markers_x_m_z_p[:, -1] = -2.0 - self._markers_x_p_z_m[:, -1] = -2.0 - self._markers_x_p_z_p[:, -1] = -2.0 + self._markers_x_m_z_m[-1] = -2.0 + self._markers_x_m_z_p[-1] = -2.0 + self._markers_x_p_z_m[-1] = -2.0 + self._markers_x_p_z_p[-1] = -2.0 # Adjust box number - self._markers_x_m_z_m[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["z_m"] - self._markers_x_m_z_p[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["z_p"] - self._markers_x_p_z_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["z_m"] - self._markers_x_p_z_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["z_p"] + self._markers_x_m_z_m[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["z_m"] + self._markers_x_m_z_p[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["z_p"] + self._markers_x_p_z_m[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["z_m"] + self._markers_x_p_z_p[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["z_p"] # Mirror position for boundary condition if self.bc_sph[0] in ("mirror", "fixed", "noslip") or self.bc_sph[2] in ("mirror", "fixed", "noslip"): @@ -3352,16 +3335,16 @@ def _prepare_ghost_particles(self): self._markers_y_p_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_y_p_z_p) # Put last index to -2 to indicate that they are ghosts on the new process - self._markers_y_m_z_m[:, -1] = -2.0 - self._markers_y_m_z_p[:, -1] = -2.0 - self._markers_y_p_z_m[:, -1] = -2.0 - self._markers_y_p_z_p[:, -1] = -2.0 + self._markers_y_m_z_m[-1] = -2.0 + self._markers_y_m_z_p[-1] = -2.0 + self._markers_y_p_z_m[-1] = -2.0 + self._markers_y_p_z_p[-1] = -2.0 # Adjust box number - self._markers_y_m_z_m[:, self._sorting_boxes.box_index] += shifts["y_m"] + shifts["z_m"] - self._markers_y_m_z_p[:, self._sorting_boxes.box_index] += shifts["y_m"] - shifts["z_p"] - self._markers_y_p_z_m[:, self._sorting_boxes.box_index] += -shifts["y_p"] + shifts["z_m"] - self._markers_y_p_z_p[:, self._sorting_boxes.box_index] += -shifts["y_p"] - shifts["z_p"] + self._markers_y_m_z_m[self._sorting_boxes.box_index] += shifts["y_m"] + shifts["z_m"] + self._markers_y_m_z_p[self._sorting_boxes.box_index] += shifts["y_m"] - shifts["z_p"] + self._markers_y_p_z_m[self._sorting_boxes.box_index] += -shifts["y_p"] + shifts["z_m"] + self._markers_y_p_z_p[self._sorting_boxes.box_index] += -shifts["y_p"] - shifts["z_p"] # Mirror position for boundary condition if self.bc_sph[1] in ("mirror", "fixed", "noslip") or self.bc_sph[2] in ("mirror", "fixed", "noslip"): @@ -3387,24 +3370,24 @@ def _prepare_ghost_particles(self): self._markers_x_p_y_p_z_p = self._determine_markers_in_box(self._sorting_boxes._bnd_boxes_x_p_y_p_z_p) # Put last index to -2 to indicate that they are ghosts on the new process - self._markers_x_m_y_m_z_m[:, -1] = -2.0 - self._markers_x_m_y_m_z_p[:, -1] = -2.0 - self._markers_x_m_y_p_z_m[:, -1] = -2.0 - self._markers_x_m_y_p_z_p[:, -1] = -2.0 - self._markers_x_p_y_m_z_m[:, -1] = -2.0 - self._markers_x_p_y_m_z_p[:, -1] = -2.0 - self._markers_x_p_y_p_z_m[:, -1] = -2.0 - self._markers_x_p_y_p_z_p[:, -1] = -2.0 + self._markers_x_m_y_m_z_m[-1] = -2.0 + self._markers_x_m_y_m_z_p[-1] = -2.0 + self._markers_x_m_y_p_z_m[-1] = -2.0 + self._markers_x_m_y_p_z_p[-1] = -2.0 + self._markers_x_p_y_m_z_m[-1] = -2.0 + self._markers_x_p_y_m_z_p[-1] = -2.0 + self._markers_x_p_y_p_z_m[-1] = -2.0 + self._markers_x_p_y_p_z_p[-1] = -2.0 # Adjust box number - self._markers_x_m_y_m_z_m[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] + shifts["z_m"] - self._markers_x_m_y_m_z_p[:, self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] - shifts["z_p"] - self._markers_x_m_y_p_z_m[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] + shifts["z_m"] - self._markers_x_m_y_p_z_p[:, self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] - shifts["z_p"] - self._markers_x_p_y_m_z_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] + shifts["z_m"] - self._markers_x_p_y_m_z_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] - shifts["z_p"] - self._markers_x_p_y_p_z_m[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] + shifts["z_m"] - self._markers_x_p_y_p_z_p[:, self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] - shifts["z_p"] + self._markers_x_m_y_m_z_m[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] + shifts["z_m"] + self._markers_x_m_y_m_z_p[self._sorting_boxes.box_index] += shifts["x_m"] + shifts["y_m"] - shifts["z_p"] + self._markers_x_m_y_p_z_m[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] + shifts["z_m"] + self._markers_x_m_y_p_z_p[self._sorting_boxes.box_index] += shifts["x_m"] - shifts["y_p"] - shifts["z_p"] + self._markers_x_p_y_m_z_m[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] + shifts["z_m"] + self._markers_x_p_y_m_z_p[self._sorting_boxes.box_index] += -shifts["x_p"] + shifts["y_m"] - shifts["z_p"] + self._markers_x_p_y_p_z_m[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] + shifts["z_m"] + self._markers_x_p_y_p_z_p[self._sorting_boxes.box_index] += -shifts["x_p"] - shifts["y_p"] - shifts["z_p"] # Mirror position for boundary condition if any([bci in ("mirror", "fixed", "noslip") for bci in self.bc_sph]): @@ -3453,16 +3436,16 @@ def _mirror_particles( # x-direction if self.bc_sph[0] in ("mirror", "fixed", "noslip"): if "x_m" in arr_name and is_domain_boundary["x_m"]: - arr[:, 0] *= -1.0 + arr[0] *= -1.0 if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, + *arr[:3], flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! - arr[:, self.index["weights"]] = ( + arr[self.index["weights"]] = ( -boundary_values / self.s0( - *arr[:, :3].T, + *arr[:3], flat_eval=True, remove_holes=False, ) @@ -3471,25 +3454,25 @@ def _mirror_particles( self._fixed_markers_set[arr_name] = True elif self.bc_sph[0] == "noslip": # invert the velocities to have zero velocity at the boundary - arr[:, 3] *= -1.0 - arr[:, 4] *= -1.0 - arr[:, 5] *= -1.0 + arr[3] *= -1.0 + arr[4] *= -1.0 + arr[5] *= -1.0 if mean_velocity_index is not None: - arr[:, mean_velocity_index] *= -1.0 - arr[:, mean_velocity_index + 1] *= -1.0 - arr[:, mean_velocity_index + 2] *= -1.0 + arr[mean_velocity_index] *= -1.0 + arr[mean_velocity_index + 1] *= -1.0 + arr[mean_velocity_index + 2] *= -1.0 elif "x_p" in arr_name and is_domain_boundary["x_p"]: - arr[:, 0] = 2.0 - arr[:, 0] + arr[0] = 2.0 - arr[0] if self.bc_sph[0] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, + *arr[:3], flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! - arr[:, self.index["weights"]] = ( + arr[self.index["weights"]] = ( -boundary_values / self.s0( - *arr[:, :3].T, + *arr[:3], flat_eval=True, remove_holes=False, ) @@ -3498,27 +3481,27 @@ def _mirror_particles( self._fixed_markers_set[arr_name] = True elif self.bc_sph[0] == "noslip": # invert the velocities to have zero velocity at the boundary - arr[:, 3] *= -1.0 - arr[:, 4] *= -1.0 - arr[:, 5] *= -1.0 + arr[3] *= -1.0 + arr[4] *= -1.0 + arr[5] *= -1.0 if mean_velocity_index is not None: - arr[:, mean_velocity_index] *= -1.0 - arr[:, mean_velocity_index + 1] *= -1.0 - arr[:, mean_velocity_index + 2] *= -1.0 + arr[mean_velocity_index] *= -1.0 + arr[mean_velocity_index + 1] *= -1.0 + arr[mean_velocity_index + 2] *= -1.0 # y-direction if self.bc_sph[1] in ("mirror", "fixed", "noslip"): if "y_m" in arr_name and is_domain_boundary["y_m"]: - arr[:, 1] *= -1.0 + arr[1] *= -1.0 if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, + *arr[:3], flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! - arr[:, self.index["weights"]] = ( + arr[self.index["weights"]] = ( -boundary_values / self.s0( - *arr[:, :3].T, + *arr[:3], flat_eval=True, remove_holes=False, ) @@ -3527,25 +3510,25 @@ def _mirror_particles( self._fixed_markers_set[arr_name] = True elif self.bc_sph[1] == "noslip": # invert the velocities to have zero velocity at the boundary - arr[:, 3] *= -1.0 - arr[:, 4] *= -1.0 - arr[:, 5] *= -1.0 + arr[3] *= -1.0 + arr[4] *= -1.0 + arr[5] *= -1.0 if mean_velocity_index is not None: - arr[:, mean_velocity_index] *= -1.0 - arr[:, mean_velocity_index + 1] *= -1.0 - arr[:, mean_velocity_index + 2] *= -1.0 + arr[mean_velocity_index] *= -1.0 + arr[mean_velocity_index + 1] *= -1.0 + arr[mean_velocity_index + 2] *= -1.0 elif "y_p" in arr_name and is_domain_boundary["y_p"]: - arr[:, 1] = 2.0 - arr[:, 1] + arr[1] = 2.0 - arr[1] if self.bc_sph[1] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, + *arr[:3], flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! - arr[:, self.index["weights"]] = ( + arr[self.index["weights"]] = ( -boundary_values / self.s0( - *arr[:, :3].T, + *arr[:3], flat_eval=True, remove_holes=False, ) @@ -3554,27 +3537,27 @@ def _mirror_particles( self._fixed_markers_set[arr_name] = True elif self.bc_sph[1] == "noslip": # invert the velocities to have zero velocity at the boundary - arr[:, 3] *= -1.0 - arr[:, 4] *= -1.0 - arr[:, 5] *= -1.0 + arr[3] *= -1.0 + arr[4] *= -1.0 + arr[5] *= -1.0 if mean_velocity_index is not None: - arr[:, mean_velocity_index] *= -1.0 - arr[:, mean_velocity_index + 1] *= -1.0 - arr[:, mean_velocity_index + 2] *= -1.0 + arr[mean_velocity_index] *= -1.0 + arr[mean_velocity_index + 1] *= -1.0 + arr[mean_velocity_index + 2] *= -1.0 # z-direction if self.bc_sph[2] in ("mirror", "fixed", "noslip"): if "z_m" in arr_name and is_domain_boundary["z_m"]: - arr[:, 2] *= -1.0 + arr[2] *= -1.0 if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, + *arr[:3], flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! - arr[:, self.index["weights"]] = ( + arr[self.index["weights"]] = ( -boundary_values / self.s0( - *arr[:, :3].T, + *arr[:3], flat_eval=True, remove_holes=False, ) @@ -3583,25 +3566,25 @@ def _mirror_particles( self._fixed_markers_set[arr_name] = True elif self.bc_sph[2] == "noslip": # invert the velocities to have zero velocity at the boundary - arr[:, 3] *= -1.0 - arr[:, 4] *= -1.0 - arr[:, 5] *= -1.0 + arr[3] *= -1.0 + arr[4] *= -1.0 + arr[5] *= -1.0 if mean_velocity_index is not None: - arr[:, mean_velocity_index] *= -1.0 - arr[:, mean_velocity_index + 1] *= -1.0 - arr[:, mean_velocity_index + 2] *= -1.0 + arr[mean_velocity_index] *= -1.0 + arr[mean_velocity_index + 1] *= -1.0 + arr[mean_velocity_index + 2] *= -1.0 elif "z_p" in arr_name and is_domain_boundary["z_p"]: - arr[:, 2] = 2.0 - arr[:, 2] + arr[2] = 2.0 - arr[2] if self.bc_sph[2] == "fixed" and arr_name not in self._fixed_markers_set: boundary_values = self.f_init( - *arr[:, :3].T, + *arr[:3], flat_eval=True, ) # evaluation outside of the unit cube - maybe not working for all f_init! - arr[:, self.index["weights"]] = ( + arr[self.index["weights"]] = ( -boundary_values / self.s0( - *arr[:, :3].T, + *arr[:3], flat_eval=True, remove_holes=False, ) @@ -3610,13 +3593,13 @@ def _mirror_particles( self._fixed_markers_set[arr_name] = True elif self.bc_sph[2] == "noslip": # invert the velocities to have zero velocity at the boundary - arr[:, 3] *= -1.0 - arr[:, 4] *= -1.0 - arr[:, 5] *= -1.0 + arr[3] *= -1.0 + arr[4] *= -1.0 + arr[5] *= -1.0 if mean_velocity_index is not None: - arr[:, mean_velocity_index] *= -1.0 - arr[:, mean_velocity_index + 1] *= -1.0 - arr[:, mean_velocity_index + 2] *= -1.0 + arr[mean_velocity_index] *= -1.0 + arr[mean_velocity_index + 1] *= -1.0 + arr[mean_velocity_index + 2] *= -1.0 def _determine_markers_in_box(self, list_boxes): """Gather the markers currently sorted into any of the given boxes into a new array @@ -3632,14 +3615,16 @@ def _determine_markers_in_box(self, list_boxes): Returns ------- markers_in_box : xp.ndarray - Copy of the marker-array rows belonging to any of ``list_boxes``. + Copy of the marker-array columns (markers) belonging to any of ``list_boxes``. """ indices = [] for i in list_boxes: indices += list(self._sorting_boxes._boxes[i][self._sorting_boxes._boxes[i] != -1]) indices = xp.array(indices, dtype=int) - markers_in_box = self.markers[indices] + # C-contiguous, as the receive buffers (n_cols, n) of _sendrecv_markers_boxes: selecting columns gives a + # Fortran-ordered copy, which MPI would send marker by marker + markers_in_box = xp.ascontiguousarray(self.markers[:, indices]) return markers_in_box def _get_destinations_box(self): @@ -3650,155 +3635,187 @@ def _get_destinations_box(self): (:attr:`_send_list_box`, used by :meth:`_self_communication_boxes` and :meth:`_sendrecv_markers_boxes`).""" self._send_info_box = xp.zeros(self.mpi_size, dtype=int) - self._send_list_box = [xp.zeros((0, self.n_cols))] * self.mpi_size + self._send_list_box = [xp.zeros((self.n_cols, 0))] * self.mpi_size # Faces # if self._x_m_proc is not None: - self._send_info_box[self._x_m_proc] += len(self._markers_x_m) - self._send_list_box[self._x_m_proc] = xp.concatenate((self._send_list_box[self._x_m_proc], self._markers_x_m)) + self._send_info_box[self._x_m_proc] += self._markers_x_m.shape[1] + self._send_list_box[self._x_m_proc] = xp.concatenate( + (self._send_list_box[self._x_m_proc], self._markers_x_m), axis=1 + ) # if self._x_p_proc is not None: - self._send_info_box[self._x_p_proc] += len(self._markers_x_p) - self._send_list_box[self._x_p_proc] = xp.concatenate((self._send_list_box[self._x_p_proc], self._markers_x_p)) + self._send_info_box[self._x_p_proc] += self._markers_x_p.shape[1] + self._send_list_box[self._x_p_proc] = xp.concatenate( + (self._send_list_box[self._x_p_proc], self._markers_x_p), axis=1 + ) # if self._y_m_proc is not None: - self._send_info_box[self._y_m_proc] += len(self._markers_y_m) - self._send_list_box[self._y_m_proc] = xp.concatenate((self._send_list_box[self._y_m_proc], self._markers_y_m)) + self._send_info_box[self._y_m_proc] += self._markers_y_m.shape[1] + self._send_list_box[self._y_m_proc] = xp.concatenate( + (self._send_list_box[self._y_m_proc], self._markers_y_m), axis=1 + ) # if self._y_p_proc is not None: - self._send_info_box[self._y_p_proc] += len(self._markers_y_p) - self._send_list_box[self._y_p_proc] = xp.concatenate((self._send_list_box[self._y_p_proc], self._markers_y_p)) + self._send_info_box[self._y_p_proc] += self._markers_y_p.shape[1] + self._send_list_box[self._y_p_proc] = xp.concatenate( + (self._send_list_box[self._y_p_proc], self._markers_y_p), axis=1 + ) # if self._z_m_proc is not None: - self._send_info_box[self._z_m_proc] += len(self._markers_z_m) - self._send_list_box[self._z_m_proc] = xp.concatenate((self._send_list_box[self._z_m_proc], self._markers_z_m)) + self._send_info_box[self._z_m_proc] += self._markers_z_m.shape[1] + self._send_list_box[self._z_m_proc] = xp.concatenate( + (self._send_list_box[self._z_m_proc], self._markers_z_m), axis=1 + ) # if self._z_p_proc is not None: - self._send_info_box[self._z_p_proc] += len(self._markers_z_p) - self._send_list_box[self._z_p_proc] = xp.concatenate((self._send_list_box[self._z_p_proc], self._markers_z_p)) + self._send_info_box[self._z_p_proc] += self._markers_z_p.shape[1] + self._send_list_box[self._z_p_proc] = xp.concatenate( + (self._send_list_box[self._z_p_proc], self._markers_z_p), axis=1 + ) # x-y edges # if self._x_m_y_m_proc is not None: - self._send_info_box[self._x_m_y_m_proc] += len(self._markers_x_m_y_m) + self._send_info_box[self._x_m_y_m_proc] += self._markers_x_m_y_m.shape[1] self._send_list_box[self._x_m_y_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_m_proc], self._markers_x_m_y_m), + axis=1, ) # if self._x_m_y_p_proc is not None: - self._send_info_box[self._x_m_y_p_proc] += len(self._markers_x_m_y_p) + self._send_info_box[self._x_m_y_p_proc] += self._markers_x_m_y_p.shape[1] self._send_list_box[self._x_m_y_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_p_proc], self._markers_x_m_y_p), + axis=1, ) # if self._x_p_y_m_proc is not None: - self._send_info_box[self._x_p_y_m_proc] += len(self._markers_x_p_y_m) + self._send_info_box[self._x_p_y_m_proc] += self._markers_x_p_y_m.shape[1] self._send_list_box[self._x_p_y_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_m_proc], self._markers_x_p_y_m), + axis=1, ) # if self._x_p_y_p_proc is not None: - self._send_info_box[self._x_p_y_p_proc] += len(self._markers_x_p_y_p) + self._send_info_box[self._x_p_y_p_proc] += self._markers_x_p_y_p.shape[1] self._send_list_box[self._x_p_y_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_p_proc], self._markers_x_p_y_p), + axis=1, ) # x-z edges # if self._x_m_z_m_proc is not None: - self._send_info_box[self._x_m_z_m_proc] += len(self._markers_x_m_z_m) + self._send_info_box[self._x_m_z_m_proc] += self._markers_x_m_z_m.shape[1] self._send_list_box[self._x_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_z_m_proc], self._markers_x_m_z_m), + axis=1, ) # if self._x_m_z_p_proc is not None: - self._send_info_box[self._x_m_z_p_proc] += len(self._markers_x_m_z_p) + self._send_info_box[self._x_m_z_p_proc] += self._markers_x_m_z_p.shape[1] self._send_list_box[self._x_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_z_p_proc], self._markers_x_m_z_p), + axis=1, ) # if self._x_p_z_m_proc is not None: - self._send_info_box[self._x_p_z_m_proc] += len(self._markers_x_p_z_m) + self._send_info_box[self._x_p_z_m_proc] += self._markers_x_p_z_m.shape[1] self._send_list_box[self._x_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_z_m_proc], self._markers_x_p_z_m), + axis=1, ) # if self._x_p_z_p_proc is not None: - self._send_info_box[self._x_p_z_p_proc] += len(self._markers_x_p_z_p) + self._send_info_box[self._x_p_z_p_proc] += self._markers_x_p_z_p.shape[1] self._send_list_box[self._x_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_z_p_proc], self._markers_x_p_z_p), + axis=1, ) # y-z edges # if self._y_m_z_m_proc is not None: - self._send_info_box[self._y_m_z_m_proc] += len(self._markers_y_m_z_m) + self._send_info_box[self._y_m_z_m_proc] += self._markers_y_m_z_m.shape[1] self._send_list_box[self._y_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._y_m_z_m_proc], self._markers_y_m_z_m), + axis=1, ) # if self._y_m_z_p_proc is not None: - self._send_info_box[self._y_m_z_p_proc] += len(self._markers_y_m_z_p) + self._send_info_box[self._y_m_z_p_proc] += self._markers_y_m_z_p.shape[1] self._send_list_box[self._y_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._y_m_z_p_proc], self._markers_y_m_z_p), + axis=1, ) # if self._y_p_z_m_proc is not None: - self._send_info_box[self._y_p_z_m_proc] += len(self._markers_y_p_z_m) + self._send_info_box[self._y_p_z_m_proc] += self._markers_y_p_z_m.shape[1] self._send_list_box[self._y_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._y_p_z_m_proc], self._markers_y_p_z_m), + axis=1, ) # if self._y_p_z_p_proc is not None: - self._send_info_box[self._y_p_z_p_proc] += len(self._markers_y_p_z_p) + self._send_info_box[self._y_p_z_p_proc] += self._markers_y_p_z_p.shape[1] self._send_list_box[self._y_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._y_p_z_p_proc], self._markers_y_p_z_p), + axis=1, ) # corners # if self._x_m_y_m_z_m_proc is not None: - self._send_info_box[self._x_m_y_m_z_m_proc] += len(self._markers_x_m_y_m_z_m) + self._send_info_box[self._x_m_y_m_z_m_proc] += self._markers_x_m_y_m_z_m.shape[1] self._send_list_box[self._x_m_y_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_m_z_m_proc], self._markers_x_m_y_m_z_m), + axis=1, ) # if self._x_m_y_m_z_p_proc is not None: - self._send_info_box[self._x_m_y_m_z_p_proc] += len(self._markers_x_m_y_m_z_p) + self._send_info_box[self._x_m_y_m_z_p_proc] += self._markers_x_m_y_m_z_p.shape[1] self._send_list_box[self._x_m_y_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_m_z_p_proc], self._markers_x_m_y_m_z_p), + axis=1, ) # if self._x_m_y_p_z_m_proc is not None: - self._send_info_box[self._x_m_y_p_z_m_proc] += len(self._markers_x_m_y_p_z_m) + self._send_info_box[self._x_m_y_p_z_m_proc] += self._markers_x_m_y_p_z_m.shape[1] self._send_list_box[self._x_m_y_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_p_z_m_proc], self._markers_x_m_y_p_z_m), + axis=1, ) # if self._x_m_y_p_z_p_proc is not None: - self._send_info_box[self._x_m_y_p_z_p_proc] += len(self._markers_x_m_y_p_z_p) + self._send_info_box[self._x_m_y_p_z_p_proc] += self._markers_x_m_y_p_z_p.shape[1] self._send_list_box[self._x_m_y_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_m_y_p_z_p_proc], self._markers_x_m_y_p_z_p), + axis=1, ) # if self._x_p_y_m_z_m_proc is not None: - self._send_info_box[self._x_p_y_m_z_m_proc] += len(self._markers_x_p_y_m_z_m) + self._send_info_box[self._x_p_y_m_z_m_proc] += self._markers_x_p_y_m_z_m.shape[1] self._send_list_box[self._x_p_y_m_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_m_z_m_proc], self._markers_x_p_y_m_z_m), + axis=1, ) # if self._x_p_y_m_z_p_proc is not None: - self._send_info_box[self._x_p_y_m_z_p_proc] += len(self._markers_x_p_y_m_z_p) + self._send_info_box[self._x_p_y_m_z_p_proc] += self._markers_x_p_y_m_z_p.shape[1] self._send_list_box[self._x_p_y_m_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_m_z_p_proc], self._markers_x_p_y_m_z_p), + axis=1, ) # if self._x_p_y_p_z_m_proc is not None: - self._send_info_box[self._x_p_y_p_z_m_proc] += len(self._markers_x_p_y_p_z_m) + self._send_info_box[self._x_p_y_p_z_m_proc] += self._markers_x_p_y_p_z_m.shape[1] self._send_list_box[self._x_p_y_p_z_m_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_p_z_m_proc], self._markers_x_p_y_p_z_m), + axis=1, ) # if self._x_p_y_p_z_p_proc is not None: - self._send_info_box[self._x_p_y_p_z_p_proc] += len(self._markers_x_p_y_p_z_p) + self._send_info_box[self._x_p_y_p_z_p_proc] += self._markers_x_p_y_p_z_p.shape[1] self._send_list_box[self._x_p_y_p_z_p_proc] = xp.concatenate( (self._send_list_box[self._x_p_y_p_z_p_proc], self._markers_x_p_y_p_z_p), + axis=1, ) def _self_communication_boxes(self): @@ -3824,14 +3841,16 @@ def _self_communication_boxes(self): # self._bufsize *= 2.0 # self._allocate_marker_array() # logger.info(f"new: {self.markers.shape = }\n") - # self.markers[:] = -1.0 - # self.markers[:_n_rows_old] = _tmp + # self.markers[:, :] = -1.0 + # self.markers[:, :_n_rows_old] = _tmp # self.update_holes() # self._update_ghost_particles() # self._update_valid_mks() # holes_inds = xp.nonzero(self.holes)[0] - self.markers[holes_inds[xp.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[self.mpi_rank] + self.markers[:, holes_inds[xp.arange(self._send_info_box[self.mpi_rank])]] = self._send_list_box[ + self.mpi_rank + ] def _sendrecv_all_to_all_boxes(self): """ @@ -3863,7 +3882,7 @@ def _sendrecv_markers_boxes(self): else: send_reqs += [self.mpi_comm.Isend(data, dest=i, tag=self.mpi_comm.Get_rank())] - recvbufs += [xp.zeros((N_recv, self._markers.shape[1]), dtype=float)] + recvbufs += [xp.zeros((self.n_cols, N_recv), dtype=float)] reqs += [self.mpi_comm.Irecv(recvbufs[-1], source=i, tag=i)] # Wait for buffer, then put markers into holes @@ -3886,7 +3905,7 @@ def _sendrecv_markers_boxes(self): self.mpi_comm.Abort() # exit() - self._markers[hole_inds[first_hole[i] + xp.arange(self._recv_info_box[i])]] = recvbufs[i] + self._markers[:, hole_inds[first_hole[i] + xp.arange(self._recv_info_box[i])]] = recvbufs[i] test_reqs.pop() reqs[i] = None @@ -4247,7 +4266,7 @@ def _communicate_boxes(self): (:meth:`_get_destinations_box`), deliver the ones staying on this process (:meth:`_self_communication_boxes`), and, if running under MPI, exchange the rest with neighbouring processes (:meth:`_sendrecv_all_to_all_boxes`, - :meth:`_sendrecv_markers_boxes`) before marking the received rows as ghost + :meth:`_sendrecv_markers_boxes`) before marking the received markers as ghost particles (:meth:`_update_ghost_particles`).""" # if verbose: # n_valid = xp.count_nonzero(self.valid_mks) @@ -4292,7 +4311,7 @@ def _eval_sph( \rho(\boldsymbol \eta_i) = \sum_{j=0}^{N-1} \rho_j\, W_h(\boldsymbol \eta_i - \boldsymbol \eta_j)\,. - The coefficients :math:`\rho_j` must be available in the marker array, stored at some index ``self.markers[j, index]``. + The coefficients :math:`\rho_j` must be available in the marker array, stored at some index ``self.markers[index, j]``. In case that `derivative=k` where `k` is not zero, the `k`-th component of the gradient of :math:`\rho` is computed: .. math:: @@ -4417,7 +4436,7 @@ def _is_on_domain_of_rank(self, etas, rank): Parameters ---------- etas : array[float] - Logical positions of shape (n, 3). + Logical positions of shape (3, n), one column per position. rank : int Process rank (row index of domain_array). @@ -4425,10 +4444,10 @@ def _is_on_domain_of_rank(self, etas, rank): Returns ------- is_on_domain : array[bool] - Array of shape (n, 3), True where the position lies in the process domain in that direction. + Array of shape (3, n), True where the position lies in the process domain in that direction. """ - left = self.domain_array[rank, 0::3] - right = self.domain_array[rank, 1::3] + left = self.domain_array[rank, 0::3, None] + right = self.domain_array[rank, 1::3, None] # include eta = 1 in the last process: shift its right boundary to the next float right = xp.where(right == 1.0, xp.nextafter(right, 2.0), right) return xp.logical_and(etas >= left, etas < right) @@ -4439,13 +4458,13 @@ def _sendrecv_determine_mtbs( ): """ Determine which markers have to be sent from current process and put them in a new array. - Corresponding rows in markers array become holes and are therefore set to -1. + Corresponding columns in markers array become holes and are therefore set to -1. This can be done purely with numpy functions (fast, vectorized). Parameters ---------- alpha : list | tuple - For i=1,2,3 the sorting is according to alpha[i]*markers[:, i] + (1 - alpha[i])*markers[:, first_pusher_idx + i]. + For i=1,2,3 the sorting is according to alpha[i]*markers[i] + (1 - alpha[i])*markers[first_pusher_idx + i]. alpha[i] must be between 0 and 1. Returns @@ -4454,7 +4473,7 @@ def _sendrecv_determine_mtbs( Indices of empty columns in markers after send. sorting_etas : array[float] - Eta-values of shape (n_send, :) according to which the sorting is performed. + Eta-values of shape (3, n_rows) according to which the sorting is performed. """ # position that determines the sorting (including periodic shift of boundary conditions) if not isinstance(alpha, xp.ndarray): @@ -4467,39 +4486,38 @@ def _sendrecv_determine_mtbs( # which needs no host-to-device copy on CuPy and gives the same values on NumPy alpha = float(alpha_host[0]) else: - alpha = xp.asarray(alpha_host) + alpha = xp.asarray(alpha_host)[:, None] else: assert alpha.size == 3 assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) + alpha = alpha.reshape(3, 1) bi = self.first_pusher_idx self._sorting_etas[:] = ( - alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) - + (1.0 - alpha) * self.markers[:, bi : bi + 3] + alpha * (self.markers[:3] + self.markers[bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) + + (1.0 - alpha) * self.markers[bi : bi + 3] ) # for non-periodic axes, eta = 1 is the right end of the domain and must not be wrapped to eta = 0 non_periodic_axes = [axis for axis in range(3) if axis not in self._periodic_axes] - at_right_end = {axis: self._sorting_etas[:, axis] == 1.0 for axis in non_periodic_axes} + at_right_end = {axis: self._sorting_etas[axis] == 1.0 for axis in non_periodic_axes} xp.mod(self._sorting_etas, 1.0, out=self._sorting_etas) for axis, mask in at_right_end.items(): - self._sorting_etas[mask, axis] = 1.0 + self._sorting_etas[axis, mask] = 1.0 # check which particles are on the current process domain self._is_on_proc_domain = self._is_on_domain_of_rank(self._sorting_etas, self.mpi_rank) - # to stay on the current process, all three columns must be True. - # Reducing over a size-3 trailing axis of an array with many rows is slow - # This is faster - # Improvement of approximately 10x (both numpy and cupy) - self._can_stay = self._is_on_proc_domain[:, 0] & self._is_on_proc_domain[:, 1] & self._is_on_proc_domain[:, 2] + # to stay on the current process, all three directions must be True. + # Combining the three rows is faster than reducing over the axis of size 3 + self._can_stay = self._is_on_proc_domain[0] & self._is_on_proc_domain[1] & self._is_on_proc_domain[2] # holes and ghosts can stay, too self._can_stay[self.holes] = True self._can_stay[self.ghost_particles] = True - # True values can stay on the process, False must be sent, already empty rows (-1) cannot be sent + # True values can stay on the process, False must be sent, already empty markers (-1) cannot be sent send_inds = xp.nonzero(~self._can_stay)[0] hole_inds_after_send = xp.nonzero(xp.logical_or(~self._can_stay, self.holes))[0] @@ -4577,10 +4595,10 @@ def _sendrecv_get_destinations(self, send_inds): send_info = np.zeros(self.mpi_size, dtype=int) # Gathered once and reused for every rank below, instead of re-gathering - # self.markers[send_inds] and self._sorting_etas[send_inds] fresh on every + # self.markers[:, send_inds] and self._sorting_etas[:, send_inds] fresh on every # iteration of the rank loop (as the previous version did). - candidates = self.markers[send_inds] - etas_to_send = self._sorting_etas[send_inds] + candidates = self.markers[:, send_inds] + etas_to_send = self._sorting_etas[:, send_inds] # Reset every rank's send buffer to empty first. The neighbour/non-neighbour # loop below can `break` before visiting every rank once all candidates are @@ -4589,10 +4607,10 @@ def _sendrecv_get_destinations(self, send_inds): # non-empty buffer from a previous call (send/recv size would then disagree # with send_info, which is always correct since it defaults to 0 above). empty_local = xp.empty(0, dtype=int) - empty_rows = candidates[:0] + empty_markers = candidates[:, :0] for i in range(self.mpi_size): self._send_to_i[i] = empty_local - self._send_list[i] = empty_rows + self._send_list[i] = empty_markers # A marker leaving this rank's domain is overwhelmingly likely to land in a # geometrically adjacent sub-domain -- one push (sub-)step is small compared to @@ -4607,17 +4625,19 @@ def _sendrecv_get_destinations(self, send_inds): if remaining.size == 0: break - etas_remaining = etas_to_send[remaining] + etas_remaining = etas_to_send[:, remaining] still_remaining = xp.ones(remaining.shape[0], dtype=bool) for i in rank_group: conds = self._is_on_domain_of_rank(etas_remaining, i) - matched_local = xp.nonzero(xp.all(conds, axis=1))[0] + matched_local = xp.nonzero(xp.all(conds, axis=0))[0] matched = remaining[matched_local] self._send_to_i[i] = matched send_info[i] = matched.size - self._send_list[i] = candidates[matched] + # C-contiguous, as the receive buffers (n_cols, n): selecting columns gives a Fortran-ordered copy, + # which MPI would send marker by marker + self._send_list[i] = xp.ascontiguousarray(candidates[:, matched]) still_remaining[matched_local] = False @@ -4643,7 +4663,7 @@ def _sendrecv_all_to_all(self, send_info): ----- ``send_info`` and ``recv_info`` are small NumPy arrays (one integer per rank) on every backend: the counts size the receive buffers and the holes to fill, which Python needs - on the host anyway. On CuPy, only the marker rows themselves are device buffers + on the host anyway. On CuPy, only the markers themselves are device buffers (:meth:`_sendrecv_markers`). """ @@ -4663,7 +4683,7 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): Amount of markers to be received from i-th process. hole_inds_after_send : array[int] - Indices of empty rows in markers after send. + Indices of empty columns (holes) in markers after send. """ if xp.get_array_backend(self._markers) == "cupy": @@ -4682,7 +4702,7 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): else: send_reqs += [self.mpi_comm.Isend(data, dest=i, tag=self.mpi_rank)] - self._recvbufs[i] = xp.zeros((N_recv, self.markers.shape[1]), dtype=float) + self._recvbufs[i] = xp.zeros((self.n_cols, N_recv), dtype=float) self._reqs[i] = self.mpi_comm.Irecv(self._recvbufs[i], source=i, tag=i) # Wait for buffer, then put markers into holes @@ -4704,7 +4724,9 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): ) self.mpi_comm.Abort() - self.markers[hole_inds_after_send[first_hole[i] + xp.arange(recv_info[i])]] = self._recvbufs[i] + self.markers[:, hole_inds_after_send[first_hole[i] + xp.arange(recv_info[i])]] = self._recvbufs[ + i + ] test_reqs.pop() self._reqs[i] = None @@ -4727,14 +4749,14 @@ def _mpi_cuda_aware(self) -> bool: if not cuda_aware and not Particles._warned_mpi_staging: Particles._warned_mpi_staging = True logger.warning( - "MPI is not CUDA-aware: marker rows exchanged between ranks are staged through host memory.", + "MPI is not CUDA-aware: markers exchanged between ranks are staged through host memory.", ) return cuda_aware _warned_mpi_staging = False def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): - """Exchange marker rows and scatter incoming rows into empty slots. + """Exchange markers and scatter the incoming ones into empty slots (holes). :func:`cunumpy.mpi.exchange` handles GPU synchronization, optional host staging, and completion of all MPI transfers. @@ -4744,7 +4766,7 @@ def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): recv_info : array[int] Number of markers to receive from each process (NumPy array). hole_inds_after_send : array[int] - Indices of empty rows in markers after send (device array). + Indices of empty columns (holes) in markers after send (device array). """ recv_info = np.asarray(recv_info) n_recv = int(recv_info.sum()) @@ -4759,8 +4781,9 @@ def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): ) self.mpi_comm.Abort() - recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=self.markers.dtype) - first_row = np.cumsum(recv_info) - recv_info + # one receive buffer (n_cols, n) per source rank: a slice of a common (n_cols, n_recv) buffer would not be + # contiguous + first_hole = np.cumsum(recv_info) - recv_info sends = [] receives = [] @@ -4769,12 +4792,11 @@ def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): continue if recv_info[rank] > 0: - start = int(first_row[rank]) - stop = start + int(recv_info[rank]) - receives.append((recvbuf[start:stop], rank, rank)) + recvbuf = xp.empty((self.n_cols, int(recv_info[rank])), dtype=self.markers.dtype) + receives.append((recvbuf, rank, rank)) data = self._send_list[rank] - if data.shape[0] > 0: + if data.shape[1] > 0: sends.append((data, rank, self.mpi_rank)) xp.mpi.exchange( @@ -4784,8 +4806,10 @@ def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): cuda_aware=self._mpi_cuda_aware(), ) - if n_recv > 0: - self._markers[hole_inds_after_send[:n_recv]] = recvbuf + for recvbuf, rank, _ in receives: + start = int(first_hole[rank]) + stop = start + int(recv_info[rank]) + self._markers[:, hole_inds_after_send[start:stop]] = recvbuf class Tesselation: diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index da77cd2a2..1defb9d5f 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -147,10 +147,7 @@ def save_constants_of_motion(self): idx_can_momentum = self.first_diagnostics_idx + 5 # save cartesian positions - self.markers[~self.holes, slice_gc] = self.domain( - self.positions, - change_out_order=True, - ) + self.markers[slice_gc, ~self.holes] = self.domain(self.positions) # eval guiding center phase space eval_guiding_center_from_6d( @@ -169,14 +166,11 @@ def save_constants_of_motion(self): # TODO: currently only possible with the geometry where its inverse map is defined. assert hasattr(self.domain, "inverse_map") - self.markers[~self.holes, slice_gc] = self.domain.inverse_map( - *self.markers[~self.holes, slice_gc].T, - change_out_order=True, - ) + self.markers[slice_gc, ~self.holes] = xp.stack(self.domain.inverse_map(*self.markers[slice_gc, ~self.holes])) # eval energy - self.markers[~self.holes, idx_energy] = ( - self.markers[~self.holes, 3] ** 2 + self.markers[~self.holes, 4] ** 2 + self.markers[~self.holes, 5] ** 2 + self.markers[idx_energy, ~self.holes] = ( + self.markers[3, ~self.holes] ** 2 + self.markers[4, ~self.holes] ** 2 + self.markers[5, ~self.holes] ** 2 ) / (2) # eval psi at etas @@ -184,13 +178,12 @@ def save_constants_of_motion(self): R0 = self.equil.params["R0"] B0 = self.equil.params["B0"] - r = self.markers[~self.holes, idx_gc_r] * (1 - a1) + a1 - self.markers[~self.holes, idx_can_momentum] = self.equil.psi_r(r) + r = self.markers[idx_gc_r, ~self.holes] * (1 - a1) + a1 + self.markers[idx_can_momentum, ~self.holes] = self.equil.psi_r(r) # send particles to the guiding center positions - self.markers[~self.holes, self.first_pusher_idx : self.first_pusher_idx + 3] = self.markers[ - ~self.holes, - slice_gc, + self.markers[self.first_pusher_idx : self.first_pusher_idx + 3, ~self.holes] = self.markers[ + slice_gc, ~self.holes ] if self.mpi_comm is not None: self.mpi_sort_markers(alpha=1) @@ -208,7 +201,7 @@ def save_constants_of_motion(self): # send back and clear buffer if self.mpi_comm is not None: self.mpi_sort_markers() - self.markers[~self.holes, self.first_pusher_idx : self.first_pusher_idx + 3] = 0 + self.markers[self.first_pusher_idx : self.first_pusher_idx + 3, ~self.holes] = 0 class DeltaFParticles6D(Particles6D): @@ -457,8 +450,8 @@ def save_constants_of_motion(self): R0 = self.equil.params["R0"] B0 = self.equil.params["B0"] - r = self.markers[~self.holes, 0] * (1 - a1) + a1 - self.markers[~self.holes, idx_can_momentum] = self.equil.psi_r(r) + r = self.markers[0, ~self.holes] * (1 - a1) + a1 + self.markers[idx_can_momentum, ~self.holes] = self.equil.psi_r(r) eval_canonical_toroidal_moment_5d( self.markers, @@ -742,8 +735,8 @@ def save_constants_of_motion(self): R0 = self.equil.params["R0"] B0 = self.equil.params["B0"] - r = self.markers[~self.holes, 0] * (1 - a1) + a1 - self.markers[~self.holes, idx_can_momentum] = self.equil.psi_r(r) + r = self.markers[0, ~self.holes] * (1 - a1) + a1 + self.markers[idx_can_momentum, ~self.holes] = self.equil.psi_r(r) eval_canonical_toroidal_moment_5d( self.markers, diff --git a/src/struphy/pic/sorting.py b/src/struphy/pic/sorting.py index 74db8c182..5c48e9219 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -224,7 +224,7 @@ def _set_boxes(self): self._n_boxes = (self._nx + 2) * (self._ny + 2) * (self._nz + 2) n_box_in = self._nx * self._ny * self._nz - n_particles = self._markers_shape[0] + n_particles = self._markers_shape[1] n_mkr = int(n_particles / n_box_in) + 1 n_cols = round( n_mkr * (1 + 1 / math.sqrt(n_mkr) + self._box_bufsize), @@ -240,8 +240,8 @@ def _set_boxes(self): initialize_neighbours(self._neighbours, self.nx, self.ny, self.nz) # logger.info(f"{self._rank = }\n{self._neighbours = }") - self._swap_line_1 = xp.zeros(self._markers_shape[1]) - self._swap_line_2 = xp.zeros(self._markers_shape[1]) + self._swap_line_1 = xp.zeros(self._markers_shape[0]) + self._swap_line_2 = xp.zeros(self._markers_shape[0]) def _set_boundary_boxes(self): """Collect the (flat) indices of all non-ghost boxes that lie on the outer surface diff --git a/src/struphy/pic/utilities.py b/src/struphy/pic/utilities.py index 229361d2e..fc3e4725b 100644 --- a/src/struphy/pic/utilities.py +++ b/src/struphy/pic/utilities.py @@ -3,7 +3,6 @@ from struphy.feec.utilities import get_quad_grids from struphy.pic.diagnostics.kernels.canonical_kinetic_particles import canonical_kinetic_particles from struphy.pic.diagnostics.kernels.thermal_energy import thermal_energy -from struphy.pic.marker_layout import as_marker_array def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): @@ -25,7 +24,7 @@ def get_kinetic_energy_particles(fe_coeffs, derham, domain, particles): res = xp.empty(1, dtype=float) canonical_kinetic_particles( res, - as_marker_array(particles.markers), + particles.markers, derham.args_derham, domain.args_domain, fe_coeffs.blocks[0]._data, From 9bf613f490df806ba7a1b64010a99e9acc057a07 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 15 Oct 2026 17:18:45 +0200 Subject: [PATCH 06/12] Marker evaluation takes (k, n) arrays, like the marker buffer Evaluation at markers (a single 2D array) now takes one column per point, eta1, eta2, eta3 in rows 0 to 2, so the marker buffer and its selections go in without a transpose: Domain (metric coefficients, pull/push/transform, inside_logical_cube, show), SplineFunction and _flag_pts_not_on_proc, the GVEC equilibria, the kinetic-background moments and B0 callables, and the canonical Maxwellian (coords[0] is (5, n)). Output shapes are unchanged. Domain and SplineFunction make the input C-contiguous for the kernels (no copy for the buffer or its rows). The user callbacks of the examples read etas[0][i]; orbit post-processing passes its (n, 3) arrays transposed. Co-Authored-By: Claude Opus 5.5 --- .../cyclone/params_cyclone.py | 6 +- .../itg_cylindre/params_drift_kinetic.py | 8 +- .../params_screw_pinch_equilibrium.py | 6 +- .../params_tokamak_equilibrium.py | 4 +- .../diocotron_instability/params_diocotron.py | 2 +- src/struphy/feec/psydac_derham.py | 23 ++-- src/struphy/fields_background/equils.py | 42 +++---- src/struphy/geometry/base.py | 104 ++++++++---------- src/struphy/kinetic_background/base.py | 13 +-- src/struphy/kinetic_background/maxwellians.py | 33 +++--- .../post_processing/orbits/orbits_tools.py | 2 +- 11 files changed, 111 insertions(+), 132 deletions(-) diff --git a/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py b/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py index 59e27d63e..c1eb7dcb9 100644 --- a/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py +++ b/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py @@ -199,7 +199,7 @@ def n_r(r): def n_init(*etas): if len(etas) == 1: - eta1 = etas[0][:, 0] + eta1 = etas[0][0] else: eta1 = etas[0] r = r_min + (a - r_min) * eta1 @@ -212,7 +212,7 @@ def Ti_r(r): def vth_init(*etas): if len(etas) == 1: - eta1 = etas[0][:, 0] + eta1 = etas[0][0] else: eta1 = etas[0] r = r_min + (a - r_min) * eta1 @@ -235,7 +235,7 @@ def p_xyz(x, y, z): def pert_func(*etas): if len(etas) == 1: - e1, e2, e3 = etas[0][:, 0], etas[0][:, 1], etas[0][:, 2] + e1, e2, e3 = etas[0][0], etas[0][1], etas[0][2] else: e1, e2, e3 = etas[0], etas[1], etas[2] r = (a - r_min) * e1 + r_min diff --git a/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py b/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py index 1e32cb509..2d6fbf40b 100644 --- a/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py +++ b/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py @@ -198,7 +198,7 @@ def n0(r): def n_init(*etas): if len(etas) == 1: - eta1 = etas[0][:, 0] + eta1 = etas[0][0] else: eta1 = etas[0] r = a1 + (a2 - a1) * eta1 @@ -207,7 +207,7 @@ def n_init(*etas): def pert_func(*etas): if len(etas) == 1: - e1, e2, e3 = etas[0][:, 0], etas[0][:, 1], etas[0][:, 2] + e1, e2, e3 = etas[0][0], etas[0][1], etas[0][2] else: e1, e2, e3 = etas[0], etas[1], etas[2] r = a1 + (a2 - a1) * e1 @@ -222,7 +222,7 @@ def Ti(r): def vth_i(*etas): if len(etas) == 1: - eta1 = etas[0][:, 0] + eta1 = etas[0][0] else: eta1 = etas[0] r = a1 + (a2 - a1) * eta1 @@ -231,7 +231,7 @@ def vth_i(*etas): def vth_e(*etas): if len(etas) == 1: - eta1 = etas[0][:, 0] + eta1 = etas[0][0] else: eta1 = etas[0] r = a1 + (a2 - a1) * eta1 diff --git a/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py b/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py index 6ec2bfada..8e731b108 100644 --- a/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py +++ b/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py @@ -163,15 +163,15 @@ def n_profile(*etas): - """Density n(r) of the screw pinch, r = a * eta1 (called with eta1, eta2, eta3 or one (n, 3) array).""" - eta1 = etas[0] if len(etas) == 3 else etas[0][:, 0] + """Density n(r) of the screw pinch, r = a * eta1 (called with eta1, eta2, eta3 or one (3, n) array).""" + eta1 = etas[0] if len(etas) == 3 else etas[0][0] r = a * eta1 return 0.2 + 0.8 * (1.0 - r**2 / a**2) def vth_profile(*etas): """Thermal speed sqrt(p(r)/n(r)) of the screw pinch.""" - eta1 = etas[0] if len(etas) == 3 else etas[0][:, 0] + eta1 = etas[0] if len(etas) == 3 else etas[0][0] r = a * eta1 q = 1.05 + (1.8 - 1.05) * r**2 / a**2 p = 0.005 + 1.0**2 * a**2 * 1.05 / (2 * R0**2 * (1.8 - 1.05)) * (1 / q**2 - 1 / 1.8**2) diff --git a/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py b/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py index 2f371fb76..9f2a2f591 100644 --- a/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py +++ b/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py @@ -195,13 +195,13 @@ def vth_canonical(psi_c): def n_position(*etas): """Density at the position, with the flux label s = eta1**2 (r = a * eta1, nearly constant q).""" - eta1 = etas[0] if len(etas) == 3 else etas[0][:, 0] + eta1 = etas[0] if len(etas) == 3 else etas[0][0] return n_of_s(eta1**2) def vth_position(*etas): """Thermal speed at the position, with the flux label s = eta1**2.""" - eta1 = etas[0] if len(etas) == 3 else etas[0][:, 0] + eta1 = etas[0] if len(etas) == 3 else etas[0][0] return vth_of_s(eta1**2) diff --git a/examples/ToyDrift/diocotron_instability/params_diocotron.py b/examples/ToyDrift/diocotron_instability/params_diocotron.py index aadc5322d..357081831 100644 --- a/examples/ToyDrift/diocotron_instability/params_diocotron.py +++ b/examples/ToyDrift/diocotron_instability/params_diocotron.py @@ -194,7 +194,7 @@ def n_init(etas, r_minus=r_minus, r_plus=r_plus): # transform logical coordinate to polar a1, a2 = domain.params["a1"], domain.params["a2"] - radial = a1 + (a2 - a1) * etas[:, 0] + radial = a1 + (a2 - a1) * etas[0] return 1.0 * ((r_minus <= radial) & (radial < r_plus)) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 5ae7d73f3..93dbe684e 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -49,7 +49,6 @@ from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.kernel_arguments.spline_args_kernels import SplineArguments -from struphy.pic.marker_layout import MARKER_ORDER from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator @@ -2901,7 +2900,7 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. out : array[float] or list @@ -2930,11 +2929,10 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): marker_evaluation = True is_sparse_meshgrid = False assert etas[0].ndim == 2 - # copy positions, such that flagging does not modify the caller's array; in the column-major marker - # layout that the kernel takes - markers = xp.array(etas[0][:, :3], dtype=float, order=MARKER_ORDER) + # copy positions, such that flagging does not modify the caller's array; C order, as the kernel takes + markers = xp.array(etas[0][:3], dtype=float, order="C") self._flag_pts_not_on_proc(markers) - tmp_shape = (markers.shape[0],) + tmp_shape = (markers.shape[1],) # 3D meshgrid evaluation else: marker_evaluation = False @@ -3104,7 +3102,7 @@ def _flag_pts_not_on_proc(self, *etas): *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes.""" # get domain decompoistion info @@ -3120,14 +3118,15 @@ def _flag_pts_not_on_proc(self, *etas): # check which particles are on the current process domain; # intervals are half-open [start, end) such that points on internal breaks are counted once - ends = dom_arr[rank, 1::3] + starts = dom_arr[rank, 0::3, None] + ends = dom_arr[rank, 1::3, None] is_on_proc_domain = xp.logical_and( - markers[:, :3] >= dom_arr[rank, 0::3], - xp.where(ends == 1.0, markers[:, :3] <= ends, markers[:, :3] < ends), + markers[:3] >= starts, + xp.where(ends == 1.0, markers[:3] <= ends, markers[:3] < ends), ) - on_proc = xp.all(is_on_proc_domain, axis=1) + on_proc = xp.all(is_on_proc_domain, axis=0) - xp.copyto(markers, -1.0, where=~on_proc[:, None]) + xp.copyto(markers, -1.0, where=~on_proc[None, :]) # 3D meshgrid evaluation else: diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index cde0e9ec7..5d2ac7b02 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -2178,9 +2178,9 @@ def n0(self, *etas, squeeze_out=False): # flat (marker) evaluation if len(etas) == 1: assert etas[0].ndim == 2 - eta1 = etas[0][:, 0] - eta2 = etas[0][:, 1] - eta3 = etas[0][:, 2] + eta1 = etas[0][0] + eta2 = etas[0][1] + eta3 = etas[0][2] flat_eval = True # meshgrid evaluation else: @@ -2215,9 +2215,9 @@ def _gvec_evaluations(self, *etas): # flat (marker) evaluation if len(etas) == 1: assert etas[0].ndim == 2 - eta1 = etas[0][:, 0] - eta2 = etas[0][:, 1] - eta3 = etas[0][:, 2] + eta1 = etas[0][0] + eta2 = etas[0][1] + eta3 = etas[0][2] flat_eval = True # meshgrid evaluation else: @@ -2434,9 +2434,9 @@ def _eval_bv(self, *etas, squeeze_out=False): # flat (marker) evaluation if len(etas) == 1: assert etas[0].ndim == 2 - eta1 = etas[0][:, 0] - eta2 = etas[0][:, 1] - eta3 = etas[0][:, 2] + eta1 = etas[0][0] + eta2 = etas[0][1] + eta3 = etas[0][2] flat_eval = True # meshgrid evaluation else: @@ -2510,9 +2510,9 @@ def _eval_jv(self, *etas, squeeze_out=False): # flat (marker) evaluation if len(etas) == 1: assert etas[0].ndim == 2 - eta1 = etas[0][:, 0] - eta2 = etas[0][:, 1] - eta3 = etas[0][:, 2] + eta1 = etas[0][0] + eta2 = etas[0][1] + eta3 = etas[0][2] flat_eval = True # meshgrid evaluation else: @@ -2552,9 +2552,9 @@ def p0(self, *etas, squeeze_out=False): # flat (marker) evaluation if len(etas) == 1: assert etas[0].ndim == 2 - eta1 = etas[0][:, 0] - eta2 = etas[0][:, 1] - eta3 = etas[0][:, 2] + eta1 = etas[0][0] + eta2 = etas[0][1] + eta3 = etas[0][2] flat_eval = True # meshgrid evaluation else: @@ -2584,9 +2584,9 @@ def n0(self, *etas, squeeze_out=False): # flat (marker) evaluation if len(etas) == 1: assert etas[0].ndim == 2 - eta1 = etas[0][:, 0] - eta2 = etas[0][:, 1] - eta3 = etas[0][:, 2] + eta1 = etas[0][0] + eta2 = etas[0][1] + eta3 = etas[0][2] flat_eval = True # meshgrid evaluation else: @@ -2640,9 +2640,9 @@ def _eval_gradB1(self, *etas, squeeze_out=False): # flat (marker) evaluation if len(etas) == 1: assert etas[0].ndim == 2 - eta1 = etas[0][:, 0] - eta2 = etas[0][:, 1] - eta3 = etas[0][:, 2] + eta1 = etas[0][0] + eta2 = etas[0][1] + eta3 = etas[0][2] flat_eval = True # meshgrid evaluation else: diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 67149166a..0ed8ebb90 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -19,7 +19,6 @@ from struphy.kernel_arguments.pusher_args_cuda import CudaDomainArguments from struphy.kernel_arguments.pusher_args_kernels import DomainArguments from struphy.linear_algebra import linalg_kron -from struphy.pic.marker_layout import as_marker_array from struphy.utils.cuda_arguments import check_mapping_on_device from struphy.utils.docstring_converter import rst_to_html, rst_to_latex, rst_to_markdown from struphy.utils.ipython_compat import HTML, display @@ -29,14 +28,14 @@ def inside_logical_cube(markers): - """Mask of the markers whose logical position ``markers[:, :3]`` lies in [0, 1]^3. + """Mask of the markers whose logical position ``markers[:3]`` lies in [0, 1]^3. The test of the pyccel and CUDA geometry kernels (``e1 < 0.0 or e1 > 1.0 or ...`` marks a marker as outside, so a - NaN coordinate counts as inside). The geometry kernels are called with ``remove_outside=False`` and the rows of + NaN coordinate counts as inside). The geometry kernels are called with ``remove_outside=False`` and the values of outside markers are removed with this mask, on both backends. """ - positions = markers[:, :3] - return ~xp.any((positions < 0.0) | (positions > 1.0), axis=1) + positions = markers[:3] + return ~xp.any((positions < 0.0) | (positions > 1.0), axis=0) class DomainMeta(ABCMeta): @@ -500,7 +499,7 @@ def __call__( ---------- *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. change_out_order : bool @@ -589,7 +588,7 @@ def jacobian_det( ---------- *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. squeeze_out : bool @@ -628,7 +627,7 @@ def jacobian_inv( ---------- *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. transposed : bool @@ -679,7 +678,7 @@ def metric( ---------- *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. transposed : bool @@ -730,7 +729,7 @@ def metric_inv( ---------- *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. transposed : bool @@ -786,7 +785,7 @@ def pull( *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. flat_eval : bool @@ -820,7 +819,7 @@ def pull( - ``kind in ('0', '3')`` (scalar p-form): no component axis is added. - - Markers / flat eval (``etas`` is a single ``(N, 3)`` array, or three 1d arrays of length ``N`` + - Markers / flat eval (``etas`` is a single ``(3, N)`` array, one column per point, or three 1d arrays of length ``N`` with ``flat_eval=True``): ``out`` has shape ``(n,)``, where ``n <= N`` if ``remove_outside=True`` removed points outside :math:`(0, 1)^3`, else ``n == N``. - Tensor-product eval (``etas`` is a tuple ``(eta1, eta2, eta3)`` broadcast to shape @@ -872,7 +871,7 @@ def push( *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. flat_eval : bool @@ -902,7 +901,7 @@ def push( - ``kind in ('0', '3')`` (scalar field): no component axis is added. - - Markers / flat eval (``etas`` is a single ``(N, 3)`` array, or three 1d arrays of length ``N`` + - Markers / flat eval (``etas`` is a single ``(3, N)`` array, one column per point, or three 1d arrays of length ``N`` with ``flat_eval=True``): ``out`` has shape ``(n,)``, where ``n <= N`` if ``remove_outside=True`` removed points outside :math:`(0, 1)^3`, else ``n == N``. - Tensor-product eval (``etas`` is a tuple ``(eta1, eta2, eta3)`` broadcast to shape @@ -953,7 +952,7 @@ def transform( *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. flat_eval : bool @@ -983,7 +982,7 @@ def transform( - ``kind in ('0_to_3', '3_to_0')`` (scalar-to-scalar): no component axis is added. - - Markers / flat eval (``etas`` is a single ``(N, 3)`` array, or three 1d arrays of length ``N`` + - Markers / flat eval (``etas`` is a single ``(3, N)`` array, one column per point, or three 1d arrays of length ``N`` with ``flat_eval=True``): ``out`` has shape ``(n,)``, where ``n <= N`` if ``remove_outside=True`` removed points outside :math:`(0, 1)^3`, else ``n == N``. - Tensor-product eval (``etas`` is a tuple ``(eta1, eta2, eta3)`` broadcast to shape @@ -1031,7 +1030,7 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): *etas : array-like | tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. which : int @@ -1055,11 +1054,11 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): # markers evaluation if len(etas) == 1: - # the kernel takes markers in the column-major marker layout (no copy for the marker buffer) - markers = as_marker_array(etas[0]) + # one column per marker (see struphy.pic.marker_layout); no copy for a contiguous array + markers = xp.ascontiguousarray(etas[0]) # to keep C-ordering the (3, 3)-part is in the last indices - out = xp.empty((markers.shape[0], 3, 3), dtype=float) + out = xp.empty((markers.shape[1], 3, 3), dtype=float) # one row per marker (remove_outside=False): a CUDA kernel cannot return the number of inside # markers, so the rows of outside markers are removed below, the same way on both backends kernel_evaluate_pic( @@ -1069,6 +1068,7 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): out, False, avoid_round_off, + n_threads=markers.shape[1], # CUDA: one thread per marker ) # move the (3, 3)-part to front @@ -1166,7 +1166,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa *etas : array-like| tuple Logical coordinates at which to evaluate. Two cases are possible: - 1. 2d numpy array, where coordinates are taken from eta1 = etas[:, 0], eta2 = etas[:, 1], etc. (like markers). + 1. 2d numpy array, where coordinates are taken from eta1 = etas[0], eta2 = etas[1], etc. (like markers, one column per point). 2. list/tuple (eta1, eta2, ...), where eta1, eta2, ... can be float or array-like of various shapes. flat_eval : bool @@ -1185,7 +1185,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa - ``kind_int < 10`` (scalar-valued result, e.g. kinds ``'0'``, ``'3'``, ``'0_to_3'``, ``'3_to_0'``): no component axis is added. - - Markers / flat eval (``etas`` is a single ``(N, 3)`` array, or three 1d arrays of length ``N`` + - Markers / flat eval (``etas`` is a single ``(3, N)`` array, one column per point, or three 1d arrays of length ``N`` with ``flat_eval=True``): ``out`` has shape ``(n,)``, where ``n <= N`` if ``remove_outside=True`` removed points outside :math:`(0, 1)^3`, else ``n == N``. - Tensor-product eval (``etas`` is a tuple ``(eta1, eta2, eta3)`` broadcast to shape @@ -1219,20 +1219,18 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa assert len(etas) == 3 assert etas[0].shape == etas[1].shape == etas[2].shape assert etas[0].ndim == 1 - markers = xp.stack(etas, axis=1) + markers = xp.stack(etas, axis=0) else: - markers = etas[0] - # the kernel takes markers in the column-major marker layout (no copy for the marker buffer) - markers = as_marker_array(markers) + # one column per marker (see struphy.pic.marker_layout); no copy for a contiguous array + markers = xp.ascontiguousarray(etas[0]) - # coordinates (:, 3) and argument evaluation (without holes) + # coordinates (3, :) and argument evaluation (without holes) if callable(a): if coordinates == "logical": A = Domain.prepare_arg( a, self( markers, - change_out_order=True, remove_outside=remove_outside, identity_map=True, ), @@ -1241,7 +1239,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa else: A = Domain.prepare_arg( a, - self(markers, change_out_order=True, remove_outside=remove_outside), + self(markers, remove_outside=remove_outside), a_kwargs=a_kwargs, ) @@ -1252,7 +1250,6 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa a, self( markers, - change_out_order=True, remove_outside=remove_outside, identity_map=True, ), @@ -1261,7 +1258,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa else: A = Domain.prepare_arg( a, - self(markers, change_out_order=True, remove_outside=remove_outside), + self(markers, remove_outside=remove_outside), a_kwargs=a_kwargs, ) else: @@ -1271,13 +1268,13 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa A = Domain.prepare_arg(a, markers) # check if A includes holes or not - if A.shape[0] == markers.shape[0]: + if A.shape[0] == markers.shape[1]: A_has_holes = True else: A_has_holes = False # call evaluation kernel - out = xp.empty((markers.shape[0], 3), dtype=float) + out = xp.empty((markers.shape[1], 3), dtype=float) inside = inside_logical_cube(markers) # make sure we don't have stride = 0; give A one row per marker, as the kernel is called with @@ -1287,7 +1284,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa A = A.copy() else: A_values = A - A = xp.zeros((markers.shape[0], A_values.shape[1]), dtype=float) + A = xp.zeros((markers.shape[1], A_values.shape[1]), dtype=float) A[inside] = A_values kernel_pullpush_pic( @@ -1298,7 +1295,7 @@ def _pull_push_transform(self, which, a, kind_fun, *etas, flat_eval=False, **kwa self.args_domain, out, False, - n_threads=markers.shape[0], # CUDA: one thread per marker (the first array is A) + n_threads=markers.shape[1], # CUDA: one thread per marker ) # move the (3, 3)-part to front @@ -1564,7 +1561,7 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): # vector-valued function -> must return a 4d array of shape (3,:,:,:) elif callable(a_in): if flat_eval: - a_out = a_in(Xs[0][:, 0], Xs[0][:, 1], Xs[0][:, 2], **a_kwargs) + a_out = a_in(Xs[0][0], Xs[0][1], Xs[0][2], **a_kwargs) if a_out.ndim == 1: a_out = a_out[None, :] else: @@ -1596,9 +1593,9 @@ def prepare_arg(a_in, *Xs, is_sparse_meshgrid=False, a_kwargs={}): if flat_eval: a_out += [ component( - Xs[0][:, 0], - Xs[0][:, 1], - Xs[0][:, 2], + Xs[0][0], + Xs[0][1], + Xs[0][2], **a_kwargs, ), ] @@ -1830,7 +1827,7 @@ def show( * a 2d array with information about MPI decomposition. markers : array-like - Markers to be plotted. Can be of shape (Np, 3) or (:, Np, 3). For the former, all markers are plotted with the same color. For the latter, with different colors (are interpreted as orbits in time). + Markers to be plotted, one column per marker. Can be of shape (3, Np) or (:, 3, Np). For the former, all markers are plotted with the same color. For the latter, with different colors (are interpreted as orbits in time). marker_coords : bool Whether the marker coordinates are logical or physical. @@ -2135,16 +2132,13 @@ def show( if markers.ndim == 2: if not logical and marker_coords == "logical": tmp = markers.copy() # TODO: needed for eta3=0 - tmp[:, 2] = 0.0 # TODO: needed for eta3=0 + tmp[2] = 0.0 # TODO: needed for eta3=0 X = self(tmp, remove_outside=True) else: X = ( - markers[:, 0].copy(), - markers[ - :, - 1, - ].copy(), - markers[:, 2].copy(), + markers[0].copy(), + markers[1].copy(), + markers[2].copy(), ) ax.scatter(X[co1], X[co2], s=1, color="b") @@ -2153,21 +2147,17 @@ def show( elif markers.ndim == 3: colors = ["k", "m", "b", "g", "r", "c", "y"] - for i in range(markers.shape[1]): + for i in range(markers.shape[2]): if not logical and marker_coords == "logical": # TODO: needed for eta3=0 - tmp = markers[:, i, :].copy() - tmp[:, 2] = 0.0 # TODO: needed for eta3 = 0 + tmp = markers[:, :, i].T.copy() + tmp[2] = 0.0 # TODO: needed for eta3 = 0 X = self(tmp, remove_outside=True, squeeze_out=True) else: X = ( - markers[:, i, 0].copy(), - markers[ - :, - i, - 1, - ].copy(), - markers[:, i, 2].copy(), + markers[:, 0, i].copy(), + markers[:, 1, i].copy(), + markers[:, 2, i].copy(), ) # ax.scatter(X[co1], X[co2], s=2, color=colors[i%len(colors)]) diff --git a/src/struphy/kinetic_background/base.py b/src/struphy/kinetic_background/base.py index 321b589be..f542efaa2 100644 --- a/src/struphy/kinetic_background/base.py +++ b/src/struphy/kinetic_background/base.py @@ -913,8 +913,8 @@ def __call__(self, *phase_space_coords): if self.gauss_types[i] == "mu": B0_param = self.params.get("B0", 2.0) if callable(B0_param): - # callable expects flat (N, 3) positions; reshape back for meshgrid input - etas = xp.concatenate([xp.ravel(a)[:, None] for a in args[:3]], axis=1) + # callable expects flat (3, N) positions (one column per point); reshape back for meshgrid input + etas = xp.stack([xp.ravel(a) for a in args[:3]], axis=0) B0 = xp.reshape(B0_param(etas), shape0) else: B0 = B0_param @@ -954,14 +954,9 @@ def _evaluate_moment(self, eta1, eta2, eta3, *, name: str = "n", add_perturbatio assert isinstance(params, tuple) assert len(params) == 2 - # flat evaluation for markers + # flat evaluation for markers (one column per marker) if eta1.ndim == 1: - etas = [ - xp.concatenate( - (eta1[:, None], eta2[:, None], eta3[:, None]), - axis=1, - ), - ] + etas = [xp.stack((eta1, eta2, eta3), axis=0)] # assuming that input comes from meshgrid. elif eta1.ndim == 4: etas = ( diff --git a/src/struphy/kinetic_background/maxwellians.py b/src/struphy/kinetic_background/maxwellians.py index 9293d2c79..cf0fa5a2e 100644 --- a/src/struphy/kinetic_background/maxwellians.py +++ b/src/struphy/kinetic_background/maxwellians.py @@ -66,7 +66,7 @@ def velocity_coords(self) -> LiteralOptions.VelocityCoordinates: def velocity_jacobian_det(self, eta1, eta2, eta3, vx, vy, vz): """Jacobian determinant is 1 (Cartesian velocity coordinates). - Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords.T) + Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords) Parameters ---------- @@ -194,7 +194,7 @@ def velocity_coords(self) -> LiteralOptions.VelocityCoordinates: def velocity_jacobian_det(self, eta1, eta2, eta3, v_para, mu): r"""Jacobian determinant of the velocity coordinate transformation to :math:`(v_\parallel, mu)`, is :math:`B_0`. - Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords.T) + Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords) Parameters ---------- @@ -216,7 +216,7 @@ def velocity_jacobian_det(self, eta1, eta2, eta3, v_para, mu): B0 = self.params["B0"] if callable(B0): - etas = xp.stack((eta1, eta2, eta3), axis=1) + etas = xp.stack((eta1, eta2, eta3), axis=0) return B0(etas) return B0 + 0 * eta1 @@ -324,7 +324,7 @@ def velocity_coords(self) -> LiteralOptions.VelocityCoordinates: def velocity_jacobian_det(self, eta1, eta2, eta3, v_para, v_perp): r"""Jacobian determinant of the velocity coordinate transformation to :math:`(v_\parallel, v_\perp)`, is :math:`v_\perp`. - Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords.T) + Input parameters should be slice of 2d numpy marker array. (i.e. *self.phasespace_coords) Parameters ---------- @@ -562,14 +562,9 @@ def _evaluate_moment( assert isinstance(params, tuple) assert len(params) == 2 - # flat evaluation for markers + # flat evaluation for markers (one column per marker) if eta1.ndim == 1: - coords = [ - xp.concatenate( - (eta1[:, None], eta2[:, None], eta3[:, None], v_parallel[:, None], mu[:, None]), - axis=1, - ), - ] + coords = [xp.stack((eta1, eta2, eta3, v_parallel, mu), axis=0)] # assuming that input comes from meshgrid. elif eta1.ndim == 5: coords = (eta1, eta2, eta3, v_parallel, mu) @@ -625,16 +620,16 @@ def _check_psi_c_cached(self, *coords): """Check if psi_c has been cached for the given coordinates.""" cached = False if getattr(self, "psi_c", None) is not None: - # markers: coords[0] is (n_markers, 5) and psi_c is (n_markers,); meshgrid: all have the same 5d shape - shape = coords[0].shape[:1] if len(coords) == 1 else coords[0].shape + # markers: coords[0] is (5, n_markers) and psi_c is (n_markers,); meshgrid: all have the same 5d shape + shape = coords[0].shape[1:] if len(coords) == 1 else coords[0].shape if self.psi_c.shape == shape and self.psi_c.size > 0: # compare the cached values with a fresh evaluation at a few test points size = self.psi_c.size test_idx = xp.unique(xp.array([0, size // 3, 2 * size // 3, size - 1])) if len(coords) == 1: - test_coords = coords[0][test_idx] + test_coords = coords[0][:, test_idx] else: - test_coords = xp.stack([c.reshape(-1)[test_idx] for c in coords], axis=1) + test_coords = xp.stack([c.reshape(-1)[test_idx] for c in coords], axis=0) # do not use the cache buffers here, self.psi_c may point to them cached_psi_c = self.eval_psic(test_coords, use_cbufs=False) cached = xp.allclose(self.psi_c.reshape(-1)[test_idx], cached_psi_c) @@ -655,10 +650,10 @@ def eval_psic(self, *coords, use_cbufs: bool = True): logger.warning( f"Initialize {self.__class__.__name__} with `cache_size` for faster psi_c evaluation for markers!" ) - etas = coords[0][:, :3] # these are views (no mem allocation) - vparallel = coords[0][:, 3] - mu = coords[0][:, 4] - n_markers = etas.shape[0] + etas = coords[0][:3] # these are views (no mem allocation) + vparallel = coords[0][3] + mu = coords[0][4] + n_markers = etas.shape[1] if cbufs is None: absB0 = self.equil.absB0(etas) x, y, z = self.equil.domain(etas) diff --git a/src/struphy/post_processing/orbits/orbits_tools.py b/src/struphy/post_processing/orbits/orbits_tools.py index f26af9f40..b1315ed9e 100644 --- a/src/struphy/post_processing/orbits/orbits_tools.py +++ b/src/struphy/post_processing/orbits/orbits_tools.py @@ -109,7 +109,7 @@ def post_process_orbit_guiding_center(domain, equil, path_kinetics_species, spec ) # eval cartesian magnetic filed at marker positions - B_cart[~lost_particles_mask, :] = equil.b_cart(etas[~lost_particles_mask])[0].T + B_cart[~lost_particles_mask, :] = equil.b_cart(etas[~lost_particles_mask].T)[0].T # calculate guiding center positions calculate_guiding_center_from_6d(temp, B_cart) From fd0533ab9a9cd963f1868b273b872f3498b0ecd8 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 15 Oct 2026 17:18:49 +0200 Subject: [PATCH 07/12] Propagators, models and output: one column per marker The propagators, model diagnostics and scalars index the marker buffer as markers[j, ip] and reduce over axis 0 (velocities are (vdim, n)). The saved markers are (n_cols, n_to_save), so the hdf5 marker output and the restart data are (nt, n_cols, n); post-processing of the orbits reads them in that layout and keeps its (n, q) orbit products. Co-Authored-By: Claude Opus 5.5 --- src/struphy/examples/_draw_parallel.py | 16 +++++++------- src/struphy/models/base.py | 10 ++++----- .../drift_kinetic_electrostatic_adiabatic.py | 2 +- src/struphy/models/guiding_center.py | 4 ++-- .../models/linear_mhd_driftkinetic_cc.py | 4 ++-- .../linear_vlasov_ampere_one_species.py | 4 ++-- src/struphy/models/scalars.py | 6 ++--- src/struphy/models/variables.py | 4 ++-- src/struphy/post_processing/output.py | 11 +++++----- .../propagators/current_coupling_5d_gradb.py | 22 +++++++++---------- .../propagators/efield_weights_coupling.py | 20 ++++++++--------- src/struphy/propagators/implicit_diffusion.py | 2 +- .../propagators/push_random_diffusion.py | 8 +++---- .../propagators/vlasov_ampere_coupling.py | 8 +++---- src/struphy/simulation/sim.py | 2 +- 15 files changed, 62 insertions(+), 61 deletions(-) diff --git a/src/struphy/examples/_draw_parallel.py b/src/struphy/examples/_draw_parallel.py index 6c7c1632d..7755219cf 100644 --- a/src/struphy/examples/_draw_parallel.py +++ b/src/struphy/examples/_draw_parallel.py @@ -57,7 +57,7 @@ def main(): comm.Barrier() logger.info("Number of particles w/wo holes on each process before sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.markers.shape[0]}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") domain.show( grid_info=derham.domain_array, @@ -70,7 +70,7 @@ def main(): comm.Barrier() logger.info("Number of particles w/wo holes on each process after sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.markers.shape[0]}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") domain.show( grid_info=derham.domain_array, @@ -79,18 +79,18 @@ def main(): # are all markers in the correct domain? conds = xp.logical_and( - particles.markers[:, :3] > derham.domain_array[rank, 0::3], - particles.markers[:, :3] < derham.domain_array[rank, 1::3], + particles.markers[:3] > derham.domain_array[rank, 0::3, None], + particles.markers[:3] < derham.domain_array[rank, 1::3, None], ) - holes = particles.markers[:, 0] == -1.0 - stay = xp.all(conds, axis=1) + holes = particles.markers[0] == -1.0 + stay = xp.all(conds, axis=0) - error_mks = particles.markers[xp.logical_and(~stay, ~holes)] + error_mks = particles.markers[:, xp.logical_and(~stay, ~holes)] logger.info( f"rank {rank} | markers not on correct process: {xp.nonzero(xp.logical_and(~stay, ~holes))} \ - \n corresponding positions:\n {error_mks[:, :3]}", + \n corresponding positions:\n {error_mks[:3]}", ) assert error_mks.size == 0 diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 4b0b6219c..0abc45730 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -560,14 +560,14 @@ def update_markers_to_be_saved(self): if var.n_to_save > 0: markers_on_proc = xp.logical_and( - obj.markers[:, -1] >= 0.0, - obj.markers[:, -1] < var.n_to_save, + obj.markers[-1] >= 0.0, + obj.markers[-1] < var.n_to_save, ) - # the number of rows is known on the host after the gather; counting the mask first + # the number of markers is known on the host after the gather; counting the mask first # (xp.count_nonzero) would wait for the device a second time - rows = obj.markers[markers_on_proc] + saved = obj.markers[:, markers_on_proc] var.saved_markers[:] = -1.0 - var.saved_markers[: rows.shape[0]] = rows + var.saved_markers[:, : saved.shape[1]] = saved @profile def update_distr_functions(self): diff --git a/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py b/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py index 8bd3c7779..437ac7780 100644 --- a/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py +++ b/src/struphy/models/drift_kinetic_electrostatic_adiabatic.py @@ -205,7 +205,7 @@ def _compute_en_particle_magnetic(self): # The weights (column 5) already carry the 1/N_p of the Monte-Carlo estimate, see KineticEnergyPIC. # only the two columns needed, not the whole rows of the valid markers valid_mks = particles.valid_mks - return particles.markers[valid_mks, 5].dot(particles.markers[valid_mks, 8]) + return particles.markers[5, valid_mks].dot(particles.markers[8, valid_mks]) ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): diff --git a/src/struphy/models/guiding_center.py b/src/struphy/models/guiding_center.py index 2e067a5a1..9c3147234 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[5, particles.valid_mks].dot( + particles.markers[8, particles.valid_mks], ) return energy diff --git a/src/struphy/models/linear_mhd_driftkinetic_cc.py b/src/struphy/models/linear_mhd_driftkinetic_cc.py index 5e7e96a4f..e826f94be 100644 --- a/src/struphy/models/linear_mhd_driftkinetic_cc.py +++ b/src/struphy/models/linear_mhd_driftkinetic_cc.py @@ -208,8 +208,8 @@ def _compute_en_fB(self): particles.save_magnetic_energy(self._PBb) return ( - particles.markers[particles.valid_mks, 5].dot( - particles.markers[particles.valid_mks, 8], + particles.markers[5, particles.valid_mks].dot( + particles.markers[8, particles.valid_mks], ) * Ah / Ab diff --git a/src/struphy/models/linear_vlasov_ampere_one_species.py b/src/struphy/models/linear_vlasov_ampere_one_species.py index 8644784de..9649b0276 100644 --- a/src/struphy/models/linear_vlasov_ampere_one_species.py +++ b/src/struphy/models/linear_vlasov_ampere_one_species.py @@ -202,12 +202,12 @@ def _compute_en_w(self): else: self._f0 = backgrounds self._f0_values = xp.zeros( - self.kinetic_ions.var.particles.markers.shape[0], + self.kinetic_ions.var.particles.n_rows, dtype=float, ) assert isinstance(self._f0, Maxwellian3D) - self._f0_values[particles.valid_mks] = self._f0(*particles.phasespace_coords.T) + self._f0_values[particles.valid_mks] = self._f0(*particles.phasespace_coords) # alpha^2 * v_th^2 * N / 2 * sum_p s_0 * w_p^2 / f_{0,p} (weights are w_p = delta f_p / (N * s_{0,p})) alpha = self.kinetic_ions.equation_params.alpha diff --git a/src/struphy/models/scalars.py b/src/struphy/models/scalars.py index d614b3a0f..51c13109f 100644 --- a/src/struphy/models/scalars.py +++ b/src/struphy/models/scalars.py @@ -351,9 +351,9 @@ def _local_update(self): # Particles5D velocities are (v_par, mu): mu is not a velocity, only v_par contributes here. if isinstance(particles, Particles5D): - velocities = velocities[:, :1] + velocities = velocities[:1] - energy = self.normalization * 0.5 * xp.sum(weights * xp.sum(velocities**2, axis=1)) + energy = self.normalization * 0.5 * xp.sum(weights * xp.sum(velocities**2, axis=0)) self.local_value[0] = energy @@ -399,7 +399,7 @@ def _local_update(self): velocities = particles.velocities weights = particles.weights - energy = self.normalization * 0.5 * xp.sum(weights * xp.sum(velocities**2, axis=1)) + energy = self.normalization * 0.5 * xp.sum(weights * xp.sum(velocities**2, axis=0)) self.local_value[0] = energy diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index c65bf00b5..a183173b8 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -659,7 +659,7 @@ def allocate( ) if self._n_to_save > 0: self._saved_markers = xp.zeros( - (self._n_to_save, self.particles.markers.shape[1]), + (self.particles.n_cols, self._n_to_save), dtype=float, ) @@ -955,7 +955,7 @@ def allocate( ) if self._n_to_save > 0: self._saved_markers = xp.zeros( - (self._n_to_save, self.particles.markers.shape[1]), + (self.particles.n_cols, self._n_to_save), dtype=float, ) diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index 57e3c0e32..bfcf5fd58 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -1834,8 +1834,9 @@ def _post_process_markers( # open hdf5 files and get names and number of saved markers of kinetic species with h5py.File(os.path.join(self.path_out, "data/data_proc0.hdf5"), "r") as file_0: - # get number of time steps and markers - nt, n_markers, n_cols = file_0["kinetic/" + species + "/markers"].shape + # get number of time steps and markers (the saved markers have the layout of the marker array, one + # column per marker) + nt, n_cols, n_markers = file_0["kinetic/" + species + "/markers"].shape # get the saved orbit quantities from the particle class of the species for _, var in species_obj.variables.items(): @@ -1883,9 +1884,9 @@ def _post_process_markers( for rank in self._pproc_ranks: with h5py.File(os.path.join(self.path_out, "data/", f"data_proc{rank}.hdf5"), "r") as file: markers = file["kinetic/" + species + "/markers"] - ids = markers[n * step, :, -1].astype("int") + ids = markers[n * step, -1, :].astype("int") ids = ids[ids != -1] # exclude holes - temp[ids] = markers[n * step, : ids.size, save_index] + temp[ids] = markers[n * step, save_index, : ids.size].T if self._pproc_parallel: if self._pproc_rank == 0: @@ -1909,7 +1910,7 @@ def _post_process_markers( assert xp.all(sorted(ids) == xp.arange(n_markers)) # compute physical positions (x, y, z) - pos_phys = self.domain(xp.array(temp[~lost_particles_mask, :3]), change_out_order=True) + pos_phys = self.domain(xp.array(temp[~lost_particles_mask, :3].T), change_out_order=True) temp[~lost_particles_mask, :3] = pos_phys if self._pproc_rank == 0: diff --git a/src/struphy/propagators/current_coupling_5d_gradb.py b/src/struphy/propagators/current_coupling_5d_gradb.py index a596074df..c5ab5ef66 100644 --- a/src/struphy/propagators/current_coupling_5d_gradb.py +++ b/src/struphy/propagators/current_coupling_5d_gradb.py @@ -463,10 +463,10 @@ def __call__(self, dt): first_free_idx = args_markers.first_free_idx # clear buffer - markers[:, first_pusher_idx:-2] = 0.0 + markers[first_pusher_idx:-2] = 0.0 # save old marker positions - markers[:, first_pusher_idx : first_pusher_idx + 3] = markers[:, :3] + markers[first_pusher_idx : first_pusher_idx + 3] = markers[:3] # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) b_full = self._b2.copy(out=self._b_full) @@ -530,7 +530,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_pusher_idx:-2] = 0.0 + markers[first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: @@ -565,7 +565,7 @@ def __call__(self, dt): # save en_fB_old particles.save_magnetic_energy(PB_b) - en_fB_old = xp.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_old = xp.sum(markers[8, ~holes].dot(markers[5, ~holes])) * self.options.ep_scale en_fB_old /= n_mks_tot buffer_array = xp.array([en_fB_old]) @@ -611,7 +611,7 @@ def __call__(self, dt): # save en_fB_new particles.save_magnetic_energy(PB_b) - en_fB_new = xp.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_new = xp.sum(markers[8, ~holes].dot(markers[5, ~holes])) * self.options.ep_scale en_fB_new /= n_mks_tot buffer_array = xp.array([en_fB_new]) @@ -655,13 +655,13 @@ def __call__(self, dt): u_mid.update_ghost_regions() # save H^{n+1, k} - markers[~holes, first_free_idx : first_free_idx + 3] = markers[~holes, 0:3] + markers[first_free_idx : first_free_idx + 3, ~holes] = markers[0:3, ~holes] # calculate denominator ||z^{n+1, k} - z^n||^2 sum_u_diff_loc = xp.sum((u_diff.toarray() ** 2)) sum_H_diff_loc = xp.sum( - (markers[~holes, :3] - markers[~holes, first_pusher_idx : first_pusher_idx + 3]) ** 2, + (markers[:3, ~holes] - markers[first_pusher_idx : first_pusher_idx + 3, ~holes]) ** 2, ) buffer_array = xp.array([sum_u_diff_loc]) @@ -702,7 +702,7 @@ def __call__(self, dt): *self._args_accum_kernel_en_fB_mid, first_free_idx + 3, ) - en_fB_mid = xp.sum(markers[~holes, first_free_idx + 3].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_mid = xp.sum(markers[first_free_idx + 3, ~holes].dot(markers[5, ~holes])) * self.options.ep_scale en_fB_mid /= n_mks_tot @@ -755,7 +755,7 @@ def __call__(self, dt): ) sum_H_diff_loc = xp.sum( - xp.abs(markers[~holes, 0:3] - markers[~holes, first_free_idx : first_free_idx + 3]), + xp.abs(markers[0:3, ~holes] - markers[first_free_idx : first_free_idx + 3, ~holes]), ) if particles.mpi_comm is not None: @@ -763,7 +763,7 @@ def __call__(self, dt): # update en_fB_new particles.save_magnetic_energy(PB_b) - en_fB_new = xp.sum(markers[~holes, 8].dot(markers[~holes, 5])) * self.options.ep_scale + en_fB_new = xp.sum(markers[8, ~holes].dot(markers[5, ~holes])) * self.options.ep_scale en_fB_new /= n_mks_tot buffer_array = xp.array([en_fB_new]) @@ -853,7 +853,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_pusher_idx:-2] = 0.0 + markers[first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: diff --git a/src/struphy/propagators/efield_weights_coupling.py b/src/struphy/propagators/efield_weights_coupling.py index 5604dfcf4..2990b7e62 100644 --- a/src/struphy/propagators/efield_weights_coupling.py +++ b/src/struphy/propagators/efield_weights_coupling.py @@ -204,8 +204,8 @@ def allocate(self): self._e_sum = e.space.zeros() # marker storage - self._f0_values = xp.zeros(particles.markers.shape[0], dtype=float) - self._old_weights = xp.empty(particles.markers.shape[0], dtype=float) + self._f0_values = xp.zeros(particles.n_rows, dtype=float) + self._old_weights = xp.empty(particles.n_rows, dtype=float) # ================================ # ========= Schur Solver ========= @@ -257,12 +257,12 @@ def __call__(self, dt): # evaluate f0 and accumulate self._f0_values[:] = self._f0( - particles.markers[:, 0], - particles.markers[:, 1], - particles.markers[:, 2], - particles.markers[:, 3], - particles.markers[:, 4], - particles.markers[:, 5], + particles.markers[0], + particles.markers[1], + particles.markers[2], + particles.markers[3], + particles.markers[4], + particles.markers[5], ) self._accum(self._f0_values) @@ -285,7 +285,7 @@ def __call__(self, dt): ) # Store old weights - self._old_weights[~particles.holes] = particles.markers_wo_holes[:, 6] + self._old_weights[~particles.holes] = particles.markers_wo_holes[6] # Compute (e^{n+1} + e^n) (no tmps created here) self._e_sum *= 0.0 @@ -306,7 +306,7 @@ def __call__(self, dt): logger.info(f"Maxdiff e1 for StepEfieldWeights: {max_de}") max_diff = xp.max( xp.abs( - self._old_weights[~particles.holes] - particles.markers[~particles.holes, 6], + self._old_weights[~particles.holes] - particles.markers[6, ~particles.holes], ), ) logger.info(f"Maxdiff weights for StepEfieldWeights: {max_diff}") diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index c6e01c065..f05598256 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -561,7 +561,7 @@ def __call__(self, dt): first_free_idx = src.particles.first_free_idx density = src.particles.f0.n0(eta) - src.particles.markers[valid_mks, first_free_idx] = density + src.particles.markers[first_free_idx, valid_mks] = density # 2. accumulate src() # 3. take weak divergence diff --git a/src/struphy/propagators/push_random_diffusion.py b/src/struphy/propagators/push_random_diffusion.py index dce8153e7..1d1881627 100644 --- a/src/struphy/propagators/push_random_diffusion.py +++ b/src/struphy/propagators/push_random_diffusion.py @@ -110,8 +110,8 @@ def allocate(self): particles = self.variables.var.particles - # C-ordered, as the kernel declares it (a copy of the column-major markers would keep their order) - self._noise = array(particles.markers[:, :3], order="C") + # one column per marker, like the markers + self._noise = array(particles.markers[:3]) self._butcher = self.options.butcher # temp fix due to refactoring of ButcherTableau: @@ -150,8 +150,8 @@ def __call__(self, dt): self._noise[:] = random.multivariate_normal( self._mean, self._cov, - len(particles.markers), - ) + particles.n_rows, + ).T # push markers self._pusher(dt) diff --git a/src/struphy/propagators/vlasov_ampere_coupling.py b/src/struphy/propagators/vlasov_ampere_coupling.py index 18dbe2884..02bf5ed52 100644 --- a/src/struphy/propagators/vlasov_ampere_coupling.py +++ b/src/struphy/propagators/vlasov_ampere_coupling.py @@ -234,10 +234,10 @@ def __call__(self, dt): logger.info(f"Maxdiff e1 for VlasovMaxwell: {max_de}") particles = self.variables.ions.particles # the velocities before the push are saved at first_pusher_idx + 3 (see Pusher._push); - # gather only the six columns needed, not the whole rows + # gather only the six rows (quantities) needed, not all quantities of the markers not_holes = ~particles.holes - v_new = particles.markers[not_holes, 3:6] - v_old = particles.markers[not_holes, particles.first_pusher_idx + 3 : particles.first_pusher_idx + 6] - max_diff = xp.max(xp.abs(xp.linalg.norm(v_new, axis=1) - xp.linalg.norm(v_old, axis=1))) + v_new = particles.markers[3:6, not_holes] + v_old = particles.markers[particles.first_pusher_idx + 3 : particles.first_pusher_idx + 6, not_holes] + max_diff = xp.max(xp.abs(xp.linalg.norm(v_new, axis=0) - xp.linalg.norm(v_old, axis=0))) logger.info(f"Maxdiff |v| for VlasovMaxwell: {max_diff}") logger.info("") diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 2719095f5..6567c0d46 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -1624,7 +1624,7 @@ def _initialize_from_restart(self, data: DataContainer): elif isinstance(subval, (PICVariable, SPHVariable)): key_restart = os.path.join("restart", species) subval.particles._markers[:, :] = file[key_restart][-1, :, :] - # refresh the cached hole/ghost masks for the restored rows + # refresh the cached hole/ghost masks for the restored markers subval.particles.update_holes() subval.particles._update_ghost_particles() From feeccbe9e0704306a6abb35b039214a09f55b8d7 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 15 Oct 2026 17:18:54 +0200 Subject: [PATCH 08/12] Tests for the (n_cols, n) marker layout The tests build and index marker arrays with one column per marker and pass (k, n) arrays to the marker evaluation. test_marker_layout checks the new contract: the buffer is C-ordered (n_cols, n_rows) with n_rows a multiple of 32, selections of markers are Fortran-ordered copies that kernels accept after ascontiguousarray, and evaluation on the buffer and on selections agrees. The CUDA parity cases pass n_threads explicitly. test_mhd_equils checked the eta1 case with kind="markers" left over from the previous case. Co-Authored-By: Claude Opus 5.5 --- .../bsplines/tests/test_evaluation_cuda.py | 6 +- src/struphy/feec/tests/test_eval_field.py | 20 +-- .../tests/test_equils_cupy.py | 2 +- .../tests/test_gvec_equil.py | 8 +- .../tests/test_mhd_equils.py | 8 +- src/struphy/geometry/tests/test_domain.py | 66 ++++---- .../kinetic_background/tests/test_base.py | 4 +- .../tests/test_maxwellians.py | 14 +- src/struphy/models/tests/gpu_e2e.py | 8 +- .../tests/test_diagnostics_on_device.py | 31 ++-- .../test_gpu_LinearVlasovAmpereOneSpecies.py | 10 +- .../tests/test_gpu_VlasovAmpereOneSpecies.py | 10 +- src/struphy/models/tests/test_scalars.py | 50 +++--- ...est_verif_IncompressibleNavierStokesSPH.py | 8 +- src/struphy/pic/tests/bench_marker_layout.py | 36 +++-- src/struphy/pic/tests/cuda_parity_cases.py | 67 +++++--- src/struphy/pic/tests/kernel_test_args.py | 37 +++-- src/struphy/pic/tests/test_accum_gc_gradB.py | 10 +- .../pic/tests/test_accum_pressure_coupling.py | 4 +- src/struphy/pic/tests/test_accum_vec_H1.py | 2 +- src/struphy/pic/tests/test_cuda_emulation.py | 12 +- src/struphy/pic/tests/test_device_helpers.py | 6 +- src/struphy/pic/tests/test_draw_parallel.py | 32 ++-- src/struphy/pic/tests/test_kernel_backends.py | 15 +- src/struphy/pic/tests/test_kernel_setup.py | 55 +++---- src/struphy/pic/tests/test_load_external.py | 6 +- src/struphy/pic/tests/test_marker_layout.py | 150 ++++++++++-------- src/struphy/pic/tests/test_pole.py | 76 ++++----- src/struphy/pic/tests/test_pushers.py | 26 +-- src/struphy/pic/tests/test_reflect.py | 22 +-- src/struphy/pic/tests/test_sorting.py | 12 +- src/struphy/pic/tests/test_sorting_device.py | 22 +-- src/struphy/pic/tests/test_sph.py | 8 +- src/struphy/pic/tests/test_tesselation.py | 14 +- src/struphy/pic/tests/test_utilities.py | 12 +- .../tests/test_orbits_tools.py | 2 +- .../tests/test_push_random_diffusion.py | 8 +- src/struphy/simulation/tests/test_output.py | 4 +- 38 files changed, 463 insertions(+), 420 deletions(-) diff --git a/src/struphy/bsplines/tests/test_evaluation_cuda.py b/src/struphy/bsplines/tests/test_evaluation_cuda.py index b2c4b0ed8..3fba7c53e 100644 --- a/src/struphy/bsplines/tests/test_evaluation_cuda.py +++ b/src/struphy/bsplines/tests/test_evaluation_cuda.py @@ -37,11 +37,11 @@ def test_evaluation_parity(mode, kind, empty): xp.asarray([1, 1, 1], dtype=xp.int64), ) if mode == "markers": - points = xp.full((0 if empty else 129, 6), 0.5)[:, ::2] + points = xp.full((6, 0 if empty else 129), 0.5)[::2] if not empty: points[0, 0] = -1.0 coords = (points,) - storage = xp.full((points.shape[0] * 2,), 17.0) + storage = xp.full((points.shape[1] * 2,), 17.0) out = storage[::2] else: axes = [xp.linspace(0.3, 0.7, n) for n in (0 if empty else 7, 5, 4)] @@ -97,7 +97,7 @@ def test_spline_function_backends(space, mode, local, target_backend): component._data[:] = i + 1.0 axes = [np.array([0.2, 0.6]), np.array([0.3, 0.7]), np.array([0.4])] if mode == "markers": - coords = (xp.asarray([[0.2, 0.3, 0.4], [0.6, 0.7, 0.4], [-0.1, 0.3, 0.4]]),) + coords = (xp.asarray([[0.2, 0.6, -0.1], [0.3, 0.7, 0.3], [0.4, 0.4, 0.4]]),) shape = (3,) elif mode == "matrix": coords = tuple(xp.asarray(a) for a in np.meshgrid(*axes, indexing="ij")) diff --git a/src/struphy/feec/tests/test_eval_field.py b/src/struphy/feec/tests/test_eval_field.py index c60f36c76..5f6e48c57 100644 --- a/src/struphy/feec/tests/test_eval_field.py +++ b/src/struphy/feec/tests/test_eval_field.py @@ -73,13 +73,13 @@ def test_eval_field(num_elements, degree, bcs): # evaluation points for markers Np = 33 - markers = xp.random.rand(Np, 3) - markers_1 = xp.zeros((eta1.size, 3)) - markers_1[:, 0] = eta1 - markers_2 = xp.zeros((eta2.size, 3)) - markers_2[:, 1] = eta2 - markers_3 = xp.zeros((eta3.size, 3)) - markers_3[:, 2] = eta3 + markers = xp.random.rand(3, Np) + markers_1 = xp.zeros((3, eta1.size)) + markers_1[0] = eta1 + markers_2 = xp.zeros((3, eta2.size)) + markers_2[1] = eta2 + markers_3 = xp.zeros((3, eta3.size)) + markers_3[2] = eta3 # arrays for legacy evaluation arr1, arr2, arr3, is_sparse_meshgrid = Domain.prepare_eval_pts(eta1, eta2, eta3) @@ -575,11 +575,11 @@ def test_eval_markers_on_breaks_and_squeeze_out(): ], ) - # markers on all process breaks, with extra (non-position) columns + # markers on all process breaks, with extra (non-position) rows breaks = [xp.unique(xp.concatenate([dom[:, 3 * d], dom[:, 3 * d + 1]])) for d in range(3)] pts = xp.array([[b1, b2, b3] for b1 in breaks[0] for b2 in breaks[1] for b3 in breaks[2]]) - markers = xp.full((pts.shape[0], 5), 7.0) - markers[:, :3] = pts + markers = xp.full((5, pts.shape[0]), 7.0) + markers[:3] = pts.T markers_orig = markers.copy() vals = p0(markers) diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py index 8c51674bc..2a186bc99 100644 --- a/src/struphy/fields_background/tests/test_equils_cupy.py +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -192,7 +192,7 @@ def check_equil_on_cupy(case): e1 = np.sort(rng.uniform(0.1, 0.9, 4)) e2 = np.sort(rng.uniform(0.0, 1.0, 5)) e3 = np.sort(rng.uniform(0.0, 1.0, 3)) - markers = np.column_stack([rng.uniform(0.1, 0.9, 7), rng.uniform(0.0, 1.0, 7), rng.uniform(0.0, 1.0, 7)]) + markers = np.stack([rng.uniform(0.1, 0.9, 7), rng.uniform(0.0, 1.0, 7), rng.uniform(0.0, 1.0, 7)]) with cunumpy.use_backend("numpy"): equil = _build(case) diff --git a/src/struphy/fields_background/tests/test_gvec_equil.py b/src/struphy/fields_background/tests/test_gvec_equil.py index d21017ede..1ee16e5ae 100644 --- a/src/struphy/fields_background/tests/test_gvec_equil.py +++ b/src/struphy/fields_background/tests/test_gvec_equil.py @@ -11,8 +11,8 @@ def _markers(n=7, seed=3): """``n`` markers at random logical positions, away from the hole at the axis.""" rng = np.random.default_rng(seed) - markers = rng.uniform(0.0, 1.0, (n, 3)) - markers[:, 0] = rng.uniform(0.1, 0.9, n) + markers = rng.uniform(0.0, 1.0, (3, n)) + markers[0] = rng.uniform(0.1, 0.9, n) return markers @@ -33,11 +33,11 @@ def test_gvec_marker_evaluation_matches_grid(name, with_gvec=False): flat = getattr(equil, name)(markers) flat = flat if isinstance(flat, tuple) else (flat,) - for ip, (e1, e2, e3) in enumerate(markers): + for ip, (e1, e2, e3) in enumerate(markers.T): grid = getattr(equil, name)(np.array([e1]), np.array([e2]), np.array([e3])) grid = grid if isinstance(grid, tuple) else (grid,) for f, g in zip(flat, grid): - assert f.shape == (len(markers),) + assert f.shape == (markers.shape[1],) assert np.isclose(f[ip], g.ravel()[0], rtol=1e-12, atol=1e-14) diff --git a/src/struphy/fields_background/tests/test_mhd_equils.py b/src/struphy/fields_background/tests/test_mhd_equils.py index aded64343..dfe0c708b 100644 --- a/src/struphy/fields_background/tests/test_mhd_equils.py +++ b/src/struphy/fields_background/tests/test_mhd_equils.py @@ -161,7 +161,7 @@ def test_equils(equil_domain_pair): mat_123_1_sp, mat_123_2_sp, mat_123_3_sp = xp.meshgrid(e1, e2, e3, indexing="ij", sparse=True) # markers - markers = xp.random.rand(33, 10) + markers = xp.random.rand(10, 33) # create MHD equilibrium eq_mhd = getattr(equils, equil_domain_pair[0])(**equil_domain_pair[1]) @@ -307,6 +307,8 @@ def test_equils(equil_domain_pair): results.append(eq_mhd.unit_b_cart(e1, e2_pt, e3_pt, squeeze_out=True)) # asserts + kind = "e1" + for i in range(0, 5): assert_scalar(results[i], kind, e1, e2_pt, e3_pt) @@ -858,7 +860,7 @@ def test_current_sheet_plot_profiles(monkeypatch): def assert_scalar(result, kind, *etas): if kind == "markers": markers = etas[0] - n_p = markers.shape[0] + n_p = markers.shape[1] assert isinstance(result, xp.ndarray) assert result.shape == (n_p,) @@ -925,7 +927,7 @@ def assert_scalar(result, kind, *etas): def assert_vector(result, kind, *etas): if kind == "markers": markers = etas[0] - n_p = markers.shape[0] + n_p = markers.shape[1] assert isinstance(result, xp.ndarray) assert result.shape == (3, n_p) diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index 466d138d9..7ff33dc69 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -32,7 +32,7 @@ def test_domain_deepcopy(mapping): assert domain_copy.args_domain.params is not domain.args_domain.params assert domain_copy.args_domain.degree is not domain.args_domain.degree - markers = xp.array([[0.2, 0.4, 0.6]]) + markers = xp.array([[0.2], [0.4], [0.6]]) assert xp.allclose(domain(markers), domain_copy(markers)) assert xp.allclose(domain.jacobian(markers), domain_copy.jacobian(markers)) @@ -53,7 +53,7 @@ def test_domain_pickle_roundtrip(mapping): assert restored.args_domain is not None assert restored.args_domain.kind_map == domain.args_domain.kind_map - markers = xp.array([[0.3, 0.5, 0.7]]) + markers = xp.array([[0.3], [0.5], [0.7]]) assert xp.allclose(domain(markers), restored(markers)) assert xp.allclose(domain.jacobian_det(markers), restored.jacobian_det(markers)) @@ -77,7 +77,7 @@ def test_spline_default(with_gvec=False): assert xp.allclose(domain.cy, ref.cy) assert xp.allclose(domain.cz, ref.cz) - markers = xp.array([[0.3, 0.5, 0.7]]) + markers = xp.array([[0.3], [0.5], [0.7]]) assert xp.allclose(domain(markers), ref(markers)) assert xp.allclose(domain.jacobian_det(markers), ref.jacobian_det(markers)) @@ -166,10 +166,10 @@ def a_vec(e1, e2, e3): assert Domain.prepare_arg([A1, A2, A3], E1, E2, E3).shape == shape_vector # ============== markers evaluation ========================== - markers = xp.random.rand(10, 6) + markers = xp.random.rand(6, 10) - shape_scalar = (markers.shape[0], 1) - shape_vector = (markers.shape[0], 3) + shape_scalar = (markers.shape[1], 1) + shape_vector = (markers.shape[1], 3) # ======== callables ============ @@ -193,11 +193,11 @@ def a_vec(e1, e2, e3): # ======== arrays =============== - A1 = a1(markers[:, 0], markers[:, 1], markers[:, 2]) - A2 = a2(markers[:, 0], markers[:, 1], markers[:, 2]) - A3 = a3(markers[:, 0], markers[:, 1], markers[:, 2]) + A1 = a1(markers[0], markers[1], markers[2]) + A2 = a2(markers[0], markers[1], markers[2]) + A3 = a3(markers[0], markers[1], markers[2]) - A = a_vec(markers[:, 0], markers[:, 1], markers[:, 2]) + A = a_vec(markers[0], markers[1], markers[2]) # scalar function assert Domain.prepare_arg(A1, markers).shape == shape_scalar @@ -229,18 +229,18 @@ def test_hollow_torus_inverse_map(sfl, pol_period): domain = HollowTorus(sfl=sfl, pol_period=pol_period) xp.random.seed(1234) - etas = xp.random.rand(50, 3) + etas = xp.random.rand(3, 50) # eta2 = 0 and eta2 = 0.5 lie on the midplane - etas[:10, 1] = 0.0 - etas[10:20, 1] = 0.5 + etas[1, :10] = 0.0 + etas[1, 10:20] = 0.5 x, y, z = domain(etas, remove_outside=False) e1, e2, e3 = domain.inverse_map(x, y, z) - assert xp.allclose(domain(xp.stack([e1, e2, e3], axis=1), remove_outside=False), xp.stack([x, y, z])) - assert xp.allclose(e1, etas[:, 0]) - assert xp.allclose(xp.minimum(xp.abs(e2 - etas[:, 1]), 1 - xp.abs(e2 - etas[:, 1])), 0.0) - assert xp.allclose(e3, etas[:, 2]) + assert xp.allclose(domain(xp.stack([e1, e2, e3]), remove_outside=False), xp.stack([x, y, z])) + assert xp.allclose(e1, etas[0]) + assert xp.allclose(xp.minimum(xp.abs(e2 - etas[1]), 1 - xp.abs(e2 - etas[1])), 0.0) + assert xp.allclose(e3, etas[2]) @pytest.mark.parametrize( @@ -274,7 +274,7 @@ def test_evaluation_mappings(mapping): arr1 = xp.linspace(0.0, 1.0, 4) arr2 = xp.linspace(0.0, 1.0, 5) arr3 = xp.linspace(0.0, 1.0, 6) - arrm = xp.random.rand(10, 8) + arrm = xp.random.rand(8, 10) logger.info("") logger.info('Testing "evaluate"...') logger.info(f"array shapes: {arr1.shape} {arr2.shape} {arr3.shape} {arrm.shape}") @@ -306,12 +306,12 @@ def test_evaluation_mappings(mapping): # markers evaluation: logger.info(f"markers evaluation, shape: {domain(arrm).shape}") - assert domain(arrm).shape == (3, arrm.shape[0]) - assert domain.jacobian(arrm).shape == (3, 3, arrm.shape[0]) - assert domain.jacobian_det(arrm).shape == (arrm.shape[0],) - assert domain.jacobian_inv(arrm).shape == (3, 3, arrm.shape[0]) - assert domain.metric(arrm).shape == (3, 3, arrm.shape[0]) - assert domain.metric_inv(arrm).shape == (3, 3, arrm.shape[0]) + assert domain(arrm).shape == (3, arrm.shape[1]) + assert domain.jacobian(arrm).shape == (3, 3, arrm.shape[1]) + assert domain.jacobian_det(arrm).shape == (arrm.shape[1],) + assert domain.jacobian_inv(arrm).shape == (3, 3, arrm.shape[1]) + assert domain.metric(arrm).shape == (3, 3, arrm.shape[1]) + assert domain.metric_inv(arrm).shape == (3, 3, arrm.shape[1]) # eta1-array evaluation: logger.info(f"eta1 array evaluation, shape: {domain(arr1, 0.5, 0.5, squeeze_out=True).shape}") @@ -473,7 +473,7 @@ def test_pullback(): logger.info('Testing "pull"...') logger.info(f"array shapes: {arr1.shape} {arr2.shape} {arr3.shape}") - markers = xp.random.rand(13, 6) + markers = xp.random.rand(6, 13) # physical function to pull back (used as components of forms too): def fun(x, y, z): @@ -504,9 +504,9 @@ def fun(x, y, z): # markers pullback: if p_str == "0" or p_str == "3": - assert domain.pull(fun_form, markers, kind=p_str, squeeze_out=True).shape == (markers.shape[0],) + assert domain.pull(fun_form, markers, kind=p_str, squeeze_out=True).shape == (markers.shape[1],) else: - assert domain.pull(fun_form, markers, kind=p_str, squeeze_out=True).shape == (3, markers.shape[0]) + assert domain.pull(fun_form, markers, kind=p_str, squeeze_out=True).shape == (3, markers.shape[1]) # eta1-array pullback: # logger.info('eta1 array pullback, shape:', domain.pull(fun_form, arr1, .5, .5, p_str).shape) @@ -633,7 +633,7 @@ def test_pushforward(): logger.info('Testing "push"...') logger.info(f"array shapes: {arr1.shape} {arr2.shape} {arr3.shape}") - markers = xp.random.rand(13, 6) + markers = xp.random.rand(6, 13) # logical function to push (used as components of forms too): def fun(e1, e2, e3): @@ -664,9 +664,9 @@ def fun(e1, e2, e3): # markers push: if p_str == "0" or p_str == "3": - assert domain.push(fun_form, markers, kind=p_str).shape == (markers.shape[0],) + assert domain.push(fun_form, markers, kind=p_str).shape == (markers.shape[1],) else: - assert domain.push(fun_form, markers, kind=p_str).shape == (3, markers.shape[0]) + assert domain.push(fun_form, markers, kind=p_str).shape == (3, markers.shape[1]) # eta1-array push: # logger.info('eta1 array push, shape:', domain.push(fun_form, arr1, .5, .5, p_str).shape) @@ -793,7 +793,7 @@ def test_transform(): logger.info('Testing "transform"...') logger.info(f"array shapes: {arr1.shape} {arr2.shape} {arr3.shape}") - markers = xp.random.rand(13, 6) + markers = xp.random.rand(6, 13) # logical function to push (used as components of forms too): def fun(e1, e2, e3): @@ -824,9 +824,9 @@ def fun(e1, e2, e3): # markers transform: if p_str == "0_to_3" or p_str == "3_to_0": - assert domain.transform(fun_form, markers, kind=p_str).shape == (markers.shape[0],) + assert domain.transform(fun_form, markers, kind=p_str).shape == (markers.shape[1],) else: - assert domain.transform(fun_form, markers, kind=p_str).shape == (3, markers.shape[0]) + assert domain.transform(fun_form, markers, kind=p_str).shape == (3, markers.shape[1]) # eta1-array transform: # logger.info('eta1 array transform, shape:', domain.transform(fun_form, arr1, .5, .5, p_str).shape) diff --git a/src/struphy/kinetic_background/tests/test_base.py b/src/struphy/kinetic_background/tests/test_base.py index ff8d89190..895a8ab2b 100644 --- a/src/struphy/kinetic_background/tests/test_base.py +++ b/src/struphy/kinetic_background/tests/test_base.py @@ -102,7 +102,7 @@ def test_plotting_function(show=False): def n_init(*etas): if len(etas) == 1: - e1, e2, e3 = etas[0][:, 0], etas[0][:, 1], etas[0][:, 1] + e1, e2, e3 = etas[0][0], etas[0][1], etas[0][1] else: assert len(etas) == 3 e1, e2, e3 = etas[0], etas[1], etas[2] @@ -110,7 +110,7 @@ def n_init(*etas): def vth(*etas): if len(etas) == 1: - e1, e2, e3 = etas[0][:, 0], etas[0][:, 1], etas[0][:, 1] + e1, e2, e3 = etas[0][0], etas[0][1], etas[0][1] else: assert len(etas) == 3 e1, e2, e3 = etas[0], etas[1], etas[2] diff --git a/src/struphy/kinetic_background/tests/test_maxwellians.py b/src/struphy/kinetic_background/tests/test_maxwellians.py index a6fbffab7..395c8ad37 100644 --- a/src/struphy/kinetic_background/tests/test_maxwellians.py +++ b/src/struphy/kinetic_background/tests/test_maxwellians.py @@ -285,7 +285,7 @@ def test_maxwellian_3d_mhd(num_elements, with_desc, show_plot=False): v2_fl = xp.random.randn(n_mks) v3_fl = xp.random.randn(n_mks) args_fl = [e1_fl, e2_fl, e3_fl, v1_fl, v2_fl, v3_fl] - e_args_fl = xp.concatenate((e1_fl[:, None], e2_fl[:, None], e3_fl[:, None]), axis=1) + e_args_fl = xp.stack((e1_fl, e2_fl, e3_fl), axis=0) for key, val in inspect.getmembers(equils): if inspect.isclass(val) and val.__module__ == equils.__name__: @@ -806,7 +806,7 @@ def test_maxwellian_2d_uniform(num_elements, show_plot=False): assert xp.allclose(res, res_ana, atol=10e-10), f"{res=},\n {res_ana =}" -@pytest.mark.parametrize("B0", [5.0, lambda etas: 1.0 + etas[:, 0] + 0.5 * etas[:, 2]]) +@pytest.mark.parametrize("B0", [5.0, lambda etas: 1.0 + etas[0] + 0.5 * etas[2]]) @pytest.mark.parametrize("volume_form", [False, True]) def test_gyro_maxwellian_2d_B0_meshgrid(B0, volume_form): """Meshgrid and flat evaluation of GyroMaxwellian2D must agree for scalar and callable B0.""" @@ -1075,7 +1075,7 @@ def test_maxwellian_2d_mhd(num_elements, with_desc, show_plot=False): v1_fl = xp.random.randn(n_mks) v2_fl = xp.random.rand(n_mks) args_fl = [e1_fl, e2_fl, e3_fl, v1_fl, v2_fl] - e_args_fl = xp.concatenate((e1_fl[:, None], e2_fl[:, None], e3_fl[:, None]), axis=1) + e_args_fl = xp.stack((e1_fl, e2_fl, e3_fl), axis=0) for key, val in inspect.getmembers(equils): if inspect.isclass(val) and val.__module__ == equils.__name__: @@ -1542,7 +1542,7 @@ def test_canonical_maxwellian_uniform(num_markers, show_plot=False): def ref_psic_and_absB(eta1, eta2, eta3, v_para, mu): """Independent reference computation of psi_c and |B0|, mirroring CanonicalMaxwellian2D.eval_psic but without using the class under test.""" - etas = xp.concatenate((eta1[:, None], eta2[:, None], eta3[:, None]), axis=1) + etas = xp.stack((eta1, eta2, eta3), axis=0) absB = mhd_equil.absB0(etas) x, y, z = mhd_equil.domain(etas) R, P, Z = mhd_equil.inverse_map(x, y, z) @@ -1643,7 +1643,7 @@ def n_of_psic(psic): # calling again with the same markers must hit the internal psi_c cache # and still return the (correct) result - markers = xp.concatenate((eta1[:, None], eta2[:, None], eta3[:, None], v_para[:, None], mu[:, None]), axis=1) + markers = xp.stack((eta1, eta2, eta3, v_para, mu), axis=0) assert maxwellian_nc._check_psi_c_cached(markers) res_cached = maxwellian_nc(eta1, eta2, eta3, v_para, mu) assert xp.allclose(res_cached, res_ana, atol=10e-10), f"{res_cached=},\n {res_ana=}" @@ -1721,9 +1721,7 @@ def _random_markers(seed): assert xp.allclose(res_cache, res_ana_c, atol=10e-10), f"{res_cache=},\n {res_cache=}" # a cache hit must not overwrite the buffer that the cached psi_c points to - markers_c = xp.concatenate( - (eta1_c[:, None], eta2_c[:, None], eta3_c[:, None], v_para_c[:, None], mu_c[:, None]), axis=1 - ) + markers_c = xp.stack((eta1_c, eta2_c, eta3_c, v_para_c, mu_c), axis=0) assert maxwellian_cache._check_psi_c_cached(markers_c) res_cache = maxwellian_cache(eta1_c, eta2_c, eta3_c, v_para_c, mu_c) assert xp.allclose(res_cache, res_ana_c, atol=10e-10), f"{res_cache=},\n {res_ana_c=}" diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py index 332a48492..ecfa0c9a2 100644 --- a/src/struphy/models/tests/gpu_e2e.py +++ b/src/struphy/models/tests/gpu_e2e.py @@ -199,18 +199,18 @@ class RunResult: scalars: dict # name -> time series (rank 0 only, else empty) fields: dict # FEEC variable name -> local coefficients at the end - markers: np.ndarray # all valid markers (gathered from all ranks), sorted by marker ID + markers: np.ndarray # all valid markers (gathered from all ranks, one column each), sorted by marker ID n_scalars: int budget: TransferBudget def gather_markers(particles) -> np.ndarray: - """The valid markers of all ranks on the host, sorted by marker ID (the last column).""" + """The valid markers of all ranks on the host (one column each), sorted by marker ID (the last row).""" comm = MPI.COMM_WORLD local = np.asarray(xp.to_numpy(particles.markers_wo_holes_and_ghost)) parts = comm.allgather(local) if comm.Get_size() > 1 else [local] - markers = np.concatenate(parts, axis=0) - return markers[np.argsort(markers[:, -1])] + markers = np.concatenate(parts, axis=1) + return markers[:, np.argsort(markers[-1])] def run( diff --git a/src/struphy/models/tests/test_diagnostics_on_device.py b/src/struphy/models/tests/test_diagnostics_on_device.py index b29ce81b0..733918f4b 100644 --- a/src/struphy/models/tests/test_diagnostics_on_device.py +++ b/src/struphy/models/tests/test_diagnostics_on_device.py @@ -26,18 +26,18 @@ from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env -N_ROWS = 12 # marker rows: 8 valid markers, 2 ghosts, 2 holes -N_COLS = 13 # Particles6D columns: eta (3), v (3), w, s0, w0, 2 auxiliary, box, ID +N_ROWS = 12 # markers (one column each): 8 valid markers, 2 ghosts, 2 holes +N_COLS = 13 # Particles6D quantities (one row each): eta (3), v (3), w, s0, w0, 2 auxiliary, box, ID def _markers(): """Particles6D-like markers on the host: 8 valid markers, 2 ghosts (ID -2) and 2 holes (all -1).""" rng = np.random.default_rng(697) - markers = rng.random((N_ROWS, N_COLS)) - markers[:, 3:6] = rng.normal(size=(N_ROWS, 3)) - markers[:, -1] = np.arange(N_ROWS) - markers[8:10, -1] = -2.0 # ghosts - markers[10:] = -1.0 # holes + markers = rng.random((N_COLS, N_ROWS)) + markers[3:6] = rng.normal(size=(3, N_ROWS)) + markers[-1] = np.arange(N_ROWS) + markers[-1, 8:10] = -2.0 # ghosts + markers[:, 10:] = -1.0 # holes return markers @@ -45,6 +45,7 @@ class _MockParticles: """The parts of :class:`~struphy.pic.base.Particles` used by the diagnostics, on the active backend.""" vdim = 3 + n_cols = N_COLS Np = 8 n_lost_markers = 3 clone_config = None @@ -52,7 +53,7 @@ class _MockParticles: def __init__(self): self._markers = xp.asarray(_markers()) - self._valid_mks = ~xp.logical_or(self._markers[:, 0] == -1.0, self._markers[:, -1] == -2.0) + self._valid_mks = ~xp.logical_or(self._markers[0] == -1.0, self._markers[-1] == -2.0) @property def markers(self): @@ -64,15 +65,15 @@ def valid_mks(self): @property def markers_wo_holes_and_ghost(self): - return self.markers[self.valid_mks] + return self.markers[:, self.valid_mks] @property def velocities(self): - return self.markers[self.valid_mks, self.index["vel"]] + return self.markers[self.index["vel"], self.valid_mks] @property def weights(self): - return self.markers[self.valid_mks, self.index["weights"]] + return self.markers[self.index["weights"], self.valid_mks] def _fake_cupy_active(): @@ -142,7 +143,7 @@ def _scalar_values(backend): values = {key: scalars.host_value(key).copy() for key in scalars.dct} # the host values follow the next update, and are views of one host buffer - particles.markers[:, 6] *= 2.0 + particles.markers[6] *= 2.0 scalars.update() doubled = scalars.host_value("en_kin")[0] return values, counter, doubled @@ -214,7 +215,7 @@ class _Variable(PICVariable): def __init__(self): self._particles = particles self._n_to_save = 5 - self._saved_markers = xp.zeros((5, N_COLS)) + self._saved_markers = xp.zeros((N_COLS, 5)) class _Species(ParticleSpecies): def __init__(self): @@ -233,7 +234,7 @@ def check_saved_markers(): reference, _ = _saved_markers("numpy") assert counter.events == [], counter.report() # IDs 0-4 are valid markers here; the ghosts and holes have no ID below n_to_save - np.testing.assert_array_equal(reference[:, -1], [0.0, 1.0, 2.0, 3.0, 4.0]) + np.testing.assert_array_equal(reference[-1], [0.0, 1.0, 2.0, 3.0, 4.0]) np.testing.assert_array_equal(saved, reference) @@ -271,7 +272,7 @@ def _flagged(backend): spline = SplineFunction.__new__(SplineFunction) spline._derham = SimpleNamespace(domain_array=xp.asarray(domain_array), comm=None) E1, E2, E3 = (xp.asarray(a) for a in np.meshgrid(*[np.linspace(0.0, 1.0, 5)] * 3, indexing="ij")) - markers = xp.asarray(np.random.default_rng(1).random((10, 3))) + markers = xp.asarray(np.random.default_rng(1).random((3, 10))) spline._flag_pts_not_on_proc(E1, E2, E3) # first call: caches the host bounds with xp.profiling.count_transfers() as counter, no_implicit_syncs(): spline._flag_pts_not_on_proc(E1, E2, E3) diff --git a/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py index b6a9a6f51..e23a38ae5 100644 --- a/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py @@ -120,7 +120,7 @@ def run_model(backend: str, out_folders: str) -> RunResult: def check_reference(ref: RunResult): """Physics sanity of the NumPy run: the comparison is only meaningful if the reference is right.""" - assert ref.markers.shape[0] == PPC * int(np.prod(NUM_ELEMENTS)) + assert ref.markers.shape[1] == PPC * int(np.prod(NUM_ELEMENTS)) if MPI.COMM_WORLD.Get_rank() == 0: en_E, en_tot = ref.scalars["en_E"], ref.scalars["en_tot"] assert en_tot.size == N_STEPS + 1 @@ -148,11 +148,11 @@ def compare_runs(ref: RunResult, gpu: RunResult): entry) leaves a factor 100 against that and is still far below the change per step (~1e-2). """ rtol = 1e-8 - np.testing.assert_array_equal(gpu.markers[:, -1], ref.markers[:, -1]) # same markers - np.testing.assert_allclose(gpu.markers[:, 3:6], ref.markers[:, 3:6], rtol=1e-13, atol=1e-13) - np.testing.assert_allclose(gpu.markers[:, 0:3], ref.markers[:, 0:3], rtol=1e-13, atol=1e-13) + np.testing.assert_array_equal(gpu.markers[-1], ref.markers[-1]) # same markers + np.testing.assert_allclose(gpu.markers[3:6], ref.markers[3:6], rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(gpu.markers[0:3], ref.markers[0:3], rtol=1e-13, atol=1e-13) - w_ref, w_gpu = ref.markers[:, 6], gpu.markers[:, 6] + w_ref, w_gpu = ref.markers[6], gpu.markers[6] np.testing.assert_allclose(w_gpu, w_ref, rtol=0, atol=rtol * np.max(np.abs(w_ref))) e_ref, e_gpu = ref.fields["e_field"], gpu.fields["e_field"] diff --git a/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py index c647847e1..e3650c0f6 100644 --- a/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py @@ -130,7 +130,7 @@ def run_model(backend: str, out_folders: str) -> RunResult: def check_reference(ref: RunResult): """Physics sanity of the NumPy run: the comparison is only meaningful if the reference is right.""" - assert ref.markers.shape[0] == PPC * int(np.prod(NUM_ELEMENTS)) + assert ref.markers.shape[1] == PPC * int(np.prod(NUM_ELEMENTS)) if MPI.COMM_WORLD.Get_rank() == 0: en_E = ref.scalars["electric_energy"] en_tot = ref.scalars["total_energy"] @@ -156,14 +156,14 @@ def compare_runs(ref: RunResult, gpu: RunResult): the change per step (~5e-2 in the electric energy). Full-f weights do not change, so they must agree to round-off. """ rtol = 1e-8 - np.testing.assert_array_equal(gpu.markers[:, -1], ref.markers[:, -1]) # same markers + np.testing.assert_array_equal(gpu.markers[-1], ref.markers[-1]) # same markers def close(a, b, name): np.testing.assert_allclose(a, b, rtol=0, atol=rtol * max(np.max(np.abs(b)), 1e-300), err_msg=name) - close(gpu.markers[:, 0:3], ref.markers[:, 0:3], "positions") - close(gpu.markers[:, 3:6], ref.markers[:, 3:6], "velocities") - np.testing.assert_allclose(gpu.markers[:, 6], ref.markers[:, 6], rtol=1e-13, atol=0, err_msg="weights") + close(gpu.markers[0:3], ref.markers[0:3], "positions") + close(gpu.markers[3:6], ref.markers[3:6], "velocities") + np.testing.assert_allclose(gpu.markers[6], ref.markers[6], rtol=1e-13, atol=0, err_msg="weights") for name, coeffs in ref.fields.items(): assert gpu.fields[name].shape == coeffs.shape, name close(gpu.fields[name], coeffs, name) diff --git a/src/struphy/models/tests/test_scalars.py b/src/struphy/models/tests/test_scalars.py index cace1a698..01d888b93 100644 --- a/src/struphy/models/tests/test_scalars.py +++ b/src/struphy/models/tests/test_scalars.py @@ -33,20 +33,20 @@ def test_nested_sum_follows_its_summands(): class _DriftKineticMarkers: - """Particles5D-like markers: weight in column 5, mu * |B_0| in column 8, plus holes and ghosts.""" + """Particles5D-like markers: weight in row 5, mu * |B_0| in row 8, plus holes and ghosts.""" def __init__(self, Np): self.Np = Np - self.markers = xp.zeros((Np + 4, 10), dtype=float) - self.markers[:Np, 5] = 2.0 / Np # density 2; the weights carry the 1/Np of the Monte-Carlo estimate - self.markers[:Np, 8] = 3.0 - self.markers[Np:, 5:9] = 1.0 # holes and ghosts, which must not be counted + self.markers = xp.zeros((10, Np + 4), dtype=float) + self.markers[5, :Np] = 2.0 / Np # density 2; the weights carry the 1/Np of the Monte-Carlo estimate + self.markers[8, :Np] = 3.0 + self.markers[5:9, Np:] = 1.0 # holes and ghosts, which must not be counted self.valid_mks = xp.zeros(Np + 4, dtype=bool) self.valid_mks[:Np] = True @property def markers_wo_holes_and_ghost(self): - return self.markers[self.valid_mks] + return self.markers[:, self.valid_mks] def save_magnetic_background_energy(self): pass @@ -69,14 +69,14 @@ class _MovingParticles: """Markers whose accessors return copies, as `Particles.velocities` does (it is fancy-indexed).""" def __init__(self): - self._markers = xp.ones((4, 3), dtype=float) + self._markers = xp.ones((3, 4), dtype=float) self._weights = xp.ones(4, dtype=float) self._valid = xp.ones(4, dtype=bool) self.Np = 4 @property def velocities(self): - return self._markers[self._valid] # boolean indexing, as Particles does: a copy, not a view + return self._markers[:, self._valid] # boolean indexing, as Particles does: a copy, not a view @property def weights(self): @@ -116,10 +116,10 @@ def test_kinetic_energy_follows_the_markers(): class _Fake5DParticles(Particles5D): - """Particles5D markers with velocity columns (v_par, mu), without the equilibrium needed to set them up.""" + """Particles5D markers with velocity rows (v_par, mu), without the equilibrium needed to set them up.""" def __init__(self, v_par, mu, weights): - self._velocities = xp.stack((v_par, mu), axis=1) + self._velocities = xp.stack((v_par, mu)) self._weights = weights @property @@ -132,7 +132,7 @@ def weights(self): def test_kinetic_energy_ignores_magnetic_moment_of_5d_markers(): - """For Particles5D the second velocity column is mu, which must not enter 0.5 * sum(w v^2).""" + """For Particles5D the second velocity row is mu, which must not enter 0.5 * sum(w v^2).""" v_par = xp.array([1.0, -2.0, 3.0]) mu = xp.array([5.0, 7.0, 11.0]) weights = xp.array([0.5, 1.0, 2.0]) @@ -145,19 +145,19 @@ def test_kinetic_energy_ignores_magnetic_moment_of_5d_markers(): class _MarkersWithGhosts: - """Particles5D-like markers (weight in column 5, mu * |B| in column 8), plus holes and ghost copies of the first markers.""" + """Particles5D-like markers (weight in row 5, mu * |B| in row 8), plus holes and ghost copies of the first markers.""" def __init__(self, n_valid, n_ghosts): n_rows = n_valid + n_ghosts + 2 self.Np = n_valid - self.markers = xp.zeros((n_rows, 10), dtype=float) - self.markers[:n_valid, 5] = xp.linspace(0.5, 1.5, n_valid) - self.markers[:n_valid, 8] = 3.0 - self.markers[n_valid : n_valid + n_ghosts] = self.markers[:n_ghosts] - self.markers[n_valid : n_valid + n_ghosts, -1] = -2.0 # ghosts - self.markers[n_valid + n_ghosts :, :-1] = -1.0 # holes - self.holes = self.markers[:, 0] == -1.0 - self.valid_mks = ~xp.logical_or(self.holes, self.markers[:, -1] == -2.0) + self.markers = xp.zeros((10, n_rows), dtype=float) + self.markers[5, :n_valid] = xp.linspace(0.5, 1.5, n_valid) + self.markers[8, :n_valid] = 3.0 + self.markers[:, n_valid : n_valid + n_ghosts] = self.markers[:, :n_ghosts] + self.markers[-1, n_valid : n_valid + n_ghosts] = -2.0 # ghosts + self.markers[:-1, n_valid + n_ghosts :] = -1.0 # holes + self.holes = self.markers[0] == -1.0 + self.valid_mks = ~xp.logical_or(self.holes, self.markers[-1] == -2.0) def save_magnetic_background_energy(self): pass @@ -167,14 +167,14 @@ def save_magnetic_energy(self, PBb): class _GuidingCenterMarkers: - """Particles5D-like markers: weight in column 5, mu * |B_0| in column 8 (the first diagnostics column).""" + """Particles5D-like markers: weight in row 5, mu * |B_0| in row 8 (the first diagnostics row).""" def __init__(self, Np): self.Np = Np - self.markers = xp.zeros((Np + 2, 10), dtype=float) - self.markers[:Np, 5] = 1.0 / Np # the weights carry the 1/Np of the Monte-Carlo estimate - self.markers[:Np, 8] = 3.0 - self.markers[Np:, 5:9] = -1.0 # holes + self.markers = xp.zeros((10, Np + 2), dtype=float) + self.markers[5, :Np] = 1.0 / Np # the weights carry the 1/Np of the Monte-Carlo estimate + self.markers[8, :Np] = 3.0 + self.markers[5:9, Np:] = -1.0 # holes self.holes = xp.zeros(Np + 2, dtype=bool) self.holes[Np:] = True diff --git a/src/struphy/models/tests/verification/test_verif_IncompressibleNavierStokesSPH.py b/src/struphy/models/tests/verification/test_verif_IncompressibleNavierStokesSPH.py index 76709247c..584e5c641 100644 --- a/src/struphy/models/tests/verification/test_verif_IncompressibleNavierStokesSPH.py +++ b/src/struphy/models/tests/verification/test_verif_IncompressibleNavierStokesSPH.py @@ -363,15 +363,15 @@ def test_channel_noslip_shear_relaxation(nx: int, do_plot: bool = False): # --- DEBUG: Check marker velocities --- markers = model.fluid.density.particles.markers - # markers columns: 0:eta1, 1:eta2, 2:eta3, 3:v1, 4:v2, 5:v3, 6:weight, ... + # markers rows: 0:eta1, 1:eta2, 2:eta3, 3:v1, 4:v2, 5:v3, 6:weight, ... logger.debug( - f"Marker velocity v1: min = {markers[:, 3].min()}, max = {markers[:, 3].max()}, mean = {markers[:, 3].mean()}", + f"Marker velocity v1: min = {markers[3].min()}, max = {markers[3].max()}, mean = {markers[3].mean()}", ) logger.debug( - f"Marker velocity v2: min = {markers[:, 4].min()}, max = {markers[:, 4].max()}, mean = {markers[:, 4].mean()}", + f"Marker velocity v2: min = {markers[4].min()}, max = {markers[4].max()}, mean = {markers[4].mean()}", ) logger.debug( - f"Marker velocity v3: min = {markers[:, 5].min()}, max = {markers[:, 5].max()}, mean = {markers[:, 5].mean()}", + f"Marker velocity v3: min = {markers[5].min()}, max = {markers[5].max()}, mean = {markers[5].mean()}", ) max_j2_initial = np.max(np.abs(j2_binned[0])) diff --git a/src/struphy/pic/tests/bench_marker_layout.py b/src/struphy/pic/tests/bench_marker_layout.py index db60bdff9..74a1acc6f 100644 --- a/src/struphy/pic/tests/bench_marker_layout.py +++ b/src/struphy/pic/tests/bench_marker_layout.py @@ -1,7 +1,7 @@ """ Standalone benchmark (not a pytest test) of the operations on the marker buffer that depend on its memory layout -(see :mod:`struphy.pic.marker_layout` and struphy-hub/struphy#707): the column assignments at the start of every -push, two pushers, one accumulation kernel, hole detection, boundary conditions, the MPI sort, binning and the +(see :mod:`struphy.pic.marker_layout` and struphy-hub/struphy#707; one column per marker, C-ordered): the row +assignments at the start of every push, two pushers, one accumulation kernel, hole detection, boundary conditions, the MPI sort, binning and the gathers of valid markers used by the diagnostics. Runs on the active cunumpy backend (``CUNUMPY_BACKEND=numpy`` or ``cupy``); on the GPU every timing ends with a @@ -76,9 +76,9 @@ def run(Np: int, repeats: int, num_elements: int, degree: int, save_markers: str n_init = 3 + particles.vdim def preamble(): - """The column assignments at the start of Pusher._push.""" + """The row assignments at the start of Pusher._push.""" if xp.get_backend() == "cupy": - n_rows, n_cols = markers.shape + n_cols, n_rows = markers.shape _prepare_push_cuda()( markers, first_pusher_idx, @@ -86,9 +86,9 @@ def preamble(): n_threads=n_rows * (n_cols - 2 - first_pusher_idx), ) else: - markers[:, first_pusher_idx:first_shift_idx] = markers[:, :n_init] - markers[:, first_shift_idx:residual_idx] = 0.0 - markers[:, residual_idx:-2] = 0.0 + markers[first_pusher_idx:first_shift_idx] = markers[:n_init] + markers[first_shift_idx:residual_idx] = 0.0 + markers[residual_idx:-2] = 0.0 butcher = ButcherTableau("rk4") push_eta = Pusher( @@ -122,7 +122,7 @@ def accumulate(): components = (True, False, False, True, False, False) ops = { - "preamble (column assignments)": preamble, + "preamble (row assignments)": preamble, "push_eta_stage (RK4 pusher)": lambda: push_eta(1e-3), "push_vxb_analytic (pusher)": lambda: push_vxb(1e-3), "charge_density_0form (accumulation)": accumulate, @@ -130,30 +130,34 @@ def accumulate(): "apply_kinetic_bc (periodic)": particles.apply_kinetic_bc, "mpi_sort_markers": particles.mpi_sort_markers, "binning e1_v1": lambda: particles.binning(components, bin_edges), - "gather valid rows (markers_wo_holes_and_ghost)": lambda: particles.markers_wo_holes_and_ghost, + "gather valid markers (markers_wo_holes_and_ghost)": lambda: particles.markers_wo_holes_and_ghost, "gather weights of valid markers": lambda: particles.weights, } info = { "backend": xp.get_backend(), - "n_rows": int(markers.shape[0]), - "n_cols": int(markers.shape[1]), + "n_rows": particles.n_rows, + "n_cols": particles.n_cols, "Np": Np, - "f_contiguous": bool(markers.flags.f_contiguous), + "c_contiguous": bool(markers.flags.c_contiguous), } result = {"info": info, "ms": {name: _time(fun, repeats) for name, fun in ops.items()}} if save_markers is not None: - # the valid markers after all timed calls, sorted by ID, for a bit-for-bit comparison of two runs + # the valid markers after all timed calls, sorted by ID, for a bit-for-bit comparison of two runs; saved with + # one row per marker, as by runs with the earlier (n_rows, n_cols) buffer valid = xp.to_numpy(particles.markers_wo_holes_and_ghost) - np.save(save_markers, valid[np.argsort(valid[:, -1])]) + np.save(save_markers, valid[:, np.argsort(valid[-1])].T) return result def _print(result: dict): info = result["info"] - order = "F (column-major)" if info["f_contiguous"] else "C (row-major)" - print(f"backend {info['backend']}, markers {info['n_rows']} x {info['n_cols']} in {order}, Np = {info['Np']}") + order = "C-ordered" if info["c_contiguous"] else "not C-contiguous" + print( + f"backend {info['backend']}, markers {info['n_cols']} x {info['n_rows']} (one column per marker, {order}), " + f"Np = {info['Np']}" + ) width = max(len(name) for name in result["ms"]) for name, ms in result["ms"].items(): print(f" {name:<{width}} {ms:10.2f} ms") diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index 1f28821e5..a2396ee5b 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -47,7 +47,7 @@ class ParityCases: rtol: float = 1e-12 atol: float = 0.0 n_threads: Callable[[tuple], int] | None = None - """Launch size, if it is not one thread per row of the first array.""" + """Launch size, if it is not one thread per row (first axis) of the first array.""" # pyccel argument class name -> the fields of its CUDA struct, the data of an argument object (the pyccel scratch @@ -103,7 +103,7 @@ def push_weights_with_efield_lin_va_args(bc): def reflect_args(axis): """Reverse one logical velocity component of 10 markers in a Cuboid.""" - markers = np.asfortranarray(np.random.default_rng(32).random((140, 8))) + markers = np.random.default_rng(32).random((8, 140)) outside_inds = np.arange(139, 129, -1, dtype=np.int64) return (xp.asarray(markers), Cuboid(r1=2.0, r2=3.0, r3=4.0).args_domain, xp.asarray(outside_inds), axis) @@ -117,11 +117,12 @@ def reflect_args(axis): def bstar_parallel_3form_args(case): - """Markers with shifts in [0, 1), so eta + shift wraps around, a negative eta_n in row 0 and a hole in row 2.""" + """Markers with shifts in [0, 1), so eta + shift wraps around, a negative eta_n in marker 0 and a hole in + marker 2.""" domain, alpha, output_column = case args_markers, _ = marker_arguments((0, 0, 0)) - args_markers.markers[0, 8] = -0.75 # eta_n[0] = mod(-0.75, 1.0) = 0.25 (not 0, where some Jacobians are singular) - args_markers.markers[2, 0] = -1.0 # a hole: pyccel tests markers[ip, 0] == -1, not valid_mks + args_markers.markers[8, 0] = -0.75 # eta_n[0] = mod(-0.75, 1.0) = 0.25 (not 0, where some Jacobians are singular) + args_markers.markers[0, 2] = -1.0 # a hole: pyccel tests markers[0, ip] == -1, not valid_mks return ( xp.asarray(alpha), xp.asarray([output_column], dtype=np.int64), @@ -155,7 +156,7 @@ def linear_vlasov_ampere_args(domain): def vlasov_maxwell_args(domain): - """129 markers (a hole in row 0, a boundary particle in row 1, which is accumulated as in pyccel) in domain + """129 markers (a hole, marker 0, and a boundary particle, marker 1, which is accumulated as in pyccel) in domain `domain` of geometry_domain(); zeroed V1 data.""" args_markers, _ = marker_arguments((0, 0, 0)) args_markers.markers[0, 0] = -1.0 @@ -167,7 +168,7 @@ def vlasov_maxwell_args(domain): def eval_spline_mpi_markers_args(kind): """Spline values at 129 markers, one of them flagged.""" - markers = np.asfortranarray(np.random.default_rng(5).random((129, 3))) + markers = np.random.default_rng(5).random((3, 129)) markers[0, 0] = -1.0 # not on the process domain: skipped return (xp.asarray(markers), *spline_evaluation_arguments(kind), xp.zeros(129)) @@ -204,7 +205,7 @@ def eval_spline_mpi_tensor_product_fixed_args(kind): def kernel_evaluate_pic_args(case): domain, kind_coeff, remove_outside, avoid_round_off = case markers = logical_markers() - mat_f = np.random.default_rng(6).random((markers.shape[0], 3, 3)) + mat_f = np.random.default_rng(6).random((markers.shape[1], 3, 3)) return ( xp.asarray(markers), kind_coeff, @@ -263,15 +264,16 @@ def kernel_evaluate_args(case): def kernel_pullpush_pic_args(case): - """Every third case passes `a` without holes (one row per inside marker).""" + """Every third case passes `a` without holes (one row per inside marker); `a` and `out` have one row per + marker.""" domain, kind_transform, kind_fun, a_has_holes, remove_outside = case markers = logical_markers() rng = np.random.default_rng(9) n_comp = 1 if kind_fun < 10 else 3 - a = rng.normal(size=(markers.shape[0], n_comp)) + a = rng.normal(size=(markers.shape[1], n_comp)) if not a_has_holes: a = a[inside_logical_cube(markers)] - out = rng.random((markers.shape[0], 3)) + out = rng.random((markers.shape[1], 3)) return ( xp.asarray(a), xp.asarray(markers), @@ -320,9 +322,14 @@ def kernel_pullpush_args(case): SPLINE_KINDS = ((0, 0, 0), (1, 0, 1)) -def rows_of(index): - """One thread per row of argument `index`.""" - return lambda args: args[index].shape[0] +def markers_of(index): + """One thread per marker of argument `index`, a marker argument object.""" + return lambda args: args[index].n_markers + + +def columns_of(index): + """One thread per column of argument `index`, an array of markers (one column per marker).""" + return lambda args: args[index].shape[1] def size_of(index, per_thread=1): @@ -339,44 +346,54 @@ def size_of(index, per_thread=1): for stage in range(ButcherTableau(method).n_stages) ], push_eta_stage_args, + n_threads=markers_of(2), **PUSHER_TOLERANCES, ), "push_v_with_efield": ParityCases( BOUNDARY_CONDITIONS, lambda bc: (0.2, 0, *marker_arguments(bc), derham_arguments(), *spline_coefficients(), 0.7), + n_threads=markers_of(2), **PUSHER_TOLERANCES, ), "push_vxb_analytic": ParityCases( BOUNDARY_CONDITIONS, lambda bc: (0.2, 0, *marker_arguments(bc), derham_arguments(), *spline_coefficients()), + n_threads=markers_of(2), **PUSHER_TOLERANCES, ), "push_vxb_implicit": ParityCases( BOUNDARY_CONDITIONS, lambda bc: (0.2, 0, *marker_arguments(bc), derham_arguments(), *spline_coefficients()), + n_threads=markers_of(2), **PUSHER_TOLERANCES, ), "push_weights_with_efield_lin_va": ParityCases( - BOUNDARY_CONDITIONS, push_weights_with_efield_lin_va_args, **PUSHER_TOLERANCES + BOUNDARY_CONDITIONS, push_weights_with_efield_lin_va_args, n_threads=markers_of(2), **PUSHER_TOLERANCES ), "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), - # one thread per marker row; the first array, alpha, does not tell the launch size + # one thread per marker "bstar_parallel_3form": ParityCases( BSTAR_PARALLEL_3FORM_CASES, bstar_parallel_3form_args, - n_threads=lambda args: args[2].n_markers, + n_threads=markers_of(2), **GEOMETRY_TOLERANCES, ), # atomics add in another order than the serial loop - "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), + "charge_density_0form": ParityCases( + ((0, 0, 0), (2, 0, 1)), charge_density_0form_args, n_threads=markers_of(0), rtol=1e-12, atol=1e-13 + ), # Cuboid, Colella and HollowTorus (non-diagonal DF) and a 3d spline mapping, by index into geometry_domain(); # atomics add in another order than the serial loop, and the entries reach 1e4 (atol for cancelling sums) - "linear_vlasov_ampere": ParityCases((0, 2, 5, 12), linear_vlasov_ampere_args, rtol=1e-12, atol=1e-10), + "linear_vlasov_ampere": ParityCases( + (0, 2, 5, 12), linear_vlasov_ampere_args, n_threads=markers_of(0), rtol=1e-12, atol=1e-10 + ), # Cuboid, then Colella, HollowTorus, ShafranovDshapedCylinder (non-diagonal DF) and a 3d spline mapping, by index into # geometry_domain(); atomics add in another order than the serial loop, and G^{-1} of the spline mapping makes # entries of 6e5 (atol for cancelling sums) - "vlasov_maxwell": ParityCases((0, 2, 5, 9, 12), vlasov_maxwell_args, rtol=1e-12, atol=1e-8), - "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), + "vlasov_maxwell": ParityCases( + (0, 2, 5, 9, 12), vlasov_maxwell_args, n_threads=markers_of(0), rtol=1e-12, atol=1e-8 + ), + "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args, n_threads=columns_of(0)), "eval_spline_mpi_matrix": ParityCases(SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=False), n_threads=size_of(-1)), "eval_spline_mpi_sparse_meshgrid": ParityCases( SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=True), n_threads=size_of(-1) @@ -384,14 +401,16 @@ def size_of(index, per_thread=1): "eval_spline_mpi_tensor_product_fixed": ParityCases( SPLINE_KINDS, eval_spline_mpi_tensor_product_fixed_args, n_threads=size_of(-1) ), - "kernel_evaluate_pic": ParityCases(KERNEL_EVALUATE_PIC_CASES, kernel_evaluate_pic_args, **GEOMETRY_TOLERANCES), + "kernel_evaluate_pic": ParityCases( + KERNEL_EVALUATE_PIC_CASES, kernel_evaluate_pic_args, n_threads=columns_of(0), **GEOMETRY_TOLERANCES + ), # one thread per grid point: the first three axes of mat_f, shape (n1, n2, n3, 3, 3) "kernel_evaluate": ParityCases( KERNEL_EVALUATE_CASES, kernel_evaluate_args, n_threads=size_of(5, per_thread=9), **GEOMETRY_TOLERANCES ), - # one thread per marker row; the first array, a, has fewer rows when it has no holes + # one thread per marker (column of the markers); the first array, a, has fewer rows when it has no holes "kernel_pullpush_pic": ParityCases( - KERNEL_PULLPUSH_PIC_CASES, kernel_pullpush_pic_args, n_threads=rows_of(1), **GEOMETRY_TOLERANCES + KERNEL_PULLPUSH_PIC_CASES, kernel_pullpush_pic_args, n_threads=columns_of(1), **GEOMETRY_TOLERANCES ), # one thread per grid point: the first three axes of out, shape (n1, n2, n3, 3) "kernel_pullpush": ParityCases( diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index ae30c272f..10d28bea4 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -17,7 +17,6 @@ from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.ode.utils import ButcherTableau -from struphy.pic.marker_layout import as_marker_array N_MARKERS = 129 # not a multiple of the block size N_COLS = 25 @@ -26,21 +25,21 @@ def marker_arguments(bc): - """Markers with a hole (row 0) and a boundary particle (row 1), in a Cuboid.""" + """Markers with a hole (marker 0) and a boundary particle (marker 1), in a Cuboid.""" rng = np.random.default_rng(7) - markers = rng.random((N_MARKERS, N_COLS)) - markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) - markers[:, 8:14] = markers[:, :6] - markers[:, 18:21] = 0.0 - # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_pusher_idx (8), - # accumulation kernels test column 0. - markers[0, :] = -1.0 - markers[1, -1] = -2.0 + markers = rng.random((N_COLS, N_MARKERS)) + markers[3:6] = rng.uniform(-2, 2, (3, N_MARKERS)) + markers[8:14] = markers[:6] + markers[18:21] = 0.0 + # Marker 0 is a hole. Struphy marks every entry of a hole with -1: pushers test row first_pusher_idx (8), + # accumulation kernels test row 0. + markers[:, 0] = -1.0 + markers[-1, 1] = -2.0 valid = np.ones(N_MARKERS, dtype=bool) valid[:2] = False args_class = CudaMarkerArguments if xp.get_backend() == "cupy" else MarkerArguments args_markers = args_class( - as_marker_array(xp.asarray(markers)), + xp.asarray(markers), xp.asarray(valid), N_MARKERS, *MARKER_INDICES, @@ -214,15 +213,15 @@ def geometry_domain(index): def logical_markers(n=N_MARKERS, seed=5): - """Logical points in marker format, shape (n, 7), column-major like the marker buffer, away from the poles - (eta1 >= 0.1). + """Logical points in marker format, shape (7, n), one column per point like the marker buffer, away from the + poles (eta1 >= 0.1). - Row 0 is a hole (-1), row 1 is outside in eta1 and row 2 in eta3, so ``remove_outside`` changes the result. + Point 0 is a hole (-1), point 1 is outside in eta1 and point 2 in eta3, so ``remove_outside`` changes the result. """ rng = np.random.default_rng(seed) - markers = rng.random((n, 7)) - markers[:, 0] = rng.uniform(0.1, 0.95, n) - markers[0, :] = -1.0 - markers[1, 0] = 1.2 + markers = rng.random((7, n)) + markers[0] = rng.uniform(0.1, 0.95, n) + markers[:, 0] = -1.0 + markers[0, 1] = 1.2 markers[2, 2] = -0.3 - return np.asfortranarray(markers) + return markers diff --git a/src/struphy/pic/tests/test_accum_gc_gradB.py b/src/struphy/pic/tests/test_accum_gc_gradB.py index 0fce51a4a..7348e27da 100644 --- a/src/struphy/pic/tests/test_accum_gc_gradB.py +++ b/src/struphy/pic/tests/test_accum_gc_gradB.py @@ -36,11 +36,11 @@ def rand(space): zero = derham.coeff_spaces["Hcurl"].zeros() n_mk = 5 - markers = np.zeros((n_mk, 40), order="F") - markers[:, :3] = rng.uniform(0.05, 0.95, (n_mk, 3)) - markers[:, 3] = rng.standard_normal(n_mk) - markers[:, 4] = rng.uniform(0.1, 1.0, n_mk) - markers[:, 5] = rng.uniform(0.5, 1.5, n_mk) + markers = np.zeros((40, n_mk)) + markers[:3] = rng.uniform(0.05, 0.95, (3, n_mk)) + markers[3] = rng.standard_normal(n_mk) + 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), False ) diff --git a/src/struphy/pic/tests/test_accum_pressure_coupling.py b/src/struphy/pic/tests/test_accum_pressure_coupling.py index b4abef14d..3f6ca54c3 100644 --- a/src/struphy/pic/tests/test_accum_pressure_coupling.py +++ b/src/struphy/pic/tests/test_accum_pressure_coupling.py @@ -54,11 +54,11 @@ def accumulate(): valid = particles.valid_mks # make the current weights distinct from w0 (as in delta-f) - particles.markers[valid, idx_w] *= 1.0 + xp.arange(xp.count_nonzero(valid)) / Np + particles.markers[idx_w, valid] *= 1.0 + xp.arange(xp.count_nonzero(valid)) / Np mats_ref, vecs_ref = accumulate() # garbage in w0 must not change the result - particles.markers[valid, idx_w0] = -17.0 + particles.markers[idx_w0, valid] = -17.0 mats, vecs = accumulate() assert any(xp.max(xp.abs(v)) > 0.0 for v in vecs_ref) diff --git a/src/struphy/pic/tests/test_accum_vec_H1.py b/src/struphy/pic/tests/test_accum_vec_H1.py index c98436a81..fbfa95153 100644 --- a/src/struphy/pic/tests/test_accum_vec_H1.py +++ b/src/struphy/pic/tests/test_accum_vec_H1.py @@ -466,7 +466,7 @@ def u_xyz(x, y, z): # correct valid_mks mask (excludes holes and ghosts). # # ------------------------------------------------------------------ # eta = particles.positions - particles.markers[particles.valid_mks, particles.first_free_idx] = n_xyz(eta[:, 0], eta[:, 1], eta[:, 2]) + particles.markers[particles.first_free_idx, particles.valid_mks] = n_xyz(eta[0], eta[1], eta[2]) # ------------------------------------------------------------------ # # Accumulate the weak-divergence-1form RHS vector V^1. # diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index a2548df03..7a8620a57 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -116,17 +116,17 @@ def test_emulated_v1_symm_filler(): @requires_compiler @pytest.mark.parametrize("n_rows", [96, 129]) def test_emulated_prepare_push(n_rows): - """The fused preamble of Pusher._push (prepare_push_cuda.cu) writes the column-major buffer like the three column - assignments, and leaves the columns before first_pusher_idx and the last two columns alone.""" + """The fused preamble of Pusher._push (prepare_push_cuda.cu) writes the marker buffer (one column per marker) like + the three row assignments, and leaves the rows before first_pusher_idx and the last two rows alone.""" from struphy.pic.marker_layout import marker_array from struphy.pic.pushing.pusher import _prepare_push_cuda n_cols, first_pusher_idx, n_init = 26, 8, 6 - markers = marker_array(n_rows, n_cols) - markers[:] = np.random.default_rng(11).random((n_rows, n_cols)) + markers = marker_array(n_cols, n_rows) + markers[:] = np.random.default_rng(11).random((n_cols, n_rows)) expected = markers.copy() - expected[:, first_pusher_idx : first_pusher_idx + n_init] = expected[:, :n_init] - expected[:, first_pusher_idx + n_init : -2] = 0.0 + expected[first_pusher_idx : first_pusher_idx + n_init] = expected[:n_init] + expected[first_pusher_idx + n_init : -2] = 0.0 width = n_cols - 2 - first_pusher_idx emulate_struct_kernel(_prepare_push_cuda(), markers, first_pusher_idx, n_init, n_threads=n_rows * width) diff --git a/src/struphy/pic/tests/test_device_helpers.py b/src/struphy/pic/tests/test_device_helpers.py index d0bfaaaa7..8c8f7675b 100644 --- a/src/struphy/pic/tests/test_device_helpers.py +++ b/src/struphy/pic/tests/test_device_helpers.py @@ -94,7 +94,7 @@ def test_boundary_helpers(bc, newton): with cunumpy.use_backend(backend): m, d = make_arguments(129) m.bc_type[:] = cunumpy.asarray(bc) - m.markers[:, :3] = cunumpy.asarray(np.random.default_rng(2).uniform(-0.5, 1.5, (129, 3))) + m.markers[:3] = cunumpy.asarray(np.random.default_rng(2).uniform(-0.5, 1.5, (3, 129))) if backend == "numpy": for i in range(129): apply_kinetic_bc_marker(i, m, d, newton) @@ -291,7 +291,7 @@ def metric_helper_case(domain_index, avoid_round_off): "bool avoid_round_off, const DomainArgs& args)", **CUDA_OPTIONS, ) - points = logical_markers()[3:, :3] # inside the logical cube + points = logical_markers()[:3, 3:].T # inside the logical cube, one row per point rows = [(p, kind, entry) for p in points for kind in range(6) for entry in range(9)] eta = np.array([r[0] for r in rows]) kinds = np.array([r[1] for r in rows], dtype=np.int32) @@ -339,7 +339,7 @@ def transform_helper_case(domain_index): ) kinds = [(0, k) for k in (0, 1, 10, 11, 12)] + [(1, k) for k in (0, 1, 10, 11, 12)] kinds += [(2, k) for k in (0, 1, *range(10, 22))] - points = logical_markers()[3:20, :3] + points = logical_markers()[:3, 3:20].T # one row per point values = np.random.default_rng(12).normal(size=(len(points), 3)) rows = [(i, t, k, e) for i in range(len(points)) for t, k in kinds for e in range(3 if k >= 10 else 1)] diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 1de1a0f4b..1a6ca591e 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -105,7 +105,7 @@ def test_draw(num_elements, degree, bcs, mapping, ppc=10): comm.Barrier() logger.info("Number of particles w/wo holes on each process before sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.markers.shape[0]}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") # sort particles according to domain decomposition comm.Barrier() @@ -113,20 +113,20 @@ def test_draw(num_elements, degree, bcs, mapping, ppc=10): comm.Barrier() logger.info("Number of particles w/wo holes on each process after sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.markers.shape[0]}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") # are all markers in the correct domain? conds = xp.logical_and( - particles.markers[:, :3] > derham.domain_array[rank, 0::3], - particles.markers[:, :3] < derham.domain_array[rank, 1::3], + particles.markers[:3] > derham.domain_array[rank, 0::3, None], + particles.markers[:3] < derham.domain_array[rank, 1::3, None], ) - holes = particles.markers[:, 0] == -1.0 - stay = xp.all(conds, axis=1) + holes = particles.markers[0] == -1.0 + stay = xp.all(conds, axis=0) - error_mks = particles.markers[xp.logical_and(~stay, ~holes)] + error_mks = particles.markers[:, xp.logical_and(~stay, ~holes)] assert error_mks.size == 0, ( - f"rank {rank} | markers not on correct process: {xp.nonzero(xp.logical_and(~stay, ~holes))} \n corresponding positions:\n {error_mks[:, :3]}" + f"rank {rank} | markers not on correct process: {xp.nonzero(xp.logical_and(~stay, ~holes))} \n corresponding positions:\n {error_mks[:3]}" ) @@ -166,9 +166,9 @@ def test_marker_ids(loading, Np=1000): particles.draw_markers(sort=False) - valid = particles.markers[~particles.holes] - ids = valid[:, -1] - n_at_origin = int(xp.sum(xp.all(valid[:, :3] == 0.0, axis=1))) + valid = particles.markers[:, ~particles.holes] + ids = valid[-1] + n_at_origin = int(xp.sum(xp.all(valid[:3] == 0.0, axis=0))) if particles.mpi_size > 1: ids = xp.concatenate(particles.mpi_comm.allgather(ids)) n_at_origin = particles.mpi_comm.allreduce(n_at_origin) @@ -216,10 +216,10 @@ def test_draw_without_boxes_on_rank_domain(loading, sort, Np=1000): particles.draw_markers(sort=sort) - positions = particles.markers_wo_holes[:, :3] - on_rank = xp.all(particles._is_on_domain_of_rank(positions, particles.mpi_rank), axis=1) + positions = particles.markers_wo_holes[:3] + on_rank = xp.all(particles._is_on_domain_of_rank(positions, particles.mpi_rank), axis=0) n_off_rank = int(xp.sum(~on_rank)) - n_markers = positions.shape[0] + n_markers = positions.shape[1] if particles.mpi_size > 1: n_markers = particles.mpi_comm.allreduce(n_markers) @@ -260,10 +260,10 @@ def test_draw_disc_on_own_process(loading, ppc=50): 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 + assert comm.allreduce(positions.shape[1]) == 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 + 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) diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index cccb16898..164f53d05 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -20,7 +20,6 @@ from struphy.kernel_arguments.pusher_args_kernels import DerhamArguments, DomainArguments, MarkerArguments from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.kernel_arguments.spline_args_kernels import SplineArguments -from struphy.pic.marker_layout import as_marker_array from struphy.pic.tests.kernel_test_args import N_GEOMETRY_DOMAINS from struphy.utils.cuda_arguments import ( CUDA_OPTIONS, @@ -54,7 +53,7 @@ def make_arguments(n_markers, seed=0): rng = np.random.default_rng(seed) - markers = as_marker_array(cunumpy.asarray(rng.random((n_markers, N_COLS)))) + markers = cunumpy.asarray(rng.random((N_COLS, n_markers))) 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 @@ -108,17 +107,13 @@ def test_generated_spline_header(tmp_path): @pytest.mark.parametrize("cuda_class, source, class_name, cuda_only, header", ARGUMENT_PAIRS) -def test_argument_classes_correspond(cuda_class, source, class_name, cuda_only, header, tmp_path): +def test_argument_classes_correspond(cuda_class, source, class_name, cuda_only, header): """Each CUDA argument class mirrors its pyccel class: same constructor, same attributes in the same order.""" assert cuda_class.__name__ == "Cuda" + class_name cuda_parameters = list(inspect.signature(cuda_class.__init__).parameters)[1:] assert cuda_parameters == pyccel_init_parameters(source, class_name) # attributes set from constructor arguments; derived ones (e.g. n_markers, bn1) are skipped by the parser. - # cunumpy's parser does not read pyccel's memory-order suffix ("float[:, :](order=F)" of the column-major - # markers), so it is removed in a copy of the source; only the names are compared here. - pyccel_source = tmp_path / source - pyccel_source.write_text((ARGS_DIR / source).read_text().replace("(order=F)", "")) - pyccel_fields = [field.name for field in CudaStruct.from_pyccel_class(pyccel_source, class_name).fields] + pyccel_fields = [field.name for field in CudaStruct.from_pyccel_class(ARGS_DIR / source, class_name).fields] cuda_fields = [field.name for field in cuda_class.struct.fields] assert [name for name in cuda_fields if name not in cuda_only] == pyccel_fields @@ -154,7 +149,7 @@ def test_cuda_struct_layout(cuda_class, source, class_name, cuda_only, header): @requires_cupy def test_cuda_classes_reject_host_arrays(): with cunumpy.use_backend("numpy"): - markers = np.zeros((10, N_COLS)) + markers = np.zeros((N_COLS, 10)) with pytest.raises(TypeError, match="CuPy array"): CudaMarkerArguments( markers, np.ones(10, dtype=bool), 10, *MARKER_INDICES, np.zeros(3, dtype=np.int64), False @@ -235,7 +230,7 @@ def test_geometry_evaluation_on_cupy(domain_index): """Geometry evaluations run on CuPy for every analytic mapping and four spline mappings and agree with NumPy.""" from struphy.pic.tests.kernel_test_args import geometry_domain - markers = np.random.default_rng(3).uniform(-0.1, 1.0, (50, 3)) + markers = np.random.default_rng(3).uniform(-0.1, 1.0, (3, 50)) eta = (np.linspace(0.1, 0.9, 5), np.linspace(0.0, 1.0, 4), np.linspace(0.2, 0.8, 3)) results = [] for backend in ("numpy", "cupy"): diff --git a/src/struphy/pic/tests/test_kernel_setup.py b/src/struphy/pic/tests/test_kernel_setup.py index ae3de77f1..bbd8e2bf7 100644 --- a/src/struphy/pic/tests/test_kernel_setup.py +++ b/src/struphy/pic/tests/test_kernel_setup.py @@ -33,11 +33,11 @@ @pytest.fixture def marker_args(): - markers = np.zeros((2, 40), order="F") - markers[0, :7] = (0.4, 0.3, 0.2, 2.0, 4.0, 6.0, 2.0) - 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. + markers = np.zeros((40, 2)) + markers[:7, 0] = (0.4, 0.3, 0.2, 2.0, 4.0, 6.0, 2.0) + markers[8:14, 0] = markers[:6, 0] + markers[11, 0] = 1.0 # Initial parallel velocity. + markers[0, 1] = -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), False ) @@ -81,7 +81,7 @@ def test_guiding_center_destinations(marker_args, derham_args, kernel, values, e reference = before.copy() for component, index in enumerate(indices): if index is not None: - reference[0, index] = expected[component] + reference[index, 0] = expected[component] np.testing.assert_allclose(marker_args.markers, reference) @@ -100,17 +100,17 @@ def test_sph_vector_destinations(marker_args, kernel): reference = before.copy() # One particle: density = weight * W(0) = 2 * (1 / h) = 4. if kernel.name == "sph_pressure_coeffs": - reference[0, 30], reference[0, 24] = 4.0, 2.0 * 4.0 ** (-1.0 / 3.0) + reference[30, 0], reference[24, 0] = 4.0, 2.0 * 4.0 ** (-1.0 / 3.0) else: - reference[0, 30], reference[0, 24] = 1.0, 3.0 + reference[30, 0], reference[24, 0] = 1.0, 3.0 np.testing.assert_allclose(marker_args.markers, reference) def test_sph_tensor_destinations(marker_args): # Second particle at eta1 + 0.2 (inside h = 0.5) with the same velocity coefficients. - marker_args.markers[1] = marker_args.markers[0] - marker_args.markers[1, 0] = 0.6 - marker_args.markers[:, 18:21] = (1.0, 2.0, 3.0) + marker_args.markers[:, 1] = marker_args.markers[:, 0] + marker_args.markers[0, 1] = 0.6 + marker_args.markers[18:21] = np.array([1.0, 2.0, 3.0])[:, None] marker_args.valid_mks[1] = True boxes = np.array([[0, 1, -1], [-1, -1, -1]]) neighbours = np.ones((1, 27), dtype=int) @@ -131,8 +131,8 @@ def test_sph_tensor_destinations(marker_args): reference = before.copy() for value, index in zip(tensor.flat, indices): if index is not None: - reference[0, index] = value - reference[1, index] = -value + reference[index, 0] = value + reference[index, 1] = -value np.testing.assert_allclose(marker_args.markers, reference) @@ -161,16 +161,16 @@ def test_sph_kernel_gradients_vanish_at_zero(kernel_type, point): ) def test_gc_discrete_gradient_residual_at_zero_velocity(derham_args, kernel, n_fields): # A marker at rest (v = mu = 0, no E-field) stays put; the residual must be 0, not NaN (issue #589). - markers = np.zeros((1, 40), order="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. + markers = np.zeros((40, 1)) + markers[:3, 0] = (0.4, 0.3, 0.2) + markers[8:11, 0] = markers[:3, 0] + markers[18, 0] = markers[21, 0] = 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), 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)) + assert markers[16, 0] == 0.0 + np.testing.assert_array_equal(markers[:4, 0], (0.4, 0.3, 0.2, 0.0)) def test_sph_tensor_mapped_domain(marker_args): @@ -178,9 +178,9 @@ def test_sph_tensor_mapped_domain(marker_args): # stored columns are in Piola form det(DF) * tensor @ DF^{-T}. # Only the neighbour at eta1 + 0.2 carries velocity coefficients, so the result does not # depend on the kernel gradient at zero; its linear 1D kernel gradient is 1 / h**2 = 4. - marker_args.markers[1] = marker_args.markers[0] - marker_args.markers[1, 0] = 0.6 - marker_args.markers[1, 18:21] = (1.0, 2.0, 3.0) + marker_args.markers[:, 1] = marker_args.markers[:, 0] + marker_args.markers[0, 1] = 0.6 + marker_args.markers[18:21, 1] = (1.0, 2.0, 3.0) marker_args.valid_mks[1] = True boxes = np.array([[0, 1, -1], [-1, -1, -1]]) neighbours = np.ones((1, 27), dtype=int) @@ -200,7 +200,7 @@ def test_sph_tensor_mapped_domain(marker_args): symmetric = (gradient + gradient.T) / 2.0 tensor = -2.0 * 3.0 * (2.0 / density) * (symmetric - np.eye(3) * np.trace(symmetric) / 3.0) piola = np.linalg.det(jac) * tensor @ jac_inv.T - np.testing.assert_allclose(marker_args.markers[0, 22:31], piola.flat) + np.testing.assert_allclose(marker_args.markers[22:31, 0], piola.flat) @pytest.mark.parametrize("indices", [(), (1, 2), (-1,), (True,), (1.5,), (None,), (1, 1, 2)]) @@ -242,16 +242,16 @@ def test_pusher_evaluation_order_and_live_arguments(marker_args): def initialize(alpha, indices, markers, domain, values): events.append("init") - markers.markers[0, indices[0]] = values[0] + markers.markers[indices[0], 0] = values[0] def evaluate(alpha, indices, markers, domain): events.append("eval") - markers.markers[0, indices[0]] = markers.markers[0, 20] + alpha[0] + markers.markers[indices[0], 0] = markers.markers[20, 0] + alpha[0] def push(dt, stage, markers, domain): events.append("push") - assert markers.markers[0, 21] == field[0] + 0.5 - markers.markers[0, markers.residual_idx] = 1.0 # Force both iterations. + assert markers.markers[21, 0] == field[0] + 0.5 + markers.markers[markers.residual_idx, 0] = 1.0 # Force both iterations. def sort(**kwargs): # Final sort after the last stage uses current positions (no alpha). @@ -263,6 +263,7 @@ def sort(**kwargs): particles = SimpleNamespace( markers=marker_args.markers, args_markers=marker_args, + n_rows=2, n_cols=40, vdim=3, first_pusher_idx=8, diff --git a/src/struphy/pic/tests/test_load_external.py b/src/struphy/pic/tests/test_load_external.py index 32b0d36bc..25c70fe7f 100644 --- a/src/struphy/pic/tests/test_load_external.py +++ b/src/struphy/pic/tests/test_load_external.py @@ -22,7 +22,7 @@ def test_load_external(tmp_path): path = tmp_path / "markers.hdf5" with h5py.File(path, "w") as file: - file.create_dataset("markers", data=ref.markers[:n_mks]) + file.create_dataset("markers", data=ref.markers[:, :n_mks]) ext = Particles6D( loading_params=LoadingParameters(Np=1000, loading="external", dir_external=str(path)), @@ -31,8 +31,8 @@ def test_load_external(tmp_path): ) ext.draw_markers(sort=False) - assert xp.array_equal(ext.markers[:n_mks], ref.markers[:n_mks]) - assert xp.all(ext.markers[n_mks:] == -1.0) + assert xp.array_equal(ext.markers[:, :n_mks], ref.markers[:, :n_mks]) + assert xp.all(ext.markers[:, n_mks:] == -1.0) if __name__ == "__main__": diff --git a/src/struphy/pic/tests/test_marker_layout.py b/src/struphy/pic/tests/test_marker_layout.py index 92372672d..0cd453b2c 100644 --- a/src/struphy/pic/tests/test_marker_layout.py +++ b/src/struphy/pic/tests/test_marker_layout.py @@ -1,4 +1,5 @@ -"""The marker buffer is column-major (struct-of-arrays), see :mod:`struphy.pic.marker_layout`. +"""The marker buffer has one column per marker and is C-ordered (struct-of-arrays), see +:mod:`struphy.pic.marker_layout`. The CuPy cases run on a GPU, or without one on cunumpy's fake CuPy with every CUDA launch emulated on the CPU:: @@ -13,18 +14,14 @@ import pytest from struphy.models.tests.gpu_e2e import device_context, requires_cupy_backend -from struphy.pic.marker_layout import ( - ROW_ALIGNMENT, - as_marker_array, - check_marker_layout, - has_marker_layout, - marker_array, - padded_n_rows, -) +from struphy.pic.marker_layout import ROW_ALIGNMENT, marker_array, padded_n_rows # CuPy on a GPU, or cunumpy's fake CuPy BACKENDS = ["numpy", pytest.param("cupy", marks=requires_cupy_backend)] +N_COLS = 25 +MARKER_INDICES = (3, 6, 7, 8, 14, 17, 18, 4) # vdim, weight_idx, ..., mu_idx + def test_padded_n_rows(): assert [padded_n_rows(n) for n in (0, 1, 31, 32, 33, 64, 65)] == [0, 32, 32, 32, 64, 64, 96] @@ -32,44 +29,22 @@ def test_padded_n_rows(): @pytest.mark.parametrize("backend", BACKENDS) -def test_layout_helpers(backend): +def test_marker_array(backend): with xp.use_backend(backend): - markers = marker_array(64, 7) - assert markers.shape == (64, 7) and has_marker_layout(markers) - check_marker_layout(markers) - # no copy for an array in the marker layout, nor for a block of its columns - assert as_marker_array(markers) is markers - assert has_marker_layout(markers[:, 2:5]) - - # row gathers and .copy() are C-ordered: converted (with a copy) before they go to a kernel - rows = markers[xp.arange(0, 64, 2)] - assert not has_marker_layout(rows) - with pytest.raises(TypeError, match="Fortran-ordered"): - check_marker_layout(rows) - converted = as_marker_array(rows) - assert has_marker_layout(converted) - np.testing.assert_array_equal(xp.to_numpy(converted), xp.to_numpy(rows)) - - # a single row in C order is F-contiguous for NumPy, but its strides decrease, which pyccel rejects - row = xp.asarray([[0.2, 0.4, 0.6]]) - assert row.flags.f_contiguous and not has_marker_layout(row) - assert has_marker_layout(as_marker_array(row)) - - -def test_pyccel_rejects_row_major_markers(): - """pyccel checks the memory order of the marker arrays at every call, so a missed conversion fails loudly.""" - from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments + markers = marker_array(7, 64) + assert markers.shape == (7, 64) and markers.flags.c_contiguous + assert xp.all(markers == 0.0) + # one quantity of all markers is contiguous + assert markers[2].flags.c_contiguous and markers[2:5].flags.c_contiguous + - args = (np.ones(4, dtype=bool), 4, 3, 6, 7, 8, 14, 17, 18, 4, np.zeros(3, dtype=np.int64), False) - MarkerArguments(np.zeros((4, 25), order="F"), *args) - with pytest.raises(TypeError, match="ordering"): - MarkerArguments(np.zeros((4, 25)), *args) +def _marker_arguments(markers, valid): + from struphy.kernel_arguments.pusher_args_cuda import CudaMarkerArguments + from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments - # a single row: as_marker_array gives it the strides pyccel accepts - one = (np.ones(1, dtype=bool), 1, *args[2:]) - with pytest.raises(TypeError, match="ordering"): - MarkerArguments(np.zeros((1, 25)), *one) - MarkerArguments(as_marker_array(np.zeros((1, 25))), *one) + args_class = CudaMarkerArguments if xp.get_backend() == "cupy" else MarkerArguments + n = markers.shape[1] + return args_class(markers, valid, n, *MARKER_INDICES, xp.zeros(3, dtype=np.int64), False) def _device_context(backend): @@ -78,9 +53,52 @@ def _device_context(backend): @pytest.mark.parametrize("backend", BACKENDS) -def test_markers_stay_column_major(backend): +def test_kernel_accepts_marker_copies(backend): + """A C-ordered copy of the valid markers (``xp.ascontiguousarray(markers[:, valid])``) goes to MarkerArguments and + a kernel like the buffer, and the kernel moves the copied markers like the markers of the buffer.""" + from struphy.geometry.domains import Cuboid + from struphy.ode.utils import ButcherTableau + from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage + + rng = np.random.default_rng(7) + host = rng.random((N_COLS, 64)) + host[3:6] = rng.uniform(-2, 2, (3, 64)) + host[8:14] = host[:6] + host[18:21] = 0.0 + host[:, 0] = -1.0 # a hole + host_valid = np.ones(64, dtype=bool) + host_valid[0] = False + + with xp.use_backend(backend), _device_context(backend): + markers, valid = xp.array(host), xp.asarray(host_valid) + # selecting markers gives a Fortran-ordered copy on NumPy; the kernels take C-ordered arrays + if backend == "numpy": + assert markers[:, valid].flags.f_contiguous and not markers[:, valid].flags.c_contiguous + copy = xp.ascontiguousarray(markers[:, valid]) + single = xp.ascontiguousarray(markers[:, 5:6]) + assert copy.flags.c_contiguous and single.flags.c_contiguous + + args_buffer = _marker_arguments(markers, valid) + args_copy = _marker_arguments(copy, xp.ones(copy.shape[1], dtype=bool)) + args_single = _marker_arguments(single, xp.ones(1, dtype=bool)) + assert args_buffer.n_markers == 64 and args_copy.n_markers == 63 and args_single.n_markers == 1 + + butcher = ButcherTableau("forward_euler") + butcher_args = (xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(butcher.c), butcher.n_stages) + args_domain = Cuboid(r1=2.0, r2=3.0, r3=4.0).args_domain + for args in (args_buffer, args_copy, args_single): + push_eta_stage(0.1, 0, args, args_domain, *butcher_args, n_threads=args.n_markers) + + np.testing.assert_array_equal(xp.to_numpy(copy), xp.to_numpy(markers[:, valid])) + np.testing.assert_array_equal(xp.to_numpy(single), xp.to_numpy(markers[:, 5:6])) + assert not np.array_equal(xp.to_numpy(copy[:3]), host[:3, 1:]) + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_marker_buffer_layout(backend): """Loading, pushing, boundary conditions, sorting between ranks, evaluation at markers and binning keep the - buffer column-major, and evaluation at row-major copies of the markers gives the same values.""" + buffer of shape (n_cols, n_rows) and C-ordered, and evaluation at copies of the valid markers gives the values + at the buffer.""" from maybempi import MPI from struphy import BoundaryParameters, LoadingParameters, WeightsParameters, domains, perturbations @@ -111,8 +129,8 @@ def test_markers_stay_column_major(backend): particles.initialize_weights() buffer = particles.markers - check_marker_layout(buffer) - assert buffer.shape[0] % ROW_ALIGNMENT == 0 + assert buffer.shape == (particles.n_cols, particles.n_rows) and buffer.flags.c_contiguous + assert particles.n_rows % ROW_ALIGNMENT == 0 butcher = ButcherTableau("rk4") push_eta = Pusher( @@ -140,29 +158,35 @@ def test_markers_stay_column_major(backend): # the buffer is never replaced, only written to assert particles.markers is buffer - check_marker_layout(particles.markers) - - # evaluation at row-major copies of the markers equals evaluation at the buffer (converted at the entry) - rows = particles.markers_wo_holes_and_ghost - assert rows.flags.c_contiguous - n_valid = rows.shape[0] - det_buffer = domain.jacobian_det(particles.markers, remove_outside=False)[particles.valid_mks] - np.testing.assert_array_equal(xp.to_numpy(domain.jacobian_det(rows)), xp.to_numpy(det_buffer)) + assert buffer.shape == (particles.n_cols, particles.n_rows) and buffer.flags.c_contiguous + + # evaluation at copies of the valid markers (Fortran-ordered on NumPy, converted by Domain and SplineFunction) + # equals evaluation at the buffer + valid = particles.valid_mks + columns = particles.markers_wo_holes_and_ghost + if backend == "numpy": + assert columns.flags.f_contiguous and not columns.flags.c_contiguous + n_valid = columns.shape[1] + assert n_valid == int(xp.count_nonzero(valid)) + det_buffer = domain.jacobian_det(buffer, remove_outside=False)[valid] + np.testing.assert_array_equal(xp.to_numpy(domain.jacobian_det(columns)), xp.to_numpy(det_buffer)) + pulled = domain.pull(lambda x, y, z: x * y + z, columns, kind="3") + assert pulled.shape == (n_valid,) np.testing.assert_array_equal( - xp.to_numpy(domain.pull(lambda x, y, z: x * y + z, rows, kind="3")), - xp.to_numpy(domain.pull(lambda x, y, z: x * y + z, as_marker_array(rows), kind="3")), + xp.to_numpy(pulled), + xp.to_numpy(domain.pull(lambda x, y, z: x * y + z, buffer, kind="3", remove_outside=False)[valid]), ) - # spline evaluation at row-major markers (local: every rank has its own markers) + # spline evaluation at the markers (local: every rank has its own markers) field = derham.create_spline_function("f", "H1") field.initialize_coeffs(perturbations=perturbations.ModesCos(ls=(1,), ms=(1,), ns=(1,), amps=(1.0,))) - values = field(rows, local=True) + values = field(columns, local=True) assert values.shape == (n_valid,) - np.testing.assert_array_equal(xp.to_numpy(values), xp.to_numpy(field(as_marker_array(rows), local=True))) + np.testing.assert_array_equal(xp.to_numpy(values), xp.to_numpy(field(buffer, local=True)[valid])) - # binning reads single columns of the valid markers + # binning reads single rows of the valid markers f_slice, _ = particles.binning( (True, False, False, True, False, False), (xp.linspace(0.0, 1.0, 9), xp.linspace(-4.0, 4.0, 9)) ) assert f_slice.shape == (8, 8) - check_marker_layout(particles.markers) + assert particles.markers is buffer and buffer.flags.c_contiguous diff --git a/src/struphy/pic/tests/test_pole.py b/src/struphy/pic/tests/test_pole.py index 31a2b1874..13d5be106 100644 --- a/src/struphy/pic/tests/test_pole.py +++ b/src/struphy/pic/tests/test_pole.py @@ -22,7 +22,7 @@ def _marker_args(markers, vdim, bc_type, pole): from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments - n_mks = markers.shape[0] + n_mks = markers.shape[1] return MarkerArguments( markers, np.ones(n_mks, dtype=bool), @@ -47,26 +47,26 @@ def test_bc_kernel_axis_map(): domain = domains.HollowCylinder(a1=0.0, a2=1.0, Lz=4.0) - markers = np.zeros((3, N_COLS), order="F") - 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) + markers = np.zeros((N_COLS, 3)) + markers[-1] = np.arange(3) + markers[:6, 0] = (-0.2, 0.3, 0.5, 0.1, 0.2, 0.3) + markers[:6, 1] = (-0.1, 0.7, 0.5, 0.1, 0.2, 0.3) + markers[:6, 2] = (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 + assert np.allclose(markers[:3, 0], (0.2, 0.8, 0.5)) + assert np.allclose(markers[:3, 1], (0.1, 0.2, 0.5)) + assert markers[0, 2] == 1e-10 # Cartesian velocities are unchanged, the marker is not removed - assert np.all(markers[:, 3:6] == (0.1, 0.2, 0.3)) + assert np.all(markers[3:6].T == (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) + markers[:6, 0] = (-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) + assert np.all(markers[:-1, 0] == -1.0) @pytest.mark.mpi_skip @@ -77,19 +77,19 @@ def test_bc_kernel_reflect_vdim(vdim): domain = domains.HollowCylinder(a1=0.2, a2=1.0, Lz=4.0) - markers = np.zeros((1, N_COLS), order="F") - markers[0, :6] = (1.1, 0.3, 0.5, 0.4, 0.2, 0.7) + markers = np.zeros((N_COLS, 1)) + markers[:6, 0] = (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 + assert markers[FIRST_PUSHER_IDX, 0] == -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))) + assert not np.allclose(markers[3:6, 0], (0.4, 0.2, 0.7)) + assert np.isclose(np.linalg.norm(markers[3:6, 0]), 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)) + assert np.all(markers[3:6, 0] == (0.4, 0.2, 0.7)) def _particles_on_pole_domain(comm, domain, ppc=10, num_elements=(8, 9, 4), bufsize=0.25): @@ -135,21 +135,21 @@ def test_python_bc_axis_map(): particles = _particles_on_pole_domain(comm, domain) assert particles.args_markers.pole - assert np.all(particles.positions[:, 0] >= 1e-10) + 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 + markers[0, valid[:3]] = -0.25 + markers[1, valid[:3]] = 0.6 # outside at eta1=1: removed (boundary condition of axis 0), not mapped through the axis - markers[valid[3], 0] = 1.05 + markers[0, valid[3]] = 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 np.allclose(markers[0, valid[:3]], 0.25) + assert np.allclose(markers[1, valid[:3]], 0.1) + assert np.all(markers[:-1, valid[3]] == -1.0) assert particles.n_lost_markers == 1 assert np.count_nonzero(particles.valid_mks) == n_valid - 1 @@ -174,14 +174,14 @@ def test_push_eta_through_pole(algo): 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 + x0, y0, z0 = domain(markers[:3, valid].copy()) + markers[3, valid] = (-1.8 * x0 + 0.01) / T + markers[4, valid] = (-1.8 * y0 - 0.02) / T + markers[5, valid] = 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) + expected = np.stack((x0 + markers[3, valid] * T, y0 + markers[4, valid] * T, z0 + markers[5, valid] * T), axis=1) + expected_by_id = _by_unique_id(comm, markers[-1, valid], expected) n_valid_0 = comm.allreduce(np.count_nonzero(valid)) butcher = ButcherTableau(algo=algo) @@ -202,10 +202,10 @@ def test_push_eta_through_pole(algo): assert comm.allreduce(np.count_nonzero(valid)) == n_valid_0 assert particles.n_lost_markers == 0 - ids = markers[valid, -1] + ids = markers[-1, valid] 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, y, z = domain(markers[:3, valid][:, 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) @@ -259,8 +259,8 @@ def _read_scalars(sim, 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] + # called with (eta1, eta2, eta3) or with one (3, n) array of positions + eta1 = etas[0] if len(etas) == 3 else etas[0][0] return np.where(eta1 < 0.6, 1.0, 0.0) @@ -304,9 +304,9 @@ def _gc_energies(sim, particles): 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 + eta = markers[:3, valid].T.copy() # one row per marker + energy = 0.5 * markers[3, valid] ** 2 + markers[particles.first_diagnostics_idx, valid] + return markers[-1, valid].copy(), eta, energy @pytest.mark.mpi(min_size=1) diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index 9f5efe75c..0e4b5024b 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -578,7 +578,7 @@ def test_push_bxu_Hdiv_pauli(num_elements, degree, bcs, mapping, show_plots=Fals flattened=True, ) - mu0 = xp.zeros(particles.markers.copy().T.shape[1], dtype=float) + mu0 = xp.zeros(particles.n_rows, dtype=float) pusher_psy = Pusher_psy( particles, @@ -708,7 +708,7 @@ def test_push_eta_rk4(num_elements, degree, bcs, mapping, show_plots=False): n_mks_load = xp.zeros(size, dtype=int) - comm.Allgather(xp.array(xp.shape(particles.markers)[0]), n_mks_load) + comm.Allgather(xp.array(particles.n_rows), n_mks_load) sendcounts = xp.zeros(size, dtype=int) displacements = xp.zeros(size, dtype=int) @@ -722,7 +722,7 @@ def test_push_eta_rk4(num_elements, degree, bcs, mapping, show_plots=False): all_particles_psy = xp.zeros((int(accum_sendcounts) * 3,), dtype=float) comm.Barrier() - comm.Allgatherv(xp.array(particles.markers[:, :3]), [all_particles_psy, sendcounts, displacements, MPI.DOUBLE]) + comm.Allgatherv(xp.array(particles.markers[:3]), [all_particles_psy, sendcounts, displacements, MPI.DOUBLE]) comm.Barrier() @@ -758,8 +758,8 @@ def test_kinetic_bc_in_kernel(bc, mapping): bc_type_kernel = bc_type.copy() # large time step such that many markers leave the unit cube - markers[:, first_pusher_idx:first_shift_idx] = markers[:, :6] - markers[:, first_shift_idx:-2] = 0.0 + markers[first_pusher_idx:first_shift_idx] = markers[:6] + markers[first_shift_idx:-2] = 0.0 markers_init = markers.copy() n_holes_init = xp.count_nonzero(particles.holes) @@ -769,7 +769,7 @@ def test_kinetic_bc_in_kernel(bc, mapping): 0.2, 0, particles.args_markers, domain.args_domain, butcher.a_stage, butcher.b, butcher.c, butcher.n_stages ) n_outside = xp.count_nonzero( - xp.logical_or(markers[~particles.holes, :3] > 1.0, markers[~particles.holes, :3] < 0.0) + xp.logical_or(markers[:3, ~particles.holes] > 1.0, markers[:3, ~particles.holes] < 0.0) ) assert n_outside > 0 particles.apply_kinetic_bc() @@ -785,12 +785,12 @@ def test_kinetic_bc_in_kernel(bc, mapping): ) particles.finish_kernel_bc() - assert xp.array_equal(markers[:, :first_shift_idx], markers_ref[:, :first_shift_idx]) - assert xp.all(markers[~particles.holes, :3] >= 0.0) - assert xp.all(markers[~particles.holes, :3] <= 1.0) + assert xp.array_equal(markers[:first_shift_idx], markers_ref[:first_shift_idx]) + assert xp.all(markers[:3, ~particles.holes] >= 0.0) + assert xp.all(markers[:3, ~particles.holes] <= 1.0) if bc == "periodic": shift_slice = slice(first_shift_idx, first_shift_idx + 3) - assert xp.array_equal(markers[:, shift_slice], markers_ref[:, shift_slice]) + assert xp.array_equal(markers[shift_slice], markers_ref[shift_slice]) if bc == "remove": n_new_holes = xp.count_nonzero(particles.holes) - n_holes_init assert n_new_holes > 0 @@ -816,9 +816,9 @@ def test_kinetic_bc_remove_counts_each_marker_once(): valid = xp.nonzero(particles.valid_mks)[0] n_valid = valid.size - particles.markers[valid[0], :3] = [1.5, -0.5, 0.5] # outside on two axes - particles.markers[valid[1], :3] = [1.5, 1.5, 1.5] # outside on all three axes - particles.markers[valid[2], :3] = [0.5, 0.5, -0.1] # outside on one axis + particles.markers[:3, valid[0]] = [1.5, -0.5, 0.5] # outside on two axes + particles.markers[:3, valid[1]] = [1.5, 1.5, 1.5] # outside on all three axes + particles.markers[:3, valid[2]] = [0.5, 0.5, -0.1] # outside on one axis n_lost_before = particles.n_lost_markers particles.apply_kinetic_bc() diff --git a/src/struphy/pic/tests/test_reflect.py b/src/struphy/pic/tests/test_reflect.py index 5830346e6..742311e08 100644 --- a/src/struphy/pic/tests/test_reflect.py +++ b/src/struphy/pic/tests/test_reflect.py @@ -22,14 +22,14 @@ def test_reflect_in_place(backend, axis, count): # SIGFPE on x86-64). Particles.apply_kinetic_bc never passes an empty index array; the CUDA launch of # zero threads is a no-op and stays tested. pytest.skip("pyccel kernels must not be called with zero-size arrays") - original = np.random.default_rng(32).random((140, 16)) + original = np.random.default_rng(32).random((16, 140)) indices = np.arange(139, 139 - count, -1, dtype=np.int64) expected = original.copy() - expected[indices, 2 * (3 + axis)] *= -1 + expected[2 * (3 + axis), indices] *= -1 with xp.use_backend(backend): - # every other column of a column-major buffer: strided, in the marker layout - storage = xp.asfortranarray(xp.asarray(original)) - markers = storage[:, ::2] + # every other row of a marker buffer (one column per marker): strided, in the marker layout + storage = xp.asarray(original) + markers = storage[::2] outside_inds = xp.asarray(np.repeat(indices, 2))[::2] domain = Cuboid(r1=2.0, r2=3.0, r3=4.0) reflect(markers, domain.args_domain, outside_inds, axis, n_threads=outside_inds.size) @@ -69,7 +69,7 @@ def test_reflect_rejects_unsupported_mapping(): def test_reflect_rejects_wrong_index_dtype(dtype): with xp.use_backend("cupy"): with pytest.raises(TypeError, match="Array1D"): - reflect(xp.zeros((2, 6)), Cuboid().args_domain, xp.asarray([0], dtype=dtype), 0, n_threads=1) + reflect(xp.zeros((6, 2)), Cuboid().args_domain, xp.asarray([0], dtype=dtype), 0, n_threads=1) @pytest.mark.parametrize("backend", BACKENDS) @@ -77,13 +77,13 @@ def test_particle_boundary_reflection(backend): with xp.use_backend(backend): particles = reflecting_particles(Cuboid()) selected = xp.nonzero(particles.valid_mks)[0][:3] - particles.markers[selected, :3] = xp.asarray([[-0.1, 0.5, 0.5], [0.5, 1.1, 0.5], [-0.2, 1.2, -0.3]]) - particles.markers[selected, 3:6] = xp.asarray([1.0, 2.0, 3.0]) + particles.markers[:3, selected] = xp.asarray([[-0.1, 0.5, 0.5], [0.5, 1.1, 0.5], [-0.2, 1.2, -0.3]]).T + particles.markers[3:6, selected] = xp.asarray([1.0, 2.0, 3.0])[:, None] particles.apply_kinetic_bc() np.testing.assert_allclose( - xp.to_numpy(particles.markers[selected, :3]), [[0.1, 0.5, 0.5], [0.5, 0.9, 0.5], [0.2, 0.8, 0.3]] + xp.to_numpy(particles.markers[:3, selected]).T, [[0.1, 0.5, 0.5], [0.5, 0.9, 0.5], [0.2, 0.8, 0.3]] ) np.testing.assert_allclose( - xp.to_numpy(particles.markers[selected, 3:6]), [[-1.0, 2.0, 3.0], [1.0, -2.0, 3.0], [-1.0, -2.0, -3.0]] + xp.to_numpy(particles.markers[3:6, selected]).T, [[-1.0, 2.0, 3.0], [1.0, -2.0, 3.0], [-1.0, -2.0, -3.0]] ) - assert bool(xp.all(particles.markers[selected, particles.first_pusher_idx] == -1.0)) + assert bool(xp.all(particles.markers[particles.first_pusher_idx, selected] == -1.0)) diff --git a/src/struphy/pic/tests/test_sorting.py b/src/struphy/pic/tests/test_sorting.py index c3748a722..3f40a73c1 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -149,7 +149,7 @@ def test_sorting(num_elements, degree, bcs, mapping, Np): logger.info("Rank : {0} | Sorting time : {1:8.6f}".format(rank, time_sorting)) - box_markers = particles.markers[:, -2] + box_markers = particles.markers[-2] assert all(box_markers[i] <= box_markers[i + 1] for i in range(len(box_markers) - 1)) @@ -177,10 +177,10 @@ def test_mpi_sort_markers_on_rank_boundary(bc): tags = 1000.0 + xp.arange(len(special)) if mpi_comm.Get_rank() == 0: - rows = xp.nonzero(particles.holes)[0][: len(special)] - particles.markers[rows] = 0.0 - particles.markers[rows, :3] = xp.array(special) - particles.markers[rows, 3] = tags + holes = xp.nonzero(particles.holes)[0][: len(special)] + particles.markers[:, holes] = 0.0 + particles.markers[:3, holes] = xp.array(special).T + particles.markers[3, holes] = tags particles.update_holes() n_before = mpi_comm.allreduce(particles.n_mks_loc) @@ -188,7 +188,7 @@ def test_mpi_sort_markers_on_rank_boundary(bc): n_after = mpi_comm.allreduce(particles.n_mks_loc) assert n_after == n_before - v1 = particles.markers[particles.valid_mks, 3] + v1 = particles.markers[3, particles.valid_mks] for tag, eta in zip(tags, special): n_found = mpi_comm.allreduce(int(xp.count_nonzero(v1 == tag))) assert n_found == 1, f"marker at {eta} found on {n_found} processes" diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index dc8fa9462..13f354ff8 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -81,10 +81,10 @@ def _make_particles(comm, bc): def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.ndarray: - """New positions in [0, 1) for the valid rows, far enough away that most markers change rank.""" + """New positions in [0, 1) for the valid markers, far enough away that most markers change rank.""" new = markers.copy() - shift = np.array([0.37, 0.21, 0.53]) * (step + 1) - new[valid, :3] = (markers[valid, :3] + shift) % 1.0 + shift = np.array([0.37, 0.21, 0.53])[:, None] * (step + 1) + new[:3, valid] = (markers[:3, valid] + shift) % 1.0 return new @@ -120,11 +120,11 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", for step in range(2): new = _shift_positions(ref.markers, ref.valid_mks, step) - # number of rows this rank has to send, from the host reference - etas = new[ref.valid_mks, :3] - left, right = ref.domain_array[rank, 0::3], ref.domain_array[rank, 1::3] + # number of markers this rank has to send, from the host reference + etas = new[:3, ref.valid_mks] + left, right = ref.domain_array[rank, 0::3, None], ref.domain_array[rank, 1::3, None] right = np.where(right == 1.0, np.nextafter(right, 2.0), right) - n_send = int(np.count_nonzero(~np.all((etas >= left) & (etas < right), axis=1))) + n_send = int(np.count_nonzero(~np.all((etas >= left) & (etas < right), axis=0))) n_send_total = comm.allreduce(n_send) assert size == 1 or n_send_total > 0 @@ -144,14 +144,14 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", assert transfers.total == 0, transfers.report() else: # exchange stages each non-empty send and receive buffer separately - n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) - n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[0] > 0) + n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[1] > 0) + n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[1] > 0) assert transfers.to_host == n_dest, transfers.report() assert transfers.to_device == n_sources, transfers.report() sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") - assert sent_bytes == n_send * ref.markers.shape[1] * 8 + assert sent_bytes == n_send * ref.n_cols * 8 received_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_device") - assert received_bytes == n_recv * ref.markers.shape[1] * 8 + assert received_bytes == n_recv * ref.n_cols * 8 if count_reads and cuda_aware: # (the staged exchange reads the send buffers with .get(), counted above) diff --git a/src/struphy/pic/tests/test_sph.py b/src/struphy/pic/tests/test_sph.py index 15d77e99a..611ebc0be 100644 --- a/src/struphy/pic/tests/test_sph.py +++ b/src/struphy/pic/tests/test_sph.py @@ -1854,7 +1854,7 @@ def u_xyz(x, y, z): ghost_inds = xp.where(particles.ghost_particles)[0] logger.info(f"After do_sort: {len(ghost_inds)} ghosts") if len(ghost_inds) > 0: - logger.info(f"First 10 ghost eta1: {particles.markers[ghost_inds[:10], 0]}") + logger.info(f"First 10 ghost eta1: {particles.markers[0, ghost_inds[:10]]}") particles.initialize_weights() if direction == "x": @@ -1888,8 +1888,8 @@ def u_xyz(x, y, z): ) # if rank == 0 and len(ghost_inds) > 0: - # logger.info("Ghost coefficients after eval:", particles.markers[ghost_inds[:10], particles.first_free_idx]) - # logger.info("Ghost positions after eval:", particles.markers[ghost_inds[:10], 0]) + # logger.info("Ghost coefficients after eval:", particles.markers[particles.first_free_idx, ghost_inds[:10]]) + # logger.info("Ghost positions after eval:", particles.markers[0, ghost_inds[:10]]) if comm is not None: all_v1 = xp.zeros_like(v1) @@ -2158,7 +2158,7 @@ def test_sph_full_box_row_detected(): particles = _make_sph_particles_1d("periodic", None) particles.put_particles_in_boxes() - max_in_box = int(xp.max(xp.bincount(xp.int64(particles.markers_wo_holes[:, -2])))) + max_in_box = int(xp.max(xp.bincount(xp.int64(particles.markers_wo_holes[-2])))) # one spare column: fits, the -1 terminator is kept n_boxes = particles.sorting_boxes.boxes.shape[0] diff --git a/src/struphy/pic/tests/test_tesselation.py b/src/struphy/pic/tests/test_tesselation.py index c4ab72f2e..55a1dd5a4 100644 --- a/src/struphy/pic/tests/test_tesselation.py +++ b/src/struphy/pic/tests/test_tesselation.py @@ -42,7 +42,7 @@ def test_draw(ppb, nx, ny, nz): ) particles.draw_markers(sort=False) - # logger.info(f'{particles.markers[:, :3] = }') + # logger.info(f'{particles.markers[:3] = }') # logger.info(f'{rank = }, {particles.positions = }') # test @@ -68,9 +68,9 @@ def test_draw(ppb, nx, ny, nz): # logger.info(f'\n{rank = }, {e1 = }') - assert xp.allclose(particles.positions[:, 0], e1) - assert xp.allclose(particles.positions[:, 1], e2) - assert xp.allclose(particles.positions[:, 2], e3) + assert xp.allclose(particles.positions[0], e1) + assert xp.allclose(particles.positions[1], e2) + assert xp.allclose(particles.positions[2], e3) @pytest.mark.parametrize("ppb", [8, 12]) @@ -126,8 +126,8 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): if ny == nz == 1: plt.figure(figsize=(15, 10)) - plt.plot(particles.positions[:, 0], xp.zeros_like(particles.weights), "o", label="markers") - plt.plot(particles.positions[:, 0], particles.weights, "-o", label="weights") + plt.plot(particles.positions[0], xp.zeros_like(particles.weights), "o", label="markers") + plt.plot(particles.positions[0], particles.weights, "-o", label="weights") plt.plot( xp.linspace(xl, xr, 100), particles.f_init(xp.linspace(xl, xr, 100), 0.5, 0.5).squeeze(), @@ -151,7 +151,7 @@ def test_cell_average(ppb, nx, ny, nz, n_quad, show_plot=False): ax.set_xticks(xp.linspace(0, 1, nx + 1)) ax.set_yticks(xp.linspace(0, 1, ny + 1)) coloring = particles.weights - plt.scatter(particles.positions[:, 0], particles.positions[:, 1], c=coloring, s=40) + plt.scatter(particles.positions[0], particles.positions[1], c=coloring, s=40) plt.grid(c="k") plt.axis("square") plt.title("initial markers") diff --git a/src/struphy/pic/tests/test_utilities.py b/src/struphy/pic/tests/test_utilities.py index d532eabb5..5343b0728 100644 --- a/src/struphy/pic/tests/test_utilities.py +++ b/src/struphy/pic/tests/test_utilities.py @@ -22,14 +22,14 @@ def test_energy_utilities(mapping): # kinetic energy with zero vector potential: 0.5 * sum_p w_p |v_p|^2 rng = xp.random.default_rng(0) - markers = xp.zeros((5, 7)) - markers[:, :3] = rng.uniform(0.0, 1.0, (5, 3)) - markers[:, 3:6] = rng.normal(size=(5, 3)) - markers[:, 6] = rng.uniform(0.5, 1.5, 5) - markers[-1, 0] = -1.0 # hole + markers = xp.zeros((7, 5)) + markers[:3] = rng.uniform(0.0, 1.0, (3, 5)) + markers[3:6] = rng.normal(size=(3, 5)) + markers[6] = rng.uniform(0.5, 1.5, 5) + markers[0, -1] = -1.0 # hole a = derham.V1.zeros() res = get_kinetic_energy_particles(a, derham, domain, SimpleNamespace(markers=markers)) - ref = 0.5 * xp.sum(markers[:-1, 6] * xp.sum(markers[:-1, 3:6] ** 2, axis=1)) + ref = 0.5 * xp.sum(markers[6, :-1] * xp.sum(markers[3:6, :-1] ** 2, axis=0)) assert xp.isclose(res[0], ref) # thermal energy of a constant density n = e: int n ln(n) sqrt(g) = e * volume diff --git a/src/struphy/post_processing/tests/test_orbits_tools.py b/src/struphy/post_processing/tests/test_orbits_tools.py index ec4456526..fc38b13c7 100644 --- a/src/struphy/post_processing/tests/test_orbits_tools.py +++ b/src/struphy/post_processing/tests/test_orbits_tools.py @@ -17,7 +17,7 @@ def test_guiding_center_and_classification(tmp_path): equil.domain = domain rng = np.random.default_rng(0) - etas = np.column_stack((rng.uniform(0.3, 0.9, N_MARKERS), rng.random(N_MARKERS), rng.random(N_MARKERS))) + etas = np.stack((rng.uniform(0.3, 0.9, N_MARKERS), rng.random(N_MARKERS), rng.random(N_MARKERS))) x = domain(etas, change_out_order=True) b = equil.b_cart(etas)[0].T absB = np.linalg.norm(b, axis=1) diff --git a/src/struphy/propagators/tests/test_push_random_diffusion.py b/src/struphy/propagators/tests/test_push_random_diffusion.py index d75bf0d95..0b69dd567 100644 --- a/src/struphy/propagators/tests/test_push_random_diffusion.py +++ b/src/struphy/propagators/tests/test_push_random_diffusion.py @@ -34,13 +34,13 @@ def test_random_diffusion_msd(algo): prop.allocate() np.random.seed(0) - valid = particles.markers[:, 0] != -1.0 - x0 = particles.markers[valid, :3].copy() + valid = particles.markers[0] != -1.0 + x0 = particles.markers[:3, valid].copy() prop(dt) - dx = particles.markers[valid, :3] - x0 + dx = particles.markers[:3, valid] - x0 dx -= np.round(dx) # minimal image in the periodic unit cube - msd = np.mean(dx**2, axis=0) / (2 * diff_coeff * dt) + msd = np.mean(dx**2, axis=1) / (2 * diff_coeff * dt) # statistical error of the MSD estimate is sqrt(2 / 4000) ~ 0.022 assert np.all(np.abs(msd - 1.0) < 0.1), msd diff --git a/src/struphy/simulation/tests/test_output.py b/src/struphy/simulation/tests/test_output.py index f17702e4a..e33bafbf5 100644 --- a/src/struphy/simulation/tests/test_output.py +++ b/src/struphy/simulation/tests/test_output.py @@ -360,8 +360,8 @@ def spy(self, data): sim.time_opts.Tend += sim.time_opts.dt sim.run() - valid = saved[:, 0] != -1.0 - assert np.array_equal(restored["markers"][valid, :6], saved[valid, :6]) + valid = saved[0] != -1.0 + assert np.array_equal(restored["markers"][:6, valid], saved[:6, valid]) def test_run_metadata_names_variable_keys_in_propagator_options(tmp_path): From ff71d9999ddf0e8645843404898e845304504760 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 15 Oct 2026 17:18:54 +0200 Subject: [PATCH 09/12] CUDA strategy: marker layout (n_cols, n_markers) in C order Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 6af88b79e..f3b1d89a1 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -108,7 +108,7 @@ Conventions: - The device functions of a header live in the namespace `struphy_cuda::`, where `` is the name of the pyccel module they port: `struphy_cuda::bsplines_kernels` (`bsplines/bsplines_kernels.cuh`), `struphy_cuda::evaluation_kernels` (`geometry/evaluation_kernels.cuh`), `struphy_cuda::cuboid_kernels` (`geometry/domains/cuboid/cuboid_cuda.cuh`, which ports `cuboid_kernels.py`), and so on. CUDA-only helpers and constants go into the namespace of the module whose calculation they extract (`bsplines_kernels::MAX_SPLINE_DEGREE`, `evaluation_kernels_3d::SplineScratch`). Calls into another module are qualified with the module name, as the pyccel code calls `linalg_kernels.matrix_inv`: `linalg_kernels::matrix_inv(...)`; calls within the same module are not. A `.cu` file writes `using namespace struphy_cuda;` after its includes and qualifies every helper call (`evaluation_kernels_3d::get_spans(...)`). Only `struphy_cuda::pi` (`geometry/domains/constants_cuda.cuh`) and the generated argument structs (`MarkerArgs`, `DomainArgs`, `DerhamArgs`, `LocalProjectorsArgs`, which are global because cunumpy matches the kernel signature against their names) are outside a module namespace. - Every function in a `.cu` or `.cuh` file has a documentation comment (`/** ... */`, the C++ equivalent of a docstring) immediately above its definition. Identify the pyccel counterpart and explain the operation, parameters, outputs or return value, array shapes/layouts, and any CUDA-specific limits or storage requirements. Document CUDA-only helpers by identifying the pyccel calculation they extract. - CUDA parameter, local variable and scratch-field names exactly match the corresponding pyccel names, including spelling and capitalization (e.g. `eta1`, `args_markers`, `args_domain`, `args_derham`, `df_out`, `dfm`, `dfinv`, `v_logical`, `tn`, `pn`, `pd`, `il`, `det_a`, `span1`, `bn1`, `bd1`). Do not abbreviate them to `x`, `m`, `d`, `a`, or `out` when the pyccel counterpart uses a different name. Additional CUDA-only variables or arguments, such as pointer lengths, thread indices or per-thread scratch storage, must have descriptive names and documented purposes. -- Kernels read and write markers through the `Array2D` view: `args.markers(ip, j)`. The view carries the existing device pointer, shape and element strides; no array data is copied. The column count is `args.markers.shape[1]`. +- Kernels read and write markers through the `Array2D` view: `args.markers(j, ip)` (one column per marker, see the marker layout below). The view carries the existing device pointer, shape and element strides; no array data is copied. The number of marker quantities (`n_cols`) is `args.markers.shape[0]`. The geometry domains already follow a similar layout (`geometry/domains/cuboid/cuboid_kernels.py`), which can be extended with `cuboid_cuda.cuh` for the device version of the mapping. @@ -181,7 +181,7 @@ The pusher kernels call helpers from other pyccel modules: B-spline evaluation ( - First: run the PR 8 GPU tests (`test_cuda_struct_layout`, `test_cuda_struct_scalars_are_checked`, the demo kernels) on an H100. Everything below assumes structs arrive by value correctly. - Port only what `push_eta_stage` needs: `get_spans` and the basis functions for the spline degrees in use, the analytic mapping evaluation (`kind_map` switch, starting with `Cuboid`), the 3×3 linear algebra, and the kinetic boundary conditions. Everything else comes with the kernel that needs it. -- Device helpers take the argument structs by reference (`const DomainArgs&`), exactly like the pyccel helpers take `args_domain`, and read markers through `args.markers(ip, j)`. +- Device helpers take the argument structs by reference (`const DomainArgs&`), exactly like the pyccel helpers take `args_domain`, and read markers through `args.markers(j, ip)`. - Scratch arrays that the pyccel classes allocate once (e.g. `DerhamArguments.bn1`, ..., `bd3`) become per-thread local arrays in CUDA with a fixed maximum spline degree (compile-time constant, checked against `pn` when the arguments are built) or a template parameter. - Tests: each device helper is called from a small elementwise test kernel (one thread per evaluation point, inputs and outputs as arrays) and compared with its pyccel version from Python. One generic test-kernel template in `test_kernel_backends.py` (or a new `test_device_helpers.py`) serves all scalar-returning helpers. Tests are skipped without a GPU. - Headers are included through the source root (`#include "struphy/bsplines/bsplines_kernels.cuh"`). Remember the CuPy cache caveat from PR 2 when editing headers. @@ -190,7 +190,7 @@ The pusher kernels call helpers from other pyccel modules: B-spline evaluation ( - Write `push_eta_stage_cuda.cu`, using the device helpers from PR 10; one thread per marker, `if (ip >= args_markers.n_markers) return;`. - Parity: same markers on both backends, results agree to round-off (`rtol ~ 1e-13`). The test is written once against the catalog: it parametrizes over every kernel that has a CUDA version (`[n for n in catalog.names if n not in catalog.missing_cuda]`), builds the arguments with a per-kernel factory, runs the pyccel and the CUDA kernel, and compares every array argument. Later kernels then only add their factory. -- `Pusher` on the CuPy backend passes `n_threads=particles.markers.shape[0]` and device arrays in `args_kernel` (the Butcher tableau arrays `a_stage`, `b`, `c` are tiny NumPy arrays today; `PushEta` makes device copies once at setup). +- `Pusher` on the CuPy backend passes `n_threads=args_markers.n_markers` (`particles.markers.shape[1]`) and device arrays in `args_kernel` (the Butcher tableau arrays `a_stage`, `b`, `c` are tiny NumPy arrays today; `PushEta` makes device copies once at setup). - End-to-end: run `PushEta` (a propagator that only needs this kernel) with `CUNUMPY_BACKEND=cupy` on 1 and 2 ranks, and check that no host/device transfers happen inside the time loop: `nsys profile` (look for `cudaMemcpy` between the kernel launches) or monkeypatch `cupy.ndarray.get`/`cupy.asarray` to count calls during the loop. ### PR 12+: Port kernels and particle boundary handling @@ -224,7 +224,7 @@ cunumpy 0.5.0 contains equivalents of everything struphy built in PRs 1–8, wit | `utils.kernel_backends.Kernel`, `KernelCatalog`, `CudaKernel` | `xp.Kernel`, `xp.KernelCatalog.from_package` (same folder convention, source root on the include path by default), `xp.CudaKernel` (parses the `__global__` signature once and checks every call: argument count, array dtypes and contiguity, scalars cast and range-checked; 1D–3D launches via `n_threads` or `grid`; header-aware compile cache; `compile()` at setup; debug mode with `-lineinfo`, bounds checks and synchronized launches) | | `utils.cuda_arguments.Argument` + `_pack()` | `xp.CudaStructArguments` (`struct_name`, `fields`, `pack()`, repack on copy/unpickle, `cls.struct.declaration`/`to_header()`), `xp.CudaStruct.from_signature` to derive a struct from pyccel-style annotations | | the two bundles per owner (`args_*` and `_pyccel_args_*`) | `xp.KernelArguments`: one object with `__host_args__()` (the pyccel instance) and `__cuda_args__()` (the packed struct), resolved by `Kernel` and `PyccelKernel`, so call sites never branch on the backend | -| Local `Array2D` marker view | `Array2D` view (pointer, shape, strides) from `cunumpy/array_view.cuh`; kernels write `markers(ip, j)` | +| Local `Array2D` marker view | `Array2D` view (pointer, shape, strides) from `cunumpy/array_view.cuh`; kernels write `markers(j, ip)` | | the "reference or copy once" rule, written by hand | `xp.as_device_array(value, dtype, ndim)` | | the parity test of PR 11 | `cunumpy.testing.assert_kernels_agree`, `KernelCatalog.parity_cases()`, `requires_cupy`, `backend` fixture | | the device-helper test kernels of PR 10 | `cunumpy.testing.device_function_kernel(header, prototype)` | @@ -378,7 +378,7 @@ stencil matrix data) are in cunumpy since 0.6.1 (`Array6D`); `linear_vlasov_ampe - **Scalar types.** Members of the argument structs are checked when they are packed (PR 8). Scalars passed directly to a kernel (`dt`, `stage`) are not checked against the kernel signature: a Python `float` for an `int` parameter arrives wrong without an error. Either `CudaKernel` parses the `extern "C"` signature once and casts/checks (what cunumpy's `CudaKernel` does), or we accept it until the move to cunumpy. Until then: kernels take `double dt, int stage` in exactly this order, and the pyccel signature is the reference. - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. -- **Marker layout.** Decided (#707): the markers array (`n_markers × n_cols`) is column-major (struct-of-arrays, `order="F"`) on both backends, see `struphy.pic.marker_layout`. Neighbouring threads (markers) read neighbouring addresses, and column operations (the pusher preamble, hole detection, masks in the diagnostics) touch only the columns involved. Indexing is unchanged (`markers[ip, j]`); the pyccel kernels declare `float[:, :](order=F)`, and the CUDA kernels read the buffer through `Array2D`, which carries the strides. `n_rows` is a multiple of 32, so every column starts on a 256-byte boundary. +- **Marker layout.** Decided (#707): the markers array is a C-ordered `n_cols × n_markers` array, one column per marker (struct-of-arrays), on both backends, see `struphy.pic.marker_layout`; it is indexed `markers[j, ip]` (pyccel) and `markers(j, ip)` (CUDA). Neighbouring threads (markers) read neighbouring addresses, and operations on one quantity (the pusher preamble, hole detection, masks in the diagnostics) touch only the rows involved. The pyccel kernels declare a plain `float[:, :]`; the marker count is `markers.shape[1]`, so every launch over markers passes `n_threads=args_markers.n_markers` (cunumpy's default inference takes the first axis of the first array). Arrays of evaluation points (`Domain`, `SplineFunction`, the equilibria) follow the same convention, `(k, n)` with the logical coordinates in rows 0 to 2. `n_rows` (the number of markers including holes) is a multiple of 32, so every row starts on a 256-byte boundary. - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet. - **Single-source alternatives.** Decided (#690): no code generation; every CUDA kernel is written by hand, see [Hand-written CUDA](#hand-written-cuda-decision-for-690). - **Polar splines and MHD equilibria on the GPU** (left after PR 19): both run on CuPy now, polar splines since #695 @@ -547,7 +547,8 @@ at setup instead of `get_kernel()`: on CuPy a kernel without CUDA version raises is compiled. Launch sizes are not passed any more where cunumpy infers them: the first array of a pusher or accumulation call is `args_markers.markers`, so one thread per marker row. `KernelSetup` passes `n_threads=args_markers.n_markers` because its first array is `alpha`; `reflect` and the spline evaluation -pass theirs explicitly. +pass theirs explicitly. (Since #707 the marker array has one column per marker, and pushers and accumulators +pass `n_threads=args_markers.n_markers` too.) Parity tests follow cunumpy's convention (replaced in PR 16 by `pic/tests/cuda_parity_cases.py`): `_test_args.py` in the kernel folder defines `make_args(backend, seed)`; struphy adds `CASES` (the test cases, selected by `seed`) and the tolerances. @@ -666,9 +667,8 @@ Geometry evaluation runs on the GPU for every analytic mapping (`kind_map` 10– the determinant with `det()` like pyccel (it used a different expansion, which differed in round-off). Spline mappings (`kind_map` 0–2) trap on the device. - **Entry kernels.** `_cuda.cu` in the four `geometry/kernels//` folders, same arguments in the same order. - `kernel_evaluate_pic` launches one thread per marker row (the default: its first array is `markers`). - `kernel_pullpush_pic` launches one thread per marker row, passed as `n_threads=markers.shape[0]` because its first - array is `a`. `kernel_evaluate` and `kernel_pullpush` launch one thread per grid point (`n_threads=n1 * n2 * n3`). + `kernel_evaluate_pic` and `kernel_pullpush_pic` launch one thread per marker, passed as + `n_threads=markers.shape[1]` (the markers have one column per marker since #707). `kernel_evaluate` and `kernel_pullpush` launch one thread per grid point (`n_threads=n1 * n2 * n3`). `kernel_evaluate`'s `mat_f` has five dimensions, more than cunumpy's views, so it is a C-contiguous `double*`. - **`n_inside`.** The pyccel `_pic` kernels compact the rows of inside markers (`remove_outside=True`) and return their number; a CUDA kernel returns nothing, and compaction needs a prefix count over all markers. `Domain` therefore @@ -777,7 +777,7 @@ The first matrix accumulation on the GPU (part of #688, tracked in #650): `linea fillers, each one is ported with the first kernel that calls it. - **Kernel.** `linear_vlasov_ampere_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row, holes and boundary particles skipped; `df`, `matrix_inv`, `matrix_vector`, `outer`, then `m_v_fill_b_v1_symm`. The - 3 x 3 `filling_m` is nine row-major entries; the CUDA-only `factor` is pyccel's `f0_values[ip] / (markers[ip, 7] * + 3 x 3 `filling_m` is nine row-major entries; the CUDA-only `factor` is pyccel's `f0_values[ip] / (markers[7, ip] * n_markers_tot)`, computed once as pyccel does. - **`Accumulator` on CuPy.** Nothing in `Accumulator` had to change: on the CuPy backend `StencilMatrix._data` and `StencilVector._data` are device arrays (feectools allocates them with `xp`), the accumulator zeroes them in place and @@ -812,7 +812,7 @@ the [porting order](#porting-order); part of #688). - **Kernel.** `vlasov_maxwell_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row; `df`, `matrix_inv`, `transpose`, `matrix_matrix`, `matrix_vector`, then `m_v_fill_b_v1_symm` with `A_p = w_p DF^{-1} DF^{-T}` and `B_p = w_p DF^{-1} v_p`. As in pyccel, only holes are skipped: boundary particles - (`markers[ip, -1] == -2`) are accumulated (unlike `linear_vlasov_ampere`). + (`markers[-1, ip] == -2`) are accumulated (unlike `linear_vlasov_ampere`). - **Device helpers.** `fill_mat`, `fill_mat_vec` (`filler_kernels.cuh`) and `m_v_fill_b_v1_symm` (`particle_to_mat_kernels.cuh`) with atomic adds into `Array6D`/`Array3D` data, byte-identical to the ones of the `linear_vlasov_ampere` PR (#705), as are `v1_symm_accumulation_data()` and the `m_v_fill_b_v1_symm` From 8386ecb9a4fad97a9adb0f4cf720ddb5a7b5aadd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 08:37:25 +0200 Subject: [PATCH 10/12] Screw pinch example: evaluate |B0| at (3, n) points The exact radial profile for the binning plots evaluated equil.absB0 at an (n, 3) array of points; marker evaluation takes one column per point. Co-Authored-By: Claude Opus 5.5 --- .../screw_pinch_equilibrium/params_screw_pinch_equilibrium.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py b/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py index 5b7f79fb4..ee397af37 100644 --- a/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py +++ b/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py @@ -255,7 +255,7 @@ def vth_profile(*etas): # exact radial profile of the binned distribution (normalized to the binned data): the binning divides by the # spatial Jacobian only, the velocity-space Jacobian |B0| of (v_parallel, mu) remains -exact = {"e1_density": lambda e1: n_profile(e1, 0 * e1, 0 * e1) / equil.absB0(np.stack((e1, 0 * e1, 0 * e1), axis=1))} +exact = {"e1_density": lambda e1: n_profile(e1, 0 * e1, 0 * e1) / equil.absB0(np.stack((e1, 0 * e1, 0 * e1), axis=0))} if __name__ == "__main__": # post-processing and plots: pproc_screw_pinch_equilibrium.py From 67eb34e74cd545042dd9a10554526b07eb897c69 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:26 +0200 Subject: [PATCH 11/12] Name the marker slots after what they count: n_markers_max With one column per marker, n_rows (the number of marker slots, the second axis of the buffer) no longer counts rows. Renamed after its meaning: Particles.n_rows -> n_markers_max, padded_n_rows -> padded_n_markers and ROW_ALIGNMENT -> MARKER_ALIGNMENT; local marker counts are n_markers. n_cols stays the number of quantities per marker. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 2 +- src/struphy/examples/_draw_parallel.py | 4 +- .../kernel_arguments/pusher_args_cuda.py | 2 +- .../kernel_arguments/pusher_args_kernels.py | 2 +- .../linear_vlasov_ampere_one_species.py | 2 +- src/struphy/models/tests/test_scalars.py | 4 +- src/struphy/models/variables.py | 4 +- src/struphy/pic/base.py | 83 ++++++++++--------- src/struphy/pic/marker_layout.py | 20 ++--- src/struphy/pic/pushing/prepare_push_cuda.cu | 12 +-- src/struphy/pic/pushing/pusher.py | 6 +- src/struphy/pic/tests/bench_marker_layout.py | 8 +- src/struphy/pic/tests/test_cuda_emulation.py | 10 +-- src/struphy/pic/tests/test_draw_parallel.py | 4 +- src/struphy/pic/tests/test_estimate_mem.py | 8 +- src/struphy/pic/tests/test_kernel_setup.py | 2 +- src/struphy/pic/tests/test_marker_layout.py | 14 ++-- src/struphy/pic/tests/test_pushers.py | 4 +- .../propagators/efield_weights_coupling.py | 4 +- .../propagators/push_random_diffusion.py | 2 +- 20 files changed, 99 insertions(+), 98 deletions(-) diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index f3b1d89a1..f06e0c1d1 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -378,7 +378,7 @@ stencil matrix data) are in cunumpy since 0.6.1 (`Array6D`); `linear_vlasov_ampe - **Scalar types.** Members of the argument structs are checked when they are packed (PR 8). Scalars passed directly to a kernel (`dt`, `stage`) are not checked against the kernel signature: a Python `float` for an `int` parameter arrives wrong without an error. Either `CudaKernel` parses the `extern "C"` signature once and casts/checks (what cunumpy's `CudaKernel` does), or we accept it until the move to cunumpy. Until then: kernels take `double dt, int stage` in exactly this order, and the pyccel signature is the reference. - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. -- **Marker layout.** Decided (#707): the markers array is a C-ordered `n_cols × n_markers` array, one column per marker (struct-of-arrays), on both backends, see `struphy.pic.marker_layout`; it is indexed `markers[j, ip]` (pyccel) and `markers(j, ip)` (CUDA). Neighbouring threads (markers) read neighbouring addresses, and operations on one quantity (the pusher preamble, hole detection, masks in the diagnostics) touch only the rows involved. The pyccel kernels declare a plain `float[:, :]`; the marker count is `markers.shape[1]`, so every launch over markers passes `n_threads=args_markers.n_markers` (cunumpy's default inference takes the first axis of the first array). Arrays of evaluation points (`Domain`, `SplineFunction`, the equilibria) follow the same convention, `(k, n)` with the logical coordinates in rows 0 to 2. `n_rows` (the number of markers including holes) is a multiple of 32, so every row starts on a 256-byte boundary. +- **Marker layout.** Decided (#707): the markers array is a C-ordered `n_cols × n_markers` array, one column per marker (struct-of-arrays), on both backends, see `struphy.pic.marker_layout`; it is indexed `markers[j, ip]` (pyccel) and `markers(j, ip)` (CUDA). Neighbouring threads (markers) read neighbouring addresses, and operations on one quantity (the pusher preamble, hole detection, masks in the diagnostics) touch only the rows involved. The pyccel kernels declare a plain `float[:, :]`; the marker count is `markers.shape[1]`, so every launch over markers passes `n_threads=args_markers.n_markers` (cunumpy's default inference takes the first axis of the first array). Arrays of evaluation points (`Domain`, `SplineFunction`, the equilibria) follow the same convention, `(k, n)` with the logical coordinates in rows 0 to 2. `n_markers_max` (the number of markers including holes) is a multiple of 32, so every row starts on a 256-byte boundary. - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet. - **Single-source alternatives.** Decided (#690): no code generation; every CUDA kernel is written by hand, see [Hand-written CUDA](#hand-written-cuda-decision-for-690). - **Polar splines and MHD equilibria on the GPU** (left after PR 19): both run on CuPy now, polar splines since #695 diff --git a/src/struphy/examples/_draw_parallel.py b/src/struphy/examples/_draw_parallel.py index 7755219cf..cb2b95a4e 100644 --- a/src/struphy/examples/_draw_parallel.py +++ b/src/struphy/examples/_draw_parallel.py @@ -57,7 +57,7 @@ def main(): comm.Barrier() logger.info("Number of particles w/wo holes on each process before sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_markers_max}") domain.show( grid_info=derham.domain_array, @@ -70,7 +70,7 @@ def main(): comm.Barrier() logger.info("Number of particles w/wo holes on each process after sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_markers_max}") domain.show( grid_info=derham.domain_array, diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index 16e4361a7..879ee4e8f 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -64,7 +64,7 @@ def __init__( bc_type, pole: bool, ): - # shape (n_cols, n_rows), one column per marker (struphy.pic.marker_layout) + # shape (n_cols, n_markers_max), one column per marker (struphy.pic.marker_layout) self.markers = _device_array("markers", markers, np.float64, ndim=2) self.valid_mks = _device_array("valid_mks", valid_mks, np.bool_) self.Np = Np diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index 25b9cf7b7..0a35de380 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -11,7 +11,7 @@ class MarkerArguments: Paramaters ---------- markers : array[float] - Markers array of shape (n_cols, n_rows), one column per marker (see :mod:`struphy.pic.marker_layout`). + Markers array of shape (n_cols, n_markers_max), one column per marker (see :mod:`struphy.pic.marker_layout`). Np : int Total number of particles. diff --git a/src/struphy/models/linear_vlasov_ampere_one_species.py b/src/struphy/models/linear_vlasov_ampere_one_species.py index 9649b0276..b60d0f3f2 100644 --- a/src/struphy/models/linear_vlasov_ampere_one_species.py +++ b/src/struphy/models/linear_vlasov_ampere_one_species.py @@ -202,7 +202,7 @@ def _compute_en_w(self): else: self._f0 = backgrounds self._f0_values = xp.zeros( - self.kinetic_ions.var.particles.n_rows, + self.kinetic_ions.var.particles.n_markers_max, dtype=float, ) assert isinstance(self._f0, Maxwellian3D) diff --git a/src/struphy/models/tests/test_scalars.py b/src/struphy/models/tests/test_scalars.py index 01d888b93..30be2a162 100644 --- a/src/struphy/models/tests/test_scalars.py +++ b/src/struphy/models/tests/test_scalars.py @@ -148,9 +148,9 @@ class _MarkersWithGhosts: """Particles5D-like markers (weight in row 5, mu * |B| in row 8), plus holes and ghost copies of the first markers.""" def __init__(self, n_valid, n_ghosts): - n_rows = n_valid + n_ghosts + 2 + n_markers_max = n_valid + n_ghosts + 2 self.Np = n_valid - self.markers = xp.zeros((10, n_rows), dtype=float) + self.markers = xp.zeros((10, n_markers_max), dtype=float) self.markers[5, :n_valid] = xp.linspace(0.5, 1.5, n_valid) self.markers[8, :n_valid] = 3.0 self.markers[:, n_valid : n_valid + n_ghosts] = self.markers[:, :n_ghosts] diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index a183173b8..5747993cf 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -678,7 +678,7 @@ def estimate_mem( without allocating them. Constructs the same :class:`~struphy.pic.base.Particles` object as :meth:`allocate` would, but with - ``dry_run=True`` so that only the marker array sizing (:attr:`~struphy.pic.base.Particles.n_rows`, + ``dry_run=True`` so that only the marker array sizing (:attr:`~struphy.pic.base.Particles.n_markers_max`, :attr:`~struphy.pic.base.Particles.n_cols`) is computed.""" assert isinstance(self.backgrounds, KineticBackground), ( "List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." @@ -974,7 +974,7 @@ def estimate_mem( Constructs the same :class:`~struphy.pic.particles.ParticlesSPH` object as :meth:`allocate` would, but with ``dry_run=True`` so that only the marker array sizing - (:attr:`~struphy.pic.base.Particles.n_rows`, :attr:`~struphy.pic.base.Particles.n_cols`) is computed.""" + (:attr:`~struphy.pic.base.Particles.n_markers_max`, :attr:`~struphy.pic.base.Particles.n_cols`) is computed.""" assert isinstance(self.backgrounds, FluidEquilibrium), ( "List input not allowed; you can sum FluidEquilibrium objects before passing them to add_background." ) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 3ba5b9c6a..c4feb4d23 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -43,7 +43,7 @@ from struphy.pic.kernels.box_particles import box_particles from struphy.pic.kernels.fill_boxes import fill_boxes from struphy.pic.kernels.sort_and_fill_boxes import sort_and_fill_boxes -from struphy.pic.marker_layout import marker_array, padded_n_rows +from struphy.pic.marker_layout import marker_array, padded_n_markers from struphy.pic.pushing.kernels.reflect import reflect from struphy.pic.pushing.kernels.sph_mean_velocity_coeffs import sph_mean_velocity_coeffs from struphy.pic.pushing.kernels.sph_viscosity_tensor import sph_viscosity_tensor @@ -80,10 +80,10 @@ class Particles(metaclass=ABCMeta): r""" Base class for particle species. - The marker information is stored in a 2D array of shape ``(n_cols, n_rows)``, one column per marker + The marker information is stored in a 2D array of shape ``(n_cols, n_markers_max)``, one column per marker (see :mod:`struphy.pic.marker_layout`). In ``markers[j, ip]`` the index ``ip`` refers to a specific particle, - the index ``j`` to its attributes; ``n_cols`` is the number of attributes and ``n_rows`` the number of markers + the index ``j`` to its attributes; ``n_cols`` is the number of attributes and ``n_markers_max`` the number of markers (including holes), and the attribute indices are called "columns" below and in the rest of the code. They are as follows: @@ -185,7 +185,7 @@ class Particles(metaclass=ABCMeta): Normalization parameters (epsilon, alpha, ...) dry_run : bool - If True, only compute the sizing of the marker array (:attr:`n_rows`, :attr:`n_cols`, ...) + If True, only compute the sizing of the marker array (:attr:`n_markers_max`, :attr:`n_cols`, ...) and return early, without allocating any of the (potentially large) marker/sorting/buffer arrays. Used by :attr:`nbytes_local` to estimate the memory footprint before actually allocating the particles, see :meth:`~struphy.models.variables.PICVariable.estimate_mem`. @@ -417,9 +417,9 @@ def __init__( self._generate_sampling_moments() # create buffers for mpi_sort_markers - self._sorting_etas = xp.zeros((3, self.n_rows), dtype=float) - self._is_on_proc_domain = xp.zeros((3, self.n_rows), dtype=bool) - self._can_stay = xp.zeros(self.n_rows, dtype=bool) + self._sorting_etas = xp.zeros((3, self.n_markers_max), dtype=float) + self._is_on_proc_domain = xp.zeros((3, self.n_markers_max), dtype=bool) + self._can_stay = xp.zeros(self.n_markers_max, dtype=bool) self._reqs = [None] * self.mpi_size self._recvbufs = [None] * self.mpi_size self._send_to_i = [None] * self.mpi_size @@ -550,12 +550,13 @@ def n_cols(self): return self.first_free_idx + self.n_cols_aux + 2 @property - def n_rows(self): - """Total number of markers (columns) in markers array, including holes.""" - if not hasattr(self, "_n_rows"): + def n_markers_max(self): + """Number of marker slots (columns of :attr:`markers`), holes included: the most markers this process can + hold.""" + if not hasattr(self, "_n_markers_max"): input("\nWarning: marker array not yet created, creating now ...") self._allocate_marker_array() - return self._n_rows + return self._n_markers_max @property def mean_velocity_index(self): @@ -565,22 +566,22 @@ def mean_velocity_index(self): @property def nbytes_local(self) -> int: """Estimated local (per-MPI-rank) memory footprint, in bytes, of all marker-related arrays - (markers, sorting buffers, lost-marker container). Only depends on :attr:`n_rows` and + (markers, sorting buffers, lost-marker container). Only depends on :attr:`n_markers_max` and :attr:`n_cols`, so it is valid whether or not the arrays were actually allocated (see the ``dry_run`` argument of :meth:`__init__`).""" float_size = 8 # dtype=float bool_size = 1 # dtype=bool - n_rows = self.n_rows + n_markers_max = self.n_markers_max n_cols = self.n_cols nbytes = 0 - nbytes += n_rows * n_cols * float_size # markers - nbytes += n_rows * 3 * float_size # sorting_etas (mpi_sort_markers buffer) - nbytes += n_rows * 3 * bool_size # is_on_proc_domain - nbytes += n_rows * bool_size # can_stay + nbytes += n_markers_max * n_cols * float_size # markers + nbytes += n_markers_max * 3 * float_size # sorting_etas (mpi_sort_markers buffer) + nbytes += n_markers_max * 3 * bool_size # is_on_proc_domain + nbytes += n_markers_max * bool_size # can_stay # holes, ghost_particles, valid_mks, is_outside_right, is_outside_left, is_outside - nbytes += n_rows * bool_size * 6 - nbytes += int(n_rows * 0.5) * 10 * float_size # lost_markers + nbytes += n_markers_max * bool_size * 6 + nbytes += int(n_markers_max * 0.5) * 10 * float_size # lost_markers return int(nbytes) @property @@ -761,7 +762,7 @@ def n_mks_load(self): @property def markers(self): - """2D array of shape ``(n_cols, n_rows)`` holding the marker information, including holes. + """2D array of shape ``(n_cols, n_markers_max)`` holding the marker information, including holes. The i-th column holds the i-th marker info (see :mod:`struphy.pic.marker_layout`). ===== ============== ======================= ======= ====== ====== ========== === === === @@ -1890,7 +1891,7 @@ def apply_kinetic_bc(self, newton=False): # apply boundary conditions if self._remove_axes: - # snapshot of the (3, n_rows) logical coordinates for the loop over the axes + # snapshot of the (3, n_markers_max) logical coordinates for the loop over the axes # (see _find_outside_particles) self._eta_bc_buf[:] = self.markers[:3] @@ -2600,7 +2601,7 @@ def _n_mks_load_and_Np_per_clone(self): def _allocate_marker_array(self, dry_run: bool = False): """Create marker array :attr:`~struphy.pic.base.Particles.markers`. - If dry_run is True, only :attr:`n_rows` (and :attr:`n_cols`) are computed and no array + If dry_run is True, only :attr:`n_markers_max` (and :attr:`n_cols`) are computed and no array is actually allocated; see :attr:`nbytes_local`.""" if not hasattr(self, "_n_mks_load"): self._n_mks_load, self._Np_per_clone = self._n_mks_load_and_Np_per_clone() @@ -2611,7 +2612,7 @@ def _allocate_marker_array(self, dry_run: bool = False): # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer; # one column per marker, every row starting on a 256-byte boundary (see struphy.pic.marker_layout) - self._n_rows = padded_n_rows(round(float(n_mks_load_loc * (1 + bufsize)))) + self._n_markers_max = padded_n_markers(round(float(n_mks_load_loc * (1 + bufsize)))) # Have at least 3 spare places in markers array assert self.first_free_idx + 2 < self.n_cols - 2, ( @@ -2621,22 +2622,22 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return - self._markers = marker_array(self.n_cols, self.n_rows) + self._markers = marker_array(self.n_cols, self.n_markers_max) # allocate auxiliary arrays - self._holes = xp.zeros(self.n_rows, dtype=bool) - self._ghost_particles = xp.zeros(self.n_rows, dtype=bool) - self._valid_mks = xp.zeros(self.n_rows, dtype=bool) - self._is_outside_right = xp.zeros(self.n_rows, dtype=bool) - self._is_outside_left = xp.zeros(self.n_rows, dtype=bool) - self._is_outside = xp.zeros(self.n_rows, dtype=bool) + self._holes = xp.zeros(self.n_markers_max, dtype=bool) + self._ghost_particles = xp.zeros(self.n_markers_max, dtype=bool) + self._valid_mks = xp.zeros(self.n_markers_max, dtype=bool) + self._is_outside_right = xp.zeros(self.n_markers_max, dtype=bool) + self._is_outside_left = xp.zeros(self.n_markers_max, dtype=bool) + self._is_outside = xp.zeros(self.n_markers_max, dtype=bool) # scratch copy of markers[:3], refreshed once per apply_kinetic_bc boundary-condition-type # loop (see there) - self._eta_bc_buf = marker_array(3, self.n_rows) + self._eta_bc_buf = marker_array(3, self.n_markers_max) # create array container (3 x positions, vdim x velocities, weight, s0, w0, ID) for removed markers self._n_lost_markers = 0 - self._lost_markers = xp.zeros((10, int(self.n_rows * 0.5)), dtype=float) + self._lost_markers = xp.zeros((10, int(self.n_markers_max * 0.5)), dtype=float) # arguments for kernels: the pyccel class on the NumPy backend, its CUDA version on the CuPy backend args_class = CudaMarkerArguments if self._args_backend == "cupy" else MarkerArguments @@ -2938,15 +2939,15 @@ def _load_restart(self): with h5py.File(data.file_path, "a") as file: saved = file["restart/" + self.loading_params.restart_key][-1, :, :] - # the number of markers n_rows is padded (see struphy.pic.marker_layout), so a file written with another - # padding has a different number of markers: copy the common ones; the markers after them are holes, in the + # the number of marker slots n_markers_max is padded (see struphy.pic.marker_layout), so a file written with + # another padding has a different number of slots: copy the common ones; the slots after them are holes, in the # buffer (set before loading) and in the file - n_rows = min(saved.shape[1], self.n_rows) - assert np.all(saved[0, n_rows:] == -1.0), ( + n_common = min(saved.shape[1], self.n_markers_max) + assert np.all(saved[0, n_common:] == -1.0), ( f"the restart file has {saved.shape[1]} markers on rank {self.mpi_rank}, but there are only " - f"{self.n_rows} here, and the extra ones are not holes" + f"{self.n_markers_max} here, and the extra ones are not holes" ) - self._markers[:, :n_rows] = saved[:, :n_rows] + self._markers[:, :n_common] = saved[:, :n_common] def _load_tesselation(self, n_quad: int = 1): """ @@ -2990,7 +2991,7 @@ def _find_outside_particles(self, axis, eta=None): Row of the markers array (0, 1 or 2) holding the logical coordinate to check. eta : xp.ndarray[float], optional - Snapshot ``(3, n_rows)`` of ``markers[:3]`` (see :meth:`apply_kinetic_bc`). + Snapshot ``(3, n_markers_max)`` of ``markers[:3]`` (see :meth:`apply_kinetic_bc`). If ``None``, reads straight from ``markers``. Returns @@ -3162,7 +3163,7 @@ def _sort_boxed_particles_numpy(self): sorting_axis = self._sorting_boxes.box_index if not hasattr(self, "_argsort_array"): - self._argsort_array = xp.zeros(self.n_rows, dtype=int) + self._argsort_array = xp.zeros(self.n_markers_max, dtype=int) self._argsort_array[:] = self._markers[sorting_axis].argsort() self._markers[:] = self._markers[:, self._argsort_array] @@ -4481,7 +4482,7 @@ def _sendrecv_determine_mtbs( Indices of empty columns in markers after send. sorting_etas : array[float] - Eta-values of shape (3, n_rows) according to which the sorting is performed. + Eta-values of shape (3, n_markers_max) according to which the sorting is performed. """ # position that determines the sorting (including periodic shift of boundary conditions) if not isinstance(alpha, xp.ndarray): diff --git a/src/struphy/pic/marker_layout.py b/src/struphy/pic/marker_layout.py index d12d5ada2..7b2afcddc 100644 --- a/src/struphy/pic/marker_layout.py +++ b/src/struphy/pic/marker_layout.py @@ -1,7 +1,7 @@ """Memory layout of the marker buffer. -:attr:`Particles.markers ` has shape ``(n_cols, n_rows)``, one column per marker, -and is stored in C order (struct-of-arrays): every quantity (row ``j``) is one contiguous block of ``n_rows`` values. +:attr:`Particles.markers ` has shape ``(n_cols, n_markers_max)``, one column per marker, +and is stored in C order (struct-of-arrays): every quantity (row ``j``) is one contiguous block of ``n_markers_max`` values. ``markers[j, ip]`` is the ``j``-th entry of marker ``ip``; ``markers[j]`` is quantity ``j`` of all markers, and ``markers[:, ip]`` all quantities of marker ``ip``. Operations on quantities (``markers[j]``, ``markers[a:b] = ...``, a mask on one row) then touch only the rows involved, and on the GPU neighbouring threads (markers) read neighbouring @@ -20,16 +20,16 @@ import cunumpy as xp -ROW_ALIGNMENT = 32 -"""The number of markers ``n_rows`` is a multiple of this, so that every row of the buffer starts on a 256-byte +MARKER_ALIGNMENT = 32 +"""The number of marker slots ``n_markers_max`` is a multiple of this, so that every row of the buffer starts on a 256-byte boundary (32 doubles).""" -def padded_n_rows(n_rows: int) -> int: - """Smallest multiple of :data:`ROW_ALIGNMENT` that is at least ``n_rows``.""" - return -(-int(n_rows) // ROW_ALIGNMENT) * ROW_ALIGNMENT +def padded_n_markers(n_markers_max: int) -> int: + """Smallest multiple of :data:`MARKER_ALIGNMENT` that is at least ``n_markers_max``.""" + return -(-int(n_markers_max) // MARKER_ALIGNMENT) * MARKER_ALIGNMENT -def marker_array(n_cols: int, n_rows: int): - """Zero marker array of shape ``(n_cols, n_rows)`` on the active backend.""" - return xp.zeros((n_cols, n_rows), dtype=float) +def marker_array(n_cols: int, n_markers_max: int): + """Zero marker array of shape ``(n_cols, n_markers_max)`` on the active backend.""" + return xp.zeros((n_cols, n_markers_max), dtype=float) diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index c19e15919..a6d1126e4 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -8,21 +8,21 @@ * markers[residual_idx:-2] = 0.0 * * One thread per entry of the row range [first_pusher_idx, n_cols - 2) of every marker, numbered row by row. In the - * marker buffer (n_cols x n_rows, one column per marker, struphy.pic.marker_layout) neighbouring threads then read + * marker buffer (n_cols x n_markers, one column per marker, struphy.pic.marker_layout) neighbouring threads then read * and write neighbouring markers of the same row, and the whole range is written in a single launch (instead of one * per slice). * - * @param markers Marker buffer (n_cols x n_rows), every marker including holes. + * @param markers Marker buffer (n_cols x n_markers), every marker including holes. * @param first_pusher_idx First row of the saved initial phase space coordinates. * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ extern "C" __global__ void prepare_push(Array2D markers, int first_pusher_idx, int n_init) { - long long n_rows = markers.shape[1]; + long long n_markers = markers.shape[1]; long long width = markers.shape[0] - 2 - first_pusher_idx; long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; - if (t >= n_rows * width) return; + if (t >= n_markers * width) return; - long long k = t / n_rows; - long long ip = t - k * n_rows; + long long k = t / n_markers; + long long ip = t - k * n_markers; markers(first_pusher_idx + k, ip) = k < n_init ? markers(k, ip) : 0.0; } diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 0b8ae5e6a..dc572e2aa 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -177,7 +177,7 @@ def __init__( # profiling region names (cached, they are looked up on every call) self._region_name = "pusher: " + self.kernel.name - self._residuals = xp.zeros(self.particles.n_rows) + self._residuals = xp.zeros(self.particles.n_markers_max) self._converged_loc = self._residuals == 1.0 self._not_converged_loc = self._residuals == 0.0 @@ -219,12 +219,12 @@ def _push(self, dt: float): if self._cuda: # the three assignments below in one launch - n_cols, n_rows = markers.shape + n_cols, n_markers = markers.shape _prepare_push_cuda()( markers, first_pusher_idx, 3 + vdim, - n_threads=n_rows * (n_cols - 2 - first_pusher_idx), + n_threads=n_markers * (n_cols - 2 - first_pusher_idx), ) else: init_slice = slice(first_pusher_idx, first_shift_idx) diff --git a/src/struphy/pic/tests/bench_marker_layout.py b/src/struphy/pic/tests/bench_marker_layout.py index 74a1acc6f..14457b66e 100644 --- a/src/struphy/pic/tests/bench_marker_layout.py +++ b/src/struphy/pic/tests/bench_marker_layout.py @@ -78,12 +78,12 @@ def run(Np: int, repeats: int, num_elements: int, degree: int, save_markers: str def preamble(): """The row assignments at the start of Pusher._push.""" if xp.get_backend() == "cupy": - n_cols, n_rows = markers.shape + n_cols, n_markers = markers.shape _prepare_push_cuda()( markers, first_pusher_idx, n_init, - n_threads=n_rows * (n_cols - 2 - first_pusher_idx), + n_threads=n_markers * (n_cols - 2 - first_pusher_idx), ) else: markers[first_pusher_idx:first_shift_idx] = markers[:n_init] @@ -136,7 +136,7 @@ def accumulate(): info = { "backend": xp.get_backend(), - "n_rows": particles.n_rows, + "n_markers_max": particles.n_markers_max, "n_cols": particles.n_cols, "Np": Np, "c_contiguous": bool(markers.flags.c_contiguous), @@ -155,7 +155,7 @@ def _print(result: dict): info = result["info"] order = "C-ordered" if info["c_contiguous"] else "not C-contiguous" print( - f"backend {info['backend']}, markers {info['n_cols']} x {info['n_rows']} (one column per marker, {order}), " + f"backend {info['backend']}, markers {info['n_cols']} x {info['n_markers_max']} (one column per marker, {order}), " f"Np = {info['Np']}" ) width = max(len(name) for name in result["ms"]) diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index 7a8620a57..5b12dce49 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -114,20 +114,20 @@ def test_emulated_v1_symm_filler(): @requires_compiler -@pytest.mark.parametrize("n_rows", [96, 129]) -def test_emulated_prepare_push(n_rows): +@pytest.mark.parametrize("n_markers", [96, 129]) +def test_emulated_prepare_push(n_markers): """The fused preamble of Pusher._push (prepare_push_cuda.cu) writes the marker buffer (one column per marker) like the three row assignments, and leaves the rows before first_pusher_idx and the last two rows alone.""" from struphy.pic.marker_layout import marker_array from struphy.pic.pushing.pusher import _prepare_push_cuda n_cols, first_pusher_idx, n_init = 26, 8, 6 - markers = marker_array(n_cols, n_rows) - markers[:] = np.random.default_rng(11).random((n_cols, n_rows)) + markers = marker_array(n_cols, n_markers) + markers[:] = np.random.default_rng(11).random((n_cols, n_markers)) expected = markers.copy() expected[first_pusher_idx : first_pusher_idx + n_init] = expected[:n_init] expected[first_pusher_idx + n_init : -2] = 0.0 width = n_cols - 2 - first_pusher_idx - emulate_struct_kernel(_prepare_push_cuda(), markers, first_pusher_idx, n_init, n_threads=n_rows * width) + emulate_struct_kernel(_prepare_push_cuda(), markers, first_pusher_idx, n_init, n_threads=n_markers * width) np.testing.assert_array_equal(markers, expected) diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index 1a6ca591e..47f418e60 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -105,7 +105,7 @@ def test_draw(num_elements, degree, bcs, mapping, ppc=10): comm.Barrier() logger.info("Number of particles w/wo holes on each process before sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_markers_max}") # sort particles according to domain decomposition comm.Barrier() @@ -113,7 +113,7 @@ def test_draw(num_elements, degree, bcs, mapping, ppc=10): comm.Barrier() logger.info("Number of particles w/wo holes on each process after sorting : ") - logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_rows}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_markers_max}") # are all markers in the correct domain? conds = xp.logical_and( diff --git a/src/struphy/pic/tests/test_estimate_mem.py b/src/struphy/pic/tests/test_estimate_mem.py index 973f7059c..1a89bfb71 100644 --- a/src/struphy/pic/tests/test_estimate_mem.py +++ b/src/struphy/pic/tests/test_estimate_mem.py @@ -17,7 +17,7 @@ def _make_domain_decomp(mpi_comm, num_elements=(8, 6, 4), degree=(2, 2, 2)): @pytest.mark.parametrize("Np", [1000, 12345]) def test_dry_run_does_not_allocate_markers(Np): - """dry_run=True must compute the marker array sizing (n_rows, n_cols) without + """dry_run=True must compute the marker array sizing (n_markers_max, n_cols) without allocating the (potentially large) marker/sorting/buffer arrays.""" mpi_comm = MPI.COMM_WORLD domain, domain_decomp = _make_domain_decomp(mpi_comm) @@ -37,7 +37,7 @@ def test_dry_run_does_not_allocate_markers(Np): ) # sizing info is available ... - assert particles.n_rows > 0 + assert particles.n_markers_max > 0 assert particles.n_cols > 0 # ... but none of the (large) arrays that allocate() would create exist @@ -61,7 +61,7 @@ def test_dry_run_does_not_allocate_markers(Np): def test_nbytes_local_matches_real_allocation(Np): """The dry_run estimate (nbytes_local) must equal the real memory footprint once the particles are actually allocated (dry_run=False), since both use the exact same - n_rows/n_cols sizing and the same (fixed) list of marker-related arrays.""" + n_markers_max/n_cols sizing and the same (fixed) list of marker-related arrays.""" mpi_comm = MPI.COMM_WORLD domain, domain_decomp = _make_domain_decomp(mpi_comm) @@ -84,7 +84,7 @@ def test_nbytes_local_matches_real_allocation(Np): real = Particles6D(**common_kwargs, dry_run=False) real.draw_markers(sort=False) - assert dry.n_rows == real.n_rows + assert dry.n_markers_max == real.n_markers_max assert dry.n_cols == real.n_cols real_nbytes = ( diff --git a/src/struphy/pic/tests/test_kernel_setup.py b/src/struphy/pic/tests/test_kernel_setup.py index bbd8e2bf7..4b71089bf 100644 --- a/src/struphy/pic/tests/test_kernel_setup.py +++ b/src/struphy/pic/tests/test_kernel_setup.py @@ -263,7 +263,7 @@ def sort(**kwargs): particles = SimpleNamespace( markers=marker_args.markers, args_markers=marker_args, - n_rows=2, + n_markers_max=2, n_cols=40, vdim=3, first_pusher_idx=8, diff --git a/src/struphy/pic/tests/test_marker_layout.py b/src/struphy/pic/tests/test_marker_layout.py index 0cd453b2c..d7509c6f3 100644 --- a/src/struphy/pic/tests/test_marker_layout.py +++ b/src/struphy/pic/tests/test_marker_layout.py @@ -14,7 +14,7 @@ import pytest from struphy.models.tests.gpu_e2e import device_context, requires_cupy_backend -from struphy.pic.marker_layout import ROW_ALIGNMENT, marker_array, padded_n_rows +from struphy.pic.marker_layout import MARKER_ALIGNMENT, marker_array, padded_n_markers # CuPy on a GPU, or cunumpy's fake CuPy BACKENDS = ["numpy", pytest.param("cupy", marks=requires_cupy_backend)] @@ -24,8 +24,8 @@ def test_padded_n_rows(): - assert [padded_n_rows(n) for n in (0, 1, 31, 32, 33, 64, 65)] == [0, 32, 32, 32, 64, 64, 96] - assert ROW_ALIGNMENT == 32 + assert [padded_n_markers(n) for n in (0, 1, 31, 32, 33, 64, 65)] == [0, 32, 32, 32, 64, 64, 96] + assert MARKER_ALIGNMENT == 32 @pytest.mark.parametrize("backend", BACKENDS) @@ -97,7 +97,7 @@ def test_kernel_accepts_marker_copies(backend): @pytest.mark.parametrize("backend", BACKENDS) def test_marker_buffer_layout(backend): """Loading, pushing, boundary conditions, sorting between ranks, evaluation at markers and binning keep the - buffer of shape (n_cols, n_rows) and C-ordered, and evaluation at copies of the valid markers gives the values + buffer of shape (n_cols, n_markers_max) and C-ordered, and evaluation at copies of the valid markers gives the values at the buffer.""" from maybempi import MPI @@ -129,8 +129,8 @@ def test_marker_buffer_layout(backend): particles.initialize_weights() buffer = particles.markers - assert buffer.shape == (particles.n_cols, particles.n_rows) and buffer.flags.c_contiguous - assert particles.n_rows % ROW_ALIGNMENT == 0 + assert buffer.shape == (particles.n_cols, particles.n_markers_max) and buffer.flags.c_contiguous + assert particles.n_markers_max % MARKER_ALIGNMENT == 0 butcher = ButcherTableau("rk4") push_eta = Pusher( @@ -158,7 +158,7 @@ def test_marker_buffer_layout(backend): # the buffer is never replaced, only written to assert particles.markers is buffer - assert buffer.shape == (particles.n_cols, particles.n_rows) and buffer.flags.c_contiguous + assert buffer.shape == (particles.n_cols, particles.n_markers_max) and buffer.flags.c_contiguous # evaluation at copies of the valid markers (Fortran-ordered on NumPy, converted by Domain and SplineFunction) # equals evaluation at the buffer diff --git a/src/struphy/pic/tests/test_pushers.py b/src/struphy/pic/tests/test_pushers.py index 0e4b5024b..e98e638bd 100644 --- a/src/struphy/pic/tests/test_pushers.py +++ b/src/struphy/pic/tests/test_pushers.py @@ -578,7 +578,7 @@ def test_push_bxu_Hdiv_pauli(num_elements, degree, bcs, mapping, show_plots=Fals flattened=True, ) - mu0 = xp.zeros(particles.n_rows, dtype=float) + mu0 = xp.zeros(particles.n_markers_max, dtype=float) pusher_psy = Pusher_psy( particles, @@ -708,7 +708,7 @@ def test_push_eta_rk4(num_elements, degree, bcs, mapping, show_plots=False): n_mks_load = xp.zeros(size, dtype=int) - comm.Allgather(xp.array(particles.n_rows), n_mks_load) + comm.Allgather(xp.array(particles.n_markers_max), n_mks_load) sendcounts = xp.zeros(size, dtype=int) displacements = xp.zeros(size, dtype=int) diff --git a/src/struphy/propagators/efield_weights_coupling.py b/src/struphy/propagators/efield_weights_coupling.py index 2990b7e62..fe4838b3e 100644 --- a/src/struphy/propagators/efield_weights_coupling.py +++ b/src/struphy/propagators/efield_weights_coupling.py @@ -204,8 +204,8 @@ def allocate(self): self._e_sum = e.space.zeros() # marker storage - self._f0_values = xp.zeros(particles.n_rows, dtype=float) - self._old_weights = xp.empty(particles.n_rows, dtype=float) + self._f0_values = xp.zeros(particles.n_markers_max, dtype=float) + self._old_weights = xp.empty(particles.n_markers_max, dtype=float) # ================================ # ========= Schur Solver ========= diff --git a/src/struphy/propagators/push_random_diffusion.py b/src/struphy/propagators/push_random_diffusion.py index 1d1881627..d31d58b69 100644 --- a/src/struphy/propagators/push_random_diffusion.py +++ b/src/struphy/propagators/push_random_diffusion.py @@ -150,7 +150,7 @@ def __call__(self, dt): self._noise[:] = random.multivariate_normal( self._mean, self._cov, - particles.n_rows, + particles.n_markers_max, ).T # push markers From 17b9cd5fa4b2995fbd4c94877a63a58bfe06f9e5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:09:23 +0200 Subject: [PATCH 12/12] test_marker_layout: test_padded_n_markers, after padded_n_markers Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/test_marker_layout.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/test_marker_layout.py b/src/struphy/pic/tests/test_marker_layout.py index d7509c6f3..08844c034 100644 --- a/src/struphy/pic/tests/test_marker_layout.py +++ b/src/struphy/pic/tests/test_marker_layout.py @@ -23,7 +23,7 @@ MARKER_INDICES = (3, 6, 7, 8, 14, 17, 18, 4) # vdim, weight_idx, ..., mu_idx -def test_padded_n_rows(): +def test_padded_n_markers(): assert [padded_n_markers(n) for n in (0, 1, 31, 32, 33, 64, 65)] == [0, 32, 32, 32, 64, 64, 96] assert MARKER_ALIGNMENT == 32