Build a large free-space launch in seconds, with the same bits

changelog.d/setup-cost.changed.md

@@ -0,0 +1 @@
+Building a large free-space launch is faster, with every bit of the launch as it was: its passes over the beam's elements run in threads, its admittance is formed once, and the transverse residual no longer takes a Python `Fraction` per node. Digesting an array not held in its own order no longer builds a whole volume, so a broadcast input costs a few megabytes. (#493, #347)

docs/dev/benchmarks.md

@@ -2493,3 +2493,22 @@ A default radial run of 256 nodes and 8192 frequency bins, which carries input n
 83.21 MB of JSON before and 0.04 MB after, and its first hash goes from 2.603 s to 0.077 s.
 A 64-node grid is about the same before and after, since its arrays sit close to the limit.
 A run into `MemoryOutput` computes no hash at all.
+
+## A free-space point's setup
+
+Building and translating a large free-space launch took longer than the steps of a short point:
+the passes over a beam's elements ran on one thread, its admittance ran twice, and the transverse
+residual took a Python `Fraction` per node.
+
+Measured on an Intel Xeon Gold 6240, one socket's 18 cores, 18 threads, on a shared machine, twice each, unprofiled:
+a 768 × 768 × 600 envelope launch through a Sellmeier glass with a Kerr response, `e49925d1` against `ece75dd7`.
+
+| | before (s) | after (s) |
+|---|---:|---:|
+| `Simulation.from_free` | 169, 179 | 14.9, 14.9 |
+| `translate` | 30.3, 27.5 | 7.2, 7.8 |
+| `runner` | 3.8, 3.0 | 3.5, 3.3 |
+
+A launch's bits don't change, which `tests/test_pulses.py` and `tests/test_cartesian_precision.py` hold across
+thread counts. The sum that scales the field stays one sum in one order for that reason, so it stays on one thread.
+

src/solprop/boundaries.py

@@ -77,6 +77,7 @@
 import numpy as np
 from jaxtyping import Bool, Complex, Float, Inexact, UInt8
 
+from solprop._process import parallel_rows
 from solprop.exceptions import GridError, SolpropWarning
 from solprop.grid import Collar, TimeGrid
 from solprop.spatial_absorber import SpatialAbsorber, SpatialAbsorberDesign
@@ -1084,15 +1085,36 @@ def occupancy(
             Shape ``(SPLIT_RATE_BUCKETS,)``, in the state's units of energy.
         """
         # In the state's own precision, without a state-sized temporary: the shares
-        # are what the limit reads.
+        # are what the limit reads, so the order of summing is free.
         spectrum = np.asarray(field_spectrum)
         re, im = np.real(spectrum), np.imag(spectrum)
-        energy = np.einsum("bmpw,bmpw->mw", re, re) + np.einsum("bmpw,bmpw->mw", im, im)
-        if columns is not None:
-            energy = energy * np.asarray(columns, dtype=np.float64)[:, None]
-        return np.bincount(
-            self.buckets.ravel(), weights=energy.ravel(), minlength=SPLIT_RATE_BUCKETS
-        ).astype(np.float64)
+        weights = None if columns is None else np.asarray(columns, dtype=np.float64)
+
+        def block(rows: slice) -> Float[np.ndarray, " n_buckets"]:
+            """Return a block of modes' energy per bucket.
+
+            Parameters
+            ----------
+            rows : slice
+                The block of modes.
+
+            Returns
+            -------
+            numpy.ndarray
+                The block's energy in each bucket.
+            """
+            energy = np.einsum("bmpw,bmpw->mw", re[:, rows], re[:, rows]) + np.einsum(
+                "bmpw,bmpw->mw", im[:, rows], im[:, rows]
+            )
+            if weights is not None:
+                energy = energy * weights[rows, None]
+            return np.bincount(
+                self.buckets[rows].ravel(),
+                weights=energy.ravel(),
+                minlength=SPLIT_RATE_BUCKETS,
+            ).astype(np.float64)
+
+        return np.sum(parallel_rows(spectrum.shape[1], spectrum.nbytes, block), axis=0)
 
     def limit(self, occupancy: Float[np.ndarray, " n_buckets"]) -> float:
         """Return the crossing limit of an occupancy, m.

src/solprop/output/schema.py

@@ -535,6 +535,10 @@ def serialise(value: object) -> JsonValue:
     return _serialise_object(value)
 
 
+_DIGEST_BLOCK = 2**16
+"""Elements a digest of an array not held in its order reads at a time."""
+
+
 @dataclass(frozen=True)
 class ArrayDigest:
     """An array recorded by the digest of its contents rather than inline.
@@ -574,21 +578,38 @@ def of(cls, values: np.ndarray) -> "ArrayDigest":  # noqa: UP037 - bound after t
         Parameters
         ----------
         values : numpy.ndarray
-            The array; read once, in little-endian byte order, without a copy when it
-            is already contiguous and little-endian.
+            The array; read once, in C order and little-endian byte order, a block at
+            a time unless it is held that way already, so that a broadcast input never
+            becomes a whole volume.
 
         Returns
         -------
         ArrayDigest
             The digest, the dtype and the shape.
         """
         array = np.asarray(values)
-        # One byte order, so that the same numbers digest alike on every platform; a
-        # copy only where the array is not contiguous in it already.
-        ordered = np.ascontiguousarray(array, dtype=array.dtype.newbyteorder("<"))
+        # One byte order, so that the same numbers digest alike on every platform.
+        little = array.dtype.newbyteorder("<")
+        digest = hashlib.sha256()
+        if array.size == 0:
+            pass
+        elif array.flags.c_contiguous and array.dtype == little:
+            digest.update(memoryview(array).cast("B"))
+        else:
+            # NumPy's stubs leave "contig" out of the flags they list.
+            for block in np.nditer(  # pyright: ignore[reportCallIssue, reportUnknownVariableType]
+                array,
+                flags=["external_loop", "buffered"],
+                op_flags=["readonly", "contig"],  # pyright: ignore[reportArgumentType]
+                op_dtypes=[little],
+                casting="equiv",
+                order="C",
+                buffersize=_DIGEST_BLOCK,
+            ):
+                digest.update(np.asarray(block).tobytes())
         return cls(
-            sha256=hashlib.sha256(memoryview(ordered).cast("B")).hexdigest(),
-            dtype=str(ordered.dtype),
+            sha256=digest.hexdigest(),
+            dtype=str(little),
             shape=tuple(int(size) for size in array.shape),
         )
 

src/solprop/pulses.py

@@ -43,6 +43,7 @@
 
 from solprop import fourier
 from solprop._display import Displayed, engineering, line, phase_text, summarise
+from solprop._process import parallel_rows
 from solprop.constants import EPS0, C
 from solprop.exceptions import PulseError, SolpropWarning
 from solprop.grid import FreeGrid, LineGrid, RadialGrid, TimeGrid
@@ -350,9 +351,8 @@ def build(
             else np.fft.fft2(profile.reshape(space.n_y, space.n_x)).ravel()
         )
         jones = _unit_jones(self.jones)
-        live = np.concatenate(
-            [flux.admittance(columns) != 0 for columns in flux.chunks()]
-        )
+        admittance = flux.whole_admittance()
+        live = admittance != 0
         exponent = None
         if self.propz != 0:
             working_frame = Frame.comoving(grid) if frame is None else frame
@@ -368,10 +368,8 @@ def build(
                 msg = "SpatialPulse requires a scalar single-column temporal pulse"
                 raise PulseError(msg)
             # A unit-peak profile keeps intermediate values near the final scale.
-            field = (
-                spectrum[:, None, None]
-                * fourier.to_internal(jones)[None, :, None]
-                * temporal
+            field = _separable_field(
+                spectrum, fourier.to_internal(jones), temporal, grid.n_omega
             )
             if exponent is not None:
                 if np.max(np.real(exponent)) > math.log(np.finfo(float).max) - math.log(
@@ -383,7 +381,8 @@ def build(
             if aligner is not None:
                 field = aligner.apply(field)
             field = _hard_masked(field, live, space.spectral_weights, grid)
-            total += _normalised(member, temporal, field, flux, grid)
+            field = _normalised(member, temporal, field, flux, admittance, grid)
+            _add_rows(total, field)
         # Each member is a finite field scaled by the finite factor of a finite,
         # positive measured flux (_normalising_factor refuses the rest), so the sum
         # is finite.
@@ -1139,17 +1138,36 @@ def _hard_masked(
     Returns
     -------
     numpy.ndarray
-        The masked beam.
+        The masked beam: `field`, masked in place.
 
     Raises
     ------
     PulseError
         If the nonpropagating share of the field norm exceeds its rounding budget.
     """
     band = grid.active_mask & (grid.omega > 0)
-    norm = weights[:, None] * np.sum(np.abs(field) ** 2, axis=1)
-    removed = float(np.sum(norm[band[None, :] & ~live]))
-    total = float(np.sum(norm[:, band]))
+    outside = band[None, :] & ~live
+
+    def shares(rows: slice) -> tuple[float, float]:
+        """Return a block's norm in the columns that don't propagate, and in all.
+
+        Parameters
+        ----------
+        rows : slice
+            The block of the first axis.
+
+        Returns
+        -------
+        tuple of float
+            The block's two sums.
+        """
+        norm = weights[rows, None] * np.sum(np.abs(field[rows]) ** 2, axis=1)
+        return float(np.sum(norm[outside[rows]])), float(np.sum(norm[:, band]))
+
+    # Judged against a budget of rounding, so that the blocks' order of summing is free.
+    parts = parallel_rows(field.shape[0], field.nbytes, shares)
+    removed = math.fsum(part[0] for part in parts)
+    total = math.fsum(part[1] for part in parts)
     # The transforms that formed the beam carry a few rounding units of its norm
     # per bin they sum; content above that is physical.
     budget = 8 * field.shape[0] * grid.n_omega * np.finfo(float).eps
@@ -1160,7 +1178,80 @@ def _hard_masked(
             "widen the waist or refine the transverse grid"
         )
         raise PulseError(msg)
-    return field * live[:, None, :]
+
+    def mask(rows: slice) -> None:
+        """Zero one block's dead bins.
+
+        Parameters
+        ----------
+        rows : slice
+            The block of the first axis.
+        """
+        field[rows] *= live[rows, None, :]
+
+    parallel_rows(field.shape[0], field.nbytes, mask)
+    return field
+
+
+def _separable_field(
+    spectrum: np.ndarray, jones: np.ndarray, temporal: np.ndarray, n_omega: int
+) -> np.ndarray:
+    """Return ``spectrum * jones * temporal`` in a new array, its rows in threads.
+
+    Parameters
+    ----------
+    spectrum : numpy.ndarray
+        The beam's transverse spectrum, ``(n_modes,)``.
+    jones : numpy.ndarray
+        Its polarisation, ``(n_pol,)``.
+    temporal : numpy.ndarray
+        Its temporal spectrum, ``(1, 1, n_omega)``.
+    n_omega : int
+        The bins.
+
+    Returns
+    -------
+    numpy.ndarray
+        The field, ``(n_modes, n_pol, n_omega)``.
+    """
+    field = np.empty((spectrum.size, jones.size, n_omega), dtype=np.complex128)
+
+    def work(rows: slice) -> None:
+        """Form one block, whose first touch is then this thread's.
+
+        Parameters
+        ----------
+        rows : slice
+            The block of the first axis.
+        """
+        field[rows] = spectrum[rows, None, None] * jones[None, :, None] * temporal
+
+    parallel_rows(spectrum.size, field.nbytes, work)
+    return field
+
+
+def _add_rows(total: np.ndarray, part: np.ndarray) -> None:
+    """Add `part` to `total` in place, row by row, in threads.
+
+    Parameters
+    ----------
+    total : numpy.ndarray
+        The sum so far.
+    part : numpy.ndarray
+        What to add, of the same shape.
+    """
+
+    def work(rows: slice) -> None:
+        """Add one block.
+
+        Parameters
+        ----------
+        rows : slice
+            The block of the first axis.
+        """
+        total[rows] += part[rows]
+
+    parallel_rows(total.shape[0], total.nbytes, work)
 
 
 def _members(pulse: Pulse) -> list[Pulse]:
@@ -1243,15 +1334,26 @@ def _normalising_factor(
     quantity, target = _target(pulse, temporal, grid)
     if quantity == "energy":
         energy_weights = fourier.energy_weights(grid.n_omega, grid.dt, kind=grid.kind)
-        density = np.sum(np.abs(beam) ** 2, axis=1)
-        measured = plane.scale * float(
-            np.sum(
-                plane.weights[:, None]
+        terms = np.empty((beam.shape[0], beam.shape[2]))
+
+        def work(rows: slice) -> None:
+            """Form one block of the flux's terms.
+
+            Parameters
+            ----------
+            rows : slice
+                The block of the first axis.
+            """
+            terms[rows] = (
+                plane.weights[rows, None]
                 * energy_weights
-                * np.real(plane.admittance)
-                * density
+                * np.real(plane.admittance[rows])
+                * np.sum(np.abs(beam[rows]) ** 2, axis=1)
             )
-        )
+
+        parallel_rows(beam.shape[0], beam.nbytes, work)
+        # One sum over the whole array, in one order: this factor scales the field.
+        measured = plane.scale * float(np.sum(terms))
     else:
         measured = float(
             np.max(
@@ -1275,9 +1377,10 @@ def _normalised(
     temporal: np.ndarray,
     field: np.ndarray,
     flux: SpatialFlux,
+    admittance: np.ndarray,
     grid: TimeGrid,
 ) -> np.ndarray:
-    """Scale a spatial beam to its requested physical energy or peak power.
+    """Scale a spatial beam in place to its requested physical energy or peak power.
 
     Parameters
     ----------
@@ -1286,9 +1389,11 @@ def _normalised(
     temporal : numpy.ndarray
         Its built spectrum, shape ``(1, 1, n_omega)``.
     field : numpy.ndarray
-        The unnormalised beam spectrum, ``(n_modes, n_pol, n_omega)``.
+        The unnormalised beam spectrum, ``(n_modes, n_pol, n_omega)``, scaled in place.
     flux : solprop.spatial.SpatialFlux
         The launch plane's measure.
+    admittance : numpy.ndarray
+        Its admittance, ``(n_modes, n_omega)``.
     grid : solprop.grid.TimeGrid
         Temporal grid.
 
@@ -1299,12 +1404,23 @@ def _normalised(
     """
     plane = _Plane(
         weights=flux.space.spectral_weights,
-        admittance=np.concatenate(
-            [flux.admittance(columns) for columns in flux.chunks()]
-        ),
+        admittance=admittance,
         scale=flux.energy_scale,
     )
-    return field * _normalising_factor(pulse, temporal, field, plane, grid)
+    factor = _normalising_factor(pulse, temporal, field, plane, grid)
+
+    def scale(rows: slice) -> None:
+        """Scale one block.
+
+        Parameters
+        ----------
+        rows : slice
+            The block of the first axis.
+        """
+        field[rows] *= factor
+
+    parallel_rows(field.shape[0], field.nbytes, scale)
+    return field
 
 
 def _frozen_phase(phase: Sequence[float] | np.ndarray) -> Float[np.ndarray, " n_deg"]: