[WIP] Space Charge: bound the dynamically resized mesh

docs/source/usage/parameters.rst

@@ -2477,6 +2477,74 @@ See there ``nslice`` option on lattice elements for slicing.
     For instance, ``1.2`` means the mesh will span 10% above and 10% below the beam;
     ``1.0`` means the beam is exactly covered with the mesh.
 
+.. pp:param:: geometry.prob_relative_max
+    :type: ``float``
+    :unit: dimensionless
+    :optional:
+    :default: 10% above ``prob_relative[0]`` for :pp:param:`algo.poisson_solver` ``fft``, ``prob_relative[0]`` otherwise
+
+    The largest the field mesh may be, as a multiple of the maximum physical extent of
+    beam particles, while :pp:param:`geometry.prob_relative` is the smallest.
+    Setting it equal to ``prob_relative[0]`` fits the beam exactly, which is the behavior
+    when this is not used.
+
+    Leaving room between the two matters for the FFT solver. Fitting the mesh exactly
+    means it changes with every fluctuation of the beam extent, and the solver then has
+    to rebuild its Green's function on every slice step. Given a band, the mesh is
+    instead chosen from a fixed set of lengths, so a beam of nearly the same size lands
+    on exactly the same mesh and the solver reuses the Green's function it already has.
+    On a constant-focusing channel this takes the number of Green's functions built from
+    one per slice step to a handful.
+
+    The allowed lengths are
+
+    .. math::
+
+        L_k = 2^{k/m}\ \mathrm{m}, \qquad
+        m = \left\lceil \frac{1}{\log_2 r} \right\rceil, \qquad
+        r = \frac{\texttt{prob\_relative\_max}}{\texttt{prob\_relative[0]}}
+
+    where :math:`L_k` is the mesh edge length along one axis, :math:`k` is a whole number
+    chosen so that :math:`L_k` is the smallest such length still covering the padded
+    beam, :math:`m` is how many allowed lengths there are per factor of two in size, and
+    :math:`r` is the width of the band. Rounding :math:`m` up guarantees the mesh never
+    exceeds ``prob_relative_max``. The default band of 10% gives :math:`m = 8`, so the
+    mesh is at most 9.1% above ``prob_relative[0]``.
+
+    A wider band reuses the Green's function more often, at the price of a mesh that can
+    be coarser than requested. There is no generally right value, so consider a
+    convergence test if in doubt. The longitudinal direction is treated in the frame the
+    solver works in, so a changing reference energy resizes the mesh on its own.
+
+    The solver offers a second, weaker way to reuse a Green's function,
+    :pp:param:`ablastr.igf_cache_tolerance`, which lets one serve a grid slightly
+    different from the one it was built for. The two are not interchangeable, and this
+    parameter is the one to reach for first:
+
+    .. list-table::
+        :header-rows: 1
+        :widths: 22 39 39
+
+        * -
+          - ``geometry.prob_relative_max``
+          - ``ablastr.igf_cache_tolerance``
+        * - changes
+          - which mesh is solved on
+          - nothing about the mesh
+        * - Green's function
+          - matches its mesh exactly
+          - built for a slightly different mesh
+        * - price
+          - resolution, up to one allowed length of extra padding
+          - an error of the order of the tolerance
+        * - reuse is
+          - bit-identical
+          - not bit-identical
+
+    ImpactX chooses its own mesh, so it can make one come back exactly and never needs
+    the tolerance. That option is open only to a caller that owns the mesh; where the
+    quantity that varies is measured rather than chosen, the tolerance is the only way.
+
 .. pp:param:: geometry.prob_lo/hi
     :link_aliases: geometry.prob_lo geometry.prob_hi
     :type: ``3 floats``
@@ -2524,6 +2592,77 @@ See there ``nslice`` option on lattice elements for slicing.
       For the MLMG solver, we assume `Dirichlet boundary conditions <https://en.wikipedia.org/wiki/Dirichlet_boundary_condition>`__ with zero potential (a mirror charge).
       Thus, to emulate open boundaries, consider adding enough vacuum padding to the beam.
 
+.. pp:param:: ablastr.igf_cache_max_entries
+    :type: ``integer``
+    :optional:
+    :default: ``8``
+
+    Number of Green's functions the ``fft`` solver keeps for reuse, evicting the least
+    recently used one beyond that.
+
+    The Green's function depends only on the mesh, so it can be reused whenever the mesh
+    returns to a size it has had before, which is what the padding band set by
+    :pp:param:`geometry.prob_relative_max` arranges. Keeping several is what lets a beam
+    that breathes, as in a periodic lattice, cycle through its sizes without rebuilding.
+
+    ``0`` keeps only the Green's function currently in use, which is still reused for as
+    long as the mesh does not change. Each entry costs :math:`32 n^3` bytes in single and
+    :math:`64 n^3` in double precision, which is a sizable fraction of a GPU at large
+    mesh sizes, so consider lowering this or setting
+    :pp:param:`ablastr.igf_cache_max_bytes` there.
+
+.. pp:param:: ablastr.igf_cache_max_bytes
+    :type: ``integer``
+    :unit: bytes
+    :optional:
+    :default: a quarter of the free GPU memory; unlimited on CPU
+
+    Memory budget for the Green's functions kept by the ``fft`` solver.
+    Entries are evicted, least recently used first, to stay within it.
+    ``0`` means no limit, in which case only :pp:param:`ablastr.igf_cache_max_entries`
+    bounds the cache.
+
+.. pp:param:: ablastr.igf_cache_tolerance
+    :type: ``float``
+    :unit: dimensionless
+    :optional:
+    :default: 64 times the machine epsilon of the precision in use
+
+    How closely two cell sizes must agree for one Green's function to serve both.
+    The default only absorbs round-off, so it never trades accuracy for reuse.
+
+    ImpactX does not need to raise this. Because
+    :pp:param:`geometry.prob_relative_max` makes the mesh come back exactly, a reused
+    Green's function is still the right one for the mesh it is used on. Raising this
+    instead reuses one across meshes that merely nearly agree, which costs an error of
+    the order of the tolerance, so raise it only deliberately.
+
+    It exists for callers that cannot choose their mesh, where the quantity that varies
+    is measured rather than chosen. The relativistic electrostatic solver in WarpX is
+    the example: it derives the longitudinal stretch from a velocity, and recovering it
+    costs about :math:`\epsilon\gamma^2`.
+
+.. pp:param:: ablastr.igf_rebuild_always
+    :type: ``integer``
+    :optional:
+    :default: ``0``
+
+    Set to ``1`` to rebuild the Green's function on every solve, as the ``fft`` solver
+    did before it reused them.
+    This is the reference to compare against when checking that reuse leaves results
+    unchanged, and is not otherwise useful.
+
+.. pp:param:: ablastr.igf_cache_verbose
+    :type: ``integer``
+    :optional:
+    :default: ``0``
+
+    Set to ``1`` to report, at the end of the run, how many Green's functions the ``fft``
+    solver reused, built and evicted.
+    A cache that is too small for a periodic lattice keeps evicting the entry it is about
+    to need again, which costs time without ever showing up as a wrong answer, so this is
+    worth checking when tuning :pp:param:`ablastr.igf_cache_max_entries`.
+
 Multigrid-specific numerical options:
 
 .. pp:param:: algo.mlmg_relative_tolerance

docs/source/usage/python.rst

@@ -87,6 +87,38 @@ Collective Effects & Overall Simulation Parameters
 
       Use dynamic (``True``) resizing of the field mesh or static sizing (``False``).
 
+   .. py:property:: prob_relative_max
+
+      The largest the field mesh may be, as a multiple of the maximum physical extent of
+      beam particles, while ``prob_relative`` is the smallest.
+      Setting it equal to ``prob_relative[0]`` fits the beam exactly.
+
+      Leaving room between the two lets the mesh be chosen from a fixed set of sizes, so
+      that a beam of nearly the same size lands on exactly the same mesh and the FFT
+      space-charge solver reuses the Green's function it already has instead of
+      rebuilding it on every slice step.
+
+      Because the mesh comes back exactly, the reused Green's function is still the right
+      one for it and the result is bit-identical to rebuilding. The price is paid in
+      resolution, up to one allowed length of extra padding, rather than in accuracy.
+      The solver also accepts ``ablastr.igf_cache_tolerance``, which reuses a Green's
+      function across meshes that merely nearly agree, at the price of an error of that
+      order. ImpactX does not need it, since it chooses its own mesh.
+
+      Default: 10% above ``prob_relative[0]`` for ``poisson_solver = "fft"``, and
+      ``prob_relative[0]`` otherwise.
+
+   .. py:property:: igf_cache_max_entries
+
+      Number of Green's functions the FFT space-charge solver keeps for reuse, evicting
+      the least recently used one beyond that. ``0`` keeps only the one in use, which is
+      still reused for as long as the mesh does not change.
+
+      Each entry costs ``32 * n**3`` bytes in single and ``64 * n**3`` in double
+      precision, so consider lowering this on a GPU with a large mesh.
+
+      Default: ``8``.
+
    .. py:property:: space_charge
 
       The physical model of space charge used.

src/initialization/InitMeshRefinement.H

@@ -19,6 +19,7 @@
 #include <AMReX_REAL.H>
 #include <AMReX_Utility.H>
 
+#include <cmath>
 #include <limits>
 #include <stdexcept>
 #include <string>
@@ -27,7 +28,7 @@
 
 namespace impactx::initialization
 {
-    amrex::Vector<amrex::Real>
+    inline amrex::Vector<amrex::Real>
     read_mr_prob_relative ()
     {
         amrex::ParmParse pp_algo("algo");
@@ -80,4 +81,86 @@ namespace impactx::initialization
 
         return prob_relative;
     }
+
+    /** Settings that restrict the dynamically resized mesh to a discrete set of sizes.
+     *
+     * Fitting the box to the beam exactly means it moves with every fluctuation of the
+     * particle extent, and the space-charge solver then has to rebuild its Green's
+     * function each time. Allowing the box to be a little larger than requested, and
+     * choosing it from a fixed set of lengths, makes a beam of nearly the same size land
+     * on exactly the same mesh, which is what lets the Green's function be reused.
+     */
+    struct GridQuantization
+    {
+        bool enabled = false;              //!< restrict the box to the fit lengths
+        int fit_lengths_per_doubling = 8;  //!< how many fit lengths per factor of two
+    };
+
+    /** Read the padding band, geometry.prob_relative and geometry.prob_relative_max */
+    inline GridQuantization
+    read_grid_quantization ()
+    {
+        amrex::ParmParse pp_algo("algo");
+        amrex::ParmParse pp_geometry("geometry");
+
+        std::string poisson_solver = "fft";
+        pp_algo.queryAdd("poisson_solver", poisson_solver);
+
+        amrex::Real const prob_relative_min = read_mr_prob_relative()[0];
+
+        // By default allow a 10% band for the FFT solver, which reuses its Green's
+        // function when the mesh holds still. The multigrid solver gains nothing from a
+        // steady mesh, so there the box keeps fitting the beam exactly.
+        amrex::Real prob_relative_max = (poisson_solver == "fft")
+            ? prob_relative_min * amrex::Real(1.1) : prob_relative_min;
+        pp_geometry.queryAddWithParser("prob_relative_max", prob_relative_max);
+
+        if (prob_relative_max < prob_relative_min) {
+            throw std::runtime_error(
+                "geometry.prob_relative_max must be >= geometry.prob_relative[0]");
+        }
+
+        GridQuantization quant;
+
+        amrex::Real const band = prob_relative_max / prob_relative_min;
+        if (!(band > amrex::Real(1.))) {
+            return quant;  // no band: fit the beam exactly, as if this were never set
+        }
+
+        // Pick the finest set of lengths whose spacing still fits inside the band, so
+        // that rounding up can never exceed prob_relative_max. Keeping this a whole
+        // number of lengths per doubling is what makes two boxes a factor of two apart
+        // differ by exactly two, which the solver relies on to reuse a Green's function
+        // built at another scale.
+        quant.enabled = true;
+        quant.fit_lengths_per_doubling =
+            static_cast<int>(std::ceil(amrex::Real(1.) / std::log2(band)));
+
+        return quant;
+    }
+
+    /** Smallest length not below @p extent that the mesh is allowed to take.
+     *
+     * The allowed lengths are 2^(k/m) meters for whole numbers k, with m fit lengths per
+     * doubling. Anchoring them at 1 m, rather than at the first box seen, keeps them
+     * reproducible across runs and across turns of a ring, so a beam that returns to a
+     * size it had before lands on the very same mesh.
+     *
+     * @param[in] extent the length to round up, in meters
+     * @param[in] fit_lengths_per_doubling number of allowed lengths m per factor of two
+     * @return the smallest allowed length that is not below @p extent
+     */
+    inline amrex::Real
+    smallest_fit_length (amrex::Real extent, int fit_lengths_per_doubling)
+    {
+        if (!(extent > amrex::Real(0.)) || !std::isfinite(extent)) { return extent; }
+
+        auto const m = static_cast<amrex::Real>(fit_lengths_per_doubling);
+        amrex::Real const k = std::ceil(std::log2(extent) * m);
+        amrex::Real const length = std::exp2(k / m);
+
+        // Guard against the rounding of log2/exp2 landing a hair below the request,
+        // which would clip the beam.
+        return (length < extent) ? std::exp2((k + amrex::Real(1.)) / m) : length;
+    }
 } // namespace impactx::initialization