diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index c2c09d408..f06e0c1d1 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.** 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 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 @@ -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` 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/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py b/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py index fbc6fcdd3..70ef59f6d 100644 --- a/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py +++ b/examples/DriftKineticElectrostaticAdiabatic/cyclone/params_cyclone.py @@ -264,7 +264,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 @@ -277,7 +277,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 @@ -300,7 +300,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 e398a6048..c0e613ee6 100644 --- a/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py +++ b/examples/DriftKineticElectrostaticAdiabatic/itg_cylindre/params_drift_kinetic.py @@ -257,7 +257,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 @@ -266,7 +266,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 @@ -281,7 +281,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 @@ -290,7 +290,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 d71b4d832..ee397af37 100644 --- a/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py +++ b/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py @@ -192,15 +192,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) @@ -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 diff --git a/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py b/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py index a803b72f6..5cb4f818a 100644 --- a/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py +++ b/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py @@ -229,13 +229,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 3a92cdc09..dae8d5401 100644 --- a/examples/ToyDrift/diocotron_instability/params_diocotron.py +++ b/examples/ToyDrift/diocotron_instability/params_diocotron.py @@ -228,7 +228,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/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 9b166e165..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 @@ -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/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/examples/_draw_parallel.py b/src/struphy/examples/_draw_parallel.py index 6c7c1632d..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.markers.shape[0]}") + 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.markers.shape[0]}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_markers_max}") 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/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 31fca09c3..93dbe684e 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -2900,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 @@ -2929,10 +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 - markers = xp.array(etas[0][:, :3], dtype=float) + # 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 @@ -3102,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 @@ -3118,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/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/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/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/base.py b/src/struphy/geometry/base.py index 44c48ccd1..0ed8ebb90 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -28,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): @@ -499,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 @@ -588,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 @@ -627,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 @@ -678,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 @@ -729,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 @@ -785,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 @@ -819,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 @@ -871,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 @@ -901,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 @@ -952,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 @@ -982,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 @@ -1030,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 @@ -1054,10 +1054,11 @@ def _evaluate_metric_coefficient(self, *etas, which=0, **kwargs): # markers evaluation if len(etas) == 1: - markers = 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( @@ -1067,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 @@ -1164,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 @@ -1183,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 @@ -1217,18 +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] + # 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, ), @@ -1237,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, ) @@ -1248,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, ), @@ -1257,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: @@ -1267,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 @@ -1283,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( @@ -1294,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 @@ -1560,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: @@ -1592,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, ), ] @@ -1826,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. @@ -2131,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") @@ -2149,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/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 94cb2494c..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 @@ -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/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/kernel_arguments/array_view.cuh b/src/struphy/kernel_arguments/array_view.cuh index 12d8afa7c..9f07d50b0 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; 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 feac42098..879ee4e8f 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -64,14 +64,13 @@ 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.") + # 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 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 963bbbacb..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. + 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. @@ -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/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/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/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 83d3f647e..437ac7780 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[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..b60d0f3f2 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_markers_max, 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/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..30be2a162 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 + n_markers_max = 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_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] + 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/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." diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index c65bf00b5..5747993cf 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, ) @@ -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." @@ -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, ) @@ -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/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 b03038867..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, row-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 e4ce83e46..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, row-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 880538d1e..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, row-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/base.py b/src/struphy/pic/base.py index e3529b9c4..c4feb4d23 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -43,6 +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_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 @@ -79,10 +80,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_markers_max)``, one column per marker + (see :mod:`struphy.pic.marker_layout`). + In ``markers[j, ip]`` the index ``ip`` refers to a specific particle, + the index ``j`` to its attributes; ``n_cols`` is the number of attributes and ``n_markers_max`` the number of markers + (including holes), and the attribute indices are called "columns" below and in the rest of the code. + They are as follows: * ``0:3``: position in the logical unit cube (:math:`\boldsymbol \eta_p \in [0, 1]^3`) * ``3:3 + vdim``: velocities @@ -113,7 +116,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. @@ -182,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`. @@ -414,9 +417,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_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 @@ -542,17 +545,18 @@ 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.""" - 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): @@ -562,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 @@ -758,8 +762,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_markers_max)`` holding the marker information, including holes. + The i-th column holds the i-th marker info (see :mod:`struphy.pic.marker_layout`). ===== ============== ======================= ======= ====== ====== ========== === === === index | 0 | 1 | 2 | | 3 | ... | 3+(vdim-1)| 3+vdim 4+vdim 5+vdim >=6+vdim ... -2 -1 @@ -767,7 +771,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 @@ -776,25 +780,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): @@ -840,98 +844,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): @@ -939,12 +943,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: @@ -1178,7 +1182,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() @@ -1212,7 +1216,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: @@ -1236,21 +1240,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 @@ -1271,7 +1277,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 @@ -1282,7 +1288,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 @@ -1301,12 +1307,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: @@ -1317,14 +1323,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, ) @@ -1340,21 +1346,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. @@ -1365,18 +1371,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] @@ -1399,8 +1405,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}.' @@ -1413,13 +1419,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 @@ -1431,7 +1434,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 @@ -1492,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 @@ -1512,7 +1515,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]: @@ -1520,18 +1523,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( @@ -1553,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): @@ -1562,14 +1565,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 @@ -1617,42 +1620,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 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 indices of the valid markers, found once; below, only the rows (quantities) that are used are gathered + # for them, not all quantities of the markers. Everything stays on the device of the markers. + valid = xp.nonzero(self.valid_mks)[0] + if quantity != "density": + velocities = self.markers[self.index["vel"], valid] # determine histogram weights multiplier if quantity == "density": multiplier = 1 elif quantity == "current": - multiplier = velocities[:, v_axis[0]] + 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 = markers[:, self.index["w0"]] * self.Np * multiplier - _weights = markers[:, 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(markers[:, 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([markers[:, 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, @@ -1787,7 +1792,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. @@ -1798,7 +1803,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 @@ -1809,7 +1814,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`). """ @@ -1832,7 +1837,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) @@ -1843,7 +1848,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() @@ -1853,7 +1858,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)) @@ -1879,17 +1884,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: - # 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) - self._eta_bc_buf[:] = self.markers[:, :3] + # 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] for axis in self._remove_axes: outside_inds = self._find_outside_particles(axis, eta=self._eta_bc_buf) @@ -1900,16 +1904,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) @@ -1917,47 +1921,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 @@ -1992,11 +1981,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): @@ -2010,10 +1999,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") @@ -2052,15 +2039,15 @@ 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. - The box array :attr:`~struphy.pic.sorting.SortingBoxes.boxes` is rebuilt for the sorted rows, + The box array :attr:`~struphy.pic.sorting.SortingBoxes.boxes` is rebuilt for the sorted markers, so it stays valid after the sort. Parameters ---------- use_numpy_argsort : bool - If True, sort via :func:`numpy.argsort` on the box column; if False (default), + If True, sort via :func:`numpy.argsort` on the box row; if False (default), use the kernel :func:`~struphy.pic.kernels.sort_and_fill_boxes.sort_and_fill_boxes`. """ boxes = self._sorting_boxes @@ -2102,7 +2089,7 @@ def do_sort(self, use_numpy_argsort=False): ) self._check_box_capacity(max_in_box) - # 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() @@ -2388,7 +2375,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 @@ -2504,9 +2491,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. @@ -2614,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() @@ -2623,8 +2610,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; + # one column per marker, every row starting on a 256-byte boundary (see struphy.pic.marker_layout) + self._n_markers_max = padded_n_markers(round(float(n_mks_load_loc * (1 + bufsize)))) # Have at least 3 spare places in markers array assert self.first_free_idx + 2 < self.n_cols - 2, ( @@ -2634,23 +2622,22 @@ 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_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) - # 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) + self._holes = xp.zeros(self.n_markers_max, dtype=bool) + self._ghost_particles = xp.zeros(self.n_markers_max, dtype=bool) + self._valid_mks = xp.zeros(self.n_markers_max, dtype=bool) + self._is_outside_right = xp.zeros(self.n_markers_max, dtype=bool) + self._is_outside_left = xp.zeros(self.n_markers_max, dtype=bool) + self._is_outside = xp.zeros(self.n_markers_max, dtype=bool) + # scratch copy of markers[:3], refreshed once per apply_kinetic_bc boundary-condition-type + # loop (see there) + self._eta_bc_buf = marker_array(3, self.n_markers_max) # create array container (3 x positions, vdim x velocities, weight, s0, w0, ID) for removed markers self._n_lost_markers = 0 - self._lost_markers = xp.zeros((int(self.n_rows * 0.5), 10), 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 @@ -2846,7 +2833,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( @@ -2871,7 +2858,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( @@ -2902,6 +2889,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 @@ -2914,24 +2904,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.""" @@ -2950,7 +2937,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 marker slots n_markers_max is padded (see struphy.pic.marker_layout), so a file written with + # another padding has a different number of slots: copy the common ones; the slots after them are holes, in the + # buffer (set before loading) and in the file + n_common = min(saved.shape[1], self.n_markers_max) + assert np.all(saved[0, n_common:] == -1.0), ( + f"the restart file has {saved.shape[1]} markers on rank {self.mpi_rank}, but there are only " + f"{self.n_markers_max} here, and the extra ones are not holes" + ) + self._markers[:, :n_common] = saved[:, :n_common] def _load_tesselation(self, n_quad: int = 1): """ @@ -2968,9 +2965,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): @@ -2991,20 +2988,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 - 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 ``(3, n_markers_max)`` 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 @@ -3046,12 +3041,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: @@ -3062,7 +3057,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. @@ -3070,13 +3065,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): @@ -3112,13 +3107,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) @@ -3143,12 +3138,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) @@ -3159,7 +3154,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) @@ -3168,13 +3163,13 @@ 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_markers_max, 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 _fill_boxes(self): - """Assign the particles to boxes from the box column of the markers, after checking that the box + """Assign the particles to boxes from the box row of the markers, after checking that the box array has enough columns (see :meth:`_check_box_capacity`).""" max_in_box = fill_boxes( self.markers, @@ -3205,15 +3200,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): @@ -3238,20 +3233,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"): @@ -3287,16 +3282,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"): @@ -3318,16 +3313,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"): @@ -3349,16 +3344,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"): @@ -3384,24 +3379,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]): @@ -3450,16 +3445,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, ) @@ -3468,25 +3463,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, ) @@ -3495,27 +3490,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, ) @@ -3524,25 +3519,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, ) @@ -3551,27 +3546,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, ) @@ -3580,25 +3575,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, ) @@ -3607,13 +3602,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 @@ -3629,14 +3624,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): @@ -3647,155 +3644,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): @@ -3821,14 +3850,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): """ @@ -3860,7 +3891,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 @@ -3883,7 +3914,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 @@ -4244,7 +4275,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) @@ -4289,7 +4320,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:: @@ -4414,7 +4445,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). @@ -4422,10 +4453,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) @@ -4436,13 +4467,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 @@ -4451,7 +4482,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_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): @@ -4464,39 +4495,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] @@ -4574,10 +4604,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 @@ -4586,10 +4616,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 @@ -4604,17 +4634,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 @@ -4640,7 +4672,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`). """ @@ -4660,7 +4692,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": @@ -4679,7 +4711,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 @@ -4701,7 +4733,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 @@ -4724,14 +4758,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. @@ -4741,7 +4775,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()) @@ -4756,8 +4790,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 = [] @@ -4766,12 +4801,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( @@ -4781,8 +4815,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/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..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 @@ -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 350bcc803..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 @@ -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 f13d23a25..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 @@ -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 67721d35e..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 @@ -15,23 +15,23 @@ def eval_energy_5d( 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 c5c08d3d4..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 @@ -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 0f1a8201a..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 @@ -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 d079bb11b..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 @@ -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 68dc460f3..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 @@ -15,22 +15,22 @@ def eval_magnetic_moment_5d( ): """ 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/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py b/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py index cad32ca93..11d5579d9 100644 --- a/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py +++ b/src/struphy/pic/kernels/box_and_sort_particles/box_and_sort_particles_kernels.py @@ -18,9 +18,9 @@ def box_and_sort_particles( swap_line_2: "float[:]", box_index: "int" = -2, ) -> "int": - """``box_particles`` followed by ``sort_and_fill_boxes``: write the box of each particle into column ``box_index``, - count the particles per box, reorder the marker rows by box, then update ``holes`` and fill ``boxes`` for the - sorted rows. On the CuPy backend the markers cross between device and host once in each direction. + """``box_particles`` followed by ``sort_and_fill_boxes``: write the box of each particle into row ``box_index``, + count the particles per box, reorder the markers (columns) by box, then update ``holes`` and fill ``boxes`` for + the sorted markers. On the CuPy backend the markers cross between device and host once in each direction. Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, nothing is sorted, ``boxes`` is left untouched and the caller must raise. diff --git a/src/struphy/pic/kernels/box_particles/box_particles_kernels.py b/src/struphy/pic/kernels/box_particles/box_particles_kernels.py index 00d26e86a..a0943fa6f 100644 --- a/src/struphy/pic/kernels/box_particles/box_particles_kernels.py +++ b/src/struphy/pic/kernels/box_particles/box_particles_kernels.py @@ -15,9 +15,9 @@ def box_particles( next_index: "int[:]", box_index: "int" = -2, ) -> "int": - """Write the box of each particle into column ``box_index`` of the markers, count the particles per box into + """Write the box of each particle into row ``box_index`` of the markers, count the particles per box into ``next_index`` and, if every box fits into a row of ``boxes`` (the last column stays -1 as terminator), write the - particle rows into ``boxes``. + particle indices into ``boxes``. Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, ``boxes`` is left untouched and the caller must raise. diff --git a/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py b/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py index 91dc53643..681d4bfff 100644 --- a/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py +++ b/src/struphy/pic/kernels/fill_boxes/fill_boxes_kernels.py @@ -1,4 +1,4 @@ -"""Fill the box array from the box column of the markers, see :meth:`~struphy.pic.base.Particles.put_particles_in_boxes`.""" +"""Fill the box array from the box row of the markers, see :meth:`~struphy.pic.base.Particles.put_particles_in_boxes`.""" # module import (not names): the module may still be an older compiled version until `struphy compile` runs import struphy.pic.sorting_kernels as sorting_kernels @@ -11,8 +11,8 @@ def fill_boxes( next_index: "int[:]", box_index: "int" = -2, ) -> "int": - """Count the particles per box (box number in column ``box_index`` of the markers) into ``next_index`` and, if - every box fits into a row of ``boxes`` (the last column stays -1 as terminator), write the particle rows into + """Count the particles per box (box number in row ``box_index`` of the markers) into ``next_index`` and, if + every box fits into a row of ``boxes`` (the last column stays -1 as terminator), write the particle indices into ``boxes``. Returns the largest number of particles in a box; if it is ``>= boxes.shape[1]``, ``boxes`` is left untouched and diff --git a/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py b/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py index 674e47dfc..b11da56ce 100644 --- a/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py +++ b/src/struphy/pic/kernels/sort_and_fill_boxes/sort_and_fill_boxes_kernels.py @@ -14,7 +14,7 @@ def sort_and_fill_boxes( swap_line_2: "float[:]", box_index: "int" = -2, ): - """Reorder the marker rows by box, then update ``holes`` and rebuild ``boxes`` for the new rows, so that the box + """Reorder the markers (columns) by box, then update ``holes`` and rebuild ``boxes`` for the new order, so that the box array stays valid after the sort. ``next_index`` must hold the number of particles per box, as left by ``box_particles`` or ``fill_boxes``. The diff --git a/src/struphy/pic/marker_layout.py b/src/struphy/pic/marker_layout.py new file mode 100644 index 000000000..7b2afcddc --- /dev/null +++ b/src/struphy/pic/marker_layout.py @@ -0,0 +1,35 @@ +"""Memory layout of the marker buffer. + +: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 +addresses. + +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. + +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_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_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_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/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/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 7a62cb4c8..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, row-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 9a6c24d52..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, row-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 e5ede6bf2..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, row-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 7d2a270ec..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, row-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 71ef3a461..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, row-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 f3dfb6e00..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, row-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 b32974a0b..3624c0c29 100644 --- a/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py +++ b/src/struphy/pic/pushing/kernels/reflect/reflect_kernels.py @@ -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 71c02556d..a6d1126e4 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 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 row range [first_pusher_idx, n_cols - 2) of every marker, numbered row by row. In the + * 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_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 first_pusher_idx First column of the saved initial phase space coordinates. + * @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(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_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 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_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 c9271883f..dc572e2aa 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_markers_max) self._converged_loc = self._residuals == 1.0 self._not_converged_loc = self._residuals == 0.0 @@ -218,28 +218,26 @@ 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) - n_rows, n_cols = markers.shape + # the three assignments below in one launch + n_cols, n_markers = 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), + n_threads=n_markers * (n_cols - 2 - first_pusher_idx), ) else: init_slice = slice(first_pusher_idx, first_shift_idx) 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} ...") @@ -287,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 @@ -302,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 @@ -324,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.py b/src/struphy/pic/sorting.py index 4df93b0e3..bf345e86b 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -225,7 +225,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), @@ -243,8 +243,8 @@ def _set_boxes(self): initialize_neighbours(neighbours, self.nx, self.ny, self.nz) self._neighbours = xp.to_cunumpy(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/sorting_kernels.py b/src/struphy/pic/sorting_kernels.py index 4d073e807..65c1ab452 100644 --- a/src/struphy/pic/sorting_kernels.py +++ b/src/struphy/pic/sorting_kernels.py @@ -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,12 +187,12 @@ 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 @@ -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,7 +226,7 @@ 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( @@ -239,12 +239,12 @@ 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 @@ -255,13 +255,13 @@ def count_particles_in_boxes( next_index: "int[:]", box_index: "int" = -2, ) -> "int": - """Count the particles in each box (box number in column ``box_index``) into ``next_index`` and return the + """Count the particles in each box (box number in row ``box_index``) into ``next_index`` and return the largest count.""" 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 not holes[p]: - a = int(markers[p, l + box_index]) + a = int(markers[l + box_index, p]) next_index[a] += 1 max_in_box = 0 @@ -272,6 +272,6 @@ def count_particles_in_boxes( def find_holes(markers: "float[:,:]", holes: "bool[:]"): - """Mark the empty rows of the marker array (``markers[:, 0] == -1``), as ``Particles.update_holes``.""" - for p in range(markers.shape[0]): - holes[p] = markers[p, 0] == -1.0 + """Mark the empty columns of the marker array (``markers[0] == -1``), as ``Particles.update_holes``.""" + for p in range(markers.shape[1]): + holes[p] = markers[0, p] == -1.0 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/tests/bench_marker_layout.py b/src/struphy/pic/tests/bench_marker_layout.py new file mode 100644 index 000000000..14457b66e --- /dev/null +++ b/src/struphy/pic/tests/bench_marker_layout.py @@ -0,0 +1,195 @@ +""" +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; 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 +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 row assignments at the start of Pusher._push.""" + if xp.get_backend() == "cupy": + n_cols, n_markers = markers.shape + _prepare_push_cuda()( + markers, + first_pusher_idx, + n_init, + n_threads=n_markers * (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 (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, + "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 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_markers_max": particles.n_markers_max, + "n_cols": particles.n_cols, + "Np": Np, + "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; 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])].T) + return result + + +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_markers_max']} (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") + + +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..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.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.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 02f331bec..10d28bea4 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -25,16 +25,16 @@ 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 @@ -213,14 +213,15 @@ 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 (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 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..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)) - 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 0e69c148d..5b12dce49 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_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_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_markers * width) + np.testing.assert_array_equal(markers, expected) 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..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.markers.shape[0]}") + logger.info(f"Rank {rank} : {particles.n_mks_loc} {particles.n_markers_max}") # 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_markers_max}") # 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_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_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 75e3178d3..c0c0a65a8 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -53,7 +53,7 @@ def make_arguments(n_markers, seed=0): rng = np.random.default_rng(seed) - markers = 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 @@ -112,7 +112,7 @@ def test_argument_classes_correspond(cuda_class, source, class_name, cuda_only, 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 + # 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] 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 @@ -149,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 @@ -231,7 +231,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 00b9d127c..4b71089bf 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)) - 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)) - 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_markers_max=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 new file mode 100644 index 000000000..08844c034 --- /dev/null +++ b/src/struphy/pic/tests/test_marker_layout.py @@ -0,0 +1,192 @@ +"""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:: + + 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 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)] + +N_COLS = 25 +MARKER_INDICES = (3, 6, 7, 8, 14, 17, 18, 4) # vdim, weight_idx, ..., mu_idx + + +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 + + +@pytest.mark.parametrize("backend", BACKENDS) +def test_marker_array(backend): + with xp.use_backend(backend): + 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 + + +def _marker_arguments(markers, valid): + from struphy.kernel_arguments.pusher_args_cuda import CudaMarkerArguments + from struphy.kernel_arguments.pusher_args_kernels import MarkerArguments + + 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): + """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_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 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 + + 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 + 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( + 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 + 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 + 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(pulled), + xp.to_numpy(domain.pull(lambda x, y, z: x * y + z, buffer, kind="3", remove_outside=False)[valid]), + ) + + # 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(columns, local=True) + assert values.shape == (n_valid,) + np.testing.assert_array_equal(xp.to_numpy(values), xp.to_numpy(field(buffer, local=True)[valid])) + + # 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) + 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 92b682e92..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)) - 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)) - 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..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.markers.copy().T.shape[1], 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(xp.shape(particles.markers)[0]), 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) @@ -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 8b2e69de7..742311e08 100644 --- a/src/struphy/pic/tests/test_reflect.py +++ b/src/struphy/pic/tests/test_reflect.py @@ -22,13 +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): - storage = xp.asarray(original.copy()) - 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) @@ -68,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) @@ -76,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 7aa363362..5a771dfff 100644 --- a/src/struphy/pic/tests/test_sorting.py +++ b/src/struphy/pic/tests/test_sorting.py @@ -17,14 +17,14 @@ def assert_boxes_match_markers(particles): - """The box array lists, for every box, exactly the rows of the markers (not holes) whose box column is that box, - in row order and padded with -1, and ``next_index`` holds their number.""" + """The box array lists, for every box, exactly the markers (columns, not holes) whose box row is that box, + in index order and padded with -1, and ``next_index`` holds their number.""" sorting_boxes = particles.sorting_boxes boxes = xp.to_numpy(sorting_boxes.boxes) next_index = xp.to_numpy(sorting_boxes._next_index) markers = xp.to_numpy(particles.markers) holes = xp.to_numpy(particles.holes) - box_of_row = markers[:, sorting_boxes.box_index].astype(int) + box_of_row = markers[sorting_boxes.box_index].astype(int) for box in range(boxes.shape[0]): rows = np.flatnonzero(~holes & (box_of_row == box)) assert next_index[box] == rows.size, box @@ -167,7 +167,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)) assert_boxes_match_markers(particles) @@ -215,10 +215,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) @@ -226,7 +226,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 ac5c7aa88..4fd52bba0 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -82,10 +82,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 @@ -121,11 +121,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 @@ -145,14 +145,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) @@ -220,7 +220,7 @@ def make_particles(): # all markers of this rank lie in its domain, so none is in the out-of-domain box (last row of the box array): # the boxes were found with the domain_array row of this rank out_box = ref.sorting_boxes.boxes.shape[0] - 1 - assert np.all(ref.markers[ref.valid_mks, ref.sorting_boxes.box_index] < out_box) + assert np.all(ref.markers[ref.sorting_boxes.box_index, ref.valid_mks] < out_box) with xp.use_backend("cupy"): dev = make_particles() 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/pic/utilities_kernels.py b/src/struphy/pic/utilities_kernels.py index 73ec15bdb..c56cb329f 100644 --- a/src/struphy/pic/utilities_kernels.py +++ b/src/struphy/pic/utilities_kernels.py @@ -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 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) 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/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/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..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.markers.shape[0], dtype=float) - self._old_weights = xp.empty(particles.markers.shape[0], 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 ========= @@ -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 2101ef0f4..d31d58b69 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]) + # 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: @@ -149,8 +150,8 @@ def __call__(self, dt): self._noise[:] = random.multivariate_normal( self._mean, self._cov, - len(particles.markers), - ) + particles.n_markers_max, + ).T # push markers self._pusher(dt) 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/propagators/vlasov_ampere_coupling.py b/src/struphy/propagators/vlasov_ampere_coupling.py index 77a129ae0..02bf5ed52 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 rows (quantities) needed, not all quantities of the markers + not_holes = ~particles.holes + 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() 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):