Moments and binning

Download notebook View on GitHubRun when the site was built.

The velocity distribution of the markers is summarised by its moments, the sums $\sum_p w_p v_p^{i_1} \cdots v_p^{i_k}$ of every order $k$ : the total weight, the momentum over the mass, $\sum w v_i v_j$ , … velocity_moments sums them over the markers of all ranks, and phase_space_density bins the markers on one to three phase-space axes. Both are collective on the communicator they get; without one they make no MPI call. This tutorial loads two counter-streaming beams into a store and looks at them both ways.

import matplotlib.pyplot as plt
import numpy as np

from mpiparticles import (
    ParticleStore,
    moment_indices,
    moment_name,
    moment_values,
    phase_space_density,
    velocity_moments,
)

rng = np.random.default_rng(0)
count = 100_000
beam = np.where(rng.uniform(size=count) < 0.5, -1.0, 1.0)  # two beams, at -1 and +1
velocities = np.stack(
    [beam + 0.3 * rng.normal(size=count), 0.5 * rng.normal(size=count)]
)
positions = rng.uniform(0.0, 1.0, (1, count))
store = ParticleStore(ndim=1, nvel=2)
_ = store.add(positions, velocities, np.full(count, 1.0 / count))

Moments

moment_indices(nvel, max_order) lists the moments in the order the functions return them: order by order, each symmetric moment once ((0, 1) but not (1, 0)). moment_name spells an index with the letters of its velocity components.

indices = moment_indices(store.nvel, 2)
sums = velocity_moments(store.active_velocities(), store.active_weights(), max_order=2)
for index, value in zip(indices, sums, strict=True):
    print(f"{moment_name(index) or '1':>3}: {value: .4f}")

  1:  1.0000
  x: -0.0012
  y: -0.0001
 xx:  1.0922
 xy: -0.0019
 yy:  0.2506

Divided by the total weight they are averages. The mean velocity vanishes, but $\langle v_x^2 \rangle = 1 + 0.3^2$ : the spread between the beams counts as temperature in a fluid description, which is why the second moment alone cannot tell a hot plasma from two cold beams.

averages = dict(zip(map(moment_name, indices), sums / sums[0], strict=True))
temperature_x = averages["xx"] - averages["x"] ** 2
print(f"<v_x> = {averages['x']:.4f}, T_x / m = {temperature_x:.4f}")

<v_x> = -0.0012, T_x / m = 1.0922

moment_values gives the per-marker products $w v_i v_j \cdots$ instead of their sums, (num_moments, N). Deposited on a grid they are the moments on the nodes; mpifield’s Field.deposit_density(positions, moment_values(...)) does that, with the halo exchange on a split field.

print(moment_values(store.active_velocities(), store.active_weights(), 2).shape)

(6, 100000)

Phase-space density

phase_space_density sums the weights in each bin and divides by the bin volume. The axes are a mapping {name: (lower, upper, num_bins)} (or BinAxis objects), with the names x, y, z, vx, vy, vz, in the order of the dimensions of the result. Markers outside the ranges are left out, so the integral over the bins is the weight inside them.

axes = {"vx": (-2.5, 2.5, 100), "vy": (-2.0, 2.0, 80)}
density = phase_space_density(
    store.active_positions(), store.active_velocities(), store.active_weights(), axes
)
print(f"weight inside: {density.sum() * 0.05 * 0.05:.4f}")

fig, ax = plt.subplots(figsize=(6, 3.5), layout="constrained")
mesh = ax.pcolormesh(
    np.linspace(-2.5, 2.5, 101), np.linspace(-2.0, 2.0, 81), density.T, cmap="magma"
)
ax.set(xlabel="$v_x$", ylabel="$v_y$")
fig.colorbar(mesh, label="$f$")
plt.show()

weight inside: 1.0000

Under MPI

Each rank passes its own markers and the communicator: velocity_moments(..., comm=comm) and phase_space_density(..., comm=comm) return the sums over all ranks, the same on every rank (one Allreduce).