diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 1398a25..426a103 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -39,5 +39,10 @@ jobs: example-scripts/01-cloud-in-cell-density/index.html example-scripts/01-cloud-in-cell-density/plot-1.png example-scripts/12-kernel-backends/index.html + numerics/index.html + numerics/kernels-linear-cell/index.html + numerics/kernels-quadratic-cell/index.html + numerics/weighting/index.html + numerics/field/index.html api/mpifield/index.html api/mpifield/field/index.html diff --git a/CHANGELOG.md b/CHANGELOG.md index ba792ea..f4cd00e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,21 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ### Added +- 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 + exchanges, the gradient, the smoothing filters, the prescribed field and the evaluation, + with their assumptions and references. `test_method_cards.py` checks them: 19 SymPy rules + (the shape functions, partition of unity and first moment, edge folds and seams through the + compiled kernels; the operators), the second order of the two gathers, the gradient and + `electric_field_from_potential`, and the cards of the docs against those of the code. Adds + the `methodcard` dependency (`methodcard[symbolic]` in the test extra). +- A Numerics section in the docs (header and sidebar): an overview + (`docs/src/content/docs/numerics/index.mdx`) and one page per module + (starlight-theme's `methodcard: { pages: true }`, theme 0.7.4); each method's API page shows + its card. +- The cross references in the docstrings render on the API pages: code spans instead of + `:class:`, `:meth:` and `:func:` roles. - `Field`, moved out of picnic (`picnic.field`): a field on the nodes of a uniform 1D, 2D or 3D grid, split over MPI ranks with halo nodes (an mpiarray `DistributedArray` on a simgeometry `Grid`), with gather (`gather`, was `interpolate_field`) and scatter diff --git a/docs/astro.config.mjs b/docs/astro.config.mjs index 16f0eb8..401b855 100644 --- a/docs/astro.config.mjs +++ b/docs/astro.config.mjs @@ -8,9 +8,25 @@ export default defineDocs({ title: 'mpifield', description: 'Fields on uniform grids split over MPI ranks, with compiled linear gather and scatter for particle codes.', pythonPackage: 'mpifield', + // the method cards: a Numerics section with the overview (numerics/index.mdx) and one page per + // module, from the shape functions to the field operations, and a Numerics section on each + // method's API page + methodcard: { + pages: true, + index: false, + labels: { + 'mpifield.kernels.linear_cell': 'Linear weights', + 'mpifield.kernels.quadratic_cell': 'Quadratic weights', + 'mpifield.weighting': 'Deposit and gather', + 'mpifield.field': 'Halos, gradient, smoothing', + 'mpifield.prescribed': 'Prescribed fields', + 'mpifield.evaluate': 'Evaluation', + }, + }, nav: [ { label: 'Getting started', link: '/getting-started/installation/' }, { label: 'Examples', link: '/examples/' }, + { label: 'Numerics', link: '/numerics/' }, { label: 'Development', link: '/development/docs/' }, { label: 'API', link: '/api/mpifield/' }, ], diff --git a/docs/package-lock.json b/docs/package-lock.json index 92569a0..c2403b9 100644 --- a/docs/package-lock.json +++ b/docs/package-lock.json @@ -1223,8 +1223,8 @@ "license": "MIT" }, "node_modules/@max-models/starlight-theme": { - "version": "0.7.0", - "resolved": "git+ssh://git@github.com/max-models/starlight-theme.git#35ba7be44b183dd82e92f638a46d095a3fbaa23d", + "version": "0.7.4", + "resolved": "git+ssh://git@github.com/max-models/starlight-theme.git#24889b06c9553ce4af91ea2563ce9130ce190ef4", "license": "MIT", "workspaces": [ "example/docs" @@ -1232,7 +1232,7 @@ "dependencies": { "@astrojs/markdown-remark": "^7.3.2", "@mdx-js/mdx": "^3.1.1", - "katex": "^0.19.0", + "katex": "^0.16.47", "plotly.js-dist-min": "^4.1.2", "rehype-katex": "^7.0.1", "remark-math": "^6.0.0", @@ -2301,12 +2301,12 @@ } }, "node_modules/commander": { - "version": "15.0.0", - "resolved": "https://registry.npmjs.org/commander/-/commander-15.0.0.tgz", - "integrity": "sha512-z67u4ZhzCL/Tydu1lJARtEZYWbWaN7oYLHbsuzocr6y4N6WZAagG3RQ4FW61V1/0+jImpj293XfrcYnd1qxtPg==", + "version": "8.3.0", + "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz", + "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==", "license": "MIT", "engines": { - "node": ">=22.12.0" + "node": ">= 12" } }, "node_modules/common-ancestor-path": { @@ -3541,16 +3541,16 @@ "license": "MIT" }, "node_modules/katex": { - "version": "0.19.0", - "resolved": "https://registry.npmjs.org/katex/-/katex-0.19.0.tgz", - "integrity": "sha512-v6Tznz3zJ7u3niRCoDTsumM2+HA2XXcCu+WAacCeHD2z3p9A9Ks987o5FzfTGBN0e8A0vjEgIDvLbICrXpdw/Q==", + "version": "0.16.47", + "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.47.tgz", + "integrity": "sha512-Eeo8Ys1doU1z+x8AZsPpQu+p/QcZBI5PeOo7QGQdy2x2m0MU/hYagBbGOmXwr5KVbEfVuWv9LpnQWeehogurjg==", "funding": [ "https://opencollective.com/katex", "https://github.com/sponsors/katex" ], "license": "MIT", "dependencies": { - "commander": "^15.0.0" + "commander": "^8.3.0" }, "bin": { "katex": "cli.js" @@ -4456,31 +4456,6 @@ "url": "https://opencollective.com/unified" } }, - "node_modules/micromark-extension-math/node_modules/commander": { - "version": "8.3.0", - "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz", - "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==", - "license": "MIT", - "engines": { - "node": ">= 12" - } - }, - "node_modules/micromark-extension-math/node_modules/katex": { - "version": "0.16.47", - "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.47.tgz", - "integrity": "sha512-Eeo8Ys1doU1z+x8AZsPpQu+p/QcZBI5PeOo7QGQdy2x2m0MU/hYagBbGOmXwr5KVbEfVuWv9LpnQWeehogurjg==", - "funding": [ - "https://opencollective.com/katex", - "https://github.com/sponsors/katex" - ], - "license": "MIT", - "dependencies": { - "commander": "^8.3.0" - }, - "bin": { - "katex": "cli.js" - } - }, "node_modules/micromark-extension-mdx-expression": { "version": "3.0.1", "resolved": "https://registry.npmjs.org/micromark-extension-mdx-expression/-/micromark-extension-mdx-expression-3.0.1.tgz", @@ -5540,31 +5515,6 @@ "url": "https://opencollective.com/unified" } }, - "node_modules/rehype-katex/node_modules/commander": { - "version": "8.3.0", - "resolved": "https://registry.npmjs.org/commander/-/commander-8.3.0.tgz", - "integrity": "sha512-OkTL9umf+He2DZkUq8f8J9of7yL6RJKI24dVITBmNfZBmri9zYZQrKkuXiKhyfPSu8tUhnVBB1iKXevvnlR4Ww==", - "license": "MIT", - "engines": { - "node": ">= 12" - } - }, - "node_modules/rehype-katex/node_modules/katex": { - "version": "0.16.47", - "resolved": "https://registry.npmjs.org/katex/-/katex-0.16.47.tgz", - "integrity": "sha512-Eeo8Ys1doU1z+x8AZsPpQu+p/QcZBI5PeOo7QGQdy2x2m0MU/hYagBbGOmXwr5KVbEfVuWv9LpnQWeehogurjg==", - "funding": [ - "https://opencollective.com/katex", - "https://github.com/sponsors/katex" - ], - "license": "MIT", - "dependencies": { - "commander": "^8.3.0" - }, - "bin": { - "katex": "cli.js" - } - }, "node_modules/rehype-raw": { "version": "7.0.0", "resolved": "https://registry.npmjs.org/rehype-raw/-/rehype-raw-7.0.0.tgz", diff --git a/docs/src/content/docs/numerics/index.mdx b/docs/src/content/docs/numerics/index.mdx new file mode 100644 index 0000000..72984be --- /dev/null +++ b/docs/src/content/docs/numerics/index.mdx @@ -0,0 +1,102 @@ +--- +title: Overview +description: The numerics of mpifield, from the shape functions of the deposit and the gather to the halo exchange, and what it leaves out. +sidebar: + label: Overview +--- + +import MethodCard from '@max-models/starlight-theme/components/MethodCard.astro'; + +mpifield holds fields on the nodes of a uniform 1D, 2D or 3D grid, split over MPI ranks with halo +nodes, and moves values between the nodes and the markers of a particle code: a deposit +(scatter) of marker values onto the nodes and a gather (interpolation) of the field at the +markers. No field equation is solved; the Poisson or Maxwell solver, the pusher and the time +integration belong to the codes that use it. + +Each equation on these pages is declared on the code that implements it, as a method card, and +checked by the tests: SymPy rules for the shape functions and the operators, run against the +compiled kernels, and the order of convergence of the two gathers, the gradient and the +prescribed field. The pages of the section list the cards module by module; this one shows how +they fit together. + +## Deposit and gather + +Node $k$ of axis $d$ is at $x_{d,k} = x^{\mathrm{lo}}_d + k h_d$, and a marker at $\mathbf x$ has +the relative coordinates $r_d = (x_d - x^{\mathrm{lo}}_d)/h_d$. A weighting gives node +$\mathbf k$ the weight $W_{\mathbf k}(\mathbf x) = \prod_d w(r_d - k_d)$, a tensor product of +one-dimensional B-splines: + +1. **[Linear weights](/mpifield/numerics/kernels-linear-cell/)** (`LINEAR`, cloud-in-cell), + $w_1$, two nodes per axis: the default. +2. **[Quadratic weights](/mpifield/numerics/kernels-quadratic-cell/)** (`QUADRATIC`, + triangular-shaped cloud), $w_2$, three nodes per axis around the nearest one: a continuous + first derivative and less aliasing noise, for $3^n$ nodes per marker instead of $2^n$. + +Both sum to one over the nodes and reproduce the position as their first moment, so a deposit +keeps the total value and a gather is exact for constant and linear fields; both gather a +smooth field to second order in $h$. At a bounded edge the weights beyond it fold onto the edge +nodes, which keeps the total; a periodic axis wraps. + +The **[deposit and the gather](/mpifield/numerics/weighting/)** are a pair: the gather reads +the nodes and weights the scatter writes, so it is the transpose of the scatter, also at the +edges and the seams. With a symmetric field solver this keeps the momentum: a particle feels no +force of its own. + + + +The deposit is a sum of weights, not a density: divide it by $\prod_d h_d$ for one. + +## Over MPI ranks + +Each rank stores its block of the grid with halo nodes, at least as many as the weighting needs +(one for `LINEAR`, two for `QUADRATIC`). The kernels see the stored block as a grid of its own +([`Field.block_grid`](/mpifield/numerics/field/)), without the halos beyond a bounded edge of +the domain. A cycle of a particle code then runs, on every rank: + +1. select its own markers (from the block's first node up to the next rank's first node); +2. deposit them into the stored block, halo nodes included; +3. `exchange_halos`: add the halos to their owners on the neighbouring ranks; +4. set the field (a solver outside mpifield); +5. `fill_halos`: copy the owners' nodes into the halos; +6. gather the field at its own markers. + +The split deposit and gather equal the serial ones. + +## Operations on the field + +- **[Gradient](/mpifield/numerics/field/)** of a scalar field: central differences, second + order, one-sided at the ends of a bounded axis. +- **[Smoothing](/mpifield/numerics/field/)**: separable binomial, Gaussian or uniform filters, + 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. +- **[Evaluation](/mpifield/numerics/evaluate/)** as xarray: linear interpolation between the + nodes, whatever the weighting. + +## Conventions + +Positions are component-major, `(ndim, N)`, and marker values `(ncomp, N)`; a field is +`num_gridpoints + (ncomp,)`. On a bounded axis the $N$ nodes span the domain, +$h = (x^{\mathrm{hi}} - x^{\mathrm{lo}})/(N - 1)$; on a periodic one they cover $N$ cells, +$h = (x^{\mathrm{hi}} - x^{\mathrm{lo}})/N$, and the last node neighbours the first. The compiled +kernels (Pyccel with OpenMP, CUDA, a NumPy fallback, Metal in single precision) share the same +arithmetic. + +## Not covered + +- No Poisson or Maxwell solver, no particle pusher, no time integration. +- No normalisation of the deposit by the cell volume or by a species weight. +- Only the gradient of a scalar field: no divergence or curl. +- Only linear and quadratic weightings are built in, and no energy-conserving (non-adjoint) + schemes; a weighting of your own is a pair of kernels. +- Markers are not moved between ranks: each rank deposits and gathers its own, and wraps them on + a split periodic axis. +- `electric_field_from_potential` differences one-sided at both ends of every axis, also a + periodic one. + +## Verification + +Besides the checks of the cards, the tests compare the Pyccel, CUDA (emulated) and Metal +kernels with the NumPy ones, the deposit with the transpose of the gather, the OpenMP deposit +with the serial one, and, under `mpiexec`, the split deposit, gather, gradient and smoothing +with those of one rank. diff --git a/pyproject.toml b/pyproject.toml index f015ea9..960e256 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,8 @@ dependencies = [ "cunumpy>=0.7.2", # MPI only under an MPI launcher (mpirun, srun); a serial stand-in for mpi4py otherwise "maybempi>=0.1.2", + # the method cards of the numerical methods (equations, orders, references) + "methodcard", "mpiarray>=0.2.1", "numpy", # the CPU kernels: compiled on first use; their sources import Pyccel's OpenMP module @@ -56,7 +58,7 @@ optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_mach optional-dependencies.mpi = [ "mpi4py" ] # Field.plot draws with plasma-plots optional-dependencies.plot = [ "matplotlib>=3.10", "plasma-plots>=0.2.5" ] -optional-dependencies.test = [ "mpifield[plot]", "pytest", "pytest-cov", "pytest-mpi" ] +optional-dependencies.test = [ "methodcard[symbolic]", "mpifield[plot]", "pytest", "pytest-cov", "pytest-mpi" ] urls.Source = "https://github.com/max-models/mpifield" [project.entry-points."scope_profiler.opt_in_tags"] @@ -97,7 +99,10 @@ reportMissingImports = "error" ini_options.pythonpath = [ "src" ] ini_options.testpaths = [ "src" ] ini_options.addopts = "-ra" -ini_options.markers = [ "mpi: tests run under mpiexec with --with-mpi (pytest-mpi)" ] +ini_options.markers = [ + "convergence: an order-of-accuracy test of a method card", + "mpi: tests run under mpiexec with --with-mpi (pytest-mpi)", +] [tool.coverage] run.omit = [ "*/tests/*" ] @@ -111,6 +116,7 @@ report.show_missing = true # cunumpy = { git = "https://github.com/max-models/cunumpy", branch = "devel" } cunumpy = { git = "https://github.com/max-models/cunumpy", branch = "devel" } maybempi = { git = "https://github.com/max-models/maybempi", branch = "devel" } +methodcard = { git = "https://github.com/max-models/methodcard", branch = "devel" } mpiarray = { git = "https://github.com/max-models/mpiarray", branch = "devel" } plasma-plots = { git = "https://github.com/max-models/plasma-plots", branch = "devel" } scope-profiler = { git = "https://github.com/max-models/scope-profiler", branch = "devel" } diff --git a/src/mpifield/__init__.py b/src/mpifield/__init__.py index 6bfd00d..78a041d 100644 --- a/src/mpifield/__init__.py +++ b/src/mpifield/__init__.py @@ -1,11 +1,11 @@ """Fields on uniform grids split over MPI ranks, with compiled gather and scatter. -:class:`Field` holds a field on the nodes of a uniform 1D, 2D or 3D grid as an -:class:`mpiarray.DistributedArray`; it interpolates the field at marker positions -(:meth:`Field.gather`) and deposits marker values onto the nodes -(:meth:`Field.scatter`) with a :class:`Weighting`: :data:`LINEAR` (the default), -:data:`QUADRATIC` or one of your own. The built-in ones run compiled kernels of -:mod:`mpifield.kernels`, on the CPU (Pyccel, NumPy) or the GPU (CUDA, Metal). +`Field` holds a field on the nodes of a uniform 1D, 2D or 3D grid as an +`mpiarray.DistributedArray`; it interpolates the field at marker positions +(`Field.gather`) and deposits marker values onto the nodes +(`Field.scatter`) with a `Weighting`: `LINEAR` (the default), +`QUADRATIC` or one of your own. The built-in ones run compiled kernels of +`mpifield.kernels`, on the CPU (Pyccel, NumPy) or the GPU (CUDA, Metal). """ from importlib.metadata import PackageNotFoundError, version diff --git a/src/mpifield/_method_checks.py b/src/mpifield/_method_checks.py new file mode 100644 index 0000000..b888cd0 --- /dev/null +++ b/src/mpifield/_method_checks.py @@ -0,0 +1,390 @@ +"""The SymPy rules and convergence problems of mpifield's method cards (for the tests). + +The rules of the shape functions run the compiled gather and scatter kernels of +`LINEAR` and `QUADRATIC` (through `Weighting.gather_block` and +`Weighting.scatter_block`), not the NumPy cell functions that carry the cards: a +gather of the unit field on node ``k`` at a marker returns that node's weight +``W_k``. SymPy is imported only inside the rules, so the package does not need it. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import numpy as np +from methodcard import ConvergenceProblem, Rule + +#: Nodes of the 1D test grids of the shape-function rules (spacing 1, lower bound 0). +NODES = 11 + +#: The node whose weight the shape rules read: far enough from both edges for a +#: quadratic stencil around any position within 3/2 of it. +CENTRE = 5 + + +def _weighting(name: str) -> Any: + """The built-in weighting called ``name`` (``"linear"`` or ``"quadratic"``).""" + from mpifield.weighting import LINEAR, QUADRATIC + + return {"linear": LINEAR, "quadratic": QUADRATIC}[name] + + +def _grid(num_nodes: tuple[int, ...], periodic: bool) -> Any: + """A block grid of spacing 1 from 0 on.""" + from mpifield._types import BlockGrid + + ndim = len(num_nodes) + return BlockGrid.from_geometry( + (0.0,) * ndim, (1.0,) * ndim, num_nodes, (periodic,) * ndim + ) + + +def gather_1d( + name: str, rel: Any, node_values: Any, periodic: bool = False +) -> np.ndarray: + """Gather ``node_values`` (one per node, spacing 1) at the relative positions ``rel``.""" + rel = np.atleast_1d(np.asarray(rel, dtype=float)) + values = np.asarray(node_values, dtype=float) + result = np.zeros((1, rel.size)) + _weighting(name).gather_block( + rel[None, :], values[:, None], result, _grid((values.size,), periodic) + ) + return result[0] + + +def _unit(node: int, num_nodes: int = NODES) -> np.ndarray: + """The field that is 1 on ``node`` and 0 elsewhere.""" + values = np.zeros(num_nodes) + values[node] = 1.0 + return values + + +def _shape(name: str) -> Callable[[Any], Any]: + """The one-dimensional shape function ``w(xi)`` of a weighting, in SymPy.""" + import sympy + + half, three_halves = sympy.Rational(1, 2), sympy.Rational(3, 2) + + def linear(xi: Any) -> Any: + return sympy.Piecewise((1 - sympy.Abs(xi), sympy.Abs(xi) <= 1), (0, True)) + + def quadratic(xi: Any) -> Any: + return sympy.Piecewise( + (sympy.Rational(3, 4) - xi**2, sympy.Abs(xi) <= half), + (half * (three_halves - sympy.Abs(xi)) ** 2, sympy.Abs(xi) <= three_halves), + (0, True), + ) + + return {"linear": linear, "quadratic": quadratic}[name] + + +def shape_rules(name: str) -> dict[str, Rule]: + """The rules of a shape function, checked with the compiled gather kernels. + + - ``shape``: the weight of a node is ``w(r - k)``; + - ``partition``: the weights of all nodes sum to ``sum_k w(r - k)`` (= 1); + - ``first_moment``: ``sum_k W_k k = r``, so linear fields are gathered exactly; + - ``bounded``: the weight of node 0 near and beyond a bounded edge (the fold); + - ``periodic``: the weight of node 0 for markers on both sides of the seam. + """ + import sympy + + w = _shape(name) + xi, r = sympy.symbols("xi r", real=True) + reach = 1.0 if name == "linear" else 1.5 + interior = (reach, NODES - 1 - reach) + if name == "linear": + edge = sympy.Piecewise((1, r < 0), (1 - r, True)) + edge_range = (-1.0, 1.0) + else: + edge = sympy.Piecewise( + (sympy.Rational(7, 8), r < 0), + ( + sympy.Rational(1, 2) * (sympy.Rational(1, 2) - r) ** 2 + + sympy.Rational(3, 4) + - r**2, + True, + ), + ) + edge_range = (-1.0, 0.5) + + def node_weight(xi: float) -> float: + return float(gather_1d(name, CENTRE + xi, _unit(CENTRE))[0]) + + def total_weight(r: float) -> float: + return float(gather_1d(name, r, np.ones(NODES))[0]) + + def mean_node(r: float) -> float: + return float(gather_1d(name, r, np.arange(NODES, dtype=float))[0]) + + def edge_weight(r: float) -> float: + return float(gather_1d(name, r, _unit(0))[0]) + + def seam_weight(xi: float) -> float: + return float(gather_1d(name, xi, _unit(0, 8), periodic=True)[0]) + + return { + "shape": Rule(w(xi), node_weight, {"xi": (-reach - 0.25, reach + 0.25)}), + "partition": Rule( + sum(w(r - k) for k in range(NODES)), total_weight, {"r": interior} + ), + "first_moment": Rule(r, mean_node, {"r": interior}, rtol=1e-12), + "bounded": Rule(edge, edge_weight, {"r": edge_range}), + "periodic": Rule(w(xi), seam_weight, {"xi": (-reach, reach)}), + } + + +def gather_problem(name: str) -> ConvergenceProblem: + """The gather of a smooth periodic function: the error at random points. + + ``f(x) = sin(2 pi x) + cos(4 pi x) / 2`` on ``[0, 1)``, sampled on ``n`` nodes of a + periodic axis and gathered at 2000 fixed random points, for ``h = 1 / n`` from + 1/32 to 1/256; the error is the largest one. + """ + points = np.random.default_rng(0).uniform(0.0, 1.0, 2000) + + def f(x: np.ndarray) -> np.ndarray: + return np.sin(2 * np.pi * x) + 0.5 * np.cos(4 * np.pi * x) + + def error(spacing: float) -> float: + n = round(1.0 / spacing) + nodes = f(np.arange(n) / n) + gathered = gather_1d(name, points * n, nodes, periodic=True) + return float(np.max(np.abs(gathered - f(points)))) + + return ConvergenceProblem(error, steps=[1 / 32, 1 / 64, 1 / 128, 1 / 256]) + + +def scatter_rules() -> dict[str, Rule]: + """The rules of the deposit, with the compiled scatter kernels. + + - ``linear_2d``: one `LINEAR` marker at ``(i + f_x, j + f_y)`` puts + ``q (1 - f_x)(1 - f_y)`` on node ``(i, j)``; + - ``total``: three `QUADRATIC` markers, two of them beyond the bounded edges, + deposit their whole value. + """ + import sympy + + fx, fy, q = sympy.symbols("f_x f_y q", real=True) + q1, q2, q3 = sympy.symbols("q_1 q_2 q_3", real=True) + + def linear_2d(f_x: float, f_y: float, q: float) -> float: + field = np.zeros((6, 7, 1)) + position = np.array([[2.0 + f_x], [3.0 + f_y]]) + _weighting("linear").scatter_block( + position, np.array([[q]]), field, _grid((6, 7), False) + ) + return float(field[2, 3, 0]) + + def total(q_1: float, q_2: float, q_3: float) -> float: + field = np.zeros((NODES, 1)) + positions = np.array([[-0.4, 4.3, NODES - 0.8]]) + _weighting("quadratic").scatter_block( + positions, np.array([[q_1, q_2, q_3]]), field, _grid((NODES,), False) + ) + return float(field.sum()) + + return { + "linear_2d": Rule( + q * (1 - fx) * (1 - fy), + linear_2d, + {"f_x": (0.0, 1.0), "f_y": (0.0, 1.0), "q": (-2.0, 2.0)}, + ), + "total": Rule(q1 + q2 + q3, total, {"q_1": (0.5, 2.0)}, rtol=1e-12), + } + + +def gather_rules() -> dict[str, Rule]: + """The rule of the gather, with the compiled 3D `QUADRATIC` kernel. + + ``quadratic_3d``: a marker at ``(4, 4, 4) + delta`` gathers the unit field on node + ``(4, 4, 4)`` as ``prod_d (3/4 - delta_d^2)``. + """ + import sympy + + dx, dy, dz = sympy.symbols("delta_x delta_y delta_z", real=True) + rule = ( + (sympy.Rational(3, 4) - dx**2) + * (sympy.Rational(3, 4) - dy**2) + * (sympy.Rational(3, 4) - dz**2) + ) + + def quadratic_3d(delta_x: float, delta_y: float, delta_z: float) -> float: + field = np.zeros((9, 9, 9, 1)) + field[4, 4, 4, 0] = 1.0 + position = 4.0 + np.array([[delta_x], [delta_y], [delta_z]]) + result = np.zeros((1, 1)) + _weighting("quadratic").gather_block( + position, field, result, _grid((9, 9, 9), False) + ) + return float(result[0, 0]) + + half = (-0.49, 0.49) + return { + "quadratic_3d": Rule( + rule, + quadratic_3d, + {"delta_x": half, "delta_y": half, "delta_z": half}, + ) + } + + +def gradient_problem() -> ConvergenceProblem: + """`Field.gradient` of ``sin(2 pi x)`` on a periodic axis, the largest error. + + The central difference is ``sin(2 pi h) / h cos(2 pi x)``, so the error is + ``(2 pi)^3 h^2 / 6`` to leading order. + """ + from mpifield.field import Field + + def error(spacing: float) -> float: + n = round(1.0 / spacing) + field = Field( + ndim=1, + lower_bound=(0.0,), + upper_bound=(1.0,), + num_gridpoints=(n,), + periodic=True, + ) + x = np.asarray(field.get_grid_points()[0]) + field.set_field_data(np.sin(2 * np.pi * x)) + derivative = np.asarray(field.gradient().to_numpy())[:, 0] + return float(np.max(np.abs(derivative - 2 * np.pi * np.cos(2 * np.pi * x)))) + + return ConvergenceProblem(error, steps=[1 / 32, 1 / 64, 1 / 128, 1 / 256]) + + +def electric_field_problem() -> ConvergenceProblem: + """`electric_field_from_potential` of ``sin(3 x)`` on ``[0, 1]``, inside ``[1/4, 3/4]``. + + The interior error of the central difference, away from the first-order ends. + """ + from mpifield.prescribed import electric_field_from_potential + + def error(spacing: float) -> float: + n = round(1.0 / spacing) + 1 + field = electric_field_from_potential( + lambda x: np.sin(3.0 * x), (0.0,), (1.0,), n + ) + x = np.linspace(0.0, 1.0, n) + inside = (x >= 0.25 - 1e-12) & (x <= 0.75 + 1e-12) + e = np.asarray(field.to_numpy())[:, 0] + return float(np.max(np.abs(e[inside] + 3.0 * np.cos(3.0 * x[inside])))) + + return ConvergenceProblem(error, steps=[1 / 16, 1 / 32, 1 / 64, 1 / 128]) + + +def axial_potential_rules() -> dict[str, Rule]: + """`axial_potential` of the table ``(0, phi_0), (1, phi_1)`` along ``x``. + + ``table``: linear between the points, the end values outside. + """ + import sympy + + x, phi0, phi1 = sympy.symbols("x phi_0 phi_1", real=True) + rule = sympy.Piecewise( + (phi0, x < 0), (phi0 + (phi1 - phi0) * x, x <= 1), (phi1, True) + ) + + def table(x: float, phi_0: float, phi_1: float) -> float: + from mpifield.prescribed import axial_potential + + potential = axial_potential([0.0, 1.0], [phi_0, phi_1]) + return float(np.asarray(potential(np.array([x])))[0]) + + ranges = {"x": (-0.5, 1.5), "phi_0": (-5.0, 5.0), "phi_1": (-5.0, 5.0)} + return {"table": Rule(rule, table, ranges)} + + +def _field_1d(values: Any, upper: float, periodic: bool) -> Any: + """A one-component 1D `Field` on ``[0, upper]`` holding ``values``.""" + from mpifield.field import Field + + values = np.asarray(values, dtype=float) + field = Field( + ndim=1, + lower_bound=(0.0,), + upper_bound=(upper,), + num_gridpoints=(values.size,), + periodic=periodic, + ) + field.set_field_data(values) + return field + + +def smoothing_rules() -> dict[str, Rule]: + """The rules of `Field.smooth`, one binomial pass on five nodes. + + - ``filter``: node 2 becomes ``(g_1 + 2 g_2 + g_3) / 4``; + - ``boundary``: on a bounded axis node 0 becomes ``(2 g_0 + g_1) / 4`` (zero + outside); on a periodic one node 0 also takes node 4. + """ + import sympy + + g0, g1, g2, g3, g4 = sympy.symbols("g_0 g_1 g_2 g_3 g_4", real=True) + + def inner(g_1: float, g_2: float, g_3: float) -> float: + field = _field_1d([0.3, g_1, g_2, g_3, -0.7], 1.0, periodic=False) + return float(np.asarray(field.smooth("binomial").to_numpy())[2, 0]) + + def edge(g_0: float, g_1: float, g_4: float) -> float: + bounded = _field_1d([g_0, g_1, 0.5, 0.2, g_4], 1.0, periodic=False) + periodic = _field_1d([g_0, g_1, 0.5, 0.2, g_4], 1.0, periodic=True) + smooth = [ + float(np.asarray(f.smooth("binomial").to_numpy())[0, 0]) + for f in (bounded, periodic) + ] + # both edge values in one number: the periodic one weighted by pi + return smooth[0] + np.pi * smooth[1] + + edge_rule = (2 * g0 + g1) / 4 + sympy.Float(np.pi) * (g4 + 2 * g0 + g1) / 4 + return { + "filter": Rule((g1 + 2 * g2 + g3) / 4, inner, rtol=1e-13), + "boundary": Rule(edge_rule, edge, rtol=1e-13), + } + + +def gradient_rules() -> dict[str, Rule]: + """The rules of `Field.gradient` on five nodes of spacing ``h`` (a bounded axis). + + - ``central``: node 2 gets ``(g_3 - g_1) / (2 h)``; + - ``ends``: node 0 gets ``(g_1 - g_0) / h``. + """ + import sympy + + g0, g1, g3, h = sympy.symbols("g_0 g_1 g_3 h", real=True) + + def node(index: int, values: list[float], h: float) -> float: + field = _field_1d(values, 4 * h, periodic=False) + return float(np.asarray(field.gradient().to_numpy())[index, 0]) + + def central(g_1: float, g_3: float, h: float) -> float: + return node(2, [0.4, g_1, 1.3, g_3, 0.2], h) + + def ends(g_0: float, g_1: float, h: float) -> float: + return node(0, [g_0, g_1, 1.3, 0.1, 0.2], h) + + spacing = {"h": (0.1, 2.0)} + return { + "central": Rule((g3 - g1) / (2 * h), central, spacing), + "ends": Rule((g1 - g0) / h, ends, spacing), + } + + +def evaluation_rules() -> dict[str, Rule]: + """The rule of `evaluate_field` across a periodic seam. + + ``seam``: four nodes on ``[0, 1)``; between the last node (``x = 3/4``) and the + upper bound the field is linear from ``g_3`` to ``g_0``. + """ + import sympy + + x, g0, g3 = sympy.symbols("x g_0 g_3", real=True) + + def seam(x: float, g_0: float, g_3: float) -> float: + field = _field_1d([g_0, 1.0, 2.0, g_3], 1.0, periodic=True) + return float(field.evaluate(x=x)) + + rule = g3 + (g0 - g3) * (x - sympy.Rational(3, 4)) * 4 + return {"seam": Rule(rule, seam, {"x": (0.75, 1.0)})} diff --git a/src/mpifield/_profiling.py b/src/mpifield/_profiling.py index 9bef54e..a7ea712 100644 --- a/src/mpifield/_profiling.py +++ b/src/mpifield/_profiling.py @@ -1,6 +1,6 @@ """The scope-profiler tag of mpifield's profiling regions. -The regions are opt-in: :mod:`mpifield` registers the tag on import, so they +The regions are opt-in: `mpifield` registers the tag on import, so they record nothing unless the application enables it, e.g. with ``scope_profiler.setup(enable_tags=["mpifield"])``. """ diff --git a/src/mpifield/_types.py b/src/mpifield/_types.py index 44515e8..ee3ade3 100644 --- a/src/mpifield/_types.py +++ b/src/mpifield/_types.py @@ -25,12 +25,12 @@ class Bounds(NamedTuple): class BlockGrid(NamedTuple): - """The grid of one rank's stored block, as the kernels see it (:meth:`Field.block_grid`). + """The grid of one rank's stored block, as the kernels see it (`Field.block_grid`). ``geometry`` is the same grid as the kernels' trailing arguments, Python scalars: ``(*lower, *cell_size, *num_nodes, *periodic)``, one entry per axis each (``periodic`` as 0 or 1). ``index`` selects the nodes of the grid in - :attr:`Field.local_with_halos`. + `Field.local_with_halos`. """ bounds: Bounds @@ -51,8 +51,8 @@ def from_geometry( """Return the grid of a whole array of ``num_nodes`` nodes from ``lower`` on. For running a weighting on an array that is not a field's stored block - (:meth:`Weighting.scatter_block`). The arguments may be sequences or - arrays on either side (a :meth:`Field.block_grid` on the GPU holds CuPy + (`Weighting.scatter_block`). The arguments may be sequences or + arrays on either side (a `Field.block_grid` on the GPU holds CuPy arrays). ``periodic`` defaults to bounded axes. """ lower = tuple(float(v) for v in lower) diff --git a/src/mpifield/evaluate.py b/src/mpifield/evaluate.py index c5d9111..960438f 100644 --- a/src/mpifield/evaluate.py +++ b/src/mpifield/evaluate.py @@ -6,6 +6,9 @@ import cunumpy as xp import numpy as np +from methodcard import method + +from mpifield import _method_checks if TYPE_CHECKING: import xarray as xr @@ -50,6 +53,18 @@ def _check_coordinate( return float(values) if np.ndim(values) == 0 else values +@method( + equations={ + r"linear": r"G(x) = (1 - t)\, G_k + t\, G_{k+1}, \qquad t = \frac{x - x_k}{h}, \quad x_k \le x \le x_{k+1}, \qquad \text{multilinear in several axes}", + r"seam": r"G(x^{\mathrm{hi}}) = G_0 \text{ on a periodic axis}", + r"bounds": r"x \leftarrow \mathrm{clamp}(x, x^{\mathrm{lo}}, x^{\mathrm{hi}}) \text{ if } x^{\mathrm{lo}} - \epsilon \le x \le x^{\mathrm{hi}} + \epsilon, \qquad \epsilon = 10^{-12} \max\left(|x^{\mathrm{lo}}|, |x^{\mathrm{hi}}|, x^{\mathrm{hi}} - x^{\mathrm{lo}}\right)", + }, + assumptions=( + "xarray's linear `interp` on the gathered field, whatever weighting the field is used with", + "coordinates farther outside the domain, or not finite, raise `ValueError`", + ), + symbolic=_method_checks.evaluation_rules, +) def evaluate_field( field: Field, coordinates: dict[str, Any], diff --git a/src/mpifield/field.py b/src/mpifield/field.py index c588ad5..62c4f63 100644 --- a/src/mpifield/field.py +++ b/src/mpifield/field.py @@ -1,19 +1,19 @@ """A field on the nodes of a uniform grid, split over MPI ranks, with gather and scatter. -:class:`Field` is a :class:`simgeometry.grid.Grid` holding an -:class:`mpiarray.DistributedArray` of shape ``num_gridpoints + (num_components,)``: +`Field` is a `simgeometry.grid.Grid` holding an +`mpiarray.DistributedArray` of shape ``num_gridpoints + (num_components,)``: the grid axes are split over the ranks' process grid with ``num_ghost_points`` halo nodes, the component axis is never split. On it: -- :meth:`Field.gather` interpolates the field at marker positions and - :meth:`Field.scatter` deposits marker values onto the nodes, both with a - :class:`~mpifield.weighting.Weighting` (linear cloud-in-cell by default, or +- `Field.gather` interpolates the field at marker positions and + `Field.scatter` deposits marker values onto the nodes, both with a + `Weighting` (linear cloud-in-cell by default, or quadratic, or one of your own) through compiled kernels (Pyccel on the CPU, CUDA for CuPy arrays, a NumPy fallback); - elementwise arithmetic and NumPy ufuncs return new fields; ``sum``, ``min``, ... reduce over all ranks; -- :meth:`Field.evaluate` returns the field, or a line or plane of it, as an - :class:`xarray.DataArray`, and :meth:`Field.plot` draws it with plasma-plots. +- `Field.evaluate` returns the field, or a line or plane of it, as an + `xarray.DataArray`, and `Field.plot` draws it with plasma-plots. Positions are component-major, ``(ndim, N)``, and per-marker values ``(ncomp, N)``, as the kernels take them. @@ -30,9 +30,11 @@ import numpy as np import scope_profiler as sp from maybempi import MPI +from methodcard import method from mpiarray import DistributedArray from simgeometry.grid import Grid +from mpifield import _method_checks from mpifield._profiling import PROFILE_TAGS from mpifield._types import Array, BlockGrid, Bounds, IndexLike, PathLikeStr from mpifield.weighting import LINEAR, Weighting @@ -65,9 +67,9 @@ class Field(Grid): ``lower_bound`` to ``upper_bound`` on a bounded axis; on a periodic axis the last node neighbours the first, so there are ``num_gridpoints`` cells over the length. ``comm`` and ``domain_decomposition`` - split the grid axes over the ranks as :class:`simgeometry.grid.Grid` does; each + split the grid axes over the ranks as `simgeometry.grid.Grid` does; each rank stores its block with ``num_ghost_points`` halo nodes on every side - (:attr:`local_with_halos`). Without ``comm`` every rank holds the whole field. + (`local_with_halos`). Without ``comm`` every rank holds the whole field. Examples: -------- @@ -148,7 +150,7 @@ def __init__( def _make_block_grid( self, layout: mpa.Layout, cpu: bool, inside_domain: bool ) -> BlockGrid: - """Compute :meth:`block_grid`.""" + """Compute `block_grid`.""" ndim = self.ndim halo = layout.halo[:ndim] flags = _periodic_flags(self.periodic, ndim) @@ -199,23 +201,37 @@ def _make_block_grid( ), ) + @method( + equations={ + r"nodes": r"x_{d,k} = x^{\mathrm{lo}}_d + k\, h_d, \qquad h_d = \frac{x^{\mathrm{hi}}_d - x^{\mathrm{lo}}_d}{N_d - 1} \text{ (bounded)}, \quad h_d = \frac{x^{\mathrm{hi}}_d - x^{\mathrm{lo}}_d}{N_d} \text{ (periodic)}", + r"lower": r"x^{\mathrm{lo}}_{\mathrm{block},d} = x^{\mathrm{lo}}_d + (s_d - n_h + \ell_d)\, h_d, \qquad \ell_d = n_h \text{ at a bounded lower domain edge}, \ 0 \text{ otherwise}", + r"count": r"N_{\mathrm{block},d} = n_d + 2 n_h - n_h\,[\text{bounded lower edge}] - n_h\,[\text{bounded upper edge}]", + r"periodic": r"\mathrm{periodic}_{\mathrm{block},d} = \mathrm{periodic}_d \wedge (n_h = 0)", + }, + assumptions=( + "$s_d$ the first owned global node, $n_d$ the owned nodes, $n_h$ the halo width; the edge terms only with `inside_domain` (the grid of `scatter` and `gather`)", + r"on a periodic axis the last node is one spacing before $x^{\mathrm{hi}}_d$ and neighbours node 0 (simgeometry's `Grid`)", + r"with halo nodes the kernels see a bounded block, and the halos hold the wrapped-around nodes: on a split periodic axis the markers must already lie in $[x^{\mathrm{lo}}, x^{\mathrm{hi}})$", + "a split field needs at least `weighting.halo` halo nodes", + ), + ) def block_grid(self, cpu: bool = False, inside_domain: bool = False) -> BlockGrid: """The grid of this rank's stored block (halo nodes included), for the kernels. Under a spatial decomposition the kernels deposit into and interpolate - from :attr:`local_with_halos`, a block of the global grid that starts + from `local_with_halos`, a block of the global grid that starts ``halo`` nodes before the block's first node. They see it as a grid of its own: its lower bound is that node, its node count the stored extent. A periodic axis is periodic for the kernels only without halo nodes; with - them the halo nodes hold the wrapped-around nodes (:meth:`fill_halos`, - :meth:`exchange_halos`). ``cpu`` returns host arrays, for the NumPy and + them the halo nodes hold the wrapped-around nodes (`fill_halos`, + `exchange_halos`). ``cpu`` returns host arrays, for the NumPy and Pyccel kernels. Cached per layout. With ``inside_domain`` the grid leaves out the halo nodes beyond a bounded edge of the domain, which no exchange fills or collects: the kernels then clamp there as on an unsplit field, and a stencil wider than a cell neither deposits into nor reads from nodes outside the domain. This is the - grid of :meth:`scatter` and :meth:`gather`. ``local_with_halos[grid.index]`` + grid of `scatter` and `gather`. ``local_with_halos[grid.index]`` is the part of the stored block it describes; without ``inside_domain`` that is all of it. """ @@ -236,7 +252,7 @@ def on_grid(cls, grid: Grid, num_components: int = 1, split: bool = True) -> Sel 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 (:meth:`block_grid`); otherwise, or with + interpolation reads from (`block_grid`); otherwise, or with ``split=False``, every rank holds the whole grid. """ decomposition = cls._decomposition_of(grid) @@ -360,10 +376,10 @@ def set_field_data( - a global array of the field's shape ``num_gridpoints + (ncomp,)``, or of ``num_gridpoints`` for one ``component`` (or a scalar field): each rank copies its own nodes and the halos are zeroed; - - a :class:`~mpiarray.DistributedArray` with the same split (the field's + - a `DistributedArray` with the same split (the field's shape, or the grid's for one ``component``) is copied block by block, halo nodes included, without communication; - - an array of the stored block's shape (:attr:`local_with_halos`) is copied as is. + - an array of the stored block's shape (`local_with_halos`) is copied as is. """ if component is not None and not 0 <= component < self.num_components: raise IndexError( @@ -533,7 +549,7 @@ def local_with_halos(self) -> Array: """This rank's stored block, halo nodes included: a writable view. This is the array the kernels read and write, on the grid of - :meth:`block_grid`. + `block_grid`. """ return self.field_data.local_with_halos @@ -548,14 +564,34 @@ def component(self, component: int) -> DistributedArray: # Halos -------------------------------------------------------------------# + @method( + equations={ + r"update": r"G^{\mathrm{halo,lo}} \leftarrow G^{\mathrm{interior,hi}}_{\text{left}}, \qquad G^{\mathrm{halo,hi}} \leftarrow G^{\mathrm{interior,lo}}_{\text{right}} \qquad \text{axis by axis}", + }, + assumptions=( + "the halo and interior layers along an axis span the whole storage along the others, halos included, so the corner halos are filled too", + "one block along a periodic axis: copied locally; at a wall the halos are left unchanged", + "collective; before a gather near the edge of a block", + ), + ) def fill_halos(self) -> None: """Fill the halo nodes from the neighbouring ranks (``update_halos``).""" self.field_data.update_halos() + @method( + equations={ + r"accumulate": r"\rho^{\mathrm{interior,lo}} \mathrel{+}= \rho^{\mathrm{halo,hi}}_{\text{left}}, \qquad \rho^{\mathrm{interior,hi}} \mathrel{+}= \rho^{\mathrm{halo,lo}}_{\text{right}}, \qquad \rho^{\mathrm{halo}} \leftarrow 0 \qquad \text{axis by axis}", + }, + assumptions=( + "a value in a corner halo reaches the diagonal neighbour in two steps, so the result equals the serial deposit", + "one block along a periodic axis: added locally; at a wall the halo values are dropped (with `inside_domain` nothing was deposited there)", + "collective; after a deposit", + ), + ) def exchange_halos(self) -> None: """Add the halo nodes to their owners on the neighbouring ranks and clear them. - Call it after a deposit (:meth:`scatter`) into a distributed field. + Call it after a deposit (`scatter`) into a distributed field. """ self.field_data.accumulate_halos() @@ -646,6 +682,18 @@ def _convolve_axis( filtered += weight * padded[tuple(slices)] return filtered + @method( + equations={ + r"filter": r"G^{\mathrm{new}}_k = \sum_{m=-R}^{R} a_m\, G_{k+m}, \qquad a^{\mathrm{binomial}} = \tfrac14 [1, 2, 1], \quad a^{\mathrm{gaussian}} = \tfrac1{16} [1, 4, 6, 4, 1], \quad a^{\mathrm{uniform}} = \tfrac13 [1, 1, 1]", + r"boundary": r"G_k = 0 \text{ outside a bounded axis}, \qquad G_{k+N} = G_k \text{ on a periodic axis}", + }, + assumptions=( + "separable: the one-dimensional filter along each axis in turn, `passes` times", + "the sum is kept on a periodic axis, not near a bounded edge", + "under MPI with halos of width $R$ when every block is at least $R$ nodes wide (the same result as on the gathered field), gathered otherwise", + ), + symbolic=_method_checks.smoothing_rules, + ) @sp.profile("mpifield.Field.smooth", tags=PROFILE_TAGS) def smooth( self, @@ -776,13 +824,27 @@ def smooth_uniform( # Derivatives -------------------------------------------------------------# + @method( + equations={ + r"central": r"(\partial_d G)_k = \frac{G_{k+1} - G_{k-1}}{2 h_d}, \qquad G_{-1} = G_{N-1},\ G_N = G_0 \text{ on a periodic axis}", + r"ends": r"(\partial_d G)_0 = \frac{G_1 - G_0}{h_d}, \qquad (\partial_d G)_{N-1} = \frac{G_{N-1} - G_{N-2}}{h_d} \quad \text{(bounded axis)}", + }, + order=2, + assumptions=( + "a scalar field; second order inside and across a periodic seam, first order at the ends of a bounded axis (`numpy.gradient`, `edge_order=1`)", + "the order is checked on a periodic axis", + "a distributed field is gathered first, so the result does not depend on the split", + ), + symbolic=_method_checks.gradient_rules, + problem=_method_checks.gradient_problem, + ) @sp.profile("mpifield.Field.gradient", tags=PROFILE_TAGS) def gradient(self) -> Self: """Return the gradient of a scalar field: a field with ``ndim`` components. Second-order central differences inside the domain and across a periodic seam, first-order one-sided differences at the ends of a bounded axis (as - :func:`numpy.gradient`). A distributed field is gathered first (collective), + `numpy.gradient`). A distributed field is gathered first (collective), so the result does not depend on the split. """ if self.num_components != 1: @@ -874,7 +936,7 @@ def plot( """Plot the field with plasma-plots and return its ``PlotResult``. ``x``, ``y``, ``z``, ``component``, ``name`` and ``units`` are as in - :meth:`evaluate`. A one-dimensional result is drawn as a line, a + `evaluate`. A one-dimensional result is drawn as a line, a two-dimensional one as a colour map. A 3D field with no axis fixed is shown on its ``z`` midplane. Other keyword arguments (``cmap``, ``vmin``, ``title``, ``backend``, ...) go to ``plasma_plots``. Needs the ``plot`` @@ -986,7 +1048,7 @@ def _check_halo(self, weighting: Weighting) -> None: ) def _host_grid(self) -> BlockGrid: - """:meth:`block_grid` on the CPU inside the domain, cached for the current layout. + """`block_grid` on the CPU inside the domain, cached for the current layout. Checked by identity, without hashing the layout: the deposit of every step takes this path. @@ -1007,7 +1069,7 @@ def gather(self, positions: Array | float, weighting: Weighting = LINEAR) -> Arr built-in weightings a position outside the grid uses its edge nodes; on a periodic axis without halos it wraps around. Under MPI each rank interpolates from its own block, so its markers must lie in it, and the - halo nodes must be filled (:meth:`fill_halos`) for markers near its edge. + halo nodes must be filled (`fill_halos`) for markers near its edge. """ self._check_halo(weighting) positions = self._marker_positions(positions) @@ -1041,7 +1103,7 @@ def scatter( ``positions`` has shape ``(ndim, N)`` and ``values`` ``(ncomp, N)``; a 1-D ``values`` of length ``N`` is one component, a number is the same value for every marker. The values are added to the stored block, halo nodes - included; under MPI call :meth:`exchange_halos` afterwards. A marker with + included; under MPI call `exchange_halos` afterwards. A marker with a zero value contributes nothing, whatever its position. """ self._check_halo(weighting) @@ -1260,7 +1322,7 @@ def load( num_ghost_points: int = 0, domain_decomposition: Any = None, ) -> Self: - """Read a field written by :meth:`export`, split as the arguments say.""" + """Read a field written by `export`, split as the arguments say.""" import xarray as xr with xr.open_dataarray(file_path) as array: diff --git a/src/mpifield/kernels/__init__.py b/src/mpifield/kernels/__init__.py index eb0122d..67d067b 100644 --- a/src/mpifield/kernels/__init__.py +++ b/src/mpifield/kernels/__init__.py @@ -1,10 +1,10 @@ -"""mpifield's compiled kernels, one folder per kernel (see :mod:`mpifield.kernels.catalog`). +"""mpifield's compiled kernels, one folder per kernel (see `mpifield.kernels.catalog`). -Every kernel folder builds its kernel with :data:`KERNEL_PACKAGE`:: +Every kernel folder builds its kernel with `KERNEL_PACKAGE`:: kernel = KERNEL_PACKAGE.from_folder(__name__) -with the defaults of :class:`cunumpy.kernels.KernelPackage`: the Pyccel version is +with the defaults of `cunumpy.kernels.KernelPackage`: the Pyccel version is compiled on its first host call and kept under ``MPIFIELD_KERNEL_CACHE`` (default ``~/.cache/mpifield/kernels``; ``MPIFIELD_OPENMP=0`` builds serial), the CUDA version runs for CuPy arguments and the host version for NumPy ones, one CUDA thread per entry of @@ -14,7 +14,7 @@ ``KERNEL_PACKAGE.metal_kernel`` in a function with the arguments of the Pyccel version. The Apple GPU has no float64, so they compute in float32: expect a relative error of about 1e-5 against the CPU versions. Import one only where -:func:`cunumpy.kernels.metal_available` is true. +`cunumpy.kernels.metal_available` is true. """ import cunumpy as xp diff --git a/src/mpifield/kernels/_metal_grid.py b/src/mpifield/kernels/_metal_grid.py index ee8a93f..7285257 100644 --- a/src/mpifield/kernels/_metal_grid.py +++ b/src/mpifield/kernels/_metal_grid.py @@ -2,11 +2,11 @@ Positions ``(ndim, N)`` are turned into cell coordinates ``(x - lower) / cell_size`` on the host in float64; the kernels compute the cell and the weights in float32 with -:data:`CELL_HEADER`, the arithmetic of ``linear_cell.cuh`` (:data:`QUADRATIC_HEADER` and +`CELL_HEADER`, the arithmetic of ``linear_cell.cuh`` (`QUADRATIC_HEADER` and ``quadratic_cell.cuh`` for the quadratic ones). A scatter accumulates into a zeroed float32 grid with float atomics (no 64-bit atomics on the Apple GPU), which the host adds to the caller's grid in float64: the summation order, hence the -last bits, differ from the CPU. See :mod:`mpifield.kernels`. +last bits, differ from the CPU. See `mpifield.kernels`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/catalog.py b/src/mpifield/kernels/catalog.py index 180d011..7300014 100644 --- a/src/mpifield/kernels/catalog.py +++ b/src/mpifield/kernels/catalog.py @@ -16,14 +16,14 @@ runs the CUDA version when the arrays live on the GPU (CuPy) and the compiled Pyccel version otherwise (``dispatch="arrays"``). The kernel object is a -:class:`cunumpy.kernels.Kernel`, built in the folder's ``__init__.py`` by -``KERNEL_PACKAGE.from_folder`` (see :mod:`mpifield.kernels`). Both versions are +`cunumpy.kernels.Kernel`, built in the folder's ``__init__.py`` by +``KERNEL_PACKAGE.from_folder`` (see `mpifield.kernels`). Both versions are compiled on first use: importing this module needs neither Pyccel nor CuPy. CUDA sources include mpifield's header relative to the source root: ``#include "mpifield/kernels/linear_cell.cuh"``. The kernels are the linear (``*_linear_*``) and quadratic (``*_quadratic_*``) deposits -and gathers of :mod:`mpifield.weighting`. The quadratic ones, and the 1D and 2D linear +and gathers of `mpifield.weighting`. The quadratic ones, and the 1D and 2D linear gathers and the 2D linear deposit, take the arguments of the weighting kernel convention; the linear 1D deposit takes one component (``(N,)`` markers, an ``(nx,)`` grid) and the linear 3D kernels the grid as rows ``(nx * ny * nz, ncomp)``. diff --git a/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py b/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py index 7fc2b69..fe67a03 100644 --- a/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py +++ b/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``gather_linear_1d`` kernel. ``gather_linear_1d_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` ``(1, N)``, ``result`` ``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +arguments. One thread per marker (``positions`` ``(1, N)``, ``result`` ``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_pyccel.py b/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_pyccel.py index 1b1ed5e..f056aff 100644 --- a/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_pyccel.py +++ b/src/mpifield/kernels/gather_linear_1d/gather_linear_1d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``gather_linear_1d`` kernel; ``gather_linear_1d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. """ diff --git a/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py b/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py index aa959bf..c3cd5b6 100644 --- a/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py +++ b/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``gather_linear_2d`` kernel. ``gather_linear_2d_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` ``(2, N)``, ``result`` ``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +arguments. One thread per marker (``positions`` ``(2, N)``, ``result`` ``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_pyccel.py b/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_pyccel.py index 14bb13f..f99c524 100644 --- a/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_pyccel.py +++ b/src/mpifield/kernels/gather_linear_2d/gather_linear_2d_pyccel.py @@ -7,7 +7,7 @@ The Pyccel version (CPU) of the ``gather_linear_2d`` kernel; ``gather_linear_2d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. """ diff --git a/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py b/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py index 50470c8..d737c32 100644 --- a/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py +++ b/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``gather_linear_3d`` kernel. ``gather_linear_3d_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` ``(3, N)``, ``result`` ``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +arguments. One thread per marker (``positions`` ``(3, N)``, ``result`` ``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_pyccel.py b/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_pyccel.py index 9a69ebf..081399a 100644 --- a/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_pyccel.py +++ b/src/mpifield/kernels/gather_linear_3d/gather_linear_3d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``gather_linear_3d`` kernel; ``gather_linear_3d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. """ diff --git a/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py b/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py index 8da1ee4..9f0447f 100644 --- a/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py +++ b/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_metal.py @@ -2,7 +2,7 @@ ``gather_quadratic_1d_pyccel.py`` next to this file is the CPU version with the same arguments. One thread per marker (``positions`` ``(1, N)``, ``result`` -``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_pyccel.py b/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_pyccel.py index 729c90b..e1ff239 100644 --- a/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_pyccel.py +++ b/src/mpifield/kernels/gather_quadratic_1d/gather_quadratic_1d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``gather_quadratic_1d`` kernel; ``gather_quadratic_1d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. The transpose of ``scatter_quadratic_1d``: each marker reads its nearest node and the two around it along each axis (``quadratic_cell.cuh``). With OpenMP the diff --git a/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py b/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py index c59c834..0ab84c4 100644 --- a/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py +++ b/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_metal.py @@ -2,7 +2,7 @@ ``gather_quadratic_2d_pyccel.py`` next to this file is the CPU version with the same arguments. One thread per marker (``positions`` ``(2, N)``, ``result`` -``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_pyccel.py b/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_pyccel.py index 07a0a2c..10c0334 100644 --- a/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_pyccel.py +++ b/src/mpifield/kernels/gather_quadratic_2d/gather_quadratic_2d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``gather_quadratic_2d`` kernel; ``gather_quadratic_2d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. The transpose of ``scatter_quadratic_2d``: each marker reads its nearest node and the two around it along each axis (``quadratic_cell.cuh``). With OpenMP the diff --git a/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py b/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py index bd67f3c..4cb5271 100644 --- a/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py +++ b/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_metal.py @@ -2,7 +2,7 @@ ``gather_quadratic_3d_pyccel.py`` next to this file is the CPU version with the same arguments. One thread per marker (``positions`` ``(3, N)``, ``result`` -``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_pyccel.py b/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_pyccel.py index efaecbc..4cc6415 100644 --- a/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_pyccel.py +++ b/src/mpifield/kernels/gather_quadratic_3d/gather_quadratic_3d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``gather_quadratic_3d`` kernel; ``gather_quadratic_3d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. The transpose of ``scatter_quadratic_3d``: each marker reads its nearest node and the two around it along each axis (``quadratic_cell.cuh``). With OpenMP the diff --git a/src/mpifield/kernels/linear_cell.py b/src/mpifield/kernels/linear_cell.py index 29b959d..a08c4f8 100644 --- a/src/mpifield/kernels/linear_cell.py +++ b/src/mpifield/kernels/linear_cell.py @@ -6,9 +6,37 @@ from __future__ import annotations +from functools import partial + import numpy as np +from methodcard import method + +from mpifield import _method_checks +@method( + equations={ + r"shape": r"w_1(\xi) = \begin{cases} 1 - |\xi| & |\xi| \le 1 \\ 0 & \text{otherwise} \end{cases}, \qquad W_{\mathbf k}(\mathbf x) = \prod_d w_1(r_d - k_d), \qquad r_d = \frac{x_d - x^{\mathrm{lo}}_d}{h_d}", + r"weights": r"k = \lfloor r \rfloor, \quad f = r - k \in [0, 1], \qquad W_k = 1 - f, \quad W_{k+1} = f", + r"partition": r"\sum_k W_k = 1", + r"first_moment": r"\sum_k W_k\, k = r", + r"bounded": r"k = \mathrm{clamp}(\lfloor r \rfloor, 0, N - 2), \qquad f = \mathrm{clamp}(r - k, 0, 1)", + r"periodic": r"r \leftarrow r \bmod N, \qquad k = \lfloor r \rfloor, \qquad k + 1 \to (k + 1) \bmod N", + }, + order=2, + assumptions=( + "a uniform grid; the weights of several axes are the products of the one-dimensional ones (two nodes per axis, $2^n$ in $n$ dimensions)", + "the same arithmetic as the Pyccel kernels (inline) and the CUDA ones (`linear_cell.cuh`); the rules are checked with the compiled gathers", + "beyond a bounded edge all the weight is on the edge node: a deposit loses no charge, a gather returns the edge value", + "the order is that of the gather of a smooth field in the node spacing $h$", + ), + references=( + "[Birdsall and Langdon, Plasma Physics via Computer Simulation (1991), ch. 8](https://doi.org/10.1201/9781315275048)", + "[Hockney and Eastwood, Computer Simulation Using Particles (1988), ch. 5](https://www.routledge.com/products/9781138413375)", + ), + symbolic=partial(_method_checks.shape_rules, "linear"), + problem=partial(_method_checks.gather_problem, "linear"), +) def linear_cell( rel: np.ndarray, num_nodes: int, periodic: bool ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: diff --git a/src/mpifield/kernels/quadratic_cell.py b/src/mpifield/kernels/quadratic_cell.py index d723902..de431e2 100644 --- a/src/mpifield/kernels/quadratic_cell.py +++ b/src/mpifield/kernels/quadratic_cell.py @@ -7,9 +7,38 @@ from __future__ import annotations +from functools import partial + import numpy as np +from methodcard import method + +from mpifield import _method_checks +@method( + equations={ + r"shape": r"w_2(\xi) = \begin{cases} \tfrac34 - \xi^2 & |\xi| \le \tfrac12 \\ \tfrac12 \left(\tfrac32 - |\xi|\right)^2 & \tfrac12 < |\xi| \le \tfrac32 \\ 0 & \text{otherwise} \end{cases}, \qquad W_{\mathbf k}(\mathbf x) = \prod_d w_2(r_d - k_d)", + r"weights": r"c = \lfloor r + \tfrac12 \rfloor, \quad \delta = r - c \in [-\tfrac12, \tfrac12], \qquad W_{c-1} = \tfrac12 \left(\tfrac12 - \delta\right)^2, \quad W_c = \tfrac34 - \delta^2, \quad W_{c+1} = \tfrac12 \left(\tfrac12 + \delta\right)^2", + r"partition": r"\sum_k W_k = 1", + r"first_moment": r"\sum_k W_k\, k = r", + r"bounded": r"r \leftarrow \mathrm{clamp}(r, 0, N - 1), \qquad \text{nodes } \big(\max(c - 1, 0),\ c,\ \min(c + 1, N - 1)\big), \qquad W_0 = \tfrac12 \left(\tfrac12 - r\right)^2 + \tfrac34 - r^2 \ \ (0 \le r < \tfrac12)", + r"periodic": r"r \leftarrow r \bmod N, \qquad \text{nodes } (c - 1, c, c + 1) \bmod N", + }, + order=2, + assumptions=( + "a uniform grid; the weights of several axes are the products of the one-dimensional ones (three nodes per axis, $3^n$ in $n$ dimensions)", + "the same arithmetic as the Pyccel kernels (`_quadratic_cell`) and the CUDA ones (`quadratic_cell.cuh`); the rules are checked with the compiled gathers", + "on a node the weights are $(1/8, 3/4, 1/8)$; on and beyond a bounded edge $(7/8, 1/8)$ on the two edge nodes, so a deposit loses no charge and a gather is constant beyond the edge", + "within half a cell of a bounded edge node the fold breaks the first moment: a linear field is not gathered exactly there", + "the order is that of the gather of a smooth field in the node spacing $h$ (the second moment of $w_2$ is not zero)", + ), + references=( + "[Hockney and Eastwood, Computer Simulation Using Particles (1988), ch. 5](https://www.routledge.com/products/9781138413375)", + "[Birdsall and Langdon, Plasma Physics via Computer Simulation (1991), ch. 8](https://doi.org/10.1201/9781315275048)", + ), + symbolic=partial(_method_checks.shape_rules, "quadratic"), + problem=partial(_method_checks.gather_problem, "quadratic"), +) def quadratic_cell( rel: np.ndarray, num_nodes: int, periodic: bool ) -> tuple[tuple[np.ndarray, ...], tuple[np.ndarray, ...]]: diff --git a/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_metal.py b/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_metal.py index 93597b0..5b34285 100644 --- a/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_metal.py +++ b/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``scatter_linear_1d`` kernel. ``scatter_linear_1d_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker; see :mod:`mpifield.kernels._metal_grid`. +arguments. One thread per marker; see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_pyccel.py b/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_pyccel.py index cadf7e5..284b2e5 100644 --- a/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_pyccel.py +++ b/src/mpifield/kernels/scatter_linear_1d/scatter_linear_1d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``scatter_linear_1d`` kernel; ``scatter_linear_1d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. With OpenMP the markers are split over the threads, each depositing into its own copy of the grid, and the copies are summed into ``field``: no atomic updates, and diff --git a/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_metal.py b/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_metal.py index 3409cc9..ee0ca88 100644 --- a/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_metal.py +++ b/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``scatter_linear_2d`` kernel. ``scatter_linear_2d_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` ``(2, N)``, ``values`` ``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +arguments. One thread per marker (``positions`` ``(2, N)``, ``values`` ``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_pyccel.py b/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_pyccel.py index 68297a3..2cf1984 100644 --- a/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_pyccel.py +++ b/src/mpifield/kernels/scatter_linear_2d/scatter_linear_2d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``scatter_linear_2d`` kernel; ``scatter_linear_2d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. With OpenMP the markers are split over the threads in contiguous chunks, each depositing into its own buffer, and the buffers are summed into ``field``: no diff --git a/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_metal.py b/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_metal.py index e1990b4..daa0ab7 100644 --- a/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_metal.py +++ b/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_metal.py @@ -1,7 +1,7 @@ """The Metal version (Apple-silicon GPU, float32) of the ``scatter_linear_3d`` kernel. ``scatter_linear_3d_pyccel.py`` next to this file is the CPU version with the same -arguments. One thread per marker (``positions`` ``(3, N)``, ``values`` ``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +arguments. One thread per marker (``positions`` ``(3, N)``, ``values`` ``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_pyccel.py b/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_pyccel.py index e7c3933..1158180 100644 --- a/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_pyccel.py +++ b/src/mpifield/kernels/scatter_linear_3d/scatter_linear_3d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``scatter_linear_3d`` kernel; ``scatter_linear_3d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. With OpenMP the markers are split over the threads in contiguous chunks, each depositing into its own buffer, and the buffers are summed into ``field``: no diff --git a/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_metal.py b/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_metal.py index ca27ac1..cb648a9 100644 --- a/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_metal.py +++ b/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_metal.py @@ -2,7 +2,7 @@ ``scatter_quadratic_1d_pyccel.py`` next to this file is the CPU version with the same arguments. One thread per marker (``positions`` ``(1, N)``, ``values`` -``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_pyccel.py b/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_pyccel.py index 6d471d0..b8ffc93 100644 --- a/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_pyccel.py +++ b/src/mpifield/kernels/scatter_quadratic_1d/scatter_quadratic_1d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``scatter_quadratic_1d`` kernel; ``scatter_quadratic_1d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. Each marker deposits onto its nearest node and the two around it along each axis (``quadratic_cell.cuh``). With OpenMP the markers are split over the threads in diff --git a/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_metal.py b/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_metal.py index 3dc903d..353569c 100644 --- a/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_metal.py +++ b/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_metal.py @@ -2,7 +2,7 @@ ``scatter_quadratic_2d_pyccel.py`` next to this file is the CPU version with the same arguments. One thread per marker (``positions`` ``(2, N)``, ``values`` -``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_pyccel.py b/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_pyccel.py index ab31f56..14b083b 100644 --- a/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_pyccel.py +++ b/src/mpifield/kernels/scatter_quadratic_2d/scatter_quadratic_2d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``scatter_quadratic_2d`` kernel; ``scatter_quadratic_2d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. Each marker deposits onto its nearest node and the two around it along each axis (``quadratic_cell.cuh``). With OpenMP the markers are split over the threads in diff --git a/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_metal.py b/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_metal.py index f244863..35dd160 100644 --- a/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_metal.py +++ b/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_metal.py @@ -2,7 +2,7 @@ ``scatter_quadratic_3d_pyccel.py`` next to this file is the CPU version with the same arguments. One thread per marker (``positions`` ``(3, N)``, ``values`` -``(components, N)``); see :mod:`mpifield.kernels._metal_grid`. +``(components, N)``); see `mpifield.kernels._metal_grid`. """ from __future__ import annotations diff --git a/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_pyccel.py b/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_pyccel.py index 8058c5c..eac431d 100644 --- a/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_pyccel.py +++ b/src/mpifield/kernels/scatter_quadratic_3d/scatter_quadratic_3d_pyccel.py @@ -2,7 +2,7 @@ The Pyccel version (CPU) of the ``scatter_quadratic_3d`` kernel; ``scatter_quadratic_3d_cuda.cu`` next to this file is the CUDA version (GPU), with the same arguments. Both are loaded -by :mod:`mpifield.kernels.catalog`; this file is compiled with Pyccel there. +by `mpifield.kernels.catalog`; this file is compiled with Pyccel there. Each marker deposits onto its nearest node and the two around it along each axis (``quadratic_cell.cuh``). With OpenMP the markers are split over the threads in diff --git a/src/mpifield/prescribed.py b/src/mpifield/prescribed.py index 8a17372..4613a5d 100644 --- a/src/mpifield/prescribed.py +++ b/src/mpifield/prescribed.py @@ -6,13 +6,25 @@ import cunumpy as xp from maybempi import MPI +from methodcard import method +from mpifield import _method_checks from mpifield._types import Array from mpifield.field import Field Potential = Callable[..., Array] +@method( + equations={ + r"table": r"\phi(\mathbf x) = \phi_j + (\phi_{j+1} - \phi_j) \frac{x_a - z_j}{z_{j+1} - z_j} \ \ (z_j \le x_a \le z_{j+1}), \qquad \phi_0 \text{ below } z_0, \quad \phi_{n-1} \text{ above } z_{n-1}", + }, + assumptions=( + "$x_a$ the coordinate along `axis`; at least two strictly increasing table positions $z_j$", + "potentials in V, positions in m", + ), + symbolic=_method_checks.axial_potential_rules, +) def axial_potential( positions: Sequence[float], values: Sequence[float], @@ -23,7 +35,7 @@ def axial_potential( ``values`` are the potentials at the increasing ``positions`` (m); between them the potential is interpolated linearly, outside them it is constant. The result takes the coordinate arrays of the grid, as - :func:`electric_field_from_potential` calls it. + `electric_field_from_potential` calls it. """ table_positions = xp.asarray(positions, dtype=float) table_values = xp.asarray(values, dtype=float) @@ -38,6 +50,19 @@ def potential(*coordinates: Array) -> Array: return potential +@method( + equations={ + r"field": r"\mathbf E = -\nabla \phi", + r"differences": r"E_{d,k} = -\frac{\phi_{k+1} - \phi_{k-1}}{2 h_d} \text{ inside}, \qquad E_{d,0} = -\frac{\phi_1 - \phi_0}{h_d}, \quad E_{d,N-1} = -\frac{\phi_{N-1} - \phi_{N-2}}{h_d}", + }, + order=2, + assumptions=( + r"$\phi$ evaluated on the nodes; second order inside, first order at both ends of every axis, also a periodic one (no wrapping, unlike `Field.gradient`)", + "the order is that of the interior nodes", + "no Poisson equation is solved: the field is prescribed", + ), + problem=_method_checks.electric_field_problem, +) def electric_field_from_potential( potential: Potential, lower_bound: tuple[float, ...], @@ -53,7 +78,7 @@ def electric_field_from_potential( field has one component per axis, in V/m, from second-order central differences (one-sided at the grid ends). No Poisson equation is solved: the field is fixed, e.g. for test particles. - ``field_options`` go to :class:`~mpifield.Field`. + ``field_options`` go to `Field`. """ ndim = len(lower_bound) grid = Field( diff --git a/src/mpifield/tests/unit/test_method_cards.py b/src/mpifield/tests/unit/test_method_cards.py new file mode 100644 index 0000000..2d93df8 --- /dev/null +++ b/src/mpifield/tests/unit/test_method_cards.py @@ -0,0 +1,5 @@ +"""The method cards of mpifield: complete, seen by the docs, and checked.""" + +from methodcard.testing import method_suite + +TestMethodCards = method_suite("mpifield") diff --git a/src/mpifield/weighting.py b/src/mpifield/weighting.py index 7331606..0e90744 100644 --- a/src/mpifield/weighting.py +++ b/src/mpifield/weighting.py @@ -1,12 +1,12 @@ """Weightings: how marker values are deposited onto the nodes and gathered back. -A :class:`Weighting` pairs a scatter (deposit) kernel and a gather (interpolation) +A `Weighting` pairs a scatter (deposit) kernel and a gather (interpolation) kernel for each number of dimensions it covers, and says how many halo nodes a -split field needs for it. :meth:`Field.scatter ` and -:meth:`Field.gather ` take one; built in are +split field needs for it. `Field.scatter` and +`Field.gather` take one; built in are -- :data:`LINEAR`: cloud-in-cell, two nodes per axis (the default); -- :data:`QUADRATIC`: the quadratic B-spline ("triangular-shaped cloud"), three +- `LINEAR`: cloud-in-cell, two nodes per axis (the default); +- `QUADRATIC`: the quadratic B-spline ("triangular-shaped cloud"), three nodes per axis, the nearest node and its neighbours. The gather of a weighting is the transpose of its scatter, so that @@ -20,7 +20,7 @@ gather(positions, field, result, *geometry) # overwrites result with ``geometry = (*lower, *cell_size, *num_nodes, *periodic)`` -(:attr:`BlockGrid.geometry `: ``ndim`` floats, ``ndim`` +(`BlockGrid.geometry`: ``ndim`` floats, ``ndim`` floats, ``ndim`` ints, ``ndim`` ints 0 or 1). Node ``k`` of axis ``d`` is at ``lower[d] + k * cell_size[d]``; on a periodic axis the last node neighbours the first. ``positions`` is ``(ndim, N)``, ``values`` and ``result`` ``(ncomp, N)``, @@ -28,7 +28,7 @@ arrays. ``field`` is C-contiguous; the marker arrays are C-ordered but may have gaps between rows (``storage[:, :n]``). A marker whose values are all zero (a dead one left in storage) should deposit nothing and not read its position, which may -be non-finite. A :class:`cunumpy.kernels.Kernel` made with ``from_folder`` fits: +be non-finite. A `cunumpy.kernels.Kernel` made with ``from_folder`` fits: it runs its CUDA version on CuPy arrays and its Pyccel or NumPy version on NumPy ones. So does a plain Python function. """ @@ -42,7 +42,9 @@ from typing import Any import cunumpy as xp +from methodcard import method +from mpifield import _method_checks from mpifield._types import Array, BlockGrid from mpifield.kernels.gather_linear_1d import kernel as gather_linear_1d from mpifield.kernels.gather_linear_2d import kernel as gather_linear_2d @@ -73,6 +75,24 @@ def _frozen(kernels: Mapping[int, KernelFunction], role: str) -> Mapping[int, An return MappingProxyType(dict(kernels)) +@method( + equations={ + r"adjoint": r"\sum_i q_i\, g(\mathbf x_i) = \sum_i q_i \sum_{\mathbf k} W_{\mathbf k}(\mathbf x_i)\, G_{\mathbf k} = \sum_{\mathbf k} G_{\mathbf k}\, \rho_{\mathbf k}", + r"halo": r"n_{\mathrm{halo}} = \left\lfloor \frac{s + 1}{2} \right\rfloor \text{ for } s \text{ nodes per axis: 1 for LINEAR, 2 for QUADRATIC}", + }, + assumptions=( + "the gather computes the same nodes and weights as the scatter, with the same folds at bounded edges and wraps at periodic seams, so it is the transpose of the scatter", + "a symmetric field solver, for the momentum to be kept (no self-force)", + "a rank owns the markers from its first node up to the next rank's first node; a quadratic marker there reaches two nodes beyond its last owned node", + ), + references=( + "[Birdsall and Langdon, Plasma Physics via Computer Simulation (1991), ch. 8](https://doi.org/10.1201/9781315275048)", + ), + components=( + "mpifield.weighting.Weighting.scatter_block", + "mpifield.weighting.Weighting.gather_block", + ), +) @dataclass(frozen=True, eq=False, repr=False) class Weighting: """A deposit and its gather, for each number of dimensions they cover. @@ -121,6 +141,27 @@ def gather_kernel(self, ndim: int) -> KernelFunction: """Return the gather kernel for ``ndim`` dimensions.""" return self._kernel(self.gather, ndim, "gather") + @method( + equations={ + r"deposit": r"\rho_{\mathbf k} \mathrel{+}= \sum_i W_{\mathbf k}(\mathbf x_i)\, q_i, \qquad W_{\mathbf k}(\mathbf x) = \prod_d w(r_d - k_d), \qquad r_d = \frac{x_d - x^{\mathrm{lo}}_d}{h_d}", + r"linear_2d": r"W_{i,j} = (1 - f_x)(1 - f_y), \quad W_{i+1,j} = f_x (1 - f_y), \quad W_{i,j+1} = (1 - f_x) f_y, \quad W_{i+1,j+1} = f_x f_y", + r"total": r"\sum_{\mathbf k} \rho_{\mathbf k} = \sum_i q_i", + }, + assumptions=( + r"no division by the cell volume: a density is the deposit divided by $\prod_d h_d$, by the caller", + "a marker whose values are all zero is skipped before its position is read", + "the total holds at bounded edges too, where the weights fold onto the edge nodes", + "CPU threads deposit contiguous chunks of markers into buffers of their own, summed in thread order: no atomics, reproducible for a given thread count; CUDA uses atomic adds, Metal float32", + ), + references=( + "[Birdsall and Langdon, Plasma Physics via Computer Simulation (1991), ch. 8](https://doi.org/10.1201/9781315275048)", + ), + components=( + "mpifield.kernels.linear_cell.linear_cell", + "mpifield.kernels.quadratic_cell.quadratic_cell", + ), + symbolic=_method_checks.scatter_rules, + ) def scatter_block( self, positions: Array, values: Array, field: Array, grid: BlockGrid ) -> None: @@ -140,6 +181,24 @@ def scatter_block( with xp.kernels.kernel_output(field, like=field, dtype=float) as buffer: kernel(positions, values, buffer, *grid.geometry) + @method( + equations={ + r"gather": r"g(\mathbf x_i) = \sum_{\mathbf k} W_{\mathbf k}(\mathbf x_i)\, G_{\mathbf k}, \qquad W_{\mathbf k}(\mathbf x) = \prod_d w(r_d - k_d)", + r"quadratic_3d": r"W_{c_x, c_y, c_z} = \prod_d \left(\tfrac34 - \delta_d^2\right), \qquad \delta_d = r_d - c_d", + }, + assumptions=( + "overwrites the result; reads every position, also of markers whose values are zero", + "the same nodes and weights as the scatter (its transpose)", + ), + references=( + "[Birdsall and Langdon, Plasma Physics via Computer Simulation (1991), ch. 8](https://doi.org/10.1201/9781315275048)", + ), + components=( + "mpifield.kernels.linear_cell.linear_cell", + "mpifield.kernels.quadratic_cell.quadratic_cell", + ), + symbolic=_method_checks.gather_rules, + ) def gather_block( self, positions: Array, field: Array, result: Array, grid: BlockGrid ) -> None: @@ -237,7 +296,7 @@ def _gather_linear_3d( #: Quadratic B-spline weights ("triangular-shaped cloud"): each marker on its #: nearest node and the two around it, with weights ``(1/2 - d)^2 / 2``, #: ``3/4 - d^2`` and ``(1/2 + d)^2 / 2`` for the offset ``d`` (in cells) from the -#: nearest node. Smoother than :data:`LINEAR` (continuous first derivative, less +#: nearest node. Smoother than `LINEAR` (continuous first derivative, less #: aliasing noise) for 3^ndim nodes per marker instead of 2^ndim. On a bounded #: axis a position is clamped to the grid and a stencil node past the edge falls #: on the edge node: the deposit keeps all the charge, the gather extrapolates flat. diff --git a/uv.lock b/uv.lock index 84811b8..946bb0e 100644 --- a/uv.lock +++ b/uv.lock @@ -2309,6 +2309,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/57/99/36cd25f9598eef70db97577440a70a360af505f6f01c2338a18d1556079c/meson-1.12.1-py3-none-any.whl", hash = "sha256:930bc7542cbd9f57009e182fd014eba48cf1a0180a7b9006d2d32d8f168d8b02", size = 1108668, upload-time = "2026-09-22T22:58:26.047Z" }, ] +[[package]] +name = "methodcard" +version = "0.1.0" +source = { git = "https://github.com/max-models/methodcard?branch=devel#cc1c7bd98faefc860c3dee8b8843045c69990430" } +dependencies = [ + { name = "griffe" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, + { name = "numpy", version = "2.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "scope-profiler" }, +] + +[package.optional-dependencies] +symbolic = [ + { name = "sympy" }, +] + [[package]] name = "mistune" version = "3.3.4" @@ -2436,6 +2453,7 @@ source = { editable = "." } dependencies = [ { name = "cunumpy" }, { name = "maybempi" }, + { name = "methodcard" }, { name = "mpiarray" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.11.*'" }, @@ -2458,6 +2476,7 @@ dev = [ { name = "jupyterlab" }, { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "methodcard", extra = ["symbolic"] }, { name = "nbclient" }, { name = "nbformat" }, { name = "plasma-plots" }, @@ -2496,6 +2515,7 @@ plot = [ test = [ { name = "matplotlib", version = "3.10.9", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "matplotlib", version = "3.11.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "methodcard", extra = ["symbolic"] }, { name = "plasma-plots" }, { name = "pytest" }, { name = "pytest-cov" }, @@ -2512,6 +2532,8 @@ requires-dist = [ { name = "matplotlib", marker = "extra == 'docs'", specifier = ">=3.10.9" }, { name = "matplotlib", marker = "extra == 'plot'", specifier = ">=3.10" }, { name = "maybempi", git = "https://github.com/max-models/maybempi?branch=devel" }, + { name = "methodcard", git = "https://github.com/max-models/methodcard?branch=devel" }, + { name = "methodcard", extras = ["symbolic"], marker = "extra == 'test'", git = "https://github.com/max-models/methodcard?branch=devel" }, { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'metal'" }, { name = "mpi4py", marker = "extra == 'mpi'" }, { name = "mpiarray", git = "https://github.com/max-models/mpiarray?branch=devel" },