From 4c35178cbbafde7ce968d0a5b0e628069c537cc6 Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 15:03:18 +0000 Subject: [PATCH 1/8] Digest an array not held in its order block by block A broadcast input no longer becomes a whole volume to be digested, and a byte-swapped one no longer takes a copy. Every digest is as it was, which the record tests hold. Closes #493. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- src/solprop/output/schema.py | 34 +++++++++++++++++++------ tests/output/test_field_retention.py | 37 ++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 7 deletions(-) diff --git a/src/solprop/output/schema.py b/src/solprop/output/schema.py index c25420a1..01cea9f1 100644 --- a/src/solprop/output/schema.py +++ b/src/solprop/output/schema.py @@ -458,6 +458,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. @@ -497,8 +501,9 @@ 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 ------- @@ -506,12 +511,27 @@ def of(cls, values: np.ndarray) -> "ArrayDigest": # noqa: UP037 - bound after t 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: + for block in np.nditer( + array, + flags=["external_loop", "buffered"], + op_flags=[["readonly", "contig"]], + 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), ) diff --git a/tests/output/test_field_retention.py b/tests/output/test_field_retention.py index 1bb669d4..be6cd678 100644 --- a/tests/output/test_field_retention.py +++ b/tests/output/test_field_retention.py @@ -10,6 +10,7 @@ import dataclasses import json +import tracemalloc from typing import cast import h5py @@ -327,6 +328,42 @@ def test_a_field_free_file_resumes_only_from_its_checkpoint(tmp_path, metadata): # ----------------------------------------------------------- record helpers +@pytest.mark.parametrize( + "view", + [ + lambda a: a.astype(">c16"), + np.asfortranarray, + lambda a: np.repeat(a, 2, axis=0)[::2], + lambda a: a[::-1, ::-1, ::-1][::-1, ::-1, ::-1], + ], + ids=["byte_swapped", "fortran", "strided", "reversed_twice"], +) +def test_the_same_numbers_digest_alike_however_they_are_held(view): + numbers = np.arange(5 * 3 * 7, dtype=np.complex128).reshape(5, 3, 7) + assert ArrayDigest.of(view(numbers)) == ArrayDigest.of(numbers) + + +@pytest.mark.parametrize("shape", [(0,), (3, 0, 2), ()], ids=["empty", "hollow", "0d"]) +def test_an_array_without_a_contiguous_copy_digests_at_any_shape(shape): + numbers = np.ones(shape, dtype=np.float64) + assert ArrayDigest.of(numbers.astype(">f8")) == ArrayDigest.of(numbers) + + +def test_a_broadcast_input_digests_without_a_volume(): + # A simulation holds a broadcast placeholder as one repeated block: its digest is + # that of the whole volume, and its peak is a buffer's, not the volume's. + shape = (4096, 1, 1024) + view = np.broadcast_to(np.zeros((1, 1, 1), dtype=np.complex128), shape) + tracemalloc.start() + try: + digest = ArrayDigest.of(view) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert digest == ArrayDigest.of(np.zeros(shape, dtype=np.complex128)) + assert peak < 16 * 2**20 < 16 * int(np.prod(shape)) + + def test_an_array_above_the_inline_limit_is_recorded_by_its_digest(): small = np.zeros(4) assert recorded_input(small) is small From e20e28228ced790bc3ac278cc4a839ea4e56e19a Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 15:03:18 +0000 Subject: [PATCH 2/8] Form the transverse residual from exact squares and math.fsum The same bits as the rational sum, which a test holds, without a Fraction per node. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- src/solprop/spatial.py | 106 ++++++++++++++++++++++++++---- tests/test_cartesian_precision.py | 41 +++++++++++- 2 files changed, 133 insertions(+), 14 deletions(-) diff --git a/src/solprop/spatial.py b/src/solprop/spatial.py index 84e67fd6..b849c12c 100644 --- a/src/solprop/spatial.py +++ b/src/solprop/spatial.py @@ -25,6 +25,7 @@ from numpy.typing import ArrayLike from solprop import fourier +from solprop._process import parallel_rows from solprop.constants import EPS0, ROOM_TEMPERATURE, C from solprop.exceptions import GridError, MaterialError, OutputError from solprop.grid import FreeGrid, LineGrid, RadialGrid, TimeGrid @@ -519,19 +520,7 @@ def cartesian_squares( Inactive material bins are zero. No spatial-frequency table is built. """ square, low = self.material_squares(z) - x_squared = [Fraction(float(value)) ** 2 for value in space.k_x] - y_squared = [Fraction(float(value)) ** 2 for value in space.k_y] - transverse_low = np.array( - [ - float(x + y - Fraction(float(high))) - for (x, y), high in zip( - ((x, y) for y in y_squared for x in x_squared), - space.kperp_sq, - strict=True, - ) - ] - ) - return square, low, transverse_low + return square, low, _transverse_residual(space) def material_squares( self, z: float = 0.0 @@ -918,6 +907,26 @@ def admittance( columns, ) + def whole_admittance(self) -> Complex[np.ndarray, "n_modes n_omega"]: + """Return every column's admittance, its chunks formed in threads. + + Returns + ------- + numpy.ndarray + :meth:`admittance` of every column, shape ``(n_modes, n_omega)``. + """ + chunks = self.chunks() + admittance = np.empty( + (self.space.spectral_weights.size, self.grid.n_omega), dtype=np.complex128 + ) + + def work(rows: slice) -> None: + for columns in chunks[rows]: + admittance[columns] = self.admittance(columns) + + parallel_rows(len(chunks), admittance.nbytes, work) + return admittance + def chunks(self) -> list[slice]: """Return column slices bounding an integrated sum's working memory. @@ -934,6 +943,77 @@ def chunks(self) -> list[slice]: ] +def _exact_squares( + values: Float[np.ndarray, " n"], +) -> tuple[Float[np.ndarray, " n"], Float[np.ndarray, " n"]]: + """Return each value's square as the sum of two floats, exactly (Dekker, Veltkamp). + + Parameters + ---------- + values : numpy.ndarray + Each zero or of magnitude within ``[2**-450, 2**500]``, where the split can't + overflow and the low part can't underflow. + + Returns + ------- + tuple of numpy.ndarray + The rounded squares and their exact residuals. + """ + split = 134217729.0 * values # 2**27 + 1 + upper = split - (split - values) + lower = values - upper + high = values * values + return high, ((upper * upper - high) + 2.0 * upper * lower) + lower * lower + + +def _transverse_residual(space: FreeGrid) -> Float[np.ndarray, " n_modes"]: + """Return ``k_x**2 + k_y**2 - kperp_sq`` per column, rounded from its exact value. + + ``math.fsum`` rounds the exact sum of its floats, as a rational sum would. + + Parameters + ---------- + space : solprop.grid.FreeGrid + Stored unsquared transverse axes in rad/m, with x fastest. + + Returns + ------- + numpy.ndarray + The real residuals in m^-2. + """ + k_x, k_y = np.asarray(space.k_x), np.asarray(space.k_y) + moduli = np.abs(np.concatenate((k_x, k_y))) + moduli = moduli[moduli != 0.0] + if moduli.size and not (moduli.min() >= 2.0**-450 and moduli.max() <= 2.0**500): + x_squared = [Fraction(float(value)) ** 2 for value in k_x] + y_squared = [Fraction(float(value)) ** 2 for value in k_y] + return np.array( + [ + float(x + y - Fraction(float(high))) + for (x, y), high in zip( + ((x, y) for y in y_squared for x in x_squared), + space.kperp_sq, + strict=True, + ) + ] + ) + x_parts = list(zip(*map(np.ndarray.tolist, _exact_squares(k_x)), strict=True)) + y_parts = list(zip(*map(np.ndarray.tolist, _exact_squares(k_y)), strict=True)) + fsum = math.fsum + return np.fromiter( + ( + fsum((y_high, y_low, x_high, x_low, -high)) + for ((y_high, y_low), (x_high, x_low)), high in zip( + ((y, x) for y in y_parts for x in x_parts), + np.asarray(space.kperp_sq).tolist(), + strict=True, + ) + ), + dtype=np.float64, + count=k_x.size * k_y.size, + ) + + def _check_flux( spectrum: np.ndarray, flux: SpatialFlux, *, planar: bool | None = None ) -> None: diff --git a/tests/test_cartesian_precision.py b/tests/test_cartesian_precision.py index 855150f8..98887f0c 100644 --- a/tests/test_cartesian_precision.py +++ b/tests/test_cartesian_precision.py @@ -4,6 +4,8 @@ """Cartesian cutoff certificates from an independent decimal oracle.""" +from fractions import Fraction + import numpy as np import pytest @@ -11,7 +13,12 @@ from solprop.grid import EnvelopeGrid, FieldGrid, FreeGrid from solprop.linear import Frame from solprop.pulses import GaussianPulse, SpatialPulse -from solprop.spatial import SpatialDispersion, SpatialMedium +from solprop.spatial import ( + SpatialDispersion, + SpatialMedium, + _exact_squares, + _transverse_residual, +) from solprop.translate.free import free_operators from solprop.translate.spatial import separable_values from spatial_records import record_metric @@ -135,3 +142,35 @@ def test_cartesian_launch_uses_the_corrected_exact_linear_operator(): rtol=0, atol=(8 * eps + phase_error + normalisation) * float(np.max(np.abs(start))), ) + + +@pytest.mark.parametrize( + ("half_x", "n_x", "half_y", "n_y"), + [(3e-7, 9, 4e-7, 12), (2.3e-4, 64, 7.1e-5, 48), (1.0, 2, 1e-3, 33)], +) +def test_the_transverse_residual_is_the_rational_ones_bit_for_bit( + half_x, n_x, half_y, n_y +): + # The rational sum is the definition; `math.fsum` of exact squares is its fast form. + space = FreeGrid(half_x, n_x, half_y, n_y, collar_fraction=0) + x_squared = [Fraction(float(value)) ** 2 for value in space.k_x] + y_squared = [Fraction(float(value)) ** 2 for value in space.k_y] + expected = np.array( + [ + float(x + y - Fraction(float(high))) + for (x, y), high in zip( + ((x, y) for y in y_squared for x in x_squared), + space.kperp_sq, + strict=True, + ) + ] + ) + assert _transverse_residual(space).tobytes() == expected.tobytes() + + +def test_the_exact_squares_are_exact(): + rng = np.random.default_rng(61) + values = rng.standard_normal(4096) * 10.0 ** rng.integers(-100, 100, 4096) + high, low = _exact_squares(values) + for value, a, b in zip(values.tolist(), high.tolist(), low.tolist(), strict=True): + assert Fraction(a) + Fraction(b) == Fraction(value) ** 2 From 98e450fa9b703c5cf9a221f578c0c47a9069fe4f Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 15:03:19 +0000 Subject: [PATCH 3/8] Build a large launch in threads, with its admittance formed once The passes over a beam's elements ran on one thread, and the admittance ran twice. The sum that scales the field stays one sum in one order, so every bit of the launch is as it was, which a test holds across thread counts. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- src/solprop/boundaries.py | 25 +++++--- src/solprop/pulses.py | 127 ++++++++++++++++++++++++++++++-------- tests/test_pulses.py | 29 ++++++++- 3 files changed, 147 insertions(+), 34 deletions(-) diff --git a/src/solprop/boundaries.py b/src/solprop/boundaries.py index f542430b..b0211dbc 100644 --- a/src/solprop/boundaries.py +++ b/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 @@ -1086,15 +1087,25 @@ 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.""" + 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. diff --git a/src/solprop/pulses.py b/src/solprop/pulses.py index 85261e5a..33847785 100644 --- a/src/solprop/pulses.py +++ b/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,7 +1138,7 @@ def _hard_masked( Returns ------- numpy.ndarray - The masked beam. + The masked beam: `field`, masked in place. Raises ------ @@ -1147,9 +1146,17 @@ def _hard_masked( 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.""" + 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 +1167,62 @@ 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.""" + 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.""" + 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.""" + total[rows] += part[rows] + + parallel_rows(total.shape[0], total.nbytes, work) def _members(pulse: Pulse) -> list[Pulse]: @@ -1243,15 +1305,20 @@ 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.""" + 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 +1342,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 +1354,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 +1369,17 @@ 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.""" + 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"]: diff --git a/tests/test_pulses.py b/tests/test_pulses.py index 02815215..e01805e2 100644 --- a/tests/test_pulses.py +++ b/tests/test_pulses.py @@ -15,7 +15,7 @@ from hypothesis import strategies as st from numpy.testing import assert_allclose -from solprop import fourier +from solprop import _process, fourier from solprop.constants import C from solprop.exceptions import PulseError from solprop.grid import EnvelopeGrid, FieldGrid, TimeGrid @@ -1292,3 +1292,30 @@ def test_energy_and_arrival_hold_for_any_pulse(fwhm, energy, gdd, t0, kind): # a chirped pulse is a ratio of two sums with cancellation, and the strategies allow # a stretch factor of about seven, which brings the tails towards the window edge. assert_allclose(centroid(spectrum, grid), t0, rtol=0.0, atol=0.1 * grid.dt) + + +@pytest.mark.parametrize("workers", [2, 7]) +@pytest.mark.parametrize("jones", [(1,), (1, 0.4j)], ids=["one", "two"]) +def test_a_large_launch_is_the_same_on_any_number_of_threads( + workers, jones, monkeypatch +): + # The passes over a large launch take threads, one block of rows each. The sum that + # scales the field stays one sum in one order, so no bit may depend on the blocks. + from solprop.grid import FreeGrid # noqa: PLC0415 + from solprop.pulses import SpatialPulse # noqa: PLC0415 + from solprop.spatial import SpatialMedium # noqa: PLC0415 + + grid = EnvelopeGrid.from_limits( + 1e-4, 800e-9, (500e-9, 1.6e-6), 300e-15, thg=False, fft_size="smooth" + ) + space = FreeGrid(2e-4, 32, 3e-4, 48, collar_fraction=0) + temporal = GaussianPulse(wavelength=800e-9, fwhm=30e-15, energy=3e-9) + medium = SpatialMedium("Ar", 1.0) + pulse = SpatialPulse(temporal, 60e-6, jones=jones) + + serial = pulse.build(grid, space, medium=medium) + assert serial.nbytes < _process.PARALLEL_BYTES + monkeypatch.setattr(_process, "PARALLEL_BYTES", 0) + budget = dataclasses.replace(_process.thread_budget(), threads=workers) + monkeypatch.setattr(_process, "thread_budget", lambda: budget) + assert pulse.build(grid, space, medium=medium).tobytes() == serial.tobytes() From a68158f93f0d175325df30b2af8504b8f9054d90 Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 15:03:42 +0000 Subject: [PATCH 4/8] Give nditer its one operand's flags flat, as the stubs read them Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- src/solprop/output/schema.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/solprop/output/schema.py b/src/solprop/output/schema.py index 01cea9f1..36319a7d 100644 --- a/src/solprop/output/schema.py +++ b/src/solprop/output/schema.py @@ -522,7 +522,7 @@ def of(cls, values: np.ndarray) -> "ArrayDigest": # noqa: UP037 - bound after t for block in np.nditer( array, flags=["external_loop", "buffered"], - op_flags=[["readonly", "contig"]], + op_flags=["readonly", "contig"], op_dtypes=[little], casting="equiv", order="C", From 3cbb921ab5285f2e27d4e7fa923d8e96172d480b Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 15:03:54 +0000 Subject: [PATCH 5/8] Say why nditer's flag is outside the stubs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- src/solprop/output/schema.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/solprop/output/schema.py b/src/solprop/output/schema.py index 36319a7d..8d194ac1 100644 --- a/src/solprop/output/schema.py +++ b/src/solprop/output/schema.py @@ -522,7 +522,8 @@ def of(cls, values: np.ndarray) -> "ArrayDigest": # noqa: UP037 - bound after t for block in np.nditer( array, flags=["external_loop", "buffered"], - op_flags=["readonly", "contig"], + # NumPy's stubs leave "contig" out of the flags they list. + op_flags=["readonly", "contig"], # pyright: ignore[reportArgumentType] op_dtypes=[little], casting="equiv", order="C", From c8725fff67bfff6cf45441929568e39b236368ad Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 15:04:12 +0000 Subject: [PATCH 6/8] Put the stub note on the call the checker reports Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- src/solprop/output/schema.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/solprop/output/schema.py b/src/solprop/output/schema.py index 8d194ac1..2c9357e8 100644 --- a/src/solprop/output/schema.py +++ b/src/solprop/output/schema.py @@ -519,10 +519,10 @@ def of(cls, values: np.ndarray) -> "ArrayDigest": # noqa: UP037 - bound after t elif array.flags.c_contiguous and array.dtype == little: digest.update(memoryview(array).cast("B")) else: - for block in np.nditer( + # 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"], - # NumPy's stubs leave "contig" out of the flags they list. op_flags=["readonly", "contig"], # pyright: ignore[reportArgumentType] op_dtypes=[little], casting="equiv", From cd8d41f58614048724823792712c11a3d34685d4 Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 19:32:39 +0000 Subject: [PATCH 7/8] Record what a free-space point's setup costs now Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- changelog.d/setup-cost.changed.md | 1 + docs/dev/benchmarks.md | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) create mode 100644 changelog.d/setup-cost.changed.md diff --git a/changelog.d/setup-cost.changed.md b/changelog.d/setup-cost.changed.md new file mode 100644 index 00000000..19efbbb9 --- /dev/null +++ b/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) diff --git a/docs/dev/benchmarks.md b/docs/dev/benchmarks.md index 078a8d2a..382c6765 100644 --- a/docs/dev/benchmarks.md +++ b/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. + From 4cc6af0bf2b323633dfafb194e8d4cd59480650d Mon Sep 17 00:00:00 2001 From: John Travers Date: Wed, 4 Nov 2026 20:54:53 +0000 Subject: [PATCH 8/8] Give a launch's block workers their docstrings Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01XGENPtq9Gp7VCh5QmcDdzQ --- src/solprop/boundaries.py | 13 +++++++++- src/solprop/pulses.py | 53 ++++++++++++++++++++++++++++++++++----- src/solprop/spatial.py | 7 ++++++ 3 files changed, 66 insertions(+), 7 deletions(-) diff --git a/src/solprop/boundaries.py b/src/solprop/boundaries.py index 8c447b8e..a67ad33b 100644 --- a/src/solprop/boundaries.py +++ b/src/solprop/boundaries.py @@ -1091,7 +1091,18 @@ def occupancy( 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.""" + """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] ) diff --git a/src/solprop/pulses.py b/src/solprop/pulses.py index 33847785..bae51d87 100644 --- a/src/solprop/pulses.py +++ b/src/solprop/pulses.py @@ -1149,7 +1149,18 @@ def _hard_masked( 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.""" + """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])) @@ -1169,7 +1180,13 @@ def shares(rows: slice) -> tuple[float, float]: raise PulseError(msg) def mask(rows: slice) -> None: - """Zero one block's dead bins.""" + """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) @@ -1200,7 +1217,13 @@ def _separable_field( 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.""" + """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) @@ -1219,7 +1242,13 @@ def _add_rows(total: np.ndarray, part: np.ndarray) -> None: """ def work(rows: slice) -> None: - """Add one block.""" + """Add one block. + + Parameters + ---------- + rows : slice + The block of the first axis. + """ total[rows] += part[rows] parallel_rows(total.shape[0], total.nbytes, work) @@ -1308,7 +1337,13 @@ def _normalising_factor( terms = np.empty((beam.shape[0], beam.shape[2])) def work(rows: slice) -> None: - """Form one block of the flux's terms.""" + """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 @@ -1375,7 +1410,13 @@ def _normalised( factor = _normalising_factor(pulse, temporal, field, plane, grid) def scale(rows: slice) -> None: - """Scale one block.""" + """Scale one block. + + Parameters + ---------- + rows : slice + The block of the first axis. + """ field[rows] *= factor parallel_rows(field.shape[0], field.nbytes, scale) diff --git a/src/solprop/spatial.py b/src/solprop/spatial.py index b849c12c..c50e2a5d 100644 --- a/src/solprop/spatial.py +++ b/src/solprop/spatial.py @@ -921,6 +921,13 @@ def whole_admittance(self) -> Complex[np.ndarray, "n_modes n_omega"]: ) def work(rows: slice) -> None: + """Form one block of the chunks' admittance. + + Parameters + ---------- + rows : slice + The block of chunks. + """ for columns in chunks[rows]: admittance[columns] = self.admittance(columns)