Densities on the grid: Field.deposit_density and node volumes

CHANGELOG.md

@@ -16,6 +16,12 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Added
 
+- `Field.deposit_density(positions, values, weighting=LINEAR)`: adds the deposit of the
+  marker values divided by the node volumes, with the halo exchange on a split field; with
+  the products `w`, `w v`, `w v^2`, ... as components it gives the velocity moments.
+  Example 14 derives the flow, temperature and heat flux from them.
+- `Field.node_volumes()`: the volume each node collects in a deposit, half a cell on a bounded
+  edge. `Field.on_grid` takes `num_ghost_points` (default 1).
 - Method cards (methodcard) on the numerics: the linear and quadratic shape functions
   (`linear_cell`, `quadratic_cell`), the deposit and the gather (`Weighting.scatter_block`,
   `Weighting.gather_block`) and their adjointness (`Weighting`), the block grid and the halo

README.md

@@ -53,10 +53,11 @@ phi.plot(name="potential", units="V")  # through plasma-plots
   `local_with_halos`, `to_numpy`), halo exchange (`fill_halos`,
   `exchange_halos`), gather and scatter (`gather`, `scatter`),
   arithmetic and ufuncs (a scalar field broadcasts over a vector field’s
-  components), `gradient`, `cross`, `dot`, `magnitude`, global
-  reductions, smoothing filters, `set_inside` for simgeometry shapes,
-  `evaluate` to xarray, `plot` with plasma-plots, and NetCDF
-  `export`/`load`.
+  components), densities (`deposit_density`: a deposit over
+  `node_volumes`, e.g. velocity moments), `gradient`, `cross`, `dot`,
+  `magnitude`, global reductions, smoothing filters, `set_inside` for
+  simgeometry shapes, `evaluate` to xarray, `plot` with plasma-plots,
+  and NetCDF `export`/`load`.
 - `Weighting`: a deposit and its gather for each number of dimensions,
   and the halo nodes they need. `LINEAR` (the default) and `QUADRATIC`
   are built in; a weighting of your own is a pair of functions or

README.qmd

@@ -54,7 +54,8 @@ phi.plot(name="potential", units="V")  # through plasma-plots
   data access (`set_field_data`, `local`, `local_with_halos`, `to_numpy`), halo exchange
   (`fill_halos`, `exchange_halos`), gather and scatter (`gather`, `scatter`), arithmetic and
   ufuncs (a scalar field broadcasts over a vector field's
-  components), `gradient`, `cross`, `dot`, `magnitude`, global reductions, smoothing filters,
+  components), densities (`deposit_density`: a deposit over `node_volumes`, e.g. velocity
+  moments), `gradient`, `cross`, `dot`, `magnitude`, global reductions, smoothing filters,
   `set_inside` for simgeometry shapes, `evaluate` to xarray, `plot` with plasma-plots, and
   NetCDF `export`/`load`.
 - `Weighting`: a deposit and its gather for each number of dimensions, and the halo nodes they

docs/src/content/docs/numerics/index.mdx

@@ -70,6 +70,9 @@ The split deposit and gather equal the serial ones.
   zero outside a bounded axis, wrapped on a periodic one.
 - **[Prescribed fields](/mpifield/numerics/prescribed/)**: $\mathbf E = -\nabla\phi$ of a
   given potential, e.g. a table along one axis, for test particles.
+- **[Densities](/mpifield/numerics/field/)**: a deposit divided by the volume each node
+  collects, half a cell on a bounded edge (`deposit_density`), e.g. the velocity moments
+  $n$, $n u$, $n\langle v^2\rangle$ of the products $w$, $w v$, $w v^2$.
 - **[Evaluation](/mpifield/numerics/evaluate/)** as xarray: linear interpolation between the
   nodes, whatever the weighting.
 

examples/14_velocity_moments.py

@@ -0,0 +1,56 @@
+"""Velocity moments
+
+Deposit the density, flow, temperature and heat flux of markers that sample a non-uniform drifting Maxwellian.
+
+`Field.deposit_density` deposits marker values and divides them by the volume each node
+collects (`Field.node_volumes`: a cell, half of it on a bounded edge), so the result is a
+density in physical units. With the products $w$, $w v$, $w v^2$, $w v^3$ as four
+components it gives the raw velocity moments $n$, $n u$, $n\\langle v^2\\rangle$,
+$n\\langle v^3\\rangle$ in one deposit (mpiparticles' `moment_values` builds these products
+for any order and number of velocity components). The central moments follow: the
+temperature over the mass $T = \\langle v^2\\rangle - u^2$ and the heat flux
+$q = n\\langle (v - u)^3\\rangle / 2 = (M_3 - 3 u M_2 + 2 n u^3) / 2$. The markers sample
+$n(x) = 1 + x/2$, $u(x) = \\sin \\pi x$ and $T(x) = 1 + x^2$ on a bounded grid, so the heat
+flux is zero: what remains is the noise of $2 \\times 10^5$ markers.
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from mpifield import Field
+
+rng = np.random.default_rng(2)
+count = 200_000
+moments = Field(
+    ndim=1, lower_bound=(0.0,), upper_bound=(1.0,), num_gridpoints=41, num_components=4
+)
+
+# positions from n(x) = 1 + x / 2 (inverse transform), velocities from the local Maxwellian
+uniform = rng.uniform(0.0, 1.0, count)
+x = 2.0 * (np.sqrt(1.0 + 1.25 * uniform) - 1.0)
+v = np.sin(np.pi * x) + np.sqrt(1.0 + x**2) * rng.normal(size=count)
+weights = np.full(count, 1.25 / count)  # total: the integral of n
+
+moments.deposit_density(x[None, :], np.stack([weights * v**k for k in range(4)]))
+n, m1, m2, m3 = moments.to_numpy().T
+u = m1 / n
+temperature = m2 / n - u**2
+heat_flux = (m3 - 3 * u * m2 + 2 * n * u**3) / 2
+
+nodes = np.linspace(0.0, 1.0, 41)
+exact = {
+    "density $n$": (n, 1 + nodes / 2),
+    "flow $u$": (u, np.sin(np.pi * nodes)),
+    "temperature $T/m$": (temperature, 1 + nodes**2),
+    "heat flux $q$": (heat_flux, 0 * nodes),
+}
+fig, axes = plt.subplots(2, 2, figsize=(9, 6), sharex=True, layout="constrained")
+for ax, (label, (deposited, expected)) in zip(axes.flat, exact.items(), strict=True):
+    ax.plot(nodes, deposited, "o", ms=3, label="deposited")
+    ax.plot(nodes, expected, "k--", label="exact")
+    ax.set_title(label)
+for ax in axes[1]:
+    ax.set_xlabel("$x$")
+axes[0, 0].legend(frameon=False)
+print(f"max |T - exact| = {np.abs(temperature - (1 + nodes**2)).max():.3f}")
+plt.show()

src/mpifield/_method_checks.py

@@ -388,3 +388,26 @@ def seam(x: float, g_0: float, g_3: float) -> float:
 
     rule = g3 + (g0 - g3) * (x - sympy.Rational(3, 4)) * 4
     return {"seam": Rule(rule, seam, {"x": (0.75, 1.0)})}
+
+
+def density_rules() -> dict[str, Rule]:
+    """The rule of the density deposit, with the compiled linear scatter.
+
+    - ``total``: two markers, one in the boundary cell of a bounded axis; the
+      density times the node volumes sums to their values.
+    """
+    import sympy
+
+    q1, q2 = sympy.symbols("q_1 q_2", real=True)
+
+    def total(q_1: float, q_2: float) -> float:
+        from mpifield.field import Field
+
+        field = Field(
+            ndim=1, lower_bound=(0.0,), upper_bound=(1.0,), num_gridpoints=NODES
+        )
+        field.deposit_density(np.array([[0.03, 0.61]]), np.array([[q_1, q_2]]))
+        volumes = field.node_volumes().to_numpy()
+        return float((field.to_numpy() * volumes).sum())
+
+    return {"total": Rule(q1 + q2, total, {"q_1": (-2.0, 2.0)}, rtol=1e-12)}

src/mpifield/field.py

@@ -247,13 +247,20 @@ def block_grid(self, cpu: bool = False, inside_domain: bool = False) -> BlockGri
     # Construction ------------------------------------------------------------#
 
     @classmethod
-    def on_grid(cls, grid: Grid, num_components: int = 1, split: bool = True) -> Self:
+    def on_grid(
+        cls,
+        grid: Grid,
+        num_components: int = 1,
+        split: bool = True,
+        num_ghost_points: int = 1,
+    ) -> Self:
         """Return a zero field on the nodes of ``grid``, split as ``grid`` is.
 
-        On a spatially decomposed grid each rank holds its own nodes plus one
-        halo node on every side, which the deposit accumulates into and the
-        interpolation reads from (`block_grid`); otherwise, or with
-        ``split=False``, every rank holds the whole grid.
+        On a spatially decomposed grid each rank holds its own nodes plus
+        ``num_ghost_points`` halo nodes on every side (1 for `LINEAR`, 2 for
+        `QUADRATIC`), which the deposit accumulates into and the interpolation
+        reads from (`block_grid`); otherwise, or with ``split=False``, every rank
+        holds the whole grid.
         """
         decomposition = cls._decomposition_of(grid)
         distributed = (
@@ -268,7 +275,7 @@ def on_grid(cls, grid: Grid, num_components: int = 1, split: bool = True) -> Sel
             upper_bound=grid.upper_bound,
             num_gridpoints=grid.num_gridpoints,
             comm=grid.comm if distributed else None,
-            num_ghost_points=1 if distributed else 0,
+            num_ghost_points=num_ghost_points if distributed else 0,
             domain_decomposition=decomposition.decompose if distributed else None,
             periodic=grid.periodic,
         )
@@ -553,6 +560,33 @@ def local_with_halos(self) -> Array:
         """
         return self.field_data.local_with_halos
 
+    def node_volumes(self) -> Self:
+        """Return the volume each node collects in a deposit, as a one-component field.
+
+        A node collects one cell, ``prod(grid_spacing)``, and on a bounded axis
+        the first and last node half of it (the cell reaches beyond the domain
+        on one side). A deposit divided by it is a density whose sum times the
+        volumes is the total deposited. The halo nodes are filled from the
+        neighbouring ranks; those beyond a bounded domain edge stay zero.
+        """
+        volumes = self._like(num_components=1)
+        periodic = _periodic_flags(self.periodic, self.ndim)
+        local = xp.ones(volumes.local.shape[:-1], dtype=self.dtype)
+        for axis in range(self.ndim):
+            start, stop = self.proc_index_bounds[axis]
+            factor = np.full(stop - start, float(self.grid_spacing[axis]))
+            if not periodic[axis]:
+                nodes = np.arange(start, stop)
+                last = int(self.num_gridpoints[axis]) - 1
+                factor[(nodes == 0) | (nodes == last)] *= 0.5
+            shape = [1] * self.ndim
+            shape[axis] = stop - start
+            local = local * xp.asarray(factor).reshape(shape)
+        volumes.local[..., 0] = local
+        if volumes.distributed:
+            volumes.fill_halos()
+        return volumes
+
     def component(self, component: int) -> DistributedArray:
         """Return one component as a distributed array sharing this field's storage."""
         if not 0 <= component < self.num_components:
@@ -1145,6 +1179,44 @@ def scatter(
         if field_data_cpu is not data:
             data[...] = xp.asarray(field_data_cpu)
 
+    @method(
+        equations={
+            r"density": r"n_{\mathbf k} \mathrel{+}= \frac{1}{V_{\mathbf k}} \sum_i W_{\mathbf k}(\mathbf x_i)\, q_i",
+            r"volume": r"V_{\mathbf k} = \prod_d h_d \cdot \prod_{d\ \text{bounded}} \left(1 - \tfrac12 [k_d = 0] - \tfrac12 [k_d = N_d - 1]\right)",
+            r"total": r"\sum_{\mathbf k} V_{\mathbf k}\, n_{\mathbf k} = \sum_i q_i",
+        },
+        assumptions=(
+            "the weighting $W$ of `scatter`, linear by default; each value is one component, e.g. $w$, $w v_x$, $w v_x v_y$ for the velocity moments",
+            "a bounded boundary node collects half a cell, so its deposit is divided by half the volume",
+            "on a split field each rank passes the markers of its own block; the halos are exchanged here (collective)",
+        ),
+        components=("mpifield.weighting.Weighting.scatter_block",),
+        symbolic=_method_checks.density_rules,
+    )
+    @sp.profile("mpifield.Field.deposit_density", tags=PROFILE_TAGS)
+    def deposit_density(
+        self,
+        positions: Array | float,
+        values: Array | float,
+        weighting: Weighting = LINEAR,
+    ) -> None:
+        """Add the density of the marker ``values`` to the nodes: a deposit over the node volumes.
+
+        Scatters ``values`` (``(ncomp, N)``, or ``(N,)`` for one component) into
+        a zero field like this one, exchanges its halos on a split field, divides
+        it by `node_volumes` and adds it to this field's nodes; the halo nodes are
+        left as they are. With a velocity moment per component (``w``,
+        ``w v_x``, ...) this gives ``n``, ``n u_x``, ... Collective on a split
+        field. With markers spread over the ranks of a replicated field, sum the
+        result over the ranks.
+        """
+        deposit = self._like()
+        deposit.scatter(positions, values, weighting)
+        if deposit.distributed:
+            deposit.exchange_halos()
+        local = self.local
+        local += deposit.local / self.node_volumes().local
+
     # Arithmetic --------------------------------------------------------------#
 
     def _aligned_data(self, other: Field) -> DistributedArray:

src/mpifield/tests/parallel/test_field_mpi.py

@@ -520,3 +520,44 @@ def test_export_and_load_under_mpi(tmp_path_factory: pytest.TempPathFactory) ->
     )
 
     np.testing.assert_array_equal(loaded.to_numpy()[..., 0], global_data)
+
+
+@pytest.mark.parametrize(
+    ("decomposition", "periodic"),
+    [([True, False], (False, False)), ([True, True], (True, False))],
+)
+@pytest.mark.parametrize("weighting", [LINEAR, QUADRATIC], ids=lambda w: w.name)
+def test_split_density_equals_the_serial_one(
+    weighting, decomposition, periodic
+) -> None:
+    """Each rank deposits the density of its own markers; the result is the serial one."""
+    comm = MPI.COMM_WORLD
+    rng = np.random.default_rng(21)
+    positions = rng.uniform(-0.05, 1.05, (2, 3000))
+    for axis, flag in enumerate(periodic):
+        if flag:
+            positions[axis] %= 1.0
+    values = rng.normal(0.2, 1.0, (3, 3000))
+    kwargs = {
+        "ndim": 2,
+        "lower_bound": (0.0, 0.0),
+        "upper_bound": (1.0, 1.0),
+        "num_gridpoints": (24, 20),
+        "num_components": 3,
+        "periodic": periodic,
+    }
+    serial = Field(**kwargs)
+    serial.deposit_density(positions, values, weighting)
+    grid = Field(comm=comm, domain_decomposition=decomposition, **kwargs)
+    field = Field.on_grid(grid, num_components=3, num_ghost_points=weighting.halo)
+    assert field.distributed == (comm.Get_size() > 1)
+    mine = _own_markers(field, positions)
+    assert comm.allreduce(int(mine.sum())) == positions.shape[1]
+    own = np.ascontiguousarray(positions[:, mine])
+    field.deposit_density(own, np.ascontiguousarray(values[:, mine]), weighting)
+    # signed values cancel between markers: compare to the scale of the density
+    expected = serial.to_numpy()
+    scale = np.abs(expected).max(axis=(0, 1))
+    np.testing.assert_allclose(field.to_numpy() / scale, expected / scale, atol=1e-12)
+    volumes = field.node_volumes().to_numpy()
+    np.testing.assert_allclose(volumes, serial.node_volumes().to_numpy(), rtol=1e-14)

src/mpifield/tests/unit/test_density.py

@@ -0,0 +1,64 @@
+"""Tests for the density deposit and the node volumes."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from mpifield import LINEAR, QUADRATIC, Field
+
+
+def _field(periodic: tuple[bool, ...], num_components: int = 1) -> Field:
+    return Field(
+        ndim=2,
+        lower_bound=(0.0, 0.0),
+        upper_bound=(1.0, 1.0),
+        num_gridpoints=(11, 9),
+        num_components=num_components,
+        periodic=periodic,
+    )
+
+
+@pytest.mark.parametrize("weighting", [LINEAR, QUADRATIC], ids=lambda w: w.name)
+@pytest.mark.parametrize("periodic", [(False, False), (False, True), (True, True)])
+def test_the_density_times_the_node_volumes_is_the_total(weighting, periodic) -> None:
+    """Nothing is lost, also at the edges where a node collects half a cell."""
+    rng = np.random.default_rng(0)
+    positions = rng.uniform(0.0, 1.0, (2, 5000))
+    values = rng.normal(0.0, 1.0, (3, 5000))
+    field = _field(periodic, 3)
+    field.deposit_density(positions, values, weighting)
+    totals = (field.to_numpy() * field.node_volumes().to_numpy()).sum(axis=(0, 1))
+    np.testing.assert_allclose(totals, values.sum(axis=1), rtol=1e-12, atol=1e-12)
+
+
+def test_a_uniform_cloud_has_a_uniform_density_up_to_the_edges() -> None:
+    """The half volumes make the boundary nodes as dense as the inner ones."""
+    count = 200_000
+    positions = np.random.default_rng(1).uniform(0.0, 1.0, (2, count))
+    field = _field((False, False))
+    field.deposit_density(positions, np.full(count, 1.0 / count))
+    density = field.to_numpy()[..., 0]
+    assert density.mean() == pytest.approx(1.0, abs=0.02)
+    assert density[0].mean() == pytest.approx(1.0, abs=0.05)
+    assert density[:, -1].mean() == pytest.approx(1.0, abs=0.05)
+
+
+def test_deposit_density_adds_to_the_field() -> None:
+    field = _field((True, True))
+    field.set_field_data(2.0)
+    field.deposit_density(np.full((2, 1), 0.5), 0.0)
+    np.testing.assert_allclose(field.to_numpy(), 2.0)
+    density = _field((True, True))
+    density.deposit_density(np.full((2, 1), 0.5), 1.0)
+    field.deposit_density(np.full((2, 1), 0.5), 1.0)
+    np.testing.assert_allclose(field.to_numpy(), 2.0 + density.to_numpy())
+
+
+def test_node_volumes_halve_on_bounded_edges() -> None:
+    volumes = _field((False, True)).node_volumes().to_numpy()[..., 0]
+    cell = 0.1 * (1.0 / 9)
+    assert volumes.shape == (11, 9)
+    np.testing.assert_allclose(volumes[1:-1], cell)
+    np.testing.assert_allclose(volumes[[0, -1]], cell / 2)
+    assert volumes.sum() == pytest.approx(1.0)