From af25c231714965d0bc1f44e144f7340c5ed77a62 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 3 Nov 2026 19:44:50 +0100 Subject: [PATCH] Add Field.deposit_density and Field.node_volumes --- CHANGELOG.md | 6 ++ README.md | 9 +- README.qmd | 3 +- docs/src/content/docs/numerics/index.mdx | 3 + examples/14_velocity_moments.py | 56 +++++++++++++ src/mpifield/_method_checks.py | 23 +++++ src/mpifield/field.py | 84 +++++++++++++++++-- src/mpifield/tests/parallel/test_field_mpi.py | 41 +++++++++ src/mpifield/tests/unit/test_density.py | 64 ++++++++++++++ 9 files changed, 278 insertions(+), 11 deletions(-) create mode 100644 examples/14_velocity_moments.py create mode 100644 src/mpifield/tests/unit/test_density.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f4cd00e..7cc80ad 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index 08bf745..a172819 100644 --- a/README.md +++ b/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 diff --git a/README.qmd b/README.qmd index 5133481..5746836 100644 --- a/README.qmd +++ b/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 diff --git a/docs/src/content/docs/numerics/index.mdx b/docs/src/content/docs/numerics/index.mdx index 72984be..c5f35b9 100644 --- a/docs/src/content/docs/numerics/index.mdx +++ b/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. diff --git a/examples/14_velocity_moments.py b/examples/14_velocity_moments.py new file mode 100644 index 0000000..e0c5958 --- /dev/null +++ b/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() diff --git a/src/mpifield/_method_checks.py b/src/mpifield/_method_checks.py index b888cd0..cc717af 100644 --- a/src/mpifield/_method_checks.py +++ b/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)} diff --git a/src/mpifield/field.py b/src/mpifield/field.py index 62c4f63..ef8e336 100644 --- a/src/mpifield/field.py +++ b/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: diff --git a/src/mpifield/tests/parallel/test_field_mpi.py b/src/mpifield/tests/parallel/test_field_mpi.py index bdcc1de..7c053ac 100644 --- a/src/mpifield/tests/parallel/test_field_mpi.py +++ b/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) diff --git a/src/mpifield/tests/unit/test_density.py b/src/mpifield/tests/unit/test_density.py new file mode 100644 index 0000000..6d8ca76 --- /dev/null +++ b/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)